Stage 4, Phase 5 — Comparison with Published Results (4.5-Comparison_With_Published_Results)¶

def load_toml(path):
    with open(path, "rb") as f:
        return tomllib.load(f)


def numeric_result(fit_result):
    """Drop the *_bootstrap_std entries, keep only the fitted values."""
    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 test_split(df):
    gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
    _, test_idx = next(gss.split(df, groups=df["spectra_id"].values))
    return 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)))


physical_rows = []
for cfg in CONFIGS:
    df_test = test_split(dfs[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 stage4_hybrid_rmse():
    """Per-axis, held-out RMSE of the joint physical fit + independent MLP(y)+MLP(z)
    correction (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 best_single_dataset_mlp():
    """Phase 1's best independent MLP(y)/MLP(z), rgs000_0 only -- read back
    from MLflow ('residual_correction' experiment), not retrained.

    Independent (single-axis) runs log their test RMSE under the generic
    `metrics.test_rmse` column, not `test_rmse_y`/`test_rmse_z` -- those are
    only populated by the joint 2D-output model, which predicts both axes at
    once."""
    exp = mlflow.get_experiment_by_name("residual_correction")
    assert exp is not None, "run notebooks/4.1-ML_Comparison.ipynb first to populate the 'residual_correction' experiment"
    runs = mlflow.search_runs(experiment_ids=[exp.experiment_id])
    mlp = runs[runs["params.model"] == "MLP"]
    best_y = mlp[mlp["params.axis"] == "y"].sort_values("metrics.test_rmse").iloc[0]
    best_z = mlp[mlp["params.axis"] == "z"].sort_values("metrics.test_rmse").iloc[0]
    return float(best_y["metrics.test_rmse"]), float(best_z["metrics.test_rmse"])


single_ml_rmse_y, single_ml_rmse_z = best_single_dataset_mlp()
print(f"Phase 1 best independent MLP (rgs000_0 only): rmse_y={single_ml_rmse_y:.4f} mm, rmse_z={single_ml_rmse_z:.4f} mm")

PX_MM_CORRECT = 0.018  # 18 um pixel pitch (Jahnke+2024, via Euclid-NISP-Specs.md) -- physical detector pixel size
PX_MM_DECK = 0.3       # deck's stated "1 NISP pixel ~= 0.3 mm" -- see markdown above

# Reference paper (arXiv 2506.08378), Section 4 of `2506.08378v2-nisp_grism_trace_model.md`
paper_in_sample_y_mm = 7e-4   # "< 7e-4 mm (< 0.04 px)", cross-dispersion, FoV-pooled, in-sample (Fabry-Perot)
paper_in_sample_z_mm = 1e-3   # "< 1e-3 mm (< 0.06 px)", dispersion, FoV-pooled, in-sample
paper_held_out_px = 0.5       # Argon-line validation; axis not separately reported in the source

implied_y_um_per_px = paper_in_sample_y_mm / 0.04 * 1000
implied_z_um_per_px = paper_in_sample_z_mm / 0.06 * 1000
print(f"paper's own in-sample y figure (0.7e-3 mm / 0.04 px) implies {implied_y_um_per_px:.1f} um/px")
print(f"paper's own in-sample z figure (1.0e-3 mm / 0.06 px) implies {implied_z_um_per_px:.1f} um/px")
print(f"Euclid-NISP-Specs.md pixel pitch (Jahnke+2024, measured):    18.0 um/px  <- matches, confirms this is the right conversion")
print(f"deck's stated conversion:                                  {PX_MM_DECK*1000:.0f} um/px  <- ~17x too large, do not use")

paper_held_out_mm = paper_held_out_px * PX_MM_CORRECT
paper_held_out_mm_deck = paper_held_out_px * PX_MM_DECK
print(f"\nheld-out Argon validation: {paper_held_out_px} px")
print(f"  -> {paper_held_out_mm:.4f} mm using the correct 18 um/px")
print(f"  -> {paper_held_out_mm_deck:.4f} mm using the deck's conversion ({paper_held_out_mm_deck/paper_held_out_mm:.1f}x larger -- see Section 5 for why this matters)")

for cfg in CONFIGS:
    wl = dfs[cfg]["wavelength"]
    print(f"{cfg}: wavelength {wl.min():.1f}-{wl.max():.1f} nm ({len(dfs[cfg])} points)")

print("\nPaper's red-grism passband (RGS000/RGS180, all tilts, Jahnke+2024 Table 2,")
print("via the reference doc Sect 2.4): 1206-1892 nm")

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": "Task 2 (Phase 2)", "note": "mean over 6 datasets, held-out test split"},
    {"model": "Hybrid, best ML, single dataset",
     "rmse_y_mm": single_ml_rmse_y, "rmse_z_mm": single_ml_rmse_z,
     "source": "Task 1 (Phase 1)", "note": "rgs000_0 only, 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": "Task 3 (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": "Task 3 (Phase 3)", "note": "mean over 6 datasets, held-out test split"},
    {"model": "Hybrid, joint refinement (regularized)",
     "rmse_y_mm": np.nan, "rmse_z_mm": np.nan,
     "source": "Task 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); shown in both columns"},
]
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)

per_dataset = physical_test.join(joint_ml_test)
# sqrt(rmse_y^2 + rmse_z^2): a standard combined-RMSE proxy, not identical to
# mean(hypot(r_y, r_z)) used in 4.3 but close and consistent across rows here.
for strategy in ["independent", "joint", "joint_ml"]:
    per_dataset[f"{strategy}_rmse_combined_mm"] = np.hypot(
        per_dataset[f"{strategy}_rmse_y_mm"], per_dataset[f"{strategy}_rmse_z_mm"])
per_dataset.round(4)

strategies_labels = ["Physical\n(independent)", "Physical\n(joint)", "Hybrid, single-dataset\nML (rgs000_0 only)",
                     "Hybrid, joint fit\n+ ML", "Reference\n(in-sample)"]
y_vals = [physical_test["independent_rmse_y_mm"].mean(), physical_test["joint_rmse_y_mm"].mean(),
          single_ml_rmse_y, joint_ml_test["joint_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(),
          single_ml_rmse_z, joint_ml_test["joint_ml_rmse_z_mm"].mean(), paper_in_sample_z_mm]

fig, ax = plt.subplots(figsize=(10, 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, not axis-split): {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 models vs. the reference paper (arXiv 2506.08378)")
plt.tight_layout()
plt.show()