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

"""Zero-dependency ONNX inference wrapper.

Loads a serialised model from an ONNX file and exposes a ``forward`` method
that maps a 1-D state vector to a 1-D observation vector.  The only required
dependency is ``onnxruntime`` — no PyTorch, JAX, or other ML framework needed.

Example::

    from pycif.plugins.obsoperators.onnx.onnx_operator import OnnxSurrogateOperator
    op = OnnxSurrogateOperator("surrogate.onnx")
    y  = op.forward(x)   # x: np.ndarray float64, y: np.ndarray float64
"""
from __future__ import annotations

import numpy as np
from typing import List




class OnnxSurrogateOperator:
    """Drop-in replacement for a CIF obs operator backed by ONNX Runtime.

    Args:
        onnx_path: Path to the ``.onnx`` model file.
    """

    def __init__(self, onnx_path: str) -> None:
        try:
            import onnxruntime as ort
        except ImportError as exc:
            raise ImportError(
                "onnxruntime is required for zero-dependency inference: "
                "pip install onnxruntime"
            ) from exc

        providers: List[str] = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        self._session = ort.InferenceSession(onnx_path, providers=providers)
        self._input_name: str = self._session.get_inputs()[0].name
        self._output_name: str = self._session.get_outputs()[0].name



    def forward(self, x: np.ndarray) -> np.ndarray:
        """Run inference.

        Args:
            x: Input array of shape ``(N_state,)`` for a single sample or
               ``(batch, N_state)`` for a batch.  A 1-D input is temporarily
               promoted to batch size 1 and squeezed back on output.

        Returns:
            float64 output array — shape ``(N_obs,)`` for a 1-D input,
            ``(batch, N_obs)`` for batched input.
        """
        squeeze = x.ndim == 1
        if squeeze:
            x = x[np.newaxis, :]
        result: np.ndarray = self._session.run(
            [self._output_name],
            {self._input_name: x.astype(np.float32)},
        )[0]
        result = result.astype(np.float64)
        return result.squeeze(0) if squeeze else result