Sensitivities

sensitivities

Bounded forward sensitivities of smooth deterministic fixed-mesh ODEs.

The augmented equations use the existing explicit integrator. They approximate continuous sensitivities, including caller-supplied initial-condition effects. No compiler, statistical noise model or Fisher definition is imported here.

DerivativeArtifact

Bases: Protocol

Structural compiler boundary, with no runtime compiler dependency.

Matrices have shapes (N,N), (N,Q), (M,N), (M,Q). The producer declares supported operations and must represent the same smooth RHS/measurements.

model property

Identify the derivative producer's source model.

observable_names property

Give the ordered scalar measurement coordinates.

parameter_names property

Give RHS and measurement parameter coordinates.

state_names property

Give the ordered expanded state coordinates.

evaluate(t, y, **parameters)

Evaluate native derivatives without changing coordinates or graphs.

Source code in src/op_engine/sensitivities.py

def evaluate(
    self, t: object, y: object, **parameters: object
) -> DerivativeMatrices:
    """Evaluate native derivatives without changing coordinates or graphs."""
    ...

DerivativeMatrices

Bases: Protocol

Four native matrices in declared physical state/parameter coordinates.

observable_parameters property

Return dg/dtheta with shape (M,Q).

observable_state property

Return dg/dy with shape (M,N).

parameters property

Return df/dtheta with shape (N,Q).

state property

Return df/dy with shape (N,N).

SensitivityInitial(state, parameter_derivatives) dataclass

Native state (N,) and its required parameter derivative (N,P).

Pass an explicit zero matrix for parameter-independent initial states. This avoids silently dropping initial-condition derivatives.

SensitivityOptions(solver=(lambda: RunConfig(method='rk4'))(), store_history=False, max_augmented_states=16384, max_output_bytes=64 * 1024 * 1024) dataclass

Bound fixed explicit integration and retained numerical output.

Only the final output is retained by default. History adds requested output times, not internal stages. Byte limits cover returned numeric arrays, not temporary work, AD graphs, compiler memory or the whole process.

__post_init__()

Reject unsupported numerical paths instead of silently downshifting.

Raises:

Type	Description
ValueError	Solver mode or storage bounds are unsupported.
TypeError	History is not boolean.

Source code in src/op_engine/sensitivities.py

def __post_init__(self) -> None:
    """Reject unsupported numerical paths instead of silently downshifting.

    Raises:
        ValueError: Solver mode or storage bounds are unsupported.
        TypeError: History is not boolean.
    """
    cfg = self.solver
    if not isinstance(cfg, RunConfig):
        msg = "Sensitivity solver options require a RunConfig"
        raise TypeError(msg)
    if (
        cfg.adaptive
        or not cfg.strict
        or cfg.method not in {"euler", "heun", "rk4", "dopri5"}
        or cfg.forcing_breakpoints
        or any(
            value is not None
            for value in (
                cfg.state_checks,
                cfg.nonlinear,
                cfg.jacobian,
                cfg.operators.default,
                cfg.operators.tr,
                cfg.operators.bdf2,
            )
        )
    ):
        msg = "Sensitivities require a smooth fixed explicit ODE without operators"
        raise ValueError(msg)
    if type(self.store_history) is not bool:
        msg = "Sensitivity store_history must be boolean"
        raise TypeError(msg)
    if any(
        type(bound) is not int or bound < 1
        for bound in (self.max_augmented_states, self.max_output_bytes)
    ):
        msg = "Sensitivity storage bounds must be positive integers"
        raise ValueError(msg)

SensitivitySolution(source_model, state_names, parameter_names, observable_names, times, mesh, states, parameter_derivatives, observable_derivatives, augmented_dimension, options) dataclass

Native output states and parameter derivatives with static provenance.

Array shapes are (T,N), (T,N,P) and (T,M,P). Measurement sensitivities are dg/dy @ S + dg/dtheta in the same parameter order. They do not include observation noise, prior curvature or nuisance elimination. times contains only retained outputs; mesh is the requested solver mesh.

SensitivitySystem(rhs, derivatives, initial_parameters=()) dataclass

Bind one RHS to derivatives and explicitly named initial-only parameters.

rhs accepts every derivative-artifact parameter, including measurement- only parameters. Initial-only names are appended to that ordered layout and excluded from RHS/derivative calls; their effects enter through the required initial derivative matrix. The caller is responsible for binding the same RHS as the derivative producer and for declaring a smooth deterministic ODE.

parameter_names property

Return derivative parameters followed by initial-only parameters.

__post_init__()

Reject ambiguous coordinate layouts before numerical work.

Raises:

Type	Description
ValueError	Names are empty, duplicate or overlap.
TypeError	The RHS is not callable.

Source code in src/op_engine/sensitivities.py

def __post_init__(self) -> None:
    """Reject ambiguous coordinate layouts before numerical work.

    Raises:
        ValueError: Names are empty, duplicate or overlap.
        TypeError: The RHS is not callable.
    """
    if not callable(self.rhs):
        msg = "Sensitivity RHS must be callable"
        raise TypeError(msg)
    if not isinstance(self.initial_parameters, tuple):
        msg = "Initial-only parameter names must be a tuple"
        raise TypeError(msg)
    groups = (
        self.derivatives.state_names,
        self.parameter_names,
        self.derivatives.observable_names,
    )
    if (
        not self.derivatives.state_names
        or not self.parameter_names
        or any(
            not isinstance(names, tuple)
            or any(not isinstance(name, str) or not name.strip() for name in names)
            or len(set(names)) != len(names)
            for names in groups
        )
    ):
        msg = "Sensitivity layouts require nonempty unique state/parameter names"
        raise ValueError(msg)

solve_sensitivities(system, parameters, initial, time_grid, *, options=None)

Integrate S' = (df/dy) S + df/dtheta with existing explicit kernels.

The compiler is optional: any correctly ordered derivative artifact can satisfy the protocol. Callbacks/mesh are static under JIT; native scalar bindings and both initial arrays retain their graphs. Batch independent systems through native vmap. Forcing jumps, events, history dependence, stochastic paths and operator solvers need separately qualified methods.

Returns:

Type	Description
SensitivitySolution	Native retained outputs and measured sensitivities in physical coordinates.

Raises:

Type	Description
ValueError	Layout, mesh or declared storage limits disagree.

Source code in src/op_engine/sensitivities.py

def solve_sensitivities(
    system: SensitivitySystem,
    parameters: Mapping[str, object],
    initial: SensitivityInitial,
    time_grid: NDArray[np.floating],
    *,
    options: SensitivityOptions | None = None,
) -> SensitivitySolution:
    """Integrate ``S' = (df/dy) S + df/dtheta`` with existing explicit kernels.

    The compiler is optional: any correctly ordered derivative artifact can
    satisfy the protocol. Callbacks/mesh are static under JIT; native scalar
    bindings and both initial arrays retain their graphs. Batch independent
    systems through native vmap. Forcing jumps, events, history dependence,
    stochastic paths and operator solvers need separately qualified methods.

    Returns:
        Native retained outputs and measured sensitivities in physical coordinates.

    Raises:
        ValueError: Layout, mesh or declared storage limits disagree.
    """
    opts = options or SensitivityOptions()
    bindings = _validate_inputs(system, parameters, initial)
    xp = array_namespace(initial.state)
    n, p, m = (
        len(system.derivatives.state_names),
        len(system.parameter_names),
        len(system.derivatives.observable_names),
    )
    grid = np.asarray(time_grid, dtype=np.float64)
    if (
        grid.ndim != 1
        or not grid.size
        or not np.all(np.isfinite(grid))
        or np.any(np.diff(grid) <= 0)
    ):
        msg = "Sensitivity mesh must be finite and strictly increasing"
        raise ValueError(msg)
    count = grid.size if opts.store_history else 1
    dimension = n * (1 + p)
    if (
        dimension > opts.max_augmented_states
        or count * (n + n * p + m * p) * (xp.finfo(initial.state.dtype).bits // 8)
        > opts.max_output_bytes
    ):
        msg = "Sensitivity solve exceeds its augmented/output storage limit"
        raise ValueError(msg)
    core = ModelCore(
        dimension,
        1,
        grid,
        options=ModelCoreOptions(
            dtype=np.dtype(str(initial.state.dtype).removeprefix("torch.")),
            store_history=opts.store_history,
        ),
    )
    core.set_initial_state(
        xp.reshape(
            xp.concat((
                initial.state,
                xp.reshape(initial.parameter_derivatives, (-1,)),
            )),
            (-1, 1),
        )
    )
    CoreSolver(core).run(_augmented_rhs(system, bindings), config=opts.solver)
    return _solution(system, bindings, core, opts)