Differentiation and framework adapters

View guide source

Use the environment and eqiora-source checkout from Get started. The inverse problems textbook explains how a change in a parameter changes an observation. Here we work directly with the solved field and its derivatives.

Prepare a Poisson problem

This example uses the unit square, a sinusoidal source, and a constant boundary value. Run the following blocks in order, from the folder containing .venv and eqiora-source:

python

from pathlib import Path

import eqiora
import numpy as np

graph = eqiora.geometry.GeometryGraph()
rectangle = graph.rectangle(x_bounds=(0.0, 1.0), y_bounds=(0.0, 1.0))
boundaries = ("x_lower", "x_upper", "y_lower", "y_upper")
geometry = graph.build(
    rectangle,
    named_topology={
        "square": rectangle.region,
        **dict(zip(boundaries, rectangle.boundaries, strict=True)),
    },
)
mesh = eqiora.meshing.generate(eqiora.meshing.resolve(
    geometry, eqiora.meshing.CartesianMesher(cells=(8, 8)),
))
model = eqiora.compile(
    path=Path("eqiora-source/examples/inverse-poisson.eqi"),
    entry="InversePoisson",
    geometry=geometry,
    bindings={
        "square": geometry.selection("square"),
        **{
            name: (geometry.selection(name), geometry.selection("square"))
            for name in boundaries
        },
        "diffusion": 1.0,
        "wave_number": np.pi,
        "source_scale": 2.0 * np.pi**2,
        "boundary_offset": 0.0,
    },
)
plan = eqiora.resolve(
    model,
    mesh=mesh,
    spatial=eqiora.fem.Q1(),
    solve=eqiora.solve.Linear(
        algorithm=eqiora.solve.LinearSolver.BiConjugateGradientStabilized,
        preconditioner=eqiora.solve.Preconditioner.Identity,
        reduction=eqiora.solve.Reduction.Reproducible,
        provider=eqiora.solve.SolverProvider.reference(),
        relative_tolerance=1.0e-10,
        absolute_tolerance=1.0e-12,
        maximum_iterations=10_000,
    ),
)

Open eqiora-source/examples/inverse-poisson.eqi to inspect the equations and boundary conditions. source_scale multiplies the spatial source pattern; diffusion controls the diffusion coefficient.

Evaluate a point and its derivatives

Use eqiora.diff to select the parameters and output field to differentiate:

python

import numpy as np

program = eqiora.diff.compile(
    plan,
    inputs=(model.parameter("source_scale"),),
    outputs=(plan.capability.fields[0],),
)

evaluation = program.evaluate((1.5,))
primal = evaluation.primal()
jvp = evaluation.jvp((1.0,))
vjp = evaluation.vjp(
    (np.ones(program.outputs[0].shape, dtype=np.float64),)
)

primal.output[0] contains the solved field values. jvp.tangent[0] contains their directional change for a unit increase in source_scale. vjp.input_cotangent[0] contains the derivative of the sum of field values, because the output cotangent above is all ones. JVP means Jacobian–vector product; VJP is the reverse product, useful for a scalar loss with many input parameters.

Each evaluation keeps its own input point and linearization, so evaluating a new point changes neither the Model and Plan nor the primal, JVP, or VJP of an earlier evaluation. Parameters you did not select keep their model values.

Point values, selected outputs, tangents and input cotangents are ordered tuples of read-only NumPy leaves. Each leaf keeps its native real/complex dtype and shape; a scalar has shape (). Use program.inputs and program.outputs for the exact native signatures, including mathematical types, units and spatial coefficient associations. For example, outputs=(model.field("z"), model.observable("loss")) returns field, loss = evaluation.primal().output without concatenating their values. Total JVP inputs follow program.inputs; VJP seeds follow program.outputs. Their results use output and input leaves respectively. Complex covectors pair as Re(sum(conj(c) * t)); base signatures do not assign primal units to dual values. Dense map projections and residual/unknown partial actions still use real coordinates.

Static spatial differentiation uses Cartesian Q1 FEM or unconstrained TPFA FVM, dimensional system scaling and the Plan’s exact host-serial solver policy. Real problems support reference iterative solvers and Faer SparseLu with Fast reduction. Complex execution requires the selected provider’s typed complex implementation; the reference provider supplies that implementation. Primal, JVP and VJP retain the selected provider and its controls. Select an ordered tuple of complete Fields and reduced Observables. Selecting a subset retains the complete coupled state and residual in both JVP and VJP; output leaves follow the requested order. Real and complex components retain their mathematical shape; for example, a two-channel complex Field has output shape (coefficient_count, 2) and dtype complex128. Complex-linear state equations may depend nonholomorphically on Parameters, such as a conjugated diffusion coefficient. The primal uses the complex solver; JVP and VJP use its real differential and the pairing above. Mixed spatial/Parameter coefficient derivatives are supported for admitted scalar arithmetic and pointwise component expressions. Their shared component IR evaluates two input directions, including mixed derivatives of those inputs. Affine-map coefficient composition and arbitrary-order derivatives remain outside this path. Conjugated state equations remain outside this static path. Transient Plans use the Run-bound trajectory boundary described below.

Spatial integral outputs require explicit quadrature, including integrals used by a composed Observable:

python

program = eqiora.diff.compile(
    plan, inputs=(model.parameter("p"),),
    outputs=(model.observable("loss"), plan.fields[0]), quadrature_points=2,
)

The generated points and weights belong to the immutable Program identity. Primal values, output partials, total JVP/VJP and saved evaluation replay use the same rules. Derivatives include direct Parameter dependence and prescribed boundary values as well as the implicit change in the solved Field. The point count specifies a numerical rule, not an accuracy guarantee. Smooth reduced outputs are supported within the admitted static Cartesian solve; continuous point outputs and general spatial objective differentiation remain outside this path.

Unselected Parameters and prescribed values independent of the selected inputs remain fixed. For example, sqrt(q) + p*x with fixed q=0 has derivative x with respect to p. Fixed operands still require valid primal values. Activity follows source dependencies, so a zero numerical tangent does not make an input fixed: sqrt(p*p) at p=0 still rejects differentiation. Finite affine and nonlinear Plans select ordered Fields and Observables through the same accepted-point API. Array leaves use the CPU transport described in Execution, diagnostics, and arrays.

Spatial trajectories with a fixed initial State

For an admitted spatial backward-Euler Plan, pass its initial physical State and one ordinary Run schedule to eqiora.diff.compile. Select Parameters through inputs and complete Fields through outputs:

python

program = eqiora.diff.compile(
    plan, inputs=(model.parameter("p"),), outputs=(plan.fields[0],),
    state=initial_state, steps=3, output_steps=(1, 3),
)
point = program.evaluate((4.0,))
values = point.primal().output[0]
tangent = point.jvp((1.0,)).tangent[0]

Use until_s and output_times_s instead for a physical-time schedule; supply exactly one schedule family. No previously solved canonical Run is required. Each selected output has shape (sample_count, coefficient_count, *value_shape), including a leading length-one time axis when only one output time is selected. program.outputs[i].sample_times_s gives those physical times. Parameter inputs have no time axis. Mapping adds its axes before these output axes without changing the Field’s intrinsic mathematical type or coefficient entities.

The Program holds the initial physical State fixed and differentiates through all accepted steps. Intermediate steps contribute to JVP and VJP even when they are absent from the requested outputs. With steps=1, output_steps=(1,), this is the fixed-predecessor step derivative; it does not differentiate the Run that produced the supplied State. Real and complex linear/Newton products reuse the ordinary native execution and accepted derivative owners. Moving geometry, eliminated kinematics, temporal Observables and other time methods remain outside this boundary.

Use the existing evaluate, jvp, vjp and map APIs, or bind the Program with eqiora.jax.bind / eqiora.torch.bind within each adapter’s supported value types and first-order transformations. result.evidence.steps retains every accepted step in execution order, independently of the output schedule. Each step exposes its residual acceptance, primal solve summary and optional derivative solve.

Recorded evaluation collections bind the exact Run request. Reopening checks all candidate steps against the original equations and preceding physical history without solving. It validates mathematical acceptance and report consistency, not the authenticity of a producer’s iteration history. Ordinary Run interpolation retention is unchanged.

Finite nonlinear points and partial actions

A strict-interior finite nonlinear Plan from the constraints guide works with the same API. Pass its initial State and a scalar instantaneous Observable:

python

program = eqiora.diff.compile(
    plan, inputs=(model.parameter("p"),),
    outputs=(model.observable("output"),), state=state,
)
point = program.evaluate((4.0,))
zero = np.array([0.0], dtype=np.float64)
one = np.array([1.0], dtype=np.float64)
partial = point.output_partial_jvp(zero, one)
reduced = point.jvp((1.0,))

For R(w,p)=w*w-p, positive w, and O(w,p)=w+p, this point has w=2 and O=6. Holding w fixed gives partial=1; solving the implicit sensitivity gives reduced.tangent[0]=1.25. At p=9 the corresponding total is 7/6. residual_jvp(dw, dp) applies R_w dw + R_p dp, while output_partial_jvp(dw, dp) applies O_w dw + O_p dp. Their VJP counterparts return (unknown_cotangent, parameter_cotangent) without an implicit solve. Unknown coordinates follow point.accepted_unknowns; Parameter coordinates follow the ordered inputs selection.

As above, each evaluation keeps its own solution and partial derivatives. The seed, the strict-interior conditions, and the residual and regularity checks are applied again at every point. Stationary evaluations have one evidence entry: result.evidence.steps[0].primal_solve is a NonlinearSolveSummary, even when no Newton update was needed, and that entry’s derivative_solve describes the linear sensitivity solve, normal or transposed. Independent ordered maps reuse the same immutable seed at each point. This does not track a solution branch globally or provide higher implicit derivatives.

Stationary nonlinear spatial points

A stationary real Q1 Newton Plan uses the same state= argument. Supply a complete spatial State.initial with physical Field values. Each Parameter point starts from its free seed coordinates and applies that point’s prescribed boundary values. The initial State remains immutable. Affine spatial Programs consume no initial State.

For -div((1+p*u)*grad(u))=0 on [0,1], with endpoint values 1 and r, q(u)=u+p*u²/2 is affine in position. On two equal Q1 cells the midpoint satisfies q(u_mid)=(q(1)+q(r))/2. At p=0, r=2, u_mid=1.5 and du_mid/dp=1/8; at p=1, r=2, they are sqrt(6.5)-1 and 1-2.5/sqrt(6.5). These sensitivities include the constitutive dependence and the implicit change in the solution. Selecting r also differentiates the boundary lift and the prescribed output value.

The Newton solve retains its accepted Jacobian and original residual for the same residual/output partial and total JVP/VJP APIs. Direct material Parameter products use the release’s first-partial contract; they do not require a mixed second material derivative. Validity conditions still apply at the exact operating point. Saved evaluation collections retain the exact spatial seed and recheck candidates through the original equations without solving again. This path currently covers stationary real Cartesian Q1 execution; it does not establish time/history or geometry sensitivities.

Native ordered batches

Batches need no extra package. program.map prepares a batch without solving anything; execute then solves every point in native code:

python

batch = program.map((np.array([1.5, 0.5, 1.5], dtype=np.float64),), point_shape=(3,))
print(batch.point_shape, program.outputs)
result = batch.execute()
if isinstance(result, eqiora.CompleteEvaluationMap):
    states, = result.primal()  # read-only [3, *intrinsic output shape]
    tangents, = result.jvp((np.ones(3, dtype=np.float64),)).tangent
    cotangents = result.vjp((np.ones_like(states),))
else:
    print(result.stopped_index, result.statuses, result.diagnostics)

Inputs are ordered leaves following program.inputs. Each leaf has shape point_shape + signature.shape; real and complex values retain their native dtype. Repeated points are solved separately in row-major order. point_shape=() means one point, and a zero extent returns correctly shaped output without running a solver. batch.points holds the complete input leaves, and batch.occurrence_coordinates(i) maps a flat index back to grid coordinates.

To hold some inputs at one value across the whole grid, pass them as shared inputs. For a program whose inputs are (source_scale, diffusion, boundary_offset):

python

shared_program = eqiora.diff.compile(
    plan,
    inputs=(
        model.parameter("source_scale"),
        model.parameter("diffusion"),
        model.parameter("boundary_offset"),
    ),
    outputs=(plan.capability.fields[0],),
)
batch = shared_program.map(
    (np.array([3.0, 1.0, 3.0]), np.array([0.0, 0.2, 0.0])),
    point_shape=(3,),
    shared_inputs=(model.parameter("diffusion"),),
    shared=(2.0,),
)
result = batch.execute()
if isinstance(result, eqiora.CompleteEvaluationMap):
    jvp = result.jvp(
        (np.ones(3), np.ones(3)),
        shared=(0.1,),
    )
    vjp = result.vjp(tuple(np.ones_like(leaf) for leaf in result.primal()))
    # Shared covectors sum over all point occurrences; mapped ones stay separate.
    # Ordered immutable leaves; collection axes precede each intrinsic leaf shape.
    print(vjp.shared_cotangents, vjp.mapped_cotangents)

The mapped leaves are the remaining inputs, in Program order. Shared values follow the order of shared_inputs, whose entries must be distinct and belong to the Program’s Model. An input is shared across all point axes or none, and every array must have its full shape; there is no broadcasting.

JVP and VJP reuse the linearizations computed by execute. The optional seed_shape and point_axes arguments place the point axes inside a larger grid of seeds: for point shape (2, 4), seed_shape=(3,) and point_axes=(0, 2) give (2, 3, 4). Mapped tangent and output-cotangent leaves have this grid followed by their intrinsic shape; shared tangent leaves have seed_shape followed by their intrinsic shape. JVP shape describes the grid only. The products keep the Plan, the axis metadata, and the per-point results. At most 32 point and seed axes are allowed. retained_bytes_limit and numerical_bytes_limit cap the numerical data Eqiora keeps, not the process’s peak memory.

Each leaf may be a scalar, NumPy array, Eqiora Array, or CPU DLPack tensor, with the exact signature dtype (float64 or complex128). NumPy and DLPack data must be aligned, native-endian, and C-contiguous. Inputs are copied before the GIL is released, so changing them afterwards has no effect. An Eqiora Array is rank one; to use it as a grid, take its read-only NumPy view and reshape it.

You can pass eqiora.EvaluationMapCancellation() to batch.execute and cancel it from another Python thread; it is checked between points. If a batch fails or is cancelled, member(i) returns the points completed so far, with their statuses and diagnostics, but the batch has no primal, jvp, or vjp.

PyTorch

From the same working folder, install the optional adapter. Bind the program outside the function you compile:

console

uv add "eqiora[torch]"

python

import torch
import eqiora.torch as eqtorch

torch_program = eqtorch.bind(program)
theta = (torch.tensor(1.5, dtype=torch.float64, requires_grad=True),)
state, = torch_program(theta)
state.square().sum().backward()

compiled_objective = torch.compile(
    lambda point: torch_program(point)[0].square().sum(),
    fullgraph=True,
)

The adapter requires PyTorch >=2.14,<2.15. Backpropagation uses the VJP of the solved equations.

Inputs are an ordered tuple or list of strided CPU:0 tensors, one per torch_program.inputs signature. Each tensor has that leaf’s exact native shape and float64 (real) or complex128 (complex) dtype. Outputs are an ordered tuple matching torch_program.outputs; a scalar stays scalar, and a sampled Field retains its leading time axis and coefficient/value axes. Multiple output leaves and real/complex mixtures use the same native Program and VJP. Complex gradients use PyTorch’s real pairing, $\operatorname{Re}(\overline{c}\,dy)$. The adapter does not modify inputs and returns fresh tensors. Programs are kept alive for the rest of the process, because autograd and compiled graphs can outlive the wrapper; evaluations and derivatives are not cached. PyTorch and JAX share the native registration owner. The exact initial State is part of program identity, so distinct nonlinear seeds retain their own solutions and derivatives.

The adapter composes first-order torch.func.jvp, vjp, grad, jacfwd and jacrev with nested/non-leading vmap, including shared-input gradient sums, duplicate points and empty batches. Native maps own physical evaluations; Jacobian seed directions remain separate from physical, temporal and value axes. Strided views are staged at the native boundary, while saved original tensors retain autograd version checks. A failed native member rejects the operation. Torch Jacobian transforms require real arguments; expose real/imaginary coordinates in the calling function when differentiating complex-valued leaves.

In-process torch.compile(fullgraph=True) supports ordinary first-order backward and the tested first-order functional transforms on inputs that do not already require gradients. Compiling a functional gradient with requires_grad inputs can request higher-order native products and is not supported yet. Double backward and derivatives of JVP/VJP products are rejected, not replaced by zeros. AMP, CUDA, torch.export, AOT packaging and arbitrary Tensor subclasses remain unsupported. Native map resource limits apply; batching is not a speed claim.

JAX

Install the optional JAX adapter into the same environment:

console

uv add "eqiora[jax]"

python

import jax
import jax.numpy as jnp
import eqiora.jax as eqjax

jax.config.update("jax_enable_x64", True)
jax_program = eqjax.bind(program)
theta = (jnp.asarray(1.5, dtype=jnp.float64),)
direction = (jnp.asarray(0.25, dtype=jnp.float64),)

state, = jax.jit(jax_program)(theta)
_, (tangent,) = jax.jvp(
    jax_program,
    (theta,),
    (direction,),
)
gradient = jax.grad(
    lambda point: jnp.sum(jax_program(point)[0] ** 2)
)(theta)

# Request order and repeated points are preserved by the native map.
points = (jnp.stack((theta[0] + 0.1, theta[0], theta[0] + 0.1)),)
states, = jax.jit(jax.vmap(jax_program))(points)
point_gradients = jax.jit(jax.vmap(jax.grad(
    lambda point: jnp.mean(jax_program(point)[0])
)))(points)

This adapter requires Python 3.12 or newer and exactly JAX and jaxlib 0.11.0. The primal, JVP, and VJP run as separate typed FFI calls, so compiled code never calls back into Python, and derivatives come from the solved equations rather than from differentiating solver iterations.

Each call takes an ordered tuple or list of JAX leaves and returns an ordered tuple of outputs. jax_program.inputs and jax_program.outputs are the native signatures: each leaf retains its exact intrinsic shape and float64 or complex128 dtype. Scalars are zero-dimensional arrays with an explicit dtype. Concrete leaves must reside together on one CPU device. Use vmap for mapped and shared arguments, non-leading in_axes, out_axes, and nested batches. jit, first-order JVP and VJP, vmap(grad(...)), gradients of summed or averaged losses over a batch, and jacfwd/jacrev all run through the native batch and its products. When a shared input is broadcast, its reverse contributions are summed; they are divided by the batch size only if you take an explicit mean. If any point fails, the whole operation fails with a diagnostic naming that point.

For complex leaves, the adapter follows JAX’s cotangent convention: Re(sum(cotangent * tangent)). Native Eqiora covectors instead use Re(sum(conj(cotangent) * tangent)). JAX differentiates the private real/imaginary projections around the native real-linear action, so JAX VJPs and gradients retain the standard JAX representation. Conjugate complex covectors when comparing them with native results. No holomorphic assumption is made.

The program, shapes, dtype, layout, and CPU platform are fixed; only the numerical values change between calls. Point and seed axes are kept separate, with at most 32 in total. Each FFI buffer is limited to 64 MiB, and the native batch and its products each keep at most 64 MiB of numerical data; neither limit bounds peak memory. Empty and single-point batches keep their shapes. Named collective axes, sharding, pmap, higher-order derivatives, accelerators, and export are not supported, and batching is not guaranteed to be faster than a loop.

Importing eqiora does not import PyTorch or JAX.

For finite Newton Plans, the residual products use the Plan’s normalized equations: each residual and its derivatives with respect to unknowns and Parameters are divided by the same physical residual scale. The reduced output JVP and VJP use the real inner product for complex coordinates and do not depend on that normalization.