Stage 6 — Status Report: Final Models and Results¶
Assembles the project's final calibration pipeline end to end and computes
the numbers quoted in Status_Report.md. This notebook does not
introduce new methodology: every component reused here was already
validated in an earlier notebook. Two Stage 5 artifacts (the field-dependent
tier, the shared material_k) were never frozen to models/*.toml, so they
are refit here, deterministically (fixed seeds, unchanged hyperparameters),
rather than re-explored.
New here: Stage 5's field-dependent tier and Stage 4's ML residual
corrector were validated separately but never combined, and never checked
against the reference paper's held-out benchmark the way Phase 5
(4.5) checked the pre-Stage-5 pipeline. Sections 4-5 do that, refitting
the field-dependent tier on a train split only (5.2's own Section 9 fit
it on the full dataset, so its 76-83% figure is in-sample) so the held-out
RMSE reported here is genuinely comparable to Sections 2-3 and to the
paper's Argon benchmark.
Setup¶
In [1]:
Copied!
import tomllib
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import mlflow
from sklearn.model_selection import GroupShuffleSplit
from sklearn.preprocessing import StandardScaler
from sklearn.neural_network import MLPRegressor
from sklearn.metrics import mean_squared_error
from sklearn.linear_model import LinearRegression
from dispcraft.calibration import ground_test_model_from_config, predict_centroids
from dispcraft.measurement import load_spectra, median_per_spectrum, zeroth_order_separation
import dispcraft.field_calibration as fc
import dispcraft.zeroth_dispersion as zd
REPO_ROOT = Path("..")
DATA_DIR = REPO_ROOT / "data"
MODELS_DIR = REPO_ROOT / "models"
FIGURES_DIR = REPO_ROOT / "notebooks" / "figures" / "stage6"
FIGURES_DIR.mkdir(parents=True, exist_ok=True)
RNG_SEED = 42 # same seed as 4.1/4.3/4.5/5.2/5.3 -- reproduces their splits/fits exactly
CONFIGS = ["rgs000_0", "rgs000_m4", "rgs000_p4", "rgs180_0", "rgs180_m4", "rgs180_p4"]
CONFIGS_ZEROTH = CONFIGS + ["bgs000_0"]
GRISM_INSTANCES = {
"rgs000": ["rgs000_0", "rgs000_m4", "rgs000_p4"],
"rgs180": ["rgs180_0", "rgs180_m4", "rgs180_p4"],
}
with open(MODELS_DIR / "stage1_instrument.toml", "rb") as f:
base_config = tomllib.load(f)
model = ground_test_model_from_config(base_config)
mlflow.set_tracking_uri(f"sqlite:///{(REPO_ROOT / 'mlflow.db').resolve()}")
import tomllib
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import mlflow
from sklearn.model_selection import GroupShuffleSplit
from sklearn.preprocessing import StandardScaler
from sklearn.neural_network import MLPRegressor
from sklearn.metrics import mean_squared_error
from sklearn.linear_model import LinearRegression
from dispcraft.calibration import ground_test_model_from_config, predict_centroids
from dispcraft.measurement import load_spectra, median_per_spectrum, zeroth_order_separation
import dispcraft.field_calibration as fc
import dispcraft.zeroth_dispersion as zd
REPO_ROOT = Path("..")
DATA_DIR = REPO_ROOT / "data"
MODELS_DIR = REPO_ROOT / "models"
FIGURES_DIR = REPO_ROOT / "notebooks" / "figures" / "stage6"
FIGURES_DIR.mkdir(parents=True, exist_ok=True)
RNG_SEED = 42 # same seed as 4.1/4.3/4.5/5.2/5.3 -- reproduces their splits/fits exactly
CONFIGS = ["rgs000_0", "rgs000_m4", "rgs000_p4", "rgs180_0", "rgs180_m4", "rgs180_p4"]
CONFIGS_ZEROTH = CONFIGS + ["bgs000_0"]
GRISM_INSTANCES = {
"rgs000": ["rgs000_0", "rgs000_m4", "rgs000_p4"],
"rgs180": ["rgs180_0", "rgs180_m4", "rgs180_p4"],
}
with open(MODELS_DIR / "stage1_instrument.toml", "rb") as f:
base_config = tomllib.load(f)
model = ground_test_model_from_config(base_config)
mlflow.set_tracking_uri(f"sqlite:///{(REPO_ROOT / 'mlflow.db').resolve()}")
/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
1. Load frozen calibration results (Stage 3 independent, Stage 4 Tier 1-3 joint)¶
Reuses 4.5's loading pattern exactly: Stage 3's per-dataset fits and Stage
4 Phase 3's joint fits, both already frozen in models/*.toml, no
refitting.
In [2]:
Copied!
def load_toml(path):
with open(path, "rb") as f:
return tomllib.load(f)
def numeric_result(fit_result):
return {k: v for k, v in fit_result.items() if not k.endswith("_bootstrap_std")}
dfs = {}
independent_fits = {}
joint_fits = {}
for cfg in CONFIGS:
indep = load_toml(MODELS_DIR / f"stage3_fit_{cfg}.toml")
joint = load_toml(MODELS_DIR / f"joint_specific_fit_{cfg}.toml")
df = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
df = df[~df["spectra_id"].isin(indep["fit"]["excluded_spectra_ids"])].reset_index(drop=True)
assert len(df) == indep["fit"]["n_points"], f"{cfg}: row count must match stage3_fit_{cfg}.toml"
dfs[cfg] = df
independent_fits[cfg] = {
"free_names": indep["fit"]["free_parameters"],
"theta": np.array([indep["fit"]["result"][k] for k in indep["fit"]["free_parameters"]]),
}
joint_fits[cfg] = {"fixed": {**joint["fit"]["fixed"], **numeric_result(joint["fit"]["result"])}}
print(f"loaded {len(dfs)} datasets, row counts: {[len(dfs[c]) for c in CONFIGS]}")
def load_toml(path):
with open(path, "rb") as f:
return tomllib.load(f)
def numeric_result(fit_result):
return {k: v for k, v in fit_result.items() if not k.endswith("_bootstrap_std")}
dfs = {}
independent_fits = {}
joint_fits = {}
for cfg in CONFIGS:
indep = load_toml(MODELS_DIR / f"stage3_fit_{cfg}.toml")
joint = load_toml(MODELS_DIR / f"joint_specific_fit_{cfg}.toml")
df = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
df = df[~df["spectra_id"].isin(indep["fit"]["excluded_spectra_ids"])].reset_index(drop=True)
assert len(df) == indep["fit"]["n_points"], f"{cfg}: row count must match stage3_fit_{cfg}.toml"
dfs[cfg] = df
independent_fits[cfg] = {
"free_names": indep["fit"]["free_parameters"],
"theta": np.array([indep["fit"]["result"][k] for k in indep["fit"]["free_parameters"]]),
}
joint_fits[cfg] = {"fixed": {**joint["fit"]["fixed"], **numeric_result(joint["fit"]["result"])}}
print(f"loaded {len(dfs)} datasets, row counts: {[len(dfs[c]) for c in CONFIGS]}")
loaded 6 datasets, row counts: [5297, 4772, 4577, 6198, 4596, 4604]
2. Held-out physical-model RMSE: independent vs. joint (Stage 4 Tier 1-3)¶
In [3]:
Copied!
def test_split(df):
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
train_idx, test_idx = next(gss.split(df, groups=df["spectra_id"].values))
return df.iloc[train_idx], df.iloc[test_idx]
def physical_axis_rmse(df_test, theta, free_names, fixed=None):
pred = predict_centroids(df_test["y_nisp"], df_test["z_nisp"], df_test["wavelength"],
theta, free_names, model, fixed=fixed)
r_y = pred[0] - df_test["cent_y"].values
r_z = pred[1] - df_test["cent_z"].values
return float(np.sqrt(np.mean(r_y**2))), float(np.sqrt(np.mean(r_z**2)))
train_test = {cfg: test_split(dfs[cfg]) for cfg in CONFIGS}
physical_rows = []
for cfg in CONFIGS:
_, df_test = train_test[cfg]
iy, iz = physical_axis_rmse(df_test, independent_fits[cfg]["theta"], independent_fits[cfg]["free_names"])
jy, jz = physical_axis_rmse(df_test, np.array([]), [], fixed=joint_fits[cfg]["fixed"])
physical_rows.append({"cfg": cfg, "independent_rmse_y_mm": iy, "independent_rmse_z_mm": iz,
"joint_rmse_y_mm": jy, "joint_rmse_z_mm": jz})
physical_test = pd.DataFrame(physical_rows).set_index("cfg").loc[CONFIGS]
physical_test.round(5)
def test_split(df):
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
train_idx, test_idx = next(gss.split(df, groups=df["spectra_id"].values))
return df.iloc[train_idx], df.iloc[test_idx]
def physical_axis_rmse(df_test, theta, free_names, fixed=None):
pred = predict_centroids(df_test["y_nisp"], df_test["z_nisp"], df_test["wavelength"],
theta, free_names, model, fixed=fixed)
r_y = pred[0] - df_test["cent_y"].values
r_z = pred[1] - df_test["cent_z"].values
return float(np.sqrt(np.mean(r_y**2))), float(np.sqrt(np.mean(r_z**2)))
train_test = {cfg: test_split(dfs[cfg]) for cfg in CONFIGS}
physical_rows = []
for cfg in CONFIGS:
_, df_test = train_test[cfg]
iy, iz = physical_axis_rmse(df_test, independent_fits[cfg]["theta"], independent_fits[cfg]["free_names"])
jy, jz = physical_axis_rmse(df_test, np.array([]), [], fixed=joint_fits[cfg]["fixed"])
physical_rows.append({"cfg": cfg, "independent_rmse_y_mm": iy, "independent_rmse_z_mm": iz,
"joint_rmse_y_mm": jy, "joint_rmse_z_mm": jz})
physical_test = pd.DataFrame(physical_rows).set_index("cfg").loc[CONFIGS]
physical_test.round(5)
Out[3]:
| independent_rmse_y_mm | independent_rmse_z_mm | joint_rmse_y_mm | joint_rmse_z_mm | |
|---|---|---|---|---|
| cfg | ||||
| rgs000_0 | 0.14590 | 0.31715 | 0.15920 | 0.29802 |
| rgs000_m4 | 0.19935 | 0.27440 | 0.22662 | 0.26412 |
| rgs000_p4 | 0.17085 | 0.35708 | 0.18968 | 0.34455 |
| rgs180_0 | 0.17841 | 0.31507 | 0.20131 | 0.29223 |
| rgs180_m4 | 0.16886 | 0.33924 | 0.18399 | 0.31683 |
| rgs180_p4 | 0.17702 | 0.30215 | 0.20515 | 0.28717 |
3. Stage 4's frozen ML residual (pre-field), read back from MLflow¶
In [4]:
Copied!
def stage4_hybrid_rmse():
"""Per-axis, held-out RMSE of the joint physical fit + independent MLP(y)+MLP(z)
correction (Stage 4 Phase 3, Section 4) -- read back from MLflow, not retrained."""
exp = mlflow.get_experiment_by_name("stage4")
assert exp is not None, "run notebooks/4.3-Multi_Dataset_Fitting.ipynb first to populate the 'stage4' experiment"
runs = mlflow.search_runs(experiment_ids=[exp.experiment_id]).set_index("params.dataset")
return runs.loc[CONFIGS, ["metrics.hybrid_rmse_y", "metrics.hybrid_rmse_z"]].rename(
columns={"metrics.hybrid_rmse_y": "joint_ml_rmse_y_mm", "metrics.hybrid_rmse_z": "joint_ml_rmse_z_mm"}
).astype(float)
joint_ml_test = stage4_hybrid_rmse()
joint_ml_test
def stage4_hybrid_rmse():
"""Per-axis, held-out RMSE of the joint physical fit + independent MLP(y)+MLP(z)
correction (Stage 4 Phase 3, Section 4) -- read back from MLflow, not retrained."""
exp = mlflow.get_experiment_by_name("stage4")
assert exp is not None, "run notebooks/4.3-Multi_Dataset_Fitting.ipynb first to populate the 'stage4' experiment"
runs = mlflow.search_runs(experiment_ids=[exp.experiment_id]).set_index("params.dataset")
return runs.loc[CONFIGS, ["metrics.hybrid_rmse_y", "metrics.hybrid_rmse_z"]].rename(
columns={"metrics.hybrid_rmse_y": "joint_ml_rmse_y_mm", "metrics.hybrid_rmse_z": "joint_ml_rmse_z_mm"}
).astype(float)
joint_ml_test = stage4_hybrid_rmse()
joint_ml_test
Out[4]:
| joint_ml_rmse_y_mm | joint_ml_rmse_z_mm | |
|---|---|---|
| params.dataset | ||
| rgs000_0 | 0.017610 | 0.063612 |
| rgs000_m4 | 0.014082 | 0.061646 |
| rgs000_p4 | 0.017443 | 0.086905 |
| rgs180_0 | 0.016640 | 0.079652 |
| rgs180_m4 | 0.021537 | 0.074296 |
| rgs180_p4 | 0.023943 | 0.083087 |
4. Reconstructing Stage 5's field-dependent tier (train-only, held-out evaluation)¶
5.2 validated this architecture via 5-fold CV (Section 6-7) but then
trained the production tier (Section 9) on the full dataset per config --
its reported 76-83% RMSE reduction is in-sample. To compare fairly against
Sections 2-3 above and the paper's held-out benchmark (Section 7 below),
refit the same architecture (dispcraft.field_calibration .DEFAULT_FIELD_MLP_PARAMS, unchanged) on the train split only, warm-started
at the frozen tilt_deg, and evaluate on the held-out test split.
In [5]:
Copied!
def fixed_tier12_for(cfg):
"""Tier 1+2 physical params only -- offset_y_mm/offset_z_mm/tilt_deg
removed, per fit_field_dependent_tier's contract (they're replaced by
the field-dependent NN, not held fixed alongside it)."""
p = dict(joint_fits[cfg]["fixed"])
return {k: v for k, v in p.items() if k not in ("offset_y_mm", "offset_z_mm", "tilt_deg")}
field_results = {}
for cfg in CONFIGS:
df_train, df_test = train_test[cfg]
fixed_tier12 = fixed_tier12_for(cfg)
tilt_deg_init = joint_fits[cfg]["fixed"]["tilt_deg"]
tilt_deg, nn_model, history = fc.fit_field_dependent_tier(
df_train, model, fixed_tier12, tilt_deg_init, mlp_params=fc.DEFAULT_FIELD_MLP_PARAMS)
pred_test = fc.predict_centroids_field(df_test["y_nisp"], df_test["z_nisp"], df_test["wavelength"],
tilt_deg, model, fixed_tier12, nn_model)
r_y_test = pred_test[0] - df_test["cent_y"].values
r_z_test = pred_test[1] - df_test["cent_z"].values
field_results[cfg] = {
"tilt_deg": tilt_deg, "nn_model": nn_model, "fixed_tier12": fixed_tier12,
"n_iterations": len(history),
"r_y_test": r_y_test, "r_z_test": r_z_test,
"rmse_y_mm": float(np.sqrt(np.mean(r_y_test**2))),
"rmse_z_mm": float(np.sqrt(np.mean(r_z_test**2))),
}
print(f"{cfg}: converged in {len(history)} iteration(s), "
f"tilt_deg {tilt_deg_init:.4f} -> {tilt_deg:.4f}, "
f"held-out rmse_y={field_results[cfg]['rmse_y_mm']:.4f} mm, "
f"rmse_z={field_results[cfg]['rmse_z_mm']:.4f} mm")
field_test = pd.DataFrame({
cfg: {"joint_field_rmse_y_mm": field_results[cfg]["rmse_y_mm"],
"joint_field_rmse_z_mm": field_results[cfg]["rmse_z_mm"]}
for cfg in CONFIGS
}).T.loc[CONFIGS]
field_test.round(5)
def fixed_tier12_for(cfg):
"""Tier 1+2 physical params only -- offset_y_mm/offset_z_mm/tilt_deg
removed, per fit_field_dependent_tier's contract (they're replaced by
the field-dependent NN, not held fixed alongside it)."""
p = dict(joint_fits[cfg]["fixed"])
return {k: v for k, v in p.items() if k not in ("offset_y_mm", "offset_z_mm", "tilt_deg")}
field_results = {}
for cfg in CONFIGS:
df_train, df_test = train_test[cfg]
fixed_tier12 = fixed_tier12_for(cfg)
tilt_deg_init = joint_fits[cfg]["fixed"]["tilt_deg"]
tilt_deg, nn_model, history = fc.fit_field_dependent_tier(
df_train, model, fixed_tier12, tilt_deg_init, mlp_params=fc.DEFAULT_FIELD_MLP_PARAMS)
pred_test = fc.predict_centroids_field(df_test["y_nisp"], df_test["z_nisp"], df_test["wavelength"],
tilt_deg, model, fixed_tier12, nn_model)
r_y_test = pred_test[0] - df_test["cent_y"].values
r_z_test = pred_test[1] - df_test["cent_z"].values
field_results[cfg] = {
"tilt_deg": tilt_deg, "nn_model": nn_model, "fixed_tier12": fixed_tier12,
"n_iterations": len(history),
"r_y_test": r_y_test, "r_z_test": r_z_test,
"rmse_y_mm": float(np.sqrt(np.mean(r_y_test**2))),
"rmse_z_mm": float(np.sqrt(np.mean(r_z_test**2))),
}
print(f"{cfg}: converged in {len(history)} iteration(s), "
f"tilt_deg {tilt_deg_init:.4f} -> {tilt_deg:.4f}, "
f"held-out rmse_y={field_results[cfg]['rmse_y_mm']:.4f} mm, "
f"rmse_z={field_results[cfg]['rmse_z_mm']:.4f} mm")
field_test = pd.DataFrame({
cfg: {"joint_field_rmse_y_mm": field_results[cfg]["rmse_y_mm"],
"joint_field_rmse_z_mm": field_results[cfg]["rmse_z_mm"]}
for cfg in CONFIGS
}).T.loc[CONFIGS]
field_test.round(5)
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.
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.
rgs000_0: converged in 5 iteration(s), tilt_deg 0.1290 -> 0.1107, held-out rmse_y=0.0303 mm, rmse_z=0.0725 mm
/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.
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.
rgs000_m4: converged in 5 iteration(s), tilt_deg 4.1520 -> 4.1148, held-out rmse_y=0.0249 mm, rmse_z=0.0684 mm
/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.
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.
rgs000_p4: converged in 5 iteration(s), tilt_deg -3.7533 -> -3.7905, held-out rmse_y=0.0363 mm, rmse_z=0.0904 mm
/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.
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.
rgs180_0: converged in 5 iteration(s), tilt_deg 180.2389 -> 180.2335, held-out rmse_y=0.0360 mm, rmse_z=0.0800 mm
/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.
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.
rgs180_m4: converged in 5 iteration(s), tilt_deg 184.2012 -> 184.1735, held-out rmse_y=0.0478 mm, rmse_z=0.0977 mm
/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.
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.
rgs180_p4: converged in 5 iteration(s), tilt_deg 176.2084 -> 176.2215, held-out rmse_y=0.0432 mm, rmse_z=0.0867 mm
Out[5]:
| joint_field_rmse_y_mm | joint_field_rmse_z_mm | |
|---|---|---|
| rgs000_0 | 0.03027 | 0.07249 |
| rgs000_m4 | 0.02486 | 0.06844 |
| rgs000_p4 | 0.03634 | 0.09043 |
| rgs180_0 | 0.03597 | 0.08002 |
| rgs180_m4 | 0.04782 | 0.09772 |
| rgs180_p4 | 0.04316 | 0.08672 |
5. Layering Stage 4's ML residual on top of the field-dependent tier¶
Does the field-dependent tier make the generic wavelength+position MLP
residual redundant, or does it still add value? Retrain Stage 4's exact
winning MLP(y)+MLP(z) hyperparameters (residual_correction MLflow
experiment, same as 4.3 Section 4) on the post-field-tier train
residual, then evaluate held-out.
In [6]:
Copied!
import ast
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"
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(runs, "y", "metrics.test_rmse")
mlp_z_params = best_mlp_params(runs, "z", "metrics.test_rmse")
print("MLP(y):", mlp_y_params)
print("MLP(z):", mlp_z_params)
import ast
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"
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(runs, "y", "metrics.test_rmse")
mlp_z_params = best_mlp_params(runs, "z", "metrics.test_rmse")
print("MLP(y):", mlp_y_params)
print("MLP(z):", mlp_z_params)
MLP(y): {'hidden_layer_sizes': (64, 32), 'activation': 'relu', 'alpha': 0.0001, 'random_state': 42, 'max_iter': 2000, 'early_stopping': True}
MLP(z): {'hidden_layer_sizes': (64, 64), 'activation': 'relu', 'alpha': 0.001, 'random_state': 42, 'max_iter': 2000, 'early_stopping': True}
In [7]:
Copied!
field_ml_results = {}
for cfg in CONFIGS:
df_train, df_test = train_test[cfg]
fr = field_results[cfg]
pred_train = fc.predict_centroids_field(df_train["y_nisp"], df_train["z_nisp"], df_train["wavelength"],
fr["tilt_deg"], model, fr["fixed_tier12"], fr["nn_model"])
r_y_train = pred_train[0] - df_train["cent_y"].values
r_z_train = pred_train[1] - df_train["cent_z"].values
X_train = df_train[["y_nisp", "z_nisp", "wavelength"]].values
X_test = df_test[["y_nisp", "z_nisp", "wavelength"]].values
scaler = StandardScaler().fit(X_train)
X_train_s, X_test_s = scaler.transform(X_train), scaler.transform(X_test)
mlp_y = MLPRegressor(**mlp_y_params).fit(X_train_s, r_y_train)
mlp_z = MLPRegressor(**mlp_z_params).fit(X_train_s, r_z_train)
hyb_y_test = fr["r_y_test"] - mlp_y.predict(X_test_s)
hyb_z_test = fr["r_z_test"] - mlp_z.predict(X_test_s)
field_ml_results[cfg] = {
"rmse_y_mm": float(np.sqrt(mean_squared_error(fr["r_y_test"], mlp_y.predict(X_test_s)))),
"rmse_z_mm": float(np.sqrt(mean_squared_error(fr["r_z_test"], mlp_z.predict(X_test_s)))),
}
field_ml_test = pd.DataFrame({
cfg: {"joint_field_ml_rmse_y_mm": field_ml_results[cfg]["rmse_y_mm"],
"joint_field_ml_rmse_z_mm": field_ml_results[cfg]["rmse_z_mm"]}
for cfg in CONFIGS
}).T.loc[CONFIGS]
field_ml_test.round(5)
field_ml_results = {}
for cfg in CONFIGS:
df_train, df_test = train_test[cfg]
fr = field_results[cfg]
pred_train = fc.predict_centroids_field(df_train["y_nisp"], df_train["z_nisp"], df_train["wavelength"],
fr["tilt_deg"], model, fr["fixed_tier12"], fr["nn_model"])
r_y_train = pred_train[0] - df_train["cent_y"].values
r_z_train = pred_train[1] - df_train["cent_z"].values
X_train = df_train[["y_nisp", "z_nisp", "wavelength"]].values
X_test = df_test[["y_nisp", "z_nisp", "wavelength"]].values
scaler = StandardScaler().fit(X_train)
X_train_s, X_test_s = scaler.transform(X_train), scaler.transform(X_test)
mlp_y = MLPRegressor(**mlp_y_params).fit(X_train_s, r_y_train)
mlp_z = MLPRegressor(**mlp_z_params).fit(X_train_s, r_z_train)
hyb_y_test = fr["r_y_test"] - mlp_y.predict(X_test_s)
hyb_z_test = fr["r_z_test"] - mlp_z.predict(X_test_s)
field_ml_results[cfg] = {
"rmse_y_mm": float(np.sqrt(mean_squared_error(fr["r_y_test"], mlp_y.predict(X_test_s)))),
"rmse_z_mm": float(np.sqrt(mean_squared_error(fr["r_z_test"], mlp_z.predict(X_test_s)))),
}
field_ml_test = pd.DataFrame({
cfg: {"joint_field_ml_rmse_y_mm": field_ml_results[cfg]["rmse_y_mm"],
"joint_field_ml_rmse_z_mm": field_ml_results[cfg]["rmse_z_mm"]}
for cfg in CONFIGS
}).T.loc[CONFIGS]
field_ml_test.round(5)
Out[7]:
| joint_field_ml_rmse_y_mm | joint_field_ml_rmse_z_mm | |
|---|---|---|
| rgs000_0 | 0.01928 | 0.06244 |
| rgs000_m4 | 0.01586 | 0.06264 |
| rgs000_p4 | 0.01676 | 0.08538 |
| rgs180_0 | 0.01333 | 0.08462 |
| rgs180_m4 | 0.01325 | 0.07313 |
| rgs180_p4 | 0.01184 | 0.08359 |
6. RMSE progression across the whole pipeline¶
In [8]:
Copied!
progression = physical_test.join(joint_ml_test).join(field_test).join(field_ml_test)
for stage in ["independent", "joint", "joint_ml", "joint_field", "joint_field_ml"]:
progression[f"{stage}_combined_mm"] = np.hypot(
progression[f"{stage}_rmse_y_mm"], progression[f"{stage}_rmse_z_mm"])
combined_cols = [c for c in progression.columns if c.endswith("_combined_mm")]
progression_combined = progression[combined_cols].copy()
progression_combined.columns = [c.replace("_combined_mm", "") for c in combined_cols]
progression_combined.round(4)
progression = physical_test.join(joint_ml_test).join(field_test).join(field_ml_test)
for stage in ["independent", "joint", "joint_ml", "joint_field", "joint_field_ml"]:
progression[f"{stage}_combined_mm"] = np.hypot(
progression[f"{stage}_rmse_y_mm"], progression[f"{stage}_rmse_z_mm"])
combined_cols = [c for c in progression.columns if c.endswith("_combined_mm")]
progression_combined = progression[combined_cols].copy()
progression_combined.columns = [c.replace("_combined_mm", "") for c in combined_cols]
progression_combined.round(4)
Out[8]:
| independent | joint | joint_ml | joint_field | joint_field_ml | |
|---|---|---|---|---|---|
| cfg | |||||
| rgs000_0 | 0.3491 | 0.3379 | 0.0660 | 0.0786 | 0.0653 |
| rgs000_m4 | 0.3392 | 0.3480 | 0.0632 | 0.0728 | 0.0646 |
| rgs000_p4 | 0.3959 | 0.3933 | 0.0886 | 0.0975 | 0.0870 |
| rgs180_0 | 0.3621 | 0.3549 | 0.0814 | 0.0877 | 0.0857 |
| rgs180_m4 | 0.3789 | 0.3664 | 0.0774 | 0.1088 | 0.0743 |
| rgs180_p4 | 0.3502 | 0.3529 | 0.0865 | 0.0969 | 0.0844 |
In [9]:
Copied!
fig, ax = plt.subplots(figsize=(10, 5))
x = np.arange(len(CONFIGS))
width = 0.15
stages = ["independent", "joint", "joint_ml", "joint_field", "joint_field_ml"]
labels = ["independent\n(Stage 3)", "joint\n(Stage 4 T1-3)", "joint + ML\n(Stage 4)",
"joint + field\n(Stage 5)", "joint + field\n+ ML"]
for i, (stage, label) in enumerate(zip(stages, labels)):
ax.bar(x + (i - 2) * width, progression_combined[stage], width, label=label)
ax.set_xticks(x)
ax.set_xticklabels(CONFIGS, rotation=30, ha="right")
ax.set_ylabel("combined RMSE, held-out test split [mm]")
ax.legend(fontsize=8)
ax.set_title("RMSE progression across the full pipeline, per dataset")
plt.tight_layout()
plt.savefig(FIGURES_DIR / "rmse_progression.png", dpi=150)
plt.show()
print(f"mean combined RMSE: "
+ ", ".join(f"{s}={progression_combined[s].mean():.4f} mm" for s in stages))
fig, ax = plt.subplots(figsize=(10, 5))
x = np.arange(len(CONFIGS))
width = 0.15
stages = ["independent", "joint", "joint_ml", "joint_field", "joint_field_ml"]
labels = ["independent\n(Stage 3)", "joint\n(Stage 4 T1-3)", "joint + ML\n(Stage 4)",
"joint + field\n(Stage 5)", "joint + field\n+ ML"]
for i, (stage, label) in enumerate(zip(stages, labels)):
ax.bar(x + (i - 2) * width, progression_combined[stage], width, label=label)
ax.set_xticks(x)
ax.set_xticklabels(CONFIGS, rotation=30, ha="right")
ax.set_ylabel("combined RMSE, held-out test split [mm]")
ax.legend(fontsize=8)
ax.set_title("RMSE progression across the full pipeline, per dataset")
plt.tight_layout()
plt.savefig(FIGURES_DIR / "rmse_progression.png", dpi=150)
plt.show()
print(f"mean combined RMSE: "
+ ", ".join(f"{s}={progression_combined[s].mean():.4f} mm" for s in stages))
mean combined RMSE: independent=0.3626 mm, joint=0.3589 mm, joint_ml=0.0772 mm, joint_field=0.0904 mm, joint_field_ml=0.0769 mm
7. Updated comparison with the reference paper (arXiv 2506.08378)¶
Reuses 4.5's unit conversion (18 µm/px, not the deck's incorrect 0.3 mm/px)
and reference numbers, extended with the new joint+field and
joint+field+ML rows.
In [10]:
Copied!
PX_MM_CORRECT = 0.018 # 18 um pixel pitch (Jahnke+2024, via Euclid-NISP-Specs.md)
paper_in_sample_y_mm = 7e-4 # arXiv 2506.08378, Sect 4: "< 7e-4 mm (< 0.04 px)"
paper_in_sample_z_mm = 1e-3 # "< 1e-3 mm (< 0.06 px)"
paper_held_out_px = 0.5 # Argon-line validation; not axis-split in the source
paper_held_out_mm = paper_held_out_px * PX_MM_CORRECT
summary_rows = [
{"model": "Physical model, independent per-dataset fit",
"rmse_y_mm": physical_test["independent_rmse_y_mm"].mean(), "rmse_z_mm": physical_test["independent_rmse_z_mm"].mean(),
"source": "Stage 3", "note": "mean over 6 datasets, held-out test split"},
{"model": "Physical model, joint fitting (Tier 1+2+3)",
"rmse_y_mm": physical_test["joint_rmse_y_mm"].mean(), "rmse_z_mm": physical_test["joint_rmse_z_mm"].mean(),
"source": "Stage 4 Phase 3", "note": "mean over 6 datasets, held-out test split"},
{"model": "Hybrid, joint fitting + ML",
"rmse_y_mm": joint_ml_test["joint_ml_rmse_y_mm"].mean(), "rmse_z_mm": joint_ml_test["joint_ml_rmse_z_mm"].mean(),
"source": "Stage 4 Phase 3", "note": "mean over 6 datasets, held-out test split; this is Phase 5's headline row"},
{"model": "Physical model, joint + field-dependent tier",
"rmse_y_mm": field_test["joint_field_rmse_y_mm"].mean(), "rmse_z_mm": field_test["joint_field_rmse_z_mm"].mean(),
"source": "Stage 5 Phase 2 (this notebook)", "note": "mean over 6 datasets, held-out test split, field tier refit train-only"},
{"model": "Hybrid, joint + field-dependent tier + ML",
"rmse_y_mm": field_ml_test["joint_field_ml_rmse_y_mm"].mean(), "rmse_z_mm": field_ml_test["joint_field_ml_rmse_z_mm"].mean(),
"source": "Stage 5 (this notebook)", "note": "mean over 6 datasets, held-out test split -- this project's final model"},
{"model": "Hybrid, joint refinement (regularized)",
"rmse_y_mm": np.nan, "rmse_z_mm": np.nan,
"source": "Stage 4 Phase 4", "note": "NOT ATTEMPTED -- skipped by decision, deferred (see CLAUDE.md)"},
{"model": "Reference (arXiv 2506.08378), in-sample",
"rmse_y_mm": paper_in_sample_y_mm, "rmse_z_mm": paper_in_sample_z_mm,
"source": "Paper Sect 4.2 / Fig 12", "note": "FoV-pooled, all 7 configs, in-sample (Fabry-Perot)"},
{"model": "Reference (arXiv 2506.08378), held-out (Argon)",
"rmse_y_mm": paper_held_out_mm, "rmse_z_mm": paper_held_out_mm,
"source": "Paper Sect 4.2", "note": f"combined figure, not axis-split in source ({paper_held_out_px} px = {paper_held_out_mm:.4f} mm)"},
]
comparison_table = pd.DataFrame(summary_rows).set_index("model")
comparison_table["rmse_y_px"] = comparison_table["rmse_y_mm"] / PX_MM_CORRECT
comparison_table["rmse_z_px"] = comparison_table["rmse_z_mm"] / PX_MM_CORRECT
comparison_table[["rmse_y_mm", "rmse_z_mm", "rmse_y_px", "rmse_z_px", "source", "note"]].round(5)
PX_MM_CORRECT = 0.018 # 18 um pixel pitch (Jahnke+2024, via Euclid-NISP-Specs.md)
paper_in_sample_y_mm = 7e-4 # arXiv 2506.08378, Sect 4: "< 7e-4 mm (< 0.04 px)"
paper_in_sample_z_mm = 1e-3 # "< 1e-3 mm (< 0.06 px)"
paper_held_out_px = 0.5 # Argon-line validation; not axis-split in the source
paper_held_out_mm = paper_held_out_px * PX_MM_CORRECT
summary_rows = [
{"model": "Physical model, independent per-dataset fit",
"rmse_y_mm": physical_test["independent_rmse_y_mm"].mean(), "rmse_z_mm": physical_test["independent_rmse_z_mm"].mean(),
"source": "Stage 3", "note": "mean over 6 datasets, held-out test split"},
{"model": "Physical model, joint fitting (Tier 1+2+3)",
"rmse_y_mm": physical_test["joint_rmse_y_mm"].mean(), "rmse_z_mm": physical_test["joint_rmse_z_mm"].mean(),
"source": "Stage 4 Phase 3", "note": "mean over 6 datasets, held-out test split"},
{"model": "Hybrid, joint fitting + ML",
"rmse_y_mm": joint_ml_test["joint_ml_rmse_y_mm"].mean(), "rmse_z_mm": joint_ml_test["joint_ml_rmse_z_mm"].mean(),
"source": "Stage 4 Phase 3", "note": "mean over 6 datasets, held-out test split; this is Phase 5's headline row"},
{"model": "Physical model, joint + field-dependent tier",
"rmse_y_mm": field_test["joint_field_rmse_y_mm"].mean(), "rmse_z_mm": field_test["joint_field_rmse_z_mm"].mean(),
"source": "Stage 5 Phase 2 (this notebook)", "note": "mean over 6 datasets, held-out test split, field tier refit train-only"},
{"model": "Hybrid, joint + field-dependent tier + ML",
"rmse_y_mm": field_ml_test["joint_field_ml_rmse_y_mm"].mean(), "rmse_z_mm": field_ml_test["joint_field_ml_rmse_z_mm"].mean(),
"source": "Stage 5 (this notebook)", "note": "mean over 6 datasets, held-out test split -- this project's final model"},
{"model": "Hybrid, joint refinement (regularized)",
"rmse_y_mm": np.nan, "rmse_z_mm": np.nan,
"source": "Stage 4 Phase 4", "note": "NOT ATTEMPTED -- skipped by decision, deferred (see CLAUDE.md)"},
{"model": "Reference (arXiv 2506.08378), in-sample",
"rmse_y_mm": paper_in_sample_y_mm, "rmse_z_mm": paper_in_sample_z_mm,
"source": "Paper Sect 4.2 / Fig 12", "note": "FoV-pooled, all 7 configs, in-sample (Fabry-Perot)"},
{"model": "Reference (arXiv 2506.08378), held-out (Argon)",
"rmse_y_mm": paper_held_out_mm, "rmse_z_mm": paper_held_out_mm,
"source": "Paper Sect 4.2", "note": f"combined figure, not axis-split in source ({paper_held_out_px} px = {paper_held_out_mm:.4f} mm)"},
]
comparison_table = pd.DataFrame(summary_rows).set_index("model")
comparison_table["rmse_y_px"] = comparison_table["rmse_y_mm"] / PX_MM_CORRECT
comparison_table["rmse_z_px"] = comparison_table["rmse_z_mm"] / PX_MM_CORRECT
comparison_table[["rmse_y_mm", "rmse_z_mm", "rmse_y_px", "rmse_z_px", "source", "note"]].round(5)
Out[10]:
| rmse_y_mm | rmse_z_mm | rmse_y_px | rmse_z_px | source | note | |
|---|---|---|---|---|---|---|
| model | ||||||
| Physical model, independent per-dataset fit | 0.17340 | 0.31751 | 9.63320 | 17.63967 | Stage 3 | mean over 6 datasets, held-out test split |
| Physical model, joint fitting (Tier 1+2+3) | 0.19433 | 0.30049 | 10.79589 | 16.69367 | Stage 4 Phase 3 | mean over 6 datasets, held-out test split |
| Hybrid, joint fitting + ML | 0.01854 | 0.07487 | 1.03015 | 4.15923 | Stage 4 Phase 3 | mean over 6 datasets, held-out test split; thi... |
| Physical model, joint + field-dependent tier | 0.03640 | 0.08264 | 2.02245 | 4.59095 | Stage 5 Phase 2 (this notebook) | mean over 6 datasets, held-out test split, fie... |
| Hybrid, joint + field-dependent tier + ML | 0.01505 | 0.07530 | 0.83627 | 4.18321 | Stage 5 (this notebook) | mean over 6 datasets, held-out test split -- t... |
| Hybrid, joint refinement (regularized) | NaN | NaN | NaN | NaN | Stage 4 Phase 4 | NOT ATTEMPTED -- skipped by decision, deferred... |
| Reference (arXiv 2506.08378), in-sample | 0.00070 | 0.00100 | 0.03889 | 0.05556 | Paper Sect 4.2 / Fig 12 | FoV-pooled, all 7 configs, in-sample (Fabry-Pe... |
| Reference (arXiv 2506.08378), held-out (Argon) | 0.00900 | 0.00900 | 0.50000 | 0.50000 | Paper Sect 4.2 | combined figure, not axis-split in source (0.5... |
In [11]:
Copied!
final_y = field_ml_test["joint_field_ml_rmse_y_mm"].mean()
final_z = field_ml_test["joint_field_ml_rmse_z_mm"].mean()
prev_y = joint_ml_test["joint_ml_rmse_y_mm"].mean()
prev_z = joint_ml_test["joint_ml_rmse_z_mm"].mean()
print(f"Phase 5 (pre-Stage-5) hybrid vs. paper held-out: "
f"y {prev_y/paper_held_out_mm:.2f}x, z {prev_z/paper_held_out_mm:.2f}x worse")
print(f"This notebook's final hybrid (joint+field+ML) vs. paper held-out: "
f"y {final_y/paper_held_out_mm:.2f}x, z {final_z/paper_held_out_mm:.2f}x worse")
print(f"Field-dependent tier's effect on the gap: "
f"y {prev_y/final_y:.2f}x closer, z {prev_z/final_z:.2f}x closer")
final_y = field_ml_test["joint_field_ml_rmse_y_mm"].mean()
final_z = field_ml_test["joint_field_ml_rmse_z_mm"].mean()
prev_y = joint_ml_test["joint_ml_rmse_y_mm"].mean()
prev_z = joint_ml_test["joint_ml_rmse_z_mm"].mean()
print(f"Phase 5 (pre-Stage-5) hybrid vs. paper held-out: "
f"y {prev_y/paper_held_out_mm:.2f}x, z {prev_z/paper_held_out_mm:.2f}x worse")
print(f"This notebook's final hybrid (joint+field+ML) vs. paper held-out: "
f"y {final_y/paper_held_out_mm:.2f}x, z {final_z/paper_held_out_mm:.2f}x worse")
print(f"Field-dependent tier's effect on the gap: "
f"y {prev_y/final_y:.2f}x closer, z {prev_z/final_z:.2f}x closer")
Phase 5 (pre-Stage-5) hybrid vs. paper held-out: y 2.06x, z 8.32x worse This notebook's final hybrid (joint+field+ML) vs. paper held-out: y 1.67x, z 8.37x worse Field-dependent tier's effect on the gap: y 1.23x closer, z 0.99x closer
In [12]:
Copied!
strategies_labels = ["Physical\n(independent)", "Physical\n(joint)", "Hybrid\n(joint + ML)",
"Physical\n(joint + field)", "Hybrid\n(joint + field + ML)", "Reference\n(in-sample)"]
y_vals = [physical_test["independent_rmse_y_mm"].mean(), physical_test["joint_rmse_y_mm"].mean(),
joint_ml_test["joint_ml_rmse_y_mm"].mean(), field_test["joint_field_rmse_y_mm"].mean(),
field_ml_test["joint_field_ml_rmse_y_mm"].mean(), paper_in_sample_y_mm]
z_vals = [physical_test["independent_rmse_z_mm"].mean(), physical_test["joint_rmse_z_mm"].mean(),
joint_ml_test["joint_ml_rmse_z_mm"].mean(), field_test["joint_field_rmse_z_mm"].mean(),
field_ml_test["joint_field_ml_rmse_z_mm"].mean(), paper_in_sample_z_mm]
fig, ax = plt.subplots(figsize=(11, 5))
x = np.arange(len(strategies_labels))
width = 0.35
ax.bar(x - width / 2, y_vals, width, label="RMSE y (cross-dispersion)", color="tab:blue")
ax.bar(x + width / 2, z_vals, width, label="RMSE z (dispersion)", color="tab:orange")
ax.axhline(paper_held_out_mm, color="0.3", linestyle="--", linewidth=1.5,
label=f"Reference held-out, Argon (combined): {paper_held_out_mm:.4f} mm")
ax.set_yscale("log")
ax.set_xticks(x)
ax.set_xticklabels(strategies_labels)
ax.set_ylabel("RMSE [mm], held-out test split (log scale)")
ax.legend(loc="upper right", fontsize=9)
ax.set_title("This project's final models vs. the reference paper (arXiv 2506.08378)")
plt.tight_layout()
plt.savefig(FIGURES_DIR / "literature_comparison_updated.png", dpi=150)
plt.show()
strategies_labels = ["Physical\n(independent)", "Physical\n(joint)", "Hybrid\n(joint + ML)",
"Physical\n(joint + field)", "Hybrid\n(joint + field + ML)", "Reference\n(in-sample)"]
y_vals = [physical_test["independent_rmse_y_mm"].mean(), physical_test["joint_rmse_y_mm"].mean(),
joint_ml_test["joint_ml_rmse_y_mm"].mean(), field_test["joint_field_rmse_y_mm"].mean(),
field_ml_test["joint_field_ml_rmse_y_mm"].mean(), paper_in_sample_y_mm]
z_vals = [physical_test["independent_rmse_z_mm"].mean(), physical_test["joint_rmse_z_mm"].mean(),
joint_ml_test["joint_ml_rmse_z_mm"].mean(), field_test["joint_field_rmse_z_mm"].mean(),
field_ml_test["joint_field_ml_rmse_z_mm"].mean(), paper_in_sample_z_mm]
fig, ax = plt.subplots(figsize=(11, 5))
x = np.arange(len(strategies_labels))
width = 0.35
ax.bar(x - width / 2, y_vals, width, label="RMSE y (cross-dispersion)", color="tab:blue")
ax.bar(x + width / 2, z_vals, width, label="RMSE z (dispersion)", color="tab:orange")
ax.axhline(paper_held_out_mm, color="0.3", linestyle="--", linewidth=1.5,
label=f"Reference held-out, Argon (combined): {paper_held_out_mm:.4f} mm")
ax.set_yscale("log")
ax.set_xticks(x)
ax.set_xticklabels(strategies_labels)
ax.set_ylabel("RMSE [mm], held-out test split (log scale)")
ax.legend(loc="upper right", fontsize=9)
ax.set_title("This project's final models vs. the reference paper (arXiv 2506.08378)")
plt.tight_layout()
plt.savefig(FIGURES_DIR / "literature_comparison_updated.png", dpi=150)
plt.show()
8. Zeroth-order dispersion: final model¶
Reported alongside, not fused into, the first-order pipeline above -- how
the two compose is still an open design question (see CLAUDE.md).
Reconstructs Stage 5 Phase 3's baseline, shared material_k, and
per-config residual NN. Evaluated in-sample, matching 5.3's own reporting
(no held-out split was used there either).
In [13]:
Copied!
zeroth_sep = {}
for cfg in CONFIGS_ZEROTH:
sep = zeroth_order_separation(DATA_DIR / f"{cfg}_zeroth.csv")
sep = sep[zd.mad_outlier_mask(sep["dz"].values)].reset_index(drop=True)
zeroth_sep[cfg] = sep
raw_rmse = {}
for cfg in CONFIGS_ZEROTH:
sep = zeroth_sep[cfg]
r = sep[["dy", "dz"]].values # uncentered: the physical model's m=0 term predicts exactly zero
raw_rmse[cfg] = float(np.sqrt(np.mean(np.sum(r**2, axis=1))))
CONFIGS_RGS = CONFIGS # bgs000_0 has no frozen first-order fit, excluded from material_k
frozen_params_zeroth = {cfg: joint_fits[cfg]["fixed"] for cfg in CONFIGS_RGS}
fixed_by_cfg_zeroth = {cfg: {k: v for k, v in frozen_params_zeroth[cfg].items() if k != "material_k"}
for cfg in CONFIGS_RGS}
dfs_sep_rgs = {cfg: zeroth_sep[cfg] for cfg in CONFIGS_RGS}
k_fit_joint = zd.fit_shared_material_k(dfs_sep_rgs, fixed_by_cfg_zeroth, model)
print(f"shared material_k = {k_fit_joint:.5f} (nominal {base_config['material']['k']})")
zeroth_sep = {}
for cfg in CONFIGS_ZEROTH:
sep = zeroth_order_separation(DATA_DIR / f"{cfg}_zeroth.csv")
sep = sep[zd.mad_outlier_mask(sep["dz"].values)].reset_index(drop=True)
zeroth_sep[cfg] = sep
raw_rmse = {}
for cfg in CONFIGS_ZEROTH:
sep = zeroth_sep[cfg]
r = sep[["dy", "dz"]].values # uncentered: the physical model's m=0 term predicts exactly zero
raw_rmse[cfg] = float(np.sqrt(np.mean(np.sum(r**2, axis=1))))
CONFIGS_RGS = CONFIGS # bgs000_0 has no frozen first-order fit, excluded from material_k
frozen_params_zeroth = {cfg: joint_fits[cfg]["fixed"] for cfg in CONFIGS_RGS}
fixed_by_cfg_zeroth = {cfg: {k: v for k, v in frozen_params_zeroth[cfg].items() if k != "material_k"}
for cfg in CONFIGS_RGS}
dfs_sep_rgs = {cfg: zeroth_sep[cfg] for cfg in CONFIGS_RGS}
k_fit_joint = zd.fit_shared_material_k(dfs_sep_rgs, fixed_by_cfg_zeroth, model)
print(f"shared material_k = {k_fit_joint:.5f} (nominal {base_config['material']['k']})")
shared material_k = 0.01430 (nominal 0.004)
In [14]:
Copied!
zeroth_rows = []
for cfg in CONFIGS_ZEROTH:
sep = zeroth_sep[cfg].copy()
if cfg in CONFIGS_RGS:
params = {**fixed_by_cfg_zeroth[cfg], "material_k": k_fit_joint}
pred = zd.predict_zeroth_order_separation_physical(sep["y_nisp"].values, sep["z_nisp"].values, model, params)
sep["dy_resid"] = sep["dy"].values - pred[:, 0]
sep["dz_resid"] = sep["dz"].values - pred[:, 1]
r_phys = sep[["dy_resid", "dz_resid"]].values
rmse_phys = float(np.sqrt(np.mean(np.sum(r_phys**2, axis=1))))
nn_input = sep[["y_nisp", "z_nisp", "dy_resid", "dz_resid"]].rename(
columns={"dy_resid": "dy", "dz_resid": "dz"})
else:
rmse_phys = np.nan # bgs000_0: no frozen first-order fit, stays on the standalone empirical model
nn_input = sep[["y_nisp", "z_nisp", "dy", "dz"]]
nn_model = zd.fit_zeroth_dispersion_model(nn_input)
pred_nn = zd.predict_zeroth_dispersion(nn_input["y_nisp"].values, nn_input["z_nisp"].values, nn_model)
r_nn = nn_input[["dy", "dz"]].values - pred_nn
rmse_nn = float(np.sqrt(np.mean(np.sum(r_nn**2, axis=1))))
zeroth_rows.append({"config": cfg, "n": len(sep),
"rmse_raw_mm": raw_rmse[cfg],
"rmse_plus_material_k_mm": rmse_phys,
"rmse_final_nn_mm": rmse_nn})
zeroth_table = pd.DataFrame(zeroth_rows).set_index("config")
zeroth_table["material_k_reduction_pct"] = 100 * (1 - zeroth_table["rmse_plus_material_k_mm"] / zeroth_table["rmse_raw_mm"])
zeroth_table["final_reduction_vs_raw_pct"] = 100 * (1 - zeroth_table["rmse_final_nn_mm"] / zeroth_table["rmse_raw_mm"])
zeroth_table.round(5)
zeroth_rows = []
for cfg in CONFIGS_ZEROTH:
sep = zeroth_sep[cfg].copy()
if cfg in CONFIGS_RGS:
params = {**fixed_by_cfg_zeroth[cfg], "material_k": k_fit_joint}
pred = zd.predict_zeroth_order_separation_physical(sep["y_nisp"].values, sep["z_nisp"].values, model, params)
sep["dy_resid"] = sep["dy"].values - pred[:, 0]
sep["dz_resid"] = sep["dz"].values - pred[:, 1]
r_phys = sep[["dy_resid", "dz_resid"]].values
rmse_phys = float(np.sqrt(np.mean(np.sum(r_phys**2, axis=1))))
nn_input = sep[["y_nisp", "z_nisp", "dy_resid", "dz_resid"]].rename(
columns={"dy_resid": "dy", "dz_resid": "dz"})
else:
rmse_phys = np.nan # bgs000_0: no frozen first-order fit, stays on the standalone empirical model
nn_input = sep[["y_nisp", "z_nisp", "dy", "dz"]]
nn_model = zd.fit_zeroth_dispersion_model(nn_input)
pred_nn = zd.predict_zeroth_dispersion(nn_input["y_nisp"].values, nn_input["z_nisp"].values, nn_model)
r_nn = nn_input[["dy", "dz"]].values - pred_nn
rmse_nn = float(np.sqrt(np.mean(np.sum(r_nn**2, axis=1))))
zeroth_rows.append({"config": cfg, "n": len(sep),
"rmse_raw_mm": raw_rmse[cfg],
"rmse_plus_material_k_mm": rmse_phys,
"rmse_final_nn_mm": rmse_nn})
zeroth_table = pd.DataFrame(zeroth_rows).set_index("config")
zeroth_table["material_k_reduction_pct"] = 100 * (1 - zeroth_table["rmse_plus_material_k_mm"] / zeroth_table["rmse_raw_mm"])
zeroth_table["final_reduction_vs_raw_pct"] = 100 * (1 - zeroth_table["rmse_final_nn_mm"] / zeroth_table["rmse_raw_mm"])
zeroth_table.round(5)
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.
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.
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.
Out[14]:
| n | rmse_raw_mm | rmse_plus_material_k_mm | rmse_final_nn_mm | material_k_reduction_pct | final_reduction_vs_raw_pct | |
|---|---|---|---|---|---|---|
| config | ||||||
| rgs000_0 | 339 | 0.21731 | 0.01382 | 0.00623 | 93.64045 | 97.13227 |
| rgs000_m4 | 297 | 0.21433 | 0.01260 | 0.00698 | 94.12193 | 96.74521 |
| rgs000_p4 | 275 | 0.22507 | 0.01614 | 0.00678 | 92.82694 | 96.98617 |
| rgs180_0 | 394 | 0.21689 | 0.01175 | 0.00734 | 94.58161 | 96.61655 |
| rgs180_m4 | 288 | 0.22298 | 0.01375 | 0.00718 | 93.83225 | 96.78063 |
| rgs180_p4 | 286 | 0.21364 | 0.01099 | 0.00687 | 94.85713 | 96.78467 |
| bgs000_0 | 249 | 0.12276 | NaN | 0.00796 | NaN | 93.51402 |
9. Field-dependent tier: residual maps before/after (illustrative)¶
rgs000_0 and rgs180_0 (both untilted, opposite grism instances) --
illustrates the corner feature 5.2 Step 1 found, and its reduction after
the field-dependent tier.
In [15]:
Copied!
example_cfgs = ["rgs000_0", "rgs180_0"]
fig, axes = plt.subplots(len(example_cfgs), 2, figsize=(9, 3.2 * len(example_cfgs)), constrained_layout=True)
for row, cfg in enumerate(example_cfgs):
df = dfs[cfg]
pred_before = predict_centroids(df["y_nisp"], df["z_nisp"], df["wavelength"],
np.array([]), [], model, fixed=joint_fits[cfg]["fixed"])
r_before = np.hypot(pred_before[0] - df["cent_y"].values, pred_before[1] - df["cent_z"].values)
fr = field_results[cfg]
pred_after = fc.predict_centroids_field(df["y_nisp"], df["z_nisp"], df["wavelength"],
fr["tilt_deg"], model, fr["fixed_tier12"], fr["nn_model"])
r_after = np.hypot(pred_after[0] - df["cent_y"].values, pred_after[1] - df["cent_z"].values)
vmax = max(r_before.max(), r_after.max())
sc0 = axes[row, 0].scatter(df["y_nisp"], df["z_nisp"], c=r_before, cmap="magma_r", vmin=0, vmax=vmax, s=10)
sc1 = axes[row, 1].scatter(df["y_nisp"], df["z_nisp"], c=r_after, cmap="magma_r", vmin=0, vmax=vmax, s=10)
axes[row, 0].set_ylabel(f"{cfg}\nz_nisp [mm]")
for ax in axes[row]:
ax.set_xlabel("y_nisp [mm]")
ax.set_aspect("equal")
axes[0, 0].set_title("|residual| before field tier [mm]")
axes[0, 1].set_title("|residual| after field tier [mm]")
fig.colorbar(sc0, ax=axes[:, 0], shrink=0.6)
fig.colorbar(sc1, ax=axes[:, 1], shrink=0.6)
plt.savefig(FIGURES_DIR / "field_tier_before_after.png", dpi=150)
plt.show()
example_cfgs = ["rgs000_0", "rgs180_0"]
fig, axes = plt.subplots(len(example_cfgs), 2, figsize=(9, 3.2 * len(example_cfgs)), constrained_layout=True)
for row, cfg in enumerate(example_cfgs):
df = dfs[cfg]
pred_before = predict_centroids(df["y_nisp"], df["z_nisp"], df["wavelength"],
np.array([]), [], model, fixed=joint_fits[cfg]["fixed"])
r_before = np.hypot(pred_before[0] - df["cent_y"].values, pred_before[1] - df["cent_z"].values)
fr = field_results[cfg]
pred_after = fc.predict_centroids_field(df["y_nisp"], df["z_nisp"], df["wavelength"],
fr["tilt_deg"], model, fr["fixed_tier12"], fr["nn_model"])
r_after = np.hypot(pred_after[0] - df["cent_y"].values, pred_after[1] - df["cent_z"].values)
vmax = max(r_before.max(), r_after.max())
sc0 = axes[row, 0].scatter(df["y_nisp"], df["z_nisp"], c=r_before, cmap="magma_r", vmin=0, vmax=vmax, s=10)
sc1 = axes[row, 1].scatter(df["y_nisp"], df["z_nisp"], c=r_after, cmap="magma_r", vmin=0, vmax=vmax, s=10)
axes[row, 0].set_ylabel(f"{cfg}\nz_nisp [mm]")
for ax in axes[row]:
ax.set_xlabel("y_nisp [mm]")
ax.set_aspect("equal")
axes[0, 0].set_title("|residual| before field tier [mm]")
axes[0, 1].set_title("|residual| after field tier [mm]")
fig.colorbar(sc0, ax=axes[:, 0], shrink=0.6)
fig.colorbar(sc1, ax=axes[:, 1], shrink=0.6)
plt.savefig(FIGURES_DIR / "field_tier_before_after.png", dpi=150)
plt.show()
10. Summary¶
In [16]:
Copied!
print("=== Final pipeline, held-out test split, mean over 6 datasets ===")
print(f"independent physical: y={physical_test['independent_rmse_y_mm'].mean():.4f} mm "
f"z={physical_test['independent_rmse_z_mm'].mean():.4f} mm")
print(f"joint physical (Tier 1-3): y={physical_test['joint_rmse_y_mm'].mean():.4f} mm "
f"z={physical_test['joint_rmse_z_mm'].mean():.4f} mm")
print(f"joint + ML (pre-field): y={joint_ml_test['joint_ml_rmse_y_mm'].mean():.4f} mm "
f"z={joint_ml_test['joint_ml_rmse_z_mm'].mean():.4f} mm")
print(f"joint + field tier: y={field_test['joint_field_rmse_y_mm'].mean():.4f} mm "
f"z={field_test['joint_field_rmse_z_mm'].mean():.4f} mm")
print(f"joint + field tier + ML: y={field_ml_test['joint_field_ml_rmse_y_mm'].mean():.4f} mm "
f"z={field_ml_test['joint_field_ml_rmse_z_mm'].mean():.4f} mm <- final model")
print()
print(f"reference paper (arXiv 2506.08378), held-out (Argon, combined): {paper_held_out_mm:.4f} mm")
print(f"remaining gap: y {final_y/paper_held_out_mm:.2f}x, z {final_z/paper_held_out_mm:.2f}x")
print()
print("=== 0th-order dispersion (in-sample) ===")
print(zeroth_table[["rmse_raw_mm", "rmse_plus_material_k_mm", "rmse_final_nn_mm", "final_reduction_vs_raw_pct"]].round(4).to_string())
print()
print(f"shared material_k = {k_fit_joint:.5f} (nominal {base_config['material']['k']})")
print("=== Final pipeline, held-out test split, mean over 6 datasets ===")
print(f"independent physical: y={physical_test['independent_rmse_y_mm'].mean():.4f} mm "
f"z={physical_test['independent_rmse_z_mm'].mean():.4f} mm")
print(f"joint physical (Tier 1-3): y={physical_test['joint_rmse_y_mm'].mean():.4f} mm "
f"z={physical_test['joint_rmse_z_mm'].mean():.4f} mm")
print(f"joint + ML (pre-field): y={joint_ml_test['joint_ml_rmse_y_mm'].mean():.4f} mm "
f"z={joint_ml_test['joint_ml_rmse_z_mm'].mean():.4f} mm")
print(f"joint + field tier: y={field_test['joint_field_rmse_y_mm'].mean():.4f} mm "
f"z={field_test['joint_field_rmse_z_mm'].mean():.4f} mm")
print(f"joint + field tier + ML: y={field_ml_test['joint_field_ml_rmse_y_mm'].mean():.4f} mm "
f"z={field_ml_test['joint_field_ml_rmse_z_mm'].mean():.4f} mm <- final model")
print()
print(f"reference paper (arXiv 2506.08378), held-out (Argon, combined): {paper_held_out_mm:.4f} mm")
print(f"remaining gap: y {final_y/paper_held_out_mm:.2f}x, z {final_z/paper_held_out_mm:.2f}x")
print()
print("=== 0th-order dispersion (in-sample) ===")
print(zeroth_table[["rmse_raw_mm", "rmse_plus_material_k_mm", "rmse_final_nn_mm", "final_reduction_vs_raw_pct"]].round(4).to_string())
print()
print(f"shared material_k = {k_fit_joint:.5f} (nominal {base_config['material']['k']})")
=== Final pipeline, held-out test split, mean over 6 datasets ===
independent physical: y=0.1734 mm z=0.3175 mm
joint physical (Tier 1-3): y=0.1943 mm z=0.3005 mm
joint + ML (pre-field): y=0.0185 mm z=0.0749 mm
joint + field tier: y=0.0364 mm z=0.0826 mm
joint + field tier + ML: y=0.0151 mm z=0.0753 mm <- final model
reference paper (arXiv 2506.08378), held-out (Argon, combined): 0.0090 mm
remaining gap: y 1.67x, z 8.37x
=== 0th-order dispersion (in-sample) ===
rmse_raw_mm rmse_plus_material_k_mm rmse_final_nn_mm final_reduction_vs_raw_pct
config
rgs000_0 0.2173 0.0138 0.0062 97.1323
rgs000_m4 0.2143 0.0126 0.0070 96.7452
rgs000_p4 0.2251 0.0161 0.0068 96.9862
rgs180_0 0.2169 0.0118 0.0073 96.6165
rgs180_m4 0.2230 0.0138 0.0072 96.7806
rgs180_p4 0.2136 0.0110 0.0069 96.7847
bgs000_0 0.1228 NaN 0.0080 93.5140
shared material_k = 0.01430 (nominal 0.004)