Source code for pycif.plugins.modes.training.execute
"""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 :class:`OnnxSurrogateOperator` and return it
as a drop-in replacement for the CIF obsoperator.
"""
import copy
import os
import numpy as np
from logging import info
from ....utils.check.errclass import CifError
# --------------------------------------------------------------------------- #
# Helper #
# --------------------------------------------------------------------------- #
def _make_sample_fn(xb: np.ndarray, std: np.ndarray, backend: str):
"""Return a callable that draws x ~ N(xb, std) as a backend tensor."""
rng = np.random.default_rng()
def sample_fn(_: int):
x = xb + std * rng.standard_normal(len(xb))
if backend == "torch":
import torch
return torch.tensor(x, dtype=torch.float32)
else:
import jax.numpy as jnp
return jnp.array(x, dtype=jnp.float32)
return sample_fn
# --------------------------------------------------------------------------- #
# Entry point #
# --------------------------------------------------------------------------- #
[docs]
def execute(self, **kwargs):
from .trainer import SurrogateTrainer
# ------------------------------------------------------------------ #
# 1. Wrap CIF operator #
# ------------------------------------------------------------------ #
cif_op = self.obsoperator
# ------------------------------------------------------------------ #
# 2. Read state / obs dimensions #
# ------------------------------------------------------------------ #
controlvect = self.controlvect
xb = np.asarray(controlvect.xb, dtype=np.float64)
n_state = self.controlvect.dim
n_obs = self.obsvect.dim
info(f"Training mode: state dimension={n_state}, obs dimension={n_obs}")
# ------------------------------------------------------------------ #
# 3. Build default model (framework chosen by backend) #
# ------------------------------------------------------------------ #
trainer = SurrogateTrainer(cif_op, backend=self.backend)
model = trainer.build_default_model(
n_state, n_obs,
hidden_size=self.hidden_size,
n_hidden_layers=self.n_hidden_layers,
)
info(
f" model: {n_state} -> [{self.hidden_size}] x {self.n_hidden_layers} -> {n_obs}")
# ------------------------------------------------------------------ #
# 4. Sample function from prior #
# ------------------------------------------------------------------ #
std = np.asarray(getattr(controlvect, "std", np.ones(n_state)), dtype=np.float64)
sample_fn = _make_sample_fn(xb, std, self.backend)
# ------------------------------------------------------------------ #
# 5. Train #
# ------------------------------------------------------------------ #
info(f"Training surrogate: backend={self.backend!r}, n_steps={self.n_steps}, "
f"adjoint_consistency={self.adjoint_consistency}")
model, losses = trainer.train(
model, sample_fn,
n_steps=self.n_steps,
run_mode=self.run_mode,
reload_results=self.reload_results,
adjoint_consistency=self.adjoint_consistency,
lambda_adj=self.lambda_adj,
)
info(f" Training complete. Final loss: {losses[-1]:.4e}")
# ------------------------------------------------------------------ #
# 6. Export to ONNX #
# ------------------------------------------------------------------ #
onnx_path = os.path.join(self.workdir, self.onnx_file)
trainer.export(onnx_path, model, input_shape=(1, n_state))
info(f" Surrogate exported to {onnx_path}")
