Decomposition Solver

sindae.algorithms.decomp

The decomposition approach alternates between an inner NLP solve (with the network embedded as a Grey-Box Model, solved by POUNCE by default) and an outer Adam update whose gradient comes from KKT implicit differentiation of the inner solution. It supports MPI parallelism across trajectories.

DecompConfig — the Adam / KKT / slack hyperparameters.
train_decomp — runs the training loop and returns (model, trained_mlp, history).
build_decomp_model — builds the inner NLP (rarely called directly).

Usage¶

from sindae import extract_instance_data
from sindae.algorithms.decomp.train import DecompConfig, train_decomp

cfg = DecompConfig(n_steps=300, lr=5e-3, init_slack_coef=1e1, param_reg_coef=1e-3)

trained_m, mlp, history = train_decomp(
    problem, mlp, cfg,
    data=smoother_data,              # normalization statistics
    smoother_model=smoother_m,       # reuse the discretized smoother
    solver_options={'tol': 1e-6, 'max_iter': 300},   # nlp_solver='pounce' by default
)
trained_data = extract_instance_data(problem, trained_m)
# history['data_fit_history'], history['grad_norm_history'], ... for diagnostics

To parallelize across trajectories, pass an MPI communicator (one trajectory batch per rank) and launch with mpirun:

from mpi4py import MPI
trained_m, mlp, history = train_decomp(
    problem, mlp, cfg, data=smoother_data, mpi_comm=MPI.COMM_WORLD,
)

API reference¶

`DecompConfig`¶

class DecompConfig(
    n_steps: int = 100,
    lr: float = 0.01,
    grad_clip_norm: float = np.inf,
    init_slack_coef: float = 100.0,
    slack_scale: float = 2.0,
    slack_update_interval: int = np.inf,
    max_slack_coef: float = 1000.0,
    mu_target: float = 1e-10,
    param_reg_coef: float = 0.0,
    subsample_frac: float = 1.0,
    patience: int = 0,
    slack_tol: float = 1e-06,
    lr_schedule: Optional[Callable[[int], float]] = None,
)

Hyperparameters for the decomposition (Adam + KKT gradient) training loop.

Fields

n_steps (int, default 100)
lr (float, default 0.01)
grad_clip_norm (float, default np.inf)
init_slack_coef (float, default 100.0)
slack_scale (float, default 2.0)
slack_update_interval (int, default np.inf)
max_slack_coef (float, default 1000.0)
mu_target (float, default 1e-10)
param_reg_coef (float, default 0.0)
subsample_frac (float, default 1.0)
patience (int, default 0)
slack_tol (float, default 1e-06)
lr_schedule (Optional[Callable[[int], float]], default None)

`train_decomp`¶

train_decomp(
    problem: ProblemDefinition,
    mlp: SimpleMLP,
    cfg: DecompConfig,
    data: InstanceData,
    smoother_model: Optional[pyo.ConcreteModel] = None,
    mpi_comm = None,
    solver_options: Optional[dict] = None,
    nlp_solver: str = 'pounce',
    linear_solver: str = 'feral',
    unfix_io: bool = True,
) -> Tuple[pyo.ConcreteModel, SimpleMLP, dict]

Train a neural network via the decomposition (GBM + KKT gradient) approach.

Pretraining is NOT handled here — call pretrain_mlp from sindae.algorithms.pretrain before this function if desired.

Parameters

problem (ProblemDefinition)
mlp (SimpleMLP)
cfg (DecompConfig)
data (InstanceData) — Provides normalization statistics. Typically extracted from the solved smoother via extract_instance_data(problem, smoother_model).
smoother_model (Optional[pyo.ConcreteModel], default None) — When provided, reused as the decomp NLP base (no rebuild / re-discretisation); IPOPT warm-starts from the smoother solution.
mpi_comm (default None)
solver_options (Optional[dict], default None) — Options passed to the NLP backend, e.g. {'max_iter': 200, 'tol': 1e-6}.
nlp_solver (str, default 'pounce') — NLP solver for the inner grey-box solve. Must be grey-box-capable.
linear_solver (str, default 'feral') — KKT/linear solver for the decomposition gradient back-solve.
unfix_io (bool, default True) — Unfix the NN input/output variables in the decomposition model. Set False for partially observed problems.

Returns

trained_m (pyo.ConcreteModel) — The solved decomposition NLP (the rank-local model under MPI). Holds the final training iterate’s trajectory; pass to extract_instance_data to recover states/outputs. The trajectory reflects the last solve, which may differ slightly from the returned best-weights mlp; for a trajectory strictly consistent with mlp use solve_inference.
mlp (SimpleMLP (trained; rank-0 parameters are authoritative))
history (dict with keys obj_history, data_fit_history,) — grad_norm_history, diag_history, pouncetiming_history

`build_decomp_model`¶

build_decomp_model(
    problem: ProblemDefinition,
    mlp: SimpleMLP,
    traj_indices: List[int],
    data: InstanceData,
    slack_coef: float = 1.0,
    smoother_model: Optional[pyo.ConcreteModel] = None,
    unfix_io: bool = True,
) -> Tuple[pyo.ConcreteModel, NNGreyBoxModel]

Build a multi-trajectory NLP with NNGreyBoxModel for the decomposition approach.