Stage 5, Phase 1 — PyTorch/Lightning Migration (`5.1-PyTorch_Migration`)¶

Stage 4 Phases 1-3 built residual-MLP correction entirely with sklearn.neural_network.MLPRegressor, notebook-local (never promoted to dispcraft/). Stage 5 Phases 2-3 need a more expressive, physics-structured NN (field-dependent parameters as a function of (y0, z0)) that sklearn can't express. This notebook's job is narrow: port the existing residual-MLP machinery to PyTorch + Lightning and prove it reproduces Stage 4's frozen sklearn numbers, before any new modeling is attempted. It's a correctness gate, not new science -- no new architecture, no new physics.

The physical forward model (dispcraft/calibration.py) is untouched. All frozen physical fits (models/stage3_fit_*.toml, models/joint_*.toml) are reused exactly as-is -- no refitting.

Built in four checkpoints, reviewed in order:

Model skeleton (ResidualMLP + LitResidualRegressor).
Reproduce 4.1-ML_Comparison.ipynb's single-dataset (rgs000_0) numbers.
Reproduce 4.2-Per_Dataset_Pipeline.ipynb's per-dataset pipeline (6 configs).
Reproduce 4.3-Multi_Dataset_Fitting.ipynb's final hybrid step + summary.

1. Setup¶

In [1]:

Copied!





import tomllib
from pathlib import Path

import numpy as np
import pandas as pd
import mlflow
import torch
import torch.nn as nn
import pytorch_lightning as pl
from pytorch_lightning.callbacks import EarlyStopping
from pytorch_lightning.callbacks.progress import TQDMProgressBar

from dispcraft.calibration import ground_test_model_from_config, predict_centroids
from dispcraft.measurement import load_spectra, median_per_spectrum

REPO_ROOT = Path("..")
DATA_DIR = REPO_ROOT / "data"
MODELS_DIR = REPO_ROOT / "models"
RNG_SEED = 42

pl.seed_everything(RNG_SEED, workers=True)
torch.use_deterministic_algorithms(True)

with open(MODELS_DIR / "stage1_instrument.toml", "rb") as f:
    base_config = tomllib.load(f)
model = ground_test_model_from_config(base_config)
import tomllib
from pathlib import Path

import numpy as np
import pandas as pd
import mlflow
import torch
import torch.nn as nn
import pytorch_lightning as pl
from pytorch_lightning.callbacks import EarlyStopping
from pytorch_lightning.callbacks.progress import TQDMProgressBar

from dispcraft.calibration import ground_test_model_from_config, predict_centroids
from dispcraft.measurement import load_spectra, median_per_spectrum

REPO_ROOT = Path("..")
DATA_DIR = REPO_ROOT / "data"
MODELS_DIR = REPO_ROOT / "models"
RNG_SEED = 42

pl.seed_everything(RNG_SEED, workers=True)
torch.use_deterministic_algorithms(True)

with open(MODELS_DIR / "stage1_instrument.toml", "rb") as f:
    base_config = tomllib.load(f)
model = ground_test_model_from_config(base_config)

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Seed set to 42

2. Model skeleton¶

ResidualMLP/LitResidualRegressor are validated below (Checkpoints 2-4) and promoted to dispcraft.ml -- imported here, not redefined, per the project's notebook-to-library workflow. See dispcraft/ml.py's module docstring for the sklearn-compatibility details (0.5*MSE loss, alpha/batch_size weight-decay scaling) that make the port match MLPRegressor's training behavior, not just its architecture.

No MLflow or real data yet -- just a synthetic-tensor sanity check that the imported classes actually train, below.

In [2]:

Copied!

from dispcraft.ml import ResidualMLP, LitResidualRegressor, make_loaders, train_residual_mlp, predict, rmse
from dispcraft.ml import ResidualMLP, LitResidualRegressor, make_loaders, train_residual_mlp, predict, rmse

In [3]:

Copied!





# Sanity check on synthetic data: joint (n_outputs=2) model should fit a
# trivial linear target to near-zero training loss within a handful of epochs.
rng = np.random.default_rng(0)
x_synth = torch.tensor(rng.uniform(-1, 1, (200, 3)), dtype=torch.float32)
y_synth = torch.stack([x_synth[:, 0] + x_synth[:, 1], x_synth[:, 2] ** 2], dim=1)

train_ds = torch.utils.data.TensorDataset(x_synth[:160], y_synth[:160])
val_ds = torch.utils.data.TensorDataset(x_synth[160:], y_synth[160:])
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_ds, batch_size=32)

sanity_model = LitResidualRegressor(n_inputs=3, n_outputs=2, hidden_layer_sizes=(32, 32), batch_size=32)
trainer = pl.Trainer(max_epochs=200, accelerator="cpu", enable_progress_bar=False,
                      enable_model_summary=False, enable_checkpointing=False, logger=False,
                      callbacks=[EarlyStopping(monitor="val_loss", patience=20)])
trainer.fit(sanity_model, train_loader, val_loader)

with torch.no_grad():
    final_train_loss = nn.functional.mse_loss(sanity_model(x_synth[:160]), y_synth[:160]).item()
print(f"synthetic sanity check -- final train MSE: {final_train_loss:.5f} (expect << 1)")
assert final_train_loss < 0.05, "ResidualMLP/LitResidualRegressor failed to fit a trivial synthetic target"
# Sanity check on synthetic data: joint (n_outputs=2) model should fit a
# trivial linear target to near-zero training loss within a handful of epochs.
rng = np.random.default_rng(0)
x_synth = torch.tensor(rng.uniform(-1, 1, (200, 3)), dtype=torch.float32)
y_synth = torch.stack([x_synth[:, 0] + x_synth[:, 1], x_synth[:, 2] ** 2], dim=1)

train_ds = torch.utils.data.TensorDataset(x_synth[:160], y_synth[:160])
val_ds = torch.utils.data.TensorDataset(x_synth[160:], y_synth[160:])
train_loader = torch.utils.data.DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_ds, batch_size=32)

sanity_model = LitResidualRegressor(n_inputs=3, n_outputs=2, hidden_layer_sizes=(32, 32), batch_size=32)
trainer = pl.Trainer(max_epochs=200, accelerator="cpu", enable_progress_bar=False,
                      enable_model_summary=False, enable_checkpointing=False, logger=False,
                      callbacks=[EarlyStopping(monitor="val_loss", patience=20)])
trainer.fit(sanity_model, train_loader, val_loader)

with torch.no_grad():
    final_train_loss = nn.functional.mse_loss(sanity_model(x_synth[:160]), y_synth[:160]).item()
print(f"synthetic sanity check -- final train MSE: {final_train_loss:.5f} (expect << 1)")
assert final_train_loss < 0.05, "ResidualMLP/LitResidualRegressor failed to fit a trivial synthetic target"

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

`Trainer.fit` stopped: `max_epochs=200` reached.

synthetic sanity check -- final train MSE: 0.00021 (expect << 1)

3. Reproduce `4.1-ML_Comparison.ipynb`'s single-dataset (`rgs000_0`) numbers¶

Same starting point as 4.1: load the frozen Stage 3 physical fit (models/stage3_fit_rgs000_0.toml, no refit), compute residuals, and split with the identical GroupShuffleSplit(random_state=RNG_SEED) on spectra_id. Not re-running 4.1's hyperparameter search -- pulling its already-chosen winning MLP hyperparameters straight out of the frozen residual_correction MLflow experiment via the same best_mlp_params() helper 4.2/4.3 use, then training the PyTorch/Lightning equivalent on the exact same data and comparing test RMSE.

In [4]:

Copied!





from sklearn.model_selection import GroupShuffleSplit
from sklearn.preprocessing import StandardScaler

with open(MODELS_DIR / "stage3_fit_rgs000_0.toml", "rb") as f:
    fit0 = tomllib.load(f)

free_names = fit0["fit"]["free_parameters"]
theta0 = np.array([fit0["fit"]["result"][k] for k in free_names])
excluded_ids = fit0["fit"]["excluded_spectra_ids"]

df0 = median_per_spectrum(load_spectra(DATA_DIR / "rgs000_0_first.csv"))
df0 = df0[~df0["spectra_id"].isin(excluded_ids)].reset_index(drop=True)
assert len(df0) == fit0["fit"]["n_points"], "row count must match the TOML's recorded fit"

pred0 = predict_centroids(df0["y_nisp"], df0["z_nisp"], df0["wavelength"], theta0, free_names, model)
d0 = df0.assign(r_y=pred0[0] - df0["cent_y"].values, r_z=pred0[1] - df0["cent_z"].values)

X = d0[["y_nisp", "z_nisp", "wavelength"]].values
groups = d0["spectra_id"].values
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
train_idx, test_idx = next(gss.split(X, groups=groups))

scaler = StandardScaler().fit(X[train_idx])
X_train_s, X_test_s = scaler.transform(X[train_idx]), scaler.transform(X[test_idx])
r_y_train, r_z_train = d0["r_y"].values[train_idx], d0["r_z"].values[train_idx]
r_y_test, r_z_test = d0["r_y"].values[test_idx], d0["r_z"].values[test_idx]

print(f"rgs000_0: {len(d0)} rows, train {len(train_idx)} / test {len(test_idx)}")
from sklearn.model_selection import GroupShuffleSplit
from sklearn.preprocessing import StandardScaler

with open(MODELS_DIR / "stage3_fit_rgs000_0.toml", "rb") as f:
    fit0 = tomllib.load(f)

free_names = fit0["fit"]["free_parameters"]
theta0 = np.array([fit0["fit"]["result"][k] for k in free_names])
excluded_ids = fit0["fit"]["excluded_spectra_ids"]

df0 = median_per_spectrum(load_spectra(DATA_DIR / "rgs000_0_first.csv"))
df0 = df0[~df0["spectra_id"].isin(excluded_ids)].reset_index(drop=True)
assert len(df0) == fit0["fit"]["n_points"], "row count must match the TOML's recorded fit"

pred0 = predict_centroids(df0["y_nisp"], df0["z_nisp"], df0["wavelength"], theta0, free_names, model)
d0 = df0.assign(r_y=pred0[0] - df0["cent_y"].values, r_z=pred0[1] - df0["cent_z"].values)

X = d0[["y_nisp", "z_nisp", "wavelength"]].values
groups = d0["spectra_id"].values
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
train_idx, test_idx = next(gss.split(X, groups=groups))

scaler = StandardScaler().fit(X[train_idx])
X_train_s, X_test_s = scaler.transform(X[train_idx]), scaler.transform(X[test_idx])
r_y_train, r_z_train = d0["r_y"].values[train_idx], d0["r_z"].values[train_idx]
r_y_test, r_z_test = d0["r_y"].values[test_idx], d0["r_z"].values[test_idx]

print(f"rgs000_0: {len(d0)} rows, train {len(train_idx)} / test {len(test_idx)}")

rgs000_0: 5297 rows, train 4239 / test 1058

Pull the frozen sklearn-winning MLP hyperparameters¶

Same best_mlp_params() helper as 4.2-Per_Dataset_Pipeline.ipynb / 4.3-Multi_Dataset_Fitting.ipynb -- reads the already-populated residual_correction MLflow experiment (from running 4.1-ML_Comparison.ipynb) and returns the winning MLPRegressor hyperparameters per axis and for the joint model. Reused verbatim rather than re-deriving, so the PyTorch models below start from literally the same architecture choice sklearn's search landed on.

In [5]:

Copied!





import ast

mlflow.set_tracking_uri(f"sqlite:///{(REPO_ROOT / 'mlflow.db').resolve()}")

residual_exp = mlflow.get_experiment_by_name("residual_correction")
assert residual_exp is not None, "run 4.1-ML_Comparison.ipynb first to populate this experiment"
sklearn_runs = mlflow.search_runs(experiment_ids=[residual_exp.experiment_id])


def best_mlp_params(runs, axis, metric_col):
    sub = runs[(runs["params.model"] == "MLP") & (runs["params.axis"] == axis)]
    best = sub.loc[sub[metric_col].astype(float).idxmin()]
    return {
        "hidden_layer_sizes": ast.literal_eval(best["params.hidden_layer_sizes"]),
        "activation": best["params.activation"],
        "alpha": float(best["params.alpha"]),
        "random_state": int(best["params.random_state"]),
        "max_iter": int(best["params.max_iter"]),
        "early_stopping": best["params.early_stopping"] == "True",
    }


mlp_y_params = best_mlp_params(sklearn_runs, "y", "metrics.test_rmse")
mlp_z_params = best_mlp_params(sklearn_runs, "z", "metrics.test_rmse")
mlp_joint_params = best_mlp_params(sklearn_runs, "joint", "metrics.test_dist_mm")

print("MLP(y)     (independent):", mlp_y_params)
print("MLP(z)     (independent):", mlp_z_params)
print("MLP(y,z)   (joint):      ", mlp_joint_params)

# frozen sklearn reference numbers, straight from the same experiment
sklearn_ref = {
    "y": float(sklearn_runs[(sklearn_runs["params.model"] == "MLP") & (sklearn_runs["params.axis"] == "y")]
               ["metrics.test_rmse"].astype(float).min()),
    "z": float(sklearn_runs[(sklearn_runs["params.model"] == "MLP") & (sklearn_runs["params.axis"] == "z")]
               ["metrics.test_rmse"].astype(float).min()),
    "joint": float(sklearn_runs[(sklearn_runs["params.model"] == "MLP") & (sklearn_runs["params.axis"] == "joint")]
                   ["metrics.test_dist_mm"].astype(float).min()),
}
print("\nfrozen sklearn reference:", sklearn_ref)
import ast

mlflow.set_tracking_uri(f"sqlite:///{(REPO_ROOT / 'mlflow.db').resolve()}")

residual_exp = mlflow.get_experiment_by_name("residual_correction")
assert residual_exp is not None, "run 4.1-ML_Comparison.ipynb first to populate this experiment"
sklearn_runs = mlflow.search_runs(experiment_ids=[residual_exp.experiment_id])


def best_mlp_params(runs, axis, metric_col):
    sub = runs[(runs["params.model"] == "MLP") & (runs["params.axis"] == axis)]
    best = sub.loc[sub[metric_col].astype(float).idxmin()]
    return {
        "hidden_layer_sizes": ast.literal_eval(best["params.hidden_layer_sizes"]),
        "activation": best["params.activation"],
        "alpha": float(best["params.alpha"]),
        "random_state": int(best["params.random_state"]),
        "max_iter": int(best["params.max_iter"]),
        "early_stopping": best["params.early_stopping"] == "True",
    }


mlp_y_params = best_mlp_params(sklearn_runs, "y", "metrics.test_rmse")
mlp_z_params = best_mlp_params(sklearn_runs, "z", "metrics.test_rmse")
mlp_joint_params = best_mlp_params(sklearn_runs, "joint", "metrics.test_dist_mm")

print("MLP(y)     (independent):", mlp_y_params)
print("MLP(z)     (independent):", mlp_z_params)
print("MLP(y,z)   (joint):      ", mlp_joint_params)

# frozen sklearn reference numbers, straight from the same experiment
sklearn_ref = {
    "y": float(sklearn_runs[(sklearn_runs["params.model"] == "MLP") & (sklearn_runs["params.axis"] == "y")]
               ["metrics.test_rmse"].astype(float).min()),
    "z": float(sklearn_runs[(sklearn_runs["params.model"] == "MLP") & (sklearn_runs["params.axis"] == "z")]
               ["metrics.test_rmse"].astype(float).min()),
    "joint": float(sklearn_runs[(sklearn_runs["params.model"] == "MLP") & (sklearn_runs["params.axis"] == "joint")]
                   ["metrics.test_dist_mm"].astype(float).min()),
}
print("\nfrozen sklearn reference:", sklearn_ref)

MLP(y)     (independent): {'hidden_layer_sizes': (64, 32), 'activation': 'relu', 'alpha': 0.0001, 'random_state': 42, 'max_iter': 2000, 'early_stopping': True}
MLP(z)     (independent): {'hidden_layer_sizes': (64, 64), 'activation': 'relu', 'alpha': 0.001, 'random_state': 42, 'max_iter': 2000, 'early_stopping': True}
MLP(y,z)   (joint):       {'hidden_layer_sizes': (64, 64), 'activation': 'relu', 'alpha': 0.01, 'random_state': 42, 'max_iter': 2000, 'early_stopping': True}

frozen sklearn reference: {'y': 0.016473747109319505, 'z': 0.06261569821889774, 'joint': 0.056049226986950425}

Train the PyTorch/Lightning equivalents¶

train_residual_mlp (imported from dispcraft.ml) maps hidden_layer_sizes and activation directly onto ResidualMLP, and reproduces sklearn's alpha/early_stopping behavior exactly (see dispcraft/ml.py's module docstring for the alpha/batch_size weight-decay-scaling detail). New MLflow experiment residual_correction_torch keeps these runs alongside but separate from the frozen sklearn ones, so both are queryable side-by-side without ever touching the historical residual_correction experiment.

In [6]:

Copied!





TORCH_EXPERIMENT = "residual_correction_torch"
if mlflow.get_experiment_by_name(TORCH_EXPERIMENT) is None:
    mlflow.create_experiment(TORCH_EXPERIMENT, artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(TORCH_EXPERIMENT)
TORCH_EXPERIMENT = "residual_correction_torch"
if mlflow.get_experiment_by_name(TORCH_EXPERIMENT) is None:
    mlflow.create_experiment(TORCH_EXPERIMENT, artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(TORCH_EXPERIMENT)

Out[6]:

<Experiment: artifact_location='file:/home/zoubian/Workspace/dispers/dispcraft/mlruns', creation_time=1784899056181, effective_trace_archival_retention=None, experiment_id='6', last_update_time=1784899056181, lifecycle_stage='active', name='residual_correction_torch', tags={}, trace_location=None, workspace='default'>

In [7]:

Copied!





torch_models = {}
torch_metrics = {}

for axis, col, params in [("y", "r_y", mlp_y_params), ("z", "r_z", mlp_z_params)]:
    y_train = d0[col].values[train_idx]
    y_test = d0[col].values[test_idx]

    lit_model = train_residual_mlp(X_train_s, y_train, params, n_outputs=1)
    train_pred = predict(lit_model, X_train_s).ravel()
    test_pred = predict(lit_model, X_test_s).ravel()
    train_rmse, test_rmse = rmse(y_train, train_pred), rmse(y_test, test_pred)

    with mlflow.start_run(run_name=f"MLP_{axis}_torch"):
        mlflow.log_param("axis", axis)
        mlflow.log_param("model", "MLP")
        mlflow.log_param("framework", "pytorch")
        for k, v in params.items():
            mlflow.log_param(k, v)
        mlflow.log_metric("train_rmse", train_rmse)
        mlflow.log_metric("test_rmse", test_rmse)
        mlflow.log_metric("train_test_ratio", test_rmse / train_rmse if train_rmse > 0 else float("nan"))

    torch_models[axis] = lit_model
    torch_metrics[axis] = {"train_rmse": train_rmse, "test_rmse": test_rmse}

# joint (y, z) together
Y_train = d0[["r_y", "r_z"]].values[train_idx]
Y_test = d0[["r_y", "r_z"]].values[test_idx]
joint_model = train_residual_mlp(X_train_s, Y_train, mlp_joint_params, n_outputs=2)
joint_train_pred = predict(joint_model, X_train_s)
joint_test_pred = predict(joint_model, X_test_s)
joint_train_dist = float(np.mean(np.hypot(*(Y_train - joint_train_pred).T)))
joint_test_dist = float(np.mean(np.hypot(*(Y_test - joint_test_pred).T)))
joint_train_rmse_y = rmse(Y_train[:, 0], joint_train_pred[:, 0])
joint_train_rmse_z = rmse(Y_train[:, 1], joint_train_pred[:, 1])
joint_test_rmse_y = rmse(Y_test[:, 0], joint_test_pred[:, 0])
joint_test_rmse_z = rmse(Y_test[:, 1], joint_test_pred[:, 1])

with mlflow.start_run(run_name="MLP_joint_torch"):
    mlflow.log_param("axis", "joint")
    mlflow.log_param("model", "MLP")
    mlflow.log_param("framework", "pytorch")
    for k, v in mlp_joint_params.items():
        mlflow.log_param(k, v)
    mlflow.log_metric("train_dist_mm", joint_train_dist)
    mlflow.log_metric("test_dist_mm", joint_test_dist)
    mlflow.log_metric("test_rmse_y", joint_test_rmse_y)
    mlflow.log_metric("test_rmse_z", joint_test_rmse_z)

torch_metrics["joint"] = {"train_dist_mm": joint_train_dist, "test_dist_mm": joint_test_dist}
print("independent y:", torch_metrics["y"])
print("independent z:", torch_metrics["z"])
print("joint:        ", torch_metrics["joint"])
torch_models = {}
torch_metrics = {}

for axis, col, params in [("y", "r_y", mlp_y_params), ("z", "r_z", mlp_z_params)]:
    y_train = d0[col].values[train_idx]
    y_test = d0[col].values[test_idx]

    lit_model = train_residual_mlp(X_train_s, y_train, params, n_outputs=1)
    train_pred = predict(lit_model, X_train_s).ravel()
    test_pred = predict(lit_model, X_test_s).ravel()
    train_rmse, test_rmse = rmse(y_train, train_pred), rmse(y_test, test_pred)

    with mlflow.start_run(run_name=f"MLP_{axis}_torch"):
        mlflow.log_param("axis", axis)
        mlflow.log_param("model", "MLP")
        mlflow.log_param("framework", "pytorch")
        for k, v in params.items():
            mlflow.log_param(k, v)
        mlflow.log_metric("train_rmse", train_rmse)
        mlflow.log_metric("test_rmse", test_rmse)
        mlflow.log_metric("train_test_ratio", test_rmse / train_rmse if train_rmse > 0 else float("nan"))

    torch_models[axis] = lit_model
    torch_metrics[axis] = {"train_rmse": train_rmse, "test_rmse": test_rmse}

# joint (y, z) together
Y_train = d0[["r_y", "r_z"]].values[train_idx]
Y_test = d0[["r_y", "r_z"]].values[test_idx]
joint_model = train_residual_mlp(X_train_s, Y_train, mlp_joint_params, n_outputs=2)
joint_train_pred = predict(joint_model, X_train_s)
joint_test_pred = predict(joint_model, X_test_s)
joint_train_dist = float(np.mean(np.hypot(*(Y_train - joint_train_pred).T)))
joint_test_dist = float(np.mean(np.hypot(*(Y_test - joint_test_pred).T)))
joint_train_rmse_y = rmse(Y_train[:, 0], joint_train_pred[:, 0])
joint_train_rmse_z = rmse(Y_train[:, 1], joint_train_pred[:, 1])
joint_test_rmse_y = rmse(Y_test[:, 0], joint_test_pred[:, 0])
joint_test_rmse_z = rmse(Y_test[:, 1], joint_test_pred[:, 1])

with mlflow.start_run(run_name="MLP_joint_torch"):
    mlflow.log_param("axis", "joint")
    mlflow.log_param("model", "MLP")
    mlflow.log_param("framework", "pytorch")
    for k, v in mlp_joint_params.items():
        mlflow.log_param(k, v)
    mlflow.log_metric("train_dist_mm", joint_train_dist)
    mlflow.log_metric("test_dist_mm", joint_test_dist)
    mlflow.log_metric("test_rmse_y", joint_test_rmse_y)
    mlflow.log_metric("test_rmse_z", joint_test_rmse_z)

torch_metrics["joint"] = {"train_dist_mm": joint_train_dist, "test_dist_mm": joint_test_dist}
print("independent y:", torch_metrics["y"])
print("independent z:", torch_metrics["z"])
print("joint:        ", torch_metrics["joint"])

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

independent y: {'train_rmse': 0.011726579080055638, 'test_rmse': 0.014878204774501835}
independent z: {'train_rmse': 0.033300850215434, 'test_rmse': 0.06574319329681264}
joint:         {'train_dist_mm': 0.035799372041880176, 'test_dist_mm': 0.054362191518319544}

Checkpoint 2 comparison: sklearn vs. PyTorch/Lightning¶

Bit-exact reproduction isn't expected across frameworks (different Adam implementations, different internal train/val carve for early stopping, etc.) -- the bar is comparable residual-correction quality on the same held-out split, not identical weights. PASS_TOL sets how much worse (as a fraction) the PyTorch number is allowed to be than the frozen sklearn one before this is flagged as a real regression rather than framework noise.

In [8]:

Copied!





PASS_TOL = 0.25  # torch test metric may be up to 25% worse than sklearn's before flagging

checkpoint2_table = pd.DataFrame([
    {"target": "y (independent)", "sklearn_test_rmse_mm": sklearn_ref["y"], "torch_test_rmse_mm": torch_metrics["y"]["test_rmse"]},
    {"target": "z (independent)", "sklearn_test_rmse_mm": sklearn_ref["z"], "torch_test_rmse_mm": torch_metrics["z"]["test_rmse"]},
    {"target": "joint (dist)", "sklearn_test_rmse_mm": sklearn_ref["joint"], "torch_test_rmse_mm": torch_metrics["joint"]["test_dist_mm"]},
]).set_index("target")
checkpoint2_table["ratio_torch_over_sklearn"] = (
    checkpoint2_table["torch_test_rmse_mm"] / checkpoint2_table["sklearn_test_rmse_mm"]
)
checkpoint2_table["status"] = np.where(
    checkpoint2_table["ratio_torch_over_sklearn"] <= 1 + PASS_TOL, "OK", "REGRESSION"
)
display(checkpoint2_table.round(4))

n_regressions = (checkpoint2_table["status"] == "REGRESSION").sum()
print(f"\nCheckpoint 2: {len(checkpoint2_table) - n_regressions}/{len(checkpoint2_table)} targets within "
      f"{PASS_TOL:.0%} of the frozen sklearn test metric.")
if n_regressions:
    print("REGRESSION(S) FLAGGED -- reported as-is, not smoothed over:")
    print(checkpoint2_table[checkpoint2_table["status"] == "REGRESSION"])
PASS_TOL = 0.25  # torch test metric may be up to 25% worse than sklearn's before flagging

checkpoint2_table = pd.DataFrame([
    {"target": "y (independent)", "sklearn_test_rmse_mm": sklearn_ref["y"], "torch_test_rmse_mm": torch_metrics["y"]["test_rmse"]},
    {"target": "z (independent)", "sklearn_test_rmse_mm": sklearn_ref["z"], "torch_test_rmse_mm": torch_metrics["z"]["test_rmse"]},
    {"target": "joint (dist)", "sklearn_test_rmse_mm": sklearn_ref["joint"], "torch_test_rmse_mm": torch_metrics["joint"]["test_dist_mm"]},
]).set_index("target")
checkpoint2_table["ratio_torch_over_sklearn"] = (
    checkpoint2_table["torch_test_rmse_mm"] / checkpoint2_table["sklearn_test_rmse_mm"]
)
checkpoint2_table["status"] = np.where(
    checkpoint2_table["ratio_torch_over_sklearn"] <= 1 + PASS_TOL, "OK", "REGRESSION"
)
display(checkpoint2_table.round(4))

n_regressions = (checkpoint2_table["status"] == "REGRESSION").sum()
print(f"\nCheckpoint 2: {len(checkpoint2_table) - n_regressions}/{len(checkpoint2_table)} targets within "
      f"{PASS_TOL:.0%} of the frozen sklearn test metric.")
if n_regressions:
    print("REGRESSION(S) FLAGGED -- reported as-is, not smoothed over:")
    print(checkpoint2_table[checkpoint2_table["status"] == "REGRESSION"])

	sklearn_test_rmse_mm	torch_test_rmse_mm	ratio_torch_over_sklearn	status
target
y (independent)	0.0165	0.0149	0.9031	OK
z (independent)	0.0626	0.0657	1.0499	OK
joint (dist)	0.0560	0.0544	0.9699	OK

Checkpoint 2: 3/3 targets within 25% of the frozen sklearn test metric.

4. Reproduce `4.2-Per_Dataset_Pipeline.ipynb`'s per-dataset pipeline (6 configs)¶

Same reusable-pipeline structure as 4.2's run_pipeline(cfg): for each of the 6 {rgs000,rgs180}x{0,m4,p4} configs, load that config's own frozen Stage 3 physical fit (models/stage3_fit_<cfg>.toml, no refit), compute residuals, group-aware 80/20 split, then train both candidates fresh on that dataset's own train split -- independent (MLP(y)+MLP(z)) and joint (MLP(y,z)) -- reusing the same frozen-winning hyperparameters (mlp_y_params/mlp_z_params/mlp_joint_params) and the train_residual_mlp/predict/rmse helpers from Checkpoint 2 above, swapping only the estimator (LitResidualRegressor instead of MLPRegressor).

In [9]:

Copied!





CONFIGS = ["rgs000_0", "rgs000_m4", "rgs000_p4", "rgs180_0", "rgs180_m4", "rgs180_p4"]


def load_fit(cfg):
    with open(MODELS_DIR / f"stage3_fit_{cfg}.toml", "rb") as f:
        fit = tomllib.load(f)
    free_names = fit["fit"]["free_parameters"]
    theta = np.array([fit["fit"]["result"][k] for k in free_names])
    excluded_ids = fit["fit"]["excluded_spectra_ids"]
    return fit, free_names, theta, excluded_ids


def run_pipeline_torch(cfg):
    fit, free_names, theta, excluded_ids = load_fit(cfg)

    df = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
    df = df[~df["spectra_id"].isin(excluded_ids)].reset_index(drop=True)
    assert len(df) == fit["fit"]["n_points"], f"{cfg}: row count must match the TOML's recorded fit"

    pred = predict_centroids(df["y_nisp"], df["z_nisp"], df["wavelength"], theta, free_names, model)
    d = df.assign(r_y=pred[0] - df["cent_y"].values, r_z=pred[1] - df["cent_z"].values)

    X = d[["y_nisp", "z_nisp", "wavelength"]].values
    groups = d["spectra_id"].values
    gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
    train_idx, test_idx = next(gss.split(X, groups=groups))

    scaler = StandardScaler().fit(X[train_idx])
    X_train_s, X_test_s = scaler.transform(X[train_idx]), scaler.transform(X[test_idx])
    r_y_train, r_z_train = d["r_y"].values[train_idx], d["r_z"].values[train_idx]
    r_y_test, r_z_test = d["r_y"].values[test_idx], d["r_z"].values[test_idx]

    physical_dist_test = float(np.mean(np.hypot(r_y_test, r_z_test)))

    # Candidate A: independent per-axis MLPs
    lit_y = train_residual_mlp(X_train_s, r_y_train, mlp_y_params, n_outputs=1)
    lit_z = train_residual_mlp(X_train_s, r_z_train, mlp_z_params, n_outputs=1)
    hyb_y_ind = r_y_test - predict(lit_y, X_test_s).ravel()
    hyb_z_ind = r_z_test - predict(lit_z, X_test_s).ravel()

    # Candidate B: joint 2D-output MLP
    Y_train = d[["r_y", "r_z"]].values[train_idx]
    lit_joint = train_residual_mlp(X_train_s, Y_train, mlp_joint_params, n_outputs=2)
    pred_joint_test = predict(lit_joint, X_test_s)
    hyb_y_joint = r_y_test - pred_joint_test[:, 0]
    hyb_z_joint = r_z_test - pred_joint_test[:, 1]

    candidates = {
        "independent": {
            "hybrid_rmse_y": rmse(r_y_test, predict(lit_y, X_test_s).ravel()),
            "hybrid_rmse_z": rmse(r_z_test, predict(lit_z, X_test_s).ravel()),
            "hybrid_dist_test": float(np.mean(np.hypot(hyb_y_ind, hyb_z_ind))),
        },
        "joint": {
            "hybrid_rmse_y": rmse(r_y_test, pred_joint_test[:, 0]),
            "hybrid_rmse_z": rmse(r_z_test, pred_joint_test[:, 1]),
            "hybrid_dist_test": float(np.mean(np.hypot(hyb_y_joint, hyb_z_joint))),
        },
    }
    params = {name: float(v) for name, v in zip(free_names, theta)}
    return {"physical_dist_test": physical_dist_test, "candidates": candidates, "params": params}
CONFIGS = ["rgs000_0", "rgs000_m4", "rgs000_p4", "rgs180_0", "rgs180_m4", "rgs180_p4"]


def load_fit(cfg):
    with open(MODELS_DIR / f"stage3_fit_{cfg}.toml", "rb") as f:
        fit = tomllib.load(f)
    free_names = fit["fit"]["free_parameters"]
    theta = np.array([fit["fit"]["result"][k] for k in free_names])
    excluded_ids = fit["fit"]["excluded_spectra_ids"]
    return fit, free_names, theta, excluded_ids


def run_pipeline_torch(cfg):
    fit, free_names, theta, excluded_ids = load_fit(cfg)

    df = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
    df = df[~df["spectra_id"].isin(excluded_ids)].reset_index(drop=True)
    assert len(df) == fit["fit"]["n_points"], f"{cfg}: row count must match the TOML's recorded fit"

    pred = predict_centroids(df["y_nisp"], df["z_nisp"], df["wavelength"], theta, free_names, model)
    d = df.assign(r_y=pred[0] - df["cent_y"].values, r_z=pred[1] - df["cent_z"].values)

    X = d[["y_nisp", "z_nisp", "wavelength"]].values
    groups = d["spectra_id"].values
    gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
    train_idx, test_idx = next(gss.split(X, groups=groups))

    scaler = StandardScaler().fit(X[train_idx])
    X_train_s, X_test_s = scaler.transform(X[train_idx]), scaler.transform(X[test_idx])
    r_y_train, r_z_train = d["r_y"].values[train_idx], d["r_z"].values[train_idx]
    r_y_test, r_z_test = d["r_y"].values[test_idx], d["r_z"].values[test_idx]

    physical_dist_test = float(np.mean(np.hypot(r_y_test, r_z_test)))

    # Candidate A: independent per-axis MLPs
    lit_y = train_residual_mlp(X_train_s, r_y_train, mlp_y_params, n_outputs=1)
    lit_z = train_residual_mlp(X_train_s, r_z_train, mlp_z_params, n_outputs=1)
    hyb_y_ind = r_y_test - predict(lit_y, X_test_s).ravel()
    hyb_z_ind = r_z_test - predict(lit_z, X_test_s).ravel()

    # Candidate B: joint 2D-output MLP
    Y_train = d[["r_y", "r_z"]].values[train_idx]
    lit_joint = train_residual_mlp(X_train_s, Y_train, mlp_joint_params, n_outputs=2)
    pred_joint_test = predict(lit_joint, X_test_s)
    hyb_y_joint = r_y_test - pred_joint_test[:, 0]
    hyb_z_joint = r_z_test - pred_joint_test[:, 1]

    candidates = {
        "independent": {
            "hybrid_rmse_y": rmse(r_y_test, predict(lit_y, X_test_s).ravel()),
            "hybrid_rmse_z": rmse(r_z_test, predict(lit_z, X_test_s).ravel()),
            "hybrid_dist_test": float(np.mean(np.hypot(hyb_y_ind, hyb_z_ind))),
        },
        "joint": {
            "hybrid_rmse_y": rmse(r_y_test, pred_joint_test[:, 0]),
            "hybrid_rmse_z": rmse(r_z_test, pred_joint_test[:, 1]),
            "hybrid_dist_test": float(np.mean(np.hypot(hyb_y_joint, hyb_z_joint))),
        },
    }
    params = {name: float(v) for name, v in zip(free_names, theta)}
    return {"physical_dist_test": physical_dist_test, "candidates": candidates, "params": params}

MLflow experiment: `per_dataset_torch`¶

Same one-run-per-(dataset, candidate) granularity as 4.2's per_dataset experiment (12 runs total) -- kept in its own _torch-suffixed experiment, same rationale as Checkpoint 2.

In [10]:

Copied!





TORCH_PER_DATASET_EXPERIMENT = "per_dataset_torch"
if mlflow.get_experiment_by_name(TORCH_PER_DATASET_EXPERIMENT) is None:
    mlflow.create_experiment(TORCH_PER_DATASET_EXPERIMENT,
                              artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(TORCH_PER_DATASET_EXPERIMENT)

torch_results = {}
torch_rows = []

for cfg in CONFIGS:
    res = run_pipeline_torch(cfg)
    torch_results[cfg] = res

    for approach, c in res["candidates"].items():
        improvement = 1 - c["hybrid_dist_test"] / res["physical_dist_test"]
        torch_rows.append({
            "dataset": cfg, "approach": approach,
            "physical_dist_test": res["physical_dist_test"],
            "hybrid_dist_test": c["hybrid_dist_test"],
            "hybrid_rmse_y": c["hybrid_rmse_y"], "hybrid_rmse_z": c["hybrid_rmse_z"],
            "improvement_pct": 100 * improvement,
        })

        with mlflow.start_run(run_name=f"{cfg}_{approach}_torch"):
            mlflow.log_param("dataset", cfg)
            mlflow.log_param("approach", approach)
            mlflow.log_param("framework", "pytorch")
            for k, v in res["params"].items():
                mlflow.log_param(k, v)
            mlflow.log_metric("physical_dist_test", res["physical_dist_test"])
            mlflow.log_metric("hybrid_dist_test", c["hybrid_dist_test"])
            mlflow.log_metric("hybrid_rmse_y", c["hybrid_rmse_y"])
            mlflow.log_metric("hybrid_rmse_z", c["hybrid_rmse_z"])
            mlflow.log_metric("improvement_pct", 100 * improvement)

torch_results_df = pd.DataFrame(torch_rows)
print(f"logged {len(torch_rows)} runs to the '{TORCH_PER_DATASET_EXPERIMENT}' MLflow experiment")
torch_results_df.round(4)
TORCH_PER_DATASET_EXPERIMENT = "per_dataset_torch"
if mlflow.get_experiment_by_name(TORCH_PER_DATASET_EXPERIMENT) is None:
    mlflow.create_experiment(TORCH_PER_DATASET_EXPERIMENT,
                              artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(TORCH_PER_DATASET_EXPERIMENT)

torch_results = {}
torch_rows = []

for cfg in CONFIGS:
    res = run_pipeline_torch(cfg)
    torch_results[cfg] = res

    for approach, c in res["candidates"].items():
        improvement = 1 - c["hybrid_dist_test"] / res["physical_dist_test"]
        torch_rows.append({
            "dataset": cfg, "approach": approach,
            "physical_dist_test": res["physical_dist_test"],
            "hybrid_dist_test": c["hybrid_dist_test"],
            "hybrid_rmse_y": c["hybrid_rmse_y"], "hybrid_rmse_z": c["hybrid_rmse_z"],
            "improvement_pct": 100 * improvement,
        })

        with mlflow.start_run(run_name=f"{cfg}_{approach}_torch"):
            mlflow.log_param("dataset", cfg)
            mlflow.log_param("approach", approach)
            mlflow.log_param("framework", "pytorch")
            for k, v in res["params"].items():
                mlflow.log_param(k, v)
            mlflow.log_metric("physical_dist_test", res["physical_dist_test"])
            mlflow.log_metric("hybrid_dist_test", c["hybrid_dist_test"])
            mlflow.log_metric("hybrid_rmse_y", c["hybrid_rmse_y"])
            mlflow.log_metric("hybrid_rmse_z", c["hybrid_rmse_z"])
            mlflow.log_metric("improvement_pct", 100 * improvement)

torch_results_df = pd.DataFrame(torch_rows)
print(f"logged {len(torch_rows)} runs to the '{TORCH_PER_DATASET_EXPERIMENT}' MLflow experiment")
torch_results_df.round(4)

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

logged 12 runs to the 'per_dataset_torch' MLflow experiment

Out[10]:

	dataset	approach	physical_dist_test	hybrid_dist_test	hybrid_rmse_y	hybrid_rmse_z	improvement_pct
0	rgs000_0	independent	0.3078	0.0548	0.0149	0.0657	82.1991
1	rgs000_0	joint	0.3078	0.0544	0.0172	0.0627	82.3379
2	rgs000_m4	independent	0.2928	0.0498	0.0114	0.0633	82.9885
3	rgs000_m4	joint	0.2928	0.0513	0.0137	0.0610	82.4808
4	rgs000_p4	independent	0.3485	0.0788	0.0137	0.1007	77.4019
5	rgs000_p4	joint	0.3485	0.0727	0.0207	0.0851	79.1272
6	rgs180_0	independent	0.3283	0.0744	0.0131	0.1031	77.3425
7	rgs180_0	joint	0.3283	0.0645	0.0140	0.0864	80.3585
8	rgs180_m4	independent	0.3500	0.0727	0.0146	0.0945	79.2350
9	rgs180_m4	joint	0.3500	0.0600	0.0204	0.0719	82.8449
10	rgs180_p4	independent	0.3132	0.0675	0.0112	0.0891	78.4533
11	rgs180_p4	joint	0.3132	0.0682	0.0167	0.0797	78.2337

Checkpoint 3 comparison: sklearn (`per_dataset`) vs. PyTorch/Lightning (`per_dataset_torch`)¶

Per-config, per-candidate hybrid_dist_test side by side, same PASS_TOL tolerance as Checkpoint 2.

In [11]:

Copied!





sklearn_per_dataset_exp = mlflow.get_experiment_by_name("per_dataset")
assert sklearn_per_dataset_exp is not None, "run 4.2-Per_Dataset_Pipeline.ipynb first to populate this experiment"
sklearn_pd_runs = mlflow.search_runs(experiment_ids=[sklearn_per_dataset_exp.experiment_id])
sklearn_pd = sklearn_pd_runs.set_index(["params.dataset", "params.approach"])["metrics.hybrid_dist_test"].astype(float)

torch_pd = torch_results_df.set_index(["dataset", "approach"])["hybrid_dist_test"]

checkpoint3_table = pd.DataFrame({
    "sklearn_hybrid_dist_mm": sklearn_pd,
    "torch_hybrid_dist_mm": torch_pd,
}).loc[[(cfg, approach) for cfg in CONFIGS for approach in ["independent", "joint"]]]
checkpoint3_table["ratio_torch_over_sklearn"] = (
    checkpoint3_table["torch_hybrid_dist_mm"] / checkpoint3_table["sklearn_hybrid_dist_mm"]
)
checkpoint3_table["status"] = np.where(
    checkpoint3_table["ratio_torch_over_sklearn"] <= 1 + PASS_TOL, "OK", "REGRESSION"
)
display(checkpoint3_table.round(4))

n_ok = (checkpoint3_table["status"] == "OK").sum()
print(f"\nCheckpoint 3: {n_ok}/{len(checkpoint3_table)} (dataset, approach) pairs within "
      f"{PASS_TOL:.0%} of the frozen sklearn hybrid_dist_test.")
if n_ok < len(checkpoint3_table):
    print("REGRESSION(S) FLAGGED -- reported as-is, not smoothed over:")
    print(checkpoint3_table[checkpoint3_table["status"] == "REGRESSION"])
sklearn_per_dataset_exp = mlflow.get_experiment_by_name("per_dataset")
assert sklearn_per_dataset_exp is not None, "run 4.2-Per_Dataset_Pipeline.ipynb first to populate this experiment"
sklearn_pd_runs = mlflow.search_runs(experiment_ids=[sklearn_per_dataset_exp.experiment_id])
sklearn_pd = sklearn_pd_runs.set_index(["params.dataset", "params.approach"])["metrics.hybrid_dist_test"].astype(float)

torch_pd = torch_results_df.set_index(["dataset", "approach"])["hybrid_dist_test"]

checkpoint3_table = pd.DataFrame({
    "sklearn_hybrid_dist_mm": sklearn_pd,
    "torch_hybrid_dist_mm": torch_pd,
}).loc[[(cfg, approach) for cfg in CONFIGS for approach in ["independent", "joint"]]]
checkpoint3_table["ratio_torch_over_sklearn"] = (
    checkpoint3_table["torch_hybrid_dist_mm"] / checkpoint3_table["sklearn_hybrid_dist_mm"]
)
checkpoint3_table["status"] = np.where(
    checkpoint3_table["ratio_torch_over_sklearn"] <= 1 + PASS_TOL, "OK", "REGRESSION"
)
display(checkpoint3_table.round(4))

n_ok = (checkpoint3_table["status"] == "OK").sum()
print(f"\nCheckpoint 3: {n_ok}/{len(checkpoint3_table)} (dataset, approach) pairs within "
      f"{PASS_TOL:.0%} of the frozen sklearn hybrid_dist_test.")
if n_ok < len(checkpoint3_table):
    print("REGRESSION(S) FLAGGED -- reported as-is, not smoothed over:")
    print(checkpoint3_table[checkpoint3_table["status"] == "REGRESSION"])

		sklearn_hybrid_dist_mm	torch_hybrid_dist_mm	ratio_torch_over_sklearn	status
rgs000_0	independent	0.0538	0.0548	1.0184	OK
rgs000_0	joint	0.0560	0.0544	0.9699	OK
rgs000_m4	independent	0.0531	0.0498	0.9373	OK
rgs000_m4	joint	0.0531	0.0513	0.9664	OK
rgs000_p4	independent	0.0729	0.0788	1.0808	OK
rgs000_p4	joint	0.0703	0.0727	1.0343	OK
rgs180_0	independent	0.0650	0.0744	1.1452	OK
rgs180_0	joint	0.0688	0.0645	0.9378	OK
rgs180_m4	independent	0.0623	0.0727	1.1665	OK
rgs180_m4	joint	0.0582	0.0600	1.0313	OK
rgs180_p4	independent	0.0678	0.0675	0.9954	OK
rgs180_p4	joint	0.0715	0.0682	0.9530	OK

Checkpoint 3: 12/12 (dataset, approach) pairs within 25% of the frozen sklearn hybrid_dist_test.

5. Reproduce `4.3-Multi_Dataset_Fitting.ipynb`'s final hybrid step¶

4.3's Section 3 (Tier 3) produced the final joint physical fit -- global coll_f/cam_f (Tier 1), per-grism-instance A_deg/rho (Tier 2), per-dataset tilt_deg/offset_y_mm/offset_z_mm (Tier 3) -- written to models/joint_specific_fit_<cfg>.toml. Section 4 then trained the independent MLP(y)+MLP(z) residual (Phase 2's carried-forward default) on top. Reused here exactly as 4.3 does: no refitting of any physical tier, just reloading each config's frozen [fit.fixed] + [fit.result] from its TOML and swapping the residual step to PyTorch/Lightning.

predict_centroids_full (the coll_f/cam_f/material_n0-extended forward model) is copied verbatim from 4.3 -- it's explicitly notebook-local in that notebook too (not promoted to dispcraft/calibration.py pending validation, per Stage 4 Phase 3's status note), so this is reusing the same un-promoted code, not duplicating library logic.

In [12]:

Copied!





from dispcraft.optics import Camera, Collimator, Grating, Grism, Material, Prism

FULL_NOMINAL = {
    **model.nominal_params,
    "coll_f": model.coll.f,
    "cam_f": model.cam.f,
    "material_n0": model.material.n0,
}


def predict_centroids_full(y_nisp_mm, z_nisp_mm, wavelength_nm, theta, free_names, base_model, fixed=None):
    """Verbatim copy of 4.3-Multi_Dataset_Fitting.ipynb's forward-model
    extension -- allows coll_f, cam_f, material_n0 to vary on top of
    calibration.predict_centroids. Notebook-local there, stays notebook-local
    here for the same reason (not yet promoted to dispcraft/calibration.py)."""
    p = dict(FULL_NOMINAL)
    if fixed:
        p.update(fixed)
    p.update(dict(zip(free_names, theta)))

    material = Material(n0=p["material_n0"], k=base_model.material.k)
    prism = Prism(material=material, A=np.radians(p["A_deg"]))
    grating = Grating(m=base_model.m_order, rho=p["rho"])
    grism = Grism(prism, grating, tilt=np.radians(p["tilt_deg"]))
    coll, cam = Collimator(f=p["coll_f"]), Camera(f=p["cam_f"])

    pos_foc = np.stack([np.asarray(y_nisp_mm), np.asarray(z_nisp_mm)]) / 1000.0
    wavelength_um = np.asarray(wavelength_nm) / 1000.0
    angle_col = coll.forward(pos_foc)
    angle_gr = grism.forward(angle_col, wavelength_um)
    pos_cam = cam.forward(angle_gr)
    offset_mm = np.array([p["offset_y_mm"], p["offset_z_mm"]])
    return pos_cam * 1000.0 + offset_mm[:, None]


# Same self-check 4.3 runs: at nominal coll_f/cam_f/material_n0, must agree
# with dispcraft.calibration.predict_centroids exactly.
from dispcraft.calibration import cost as cost_orig, predict_centroids

_df_check = median_per_spectrum(load_spectra(DATA_DIR / "rgs000_0_first.csv"))
_free_check = ["offset_y_mm", "offset_z_mm", "tilt_deg", "rho"]
_theta_check = np.array([FULL_NOMINAL[k] for k in _free_check])
_pred_orig = predict_centroids(_df_check["y_nisp"], _df_check["z_nisp"], _df_check["wavelength"],
                                _theta_check, _free_check, model)
_pred_full = predict_centroids_full(_df_check["y_nisp"], _df_check["z_nisp"], _df_check["wavelength"],
                                     _theta_check, _free_check, model)
np.testing.assert_allclose(_pred_orig, _pred_full)
print("predict_centroids_full agrees with dispcraft.calibration.predict_centroids at nominal coll/cam/material_n0")
from dispcraft.optics import Camera, Collimator, Grating, Grism, Material, Prism

FULL_NOMINAL = {
    **model.nominal_params,
    "coll_f": model.coll.f,
    "cam_f": model.cam.f,
    "material_n0": model.material.n0,
}


def predict_centroids_full(y_nisp_mm, z_nisp_mm, wavelength_nm, theta, free_names, base_model, fixed=None):
    """Verbatim copy of 4.3-Multi_Dataset_Fitting.ipynb's forward-model
    extension -- allows coll_f, cam_f, material_n0 to vary on top of
    calibration.predict_centroids. Notebook-local there, stays notebook-local
    here for the same reason (not yet promoted to dispcraft/calibration.py)."""
    p = dict(FULL_NOMINAL)
    if fixed:
        p.update(fixed)
    p.update(dict(zip(free_names, theta)))

    material = Material(n0=p["material_n0"], k=base_model.material.k)
    prism = Prism(material=material, A=np.radians(p["A_deg"]))
    grating = Grating(m=base_model.m_order, rho=p["rho"])
    grism = Grism(prism, grating, tilt=np.radians(p["tilt_deg"]))
    coll, cam = Collimator(f=p["coll_f"]), Camera(f=p["cam_f"])

    pos_foc = np.stack([np.asarray(y_nisp_mm), np.asarray(z_nisp_mm)]) / 1000.0
    wavelength_um = np.asarray(wavelength_nm) / 1000.0
    angle_col = coll.forward(pos_foc)
    angle_gr = grism.forward(angle_col, wavelength_um)
    pos_cam = cam.forward(angle_gr)
    offset_mm = np.array([p["offset_y_mm"], p["offset_z_mm"]])
    return pos_cam * 1000.0 + offset_mm[:, None]


# Same self-check 4.3 runs: at nominal coll_f/cam_f/material_n0, must agree
# with dispcraft.calibration.predict_centroids exactly.
from dispcraft.calibration import cost as cost_orig, predict_centroids

_df_check = median_per_spectrum(load_spectra(DATA_DIR / "rgs000_0_first.csv"))
_free_check = ["offset_y_mm", "offset_z_mm", "tilt_deg", "rho"]
_theta_check = np.array([FULL_NOMINAL[k] for k in _free_check])
_pred_orig = predict_centroids(_df_check["y_nisp"], _df_check["z_nisp"], _df_check["wavelength"],
                                _theta_check, _free_check, model)
_pred_full = predict_centroids_full(_df_check["y_nisp"], _df_check["z_nisp"], _df_check["wavelength"],
                                     _theta_check, _free_check, model)
np.testing.assert_allclose(_pred_orig, _pred_full)
print("predict_centroids_full agrees with dispcraft.calibration.predict_centroids at nominal coll/cam/material_n0")

predict_centroids_full agrees with dispcraft.calibration.predict_centroids at nominal coll/cam/material_n0

Reload the frozen Tier-3 joint physical fits¶

models/joint_specific_fit_<cfg>.toml's [fit.fixed] (coll_f, cam_f, A_deg, rho, material_n0) plus [fit.result] (tilt_deg, offset_y_mm, offset_z_mm) together give exactly fixed_final_by_cfg[cfg] from 4.3's Section 4 -- no Tier 1/2/3 refitting machinery needs to be reproduced here, since 4.3 already wrote its final per-dataset state into these TOMLs.

In [13]:

Copied!





def load_final_fixed(cfg):
    with open(MODELS_DIR / f"joint_specific_fit_{cfg}.toml", "rb") as f:
        fit = tomllib.load(f)
    fixed = dict(fit["fit"]["fixed"])  # coll_f, cam_f, A_deg, rho, material_n0, material_k
    free_names = fit["fit"]["free_parameters"]  # tilt_deg, offset_y_mm, offset_z_mm
    fixed.update({k: fit["fit"]["result"][k] for k in free_names})
    return fixed


fixed_final_by_cfg = {cfg: load_final_fixed(cfg) for cfg in CONFIGS}
print("fixed_final_by_cfg[\'rgs000_0\']:", fixed_final_by_cfg["rgs000_0"])
def load_final_fixed(cfg):
    with open(MODELS_DIR / f"joint_specific_fit_{cfg}.toml", "rb") as f:
        fit = tomllib.load(f)
    fixed = dict(fit["fit"]["fixed"])  # coll_f, cam_f, A_deg, rho, material_n0, material_k
    free_names = fit["fit"]["free_parameters"]  # tilt_deg, offset_y_mm, offset_z_mm
    fixed.update({k: fit["fit"]["result"][k] for k in free_names})
    return fixed


fixed_final_by_cfg = {cfg: load_final_fixed(cfg) for cfg in CONFIGS}
print("fixed_final_by_cfg[\'rgs000_0\']:", fixed_final_by_cfg["rgs000_0"])

fixed_final_by_cfg['rgs000_0']: {'coll_f': 1.999278, 'cam_f': 1.000415, 'A_deg': 2.143503, 'rho': 13.074063, 'material_n0': 1.44, 'material_k': 0.004, 'tilt_deg': 0.128979, 'offset_y_mm': -0.549367, 'offset_z_mm': 0.735745}

Run the final hybrid pipeline (PyTorch/Lightning)¶

Same structure as 4.3's run_stage4_pipeline: predict with predict_centroids_full at the frozen Tier-1+2+3 parameters (no free theta), compute residuals, group-aware split, train independent MLP(y)+MLP(z) (only candidate 4.3 carries forward here) via train_residual_mlp, and log to a new stage4_torch MLflow experiment.

In [14]:

Copied!





STAGE4_TORCH_EXPERIMENT = "stage4_torch"
if mlflow.get_experiment_by_name(STAGE4_TORCH_EXPERIMENT) is None:
    mlflow.create_experiment(STAGE4_TORCH_EXPERIMENT, artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(STAGE4_TORCH_EXPERIMENT)


def run_stage4_pipeline_torch(cfg):
    df = dfs_raw[cfg]
    fixed = fixed_final_by_cfg[cfg]
    pred = predict_centroids_full(df["y_nisp"], df["z_nisp"], df["wavelength"], np.array([]), [], model, fixed=fixed)
    d = df.assign(r_y=pred[0] - df["cent_y"].values, r_z=pred[1] - df["cent_z"].values)

    X = d[["y_nisp", "z_nisp", "wavelength"]].values
    groups = d["spectra_id"].values
    gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
    train_idx, test_idx = next(gss.split(X, groups=groups))

    scaler = StandardScaler().fit(X[train_idx])
    X_train_s, X_test_s = scaler.transform(X[train_idx]), scaler.transform(X[test_idx])
    r_y_train, r_z_train = d["r_y"].values[train_idx], d["r_z"].values[train_idx]
    r_y_test, r_z_test = d["r_y"].values[test_idx], d["r_z"].values[test_idx]

    physical_dist_test = float(np.mean(np.hypot(r_y_test, r_z_test)))

    lit_y = train_residual_mlp(X_train_s, r_y_train, mlp_y_params, n_outputs=1)
    lit_z = train_residual_mlp(X_train_s, r_z_train, mlp_z_params, n_outputs=1)
    hyb_y = r_y_test - predict(lit_y, X_test_s).ravel()
    hyb_z = r_z_test - predict(lit_z, X_test_s).ravel()

    return {
        "cfg": cfg,
        "physical_dist_test": physical_dist_test,
        "hybrid_rmse_y": rmse(r_y_test, predict(lit_y, X_test_s).ravel()),
        "hybrid_rmse_z": rmse(r_z_test, predict(lit_z, X_test_s).ravel()),
        "hybrid_dist_test": float(np.mean(np.hypot(hyb_y, hyb_z))),
    }


# dfs_raw: same excluded-spectra-filtered per-config dataframes 4.3 loads
# (before any residual is computed -- predict_centroids_full needs the raw
# y_nisp/z_nisp/wavelength/cent_y/cent_z columns, not Checkpoint 3's
# Stage-3-fit residuals).
dfs_raw = {}
for cfg in CONFIGS:
    with open(MODELS_DIR / f"stage3_fit_{cfg}.toml", "rb") as f:
        fit = tomllib.load(f)
    excluded_ids = fit["fit"]["excluded_spectra_ids"]
    df = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
    df = df[~df["spectra_id"].isin(excluded_ids)].reset_index(drop=True)
    assert len(df) == fit["fit"]["n_points"], f"{cfg}: row count must match the TOML's recorded fit"
    dfs_raw[cfg] = df

stage4_torch_results = {}
for cfg in CONFIGS:
    res = run_stage4_pipeline_torch(cfg)
    stage4_torch_results[cfg] = res
    improvement = 100 * (1 - res["hybrid_dist_test"] / res["physical_dist_test"])

    with mlflow.start_run(run_name=f"{cfg}_hybrid_torch"):
        mlflow.log_param("dataset", cfg)
        mlflow.log_param("framework", "pytorch")
        for k, v in fixed_final_by_cfg[cfg].items():
            mlflow.log_param(k, v)
        for axis, p in [("y", mlp_y_params), ("z", mlp_z_params)]:
            for k, v in p.items():
                mlflow.log_param(f"mlp_{axis}_{k}", v)
        mlflow.log_metric("physical_dist_test", res["physical_dist_test"])
        mlflow.log_metric("hybrid_dist_test", res["hybrid_dist_test"])
        mlflow.log_metric("hybrid_rmse_y", res["hybrid_rmse_y"])
        mlflow.log_metric("hybrid_rmse_z", res["hybrid_rmse_z"])
        mlflow.log_metric("improvement_pct", improvement)

stage4_torch_df = pd.DataFrame(stage4_torch_results.values()).set_index("cfg").loc[CONFIGS]
stage4_torch_df["improvement_pct"] = 100 * (1 - stage4_torch_df["hybrid_dist_test"] / stage4_torch_df["physical_dist_test"])
print(f"logged {len(CONFIGS)} runs to the '{STAGE4_TORCH_EXPERIMENT}' MLflow experiment")
stage4_torch_df.round(4)
STAGE4_TORCH_EXPERIMENT = "stage4_torch"
if mlflow.get_experiment_by_name(STAGE4_TORCH_EXPERIMENT) is None:
    mlflow.create_experiment(STAGE4_TORCH_EXPERIMENT, artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(STAGE4_TORCH_EXPERIMENT)


def run_stage4_pipeline_torch(cfg):
    df = dfs_raw[cfg]
    fixed = fixed_final_by_cfg[cfg]
    pred = predict_centroids_full(df["y_nisp"], df["z_nisp"], df["wavelength"], np.array([]), [], model, fixed=fixed)
    d = df.assign(r_y=pred[0] - df["cent_y"].values, r_z=pred[1] - df["cent_z"].values)

    X = d[["y_nisp", "z_nisp", "wavelength"]].values
    groups = d["spectra_id"].values
    gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
    train_idx, test_idx = next(gss.split(X, groups=groups))

    scaler = StandardScaler().fit(X[train_idx])
    X_train_s, X_test_s = scaler.transform(X[train_idx]), scaler.transform(X[test_idx])
    r_y_train, r_z_train = d["r_y"].values[train_idx], d["r_z"].values[train_idx]
    r_y_test, r_z_test = d["r_y"].values[test_idx], d["r_z"].values[test_idx]

    physical_dist_test = float(np.mean(np.hypot(r_y_test, r_z_test)))

    lit_y = train_residual_mlp(X_train_s, r_y_train, mlp_y_params, n_outputs=1)
    lit_z = train_residual_mlp(X_train_s, r_z_train, mlp_z_params, n_outputs=1)
    hyb_y = r_y_test - predict(lit_y, X_test_s).ravel()
    hyb_z = r_z_test - predict(lit_z, X_test_s).ravel()

    return {
        "cfg": cfg,
        "physical_dist_test": physical_dist_test,
        "hybrid_rmse_y": rmse(r_y_test, predict(lit_y, X_test_s).ravel()),
        "hybrid_rmse_z": rmse(r_z_test, predict(lit_z, X_test_s).ravel()),
        "hybrid_dist_test": float(np.mean(np.hypot(hyb_y, hyb_z))),
    }


# dfs_raw: same excluded-spectra-filtered per-config dataframes 4.3 loads
# (before any residual is computed -- predict_centroids_full needs the raw
# y_nisp/z_nisp/wavelength/cent_y/cent_z columns, not Checkpoint 3's
# Stage-3-fit residuals).
dfs_raw = {}
for cfg in CONFIGS:
    with open(MODELS_DIR / f"stage3_fit_{cfg}.toml", "rb") as f:
        fit = tomllib.load(f)
    excluded_ids = fit["fit"]["excluded_spectra_ids"]
    df = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
    df = df[~df["spectra_id"].isin(excluded_ids)].reset_index(drop=True)
    assert len(df) == fit["fit"]["n_points"], f"{cfg}: row count must match the TOML's recorded fit"
    dfs_raw[cfg] = df

stage4_torch_results = {}
for cfg in CONFIGS:
    res = run_stage4_pipeline_torch(cfg)
    stage4_torch_results[cfg] = res
    improvement = 100 * (1 - res["hybrid_dist_test"] / res["physical_dist_test"])

    with mlflow.start_run(run_name=f"{cfg}_hybrid_torch"):
        mlflow.log_param("dataset", cfg)
        mlflow.log_param("framework", "pytorch")
        for k, v in fixed_final_by_cfg[cfg].items():
            mlflow.log_param(k, v)
        for axis, p in [("y", mlp_y_params), ("z", mlp_z_params)]:
            for k, v in p.items():
                mlflow.log_param(f"mlp_{axis}_{k}", v)
        mlflow.log_metric("physical_dist_test", res["physical_dist_test"])
        mlflow.log_metric("hybrid_dist_test", res["hybrid_dist_test"])
        mlflow.log_metric("hybrid_rmse_y", res["hybrid_rmse_y"])
        mlflow.log_metric("hybrid_rmse_z", res["hybrid_rmse_z"])
        mlflow.log_metric("improvement_pct", improvement)

stage4_torch_df = pd.DataFrame(stage4_torch_results.values()).set_index("cfg").loc[CONFIGS]
stage4_torch_df["improvement_pct"] = 100 * (1 - stage4_torch_df["hybrid_dist_test"] / stage4_torch_df["physical_dist_test"])
print(f"logged {len(CONFIGS)} runs to the '{STAGE4_TORCH_EXPERIMENT}' MLflow experiment")
stage4_torch_df.round(4)

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

GPU available: False, used: False

TPU available: False, using: 0 TPU cores

💡 Tip: For seamless cloud logging and experiment tracking, try installing [litlogger](https://pypi.org/project/litlogger/) to enable LitLogger, which logs metrics and artifacts automatically to the Lightning Experiments platform.

logged 6 runs to the 'stage4_torch' MLflow experiment

Out[14]:

	physical_dist_test	hybrid_rmse_y	hybrid_rmse_z	hybrid_dist_test	improvement_pct
cfg
rgs000_0	0.3019	0.0125	0.0696	0.0584	80.6445
rgs000_m4	0.3022	0.0090	0.0606	0.0481	84.0820
rgs000_p4	0.3458	0.0138	0.0834	0.0680	80.3229
rgs180_0	0.3224	0.0129	0.0875	0.0663	79.4276
rgs180_m4	0.3394	0.0123	0.0832	0.0629	81.4582
rgs180_p4	0.3181	0.0118	0.0835	0.0640	79.8788

Checkpoint 4 comparison: sklearn (`stage4`) vs. PyTorch/Lightning (`stage4_torch`)¶

In [15]:

Copied!





sklearn_stage4_exp = mlflow.get_experiment_by_name("stage4")
assert sklearn_stage4_exp is not None, "run 4.3-Multi_Dataset_Fitting.ipynb first to populate this experiment"
sklearn_s4_runs = mlflow.search_runs(experiment_ids=[sklearn_stage4_exp.experiment_id])
sklearn_s4 = sklearn_s4_runs.set_index("params.dataset")["metrics.hybrid_dist_test"].astype(float)

checkpoint4_table = pd.DataFrame({
    "sklearn_hybrid_dist_mm": sklearn_s4,
    "torch_hybrid_dist_mm": stage4_torch_df["hybrid_dist_test"],
}).loc[CONFIGS]
checkpoint4_table["ratio_torch_over_sklearn"] = (
    checkpoint4_table["torch_hybrid_dist_mm"] / checkpoint4_table["sklearn_hybrid_dist_mm"]
)
checkpoint4_table["status"] = np.where(
    checkpoint4_table["ratio_torch_over_sklearn"] <= 1 + PASS_TOL, "OK", "REGRESSION"
)
display(checkpoint4_table.round(4))

n_ok = (checkpoint4_table["status"] == "OK").sum()
print(f"\nCheckpoint 4: {n_ok}/{len(checkpoint4_table)} configs within {PASS_TOL:.0%} "
      f"of the frozen sklearn hybrid_dist_test.")
if n_ok < len(checkpoint4_table):
    print("REGRESSION(S) FLAGGED -- reported as-is, not smoothed over:")
    print(checkpoint4_table[checkpoint4_table["status"] == "REGRESSION"])
sklearn_stage4_exp = mlflow.get_experiment_by_name("stage4")
assert sklearn_stage4_exp is not None, "run 4.3-Multi_Dataset_Fitting.ipynb first to populate this experiment"
sklearn_s4_runs = mlflow.search_runs(experiment_ids=[sklearn_stage4_exp.experiment_id])
sklearn_s4 = sklearn_s4_runs.set_index("params.dataset")["metrics.hybrid_dist_test"].astype(float)

checkpoint4_table = pd.DataFrame({
    "sklearn_hybrid_dist_mm": sklearn_s4,
    "torch_hybrid_dist_mm": stage4_torch_df["hybrid_dist_test"],
}).loc[CONFIGS]
checkpoint4_table["ratio_torch_over_sklearn"] = (
    checkpoint4_table["torch_hybrid_dist_mm"] / checkpoint4_table["sklearn_hybrid_dist_mm"]
)
checkpoint4_table["status"] = np.where(
    checkpoint4_table["ratio_torch_over_sklearn"] <= 1 + PASS_TOL, "OK", "REGRESSION"
)
display(checkpoint4_table.round(4))

n_ok = (checkpoint4_table["status"] == "OK").sum()
print(f"\nCheckpoint 4: {n_ok}/{len(checkpoint4_table)} configs within {PASS_TOL:.0%} "
      f"of the frozen sklearn hybrid_dist_test.")
if n_ok < len(checkpoint4_table):
    print("REGRESSION(S) FLAGGED -- reported as-is, not smoothed over:")
    print(checkpoint4_table[checkpoint4_table["status"] == "REGRESSION"])

	sklearn_hybrid_dist_mm	torch_hybrid_dist_mm	ratio_torch_over_sklearn	status
rgs000_0	0.0551	0.0584	1.0598	OK
rgs000_m4	0.0532	0.0481	0.9049	OK
rgs000_p4	0.0699	0.0680	0.9727	OK
rgs180_0	0.0627	0.0663	1.0579	OK
rgs180_m4	0.0638	0.0629	0.9859	OK
rgs180_p4	0.0708	0.0640	0.9046	OK

Checkpoint 4: 6/6 configs within 25% of the frozen sklearn hybrid_dist_test.

6. Summary: go/no-go verdict¶

All three checkpoints together -- the single-dataset comparison (4.1), the per-dataset pipeline across all 6 configs and both candidates (4.2), and the final joint-fit hybrid step (4.3) -- against the same PASS_TOL tolerance throughout.

In [16]:

Copied!





def _as_case_status(table, checkpoint_name):
    idx = table.index.to_series().apply(lambda v: " / ".join(v) if isinstance(v, tuple) else str(v))
    return pd.DataFrame({"checkpoint": checkpoint_name, "case": idx.values, "status": table["status"].values})


all_checks = pd.concat([
    _as_case_status(checkpoint2_table, "4.1 single-dataset"),
    _as_case_status(checkpoint3_table, "4.2 per-dataset"),
    _as_case_status(checkpoint4_table, "4.3 final hybrid"),
], ignore_index=True)

summary = all_checks.groupby("checkpoint")["status"].apply(lambda s: f"{(s == 'OK').sum()}/{len(s)} OK")
print(summary.to_string())

n_total_ok = (all_checks["status"] == "OK").sum()
n_total = len(all_checks)
print(f"\nOverall: {n_total_ok}/{n_total} checks within {PASS_TOL:.0%} of the frozen sklearn reference.")
if n_total_ok == n_total:
    print("\nGO: the PyTorch/Lightning residual-MLP port reproduces Stage 4's frozen sklearn results "
          "within tolerance across all three checkpoints. Safe to build Stage 5 Phase 2's "
          "field-dependent parameter model on top of this framework.")
else:
    failing = all_checks[all_checks["status"] != "OK"]
    print(f"\nNO-GO: {len(failing)} case(s) still regress beyond {PASS_TOL:.0%}, reported as-is:")
    print(failing.to_string(index=False))
def _as_case_status(table, checkpoint_name):
    idx = table.index.to_series().apply(lambda v: " / ".join(v) if isinstance(v, tuple) else str(v))
    return pd.DataFrame({"checkpoint": checkpoint_name, "case": idx.values, "status": table["status"].values})


all_checks = pd.concat([
    _as_case_status(checkpoint2_table, "4.1 single-dataset"),
    _as_case_status(checkpoint3_table, "4.2 per-dataset"),
    _as_case_status(checkpoint4_table, "4.3 final hybrid"),
], ignore_index=True)

summary = all_checks.groupby("checkpoint")["status"].apply(lambda s: f"{(s == 'OK').sum()}/{len(s)} OK")
print(summary.to_string())

n_total_ok = (all_checks["status"] == "OK").sum()
n_total = len(all_checks)
print(f"\nOverall: {n_total_ok}/{n_total} checks within {PASS_TOL:.0%} of the frozen sklearn reference.")
if n_total_ok == n_total:
    print("\nGO: the PyTorch/Lightning residual-MLP port reproduces Stage 4's frozen sklearn results "
          "within tolerance across all three checkpoints. Safe to build Stage 5 Phase 2's "
          "field-dependent parameter model on top of this framework.")
else:
    failing = all_checks[all_checks["status"] != "OK"]
    print(f"\nNO-GO: {len(failing)} case(s) still regress beyond {PASS_TOL:.0%}, reported as-is:")
    print(failing.to_string(index=False))

checkpoint
4.1 single-dataset      3/3 OK
4.2 per-dataset       12/12 OK
4.3 final hybrid        6/6 OK

Overall: 21/21 checks within 25% of the frozen sklearn reference.

GO: the PyTorch/Lightning residual-MLP port reproduces Stage 4's frozen sklearn results within tolerance across all three checkpoints. Safe to build Stage 5 Phase 2's field-dependent parameter model on top of this framework.

Stage 5, Phase 1 — PyTorch/Lightning Migration (5.1-PyTorch_Migration)¶