Source code for pycif.plugins.modes.training.backends.base

"""Abstract base class for surrogate backends.

Every concrete backend must implement four methods:
    build_default_model – construct a default MLP in the backend's framework
    wrap_operator       – return a differentiable callable (framework-native)
    train               – run the training loop
    export_onnx         – serialise the trained model to ONNX
"""
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Callable, List, Tuple




class SurrogateBackend(ABC):
    """Abstract base for PyTorch / JAX surrogate backends."""

    def __init__(self, obs_operator: Any, run_mode: str = "fwd") -> None:
        self.obs_operator = obs_operator
        self.run_mode = run_mode



    @abstractmethod
    def build_default_model(
        self,
        n_state: int,
        n_obs: int,
        hidden_size: int,
        n_hidden_layers: int,
    ) -> Any:
        """Return a default untrained model for this backend's framework.

        The model must accept a 1-D state vector of length *n_state* and
        return a 1-D observation vector of length *n_obs*.
        """




    @abstractmethod
    def wrap_operator(self, run_mode: str = "fwd",
                      reload_results: bool = False) -> Callable:
        """Return a differentiable callable wrapping the CIF operator.

        Args:
            run_mode: ``"fwd"`` calls ``obs_operator.forward``; ``"tl"`` calls
                ``obs_operator.tangent_linear``.  The backward pass always
                delegates to ``obs_operator.adjoint`` regardless of *run_mode*.
            reload_results: Forwarded to every operator call; if ``True``,
                already-computed CIF simulations are reloaded from disk.
        """




    @abstractmethod
    def train(
        self,
        model: Any,
        sample_fn: Callable,
        n_steps: int,
        *,
        run_mode: str,
        reload_results: bool,
        adjoint_consistency: bool,
        lambda_adj: float,
    ) -> Tuple[Any, List[float]]:
        """Run the training loop and return ``(trained_model, losses)``."""




    @abstractmethod
    def export_onnx(self, path: str, model: Any, input_shape: Tuple[int, ...]) -> None:
        """Export *model* to ONNX at *path*."""