HybridDAE

sindae.hybrid_dae

The high-level fit/predict wrapper, and the primary way to use SiNDAE. One HybridDAE object runs the whole pipeline: fit(problem) solves the smoother, pretrains the network, and trains with the chosen method; predict(new_problem) embeds the trained network in a new problem and solves the inference NLP.

• HybridDAE selects the training approach (method='simultaneous' or 'decomposition') and the solver stack (nlp_solver=, linear_solver=). The network comes in as a prebuilt SimpleMLP (net=); defining it stays outside the wrapper.

• Each stage is configured with its config dataclass, the same objects the stage functions use: SmootherConfig, PretrainConfig, SimultaneousConfig or DecompConfig (matching method), and SolverConfig for the fit-time NLP solver options. Everything is validated at construction, so a typo fails before any solve. The inference solve in predict does not inherit the constructor’s solver_options; pass predict(..., solver_options=...) to tune it, so a bare predict matches a bare solve_inference call.

• Cross-cutting choices (nlp_solver=, linear_solver=, solver_options=, unfix_io=) live only on the wrapper, never inside a stage config, so no stage can silently override them.

After fit, the training solve’s termination condition is on model.termination (None for the decomposition method, whose per-step inner solves are tracked in model.history), and any non-optimal solve raises a UserWarning. Every intermediate stays reachable: the solved smoother model (smoother_model), the normalization data (smoother_data), the solved training model (training_model), its extracted trajectories (trained_data), and the trained network (net). The stage functions (solve_smoother, pretrain_mlp, solve_simultaneous, train_decomp, solve_inference) remain the low-level escape hatch when you need stage-level control.

Usage¶

import jax
import numpy as np
import sindae as sd

jax.config.update('jax_enable_x64', True)

problem = sd.LeslieGowerProblem(nfe=40, ncp=3)
sd.generate_data(problem, noise_std=np.array([0.05, 0.05]), obs_every=4)

mlp = sd.SimpleMLP(
    in_size=problem.input_dim, out_size=problem.z_dim,
    widths=[16, 16], activations=[jax.nn.softplus] * 2,
    key=jax.random.PRNGKey(0),
)

model = sd.HybridDAE(
    method="simultaneous",             # or "decomposition"
    nlp_solver="pounce",               # "ipopt" / "cyipopt" selectable
    linear_solver="feral",             # decomposition KKT solver; "ma27" / "scipy"
    net=mlp,
    train=sd.SimultaneousConfig(reg_coef=1e-3),  # DecompConfig for "decomposition"
    solver_options=sd.SolverConfig(tol=1e-6, max_iter=1000),
)
model.fit(problem)                     # smoother -> pretrain -> train
print(model.termination)               # "optimal"

# Predict under new initial conditions
new_problem = sd.LeslieGowerProblem(ics=np.array([[1.2, 0.15]]), nfe=40, ncp=3)
pred = model.predict(new_problem, slack_coef=1e-5)

mu_hat = model.net                     # the trained SimpleMLP

# Persist the trained network and its scaler, then reload
model.save("mu_hat.eqx")
reloaded = sd.HybridDAE.load("mu_hat.eqx")
pred = reloaded.predict(new_problem, slack_coef=1e-5)   # scaler restored

# Export for a foreign optimization tool
model.export("mu_hat.json")            # plain-text bundle, no extra dependencies
model.export("mu_hat.onnx")            # ONNX graph (normalized) + scaler sidecar
model.export("mu_hat.onnx", scaled=True)  # scaler baked in: raw-in, raw-out graph
net_def = model.to_omlt()              # in-memory OMLT NetworkDefinition

save writes the network weights, its architecture, and the four normalization vectors the inference stage needs, so a reloaded model can predict right away or warm-start a fresh fit from the loaded weights. The stage configs and training trajectories are not persisted, so load is for resuming or serving a trained network, not for reproducing the original solve bit for bit.

save/load round-trip back into SiNDAE; export and to_omlt are a one-way handoff to another modeling tool. Both carry the scaler, so the network is evaluated in the space it was trained in. export(path) writes a file (.json for a dependency-free bundle of weights, activations, scaler, input bounds, and the ordered input/output variable names; .onnx for the graph plus a scaler sidecar) and needs the matching extra (pip install 'sindae[onnx]'). to_omlt() returns an omlt.neuralnet.NetworkDefinition whose inputs and outputs are the raw physical variables (the normalization rides along as an OMLT OffsetScaling), ready to drop into your own optimization model with an OMLT formulation; it needs pip install 'sindae[omlt]'. The ONNX graph itself stays in normalized space by default because OMLT applies the scaler separately; pass export(path, scaled=True) to bake the scaler into the graph as affine layers instead, giving a self-contained model that maps raw physical inputs to raw physical outputs in any ONNX runtime (the sidecar’s scaling field records which contract applies, so a consumer never double-applies the scaler).

The smoother stage is configured the same way when its defaults are not right, for example smoother=sd.SmootherConfig(smooth_coef=10.0) for noisier data. Pretraining always runs; pretrain=None means PretrainConfig() (200 epochs), and pretrain=sd.PretrainConfig(epochs=0) disables it. For partially observed problems (unmeasured states with no data anchor), pass unfix_io=False.

API reference¶

HybridDAE¶

class HybridDAE(
    method: str = 'simultaneous',
    nlp_solver: str = 'pounce',
    linear_solver: str = 'feral',
    net: Optional[SimpleMLP] = None,
    smoother: Optional[SmootherConfig] = None,
    pretrain: Optional[PretrainConfig] = None,
    train: Union[SimultaneousConfig, DecompConfig, None] = None,
    solver_options: Optional[SolverConfig] = None,
    unfix_io: bool = True,
)

scikit-learn-style facade over the SiNDAE training pipeline.

fit(problem) runs smoother -> pretrain -> train; predict(new_problem) runs inference with the trained network.

Stage configuration comes in as the same config dataclasses the stage functions use; cross-cutting choices (nlp_solver, linear_solver, solver_options, unfix_io) live only here, so no stage config can silently override them.

Parameters

• method (str, default 'simultaneous') — Training approach: 'simultaneous' (single NLP; default) or 'decomposition' (Adam + KKT-gradient loop).

• nlp_solver (str, default 'pounce') — NLP solver used at every stage: 'pounce' (default), 'ipopt', 'cyipopt'.

• linear_solver (str, default 'feral') — KKT/linear solver for the decomposition gradient back-solve: 'feral' (default), 'ma27', 'scipy'. Unused by the simultaneous method.

• net (Optional[SimpleMLP], default None) — The network to train, constructed outside the wrapper (see :class:SimpleMLP). Its in_size / out_size must match problem.input_dim / problem.z_dim at fit.

• smoother (Optional[SmootherConfig], default None) — Smoother-stage hyperparameters. None uses SmootherConfig().

• pretrain (Optional[PretrainConfig], default None) — Supervised pretraining hyperparameters. None uses PretrainConfig() (200 epochs); pass PretrainConfig(epochs=0) to disable pretraining.

• train (Union[SimultaneousConfig, DecompConfig, None], default None) — Training hyperparameters; the config class must match method. None uses the method’s config defaults.

• solver_options (Optional[SolverConfig], default None) — NLP solver options for the fit-time solves (smoother and training). The inference solve in predict does not inherit these; pass predict(..., solver_options=...) to tune it.

• unfix_io (bool, default True) — Unfix the NN input/output variables in the smoother and training models (default True). Set False for partially observed problems: unmeasured states have no data anchor, and leaving their variables free makes the solves diverge.

Properties

• net — The trained network. Available after fit.

Methods

fit¶

fit(
    problem: ProblemDefinition,
    metrics: Optional[list[str]] = None,
    tee: bool = False,
) -> 'HybridDAE'

Run the training pipeline on problem and return self.

Stages: solve the smoother, extract normalization data, pretrain the network on the smoother arrays, then train with the configured method (simultaneous NLP or decomposition loop). Any non-optimal solve raises a UserWarning; the training solve’s termination condition lands on self.termination.

Parameters

• problem (ProblemDefinition) — Must carry observations (obs_times / obs_values), set directly or via :func:generate_data.

• metrics (Optional[list[str]], default None) — Metrics to print after training, comparing the fitted trajectories against problem’s observations per state variable and trajectory. Options: mse, rmse, mae. Each table adds a final N<METRIC> column holding the range-normalized metric per trajectory (each state’s metric divided by that state’s observed min-max range, then averaged over the states). The bottom mean row reports the per-state mean across trajectories; its last cell is the mean of the normalized column, a single overall goodness-of-fit scalar that for mse is the MSEP (mean squared error performance metric) of Industrial & Engineering Chemistry Research 61(25), 8658 (doi:10.1021/acs.iecr.1c04507).

• tee (bool, default False) — Stream solver output to stdout (simultaneous method only).

Returns

• self (HybridDAE) — Fitted wrapper; the trained network is self.net.

predict¶

predict(
    problem: ProblemDefinition,
    slack_coef: float = 0.0,
    solver_options: Optional[SolverConfig] = None,
    eval_metrics: Optional[list[str]] = None,
    tee: bool = False,
) -> InstanceData

Embed the trained network in problem and solve the inference NLP.

Normalization statistics are the ones the training stage consumed (self.smoother_data), so the network is evaluated in the space it was trained in. The solved model is kept on self.inference_model; a non-optimal solve raises a UserWarning.

Parameters

• problem (ProblemDefinition) — The problem to predict, e.g. the training system with new initial conditions. Observations are not required unless eval_metrics is set.

• slack_coef (float, default 0.0) — 0 (default) enforces the NN equality as a hard constraint; > 0 relaxes it with l1 slack variables (see :func:solve_inference).

• solver_options (Optional[SolverConfig], default None) — NLP solver options for this inference solve. Defaults to the solver’s own defaults (independent of the constructor’s fit-time solver_options), so a bare predict matches a bare :func:solve_inference call.

• eval_metrics (Optional[list[str]], default None) — Metrics to print, comparing the prediction against problem’s observations per state variable and trajectory. Options: mse, rmse, mae. The table has the same layout as :meth:fit’s metrics: a range-normalized N<METRIC> column and a mean row whose last cell is the mean normalized value (the MSEP for mse). Requires problem to carry observations.

• tee (bool, default False) — Stream solver output to stdout.

Returns

• (InstanceData) — Predicted trajectories at the collocation points.

save¶

save(path) -> None

Serialize the trained network and its scaler to path.

Writes a one-line JSON manifest (architecture, activation names, the four normalization vectors from smoother_data, plus method and termination) followed by the Equinox leaf arrays. Reload with :meth:HybridDAE.load.

Only the network and scaler are persisted, not the stage configs or the training trajectories, so a loaded model can predict or warm-start a fresh fit but cannot reproduce the original solve bit-for-bit.

Parameters

• path — Destination file (the parent directory must exist).

load¶

load(cls, path: str, verbose: bool = False) -> 'HybridDAE'

Reconstruct a fitted :class:HybridDAE from a :meth:save file.

The returned wrapper can predict immediately (the scaler is restored on smoother_data as a :class:NormStats) or fit again to warm-start from the loaded weights. trained_data and the stage configs are not restored (they are not persisted); fit would use the default configs.

Parameters

• cls

• path (str) — A file written by :meth:save.

• verbose (bool, default False) — Prints the loaded model information contained in the manifest.

Returns

• (HybridDAE) — A fitted wrapper carrying the loaded network and scaler.

export¶

export(path = None, format: Optional[str] = None, scaled: bool = False) -> str

Export the trained network to a file for a foreign optimization tool.

Unlike :meth:save (which round-trips back into SiNDAE), export is a one-way handoff. Two file targets, both carrying the scaler so the network is evaluated in the space it was trained in:

• 'onnx' — writes the network graph to path and a <path>.json sidecar with the scaler, input bounds, and IO contract. By default (scaled=False) the graph is in normalized space and the scaler is kept out of it, because every OMLT loader applies the scaler separately. With scaled=True the four normalization vectors are baked into the graph as affine layers, so the exported model consumes raw physical inputs and returns raw physical outputs (self-contained inference in any ONNX runtime, no sidecar arithmetic). Needs the onnx extra.

• 'json' — writes the whole bundle (weights, activations, scaler, bounds, IO contract) as plain text. Only use for very small MLPs since storing many weights may lead to large file sizes.

For an in-memory OMLT model (not a file), use :meth:to_omlt.

Parameters

• path (default None) — Output file. When format is omitted the target is inferred from the suffix (.onnx / .json).

• format (Optional[str], default None) — 'onnx' or 'json'. Required when path has no recognized suffix.

• scaled (bool, default False) — ONNX only. When True, bake the scaler into the exported graph so it maps raw physical inputs to raw physical outputs. Rejected for 'json' export (whose bundle always carries the scaler verbatim).

Returns

• (str) — The written path.

to_omlt¶

to_omlt()

Build an in-memory OMLT model of the trained network.

Returns an omlt.neuralnet.NetworkDefinition with the normalization attached as an OffsetScaling (so the OMLT block’s inputs/outputs are the raw physical variables, not normalized ones) and the data-derived input bounds as its scaled_input_bounds. Feed it to an OMLT formulation (e.g. FullSpaceSmoothNNFormulation) inside your own optimization model. Needs the omlt extra.

Returns

• (omlt.neuralnet.NetworkDefinition)