make getform work (without DirichletBCs) for DOLFINx

demo_dolfinx.py

@@ -0,0 +1,58 @@
+from mpi4py import MPI
+import dolfinx
+import ufl
+from irksome import GaussLegendre, Dt, MeshConstant
+from irksome.stage_derivative import getForm
+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")

irksome/__init__.py

@@ -29,6 +29,7 @@
 from .scheme import ContinuousPetrovGalerkinScheme, DiscontinuousGalerkinScheme
 from .scheme import GalerkinCollocationScheme
 
+
 __all__ = [
     "AdamsBashforth",
     "AdamsMoulton",
@@ -65,7 +66,6 @@
     from .bcs import BoundsConstrainedDirichletBC
     from .dirk_stepper import DIRKTimeStepper
     from .imex import RadauIIAIMEXMethod, DIRKIMEXMethod
-    from .stage_derivative import getForm
     from .nystrom_dirk_stepper import DIRKNystromTimeStepper, ExplicitNystromTimeStepper
     from .nystrom_stepper import (
         StageDerivativeNystromTimeStepper,
@@ -91,7 +91,6 @@
     __all__ += [
         "DIRKTimeStepper",
         "BoundsConstrainedDirichletBC",
-        "getForm",
         "RadauIIAIMEXMethod",
         "DIRKIMEXMethod",
         "DIRKNystromTimeStepper",

irksome/backend.py

@@ -1,4 +1,4 @@
-from typing import Protocol
+from typing import Protocol, Any, Sequence
 import ufl
 from importlib import import_module
 
@@ -7,6 +7,13 @@ class Backend(Protocol):
     def get_function_space(self, V: ufl.Coefficient) -> ufl.FunctionSpace:
         """Get a function space from the backend"""
 
+    def extract_bcs(bcs: Any) -> tuple[Any]:
+        """Extract boundary conditions"""
+
+    class Function: ...
+
+    class DirichletBC: ...
+
     def get_stages(self, V: ufl.FunctionSpace, num_stages: int) -> ufl.Coefficient:
         """
         Given a function space for a single time-step, get a duplicate of this space,
@@ -38,6 +45,71 @@ def ConstantOrZero(
     def get_mesh_constant(MC: MeshConstant | None) -> ufl.core.expr.Expr:
         """Get a backend class to construct a mesh constant from"""
 
+    def TestFunction(space: ufl.FunctionSpace, part: int | None = None) -> ufl.Argument:
+        """Return a test-function that can be used by forms in the backend."""
+
+    def TrialFunction(
+        space: ufl.FunctionSpace, part: int | None = None
+    ) -> ufl.Argument:
+        """Return a trial-function that can be used by forms in the backend."""
+
+    def create_nonlinearvariational_problem(
+        F: ufl.Form,
+        u: ufl.Coefficient,
+        bcs: DirichletBC | Sequence | None = None,
+        **kwargs,
+    ) -> Any:
+        """Create a non-linear variational problem in the backend language."""
+
+    def create_nonlinearvariational_solver(
+        problem: Any,
+        solver_parameters: dict | None = None,
+        **kwargs,
+    ):
+        """Create a non-linear variational solver in the backend language."""
+
+    def create_linearvariational_problem(
+        a: ufl.Form,
+        L: ufl.Form,
+        u: ufl.Coefficient | Sequence[ufl.Coefficient],
+        bcs: DirichletBC | Sequence | None = None,
+        aP: ufl.Form | None = None,
+        **kwargs,
+    ) -> Any:
+        """Create a linear variational problem in the backend language."""
+
+    def create_linearvariational_solver(
+        problem: Any,
+        solver_parameters: dict | None = None,
+        **kwargs,
+    ):
+        """Create a linear variational solver in the backend language."""
+
+    def get_stage_spaces(V: ufl.FunctionSpace, num_stages: int) -> ufl.FunctionSpace:
+        """Create a stage space with M number of components."""
+
+    def norm(
+        v: ufl.core.expr.Expr, norm_type: str = "L2", mesh: ufl.Mesh | None = None
+    ) -> float:
+        """Compute the norm of a function in the backend language."""
+
+    def assemble(expr: ufl.core.expr.Expr) -> Any:
+        """Assemble a UFL expression in the backend language."""
+
+    def replace(expr: ufl.core.expr.Expr, mapping: dict) -> ufl.core.expr.Expr:
+        """Replace sub-expressions in a UFL expression with other expressions."""
+
+    def derivative(
+        form: ufl.Form,
+        u: ufl.Coefficient,
+        du: ufl.Argument | None = None,
+        coefficient_derivatives: dict | None = None,
+    ) -> ufl.Form:
+        """Compute the derivative of a form with respect to a coefficient in the backend language."""
+
+    def invalidate_jacobian(solver: Any):
+        """Invalidate the Jacobian matrix in the backend language."""
+
 
 def get_backend(backend: str) -> Backend:
     """Get backend class from backend name.

irksome/backends/dolfinx.py

@@ -1,9 +1,76 @@
 """DOLFINx backend for Irksome"""
 
 try:
+    from mpi4py import MPI
     import basix.ufl
-    import dolfinx
+    import dolfinx.fem.petsc
     import ufl
+    import typing
+    import numpy as np
+
+    def get_stage_space(V: ufl.FunctionSpace, num_stages: int) -> ufl.FunctionSpace:
+        if num_stages == 1:
+            me = V.ufl_elemet()
+        else:
+            el = V.ufl_element()
+            if el.num_sub_elements > 0:
+                me = basix.ufl.mixed_element(
+                    np.tile(el.sub_elements, num_stages).tolist()
+                )
+            else:
+                me = basix.ufl.blocked_element(el, shape=(num_stages,))
+        return dolfinx.fem.functionspace(V.mesh, me)
+
+    def extract_bcs(bcs: typing.Any) -> tuple[typing.Any]:
+        """Extract boundary conditions"""
+        return bcs
+
+    def create_linearvariational_problem(
+        a: ufl.Form,
+        L: ufl.Form,
+        u: ufl.Coefficient | typing.Sequence[ufl.Coefficient],
+        bcs: typing.Sequence | None = None,
+        aP: ufl.Form | None = None,
+        **kwargs,
+    ) -> dolfinx.fem.petsc.LinearProblem:
+        return dolfinx.fem.petsc.LinearProblem(
+            a,
+            L,
+            u,
+            bcs=bcs,
+            petsc_options_prefix="IrkSomeLinearSolver",
+            P=aP,
+            **kwargs,
+        )
+
+    def create_linearvariational_solver(
+        problem: dolfinx.fem.petsc.LinearProblem,
+        solver_parameters: dict | None = None,
+        **kwargs,
+    ):
+        """Create a linear variational solver that uses PETSc KSP."""
+        return problem
+
+    def create_nonlinearvariational_problem(
+        F: ufl.Form,
+        g: ufl.Coefficient,
+        bcs: typing.Sequence | None = None,
+        solver_parameters: dict | None = None,
+    ) -> dolfinx.fem.petsc.NonlinearProblem:
+        return dolfinx.fem.petsc.NonlinearProblem(
+            F,
+            g,
+            petsc_options_prefix="IrkSomeNonlinearSolver",
+            bcs=bcs,
+            petsc_options=solver_parameters,
+        )
+
+    def create_nonlinearvariational_solver(
+        problem: dolfinx.fem.petsc.NonlinearProblem,
+        solver_parameters: dict | None = None,
+    ):
+        """Create a non-linear variational solver that uses PETSc SNES."""
+        return problem
 
     def get_function_space(u: ufl.Coefficient) -> ufl.FunctionSpace:
         return u.ufl_function_space()
@@ -31,11 +98,20 @@ class MeshConstant(object):
         def __init__(self, msh):
             self.msh = msh
             try:
-                import scifem
-            except ModuleNotFoundError:
-                raise RuntimeError("Scifem is required to make mesh-constants")
+                import basix.ufl
 
-            self.V = scifem.create_real_functionspace(msh, ())
+                r_el = basix.ufl.real_element(
+                    msh.basix_cell(), value_shape=(), dtype=dolfinx.default_scalar_type
+                )
+                self.V = dolfinx.fem.functionspace(msh, r_el)
+            except TypeError:
+                try:
+                    import scifem
+                except ModuleNotFoundError:
+                    raise RuntimeError(
+                        "DOLFINx with real element support or Scifem is required to make mesh-constants"
+                    )
+                self.V = scifem.create_real_functionspace(msh, ())
 
         def Constant(self, val=0.0) -> ufl.Coefficient:
             v = dolfinx.fem.Function(self.V)
@@ -45,5 +121,52 @@ def Constant(self, val=0.0) -> ufl.Coefficient:
     def get_mesh_constant(MC: MeshConstant | None) -> ufl.core.expr.Expr:
         return MC.Constant if MC is not None else ufl.constantvalue.ComplexValue
 
+    class DirichletBC(dolfinx.fem.DirichletBC):
+        pass
+
+    def norm(
+        v: ufl.core.Expr, norm_type: str = "L2", mesh: ufl.Mesh | None = None
+    ) -> float:
+        """Compute the norm of a function in the backend language."""
+        if mesh is not None:
+            dx = ufl.Mesure("dx", domain=mesh)
+        else:
+            dx = ufl.dx
+        p = 2
+        if norm_type.startswith("L"):
+            p = int(norm_type[1:])
+            if p < 1:
+                raise ValueError(f"Invalid norm type {norm_type}")
+            expr = ufl.inner(v, v) ** (p / 2)
+            form = dolfinx.fem.form(expr * dx)
+        else:
+            raise NotImplementedError(f"Norm type {norm_type} not implemented")
+        norm_loc = dolfinx.fem.assemble_scalar(form)
+        return form.mesh.comm.Allreduce(MPI.IN_PLACE, norm_loc, op=MPI.SUM) ** (1 / p)
+
+    def assemble(expr: ufl.core.Expr | float):
+        """Assemble a UFL expression in the backend language."""
+        if isinstance(expr, float):
+            return float
+        else:
+            form = dolfinx.fem.form(expr)
+            if form.rank == 0:
+                return dolfinx.fem.assemble_scalar(form)
+            elif form.rank == 1:
+                return dolfinx.fem.assemble_vector(form)
+            elif form.rank == 2:
+                return dolfinx.fem.assemble_matrix(form)
+            else:
+                raise ValueError(f"Cannot assemble form of rank {form.rank}")
+
+    replace = ufl.replace
+    derivative = ufl.derivative
+    TrialFunction = ufl.TrialFunction
+    Function = dolfinx.fem.Function
+    TestFunction = ufl.TestFunction
+
+    def invalidate_jacobian(solver: dolfinx.fem.petsc.LinearProblem):
+        """Invalidate the Jacobian matrix in the backend language."""
+        raise RuntimeError("DOLFINx does not support Jacobian invalidation")
 except ModuleNotFoundError:
     pass