This page walks through a simple training run on the Leslie-Gower predator–prey example, from data generation to trained model. The same workflow applies to all built-in and custom problems. See walkthroughs in Examples Gallery for detailed explainations.
Prerequisites¶
Follow the Installation guide first. After installation, set JAX to 64-bit precision at the top of your script:
import jax
jax.config.update('jax_enable_x64', True)The short way: HybridDAE¶
The HybridDAE wrapper runs the entire pipeline (smoother, optional pretraining,
training, inference) behind a scikit-learn-style fit/predict interface:
import jax
import numpy as np
import sindae as sd
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"
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)
new_problem = sd.LeslieGowerProblem(ics=np.array([[1.2, 0.15]]), nfe=40, ncp=3)
pred = model.predict(new_problem, slack_coef=1e-5)See the HybridDAE API page for all options, including the
solver selectors (nlp_solver=, linear_solver=, solver_options=). The rest of
this page walks the same pipeline stage by stage, which is also how you get full
control over each step.
1. Define or import a problem¶
SiNDAE represents a system via a ProblemDefinition subclass. Three built-in problems are
available in sindae.example_problems:
from sindae.example_problems import LeslieGowerProblem
problem = LeslieGowerProblem(nfe=60, ncp=3) # fine grid for data generationSee Hybrid DAE Overview for an explanation of nfe/ncp and how to subclass
ProblemDefinition for your own system.
2. Construct the neural network¶
Use SimpleMLP from sindae.nn_utils:
import jax
from sindae import SimpleMLP
mlp = SimpleMLP(
in_size=problem.input_dim, # fed from get_input_vars
out_size=problem.z_dim, # produced by get_output_vars
widths=[16, 16],
activations=[jax.nn.softplus, jax.nn.softplus],
key=jax.random.PRNGKey(0),
)3. Generate synthetic training data¶
import numpy as np
from sindae import generate_data
data = generate_data(
problem,
noise_std=np.array([0.05, 0.05]), # Gaussian noise on observations
obs_every=4, # observe every 4th collocation point
seed=0,
)generate_data solves the true model with POUNCE, adds noise, and stores
problem.obs_times / problem.obs_values in-place.
4. Solve the smoother¶
The smoother fits a smooth trajectory to the noisy observations — it produces warm-start values and normalization statistics for the subsequent training step.
from sindae.algorithms.smoother import solve_smoother
from sindae.data_utils import extract_instance_data
# Switch to the coarser training grid
problem.nfe = 40
problem.ncp = 3
smoother_m = solve_smoother(problem, mlp, smooth_coef=1.0)
smoother_data = extract_instance_data(problem, smoother_m)5a. Train (simultaneous approach)¶
from sindae.algorithms.simultaneous.train import SimultaneousConfig, solve_simultaneous
cfg = SimultaneousConfig(use_gbm=False, reg_coef=1e-3)
trained_m, mlp = solve_simultaneous(
problem=problem,
mlp=mlp,
cfg=cfg,
data=smoother_data,
smoother_model=smoother_m, # warm-start from smoother
solver_options={'tol': 1e-6, 'max_iter': 1000},
)See Simultaneous Solver for a detailed explanation.
5b. Train (decomposition approach)¶
from sindae.algorithms.pretrain import PretrainConfig, pretrain_mlp
from sindae.algorithms.decomp.train import DecompConfig, train_decomp
# Optional: supervised pretrain on smoother arrays
mlp = pretrain_mlp(mlp, smoother_data, PretrainConfig(epochs=400))
cfg = DecompConfig(n_steps=300, lr=5e-3, init_slack_coef=10.0,
param_reg_coef=1e-3)
trained_m, mlp, history = train_decomp(
problem=problem, mlp=mlp, cfg=cfg,
data=smoother_data, smoother_model=smoother_m,
)Recover the trained trajectory the same way for either approach:
trained_data = extract_instance_data(problem, trained_m)For the decomposition run, history holds the training curve
(obj_history, grad_norm_history); plot it with plot_training_history(history).
See Decomposition Solver for a detailed explanation.
6. Extract results and plot¶
from sindae.data_utils import extract_instance_data
from sindae.plot_utils import plot_instance_data
trained_data = extract_instance_data(problem, trained_m)
fig, axes = plot_instance_data(
datasets=[(trained_data, 'trained', {})],
nn_input_names=['prey', 'predator'],
nn_output_names=['z'],
obs_times=problem.obs_times,
obs_values=problem.obs_values,
obs_names=['prey', 'predator'],
)
fig.savefig('result.pdf')Next steps¶
Hybrid DAE Overview — mathematical background and problem formulation
API Reference — full API reference
Examples Gallery — complete worked examples