pycif.plugins.modes.training — API reference

Configuration reference: training plugin

Execute function for the training mode.

Full pipeline

1. Wrap the CIF obs operator as a plain-numpy object.

2. Read state / obs dimensions from controlvect and obsvect.

3. Build a default MLP surrogate (architecture from YAML config).

4. Draw training samples from the prior (xb ± std noise).

5. Train the surrogate against the live CIF operator.

6. Export the trained model to ONNX inside workdir.

7. Reload the ONNX file as an OnnxSurrogateOperator and return it as a drop-in replacement for the CIF obsoperator.

pycif.plugins.modes.training.execute.execute(self, **kwargs)

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.

class pycif.plugins.modes.training.trainer.SurrogateTrainer(obs_operator: Any, backend: str = 'torch')

Bases: object

Facade over surrogate backends; dispatches to PyTorch or JAX.

build_default_model(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.

wrap_operator(run_mode: str = 'fwd', reload_results: bool = False) → Callable

Return a framework-native differentiable callable for the CIF operator.

Parameters:

• 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.

train(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.

Parameters:

• 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)

export(path: str, model: Any, input_shape: Tuple[int, ...]) → None

Export model to ONNX at path.

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

class pycif.plugins.modes.training.backends.base.SurrogateBackend(obs_operator: Any, run_mode: str = 'fwd')

Bases: ABC

Abstract base for PyTorch / JAX surrogate backends.

abstractmethod build_default_model(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 wrap_operator(run_mode: str = 'fwd', reload_results: bool = False) → Callable

Return a differentiable callable wrapping the CIF operator.

Parameters:

• 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 train(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 export_onnx(path: str, model: Any, input_shape: Tuple[int, ...]) → None

Export model to ONNX at path.

JAX backend: wraps the CIF obs operator via jax.custom_vjp.

The forward and backward passes call the numpy CIF operator directly. This works in JAX eager mode; for JIT compilation the callbacks would need to be wrapped in jax.pure_callback (requires knowing output shapes ahead of time — see the _jit_wrap comment below).

Training uses optax for optimisation and assumes a Flax-style model API:

params = model.init(key, x) y     = model.apply(params, x)

class pycif.plugins.modes.training.backends.jax_backend.JaxBackend(obs_operator: Any)

Bases: SurrogateBackend

JAX surrogate backend.

build_default_model(n_state: int, n_obs: int, hidden_size: int, n_hidden_layers: int) → Any

Return a Flax linen.Module MLP: n_state -> hidden -> n_obs.

wrap_operator(run_mode: str = 'fwd', reload_results: bool = False) → Callable

Return a JAX-differentiable callable wrapping the CIF operator.

Parameters:

• run_mode – "fwd" uses obs_operator.forward; "tl" uses obs_operator.tangent_linear. The backward pass always calls obs_operator.adjoint.

• reload_results – Forwarded to every operator call.

Note: not JIT-compilable as-is. For JIT support replace the direct numpy calls with jax.pure_callback; see the inline comment.

train(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 (Flax-style) against the live CIF operator.

Parameters:

• model – A Flax linen.Module (or any module with .init / .apply methods).

• sample_fn – callable(batch_size) -> jnp.ndarray returning a single state vector of shape (N_state,).

• n_steps – Number of gradient steps.

• run_mode – Forwarded to wrap_operator(); "tl" trains against the tangent-linear operator instead of the full forward.

• reload_results – Forwarded to wrap_operator().

• adjoint_consistency – Add VJP-agreement loss term (Mode 3).

• lambda_adj – Weight of the adjoint consistency loss.

Returns:

(trained_params, losses)

export_onnx(path: str, model: Any, input_shape: Tuple[int, ...]) → None

Export to ONNX is not directly supported from JAX.

Recommended conversion path:

# Step 1 – JAX -> TensorFlow SavedModel
import jax
from jax.experimental import jax2tf
import tensorflow as tf

params = ...  # trained params from train()
def predict(x):
    return model.apply(params, x)

tf_fn = tf.function(
    jax2tf.convert(predict, enable_xla=False),
    input_signature=[
        tf.TensorSpec(shape=input_shape, dtype=tf.float32)
    ],
)
tf.saved_model.save(tf_fn, "/tmp/jax_saved_model")

# Step 2 – SavedModel -> ONNX  (shell)
# python -m tf2onnx.convert \
#     --saved-model /tmp/jax_saved_model \
#     --output surrogate.onnx

References

google/jax onnx/tensorflow-onnx

PyTorch backend: wraps the CIF obs operator as a torch.autograd.Function.

Training modes

Mode 2 (online)

CIF called live as teacher; adjoint_consistency=False.

Mode 3 (adjoint)

Adds a VJP-agreement loss between the surrogate Jacobian and the CIF adjoint; adjoint_consistency=True. Uses create_graph=True so the adjoint-consistency loss is itself differentiable w.r.t. model params.

ONNX export

Uses torch.onnx.export with opset 17. The batch axis is declared dynamic so the resulting file accepts any batch size at inference time.

class pycif.plugins.modes.training.backends.torch_backend.TorchBackend(obs_operator: Any, run_mode: str = 'fwd')

Bases: SurrogateBackend

PyTorch surrogate backend.

build_default_model(n_state: int, n_obs: int, hidden_size: int, n_hidden_layers: int) → Any

Return a torch.nn.Sequential MLP: n_state -> hidden -> n_obs.

wrap_operator(run_mode: str = 'fwd', reload_results: bool = False) → Callable

Return a torch.autograd.Function wrapping the CIF operator.

Parameters:

• run_mode – "fwd" uses obs_operator.forward; "tl" uses obs_operator.tangent_linear. The backward pass always calls obs_operator.adjoint.

• reload_results – Forwarded to every operator call.

train(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.

Parameters:

• model – A torch.nn.Module.

• sample_fn – callable(batch_size) -> Tensor | ndarray returning a single state vector of shape (N_state,).

• n_steps – Number of gradient steps.

• run_mode – Forwarded to wrap_operator().

• reload_results – Forwarded to wrap_operator().

• adjoint_consistency – Add VJP-agreement loss term (Mode 3).

• lambda_adj – Weight of the adjoint consistency loss.

Returns:

(trained_model, losses)

export_onnx(path: str, model: Any, input_shape: Tuple[int, ...]) → None

Export model to ONNX via torch.onnx.export (opset 17).

The batch dimension (axis 0) is made dynamic so the ONNX file accepts any batch size at inference time.

Parameters:

• path – Destination .onnx file path.

• model – A trained torch.nn.Module.

• input_shape – Shape of a single input including batch dim, e.g. (1, N_state).

ONNX export helper.

Thin wrapper over SurrogateBackend.export_onnx() so callers can use the serialisation API without holding a reference to the backend object.

pycif.plugins.modes.training.serialization.onnx_export.export_to_onnx(backend: Any, path: str, model: Any, input_shape: Tuple[int, ...]) → None

Export model to ONNX using backend.

Parameters:

• backend – A concrete SurrogateBackend.

• path – Destination file path, e.g. "surrogate.onnx".

• model – The trained framework-native model.

• input_shape – Shape of a single input tensor including the batch dim, e.g. (1, 64).

Backwards-compatibility re-export.

OnnxSurrogateOperator has moved to pycif.plugins.obsoperators.onnx.onnx_operator.

class pycif.plugins.modes.training.serialization.onnx_operator.OnnxSurrogateOperator(onnx_path: str)

Bases: object

Drop-in replacement for a CIF obs operator backed by ONNX Runtime.

Parameters:

onnx_path – Path to the .onnx model file.

forward(x: ndarray) → ndarray

Run inference.

Parameters:

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.