dispcraft.ml

ml

PyTorch/Lightning residual-correction MLP.

Replaces the sklearn.neural_network.MLPRegressor used for residual correction in Stage 4 (notebooks/4.1-4.3) with a PyTorch + Lightning equivalent, validated in notebooks/5.1-PyTorch_Migration.ipynb to reproduce Stage 4's frozen sklearn results within tolerance across all three of its notebooks. This is the generic building block Stage 5's Phases 2-3 build the field-dependent, physics-structured models on top of; the reproduction/comparison logic itself (loading frozen physical fits, sklearn-vs-pytorch tables) stays notebook-local, same as Stage 4's own per-dataset/joint-fit orchestration never left its notebooks.

ResidualMLP mirrors MLPRegressor's architecture (a stack of Linear+activation hidden layers from hidden_layer_sizes, output dim 1 for a per-axis candidate or 2 for a joint one). LitResidualRegressor wraps it to match MLPRegressor's training behavior as closely as possible, not just its architecture -- two details matter and are easy to get wrong:

  • Loss is 0.5 * MSE, matching sklearn's squared_loss convention (a bare MSELoss differs by a factor of 2 in the data-term gradient).
  • alpha (sklearn's L2 penalty) is not the same as PyTorch Adam(weight_decay=alpha). sklearn's _backprop adds alpha * w to the gradient and then divides the whole gradient (data term + penalty) by the current minibatch size (n_samples inside _backprop, i.e. the batch size, not the full training-set size) -- see sklearn.neural_network._multilayer_perceptron.BaseMultilayerPerceptron ._backprop/_compute_loss_grad. The weight_decay-equivalent is alpha / batch_size, computed by train_residual_mlp and passed to LitResidualRegressor explicitly. Using raw alpha over-regularizes by ~batch_sizex -- this was a real bug caught while validating Checkpoint 2 of 5.1-PyTorch_Migration.ipynb (the joint-model candidate, which had the largest alpha of the three configs compared, came out 2.3x worse than sklearn until fixed).

LitResidualRegressor 

LitResidualRegressor(n_inputs, n_outputs, hidden_layer_sizes, activation='relu', alpha=0.0001, lr=0.001, batch_size=200)

Bases: LightningModule

Lightning wrapper around ResidualMLP: 0.5*MSE loss + Adam, with alpha/batch_size matching sklearn MLPRegressor's L2 penalty exactly (see module docstring). batch_size must be the actual training DataLoader batch size for the weight_decay scaling to match.

Source code in dispcraft/ml.py 
70
71
72
73
74
def __init__(self, n_inputs, n_outputs, hidden_layer_sizes, activation="relu",
             alpha=1e-4, lr=1e-3, batch_size=200):
    super().__init__()
    self.save_hyperparameters()
    self.model = ResidualMLP(n_inputs, n_outputs, hidden_layer_sizes, activation)

ResidualMLP 

ResidualMLP(n_inputs, n_outputs, hidden_layer_sizes, activation='relu')

Bases: Module

Feed-forward regressor matching sklearn MLPRegressor's shape: a stack of Linear+activation hidden layers from hidden_layer_sizes, ending in a plain linear output layer of size n_outputs (1 per-axis, 2 joint).

Source code in dispcraft/ml.py 
50
51
52
53
54
55
56
57
58
def __init__(self, n_inputs, n_outputs, hidden_layer_sizes, activation="relu"):
    super().__init__()
    act_cls = ACTIVATIONS[activation]
    sizes = [n_inputs, *hidden_layer_sizes]
    layers = []
    for in_size, out_size in zip(sizes[:-1], sizes[1:]):
        layers += [nn.Linear(in_size, out_size), act_cls()]
    layers.append(nn.Linear(sizes[-1], n_outputs))
    self.net = nn.Sequential(*layers)

make_loaders 

make_loaders(X_train, y_train, val_frac=0.1, batch_size=200, seed=0)

Split (X_train, y_train) into train/val DataLoaders.

sklearn's early_stopping=True carves off validation_fraction (default 0.1) of the training data for its internal stopping criterion -- mirrored here with the same fraction so both frameworks see the same effective train/val split. batch_size should be sklearn's batch_size="auto" default, min(200, n_samples), computed by the caller from the post-validation-split training set size (see train_residual_mlp).

Source code in dispcraft/ml.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def make_loaders(X_train, y_train, val_frac=0.1, batch_size=200, seed=0):
    """Split `(X_train, y_train)` into train/val `DataLoader`s.

    sklearn's `early_stopping=True` carves off `validation_fraction`
    (default 0.1) of the training data for its internal stopping criterion
    -- mirrored here with the same fraction so both frameworks see the same
    effective train/val split. `batch_size` should be sklearn's
    `batch_size="auto"` default, `min(200, n_samples)`, computed by the
    caller from the post-validation-split training set size (see
    `train_residual_mlp`).
    """
    n_val = max(1, int(len(X_train) * val_frac))
    g = torch.Generator().manual_seed(seed)
    perm = torch.randperm(len(X_train), generator=g)
    val_idx, tr_idx = perm[:n_val], perm[n_val:]

    X_t = torch.tensor(X_train, dtype=torch.float32)
    y_t = torch.tensor(y_train, dtype=torch.float32)
    if y_t.ndim == 1:
        y_t = y_t.unsqueeze(1)

    train_ds = torch.utils.data.TensorDataset(X_t[tr_idx], y_t[tr_idx])
    val_ds = torch.utils.data.TensorDataset(X_t[val_idx], y_t[val_idx])
    train_loader = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=True,
                                                generator=torch.Generator().manual_seed(seed))
    val_loader = torch.utils.data.DataLoader(val_ds, batch_size=batch_size)
    return train_loader, val_loader

predict 

predict(lit_model, X)

Run lit_model in inference mode over X, returning a NumPy array.

Source code in dispcraft/ml.py 
162
163
164
165
def predict(lit_model, X):
    """Run `lit_model` in inference mode over `X`, returning a NumPy array."""
    with torch.no_grad():
        return lit_model(torch.tensor(X, dtype=torch.float32)).numpy()

train_residual_mlp 

train_residual_mlp(X_train, y_train, params, n_outputs, max_epochs=2000, patience=20, val_frac=0.1)

Train a LitResidualRegressor to convergence on (X_train, y_train).

params : dict with hidden_layer_sizes, activation, alpha, random_state, early_stopping -- the same shape as an MLPRegressor(**params) call, e.g. pulled from an MLflow run's logged hyperparameters (see notebooks/4.2-Per_Dataset_Pipeline .ipynb's best_mlp_params()). n_outputs : 1 for a per-axis candidate, 2 for a joint (y, z) candidate.

Returns:
  • LitResidualRegressor, in eval mode, ready for `predict()`. –
Source code in dispcraft/ml.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def train_residual_mlp(X_train, y_train, params, n_outputs, max_epochs=2000, patience=20, val_frac=0.1):
    """Train a `LitResidualRegressor` to convergence on `(X_train, y_train)`.

    `params` : dict with `hidden_layer_sizes`, `activation`, `alpha`,
        `random_state`, `early_stopping` -- the same shape as an
        `MLPRegressor(**params)` call, e.g. pulled from an MLflow run's
        logged hyperparameters (see `notebooks/4.2-Per_Dataset_Pipeline
        .ipynb`'s `best_mlp_params()`).
    n_outputs : 1 for a per-axis candidate, 2 for a joint (y, z) candidate.

    Returns
    -------
    LitResidualRegressor, in eval mode, ready for `predict()`.
    """
    pl.seed_everything(params["random_state"], workers=True, verbose=False)

    n_train_post_val = len(X_train) - max(1, int(len(X_train) * val_frac))
    batch_size = min(200, n_train_post_val)  # sklearn MLPRegressor's batch_size="auto" default
    train_loader, val_loader = make_loaders(X_train, y_train, val_frac=val_frac,
                                             batch_size=batch_size, seed=params["random_state"])

    lit_model = LitResidualRegressor(
        n_inputs=X_train.shape[1], n_outputs=n_outputs,
        hidden_layer_sizes=params["hidden_layer_sizes"], activation=params["activation"],
        alpha=params["alpha"], batch_size=batch_size,
    )
    trainer = pl.Trainer(
        max_epochs=max_epochs, accelerator="cpu", enable_progress_bar=False,
        enable_model_summary=False, enable_checkpointing=False, logger=False, deterministic=True,
        callbacks=[EarlyStopping(monitor="val_loss", patience=patience)] if params["early_stopping"] else [],
    )
    trainer.fit(lit_model, train_loader, val_loader)
    lit_model.eval()
    return lit_model