derivkit 1.0.0__py3-none-any.whl

derivkit/forecasting/fisher_xy.py

@@ -0,0 +1,357 @@
+r"""Gaussian Fisher matrix utilities for models with uncertainty in both inputs and outputs.
+
+This module implements the X–Y Gaussian Fisher formalism, where both the measured
+inputs and outputs are noisy and may be correlated. The key idea is to account for
+uncertainty in the inputs by propagating it into an effective covariance for the
+outputs through a local linearization of the model. This allows standard Gaussian
+Fisher matrix techniques to be applied without explicitly marginalizing over the
+latent input variables.
+
+Model and covariance structure
+------------------------------
+
+The model provides a mean prediction ``mu_xy(x, theta)`` for the observed output
+``y`` as a function of inputs ``x`` and parameters ``theta``. Measurement errors
+on ``x`` and ``y`` are described by a joint Gaussian covariance
+
+.. math::
+
+   C =
+   \begin{pmatrix}
+       C_{xx} & C_{xy} \\
+       C_{xy}^{\mathsf{T}} & C_{yy}
+   \end{pmatrix}.
+
+Linearizing the model mean in the inputs around the measured values ``x_obs``,
+
+.. math::
+
+   \mu_{xy}(x, \theta) \approx \mu_{xy}(x_{\mathrm{obs}}, \theta) + T (x - x_{\mathrm{obs}}),
+
+with
+
+.. math::
+
+   T = \left.\frac{\partial \mu_{xy}}{\partial x}\right|_{(x_{\mathrm{obs}}, \theta)},
+
+yields an effective output covariance
+
+.. math::
+
+   R = C_{yy}
+       - C_{xy}^{\mathsf{T}} T^{\mathsf{T}}
+       - T C_{xy}
+       + T C_{xx} T^{\mathsf{T}}.
+
+This effective covariance replaces ``C_{yy}`` in the Gaussian likelihoods and Fisher
+matrix. The covariance blocks ``Cxx``, ``Cxy``, and ``Cyy`` are treated as fixed;
+parameter dependence enters only through the local sensitivity matrix ``T``.
+
+This formalism follows the generalized Fisher matrix treatment of
+Heavens et al. (2014), https://arxiv.org/abs/1404.2854.
+"""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable
+
+import numpy as np
+from numpy.typing import NDArray
+
+from derivkit.calculus_kit import CalculusKit
+from derivkit.forecasting.fisher_gaussian import build_gaussian_fisher_matrix
+from derivkit.utils.linalg import as_1d_data_vector, split_xy_covariance
+
+MuXY = Callable[[NDArray[np.float64], NDArray[np.float64]], NDArray[np.float64] | float]
+
+__all__ = [
+    "mu_xy_given_theta",
+    "mu_xy_given_x0",
+    "build_mu_theta_from_mu_xy",
+    "build_t_matrix",
+    "build_effective_covariance_r",
+    "effective_covariance_r_theta",
+    "build_xy_gaussian_fisher_matrix",
+]
+
+
+def mu_xy_given_theta(
+    x: NDArray[np.float64],
+    *,
+    theta: NDArray[np.float64],
+    mu_xy: MuXY,
+) -> NDArray[np.float64]:
+    """Evaluates the model predicted mean as a function of ``x`` at fixed parameters.
+
+    The input ``mu_xy`` is a callable that returns the model's mean prediction for the
+    observed quantity ``y`` given input values ``x`` and parameters ``theta``. This
+    wrapper holds ``theta`` fixed, so the resulting function depends only on ``x``.
+    This form is used when varying the inputs while keeping the parameter point
+    unchanged, such as when computing sensitivities with respect to ``x``.
+
+    Args:
+        x: Input values at which to evaluate the model.
+        theta: Parameter values to hold fixed.
+        mu_xy: Function that predicts the mean of ``y`` given ``x`` and ``theta``.
+
+    Returns:
+        The models mean prediction for ``y`` at ``(x, theta)``, returned as a single
+        1D data vector.
+    """
+    x = np.atleast_1d(np.asarray(x, dtype=np.float64))
+    theta = np.atleast_1d(np.asarray(theta, dtype=np.float64))
+    return as_1d_data_vector(mu_xy(x, theta))
+
+
+def mu_xy_given_x0(
+    theta: NDArray[np.float64],
+    *,
+    x0: NDArray[np.float64],
+    mu_xy: MuXY,
+) -> NDArray[np.float64]:
+    """Evaluates the model predicted mean as a function of ``theta`` at fixed inputs.
+
+    The input ``mu_xy`` is a callable that returns the model's mean prediction for the
+    observed quantity ``y`` given input values ``x`` and parameters ``theta``. This
+    wrapper holds the inputs fixed at a chosen reference point ``x0``, so the resulting
+    function depends only on ``theta``. This form is used when varying the parameters
+    while treating the inputs as fixed.
+
+    Args:
+        theta: Parameter values at which to evaluate the model.
+        x0: Input values to hold fixed.
+        mu_xy: Function that predicts the mean of ``y`` given ``x`` and ``theta``.
+
+    Returns:
+        The models mean prediction for ``y`` at ``(x0, theta)``, returned as a single
+        1D data vector.
+    """
+    x0 = np.atleast_1d(np.asarray(x0, dtype=np.float64))
+    theta = np.atleast_1d(np.asarray(theta, dtype=np.float64))
+    return as_1d_data_vector(mu_xy(x0, theta))
+
+
+def build_mu_theta_from_mu_xy(
+    mu_xy: MuXY,
+    *,
+    x0: NDArray[np.float64],
+) -> Callable[[NDArray[np.float64]], NDArray[np.float64]]:
+    """Constructs a mean function that depends only on the model parameters.
+
+    The input ``mu_xy`` predicts the model mean for the observed quantity given
+    input values and parameters. This helper fixes the input values at a chosen
+    reference point ``x0`` and returns a callable that depends only on the
+    parameters. The resulting function represents the model mean evaluated at
+    fixed inputs and is used when building parameter derivatives or Fisher
+    matrices.
+
+    Args:
+        mu_xy: Function that predicts the mean of the observed quantity given
+            input values and parameters.
+        x0: Input values at which the model mean is evaluated.
+
+    Returns:
+        A callable that evaluates the model mean as a function of the parameters
+        with the inputs held fixed.
+    """
+    x0 = np.atleast_1d(np.asarray(x0, dtype=np.float64))
+    return partial(mu_xy_given_x0, x0=x0, mu_xy=mu_xy)
+
+
+def build_t_matrix(
+    mu_xy: MuXY,
+    *,
+    x0: NDArray[np.float64],
+    theta: NDArray[np.float64],
+    method: str | None = None,
+    n_workers: int = 1,
+    **dk_kwargs: Any,
+) -> NDArray[np.float64]:
+    """Computes the sensitivity of the model mean to changes in the inputs.
+
+    The returned matrix describes how the model mean for the observed quantity changes
+    when the input values are perturbed, evaluated at a reference input point ``x0``
+    and parameter point ``theta``. In the X–Y Gaussian formulation, this sensitivity
+    is used to propagate input uncertainty into an effective covariance for the
+    observed quantity.
+
+    Args:
+        mu_xy: Function that predicts the mean of the observed quantity given input
+            values and parameters.
+        x0: Reference input values at which the sensitivity is evaluated.
+        theta: Parameter values at which the sensitivity is evaluated.
+        method: Optional derivative method name used by the derivative backend.
+            This option is forwarded through ``CalculusKit`` to the underlying
+            derivative engine (``DerivativeKit``).
+            If ``None``, the backend default is used.
+        n_workers: Number of workers used for derivative evaluations.
+        **dk_kwargs: Additional keyword arguments forwarded to the derivative engine.
+
+    Returns:
+        A matrix of sensitivities evaluated at ``(x0, theta)``, with one row per output
+        component and one column per input component.
+    """
+    x0 = np.atleast_1d(np.asarray(x0, dtype=np.float64))
+    theta = np.atleast_1d(np.asarray(theta, dtype=np.float64))
+
+    mu_of_x = partial(mu_xy_given_theta, theta=theta, mu_xy=mu_xy)
+
+    ckit = CalculusKit(mu_of_x, x0)
+    jac = np.asarray(
+        ckit.jacobian(method=method, n_workers=n_workers, **dk_kwargs),
+        dtype=np.float64,
+    )
+
+    if jac.ndim == 1:
+        jac = jac[None, :]
+    return jac
+
+
+def build_effective_covariance_r(
+    *,
+    cov: NDArray[np.float64],
+    x0: NDArray[np.float64],
+    t: NDArray[np.float64],
+) -> NDArray[np.float64]:
+    """Computes an effective output covariance that includes input uncertainty.
+
+    The X–Y Gaussian formulation allows both the inputs and outputs to be noisy, with
+    possible correlations between input and output errors. This function combines the
+    input covariance, output covariance, and cross-covariance with a local sensitivity
+    matrix ``t`` (describing how the model mean changes with the inputs) to produce an
+    effective covariance for the outputs. The result is the covariance used in the
+    Gaussian likelihoods and Fisher matrix after input uncertainty has been propagated
+    to the output space.
+
+    Args:
+        cov: Full covariance matrix for the stacked vector ``[x, y]``.
+        x0: Input values at which the model mean is evaluated.
+        t: Sensitivity matrix of the model mean with respect to the inputs, evaluated
+            at a chosen reference point.
+
+    Returns:
+        The effective covariance matrix for the output measurements.
+    """
+    cov = np.asarray(cov, dtype=np.float64)
+    x0 = np.atleast_1d(np.asarray(x0, dtype=np.float64))
+
+    nx = int(x0.size)
+    cxx, cxy, cyy = split_xy_covariance(cov, nx=nx)
+
+    ny = cyy.shape[0]
+    t = np.asarray(t, dtype=np.float64)
+    if t.shape != (ny, nx):
+        raise ValueError(f"t must have shape ({ny}, {nx}); got {t.shape}.")
+
+    return cyy - (cxy.T @ t.T) - (t @ cxy) + (t @ cxx @ t.T)
+
+
+def effective_covariance_r_theta(
+    theta: NDArray[np.float64],
+    *,
+    mu_xy: MuXY,
+    x0: NDArray[np.float64],
+    cov: NDArray[np.float64] | None = None,
+    method: str | None,
+    n_workers: int,
+    dk_kwargs: dict[str, Any],
+) -> NDArray[np.float64]:
+    """Evaluates the effective output covariance at a given parameter point.
+
+    The block covariances for the input and output measurements are treated as fixed.
+    The effective output covariance depends on the parameters through the local
+    sensitivity of the model mean to the inputs, evaluated at the reference inputs
+    ``x0``. This function recomputes that sensitivity at the supplied ``theta`` and
+    returns the corresponding effective covariance.
+
+    Args:
+        theta: Parameter values at which the effective covariance is evaluated.
+        mu_xy: Function that predicts the mean of the observed quantity given input
+            values and parameters.
+        x0: Reference input values used for the local sensitivity evaluation.
+        cov: Full covariance matrix for the stacked vector ``[x, y]``.
+        method: Optional derivative method name passed to the derivative engine.
+        n_workers: Number of workers used for derivative evaluations.
+        dk_kwargs: Additional keyword arguments forwarded to the derivative engine.
+
+    Returns:
+        The effective covariance matrix for the output measurements at ``theta``.
+    """
+    x0 = np.atleast_1d(np.asarray(x0, dtype=np.float64))
+    theta = np.atleast_1d(np.asarray(theta, dtype=np.float64))
+
+    t_matrix = build_t_matrix(
+        mu_xy,
+        x0=x0,
+        theta=theta,
+        method=method,
+        n_workers=n_workers,
+        **dk_kwargs,
+    )
+    return build_effective_covariance_r(cov=cov, x0=x0, t=t_matrix)
+
+
+def build_xy_gaussian_fisher_matrix(
+    *,
+    theta0: NDArray[np.float64],
+    x0: NDArray[np.float64],
+    mu_xy: MuXY,
+    cov: NDArray[np.float64],
+    method: str | None = None,
+    n_workers: int = 1,
+    rcond: float = 1e-12,
+    symmetrize_dcov: bool = True,
+    **dk_kwargs: Any,
+) -> NDArray[np.float64]:
+    """Computes a Gaussian Fisher matrix when both inputs and outputs are noisy.
+
+    This function supports the X–Y Gaussian case, where measurement uncertainty is
+    present in the inputs and in the outputs, and the two may be correlated. Input
+    uncertainty is incorporated by forming an effective covariance for the output
+    measurements using a local sensitivity of the model mean to the inputs evaluated
+    at the reference inputs ``x0``. The Fisher matrix is then constructed at the
+    parameter point ``theta0`` using the model mean evaluated at ``x0`` and the
+    effective output covariance.
+
+    Args:
+        theta0: Parameter values at which the Fisher matrix is evaluated.
+        x0: Reference input values used for the local sensitivity evaluation.
+        mu_xy: Function that predicts the mean of the observed quantity given input
+            values and parameters.
+        cov: Full covariance matrix for the stacked vector ``[x, y]``.
+        method: Optional derivative method name passed to the derivative engine.
+        n_workers: Number of workers used for derivative evaluations.
+        rcond: Cutoff used when solving linear systems involving the covariance.
+        symmetrize_dcov: Whether to symmetrize numerical covariance derivatives.
+        **dk_kwargs: Additional keyword arguments forwarded to the derivative engine.
+
+    Returns:
+        The Fisher information matrix evaluated at ``theta0``.
+    """
+    theta0 = np.atleast_1d(np.asarray(theta0, dtype=np.float64))
+    x0 = np.atleast_1d(np.asarray(x0, dtype=np.float64))
+    cov = np.asarray(cov, dtype=np.float64)
+
+    mu_theta = build_mu_theta_from_mu_xy(mu_xy, x0=x0)
+
+    r_fn = partial(
+        effective_covariance_r_theta,
+        mu_xy=mu_xy,
+        x0=x0,
+        cov=cov,
+        method=method,
+        n_workers=n_workers,
+        dk_kwargs=dict(dk_kwargs),
+    )
+
+    return build_gaussian_fisher_matrix(
+        theta0=theta0,
+        cov=r_fn,
+        function=mu_theta,
+        method=method,
+        n_workers=n_workers,
+        rcond=rcond,
+        symmetrize_dcov=symmetrize_dcov,
+        **dk_kwargs,
+    )

derivkit/forecasting/forecast_core.py

@@ -0,0 +1,313 @@
+"""Core utilities for likelihoods-based forecasts.
+
+This module provides functional helpers to
+
+- compute first-, second-, and third-order derivatives of a model with
+  respect to its parameters, and
+- build Fisher, doublet-DALI, and triplet-DALI forecast tensors from those
+  derivatives and a covariance matrix.
+
+These functions are the low-level building blocks used by higher-level
+forecasting interfaces in DerivKit. For details on the DALI expansion,
+see e.g. https://doi.org/10.1103/PhysRevD.107.103506.
+"""
+
+from typing import Any, Callable
+
+import numpy as np
+from numpy.typing import NDArray
+
+from derivkit.calculus_kit import CalculusKit
+from derivkit.utils.concurrency import normalize_workers
+from derivkit.utils.linalg import invert_covariance
+from derivkit.utils.types import ArrayLike1D, ArrayLike2D
+from derivkit.utils.validate import validate_covariance_matrix_shape
+
+__all__ = [
+    "SUPPORTED_FORECAST_ORDERS",
+    "get_forecast_tensors",
+]
+
+
+#: The supported orders of the DALI expansion.
+#:
+#: A value of 1 corresponds to the Fisher matrix.
+#: A value of 2 corresponds to the DALI doublet.
+#: A value of 3 corresponds to the DALI triplet.
+SUPPORTED_FORECAST_ORDERS = (1, 2, 3)
+
+SUPPORTED_DERIVATIVE_ORDERS = (1, 2, 3)
+
+
+def get_forecast_tensors(
+    function: Callable[[ArrayLike1D], float | NDArray[np.floating]],
+    theta0: ArrayLike1D,
+    cov: ArrayLike2D,
+    *,
+    forecast_order: int = 1,
+    method: str | None = None,
+    n_workers: int = 1,
+    **dk_kwargs: Any,
+) -> dict[int, tuple[NDArray[np.float64], ...]]:
+    """Returns a set of tensors according to the requested order of the forecast.
+
+    Args:
+        function: The scalar or vector-valued function to
+            differentiate. It should accept a list or array of parameter
+            values as input and return either a scalar or a
+            :class:`np.ndarray` of observable values.
+        theta0: The points at which the
+            derivative is evaluated. A 1D array or list of parameter values
+            matching the expected input of the function.
+        cov: The covariance matrix of the observables. Should be a square
+            matrix with shape ``(n_observables, n_observables)``, where ``n_observables``
+            is the number of observables returned by the function.
+        forecast_order: The requested order of the forecast.
+            Currently supported values and their meaning are given in
+            :data:`SUPPORTED_FORECAST_ORDERS`.
+        method: Method name or alias (e.g., ``"adaptive"``, ``"finite"``).
+            If ``None``, the :class:`derivkit.derivative_kit.DerivativeKit`
+            default (``"adaptive"``) is used.
+        n_workers: Number of workers for per-parameter parallelization/threads.
+            Default ``1`` (serial). Inner batch evaluation is kept serial to
+            avoid nested pools.
+        **dk_kwargs: Additional keyword arguments passed to
+            :class:`derivkit.derivative_kit.DerivativeKit.differentiate`.
+
+    Returns:
+        A dict mapping ``order -> tensors`` for all ``order = 1..forecast_order``.
+
+        The tensors are grouped by the *forecast order at which they first appear*:
+
+        - order 1: ``(F,)``
+        - order 2: ``(D_{(2,1)}, D_{(2,2)})``
+        - order 3: ``(T_{(3,1)}, T_{(3,2)}, T_{(3,3)})``
+
+        Here ``D_{(k,l)}`` and ``T_{(k,l)}`` denote tensors obtained by contracting
+        the ``k``-th order derivative with the ``l``-th order derivative via the
+        inverse covariance.
+
+        Each tensor axis has length ``p = len(theta0)``. Shapes are:
+
+        - ``F``: ``(p, p)``
+        - ``D_{(2,1)}``: ``(p, p, p)``
+        - ``D_{(2,2)}``: ``(p, p, p, p)``
+        - ``T_{(3,1)}``: ``(p, p, p, p)``
+        - ``T_{(3,2)}``: ``(p, p, p, p, p)``
+        - ``T_{(3,3)}``: ``(p, p, p, p, p, p)``
+
+    Raises:
+        ValueError: If ``forecast_order`` is not in :data:`SUPPORTED_FORECAST_ORDERS`.
+
+    Warns:
+        RuntimeWarning: If ``cov`` is not symmetric (proceeds as-is, no symmetrization),
+            is ill-conditioned (large condition number), or inversion
+            falls back to the pseudoinverse.
+    """
+    try:
+        forecast_order = int(forecast_order)
+    except Exception as e:
+        raise TypeError(f"forecast_order must be an int;"
+                        f" got {type(forecast_order)}.") from e
+
+    if forecast_order not in SUPPORTED_FORECAST_ORDERS:
+        raise ValueError(
+            f"forecast_order={forecast_order} is not supported. "
+            f"Supported values: {SUPPORTED_FORECAST_ORDERS}."
+        )
+
+    theta0_arr = np.asarray(theta0, dtype=np.float64).reshape(-1)
+    if theta0_arr.size == 0:
+        raise ValueError("theta0 must be non-empty 1D.")
+
+    cov_arr = validate_covariance_matrix_shape(cov)
+    n_observables = cov_arr.shape[0]
+
+    y0 = np.asarray(function(theta0_arr), dtype=float)
+    y0_flat = y0.reshape(-1)
+
+    if y0_flat.size != n_observables:
+        raise ValueError(
+            f"Expected {n_observables} observables from model "
+            f"(from cov {cov_arr.shape}), "
+            f"but got {y0_flat.size} (output shape {y0.shape})."
+        )
+
+    invcov = invert_covariance(cov_arr, warn_prefix="get_forecast_tensors")
+
+    forecast_tensors: dict[int, tuple[NDArray[np.float64], ...]] = {}
+    derivatives: dict[int, NDArray[np.float64]] = {}
+
+    contractions = {
+        1: {1: "ia,ij,jb->ab"},
+        2: {1: "iab,ij,jc->abc",
+            2: "iab,ij,jcd->abcd"},
+        3: {1: "iabc,ij,jd->abcd",
+            2: "iabc,ij,jde->abcde",
+            3: "iabc,ij,jdef->abcdef"},
+    }
+
+    for order1 in range(1, 1 + forecast_order):
+        derivatives[order1] = _get_derivatives(
+            function,
+            theta0_arr,
+            cov_arr,
+            order=order1,
+            n_workers=n_workers,
+            method=method,
+            **dk_kwargs,
+        )
+
+        tensors_at_order: list[NDArray[np.float64]] = []
+        for order2 in contractions[order1]:
+            tensors_at_order.append(
+                np.einsum(
+                    contractions[order1][order2],
+                    derivatives[order1],
+                    invcov,
+                    derivatives[order2],
+                ).astype(np.float64, copy=False)
+            )
+
+        forecast_tensors[order1] = tuple(tensors_at_order)
+
+    expected_keys = set(range(1, forecast_order + 1))
+    if set(forecast_tensors.keys()) != expected_keys:
+        raise RuntimeError(
+            f"internal error: forecast_tensors keys {sorted(forecast_tensors.keys())} "
+            f"!= expected {sorted(expected_keys)}."
+        )
+
+    return forecast_tensors
+
+
+def _get_derivatives(
+    function: Callable[[ArrayLike1D], float | NDArray[np.floating]],
+    theta0: ArrayLike1D,
+    cov: ArrayLike2D,
+    *,
+    order: int,
+    method: str | None = None,
+    n_workers: int = 1,
+    **dk_kwargs: Any,
+) -> NDArray[np.float64]:
+    """Returns derivatives of the observables of the requested order.
+
+    Args:
+        function: The scalar or vector-valued function to
+            differentiate. It should accept a list or array of parameter
+            values as input and return either a scalar or a
+            :class:`np.ndarray` of observable values.
+        theta0: The points at which the
+            derivative is evaluated. A 1D array or list of parameter values
+            matching the expected input of the function.
+        cov: The covariance matrix of the observables. Should be a square
+            matrix with shape ``(n_observables, n_observables)``, where ``n_observables``
+            is the number of observables returned by the function.
+        order: The requested order of the derivatives. The value determines
+            the order of the derivative that is returned. Currently supported
+            values are given in :data:`SUPPORTED_DERIVATIVE_ORDERS`.
+        method: Method name or alias (e.g., ``"adaptive"``, ``"finite"``). If ``None``,
+            the DerivativeKit default ("adaptive") is used.
+        n_workers: Number of workers for per-parameter parallelization
+         (threads). Default ``1`` (serial).
+        **dk_kwargs: Additional keyword arguments passed to
+            :meth:`derivkit.derivative_kit.DerivativeKit.differentiate`.
+
+    Returns:
+        Array of derivative values. For ``order == 1``, the
+        shape is ``(n_observables, n_parameters)`` (first-order derivatives).
+        For ``order == 2``, the shape is
+        ``(n_observables, n_parameters, n_parameters)`` (second-order derivatives).
+        For ``order == 3``, the shape is
+        ``(n_observables, n_parameters, n_parameters, n_parameters)`` (third-order derivatives).
+
+    Raises:
+        ValueError: An error occurred if a derivative was requested of
+            higher order than 3.
+        RuntimeError: An error occurred if a ValueError was not raised
+            after calling the function.
+    """
+    if order not in SUPPORTED_DERIVATIVE_ORDERS:
+        raise ValueError(
+            f"Requested derivative order={order} is not supported. "
+            f"Supported values: {SUPPORTED_DERIVATIVE_ORDERS}."
+        )
+
+    theta0_arr = np.atleast_1d(theta0)
+    cov_arr = np.asarray(cov, dtype=float)
+
+    n_parameters = theta0_arr.shape[0]
+    n_observables = cov_arr.shape[0]
+
+    n_workers = normalize_workers(n_workers)
+
+    def _vectorize_model_output(theta: ArrayLike1D) -> NDArray[np.float64]:
+        """Returns model output as a 1D float64 vector."""
+        y = np.asarray(function(theta), dtype=np.float64)
+        if y.ndim > 1:
+            raise TypeError(
+                "model must return a scalar or 1D vector of observables; "
+                f"got shape {y.shape}."
+            )
+        return np.atleast_1d(y)
+
+    ckit = CalculusKit(_vectorize_model_output, theta0_arr)
+
+    if order == 1:
+        j_raw = np.asarray(
+            ckit.jacobian(
+                method=method,
+                n_workers=n_workers,  # allow outer parallelism across params
+                **dk_kwargs,
+            ),
+            dtype=float,
+        )
+        if j_raw.shape == (n_observables, n_parameters):
+            return j_raw
+        else:
+            raise ValueError(
+                f"jacobian returned unexpected shape {j_raw.shape}; "
+                f"expected ({n_observables},{n_parameters})."
+            )
+
+    elif order == 2:
+        # Build Hessian tensor once (shape expected (n_observables, n_parameters, n_parameters)),
+        # then return as (n_parameters, n_parameters, n_observables) for downstream einsum.
+        h_raw = np.asarray(
+            ckit.hessian(
+                method=method,
+                n_workers=n_workers,  # allow outer parallelism across params
+                **dk_kwargs,
+            ),
+            dtype=float,
+        )
+        if h_raw.shape == (n_observables, n_parameters, n_parameters):
+            return h_raw
+        else:
+            raise ValueError(
+                f"hessian returned unexpected shape {h_raw.shape}; "
+                f"expected ({n_observables},{n_parameters},{n_parameters})."
+            )
+
+
+    elif order == 3:
+        hh_raw = np.asarray(
+            ckit.hyper_hessian(
+                method=method,
+                n_workers=n_workers,
+                **dk_kwargs,
+            ),
+            dtype=float,
+        )
+        if hh_raw.shape == (n_observables, n_parameters, n_parameters, n_parameters):
+            return hh_raw
+        else:
+            raise ValueError(
+                f"hyper_hessian returned unexpected shape {hh_raw.shape}; "
+                f"expected ({n_observables},{n_parameters},{n_parameters},{n_parameters})."
+            )
+
+
+    else:
+        raise ValueError(f"Unsupported value of {order}.")