Using FEMExternalOperator in elasticity: Missing differentiation handler for type FEMExternalOperator

I’m trying to use dolfinx_external_operator.FEMExternalOperator inside a nonlinear hyperelastic variational problem. My strain energy density can not conveniently be represented in UFL form (b-spline function → de-Boor algorithm needed). Even with a dummy linear external energy, the Jacobian replacement fails with:

ValueError: Missing differentiation handler for type FEMExternalOperator. Have you added a new type?

Below is a minimal complete example (no spline, no complex model).
The external operator is simply a linear energy function with constant slope (zero second derivative).
Everything else (volumetric term, kinematics, residual/Jacobian) is standard UFL.
The error happens during:

J_replaced, J_ops = replace_external_operators(J_ext)

Minimal reproducible script

from mpi4py import MPI
import numpy as np
from petsc4py import PETSc

from dolfinx import mesh, fem
import basix
import ufl

from dolfinx_external_operator import (
    FEMExternalOperator,
    replace_external_operators,
    evaluate_operands,
    evaluate_external_operators,
)


def make_linear_iso_energy_external(slope: float):
 
    def psi_value_impl(I1_vals: np.ndarray) -> np.ndarray:
        num_cells, n_qp = I1_vals.shape
        I1_flat = I1_vals.reshape(-1)
        psi_flat = slope * I1_flat
        return psi_flat

    def dpsi_dI1_impl(I1_vals: np.ndarray) -> np.ndarray:
        num_cells, n_qp = I1_vals.shape
        N_total = num_cells * n_qp
        return slope * np.ones(N_total, dtype=I1_vals.dtype)

    def d2psi_dI12_impl(I1_vals: np.ndarray) -> np.ndarray:
        num_cells, n_qp = I1_vals.shape
        N_total = num_cells * n_qp
        return np.zeros(N_total, dtype=I1_vals.dtype)

    def psi_external(derivatives):
        """
        Dispatcher required by FEMExternalOperator.

        For a single operand I1_bar, derivatives is a tuple (d/dI1_bar,).
        """
        if derivatives == (0,):
            return psi_value_impl
        elif derivatives == (1,):
            return dpsi_dI1_impl
        elif derivatives == (2,):
            return d2psi_dI12_impl
        else:
            raise NotImplementedError(f"No implementation for derivatives={derivatives}")

    return psi_external


def main():
    comm = MPI.COMM_WORLD

    # Mesh and spaces
    domain = mesh.create_unit_square(comm, 8, 8)
    V = fem.functionspace(domain, ("CG", 1, (domain.geometry.dim,)))  # vector P1

    u = fem.Function(V)
    v = ufl.TestFunction(V)
    du = ufl.TrialFunction(V)

    # Fix left boundary: u = 0 on x = 0
    def left_boundary(x):
        return np.isclose(x[0], 0.0)

    left_dofs = fem.locate_dofs_geometrical(V, left_boundary)
    bc_left = fem.dirichletbc(PETSc.ScalarType((0.0, 0.0)), left_dofs, V)
    bcs = [bc_left]

    # Quadrature space for scalar external energy psi_iso(I1_bar)
    Qe = basix.ufl.quadrature_element(domain.topology.cell_name(),
                                    degree=3,
                                    value_shape=())  # scalar output
    Q = fem.functionspace(domain, Qe)
    dx = ufl.Measure("dx", domain=domain, metadata={"quadrature_degree": 3})

    # kineamatic quantities
    d = domain.geometry.dim
    I = ufl.Identity(d)
    F = ufl.variable(I + ufl.grad(u))
    C = F.T * F
    J = ufl.det(F)

    I1 = ufl.tr(C)
    I1_bar = J**(-2.0 / 3.0) * I1

    # External operator for isochoric energy 
    slope = 2.0  # constant slope, zero second derivative

    psi_iso_ext = FEMExternalOperator(I1_bar, function_space=Q)
    psi_iso_ext.external_function = make_linear_iso_energy_external(slope)

    # Volumetric energy
    kappa = 10.0
    psi_vol = 0.5 * kappa * (J - 1.0)**2

    # Total energy density
    psi = psi_iso_ext + psi_vol

    # First Piola–Kirchhoff stress: P = d psi / d F
    P = ufl.diff(psi, F)

    # Weak form (residual)
    Res_u = ufl.inner(P, ufl.grad(v)) * dx

    # Gateaux derivative (Jacobian)
    J_ext = ufl.derivative(Res_u, u, du)

    # Replace external operators with evaluatable versions
    Res_replaced, Res_ops = replace_external_operators(Res_u)
    J_replaced, J_ops = replace_external_operators(J_ext)

    operands = evaluate_operands(Res_ops)
    _ = evaluate_external_operators(Res_ops, operands)
    _ = evaluate_external_operators(J_ops, operands)

    F_form = fem.form(Res_replaced)
    J_form = fem.form(J_replaced)

    # Nonlinear problem and solver
    problem = fem.petsc.NonlinearProblem(F_form, u, bcs=bcs, J=J_form)
    from dolfinx.nls.petsc import NewtonSolver
    solver = NewtonSolver(comm, problem)
    solver.rtol = 1e-8
    solver.atol = 1e-10
    solver.max_it = 20
    solver.error_on_nonconvergence = False
    solver.report = True
    solver.convergence_criterion = "residual"

    its, converged = solver.solve(u)


if __name__ == "__main__":
    main()

Deatailled error message

Traceback (most recent call last):
  File "/calculate/anisotropic-fracture-soft-composites/src/scripts/dummy.py", line 137, in <module>
    main()
  File "/calculate/anisotropic-fracture-soft-composites/src/scripts/dummy.py", line 113, in main
    J_replaced, J_ops = replace_external_operators(J_ext)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/calculate/anisotropic-fracture-soft-composites/.local_fenicsx_pkgs/dolfinx_external_operator/external_operator.py", line 215, in replace_external_operators
    replaced_form, ex_ops = _replace_form(form)
                            ^^^^^^^^^^^^^^^^^^^
  File "/calculate/anisotropic-fracture-soft-composites/.local_fenicsx_pkgs/dolfinx_external_operator/external_operator.py", line 186, in _replace_form
    replaced_form = ufl.algorithms.replace(form, ex_ops_map)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/replace.py", line 94, in replace
    e = expand_derivatives(e)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/ad.py", line 34, in expand_derivatives
    form = apply_derivatives(form)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/apply_derivatives.py", line 1568, in apply_derivatives
    dexpression_dvar = map_integrand_dags(rules, expression)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/map_integrands.py", line 86, in map_integrand_dags
    return map_integrands(
           ^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/map_integrands.py", line 30, in map_integrands
    map_integrands(function, itg, only_integral_type) for itg in form.integrals()
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/map_integrands.py", line 39, in map_integrands
    return itg.reconstruct(function(itg.integrand()))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/map_integrands.py", line 87, in <lambda>
    lambda expr: map_expr_dag(function, expr, compress), form, only_integral_type
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/corealg/map_dag.py", line 35, in map_expr_dag
    (result,) = map_expr_dags(
                ^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/corealg/map_dag.py", line 103, in map_expr_dags
    r = handlers[v._ufl_typecode_](v, *[vcache[u] for u in v.ufl_operands])
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/apply_derivatives.py", line 1384, in variable_derivative
    return map_expr_dag(rules, f, vcache=self.vcaches[key], rcache=self.rcaches[key])
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/corealg/map_dag.py", line 35, in map_expr_dag
    (result,) = map_expr_dags(
                ^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/corealg/map_dag.py", line 101, in map_expr_dags
    r = handlers[v._ufl_typecode_](v)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dolfinx-env/lib/python3.12/site-packages/ufl/algorithms/apply_derivatives.py", line 99, in expr
    raise ValueError(
ValueError: Missing differentiation handler for type FEMExternalOperator. Have you added a new type?

What is the correct way to build the Jacobian when the energy depends on an external operator?

Aside, is it needed to update the external operators during the nonlinear solve?

Thank you!