Stage 8 — Chebyshev-Polynomial Residual Model¶

import tomllib
from pathlib import Path

import numpy as np
import pandas as pd
import mlflow
from sklearn.model_selection import GroupShuffleSplit

import dispcraft.chebyshev_residual as cr
from dispcraft.calibration import ground_test_model_from_config, predict_centroids

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

RNG_SEED = 42  # same seed as 4.1/4.3/4.5/5.2/5.3/6 -- reproduces their splits exactly

CONFIGS = ["rgs000_0", "rgs000_m4", "rgs000_p4", "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()}")

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")}


from dispcraft.measurement import load_spectra, median_per_spectrum

dfs = {}
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

    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)
    train_idx, test_idx = next(gss.split(df, groups=df["spectra_id"].values))
    return df.iloc[train_idx].reset_index(drop=True), df.iloc[test_idx].reset_index(drop=True)


def physical_residual(df, fixed):
    """r_y, r_z = pred_phys - data -- dispcraft.calibration.cost's own sign
    convention, and the one dispcraft.chebyshev_residual.fit_chebyshev_residual
    expects its input residual columns in (see that module's docstring)."""
    pred = predict_centroids(df["y_nisp"], df["z_nisp"], df["wavelength"], theta=[], free_names=[],
                              model=model, fixed=fixed)
    return pred[0] - df["cent_y"].values, pred[1] - df["cent_z"].values


def axis_rmse(r_y, r_z):
    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]
    r_y, r_z = physical_residual(df_test, joint_fits[cfg]["fixed"])
    y, z = axis_rmse(r_y, r_z)
    physical_rows.append({"cfg": cfg, "physical_rmse_y_mm": y, "physical_rmse_z_mm": z})

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 hybrid_joint_ml (joint physical fit +
    independent MLP(y)+MLP(z) residual, Stage 4 Phase 3 Section 4) -- read
    back from MLflow, not retrained. Mirrors notebooks/6-Status_Report
    .ipynb's own stage4_hybrid_rmse exactly."""
    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": "hybrid_ml_rmse_y_mm", "metrics.hybrid_rmse_z": "hybrid_ml_rmse_z_mm"}
    ).astype(float)


hybrid_ml_test = stage4_hybrid_rmse()
hybrid_ml_test.round(5)

def chebyshev_held_out_rmse(cfg, wave_order, spatial_order):
    df_train, df_test = train_test[cfg]
    fixed = joint_fits[cfg]["fixed"]

    r_y_train, r_z_train = physical_residual(df_train, fixed)
    d_train = df_train.assign(r_y=r_y_train, r_z=r_z_train)

    cheb_model = cr.fit_chebyshev_residual(d_train, wave_order, spatial_order, *cr.RGS_PASSBAND_NM)

    pred_test = cr.predict_centroids_chebyshev(df_test["y_nisp"], df_test["z_nisp"], df_test["wavelength"],
                                                model, fixed, cheb_model)
    r_y_test = pred_test[0] - df_test["cent_y"].values
    r_z_test = pred_test[1] - df_test["cent_z"].values
    rmse_y, rmse_z = axis_rmse(r_y_test, r_z_test)
    return rmse_y, rmse_z, cheb_model


def run_chebyshev_order(wave_order, spatial_order, log_mlflow=True):
    if log_mlflow:
        mlflow.set_experiment("chebyshev_residual")
    rows, models_by_cfg = [], {}
    for cfg in CONFIGS:
        rmse_y, rmse_z, cheb_model = chebyshev_held_out_rmse(cfg, wave_order, spatial_order)
        rows.append({"cfg": cfg, "chebyshev_rmse_y_mm": rmse_y, "chebyshev_rmse_z_mm": rmse_z})
        models_by_cfg[cfg] = cheb_model
        if log_mlflow:
            with mlflow.start_run(run_name=f"{cfg}_order{wave_order}"):
                mlflow.log_params({"dataset": cfg, "wave_order": wave_order, "spatial_order": spatial_order,
                                    "n_terms_per_axis": cr.n_chebyshev_terms(wave_order, spatial_order)})
                mlflow.log_metrics({"rmse_y": rmse_y, "rmse_z": rmse_z, "combined": float(np.hypot(rmse_y, rmse_z))})
    table = pd.DataFrame(rows).set_index("cfg").loc[CONFIGS]
    return table, models_by_cfg


chebyshev1_test, chebyshev1_models = run_chebyshev_order(wave_order=1, spatial_order=1)
chebyshev1_test.round(5)

chebyshev0_test, chebyshev0_models = run_chebyshev_order(wave_order=0, spatial_order=0)
chebyshev0_test.round(5)

PX_MM_CORRECT = 0.018  # 18 um pixel pitch (Jahnke+2024, via Euclid-NISP-Specs.md), matches notebooks/6
paper_held_out_px = 0.5
paper_held_out_mm = paper_held_out_px * PX_MM_CORRECT

summary_rows = [
    {"model": "Physical model, joint fitting (Tier 1+2+3)",
     "rmse_y_mm": physical_test["physical_rmse_y_mm"].mean(), "rmse_z_mm": physical_test["physical_rmse_z_mm"].mean(),
     "note": "mean over 6 datasets, held-out test split"},
    {"model": "Hybrid, joint fitting + ML (hybrid_joint_ml, production default)",
     "rmse_y_mm": hybrid_ml_test["hybrid_ml_rmse_y_mm"].mean(), "rmse_z_mm": hybrid_ml_test["hybrid_ml_rmse_z_mm"].mean(),
     "note": "mean over 6 datasets, held-out test split; Status_Report.md Section 5"},
    {"model": "Hybrid, joint fitting + Chebyshev residual, order 1",
     "rmse_y_mm": chebyshev1_test["chebyshev_rmse_y_mm"].mean(), "rmse_z_mm": chebyshev1_test["chebyshev_rmse_z_mm"].mean(),
     "note": "8 coefficients/axis, mean over 6 datasets, held-out test split"},
    {"model": "Hybrid, joint fitting + Chebyshev residual, order 0",
     "rmse_y_mm": chebyshev0_test["chebyshev_rmse_y_mm"].mean(), "rmse_z_mm": chebyshev0_test["chebyshev_rmse_z_mm"].mean(),
     "note": "1 coefficient/axis (sanity floor), mean over 6 datasets, held-out test split"},
    {"model": "Reference (arXiv 2506.08378), held-out (Argon)",
     "rmse_y_mm": paper_held_out_mm, "rmse_z_mm": paper_held_out_mm,
     "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["combined_mm"] = np.hypot(comparison_table["rmse_y_mm"], comparison_table["rmse_z_mm"])
comparison_table.round(5)

print("=== Chebyshev residual vs. hybrid_joint_ml, held-out, mean over 6 datasets ===")
for label, tbl, cols in [
    ("physical (joint fit alone)", physical_test, ("physical_rmse_y_mm", "physical_rmse_z_mm")),
    ("hybrid_joint_ml", hybrid_ml_test, ("hybrid_ml_rmse_y_mm", "hybrid_ml_rmse_z_mm")),
    ("chebyshev order 1", chebyshev1_test, ("chebyshev_rmse_y_mm", "chebyshev_rmse_z_mm")),
    ("chebyshev order 0", chebyshev0_test, ("chebyshev_rmse_y_mm", "chebyshev_rmse_z_mm")),
]:
    y, z = tbl[cols[0]].mean(), tbl[cols[1]].mean()
    print(f"{label:32s} y={y:.4f} mm  z={z:.4f} mm  combined={np.hypot(y, z):.4f} mm")

print()
print(f"reference paper held-out (Argon, combined): {paper_held_out_mm:.4f} mm")