17557791029

Committed 08 Sep 2025 04:38PM UTC coverage: 76.363% (+0.4%) from 75.917%

Build # 17557791029

Build Type

Pull #401

github

Committed by

web-flow

Commit Message

Merge branch 'main' into schnellerhase/remove-type-system

Pull Request Pull Request #401: Removal of custom type system

Run Details

495 of 534 new or added lines in 42 files covered. (92.7%)

6 existing lines in 2 files now uncovered.

9133 of 11960 relevant lines covered (76.36%)

0.76 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

77.29

/ufl/algorithms/apply_derivatives.py

"""Apply derivatives algorithm which computes the derivatives of a form of expression."""

# Copyright (C) 2008-2016 Martin Sandve Alnæs
#
# This file is part of UFL (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later

import warnings
from functools import singledispatchmethod
from math import pi

import numpy as np

from ufl.action import Action
from ufl.algorithms.analysis import extract_arguments, extract_coefficients
from ufl.algorithms.map_integrands import map_integrands
from ufl.algorithms.replace_derivative_nodes import replace_derivative_nodes
from ufl.argument import Argument, BaseArgument, Coargument
from ufl.averaging import CellAvg, FacetAvg
from ufl.checks import is_cellwise_constant
from ufl.classes import (
    Abs,
    CellCoordinate,
    Coefficient,
    Cofunction,
    ComponentTensor,
    Conj,
    Constant,
    ConstantValue,
    Division,
    Expr,
    ExprList,
    ExprMapping,
    FacetNormal,
    FloatValue,
    FormArgument,
    GeometricQuantity,
    Grad,
    Identity,
    Imag,
    Indexed,
    IndexSum,
    Jacobian,
    JacobianDeterminant,
    JacobianInverse,
    Label,
    ListTensor,
    Power,
    Product,
    Real,
    ReferenceGrad,
    ReferenceValue,
    SpatialCoordinate,
    Sum,
    Variable,
    Zero,
)
from ufl.conditional import BinaryCondition, Conditional, NotCondition
from ufl.constantvalue import is_true_ufl_scalar, is_ufl_scalar
from ufl.core.base_form_operator import BaseFormOperator
from ufl.core.expr import ufl_err_str
from ufl.core.external_operator import ExternalOperator
from ufl.core.interpolate import Interpolate
from ufl.core.multiindex import FixedIndex, MultiIndex, indices
from ufl.core.terminal import Terminal
from ufl.corealg.dag_traverser import DAGTraverser
from ufl.differentiation import (
    BaseFormCoordinateDerivative,
    BaseFormOperatorDerivative,
    CoefficientDerivative,
    CoordinateDerivative,
    Derivative,
    VariableDerivative,
)
from ufl.domain import MeshSequence, extract_unique_domain
from ufl.form import BaseForm, Form, ZeroBaseForm
from ufl.mathfunctions import (
    Acos,
    Asin,
    Atan,
    Atan2,
    BesselI,
    BesselJ,
    BesselK,
    BesselY,
    Cos,
    Cosh,
    Erf,
    Exp,
    Ln,
    MathFunction,
    Sin,
    Sinh,
    Sqrt,
    Tan,
    Tanh,
)
from ufl.matrix import Matrix
from ufl.operators import (
    MaxValue,
    MinValue,
    bessel_I,
    bessel_J,
    bessel_K,
    bessel_Y,
    cell_avg,
    conditional,
    cos,
    cosh,
    exp,
    facet_avg,
    ln,
    sign,
    sin,
    sinh,
    sqrt,
)
from ufl.pullback import CustomPullback, PhysicalPullback
from ufl.restriction import Restricted
from ufl.tensors import as_scalar, as_scalars, as_tensor, unit_indexed_tensor, unwrap_list_tensor

# TODO: Add more rulesets?
# - DivRuleset
# - CurlRuleset
# - ReferenceGradRuleset
# - ReferenceDivRuleset


def flatten_domain_element(domain, element):
    """Return the flattened (domain, element) pairs for mixed domain problems.

    Args:
        domain: `Mesh` or `MeshSequence`.
        element: `FiniteElement`.

    Returns:
        Nested tuples of (domain, element) pairs; just ((domain, element),)
        if domain is a `Mesh` (and not a `MeshSequence`).

    """
    if not isinstance(domain, MeshSequence):
        return ((domain, element),)
    flattened = ()
    assert len(domain) == len(element.sub_elements)
    for d, e in zip(domain, element.sub_elements):
        flattened += flatten_domain_element(d, e)
    return flattened


class GenericDerivativeRuleset(DAGTraverser):
    """A generic derivative."""

    def __init__(
        self,
        var_shape: tuple,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache)
        self._var_shape = var_shape

    def unexpected(self, o):
        """Raise error about unexpected type."""
        raise ValueError(f"Unexpected type {type(o).__name__} in AD rules.")

    def override(self, o):
        """Raise error about overriding."""
        raise ValueError(f"Type {type(o).__name__} must be overridden in specialized AD rule set.")

    # --- Some types just don't have any derivative, this is just to
    # --- make algorithm structure generic

    def non_differentiable_terminal(self, o):
        """Return the non-differentiated object.

        Labels and indices are not differentiable: it's convenient to
        return the non-differentiated object.
        """
        return o

    # --- Helper functions for creating zeros with the right shapes

    def independent_terminal(self, o):
        """A zero with correct shape for terminals independent of diff. variable."""
        return Zero(o.ufl_shape + self._var_shape)

    def independent_operator(self, o):
        """A zero with correct shape and indices for operators independent of diff. variable."""
        return Zero(o.ufl_shape + self._var_shape, o.ufl_free_indices, o.ufl_index_dimensions)

    # --- Error checking for missing handlers and unexpected types

    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    @process.register(Expr)
    def _(self, o: Expr) -> Expr:
        """Raise error."""
        raise ValueError(
            f"Missing differentiation handler for type {type(o).__name__}. "
            "Have you added a new type?"
        )

    @process.register(Derivative)
    def _(self, o: Expr) -> Expr:
        """Raise error."""
        raise ValueError(
            f"Unhandled derivative type {type(o).__name__}, nested differentiation has failed."
        )

    @process.register(Label)
    @process.register(MultiIndex)
    def _(self, o: Expr) -> Expr:
        return self.non_differentiable_terminal(o)

    # --- All derivatives need to define grad and averaging

    @process.register(Grad)
    @process.register(CellAvg)
    @process.register(FacetAvg)
    def _(self, o: Expr) -> Expr:
        return self.override(o)

    # --- Default rules for terminals

    # Literals are by definition independent of any differentiation variable
    @process.register(ConstantValue)
    # Constants are independent of any differentiation
    @process.register(Constant)
    def _(self, o: Expr) -> Expr:
        return self.independent_terminal(o)

    # Zero may have free indices
    @process.register(Zero)
    def _(self, o: Expr) -> Expr:
        return self.independent_operator(o)

    # Rules for form arguments must be specified in specialized rule set
    @process.register(FormArgument)
    # Rules for geometric quantities must be specified in specialized rule set
    @process.register(GeometricQuantity)
    def _(self, o: Expr) -> Expr:
        return self.override(o)

    # These types are currently assumed independent, but for non-affine domains
    # this no longer holds and we want to implement rules for them.
    # facet_normal = independent_terminal
    # spatial_coordinate = independent_terminal
    # cell_coordinate = independent_terminal

    # Measures of cell entities, assuming independent although
    # this will not be true for all of these for non-affine domains
    # cell_volume = independent_terminal
    # circumradius = independent_terminal
    # facet_area = independent_terminal
    # cell_surface_area = independent_terminal
    # min_cell_edge_length = independent_terminal
    # max_cell_edge_length = independent_terminal
    # min_facet_edge_length = independent_terminal
    # max_facet_edge_length = independent_terminal

    # Other stuff
    # cell_orientation = independent_terminal
    # quadrature_weigth = independent_terminal

    # These types are currently not expected to show up in AD pass.
    # To make some of these available to the end-user, they need to be
    # implemented here.
    # facet_coordinate = unexpected
    # cell_origin = unexpected
    # facet_origin = unexpected
    # cell_facet_origin = unexpected
    # jacobian = unexpected
    # jacobian_determinant = unexpected
    # jacobian_inverse = unexpected
    # facet_jacobian = unexpected
    # facet_jacobian_determinant = unexpected
    # facet_jacobian_inverse = unexpected
    # cell_facet_jacobian = unexpected
    # cell_facet_jacobian_determinant = unexpected
    # cell_facet_jacobian_inverse = unexpected
    # cell_vertices = unexpected
    # cell_edge_vectors = unexpected
    # facet_edge_vectors = unexpected
    # reference_cell_edge_vectors = unexpected
    # reference_facet_edge_vectors = unexpected
    # cell_normal = unexpected # TODO: Expecting rename
    # cell_normals = unexpected
    # facet_tangents = unexpected
    # cell_tangents = unexpected
    # cell_midpoint = unexpected
    # facet_midpoint = unexpected

    # --- Default rules for operators

    @process.register(Variable)
    def _(self, o: Expr) -> Expr:
        """Differentiate a variable."""
        op, _ = o.ufl_operands
        return self(op)

    # --- Indexing and component handling

    @process.register(Indexed)
    @DAGTraverser.postorder
    def _(self, o: Indexed, Ap: Expr, ii: MultiIndex) -> Indexed:
        """Differentiate an indexed."""
        # Propagate zeros
        if isinstance(Ap, Zero):
            return self.independent_operator(o)
        r = len(Ap.ufl_shape) - len(ii)
        if r:
            kk = indices(r)
            op = Indexed(Ap, MultiIndex(ii.indices() + kk))
            op = as_tensor(op, kk)
        else:
            op = Indexed(Ap, ii)
        return op

    @process.register(ListTensor)
    def _(self, o: Expr) -> Expr:
        """Differentiate a list_tensor."""
        return ListTensor(*(self(op) for op in o.ufl_operands))

    @process.register(ComponentTensor)
    @DAGTraverser.postorder
    def _(self, o: ComponentTensor, Ap: Expr, ii: MultiIndex) -> Expr:
        """Differentiate a component_tensor."""
        if isinstance(Ap, Zero):
            op = self.independent_operator(o)
        else:
            Ap, jj = as_scalar(Ap)
            op = as_tensor(Ap, ii.indices() + jj)
        return op

    # --- Algebra operators

    @process.register(IndexSum)
    @DAGTraverser.postorder
    def _(self, o: Expr, Ap: Expr, ii: Expr) -> Expr:
        """Differentiate an index_sum."""
        return IndexSum(Ap, ii)

    @process.register(Sum)
    @DAGTraverser.postorder
    def _(self, o: Expr, da: Expr, db: Expr) -> Expr:
        """Differentiate a sum."""
        return da + db

    @process.register(Product)
    @DAGTraverser.postorder
    def _(self, o: Expr, da: Expr, db: Expr) -> Expr:
        """Differentiate a product."""
        # Even though arguments to o are scalar, da and db may be
        # tensor valued
        a, b = o.ufl_operands
        (da, db), ii = as_scalars(da, db)
        pa = Product(da, b)
        pb = Product(a, db)
        s = Sum(pa, pb)
        if ii:
            s = as_tensor(s, ii)
        return s

    @process.register(Division)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr, gp: Expr) -> Expr:
        """Differentiate a division."""
        f, g = o.ufl_operands
        if not is_ufl_scalar(f):
            raise ValueError("Not expecting nonscalar nominator")
        if not is_true_ufl_scalar(g):
            raise ValueError("Not expecting nonscalar denominator")
        # do_df = 1/g
        # do_dg = -h/g
        # op = do_df*fp + do_df*gp
        # op = (fp - o*gp) / g
        # Get o and gp as scalars, multiply, then wrap as a tensor
        # again
        so, oi = as_scalar(o)
        sgp, gi = as_scalar(gp)
        o_gp = so * sgp
        if oi or gi:
            o_gp = as_tensor(o_gp, oi + gi)
        op = (fp - o_gp) / g
        return op

    @process.register(Power)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr, gp: Expr) -> Expr:
        """Differentiate a power."""
        f, g = o.ufl_operands
        if not is_true_ufl_scalar(f):
            raise ValueError("Expecting scalar expression f in f**g.")
        if not is_true_ufl_scalar(g):
            raise ValueError("Expecting scalar expression g in f**g.")
        # Derivation of the general case: o = f(x)**g(x)
        # do/df  = g * f**(g-1) = g / f * o
        # do/dg  = ln(f) * f**g = ln(f) * o
        # do/df * df + do/dg * dg = o * (g / f * df + ln(f) * dg)
        if isinstance(gp, Zero):
            # This probably produces better results for the common
            # case of f**constant
            op = fp * g * f ** (g - 1)
        else:
            # Note: This produces expressions like (1/w)*w**5 instead of w**4
            # op = o * (fp * g / f + gp * ln(f)) # This reuses o
            op = f ** (g - 1) * (
                g * fp + f * ln(f) * gp
            )  # This gives better accuracy in dolfin integration test
        # Example: d/dx[x**(x**3)]:
        # f = x
        # g = x**3
        # df = 1
        # dg = 3*x**2
        # op1 = o * (fp * g / f + gp * ln(f))
        #     = x**(x**3)   * (x**3/x + 3*x**2*ln(x))
        # op2 = f**(g-1) * (g*fp + f*ln(f)*gp)
        #     = x**(x**3-1) * (x**3 + x*3*x**2*ln(x))
        return op

    @process.register(Abs)
    @DAGTraverser.postorder
    def _(self, o: Expr, df: Expr) -> Expr:
        """Differentiate an abs."""
        (f,) = o.ufl_operands
        # return conditional(eq(f, 0), 0, Product(sign(f), df)) abs is
        # not complex differentiable, so we workaround the case of a
        # real F in complex mode by defensively casting to real inside
        # the sign.
        return sign(Real(f)) * df

    # --- Complex algebra

    @process.register(Conj)
    @DAGTraverser.postorder
    def _(self, o: Expr, df: Expr) -> Expr:
        """Differentiate a conj."""
        return Conj(df)

    @process.register(Real)
    @DAGTraverser.postorder
    def _(self, o: Expr, df: Expr) -> Expr:
        """Differentiate a real."""
        return Real(df)

    @process.register(Imag)
    @DAGTraverser.postorder
    def _(self, o: Expr, df: Expr) -> Expr:
        """Differentiate a imag."""
        return Imag(df)

    # --- Mathfunctions

    @process.register(MathFunction)
    @DAGTraverser.postorder
    def _(self, o: Expr, df: Expr) -> Expr:
        """Differentiate a math_function."""
        # FIXME: Introduce a UserOperator type instead of this hack
        # and define user derivative() function properly
        if hasattr(o, "derivative"):
            return df * o.derivative()
        else:
            raise ValueError("Unknown math function.")

    @process.register(Sqrt)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a sqrt."""
        return fp / (2 * o)

    @process.register(Exp)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate an exp."""
        return fp * o

    @process.register(Ln)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a ln."""
        (f,) = o.ufl_operands
        if isinstance(f, Zero):
            raise ZeroDivisionError()
        return fp / f

    @process.register(Cos)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a cos."""
        (f,) = o.ufl_operands
        return fp * -sin(f)

    @process.register(Sin)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a sin."""
        (f,) = o.ufl_operands
        return fp * cos(f)

    @process.register(Tan)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a tan."""
        (f,) = o.ufl_operands
        return 2.0 * fp / (cos(2.0 * f) + 1.0)

    @process.register(Cosh)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a cosh."""
        (f,) = o.ufl_operands
        return fp * sinh(f)

    @process.register(Sinh)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a sinh."""
        (f,) = o.ufl_operands
        return fp * cosh(f)

    @process.register(Tanh)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a tanh."""
        (f,) = o.ufl_operands

        def sech(y):
            return (2.0 * cosh(y)) / (cosh(2.0 * y) + 1.0)

        return fp * sech(f) ** 2

    @process.register(Acos)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate an acos."""
        (f,) = o.ufl_operands
        return -fp / sqrt(1.0 - f**2)

    @process.register(Asin)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate an asin."""
        (f,) = o.ufl_operands
        return fp / sqrt(1.0 - f**2)

    @process.register(Atan)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate an atan."""
        (f,) = o.ufl_operands
        return fp / (1.0 + f**2)

    @process.register(Atan2)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr, gp: Expr) -> Expr:
        """Differentiate an atan2."""
        f, g = o.ufl_operands
        return (g * fp - f * gp) / (f**2 + g**2)

    @process.register(Erf)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate an erf."""
        (f,) = o.ufl_operands
        return fp * (2.0 / sqrt(pi) * exp(-(f**2)))

    # --- Bessel functions

    @process.register(BesselJ)
    @DAGTraverser.postorder
    def _(self, o: Expr, nup: Expr, fp: Expr) -> Expr:
        """Differentiate a bessel_j."""
        nu, f = o.ufl_operands
        if not (nup is None or isinstance(nup, Zero)):
            raise NotImplementedError(
                "Differentiation of bessel function w.r.t. nu is not supported."
            )

        if isinstance(nu, Zero):
            op = -bessel_J(1, f)
        else:
            op = 0.5 * (bessel_J(nu - 1, f) - bessel_J(nu + 1, f))
        return op * fp

    @process.register(BesselY)
    @DAGTraverser.postorder
    def _(self, o: Expr, nup: Expr, fp: Expr) -> Expr:
        """Differentiate a bessel_y."""
        nu, f = o.ufl_operands
        if not (nup is None or isinstance(nup, Zero)):
            raise NotImplementedError(
                "Differentiation of bessel function w.r.t. nu is not supported."
            )

        if isinstance(nu, Zero):
            op = -bessel_Y(1, f)
        else:
            op = 0.5 * (bessel_Y(nu - 1, f) - bessel_Y(nu + 1, f))
        return op * fp

    @process.register(BesselI)
    @DAGTraverser.postorder
    def _(self, o: Expr, nup: Expr, fp: Expr) -> Expr:
        """Differentiate a bessel_i."""
        nu, f = o.ufl_operands
        if not (nup is None or isinstance(nup, Zero)):
            raise NotImplementedError(
                "Differentiation of bessel function w.r.t. nu is not supported."
            )

        if isinstance(nu, Zero):
            op = bessel_I(1, f)
        else:
            op = 0.5 * (bessel_I(nu - 1, f) + bessel_I(nu + 1, f))
        return op * fp

    @process.register(BesselK)
    @DAGTraverser.postorder
    def _(self, o: Expr, nup: Expr, fp: Expr) -> Expr:
        """Differentiate a bessel_k."""
        nu, f = o.ufl_operands
        if not (nup is None or isinstance(nup, Zero)):
            raise NotImplementedError(
                "Differentiation of bessel function w.r.t. nu is not supported."
            )

        if isinstance(nu, Zero):
            op = -bessel_K(1, f)
        else:
            op = -0.5 * (bessel_K(nu - 1, f) + bessel_K(nu + 1, f))
        return op * fp

    # --- Restrictions

    @process.register(Restricted)
    @DAGTraverser.postorder
    def _(self, o: Restricted, fp: Expr) -> Expr:
        """Differentiate a restricted."""
        # Restriction and differentiation commutes
        if isinstance(fp, ConstantValue):
            return fp  # TODO: Add simplification to Restricted instead?
        else:
            return fp(o._side)  # (f+-)' == (f')+-

    # --- Conditionals

    @process.register(BinaryCondition)
    def _(self, o: BinaryCondition) -> Expr:
        """Differentiate a binary_condition."""
        raise RuntimeError("Can not differentiate a binary_condition.")

    @process.register(NotCondition)
    def _(self, o: NotCondition) -> Expr:
        """Differentiate a not_condition."""
        raise RuntimeError("Can not differentiate a not_condition.")

    @process.register(Conditional)
    @DAGTraverser.postorder_only_children([1, 2])
    def _(self, o: Expr, dt: Expr, df: Expr) -> Expr:
        """Differentiate a conditional."""
        if isinstance(dt, Zero) and isinstance(df, Zero):
            # Assuming dt and df have the same indices here, which
            # should be the case
            return dt
        else:
            # Not placing t[1],f[1] outside, allowing arguments inside
            # conditionals.  This will make legacy ffc fail, but
            # should work with uflacs.
            c = o.ufl_operands[0]
            return conditional(c, dt, df)

    @process.register(MaxValue)
    @DAGTraverser.postorder
    def _(self, o: Expr, df: Expr, dg: Expr) -> Expr:
        """Differentiate a max_value."""
        # d/dx max(f, g) =
        # f > g: df/dx
        # f < g: dg/dx
        # Placing df,dg outside here to avoid getting arguments inside
        # conditionals
        f, g = o.ufl_operands
        dc = conditional(f > g, 1, 0)
        return dc * df + (1.0 - dc) * dg

    @process.register(MinValue)
    @DAGTraverser.postorder
    def _(self, o: Expr, df: Expr, dg: Expr) -> Expr:
        """Differentiate a min_value."""
        # d/dx min(f, g) =
        #  f < g: df/dx
        #  else: dg/dx
        #  Placing df,dg outside here to avoid getting arguments
        #  inside conditionals
        f, g = o.ufl_operands
        dc = conditional(f < g, 1, 0)
        return dc * df + (1.0 - dc) * dg


class GradRuleset(GenericDerivativeRuleset):
    """Take the grad derivative."""

    def __init__(
        self,
        geometric_dimension: int,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(
            (geometric_dimension,),
            compress=compress,
            visited_cache=visited_cache,
            result_cache=result_cache,
        )
        self._Id = Identity(geometric_dimension)

    # Work around singledispatchmethod inheritance issue;
    # see https://bugs.python.org/issue36457.
    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    # --- Specialized rules for geometric quantities

    @process.register(GeometricQuantity)
    def _(self, o: Expr) -> Expr:
        """Differentiate a geometric_quantity.

        Default for geometric quantities is do/dx = 0 if piecewise constant,
        otherwise transform derivatives to reference derivatives.
        Override for specific types if other behaviour is needed.
        """
        if is_cellwise_constant(o):
            return self.independent_terminal(o)
        else:
            domain = extract_unique_domain(o)
            K = JacobianInverse(domain)
            Do = grad_to_reference_grad(o, K)
            return Do

    @process.register(JacobianInverse)
    def _(self, o: JacobianInverse) -> Expr:
        """Differentiate a jacobian_inverse."""
        # grad(K) == K_ji rgrad(K)_rj
        if is_cellwise_constant(o):
            return self.independent_terminal(o)
        if not o._ufl_is_terminal_:
            raise ValueError("ReferenceValue can only wrap a terminal")
        Do = grad_to_reference_grad(o, o)
        return Do

    # TODO: Add more explicit geometry type handlers here, with
    # non-affine domains several should be non-zero.

    @process.register(SpatialCoordinate)
    def _(self, o: Expr) -> Expr:
        """Differentiate a spatial_coordinate.

        dx/dx = I.
        """
        return self._Id

    @process.register(CellCoordinate)
    def _(self, o: Expr) -> Expr:
        """Differentiate a cell_coordinate.

        dX/dx = inv(dx/dX) = inv(J) = K.
        """
        # FIXME: Is this true for manifolds? What about orientation?
        return JacobianInverse(extract_unique_domain(o))

    # --- Specialized rules for form arguments

    @process.register(BaseFormOperator)
    def _(self, o: Expr) -> Expr:
        """Differentiate a base_form_operator."""
        # Push the grad through the operator is not legal in most cases:
        #    -> Not enouth regularity for chain rule to hold!
        # By the time we evaluate `grad(o)`, the operator `o` will have
        # been assembled and substituted by its output.
        return Grad(o)

    @process.register(Coefficient)
    def _(self, o: Expr) -> Expr:
        """Differentiate a coefficient."""
        if is_cellwise_constant(o):
            return self.independent_terminal(o)
        return Grad(o)

    @process.register(Argument)
    def _(self, o: Expr) -> Expr:
        """Differentiate an argument."""
        # TODO: Enable this after fixing issue#13, unless we move
        # simplificat ion to a separate stage?
        # if is_cellwise_constant(o):
        #     # Collapse gradient of cellwise constant function to zero
        #     # TODO: Missing this type
        #     return AnnotatedZero(o.ufl_shape + self._var_shape, arguments=(o,))
        return Grad(o)

    # --- Rules for values or derivatives in reference frame

    @process.register(ReferenceValue)
    def _(self, o: ReferenceValue) -> Expr:
        """Differentiate a reference_value."""
        # grad(o) == grad(rv(f)) -> K_ji*rgrad(rv(f))_rj
        f = o.ufl_operands[0]
        if not f._ufl_is_terminal_:
            raise ValueError("ReferenceValue can only wrap a terminal")
        domain = extract_unique_domain(f, expand_mesh_sequence=False)
        if isinstance(domain, MeshSequence):
            element = f.ufl_function_space().ufl_element()  # type: ignore
            if element.num_sub_elements != len(domain):
                raise RuntimeError(f"{element.num_sub_elements} != {len(domain)}")
            # Get monolithic representation of rgrad(o); o might live in a mixed space.
            rgrad = ReferenceGrad(o)
            ref_dim = rgrad.ufl_shape[-1]
            # Apply K_ji(d) to the corresponding components of rgrad, store them in a list,
            # and put them back together at the end using as_tensor().
            components = []
            dofoffset = 0
            for d, e in flatten_domain_element(domain, element):
                esh = e.reference_value_shape
                ndof = int(np.prod(esh))
                assert ndof > 0
                if isinstance(e.pullback, PhysicalPullback):
                    if ref_dim != self._var_shape[0]:
                        raise NotImplementedError("""
                            PhysicalPullback not handled for immersed domain :
                            reference dim ({ref_dim}) != physical dim (self._var_shape[0])""")
                    for idx in range(ndof):
                        for i in range(ref_dim):
                            components.append(rgrad[(dofoffset + idx,) + (i,)])
                else:
                    K = JacobianInverse(d)
                    rdim, gdim = K.ufl_shape
                    if rdim != ref_dim:
                        raise RuntimeError(f"{rdim} != {ref_dim}")
                    if gdim != self._var_shape[0]:
                        raise RuntimeError(f"{gdim} != {self._var_shape[0]}")
                    # Note that rgrad[dofoffset + [0,ndof), [0,rdim)] are the components
                    # corresponding to (d, e).
                    # For each row, rgrad[dofoffset + idx, [0,rdim)], we apply
                    # K_ji(d)[[0,rdim), [0,gdim)].
                    for idx in range(ndof):
                        for i in range(gdim):
                            temp = Zero()
                            for j in range(rdim):
                                temp += rgrad[(dofoffset + idx,) + (j,)] * K[j, i]
                            components.append(temp)
                dofoffset += ndof
            return as_tensor(np.asarray(components).reshape(rgrad.ufl_shape[:-1] + self._var_shape))
        else:
            if isinstance(f.ufl_element().pullback, PhysicalPullback):  # type: ignore
                # TODO: Do we need to be more careful for immersed things?
                return ReferenceGrad(o)
            else:
                K = JacobianInverse(domain)
                return grad_to_reference_grad(o, K)

    @process.register(ReferenceGrad)
    def _(self, o: Expr) -> Expr:
        """Differentiate a reference_grad."""
        if is_cellwise_constant(o):
            return self.independent_terminal(o)
        # grad(o) == grad(rgrad(rv(f))) -> K_ji*rgrad(rgrad(rv(f)))_rj
        f = o.ufl_operands[0]
        valid_operand = f._ufl_is_in_reference_frame_ or isinstance(
            f, JacobianInverse | SpatialCoordinate | Jacobian | JacobianDeterminant | FacetNormal
        )
        if not valid_operand:
            raise ValueError("ReferenceGrad can only wrap a reference frame type!")
        domain = extract_unique_domain(f, expand_mesh_sequence=False)
        if isinstance(domain, MeshSequence):
            if not f._ufl_is_in_reference_frame_:
                raise RuntimeError("Expecting a reference frame type")
            while not f._ufl_is_terminal_:
                (f,) = f.ufl_operands  # type: ignore
            element = f.ufl_function_space().ufl_element()  # type: ignore
            if element.num_sub_elements != len(domain):
                raise RuntimeError(f"{element.num_sub_elements} != {len(domain)}")
            # Get monolithic representation of rgrad(o); o might live in a mixed space.
            rgrad = ReferenceGrad(o)
            ref_dim = rgrad.ufl_shape[-1]
            # Apply K_ji(d) to the corresponding components of rgrad, store them in a list,
            # and put them back together at the end using as_tensor().
            components = []
            dofoffset = 0
            for d, e in flatten_domain_element(domain, element):
                esh = e.reference_value_shape
                ndof = int(np.prod(esh))
                assert ndof > 0
                K = JacobianInverse(d)
                rdim, gdim = K.ufl_shape
                if rdim != ref_dim:
                    raise RuntimeError(f"{rdim} != {ref_dim}")
                if gdim != self._var_shape[0]:
                    raise RuntimeError(f"{gdim} != {self._var_shape[0]}")
                # Note that rgrad[dofoffset + [0,ndof), [0,rdim), [0,rdim)] are the components
                # corresponding to (d, e).
                # For each row, rgrad[dofoffset + idx, [0,rdim), [0,rdim)], we apply
                # K_ji(d)[[0,rdim), [0,gdim)].
                for idx in range(ndof):
                    for midx in np.ndindex(rgrad.ufl_shape[1:-1]):
                        for i in range(gdim):
                            temp = Zero()
                            for j in range(rdim):
                                temp += rgrad[(dofoffset + idx,) + midx + (j,)] * K[j, i]
                            components.append(temp)
                dofoffset += ndof
            if rgrad.ufl_shape[0] != dofoffset:
                raise RuntimeError(f"{rgrad.ufl_shape[0]} != {dofoffset}")
            return as_tensor(np.asarray(components).reshape(rgrad.ufl_shape[:-1] + self._var_shape))
        else:
            K = JacobianInverse(domain)
            return grad_to_reference_grad(o, K)

    # --- Nesting of gradients

    @process.register(Grad)
    def _(self, o: Expr) -> Expr:
        """Differentiate a grad.

        Represent grad(grad(f)) as Grad(Grad(f)).
        """
        # Check that o is a "differential terminal"
        if not isinstance(o.ufl_operands[0], Grad | Terminal):
            raise ValueError("Expecting only grads applied to a terminal.")
        return Grad(o)

    def _grad(self, o):
        """Differentiate a _grad."""
        pass
        # TODO: Not sure how to detect that gradient of f is cellwise constant.
        #       Can we trust element degrees?
        # if is_cellwise_constant(o):
        #     return self.terminal(o)
        # TODO: Maybe we can ask "f.has_derivatives_of_order(n)" to check
        #       if we should make a zero here?
        # 1) n = count number of Grads, get f
        # 2) if not f.has_derivatives(n): return zero(...)

    @process.register(CellAvg)
    @process.register(FacetAvg)
    def _(self, o: Expr) -> Expr:
        return self.independent_operator(o)


def grad_to_reference_grad(o, K):
    """Relates grad(o) to reference_grad(o) using the Jacobian inverse.

    Args:
        o: Operand
        K: Jacobian inverse
    Returns:
        grad(o) written in terms of reference_grad(o) and K
    """
    r = indices(len(o.ufl_shape))
    i, j = indices(2)
    # grad(o) == K_ji rgrad(o)_rj
    Do = as_tensor(K[j, i] * ReferenceGrad(o)[r + (j,)], r + (i,))
    return Do


class ReferenceGradRuleset(GenericDerivativeRuleset):
    """Apply the reference grad derivative."""

    def __init__(
        self,
        topological_dimension: int,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(
            (topological_dimension,),
            compress=compress,
            visited_cache=visited_cache,
            result_cache=result_cache,
        )
        self._Id = Identity(topological_dimension)

    # Work around singledispatchmethod inheritance issue;
    # see https://bugs.python.org/issue36457.
    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    # --- Specialized rules for geometric quantities

    @process.register(GeometricQuantity)
    def _(self, o: Expr) -> Expr:
        """Differentiate a geometric_quantity.

        dg/dX = 0 if piecewise constant, otherwise ReferenceGrad(g).
        """
        if is_cellwise_constant(o):
            return self.independent_terminal(o)
        else:
            # TODO: Which types does this involve? I don't think the
            # form compilers will handle this.
            return ReferenceGrad(o)

    @process.register(SpatialCoordinate)
    def _(self, o: Expr) -> Expr:
        """Differentiate a spatial_coordinate.

        dx/dX = J.
        """
        # Don't convert back to J, otherwise we get in a loop
        return ReferenceGrad(o)

    @process.register(CellCoordinate)
    def _(self, o: Expr) -> Expr:
        """Differentiate a cell_coordinate.

        dX/dX = I.
        """
        return self._Id

    # TODO: Add more geometry types here, with non-affine domains
    # several should be non-zero.

    # --- Specialized rules for form arguments

    @process.register(ReferenceValue)
    def _(self, o: Expr) -> Expr:
        """Differentiate a reference_value."""
        if not o.ufl_operands[0]._ufl_is_terminal_:
            raise ValueError("ReferenceValue can only wrap a terminal")
        return ReferenceGrad(o)

    @process.register(Coefficient)
    def _(self, o: Expr) -> Expr:
        """Differentiate a coefficient."""
        raise ValueError("Coefficient should be wrapped in ReferenceValue by now")

    @process.register(Argument)
    def _(self, o: Expr) -> Expr:
        """Differentiate an argument."""
        raise ValueError("Argument should be wrapped in ReferenceValue by now")

    # --- Nesting of gradients

    @process.register(Grad)
    def _(self, o: Expr) -> Expr:
        """Differentiate a grad."""
        raise ValueError(
            f"Grad should have been transformed by this point, but got {type(o).__name__}."
        )

    @process.register(ReferenceGrad)
    def _(self, o: Expr) -> Expr:
        """Differentiate a reference_grad.

        Represent ref_grad(ref_grad(f)) as RefGrad(RefGrad(f)).
        """
        # Check that o is a "differential terminal"
        if not isinstance(o.ufl_operands[0], ReferenceGrad | ReferenceValue | Terminal):
            raise ValueError("Expecting only grads applied to a terminal.")
        return ReferenceGrad(o)

    @process.register(CellAvg)
    @process.register(FacetAvg)
    def _(self, o: Expr) -> Expr:
        return self.independent_operator(o)


class VariableRuleset(GenericDerivativeRuleset):
    """Differentiate with respect to a variable."""

    def __init__(
        self,
        var: Expr,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(
            var.ufl_shape, compress=compress, visited_cache=visited_cache, result_cache=result_cache
        )
        if var.ufl_free_indices:
            raise ValueError("Differentiation variable cannot have free indices.")
        self._variable = var
        self._Id = self._make_identity(self._var_shape)

    def _make_identity(self, sh):
        """Differentiate a _make_identity.

        Creates a higher order identity tensor to represent dv/dv.
        """
        res = None
        if sh == ():
            # Scalar dv/dv is scalar
            return FloatValue(1.0)
        elif len(sh) == 1:
            # Vector v makes dv/dv the identity matrix
            return Identity(sh[0])
        else:
            # TODO: Add a type for this higher order identity?
            # II[i0,i1,i2,j0,j1,j2] = 1 if all((i0==j0, i1==j1, i2==j2)) else 0
            # Tensor v makes dv/dv some kind of higher rank identity tensor
            ind1 = ()
            ind2 = ()
            for d in sh:
                i, j = indices(2)
                dij = Identity(d)[i, j]
                if res is None:
                    res = dij
                else:
                    res *= dij
                ind1 += (i,)
                ind2 += (j,)
            fp = as_tensor(res, ind1 + ind2)
        return fp

    # Work around singledispatchmethod inheritance issue;
    # see https://bugs.python.org/issue36457.
    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    @process.register(GeometricQuantity)
    def _(self, o: Expr) -> Expr:
        # Explicitly defining dg/dw == 0
        return self.independent_terminal(o)

    @process.register(Argument)
    def _(self, o: Expr) -> Expr:
        # Explicitly defining da/dw == 0
        return self.independent_terminal(o)

    @process.register(Coefficient)
    def _(self, o: Expr) -> Expr:
        """Differentiate a coefficient.

        df/dv = Id if v is f else 0.

        Note that if v = variable(f), df/dv is still 0,
        but if v == f, i.e. isinstance(v, Coefficient) == True,
        then df/dv == df/df = Id.
        """
        v = self._variable
        if isinstance(v, Coefficient) and o == v:
            # dv/dv = identity of rank 2*rank(v)
            return self._Id
        else:
            # df/v = 0
            return self.independent_terminal(o)

    @process.register(Variable)
    @DAGTraverser.postorder
    def _(self, o: Expr, df: Expr, a: Expr) -> Expr:
        """Differentiate a variable."""
        v = self._variable
        if isinstance(v, Variable) and v.label() == a:
            # dv/dv = identity of rank 2*rank(v)
            return self._Id
        else:
            # df/v = df
            return df

    @process.register(Grad)
    def _(self, o: Expr) -> Expr:
        """Differentiate a grad.

        Variable derivative of a gradient of a terminal must be 0.
        """
        # Check that o is a "differential terminal"
        if not isinstance(o.ufl_operands[0], Grad | Terminal):
            raise ValueError("Expecting only grads applied to a terminal.")
        return self.independent_terminal(o)

    # --- Rules for values or derivatives in reference frame

    @process.register(ReferenceValue)
    def _(self, o: Expr) -> Expr:
        """Differentiate a reference_value."""
        # d/dv(o) == d/dv(rv(f)) = 0 if v is not f, or rv(dv/df)
        v = self._variable
        if isinstance(v, Coefficient) and o.ufl_operands[0] == v:
            if not v.ufl_element().pullback.is_identity:
                # FIXME: This is a bit tricky, instead of Identity it is
                #   actually inverse(transform), or we should rather not
                #   convert to reference frame in the first place
                raise ValueError(
                    "Missing implementation: To handle derivatives of rv(f) w.r.t. f for "
                    "mapped elements, rewriting to reference frame should not happen first..."
                )
            # dv/dv = identity of rank 2*rank(v)
            return self._Id
        else:
            # df/v = 0
            return self.independent_terminal(o)

    @process.register(ReferenceGrad)
    def _(self, o: Expr) -> Expr:
        """Differentiate a reference_grad.

        Variable derivative of a gradient of a terminal must be 0.
        """
        return self.independent_terminal(o)

    @process.register(CellAvg)
    @process.register(FacetAvg)
    def _(self, o: Expr) -> Expr:
        return self.independent_operator(o)


class GateauxDerivativeRuleset(GenericDerivativeRuleset):
    """Apply AFD (Automatic Functional Differentiation) to expression.

    Implements rules for the Gateaux derivative D_w[v](...) defined as
    D_w[v](e) = d/dtau e(w+tau v)|tau=0.
    """

    def __init__(
        self,
        coefficients: ExprList,
        arguments: ExprList,
        coefficient_derivatives: ExprMapping,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(
            (), compress=compress, visited_cache=visited_cache, result_cache=result_cache
        )
        # Type checking
        if not isinstance(coefficients, ExprList):
            raise ValueError("Expecting a ExprList of coefficients.")
        if not isinstance(arguments, ExprList):
            raise ValueError("Expecting a ExprList of arguments.")
        if not isinstance(coefficient_derivatives, ExprMapping):
            raise ValueError("Expecting a coefficient-coefficient ExprMapping.")
        # The coefficient(s) to differentiate w.r.t. and the
        # argument(s) s.t. D_w[v](e) = d/dtau e(w+tau v)|tau=0
        self._w = coefficients.ufl_operands
        self._v = arguments.ufl_operands
        self._w2v = {w: v for w, v in zip(self._w, self._v)}
        # Build more convenient dict {f: df/dw} for each coefficient f
        # where df/dw is nonzero
        cd = coefficient_derivatives.ufl_operands
        self._cd = {cd[2 * i]: cd[2 * i + 1] for i in range(len(cd) // 2)}
        # Record the operations delayed to the derivative expansion phase:
        # Example: dN(u)/du where `N` is an ExternalOperator and `u` a Coefficient
        self.pending_operations = BaseFormOperatorDerivativeRecorder(
            coefficients,
            arguments=arguments,
            coefficient_derivatives=coefficient_derivatives,
        )

    # Work around singledispatchmethod inheritance issue;
    # see https://bugs.python.org/issue36457.
    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    # --- Specialized rules for geometric quantities

    @process.register(GeometricQuantity)
    def _(self, o: Expr) -> Expr:
        # Explicitly defining dg/dw == 0
        return self.independent_terminal(o)

    @process.register(CellAvg)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a cell_avg."""
        # Cell average of a single function and differentiation
        # commutes, D_f[v](cell_avg(f)) = cell_avg(v)
        return cell_avg(fp)

    @process.register(FacetAvg)
    @DAGTraverser.postorder
    def _(self, o: Expr, fp: Expr) -> Expr:
        """Differentiate a facet_avg."""
        # Facet average of a single function and differentiation
        # commutes, D_f[v](facet_avg(f)) = facet_avg(v)
        return facet_avg(fp)

    @process.register(Argument)
    def _(self, o: Argument) -> Expr:
        # Explicitly defining da/dw == 0
        return self._process_argument(o)

    def _process_argument(self, o: Argument | Coargument) -> Zero:
        return self.independent_terminal(o)

    @process.register(Coefficient)
    def _(self, o: Coefficient) -> Expr:
        return self._process_coefficient(o)

    def _process_coefficient(self, o: Expr) -> Expr:
        """Differentiate an Expr or a BaseForm."""
        # Define dw/dw := d/ds [w + s v] = v

        # Return corresponding argument if we can find o among w
        do = self._w2v.get(o)  # type: ignore
        if do is not None:
            return do

        # Look for o among coefficient derivatives
        dos = self._cd.get(o)  # type: ignore
        if dos is None:
            # If o is not among coefficient derivatives, return
            # do/dw=0
            do = Zero(o.ufl_shape)
            return do
        else:
            # Compute do/dw_j = do/dw_h : v.
            # Since we may actually have a tuple of oprimes and vs in a
            # 'mixed' space, sum over them all to get the complete inner
            # product. Using indices to define a non-compound inner product.

            # Example:
            # (f:g) -> (dfdu:v):g + f:(dgdu:v)
            # shape(dfdu) == shape(f) + shape(v)
            # shape(f) == shape(g) == shape(dfdu : v)

            # Make sure we have a tuple to match the self._v tuple
            if not isinstance(dos, tuple):
                dos = (dos,)
            if len(dos) != len(self._v):
                raise ValueError(
                    "Got a tuple of arguments, expecting a "
                    "matching tuple of coefficient derivatives."
                )
            dosum = Zero(o.ufl_shape)
            for do, v in zip(dos, self._v):
                so, oi = as_scalar(do)
                rv = len(oi) - len(v.ufl_shape)
                oi1 = oi[:rv]
                oi2 = oi[rv:]
                prod = so * v[oi2]
                if oi1:
                    dosum += as_tensor(prod, oi1)
                else:
                    dosum += prod
            return dosum

    @process.register(ReferenceValue)
    def _(self, o: Expr) -> Expr:
        """Differentiate a reference_value."""
        raise NotImplementedError(
            "Currently no support for ReferenceValue in CoefficientDerivative."
        )
        # TODO: This is implementable for regular derivative(M(f),f,v)
        #       but too messy if customized coefficient derivative
        #       relations are given by the user.  We would only need
        #       this to allow the user to write
        #       derivative(...ReferenceValue...,...).
        # f, = o.ufl_operands
        # if not f._ufl_is_terminal_:
        #     raise ValueError("ReferenceValue can only wrap terminals directly.")
        # FIXME: check all cases like in coefficient
        # if f is w:
        #     # FIXME: requires that v is an Argument with the same element mapping!
        #     return ReferenceValue(v)
        # else:
        #     return self.independent_terminal(o)

    @process.register(ReferenceGrad)
    def _(self, o: Expr) -> Expr:
        """Differentiate a reference_grad."""
        if len(extract_coefficients(o)) > 0:
            raise NotImplementedError(
                "Currently no support for ReferenceGrad in CoefficientDerivative."
            )
        else:
            return Zero(o.ufl_shape)
        # TODO: This is implementable for regular derivative(M(f),f,v)
        #       but too messy if customized coefficient derivative
        #       relations are given by the user.  We would only need
        #       this to allow the user to write
        #       derivative(...ReferenceValue...,...).

    @process.register(Grad)
    def _(self, g: Expr) -> Expr:
        """Differentiate a grad."""
        # If we hit this type, it has already been propagated to a
        # coefficient (or grad of a coefficient) or a base form operator, # FIXME: Assert
        # this!  so we need to take the gradient of the variation or
        # return zero.  Complications occur when dealing with
        # derivatives w.r.t. single components...

        # Figure out how many gradients are around the inner terminal
        ngrads = 0
        o = g
        while isinstance(o, Grad):
            (o,) = o.ufl_operands
            ngrads += 1
        # `grad(N)` where N is a BaseFormOperator is treated as if `N` was a Coefficient.
        if not isinstance(o, FormArgument | BaseFormOperator):
            raise ValueError(f"Expecting gradient of a FormArgument, not {ufl_err_str(o)}.")

        def apply_grads(f):
            for i in range(ngrads):
                f = Grad(f)
            return f

        # Find o among all w without any indexing, which makes this
        # easy
        for w, v in zip(self._w, self._v):
            if o == w and isinstance(v, FormArgument):
                # Case: d/dt [w + t v]
                return apply_grads(v)

        # If o is not among coefficient derivatives, return do/dw=0
        gprimesum = Zero(g.ufl_shape)

        def analyse_variation_argument(v):
            # Analyse variation argument
            if isinstance(v, FormArgument):
                # Case: d/dt [w[...] + t v]
                vval, vcomp = v, ()
            elif isinstance(v, Indexed):
                # Case: d/dt [w + t v[...]]
                # Case: d/dt [w[...] + t v[...]]
                vval, vcomp = v.ufl_operands
                vcomp = tuple(vcomp)
            else:
                raise ValueError("Expecting argument or component of argument.")
            if not all(isinstance(k, FixedIndex) for k in vcomp):
                raise ValueError("Expecting only fixed indices in variation.")
            return vval, vcomp

        def compute_gprimeterm(ngrads, vval, vcomp, wshape, wcomp):
            # Apply gradients directly to argument vval, and get the
            # right indexed scalar component(s)
            kk = indices(ngrads)
            Dvkk = apply_grads(vval)[vcomp + kk]
            # Place scalar component(s) Dvkk into the right tensor
            # positions
            if wshape:
                Ejj, jj = unit_indexed_tensor(wshape, wcomp)
            else:
                Ejj, jj = 1, ()
            gprimeterm = as_tensor(Ejj * Dvkk, jj + kk)
            return gprimeterm

        # Accumulate contributions from variations in different
        # components
        for w, v in zip(self._w, self._v):
            # -- Analyse differentiation variable coefficient -- #

            # Can differentiate a Form wrt a BaseFormOperator
            if isinstance(w, FormArgument | BaseFormOperator):
                if not w == o:
                    continue
                wshape = w.ufl_shape

                if isinstance(v, FormArgument):
                    # Case: d/dt [w + t v]
                    return apply_grads(v)

                elif isinstance(v, ListTensor):
                    # Case: d/dt [w + t <...,v,...>]
                    for wcomp, vsub in unwrap_list_tensor(v):
                        if not isinstance(vsub, Zero):
                            vval, vcomp = analyse_variation_argument(vsub)
                            gprimesum = gprimesum + compute_gprimeterm(
                                ngrads, vval, vcomp, wshape, wcomp
                            )
                elif isinstance(v, Zero):
                    pass

                else:
                    if wshape != ():
                        raise ValueError("Expecting scalar coefficient in this branch.")
                    # Case: d/dt [w + t v[...]]
                    wval, wcomp = w, ()

                    vval, vcomp = analyse_variation_argument(v)
                    gprimesum = gprimesum + compute_gprimeterm(ngrads, vval, vcomp, wshape, wcomp)

            elif isinstance(
                w, Indexed
            ):  # This path is tested in unit tests, but not actually used?
                # Case: d/dt [w[...] + t v[...]]
                # Case: d/dt [w[...] + t v]
                wval, wcomp = w.ufl_operands
                if not wval == o:
                    continue
                assert isinstance(wval, FormArgument)
                if not all(isinstance(k, FixedIndex) for k in wcomp):
                    raise ValueError("Expecting only fixed indices in differentiation variable.")
                wshape = wval.ufl_shape

                vval, vcomp = analyse_variation_argument(v)
                gprimesum = gprimesum + compute_gprimeterm(ngrads, vval, vcomp, wshape, wcomp)

            else:
                raise ValueError("Expecting coefficient or component of coefficient.")

        # FIXME: Handle other coefficient derivatives: oprimes =
        # self._cd.get(o)

        if 0:
            oprimes = self._cd.get(o)
            if oprimes is None:
                if self._cd:
                    # TODO: Make it possible to silence this message
                    #       in particular?  It may be good to have for
                    #       debugging...
                    warnings.warn(f"Assuming d{{{0}}}/d{{{self._w}}} = 0.")
            else:
                # Make sure we have a tuple to match the self._v tuple
                if not isinstance(oprimes, tuple):
                    oprimes = (oprimes,)
                    if len(oprimes) != len(self._v):
                        raise ValueError(
                            "Got a tuple of arguments, expecting a"
                            " matching tuple of coefficient derivatives."
                        )

                # Compute dg/dw_j = dg/dw_h : v.
                # Since we may actually have a tuple of oprimes and vs
                # in a 'mixed' space, sum over them all to get the
                # complete inner product. Using indices to define a
                # non-compound inner product.
                for oprime, v in zip(oprimes, self._v):
                    raise NotImplementedError("FIXME: Figure out how to do this with ngrads")
                    so, oi = as_scalar(oprime)
                    rv = len(v.ufl_shape)
                    oi1 = oi[:-rv]
                    oi2 = oi[-rv:]
                    prod = so * v[oi2]
                    if oi1:
                        gprimesum += as_tensor(prod, oi1)
                    else:
                        gprimesum += prod

        return gprimesum

    @process.register(CoordinateDerivative)
    @DAGTraverser.postorder_only_children([0])
    def _(self, o: Expr, o0: Expr) -> Expr:
        """Differentiate a coordinate_derivative."""
        _, o1, o2, o3 = o.ufl_operands
        return CoordinateDerivative(o0, o1, o2, o3)

    @process.register(BaseFormOperator)
    @DAGTraverser.postorder
    def _(self, o: BaseFormOperator, *dfs) -> Expr:
        """Differentiate a base_form_operator.

        If d_coeff = 0 => BaseFormOperator's derivative is taken wrt a
        variable => we call the appropriate handler. Otherwise =>
        differentiation done wrt the BaseFormOperator (dF/dN[Nhat]) =>
        we treat o as a Coefficient.
        """
        d_coeff = self._process_coefficient(o)
        # It also handles the non-scalar case
        if d_coeff == 0:
            self.pending_operations += (o,)
        return d_coeff

    # -- Handlers for BaseForm objects -- #

    @process.register(Cofunction)
    def _(self, o: Cofunction) -> Expr:
        """Differentiate a cofunction."""
        # Same rule than for Coefficient except that we use a Coargument.
        # The coargument is already attached to the class (self._v)
        # which `self.coefficient` relies on.
        dc = self._process_coefficient(o)  # type: ignore
        if dc == 0:
            # Convert ufl.Zero into ZeroBaseForm
            return ZeroBaseForm(o.arguments() + self._v)  # type: ignore
        return dc

    @process.register(Coargument)
    def _(self, o: Coargument) -> Expr:
        """Differentiate a coargument."""
        # Same rule than for Argument (da/dw == 0).
        dc = self._process_argument(o)
        if dc == 0:
            # Convert ufl.Zero into ZeroBaseForm
            return ZeroBaseForm(o.arguments() + self._v)  # type: ignore
        return dc

    @process.register(Matrix)  # type: ignore
    def _(self, M: Matrix) -> BaseForm:
        """Differentiate a matrix."""
        # Matrix rule: D_w[v](M) = v if M == w else 0
        # We can't differentiate wrt a matrix so always return zero in
        # the appropriate space
        return ZeroBaseForm(M.arguments() + self._v)


class BaseFormOperatorDerivativeRuleset(GateauxDerivativeRuleset):
    """Apply AFD (Automatic Functional Differentiation) to BaseFormOperator.

    Implements rules for the Gateaux derivative D_w[v](...) defined as
    D_w[v](B) = d/dtau B(w+tau v)|tau=0 where B is a ufl.BaseFormOperator.
    """

    @staticmethod
    def pending_operations_recording(base_form_operator_handler):
        """Decorate a function to record pending operations."""

        def wrapper(self, base_form_op, *dfs):
            """Decorate."""
            # Get the outer `BaseFormOperator` expression, i.e. the
            # operator that is being differentiated.
            expression = self.outer_base_form_op
            # If the base form operator we observe is different from the
            # outer `BaseFormOperator`:
            # -> Record that `BaseFormOperator` so that
            # `d(expression)/d(base_form_op)` can then be computed
            # later.
            # Else:
            # -> Compute the Gateaux derivative of `base_form_ops` by
            # calling the appropriate handler.
            if expression != base_form_op:
                self.pending_operations += (base_form_op,)
                return self._process_coefficient(base_form_op)
            return base_form_operator_handler(self, base_form_op, *dfs)

        return wrapper

    def __init__(
        self,
        coefficients: ExprList,
        arguments: ExprList,
        coefficient_derivatives: ExprMapping,
        outer_base_form_op: Expr,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(
            coefficients,
            arguments,
            coefficient_derivatives,
            compress=compress,
            visited_cache=visited_cache,
            result_cache=result_cache,
        )
        self.outer_base_form_op = outer_base_form_op

    # Work around singledispatchmethod inheritance issue;
    # see https://bugs.python.org/issue36457.
    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    @process.register(Interpolate)
    @DAGTraverser.postorder
    @pending_operations_recording
    def _(self, i_op: Interpolate, dw: Expr) -> Expr:
        """Differentiate an interpolate."""
        # Interpolate rule: D_w[v](i_op(w, v*)) = i_op(v, v*), by linearity of Interpolate!
        if not dw:
            # i_op doesn't depend on w:
            #  -> It also covers the Hessian case since Interpolate is linear,
            #     e.g. D_w[v](D_w[v](i_op(w, v*))) = D_w[v](i_op(v, v*)) = 0 (since w not found).
            return ZeroBaseForm(i_op.arguments() + self._v)  # type: ignore
        return i_op._ufl_expr_reconstruct_(expr=dw)

    @process.register(ExternalOperator)
    @DAGTraverser.postorder
    @pending_operations_recording
    def external_operator(self, N: ExternalOperator, *dfs) -> Expr:
        """Differentiate an external_operator."""
        result: tuple[Expr, ...] = ()
        for i, df in enumerate(dfs):
            derivatives = tuple(dj + int(i == j) for j, dj in enumerate(N.derivatives))
            if len(extract_arguments(df)) != 0:
                # Handle the symbolic differentiation of external operators.
                # This bit returns:
                #
                #   `\sum_{i} dNdOi(..., Oi, ...; DOi(u)[v], ..., v*)`
                #
                # where we differentate wrt u, Oi is the i-th operand,
                # N(..., Oi, ...; ..., v*) an ExternalOperator and v the
                # direction (Argument). dNdOi(..., Oi, ...; DOi(u)[v])
                # is an ExternalOperator representing the
                # Gateaux-derivative of N. For example:
                #  -> From N(u) = u**2, we get `dNdu(u; uhat, v*) = 2 * u * uhat`.
                new_args = N.argument_slots() + (df,)
                extop = N._ufl_expr_reconstruct_(
                    *N.ufl_operands, derivatives=derivatives, argument_slots=new_args
                )
            elif df == 0:
                extop = ZeroBaseForm(N.arguments())
            else:
                raise NotImplementedError(
                    "Frechet derivative of external operators need to be provided!"
                )
            result += (extop,)
        return sum(result)  # type: ignore


class DerivativeRuleDispatcher(DAGTraverser):
    """Dispatch a derivative rule."""

    def __init__(
        self,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache)
        # Record the operations delayed to the derivative expansion phase:
        # Example: dN(u)/du where `N` is a BaseFormOperator and `u` a Coefficient
        self.pending_operations = ()
        # Create DAGTraverser caches.
        self._dag_traverser_cache: dict[
            tuple[type, Expr] | tuple[type, Expr, Expr, Expr] | tuple[type, Expr, Expr, Expr, Expr],
            DAGTraverser,
        ] = {}

    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    @process.register(Expr)
    @process.register(BaseForm)  # type: ignore
    def _(self, o: Expr | BaseForm) -> Expr | BaseForm:
        """Apply to expr and base form."""
        return self.reuse_if_untouched(o)

    @process.register(Terminal)
    def _(self, o: Terminal) -> Terminal:
        """Apply to a terminal."""
        return o

    @process.register(Derivative)
    def _(self, o: Derivative) -> Expr:
        """Apply to a derivative."""
        raise NotImplementedError(f"Missing derivative handler for {type(o).__name__}.")

    @process.register(Grad)
    @DAGTraverser.postorder
    def _(self, o: Grad, f: Expr | BaseForm) -> Expr:
        """Apply to a grad."""
        gdim = o.ufl_shape[-1]
        key = (GradRuleset, gdim)
        dag_traverser = self._dag_traverser_cache.setdefault(key, GradRuleset(gdim))
        return dag_traverser(f)  # type: ignore

    @process.register(ReferenceGrad)
    @DAGTraverser.postorder
    def _(self, o: ReferenceGrad, f: Expr | BaseForm) -> Expr | BaseForm:
        """Apply to a reference_grad."""
        tdim = o.ufl_shape[-1]
        key = (ReferenceGradRuleset, tdim)
        dag_traverser = self._dag_traverser_cache.setdefault(key, ReferenceGradRuleset(tdim))
        return dag_traverser(f)  # type: ignore

    @process.register(VariableDerivative)
    @DAGTraverser.postorder_only_children([0])
    def _(self, o: Expr, f: Expr | BaseForm) -> Expr | BaseForm:
        """Apply to a variable_derivative."""
        _, op = o.ufl_operands
        key = (VariableRuleset, op)
        dag_traverser = self._dag_traverser_cache.setdefault(key, VariableRuleset(op))
        return dag_traverser(f)  # type: ignore

    @process.register(CoefficientDerivative)
    @DAGTraverser.postorder_only_children([0])
    def _(self, o: CoefficientDerivative, f: Expr | BaseForm) -> Expr | BaseForm:
        """Apply to a coefficient_derivative."""
        _, w, v, cd = o.ufl_operands
        key = (GateauxDerivativeRuleset, w, v, cd)
        # We need to go through the dag first to record the pending
        # operations
        dag_traverser = self._dag_traverser_cache.setdefault(
            key,
            GateauxDerivativeRuleset(w, v, cd),  # type: ignore
        )
        # If f has been seen by the traverser, it immediately returns
        # the cached value.
        mapped_expr = dag_traverser(f)  # type: ignore
        # Need to account for pending operations that have been stored
        # in other integrands
        self.pending_operations += dag_traverser.pending_operations  # type: ignore
        return mapped_expr

    @process.register(BaseFormOperatorDerivative)
    @DAGTraverser.postorder_only_children([0])
    def _(self, o: BaseFormOperatorDerivative, f: Expr | BaseForm) -> Expr | BaseForm:
        """Apply to a base_form_operator_derivative."""
        _, w, v, cd = o.ufl_operands
        if isinstance(f, ZeroBaseForm):
            (arg,) = v.ufl_operands  # type: ignore
            arguments = f.arguments()
            # derivative(F, u, du) with `du` a Coefficient
            # is equivalent to taking the action of the derivative.
            # In that case, we don't add arguments to `ZeroBaseForm`.
            if isinstance(arg, BaseArgument):
                arguments += (arg,)
            return ZeroBaseForm(arguments)
        # Need a BaseFormOperatorDerivativeRuleset object
        # for each outer_base_form_op (= f).
        key = (BaseFormOperatorDerivativeRuleset, w, v, cd, f)
        # We need to go through the dag first to record the pending operations
        dag_traverser = self._dag_traverser_cache.setdefault(
            key,  # type: ignore
            BaseFormOperatorDerivativeRuleset(w, v, cd, f),  # type: ignore
        )
        # If f has been seen by the traverser, it immediately returns
        # the cached value.
        mapped_expr = dag_traverser(f)  # type: ignore
        mapped_f = dag_traverser._process_coefficient(f)  # type: ignore
        if mapped_f != 0:
            # If dN/dN needs to return an Argument in N space
            # with N a BaseFormOperator.
            return mapped_f
        # Need to account for pending operations that have been stored in other integrands
        self.pending_operations += dag_traverser.pending_operations  # type: ignore
        return mapped_expr

    @process.register(CoordinateDerivative)
    @DAGTraverser.postorder_only_children([0])
    def _(self, o: Expr, f: Expr | BaseForm) -> CoordinateDerivative:
        """Apply to a coordinate_derivative."""
        _, o1, o2, o3 = o.ufl_operands
        return CoordinateDerivative(f, o1, o2, o3)

    @process.register(BaseFormCoordinateDerivative)
    @DAGTraverser.postorder_only_children([0])
    def _(self, o: Expr, f: Expr | BaseForm) -> BaseFormCoordinateDerivative:
        """Apply to a base_form_coordinate_derivative."""
        _, o1, o2, o3 = o.ufl_operands
        return BaseFormCoordinateDerivative(f, o1, o2, o3)

    @process.register(Indexed)
    @DAGTraverser.postorder
    def _(self, o: Indexed, Ap: Expr, ii: MultiIndex) -> Expr | BaseForm:
        """Apply to an indexed."""
        # Reuse if untouched
        if Ap is o.ufl_operands[0]:
            return o
        r = len(Ap.ufl_shape) - len(ii)
        if r:
            kk = indices(r)
            op = Indexed(Ap, MultiIndex(ii.indices() + kk))
            op = as_tensor(op, kk)
        else:
            op = Indexed(Ap, ii)
        return op


class BaseFormOperatorDerivativeRecorder:
    """A derivative recorded for a base form operator."""

    def __init__(self, var, **kwargs):
        """Initialise."""
        base_form_ops = kwargs.pop("base_form_ops", ())

        if kwargs.keys() != {"arguments", "coefficient_derivatives"}:
            raise ValueError(
                "Only `arguments` and `coefficient_derivatives` are "
                "allowed as derivative arguments."
            )

        self.var = var
        self.der_kwargs = kwargs
        self.base_form_ops = base_form_ops

    def __len__(self):
        """Get the length."""
        return len(self.base_form_ops)

    def __bool__(self):
        """Convert to a bool."""
        return bool(self.base_form_ops)

    def __add__(self, other):
        """Add."""
        if isinstance(other, list | tuple):
            base_form_ops = self.base_form_ops + other
        elif isinstance(other, BaseFormOperatorDerivativeRecorder):
            if self.der_kwargs != other.der_kwargs:
                raise ValueError(
                    f"Derivative arguments must match when summing {type(self).__name__} objects."
                )
            base_form_ops = self.base_form_ops + other.base_form_ops
        else:
            raise NotImplementedError(
                f"Sum of {type(self)} and {type(other)} objects is not supported."
            )

        return BaseFormOperatorDerivativeRecorder(
            self.var, base_form_ops=base_form_ops, **self.der_kwargs
        )

    def __radd__(self, other):
        """Add."""
        # Recording order doesn't matter as collected
        # `BaseFormOperator`s are sorted later on.
        return self.__add__(other)

    def __iadd__(self, other):
        """Add."""
        if isinstance(other, list | tuple):
            self.base_form_ops += other
        elif isinstance(other, BaseFormOperatorDerivativeRecorder):
            self.base_form_ops += other.base_form_ops
        else:
            raise NotImplementedError
        return self


def apply_derivatives(expression):
    """Apply derivatives to an expression.

    Args:
        expression: A Form, an Expr or a BaseFormOperator to be differentiated

    Returns:
        A differentiated expression
    """
    # Notation: Let `var` be the thing we are differentating with respect to.

    dag_traverser = DerivativeRuleDispatcher()

    # If we hit a base form operator (bfo), then if `var` is:
    #    - a BaseFormOperator → Return `d(expression)/dw` where `w` is
    #      the coefficient produced by the bfo `var`.
    #    - else → Record the bfo on the DAGTraverser object and returns
    #    - 0.
    # Example:
    #    → If derivative(F(u, N(u); v), u) was taken the following line would compute `∂F/∂u`.
    dexpression_dvar = map_integrands(dag_traverser, expression)
    if (
        isinstance(expression, BaseForm)
        and isinstance(dexpression_dvar, int)
        and dexpression_dvar == 0
    ):
        # The arguments got lost, just keep an empty Form
        dexpression_dvar = Form([])

    # Get the recorded delayed operations
    pending_operations = dag_traverser.pending_operations
    if not pending_operations:
        return dexpression_dvar

    # Don't take into account empty Forms
    if isinstance(dexpression_dvar, Form) and dexpression_dvar.empty():
        dexpression_dvar = []
    else:
        dexpression_dvar = [dexpression_dvar]

    # Retrieve the base form operators, var, and the argument and
    # coefficient_derivatives for `derivative`
    var = pending_operations.var
    base_form_ops = pending_operations.base_form_ops
    der_kwargs = pending_operations.der_kwargs
    for N in sorted(set(base_form_ops), key=lambda x: x.count()):
        # -- Replace dexpr/dvar by dexpr/dN -- #
        # We don't use `apply_derivatives` since the differentiation is
        # done via `\partial` and not `d`.
        dexpr_dN = map_integrands(
            dag_traverser, replace_derivative_nodes(expression, {var.ufl_operands[0]: N})
        )
        # Don't take into account empty Forms
        if isinstance(dexpr_dN, Form) and dexpr_dN.empty():
            continue

        # -- Add the BaseFormOperatorDerivative node -- #
        (var_arg,) = der_kwargs["arguments"].ufl_operands
        cd = der_kwargs["coefficient_derivatives"]
        # Not always the case since `derivative`'s syntax enables one to
        # use a Coefficient as the Gateaux direction
        if isinstance(var_arg, BaseArgument):
            # Construct the argument number based on the
            # BaseFormOperator arguments instead of naively using
            # `var_arg`. This is critical when BaseFormOperators are
            # used inside 0-forms.
            #
            # Example: F = 0.5 * u** 2 * dx + 0.5 * N(u; v*)** 2 * dx
            #    -> dFdu[vhat] = <u, vhat> + Action(<N(u; v*), v0>, dNdu(u; v1, v*))
            # with `vhat` a 0-numbered argument, and where `v1` and
            # `vhat` have the same function space but a different
            # number. Here, applying `vhat` (`var_arg`) naively would
            # result in `dNdu(u; vhat, v*)`, i.e. the 2-forms `dNdu`
            # would have two 0-numbered arguments. Instead we increment
            # the argument number of `vhat` to form `v1`.
            var_arg = type(var_arg)(
                var_arg.ufl_function_space(), number=len(N.arguments()), part=var_arg.part()
            )
        dN_dvar = apply_derivatives(BaseFormOperatorDerivative(N, var, ExprList(var_arg), cd))
        # -- Sum the Action: dF/du = ∂F/∂u + \sum_{i=1,...} Action(∂F/∂Ni, dNi/du) -- #
        # In this case: Action <=> ufl.action since `dN_var` has 2 arguments.
        # We use Action to handle the trivial case `dN_dvar` = 0.
        dexpression_dvar.append(Action(dexpr_dN, dN_dvar))
    return sum(dexpression_dvar)


class CoordinateDerivativeRuleset(GenericDerivativeRuleset):
    """Apply AFD (Automatic Functional Differentiation) to expression.

    Implements rules for the Gateaux derivative D_w[v](...) defined as
    D_w[v](e) = d/dtau e(w+tau v)|tau=0
    where 'e' is a ufl form after pullback and w is a SpatialCoordinate.
    """

    def __init__(
        self,
        coefficients: ExprList,
        arguments: ExprList,
        coefficient_derivatives: ExprMapping,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(
            (), compress=compress, visited_cache=visited_cache, result_cache=result_cache
        )
        # Type checking
        if not isinstance(coefficients, ExprList):
            raise ValueError("Expecting a ExprList of coefficients.")
        if not isinstance(arguments, ExprList):
            raise ValueError("Expecting a ExprList of arguments.")
        if not isinstance(coefficient_derivatives, ExprMapping):
            raise ValueError("Expecting a coefficient-coefficient ExprMapping.")
        # The coefficient(s) to differentiate w.r.t. and the
        # argument(s) s.t. D_w[v](e) = d/dtau e(w+tau v)|tau=0
        self._w = coefficients.ufl_operands
        self._v = arguments.ufl_operands
        self._w2v = {w: v for w, v in zip(self._w, self._v)}
        # Build more convenient dict {f: df/dw} for each coefficient f
        # where df/dw is nonzero
        cd = coefficient_derivatives.ufl_operands
        self._cd = {cd[2 * i]: cd[2 * i + 1] for i in range(len(cd) // 2)}

    # Work around singledispatchmethod inheritance issue;
    # see https://bugs.python.org/issue36457.
    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    @process.register(GeometricQuantity)
    def _(self, o: Expr) -> Expr:
        # Explicitly defining dg/dw == 0
        return self.independent_terminal(o)

    @process.register(Argument)
    def _(self, o: Expr) -> Expr:
        # Explicitly defining da/dw == 0
        return self.independent_terminal(o)

    @process.register(Coefficient)
    def _(self, o: Expr) -> Expr:
        """Differentiate a coefficient."""
        raise NotImplementedError(
            "CoordinateDerivative of coefficient in physical space is not implemented."
        )

    @process.register(Grad)
    def _(self, o: Expr) -> Expr:
        """Differentiate a grad."""
        raise NotImplementedError("CoordinateDerivative grad in physical space is not implemented.")

    @process.register(SpatialCoordinate)
    def _(self, o: Expr) -> Expr:
        """Differentiate a spatial_coordinate."""
        do = self._w2v.get(o)  # type: ignore
        # d x /d x => Argument(x.function_space())
        if do is not None:
            return do
        else:
            raise NotImplementedError(
                "CoordinateDerivative found a SpatialCoordinate that is different "
                "from the one being differentiated."
            )

    @process.register(ReferenceValue)
    def _(self, o: Expr) -> Expr:
        """Differentiate a reference_value."""
        do = self._cd.get(o)  # type: ignore
        if do is not None:
            return do
        else:
            return self.independent_terminal(o)

    @process.register(ReferenceGrad)
    def _(self, g: Expr) -> Expr:
        """Differentiate a reference_grad."""
        # d (grad_X(...(x)) / dx => grad_X(...(Argument(x.function_space()))
        o = g
        ngrads = 0
        while isinstance(o, ReferenceGrad):
            (o,) = o.ufl_operands
            ngrads += 1
        if not (isinstance(o, SpatialCoordinate) or isinstance(o.ufl_operands[0], FormArgument)):
            raise ValueError(f"Expecting gradient of a FormArgument, not {ufl_err_str(o)}")

        def apply_grads(f):
            for i in range(ngrads):
                f = ReferenceGrad(f)
            return f

        # Find o among all w without any indexing, which makes this
        # easy
        for w, v in zip(self._w, self._v):
            if (
                o == w
                and isinstance(v, ReferenceValue)
                and isinstance(v.ufl_operands[0], FormArgument)
            ):
                # Case: d/dt [w + t v]
                return apply_grads(v)
        return self.independent_terminal(o)

    @process.register(Jacobian)
    def _(self, o: Expr) -> Expr:
        """Differentiate a jacobian."""
        # d (grad_X(x))/d x => grad_X(Argument(x.function_space())
        for w, v in zip(self._w, self._v):
            if extract_unique_domain(o) == extract_unique_domain(w) and isinstance(
                v.ufl_operands[0],  # type: ignore
                FormArgument,
            ):
                return ReferenceGrad(v)
        return self.independent_terminal(o)


class CoordinateDerivativeRuleDispatcher(DAGTraverser):
    """Dispatcher."""

    def __init__(
        self,
        compress: bool | None = True,
        visited_cache: dict[tuple, Expr] | None = None,
        result_cache: dict[Expr, Expr] | None = None,
    ) -> None:
        """Initialise."""
        super().__init__(compress=compress, visited_cache=visited_cache, result_cache=result_cache)
        self._dag_traverser_cache: dict[tuple[type, Expr, Expr, Expr], DAGTraverser] = {}

    @singledispatchmethod
    def process(self, o: Expr) -> Expr:
        """Process ``o``.

        Args:
            o: `Expr` to be processed.

        Returns:
            Processed object.

        """
        return super().process(o)

    @process.register(Expr)
    @process.register(BaseForm)  # type: ignore
    def _(self, o: Expr | BaseForm) -> Expr | BaseForm:
        """Apply to expr and base form."""
        return self.reuse_if_untouched(o)

    @process.register(Terminal)
    def _(self, o: Expr) -> Expr:
        """Apply to a terminal."""
        return o

    @process.register(Derivative)
    def _(self, o: Expr) -> Expr:
        """Apply to a derivative."""
        raise NotImplementedError(f"Missing derivative handler for {type(o).__name__}.")

    @process.register(Grad)
    def _(self, o: Expr) -> Expr:
        """Apply to a grad."""
        return o

    @process.register(ReferenceGrad)
    def _(self, o: Expr) -> Expr:
        """Apply to a reference_grad."""
        return o

    @process.register(CoefficientDerivative)
    def _(self, o: Expr) -> Expr:
        """Apply to a coefficient_derivative."""
        return o

    @process.register(CoordinateDerivative)
    @DAGTraverser.postorder_only_children([0])
    def _(self, o: Expr, f: Expr) -> Expr:
        """Apply to a coordinate_derivative."""
        from ufl.algorithms import extract_unique_elements

        for space in extract_unique_elements(o):
            if isinstance(space.pullback, CustomPullback):
                raise NotImplementedError(
                    "CoordinateDerivative is not supported for elements with custom pull back."
                )
        _, w, v, cd = o.ufl_operands
        key = (CoordinateDerivativeRuleset, w, v, cd)
        dag_traverser = self._dag_traverser_cache.setdefault(
            key,
            CoordinateDerivativeRuleset(w, v, cd),  # type: ignore
        )
        return dag_traverser(f)


def apply_coordinate_derivatives(expression):
    """Apply coordinate derivatives to an expression."""
    dag_traverser = CoordinateDerivativeRuleDispatcher()
    return map_integrands(dag_traverser, expression)

FEniCS / ufl / 17557791029

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous