make getform work (without DirichletBCs) for DOLFINx

irksome/__init__.py

@@ -29,6 +29,8 @@
 from .scheme import ContinuousPetrovGalerkinScheme, DiscontinuousGalerkinScheme
 from .scheme import GalerkinCollocationScheme
 
+from .stage_derivative import getForm
+
 __all__ = [
     "AdamsBashforth",
     "AdamsMoulton",
@@ -42,6 +44,7 @@
     "DiscontinuousGalerkinScheme",
     "Dt",
     "expand_time_derivatives",
+    "getForm",
     "GalerkinCollocationScheme",
     "GaussLegendre",
     "LobattoIIIA",
@@ -65,7 +68,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 +93,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,14 @@ 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 +46,14 @@ 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 create_nonlinearvariational_solver(F: ufl.Form, u: ufl.Coefficient, bcs: DirichletBC | Sequence | None = None, solver_parameters: dict | None = None, **kwargs):
+        """Create a non-linear variational solver that uses PETSc SNES."""
+
+    def get_stage_spaces(V: ufl.FunctionSpace, num_stages: int) -> ufl.FunctionSpace:
+        """Create a stage space with M number of components."""
 
 def get_backend(backend: str) -> Backend:
     """Get backend class from backend name.

irksome/backends/dolfinx.py

@@ -2,8 +2,32 @@
 
 try:
     import basix.ufl
-    import dolfinx
+    import dolfinx.fem.petsc
     import ufl
+    import typing
+    import numpy as np
+
+    TestFunction = ufl.TestFunction
+
+    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_nonlinearvariational_solver(F: ufl.Form, g: ufl.Coefficient, bcs: typing.Sequence | None, solver_parameters: dict):
+        """Create a non-linear variational solver that uses PETSc SNES."""
+        return dolfinx.fem.petsc.NonlinearProblem(F, g, petsc_options_prefix="IrkSomeSolver", bcs=bcs,
+                                                     petsc_options=solver_parameters)
 
     def get_function_space(u: ufl.Coefficient) -> ufl.FunctionSpace:
         return u.ufl_function_space()
@@ -31,11 +55,15 @@ class MeshConstant(object):
         def __init__(self, msh):
             self.msh = msh
             try:
-                import scifem
-            except ModuleNotFoundError:
-                raise RuntimeError("Scifem is required to make mesh-constants")
-
-            self.V = scifem.create_real_functionspace(msh, ())
+                import basix.ufl
+                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 +73,12 @@ 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 Function(dolfinx.fem.Function):
+        pass
+
+    class DirichletBC(dolfinx.fem.DirichletBC):
+        pass
+
+
 except ModuleNotFoundError:
     pass

irksome/backends/firedrake.py

@@ -1,9 +1,32 @@
 """Firedrake backend for Irksome"""
 
+
+from operator import mul
+from functools import reduce
+
 import firedrake
 import ufl
 from ..tools import get_stage_space
+import typing
+
+TestFunction = firedrake.TestFunction
+
+
+def get_stage_space(V: ufl.FunctionSpace, num_stages:int)->ufl.FunctionSpace:
+    return reduce(mul, (V for _ in range(num_stages)))
+
 
+def extract_bcs(bcs: typing.Any)->tuple[typing.Any]:
+    """Return an iterable of boundary conditions on the residual form"""
+    return tuple(bc.extract_form("F") for bc in firedrake.solving._extract_bcs(bcs))
+
+
+def create_nonlinearvariational_solver(F: ufl.Form, u: ufl.Coefficient, bcs: typing.Sequence | None = None, solver_parameters: dict | None = None, **kwargs):
+    """Create a non-linear variational solver that uses PETSc SNES."""
+    problem = firedrake.NonlinearVariationalProblem(F, u, bcs=bcs)
+    return firedrake.NonlinearVariationalSolver(
+            problem, solver_parameters=solver_parameters, **kwargs
+        )
 
 def get_function_space(u: ufl.Coefficient) -> firedrake.FunctionSpace:
     return u.function_space()
@@ -36,3 +59,10 @@ def Constant(self, val=0.0) -> ufl.Coefficient:
 
 def get_mesh_constant(MC: MeshConstant | None):
     return MC.Constant if MC else firedrake.Constant
+
+class Function(firedrake.Function):
+    pass
+
+
+class DirichletBC(firedrake.DirichletBC):
+    pass

irksome/bcs.py

@@ -1,23 +1,14 @@
-from firedrake.solving import _extract_bcs
 from firedrake import (
     DirichletBC,
     Function,
-    TestFunction,
-    NonlinearVariationalProblem,
-    NonlinearVariationalSolver,
     replace,
     inner,
     dx,
 )
-
+from .backend import get_backend
 from ufl import as_ufl
 
 
-def extract_bcs(bcs):
-    """Return an iterable of boundary conditions on the residual form"""
-    return tuple(bc.extract_form("F") for bc in _extract_bcs(bcs))
-
-
 def get_sub(u, indices):
     for i in indices:
         if i is not None:
@@ -66,25 +57,24 @@ def EmbeddedBCData(bc, butcher_tableau, t, dt, u0, stages):
 class BoundsConstrainedDirichletBC(DirichletBC):
     """A DirichletBC with bounds-constrained data."""
 
-    def __init__(self, V, g, sub_domain, bounds, solver_parameters=None):
+    def __init__(self, V, g, sub_domain, bounds, solver_parameters=None, backend:str="firedrake"):
+
+        self.g = g
+        self.solver_parameters = solver_parameters
+        self.bounds = bounds
+        self.gnew = Function(V)
+        backend_cls = get_backend(backend)
+        F = inner(self.gnew - g, backend_cls.TestFunction(V)) * dx
+
         if solver_parameters is None:
             solver_parameters = {
                 "snes_type": "vinewtonrsls",
                 "snes_max_it": 300,
                 "snes_atol": 1.0e-8,
                 "ksp_type": "preonly",
                 "mat_type": "aij",
-            }
-        self.g = g
-        self.solver_parameters = solver_parameters
-        self.bounds = bounds
-
-        self.gnew = Function(V)
-        F = inner(self.gnew - g, TestFunction(V)) * dx
-        problem = NonlinearVariationalProblem(F, self.gnew)
-        self.solver = NonlinearVariationalSolver(
-            problem, solver_parameters=self.solver_parameters
-        )
+            }       
+        self.solver = backend_cls.create_nonlinearvariational_solver(F, self.gnew, solver_parameters=solver_parameters)
         super().__init__(V, g, sub_domain)
 
     @property

irksome/galerkin_stepper.py

@@ -5,7 +5,7 @@
 from ufl import as_ufl, as_tensor
 
 from .base_time_stepper import StageCoupledTimeStepper
-from .bcs import bc2space, extract_bcs, stage2spaces4bc
+from .bcs import bc2space, stage2spaces4bc
 from .ufl.deriv import TimeDerivative, expand_time_derivatives
 from .ufl.estimate_degrees import TimeDegreeEstimator, get_degree_mapping
 from .labeling import split_quadrature, as_form
@@ -14,7 +14,7 @@
 from .constant import vecconst
 from .discontinuous_galerkin_stepper import getElement as getTestElement
 from .integrated_lagrange import IntegratedLagrange
-
+from .backends.firedrake import extract_bcs
 from .tableaux.ButcherTableaux import CollocationButcherTableau
 from .stage_derivative import getForm
 from .stage_value import getFormStage