Source code for pycif.plugins.obsoperators.onnx.obsoper

import datetime
import numpy as np
from logging import info, warning

from ....utils.check.errclass import CifIOError




def obsoper(
    self,
    controlvect,
    obsvect,
    mode,
    run_id=0,
    datei=datetime.datetime(1979, 1, 1),
    datef=datetime.datetime(2100, 1, 1),
    workdir="./",
    reload_results=False,
    check_transforms=False,
    ignore_exceptions=False,
    force_fetch_results=False,
    **kwargs,
):
    """Apply the ONNX surrogate observation operator.

    **Forward** (``mode='fwd'``):

    .. math::

        \\mathbf{y}_\\text{sim} = \\mathcal{N}(\\mathbf{x})

    Sets ``obsvect.ysim`` and returns *obsvect*.

    **Tangent-linear** (``mode='tl'``):

    .. math::

        \\delta\\mathbf{y} \\approx
        \\frac{\\mathcal{N}(\\mathbf{x}+\\varepsilon\\,\\delta\\mathbf{x})
               - \\mathcal{N}(\\mathbf{x}-\\varepsilon\\,\\delta\\mathbf{x})}
              {2\\varepsilon}

    Requires 2 forward passes (central finite differences).
    Sets both ``obsvect.ysim`` and ``obsvect.dy``; returns *obsvect*.

    **Adjoint** (``mode='adj'``):

    .. math::

        \\delta\\mathbf{x} \\approx \\mathbf{J}(\\mathbf{x})^\\top\\,\\delta\\mathbf{y}

    where :math:`\\mathbf{J}` is built column-by-column via central finite
    differences (:math:`2 n_\\text{ctl}` forward passes).  The Jacobian is
    cached and reused for repeated adjoint calls at the same
    :math:`\\mathbf{x}`.

    Args:
        self: Obs-operator plugin instance.  Must have ``_ort_operator`` set
            by a prior call to ``ini_data``.
        controlvect: Control-vector object.  ``controlvect.x`` is read in
            ``fwd`` / ``tl`` modes; ``controlvect.dx`` is read in ``tl`` mode
            and written in ``adj`` mode.
        obsvect: Observation-vector object.  ``obsvect.ysim`` is written in
            ``fwd`` / ``tl`` modes; ``obsvect.dy`` is read in ``adj`` mode
            and written in ``tl`` mode.
        mode: One of ``'fwd'``, ``'tl'``, or ``'adj'``.
        run_id: Run identifier (unused; present for interface compatibility).
        datei: Simulation start date (unused).
        datef: Simulation end date (unused).
        workdir: Working directory (unused).
        reload_results: Reload pre-computed results (unused).
        check_transforms: Validate transform consistency (unused).
        ignore_exceptions: Swallow non-fatal errors (unused).
        force_fetch_results: If ``True``, raise :exc:`CifIOError` immediately.
        **kwargs: Additional keyword arguments (ignored).

    Returns:
        *obsvect* in ``'fwd'`` and ``'tl'`` modes; *controlvect* in ``'adj'``.

    Raises:
        CifIOError: if ``force_fetch_results`` is ``True``.
    """
    if force_fetch_results:
        raise CifIOError

    H = self._ort_operator.forward
    eps = self.fd_epsilon
    x = np.asarray(controlvect.x, dtype=np.float64)

    if mode == "fwd":
        obsvect.ysim = H(x)
        # Invalidate Jacobian cache — x has changed
        self._adj_cache = None
        self._adj_cache_x = None
        return obsvect

    elif mode == "tl":
        dx = np.asarray(controlvect.dx, dtype=np.float64)
        obsvect.ysim = H(x)
        obsvect.dy = (H(x + eps * dx) - H(x - eps * dx)) / (2.0 * eps)
        return obsvect

    elif mode == "adj":
        dy = np.asarray(obsvect.dy, dtype=np.float64)
        n_state = len(x)

        # Recompute Jacobian only when x has changed
        if self._adj_cache is None or not np.array_equal(self._adj_cache_x, x):
            if n_state > 500:
                warning(
                    f"ONNX adjoint: computing {n_state}-column finite-difference "
                    f"Jacobian ({2 * n_state} forward passes). "
                    "Consider a backend that supports analytic gradients."
                )
            n_obs = len(dy)
            J = np.empty((n_obs, n_state), dtype=np.float64)
            for i in range(n_state):
                ei = np.zeros(n_state, dtype=np.float64)
                ei[i] = eps
                J[:, i] = (H(x + ei) - H(x - ei)) / (2.0 * eps)
            self._adj_cache = J
            self._adj_cache_x = x.copy()

        controlvect.dx = self._adj_cache.T @ dy
        return controlvect