make getform work (without DirichletBCs) for DOLFINx by jorgensd · Pull Request #200 · firedrakeproject/Irksome

jorgensd · 2026-03-09T20:28:24Z

Minimal example:

from mpi4py import MPI
import dolfinx
import ufl
from irksome import GaussLegendre, getForm, Dt, MeshConstant
from irksome.tools import get_stage_space
from ufl import pi, atan, div, grad, inner, dx

butcher_tableau = GaussLegendre(2)
N = 64

x0 = 0.0
x1 = 10.0
y0 = 0.0
y1 = 10.0

msh = dolfinx.mesh.create_rectangle(MPI.COMM_WORLD, [[x0,y0],[x1,y1]],[N, N])
V = dolfinx.fem.functionspace(msh, ("Lagrange", 1))
x, y = ufl.SpatialCoordinate(msh)

MC = MeshConstant(msh, backend="dolfinx")
dt = MC.Constant(10 / N)
t = MC.Constant(0.0)

Constant = lambda val: dolfinx.fem.Constant(msh, val)
S = Constant(2.0)
C = Constant(1000.0)


B = (x-Constant(x0))*(x-Constant(x1))*(y-Constant(y0))*(y-Constant(y1))/C
R = (x * x + y * y) ** 0.5
uexact = B * atan(t)*(pi / 2.0 - atan(S * (R - t)))

rhs = Dt(uexact) - div(grad(uexact))

u = dolfinx.fem.Function(V)
u.interpolate(dolfinx.fem.Expression(uexact, V.element.interpolation_points))

v = ufl.TestFunction(V)
F = inner(Dt(u), v)*dx + inner(grad(u), grad(v))*dx - inner(rhs, v)*dx

bc = []
# bc = DirichletBC(V, 0, "on_boundary")

# Get the function space for the stage-coupled problem and a function to hold the stages we're computing::

Vbig = get_stage_space(V, butcher_tableau.num_stages, backend="dolfinx")
k = dolfinx.fem.Function(Vbig)

# Get the variational form and bcs for the stage-coupled variational problem::

Fnew, bcnew = getForm(F, butcher_tableau, t, dt, u, k, bcs=bc, backend="dolfinx")

…t adding dummy classes for DOLFINx).

pbrubeck · 2026-05-04T15:09:17Z

 from .scheme import ContinuousPetrovGalerkinScheme, DiscontinuousGalerkinScheme
 from .scheme import GalerkinCollocationScheme

+from .stage_derivative import getForm


This should not be made public

Co-authored-by: Copilot <copilot@github.com>

pbrubeck · 2026-05-12T10:24:51Z

+
+
+def get_stage_space(V: ufl.FunctionSpace, num_stages: int) -> ufl.FunctionSpace:
+    return reduce(mul, (V for _ in range(num_stages)))


Suggested change

return reduce(mul, (V for _ in range(num_stages)))

return firedrake.MixedFunctionSpace(tuple(V) * num_stages)

pbrubeck · 2026-05-12T10:25:01Z

+from operator import mul
+from functools import reduce
+


Suggested change

from operator import mul

from functools import reduce

pbrubeck · 2026-05-12T10:33:03Z

-    :arg dt: a :class:`firedrake.Constant` or :class:`firedrake.Function`
-        on the Real space over the same mesh as ``u0``.  This serves as
-        a variable referring to the current time step size.
+    :arg t: a :class:`Function` or :class:`Constant` on the Real space over the same mesh as


Suggested change

:arg t: a :class:`Function` or :class:`Constant` on the Real space over the same mesh as

:arg t: a :class:`Constant` or :class:`Function` on the Real space over the same mesh as

pbrubeck · 2026-05-12T10:33:16Z

-        a variable referring to the current time step size.
+    :arg t: a :class:`Function` or :class:`Constant` on the Real space over the same mesh as
+         `u0`.  This serves as a variable referring to the current time.
+    :arg dt: a :class:`Function` or :class:`Constant` on the Real space over the same mesh as


Suggested change

:arg dt: a :class:`Function` or :class:`Constant` on the Real space over the same mesh as

:arg dt: a :class:`Constant` or :class:`Function` on the Real space over the same mesh as

pbrubeck · 2026-05-12T10:36:19Z

       - `bcnew`, a list of :class:`firedrake.DirichletBC` or :class:`EquationBC`
         objects to be posed on the stages
    """
+    backend_cls = get_backend(backend)


Should we try to consistenty set a backend variable name throught different files? In base_time_stepper we repeatedly call self._backend, but here we define backend_cls.

pbrubeck · 2026-05-12T10:39:12Z

+    assert V == backend_cls.get_function_space(u0)

-    c = vecconst(butch.c)
+    c = vecconst(butch.c, backend=backend)


Should we implement backend_cls.vecconst? The name of this function could also be improved, especially because this function returns ufl.zero for numeric values equal to zero.

pbrubeck · 2026-05-12T10:40:25Z

-                   dtu: dtusub}
+        repl[i] = {t: t + c[i] * dt, v: v_np[i], u0: usub, dtu: dtusub}

    Fnew = sum(replace(F, repl[i]) for i in range(num_stages))


Not really needed here, but we should be consistent.

Suggested change

Fnew = sum(replace(F, repl[i]) for i in range(num_stages))

Fnew = sum(backend_cls.replace(F, repl[i]) for i in range(num_stages))

maybe it'd be good to define replace = backend_cls.replace earlier in this function, and stop importing replace from .tools

pbrubeck · 2026-05-12T10:43:20Z

+            bcs,
+            self.bc_type,
+            splitting=self.splitting,
+            aux_indices=self.aux_indices,


do we need to pass the backend?

pbrubeck · 2026-05-12T10:44:47Z

I think this PR should not apply changes that are merely (ruff) formatting. This makes it a bit harder to review.

pbrubeck · 2026-05-12T10:48:59Z

+    return firedrake.NonlinearVariationalProblem(F, u, bcs=bcs, **kwargs)
+
+
+def create_nonlinearvariational_solver(


Why are we omitting an underscode here?

Do we really need to say variational?

pbrubeck · 2026-05-12T10:52:18Z

+    try:
+        from firedrake.fml import LabelledForm, Term
+        if isinstance(e, LabelledForm):
+            enew = LabelledForm(*(Term(ufl_replace(term.form, cmapping), term.labels)
+                                for term in e.terms))
+            return enew
+        else:
+            return ufl_replace(e, cmapping)
+
+    except ImportError:


Suggested change

try:

from firedrake.fml import LabelledForm, Term

if isinstance(e, LabelledForm):

enew = LabelledForm(*(Term(ufl_replace(term.form, cmapping), term.labels)

for term in e.terms))

return enew

else:

return ufl_replace(e, cmapping)

except ImportError:

try:

from firedrake.fml import LabelledForm, Term

if isinstance(e, LabelledForm):

return LabelledForm(*(Term(ufl_replace(term.form, cmapping), term.labels)

for term in e.terms))

except ImportError:

pass

finally:

pbrubeck · 2026-05-14T08:18:39Z

+Function = firedrake.Function
+
+DirichletBC = firedrake.DirichletBC
+
+norm = firedrake.norm
+
+assemble = firedrake.assemble
+
+replace = firedrake.replace
+
+derivative = firedrake.derivative
+TestFunction = firedrake.TestFunction
+TrialFunction = firedrake.TrialFunction


Do we need to define varibales? Can these not be just imports?

jorgensd added 2 commits March 9, 2026 21:27

make getform work (without bcs).

9031a1d

Some minor logic fixes in BoundsConstraindDirichletBC

f5cbeea

pbrubeck reviewed Mar 9, 2026

View reviewed changes

Comment thread irksome/stage_derivative.py

pbrubeck reviewed Mar 10, 2026

View reviewed changes

Comment thread irksome/stage_derivative.py Outdated