"""Single entry point for training surrogate models against a CIF obs operator.
Usage::
trainer = SurrogateTrainer(obs_operator, backend="torch") # or "jax"
H = trainer.wrap_operator()
model, losses = trainer.train(model, sample_fn, n_steps=5000)
trainer.export("surrogate.onnx", model, input_shape=(1, N_STATE))
Backends are loaded lazily: importing this module never triggers an
ImportError even when PyTorch or JAX are absent.
"""
from __future__ import annotations
from typing import Any, Callable, List, Optional, Tuple
[docs]
class SurrogateTrainer:
"""Facade over surrogate backends; dispatches to PyTorch or JAX."""
def __init__(self, obs_operator: Any, backend: str = "torch") -> None:
self.obs_operator = obs_operator
self.backend_name = backend
self._backend: Optional[Any] = None
# ------------------------------------------------------------------
def _load_backend(self) -> Any:
if self._backend is not None:
return self._backend
if self.backend_name == "torch":
from .backends.torch_backend import TorchBackend
self._backend = TorchBackend(self.obs_operator)
elif self.backend_name == "jax":
from .backends.jax_backend import JaxBackend
self._backend = JaxBackend(self.obs_operator)
else:
raise ValueError(
f"Unknown backend {self.backend_name!r}. Choose 'torch' or 'jax'."
)
return self._backend
# ------------------------------------------------------------------
[docs]
def build_default_model(
self,
n_state: int,
n_obs: int,
hidden_size: int = 256,
n_hidden_layers: int = 2,
) -> Any:
"""Return a default untrained model for the active backend.
Uses ``torch.nn.Sequential`` for ``"torch"`` and a Flax
``linen.Module`` for ``"jax"``. Only the active backend's framework
is imported.
"""
return self._load_backend().build_default_model(
n_state, n_obs, hidden_size, n_hidden_layers
)
# ------------------------------------------------------------------
[docs]
def wrap_operator(self, run_mode: str = "fwd",
reload_results: bool = False) -> Callable:
"""Return a framework-native differentiable callable for the CIF operator.
Args:
run_mode: ``"fwd"`` wraps the full forward; ``"tl"`` wraps the
tangent-linear operator. The backward always uses the adjoint.
reload_results: Forwarded to every operator call.
"""
return self._load_backend().wrap_operator(
run_mode=run_mode, reload_results=reload_results
)
[docs]
def train(
self,
model: Any,
sample_fn: Callable[[int], Any],
n_steps: int = 5000,
run_mode: str = "fwd",
reload_results: bool = False,
adjoint_consistency: bool = False,
lambda_adj: float = 0.1,
) -> Tuple[Any, List[float]]:
"""Train *model* against the live CIF operator.
Args:
model: Framework-native model (``nn.Module`` or Flax module).
sample_fn: ``callable(batch_size) -> tensor | ndarray`` drawing
state-space samples; shape ``(N_state,)`` per sample.
n_steps: Number of gradient steps.
reload_results: If True, reloads already-computed CIF simulations
from disk instead of re-running them.
adjoint_consistency: If True, adds a VJP-agreement loss term
(Mode 3); otherwise Mode 2 (online teacher only).
lambda_adj: Weight of the adjoint consistency loss.
Returns:
``(trained_model, losses)``
"""
return self._load_backend().train(
model, sample_fn, n_steps,
run_mode=run_mode,
reload_results=reload_results,
adjoint_consistency=adjoint_consistency,
lambda_adj=lambda_adj,
)
[docs]
def export(self, path: str, model: Any, input_shape: Tuple[int, ...]) -> None:
"""Export *model* to ONNX at *path*."""
self._load_backend().export_onnx(path, model, input_shape)
