IRKsome 2026.0.0__py3-none-any.whl

irksome/__init__.py

@@ -0,0 +1,122 @@
+from .tableaux.ButcherTableaux import (
+    Alexander,
+    BackwardEuler,
+    GaussLegendre,
+    LobattoIIIA,
+    LobattoIIIC,
+    PareschiRusso,
+    QinZhang,
+    RadauIIA,
+)
+from .tableaux.multistep_tableaux import (
+    BDF,
+    AdamsBashforth,
+    AdamsMoulton,
+)
+from .tableaux.pep_explicit_rk import PEPRK
+from .ufl.deriv import Dt, expand_time_derivatives, check_irksome_import_order
+from .ufl.lag import lag
+
+check_irksome_import_order()
+
+from .tableaux.dirk_imex_tableaux import DIRK_IMEX
+from .tableaux.ars_dirk_imex_tableaux import ARS_DIRK_IMEX
+from .tableaux.sspk_tableau import SSPK_DIRK_IMEX, SSPButcherTableau
+from .tableaux.multistep_tableaux import MultistepTableau
+
+from .constant import MeshConstant
+from .tableaux.wso_dirk_tableaux import WSODIRK
+from .scheme import create_time_quadrature
+from .scheme import ContinuousPetrovGalerkinScheme, DiscontinuousGalerkinScheme
+from .scheme import GalerkinCollocationScheme, DiscontinuousGalerkinCollocationScheme
+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,
+    ClassicNystrom4Tableau,
+)
+from .stage_value import StageValueTimeStepper
+
+from .multistep import MultistepTimeStepper
+
+
+__all__ = [
+    "AdamsBashforth",
+    "AdamsMoulton",
+    "Alexander",
+    "ARS_DIRK_IMEX",
+    "BackwardEuler",
+    "BDF",
+    "create_time_quadrature",
+    "DIRK_IMEX",
+    "DiscontinuousGalerkinCollocationScheme",
+    "Dt",
+    "expand_time_derivatives",
+    "GalerkinCollocationScheme",
+    "GaussLegendre",
+    "lag",
+    "LobattoIIIA",
+    "LobattoIIIC",
+    "MeshConstant",
+    "MultistepTableau",
+    "PareschiRusso",
+    "PEPRK",
+    "QinZhang",
+    "RadauIIA",
+    "SSPButcherTableau",
+    "SSPK_DIRK_IMEX",
+    "WSODIRK",
+    "DIRKTimeStepper",
+    "BoundsConstrainedDirichletBC",
+    "getForm",
+    "RadauIIAIMEXMethod",
+    "DIRKIMEXMethod",
+    "DIRKNystromTimeStepper",
+    "ExplicitNystromTimeStepper",
+    "StageDerivativeNystromTimeStepper",
+    "ClassicNystrom4Tableau",
+    "StageValueTimeStepper",
+    "ContinuousPetrovGalerkinTimeStepper",
+    "DiscontinuousGalerkinTimeStepper",
+    "MultistepTimeStepper",
+]
+
+
+try:
+    import importlib
+
+    importlib.import_module("firedrake")
+    from .labeling import TimeQuadratureLabel
+    from .discontinuous_galerkin_stepper import DiscontinuousGalerkinTimeStepper
+    from .galerkin_stepper import ContinuousPetrovGalerkinTimeStepper
+
+    from .pc import (
+        ClinesBase,
+        ClinesLD,
+        NystromAuxiliaryOperatorPC,
+        RanaBase,
+        RanaDU,
+        RanaLD,
+        IRKAuxiliaryOperatorPC,
+    )
+    from .stepper import TimeStepper
+
+    __all__ += [
+        "TimeQuadratureLabel",
+        "DiscontinuousGalerkinScheme",
+        "ContinuousPetrovGalerkinScheme",
+        "ClinesBase",
+        "ClinesLD",
+        "NystromAuxiliaryOperatorPC",
+        "RanaBase",
+        "RanaDU",
+        "RanaLD",
+        "IRKAuxiliaryOperatorPC",
+        "TimeStepper",
+    ]
+
+except ModuleNotFoundError:
+    pass

irksome/backend.py

@@ -0,0 +1,129 @@
+from typing import Protocol, Any, Sequence
+import ufl
+from importlib import import_module
+import types
+
+
+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,
+        repeated `num_stages` times.
+
+        Args:
+            V: Space for single step
+            num_stages: Number of stages
+
+        Returns:
+            A coefficient in the new function space
+        """
+
+    class Constant:
+        """MeshLess constant class"""
+
+    class MeshConstant:
+        def __init__(self, msh: ufl.Mesh):
+            """Initialize a mesh constant over a domain"""
+
+        def Constant(self, val: float = 0.0):
+            """Generate a constant in the backend language with a specific value"""
+
+    def ConstantOrZero(
+        x: float | complex, MC: MeshConstant | None = None
+    ) -> ufl.core.expr.Expr:
+        """
+        Create a constant with backend class if MeshConstant is not supplied.
+        Create UFL zero if `x` is sufficiently small
+        """
+
+    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_variational_problem(
+        F: ufl.Form,
+        u: ufl.Coefficient,
+        bcs: DirichletBC | Sequence | None = None,
+        **kwargs,
+    ) -> Any:
+        """Create a variational problem in the backend language."""
+
+    def create_variational_solver(
+        problem: Any,
+        solver_parameters: dict | None = None,
+        **kwargs,
+    ):
+        """Create a 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 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."""
+
+    class EquationBCSplit:
+        ...
+
+    class EquationBC:
+        ...
+
+    def create_bounds_constrained_bc(V, g, sub_domain, bounds, solver_parameters=None) -> DirichletBC:
+        ...
+
+
+def get_backend(backend: str | types.ModuleType) -> Backend:
+    """Get backend class from backend name.
+
+    Args:
+        backend: Name of the backend to get
+
+    Returns:
+        Backend class
+    """
+    if isinstance(backend, types.ModuleType):
+        return backend
+    if backend == "firedrake":
+        from .backends import firedrake as fd_backend
+
+        return fd_backend
+    elif backend == "dolfinx":
+        from .backends import dolfinx as dx_backend
+
+        return dx_backend
+    else:
+        return import_module(backend)

irksome/backends/dolfinx.py

@@ -0,0 +1,184 @@
+"""DOLFINx backend for Irksome"""
+
+try:
+    from mpi4py import MPI
+    from petsc4py import PETSc
+    import basix.ufl
+    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_variational_problem(F, u, bcs=None, aP=None, **kwargs):
+        """Create a variational problem."""
+        if len(F.arguments()) == 2:
+            a, L = ufl.system(F)
+            return dolfinx.fem.petsc.LinearProblem(
+                a,
+                L,
+                u,
+                bcs=bcs,
+                petsc_options_prefix="IrkSomeLinearSolver",
+                P=aP,
+                **kwargs,
+            )
+        else:
+            return dolfinx.fem.petsc.NonlinearProblem(
+                F,
+                u,
+                petsc_options_prefix="IrkSomeNonlinearSolver",
+                bcs=bcs,
+                petsc_options=kwargs.get("solver_parameters"),
+            )
+
+    def create_variational_solver(
+        problem: dolfinx.fem.petsc.LinearProblem | dolfinx.fem.petsc.NonlinearProblem,
+        **kwargs,
+    ):
+        """Create a variational solver that uses PETSc SNES or KSP."""
+        solver_parameters = kwargs.get("solver_parameters", {})
+        solver = problem.solver
+        solver_prefix = problem.solver.getOptionsPrefix()
+        opts = PETSc.Options(
+        )
+        opts.prefixPush(solver_prefix)
+        for k, v in solver_parameters.items():
+            opts.setValue(k, v)
+        solver.setFromOptions()
+        opts.prefixPop()
+        # For some strange reason delValue doesn't respect prefixes
+        for k, v in solver_parameters.items():
+            opts.delValue(f"{solver_prefix}{k}")
+        return problem
+
+    def get_function_space(u: ufl.Coefficient) -> ufl.FunctionSpace:
+        return u.ufl_function_space()
+
+    def get_stages(V: dolfinx.fem.FunctionSpace, num_stages: int) -> ufl.Coefficient:
+        """
+        Given a function space for a single time-step, get a duplicate of this space,
+        repeated `num_stages` times.
+
+        Args:
+            V: Space for single step
+            num_stages: Number of stages
+
+        Returns:
+            A coefficient in the new function space
+        """
+        if V.num_sub_spaces == 0:
+            el = basix.ufl.mixed_element([V.ufl_element()] * num_stages)
+        else:
+            el = basix.ufl.mixed_element(V.ufl_element().sub_elements * num_stages)
+        Vbig = dolfinx.fem.functionspace(V.mesh, el)
+        return dolfinx.fem.Function(Vbig)
+
+    class MeshConstant(object):
+        def __init__(self, msh):
+            self.msh = msh
+            try:
+                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)
+            v.value = val
+            return v
+
+    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.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.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}")
+
+    derivative = ufl.derivative
+    TrialFunction = ufl.TrialFunction
+    Function = dolfinx.fem.Function
+    TestFunction = ufl.TestFunction
+
+    class Constant(ufl.constantvalue.ScalarValue):
+        # NOTE: If dolfinx ever get's meshless constants we should change this
+        def assign(self, value):
+            self._value = value
+
+    class EquationBCSplit:
+        def __init__(self, *args, **kwargs):
+            raise NotImplementedError("DOLFINx does not support EquationBCSplit")
+
+    class EquationBC:
+        def __init__(self, *args, **kwargs):
+            raise NotImplementedError("DOLFINx does not support EquationBC")
+
+    def invalidate_jacobian(solver: dolfinx.fem.petsc.LinearProblem):
+        """Invalidate the Jacobian matrix in the backend language."""
+        raise RuntimeError("DOLFINx does not support Jacobian invalidation")
+
+    def create_bounds_constrained_bc(V, g, sub_domain, bounds, solver_parameters=None):
+        raise NotImplementedError("Bounds-constrained BCs are not implemented for DOLFINx")
+
+except ModuleNotFoundError:
+    pass

irksome/backends/firedrake.py

@@ -0,0 +1,139 @@
+"""Firedrake backend for Irksome"""
+
+import firedrake
+import ufl
+import typing
+from functools import reduce
+from operator import mul
+
+
+assemble = firedrake.assemble
+derivative = firedrake.derivative
+norm = firedrake.norm
+Function = firedrake.Function
+TestFunction = firedrake.TestFunction
+TrialFunction = firedrake.TrialFunction
+DirichletBC = firedrake.DirichletBC
+EquationBC = firedrake.bcs.EquationBC
+EquationBCSplit = firedrake.bcs.EquationBCSplit
+Constant = firedrake.Constant
+
+
+def get_stage_space(V: ufl.FunctionSpace, num_stages: int) -> ufl.FunctionSpace:
+    # NOTE: Should use firedrake.MixedFunctionSpace, but there is a bug for vector spaces.
+    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 get_function_space(u: ufl.Coefficient) -> firedrake.FunctionSpace:
+    return u.function_space()
+
+
+def get_stages(V: firedrake.FunctionSpace, num_stages: int) -> firedrake.Function:
+    """
+    Given a function space for a single time-step, get a duplicate of this space,
+    repeated `num_stages` times.
+
+    Args:
+        V: Space for single step
+        num_stages: Number of stages
+
+    Returns:
+        A coefficient in the new function space
+    """
+    Vbig = get_stage_space(V, num_stages)
+    return firedrake.Function(Vbig)
+
+
+class MeshConstant(object):
+    def __init__(self, msh: ufl.Mesh):
+        self.msh = ufl.domain.as_domain(msh)
+        self.V = firedrake.FunctionSpace(self.msh, "R", 0)
+
+    def Constant(self, val=0.0) -> ufl.Coefficient:
+        return firedrake.Function(self.V).assign(val)
+
+
+def get_mesh_constant(MC: MeshConstant | None):
+    return MC.Constant if MC else firedrake.Constant
+
+
+def create_variational_problem(F, u, bcs=None, J=None, Jp=None, **kwargs):
+    if len(F.arguments()) == 2:
+        a, L = ufl.system(F)
+        kwargs.pop("is_linear", None)
+        problem = firedrake.LinearVariationalProblem(a, L, u, bcs=bcs, aP=Jp, **kwargs)
+    else:
+        constant_jacobian = kwargs.pop("constant_jacobian", False)
+        problem = firedrake.NonlinearVariationalProblem(F, u, bcs=bcs, J=J, Jp=Jp, **kwargs)
+        if constant_jacobian:
+            problem._constant_jacobian = constant_jacobian
+    return problem
+
+
+def create_variational_solver(problem, **kwargs):
+    if isinstance(problem, firedrake.LinearVariationalProblem):
+        return firedrake.LinearVariationalSolver(problem, **kwargs)
+    else:
+        return firedrake.NonlinearVariationalSolver(problem, **kwargs)
+
+
+def invalidate_jacobian(solver):
+    return firedrake.LinearVariationalSolver.invalidate_jacobian(solver)
+
+
+class BoundsConstrainedDirichletBC(firedrake.DirichletBC):
+    def __init__(
+        self,
+        V,
+        g,
+        sub_domain,
+        bounds,
+        solver_parameters=None,
+    ):
+        self.g = g
+        self.solver_parameters = solver_parameters
+        self.bounds = bounds
+        self.gnew = Function(V)
+
+        F = ufl.inner(self.gnew - g, TestFunction(V)) * ufl.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",
+            }
+        problem = create_variational_problem(F, self.gnew)
+        self.solver = create_variational_solver(
+            problem, solver_parameters=solver_parameters
+        )
+        super().__init__(V, g, sub_domain)
+
+    @property
+    def function_arg(self):
+        """The value of this boundary condition."""
+        self.solver.solve(bounds=self.bounds)
+        return self.gnew
+
+    @function_arg.setter
+    def function_arg(self, g):
+        """Set the value of this boundary condition."""
+        self.solver.solve(bounds=self.bounds)
+        return self.gnew
+
+    def reconstruct(self, V=None, g=None, sub_domain=None):
+        V = V or self.function_space()
+        g = g or self.g
+        sub_domain = sub_domain or self.sub_domain
+        return type(self)(V, g, sub_domain, self.bounds, self.solver_parameters)
+
+
+def create_bounds_constrained_bc(V, g, sub_domain, bounds, solver_parameters=None):
+    return BoundsConstrainedDirichletBC(V, g, sub_domain, bounds, solver_parameters=solver_parameters)