Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

Quickstart Guide

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 generation

See 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