ValueError: Missing differentiation handler for type FEMExternalOperator.

[swiesheier]

@a-latyshev

I am trying to use FEMExternalOperator inside a nonlinear hyperelastic variational formulation. Even with a dummy linear energy for the operator, the Jacobian construction fails with:

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

This happens already during ufl.algorithms.expand_derivatives(…) or inside replace_external_operators(J_ext).
This looks like a missing derivative handler in this package. I thick the MWE below clearly follows the examples presented in the documentation.

Environment

dolfinx: 0.9.0
python: 3.12
dolfinx_external_operator: from tag v0.9.0

Minimal reproducible example

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
    J_ext = ufl.derivative(Res_u, u, du)

    Res_replaced, Res_ops = replace_external_operators(Res_u)
    J_replaced, J_ops     = replace_external_operators(J_ext) # <-- Error here

    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()

Full Traceback

    J_replaced, J_ops     = replace_external_operators(J_ext) # <-- Error here
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/calculate/anisotropic-fracture-soft-composites/.local_fenicsx_pkgs/dolfinx_external_operator/external_operator.py", line 238, 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 209, 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?

Question

What is the correct workflow to obtain the Jacobian for a nonlinear problem?

Or is there a missing derivative rule for FEMExternalOperator in v0.9.0 that needs to be added?

Thanks a lot!

[a-latyshev]

@swiesheier Thanks for the MWE!

This looks like a missing derivative handler in this package. I thick the MWE below clearly follows the examples presented in the documentation.

Not exactly. In the external operators tutorials, we used external operators to define the stress field directly without the use of energy density, because it looks a bit redundant: if one defines the energy density as an external operator, then extra memory is allocated to store its values in all quadrature points globally. At the same time, the stress tensor is used to define the variational problem, not the energy density. Since the user must explicitly provide how the stress tensor is evaluated (via external functions) in any case, there is no motivation for deriving the stress from the energy. If you can give an example, where it's really necessary, I'd be glad to see one.

I suggest using a workaround provided by another user of external operators: https://fenicsproject.discourse.group/t/dolfinx-external-operators-replace-form-returns-a-formsum-instead-of-form/18282/6?u=a-latyshev, where they represent the derivative of psi as a sum of two terms, F1, that can be derived with ufl.diff and F2, which depends on the external operator and is explicitly derived by the user

spring_energy = FEMExternalOperator(r,function_space=W_EO)
spring_energy.external_function = spring_external
F2 = ufl.dot(spring_energy, dr) * dx

method = 0 # 1: directional derivative of energy function, 0: differentiation of the energy function before integration

if method==0:
    dF = ufl.diff(Psi, state)
    F1 = ufl.dot(dF, dstate) * dx  #ufl.derivative(Psi,state,dstate)
    
elif method==1:
    Psi *= dx   # integrating over the domain
    F1 = ufl.algorithms.expand_derivatives(ufl.derivative(Psi,state,state))

F = F1 + F2 

It's worth mentioning that there is a bug in v0.9.0 regarding the detection of external operators in the sum of forms. We discussed this in the discourse above. Normally, it's fixed on nightly already (see PR #25), but I will make a new release v0.10.0 with the fix.

[swiesheier]

@a-latyshev

Not exactly. In the external operators tutorials, we used external operators to define the stress field directly without the use of energy density

True, in the external operator tutorials, the value of the operator was used to define the residual, not the operator's derivatives. Just to clarify: is the error message above expected because the package does not support derivatives of external operators inside the residual formulation?

At the same time, the stress tensor is used to define the variational problem, not the energy density. Since the user must explicitly provide how the stress tensor is evaluated (via external functions) in any case, there is no motivation for deriving the stress from the energy. If you can give an example, where it's really necessary, I'd be glad to see one.

In hyperelasticity, the stress is typically obtained as the derivative of the strain energy density with respect to a chosen strain measure. The resulting stress is then used to define the residual. This is also how the hyperelasticity tutorial from @jorgensd implements it. For reference, here are the relevant snippets:

P = ufl.diff(psi, F) # Psi would be the external operator
residual = (
    ufl.inner(ufl.grad(v), P) * dx - ufl.inner(v, B) * dx - ufl.inner(v, T) * ds(2)
)

And thanks also for mentioning the bug in v0.9.0 concerning the detection of external operators in sums of forms. Good to know.

As a separate question: I am planning to upgrade to a newer DOLFINx release (e.g. v0.10.0.post0). Is dolfinx-external-operator v0.9.0 still compatible with these newer versions?

[a-latyshev]

External operators don't support ufl.diff.

In hyperelasticity, the stress is typically obtained as the derivative of the strain energy density with respect to a chosen strain measure.

That's true, but there is no reason to define P via the symbolic AD ufl.diff in your case because it's defined numerically anyway. If your psi is not expressible via UFL, then you cannot apply UFL's AD to compute P.

Why not define P directly as an external operator?

def make_P_external(slope: float):

    def P_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 dP_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 P_external(derivatives):
        """
        Dispatcher required by FEMExternalOperator.

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

    return P_external

quadrature_degree = 2
cell_name = domain.topology.cell_name()

Qe = basix.ufl.quadrature_element(domain.topology.cell_name(),
                                    degree=3, 
                                    value_shape=())  # scalar output # is it for P as well?
Q = fem.functionspace(domain, Qe)

P = FEMExternalOperator(I1_bar, function_space=Q, external_function=P_impl)
residual = (
    ufl.inner(ufl.grad(v), P) * dx - ufl.inner(v, B) * dx - ufl.inner(v, T) * ds(2)
)

This avoids allocation of memory for psi and its functional space. The variational problem does not depend on psi.

For v0.10.0, you can use this dev branch https://github.com/a-latyshev/dolfinx-external-operator/tree/v0.10.0; it will be released soon.

[swiesheier]

External operators don't support

Good to know that, thanks for confirming it!

Why not define P directly as an external operator?

Yes, that would be one option. Psi typically has the form Psi = Psi_1(I1) + Psi_2(I2) + Psi_3(I3). Using the chain rule, we can write the stress as
P = ufl.diff(I1, F) * dPsi_1/dI1 + ufl.diff(I2, F) * dPsi_2/dI2 + ufl.diff(I3, F) * dPsi_3/dI3, where dPsi_j / dIj are three FEMExternalOperators. This has the advantage that we only need the derivatives of the black-box operators (in my case, the Psis are b-splines) since UFL keeps control of the kinematics (invariants). Does that make sense, or would you instead try to use as few FEMExternalOperators as possible?

[a-latyshev]

You don't need to to do this component wise.

Unfortunetly, I cannot address your issue right now, I will come back later.

I suggest you reading carefully the first part of the von Mises tutorial: https://a-latyshev.github.io/dolfinx-external-operator/demo/demo_plasticity_von_mises.html#problem-formulation. There is no much difference: you take a Gateaux derivative of your energy and there you get your P.

In UFL this may look like the following snippet

psi = FEMExternalOperator(I_expr, function_space=Q, name="psi")
E = psi * dx 
F = ufl.derivative(E, u, ufl.TestFunction(V))
J = ufl.derivative(F, u, ufl.TrialFunction(V))

which is more elegant indeed.

Actually, this should work with external operators as well but it seems that the bug with form of sums still persists, unfortunetly. I will come back to this later to enable this formulation.

For now, what I suggest is to derive manually F.

[swiesheier]

Thank you for the suggestion. I agree that deriving the Piola stress as the first variation of the energy is the cleanest approach.

In my application, however, I also perform Finite Element Model Updating and need to assemble the derivative of the residual (weak form) with respect to each material parameter. Using the implicit function theorem, I later solve linear systems with the Jacoabian of the residual as matrix. For that reason, the external operator must also provide derivatives of the energy with respect to the material parameters. I have prepared a minimal working example (MWE) below that uses two material parameters for illustration.

I have two questions regarding this setup:

1. Representation of material parameters

In the MWE, the material parameters (logically constants) are represented as quadrature-space functions with constant values. Is this the only viable option if I want to use the external-operator package for parameter sensitivities?
In my real problem I typically have 10–20 material parameters, so I am wondering:
Is there a recommended pattern for organizing multiple parameter operands efficiently?

For example, would you suggest packing parameters into a single ufl.as_vector, or should they remain passed separately to the constuctor of FEMExternalOperator similar as below?

2. Segfault during sensitivity assembly

The MWE below follows the tutorial pattern (replace external operators, evaluate operands, etc.), but it intermittently segfaults during assembly of the sensitivity forms.
From my understanding I am using the API as intended (replacement of external operators, etc), so I would like to know:
Do you see any incorrect usage in how the external operator or parameter derivatives are handled?

Sometimes, the program runs correctly without segfault, sometimes the message below gets printed.
Error Message

Sensitivity w.r.t. offset shape: (162,)
[0]PETSC ERROR: ------------------------------------------------------------------------
[0]PETSC ERROR: Caught signal number 11 SEGV: Segmentation Violation, probably memory access out of range
[0]PETSC ERROR: Try option -start_in_debugger or -on_error_attach_debugger
[0]PETSC ERROR: or see https://petsc.org/release/faq/#valgrind and https://petsc.org/release/faq/
[0]PETSC ERROR: configure using --with-debugging=yes, recompile, link, and run 
[0]PETSC ERROR: to get more information on the crash.
[0]PETSC ERROR: Run with -malloc_debug to check if memory corruption is causing the crash.
Abort(59) on node 0 (rank 0 in comm 0): application called MPI_Abort(MPI_COMM_WORLD, 59) - process 0

Minimal example two material parameters

import numpy as np
from petsc4py import PETSc
import dolfinx
from dolfinx import *
import basix
import ufl

from petsc4py import PETSc

import dolfinx.fem as femx
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.nls.petsc import NewtonSolver

default_scalar_type = PETSc.ScalarType

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

import dolfinx
import importlib.metadata as im




def make_linear_iso_energy_external():
    """
    External operator representing dPsi/dI1 for a very simple
    isochoric energy Psi(I1_bar; slope, offset).

    Here we take:
        Psi(I1_bar; slope, offset) = offset + slope * (I1_bar - 3)
    so that
        dPsi/dI1_bar = slope
    which depends on two material parameters: slope and offset.

    The external operator itself is dPsi/dI1_bar, so:
        op(I1_bar, slope, offset) = slope

    The derivatives are:
        d/dI1_bar ( dPsi/dI1_bar ) = 0
        d/dslope ( dPsi/dI1_bar ) = 1
        d/doffset( dPsi/dI1_bar ) = 0
    """

    def psi_external(derivatives):
        # derivatives is a tuple (k_I1, k_slope, k_offset)
        if derivatives == (0, 0, 0):
            # base operator: op(I1_bar, slope, offset) = slope
            return lambda I1_vals, slope_vals, offset_vals: slope_vals.reshape(-1)

        elif derivatives == (1, 0, 0):
            # derivative w.r.t. I1_bar: 0
            return lambda I1_vals, slope_vals, offset_vals: np.zeros_like(I1_vals).reshape(-1)

        elif derivatives == (0, 1, 0):
            # derivative w.r.t. slope: 1
            return lambda I1_vals, slope_vals, offset_vals: np.ones_like(slope_vals).reshape(-1)

        elif derivatives == (0, 0, 1):
            # derivative w.r.t. offset: 0
            return lambda I1_vals, slope_vals, offset_vals: np.zeros_like(offset_vals).reshape(-1)

        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)

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

    # pull on right boundary: u = (0.1, 0) on x = 1
    def right_boundary(x):
        return np.isclose(x[0], 1.0)

    left_dofs = fem.locate_dofs_geometrical(V, left_boundary)
    bc_left = fem.dirichletbc(default_scalar_type((0.0, 0.0)), left_dofs, V)

    right_dofs = fem.locate_dofs_geometrical(V, right_boundary)
    bc_right = fem.dirichletbc(default_scalar_type((0.1, 0.0)), right_dofs, V)
    bcs = [bc_left, bc_right]

    # Measure (quadrature of degree 3)
    quadrature_degree = 3
    dx = ufl.Measure(
        "dx",
        domain=domain,
        metadata={"quadrature_scheme": "default", "quadrature_degree": quadrature_degree},
    )

    # Kinematics
    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

    # Quadrature space for external operator and parameters
    Qe = basix.ufl.quadrature_element(
        domain.topology.cell_name(),
        degree=quadrature_degree,
        value_shape=(),
    )
    Q = fem.functionspace(domain, Qe)

    # --- Material parameters as quadrature fields (currently constant, but spatially resolved) ---

    # Slope parameter
    slope = fem.Function(Q)
    slope.x.array[:] = 2.0  # constant in space, but stored in quadrature field
    slope.x.scatter_forward()

    # Offset parameter
    offset = fem.Function(Q)
    offset.x.array[:] = 0.5  # constant in space, but stored in quadrature field
    offset.x.scatter_forward()

    # External operator for dPsi/dI1_bar depending on two parameters
    dPsi_dI1 = FEMExternalOperator(I1_bar, slope, offset, function_space=Q)
    dPsi_dI1.external_function = make_linear_iso_energy_external()

    # Simple Piola stress: isochoric part via external operator + volumetric penalty
    P_iso = ufl.diff(I1_bar, F) * dPsi_dI1
    P_vol = ufl.diff(J, F) * (10.0 * (J - 1.0))  # volumetric energy: 5*(J - 1)^2
    P = P_iso + P_vol

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

    # ----------------------------------------------------------------------
    # Parameter sensitivities: dRes/ d(slope) and dRes / d(offset)
    # ----------------------------------------------------------------------
    # Sensitivity w.r.t. slope field
    J_slope = ufl.derivative(Res_u, slope)
    J_slope_expanded = ufl.algorithms.expand_derivatives(J_slope)
    J_slope_replaced, J_slope_ops = replace_external_operators(J_slope_expanded)
    J_slope_form = femx.form(J_slope_replaced)

    # Sensitivity w.r.t. offset field
    J_offset = ufl.derivative(Res_u, offset)
    J_offset_expanded = ufl.algorithms.expand_derivatives(J_offset)
    J_offset_replaced, J_offset_ops = replace_external_operators(J_offset_expanded)
    J_offset_form = femx.form(J_offset_replaced)

    # Slope sensitivity
    operands_slope = evaluate_operands(J_slope_ops)
    evaluate_external_operators(J_slope_ops, operands_slope)

    b_slope = femx.petsc.assemble_vector(J_slope_form)
    b_slope.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
    print("Sensitivity w.r.t. slope  shape:", b_slope.array.shape)

    # Offset sensitivity
    operands_offset = evaluate_operands(J_offset_ops)
    evaluate_external_operators(J_offset_ops, operands_offset)

    b_offset = femx.petsc.assemble_vector(J_offset_form)
    b_offset.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
    print("Sensitivity w.r.t. offset shape:", b_offset.array.shape)
    

    # Explicit PETSc clean-up to avoid destructor-order issues
    b_slope.destroy()
    b_offset.destroy()


if __name__ == "__main__":
    main()