Stage 5 Phase 4 — BGS Grism Model¶

import tomllib
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import mlflow
from scipy.optimize import minimize
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 dispcraft.calibration import cost, ground_test_model_from_config, predict_centroids
from dispcraft.measurement import load_spectra, median_per_spectrum

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

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

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

df_raw = median_per_spectrum(load_spectra(DATA_DIR / "bgs000_0_first.csv"))
print(f"bgs000_0: {len(df_raw)} rows, {df_raw['spectra_id'].nunique()} spectra, "
      f"wavelength {df_raw['wavelength'].min():.1f}-{df_raw['wavelength'].max():.1f} nm")

FREE_OFFSET_ONLY = ["offset_y_mm", "offset_z_mm"]
theta0_offset = np.array([NOMINAL_PARAMS[k] for k in FREE_OFFSET_ONLY])

result_raw = minimize(cost, theta0_offset, args=(df_raw, FREE_OFFSET_ONLY, model), method="Nelder-Mead",
                       options={"maxiter": 5000, "xatol": 1e-8, "fatol": 1e-12})
print(f"cost before fit: {cost(theta0_offset, df_raw, FREE_OFFSET_ONLY, model):.4f} mm^2")
print(f"result.x   = {result_raw.x} mm")
print(f"result.fun = {result_raw.fun:.4f} mm^2  (RMSE = {np.sqrt(result_raw.fun):.4f} mm)")

pred0 = predict_centroids(df_raw["y_nisp"], df_raw["z_nisp"], df_raw["wavelength"],
                           result_raw.x, FREE_OFFSET_ONLY, model)
df_raw = df_raw.assign(r_y=pred0[0] - df_raw["cent_y"].values, r_z=pred0[1] - df_raw["cent_z"].values)
per_spectrum = df_raw.groupby("spectra_id").agg(
    y_nisp=("y_nisp", "first"), z_nisp=("z_nisp", "first"),
    mean_abs_r_y=("r_y", lambda s: s.abs().mean()),
    mean_abs_r_z=("r_z", lambda s: s.abs().mean()),
).sort_values("mean_abs_r_y", ascending=False)
per_spectrum.head(8)

# Decide exclusion from the ratio between the worst spectrum and the rest
# (Stage 3's rgs000_0 outlier stood out by more than an order of magnitude
# against the 2nd-worst spectrum -- the same threshold used here).
worst, second = per_spectrum["mean_abs_r_y"].iloc[0], per_spectrum["mean_abs_r_y"].iloc[1]
ratio = worst / second
print(f"worst/2nd-worst mean|r_y| ratio: {ratio:.2f}")

if ratio > 10:
    EXCLUDED_SPECTRA_IDS = [int(per_spectrum.index[0])]
    df0 = df_raw[~df_raw["spectra_id"].isin(EXCLUDED_SPECTRA_IDS)].drop(columns=["r_y", "r_z"]).reset_index(drop=True)
    print(f"excluded spectra_id={EXCLUDED_SPECTRA_IDS[0]}; "
          f"{len(df_raw) - len(df0)} rows dropped, {len(df0)} remain")
else:
    EXCLUDED_SPECTRA_IDS = []
    df0 = df_raw.drop(columns=["r_y", "r_z"]).reset_index(drop=True)
    print("no dominant outlier found (ratio <= 10) -- no exclusion applied")

def bootstrap_uncertainty(theta_star, df, free_names, model, n_boot=50, seed=RNG_SEED, fixed=None):
    rng = np.random.default_rng(seed)
    n = len(df)
    boots = []
    for _ in range(n_boot):
        sample = df.iloc[rng.integers(0, n, n)]
        r = minimize(cost, theta_star, args=(sample, free_names, model, fixed), method="Nelder-Mead",
                     options={"maxiter": 8000, "maxfev": 8000, "xatol": 1e-8, "fatol": 1e-12, "adaptive": True})
        boots.append(r.x)
    boots = np.array(boots)
    std = boots.std(axis=0)
    corr = pd.DataFrame(np.corrcoef(boots.T), index=free_names, columns=free_names)
    return std, corr


FREE_CORE = ["offset_y_mm", "offset_z_mm", "tilt_deg", "rho"]


def identifiability_check(extra_param):
    free_names = FREE_CORE + [extra_param]
    theta0 = np.array([NOMINAL_PARAMS[k] for k in free_names])
    r = minimize(cost, theta0, args=(df0, free_names, model), method="Nelder-Mead",
                 options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-14, "adaptive": True})
    std, corr = bootstrap_uncertainty(r.x, df0, free_names, model)
    print(f"=== {extra_param} identifiability check ===")
    print(f"theta* = {dict(zip(free_names, r.x))}")
    print(f"RMSE = {np.sqrt(r.fun):.4f} mm")
    print("bootstrap std:", dict(zip(free_names, std.round(5))))
    other_corrs = corr.loc[extra_param].drop(extra_param)
    print(f"max |corr({extra_param}, other)| = {other_corrs.abs().max():.3f} (with {other_corrs.abs().idxmax()})")
    print()
    return r, std, corr


result_a, std_a, corr_a = identifiability_check("A_deg")

result_k, std_k, corr_k = identifiability_check("material_k")

# Decide FREE_PROD from Section 3's numbers: material_k identifiable if its
# bootstrap std is small relative to the fitted value (well-constrained)
# and it isn't strongly correlated with rho/tilt (no degeneracy).
k_star = dict(zip(FREE_CORE + ["material_k"], result_k.x))["material_k"]
k_std = dict(zip(FREE_CORE + ["material_k"], std_k))["material_k"]
k_rel_std = abs(k_std / k_star) if k_star != 0 else np.inf
k_max_corr = corr_k.loc["material_k"].drop("material_k").abs().max()
material_k_identifiable = (k_rel_std < 0.3) and (k_max_corr < 0.9)

print(f"material_k: fitted={k_star:.5f}, bootstrap_std={k_std:.5f} (rel={k_rel_std:.2f}), "
      f"max|corr|={k_max_corr:.3f} -> identifiable={material_k_identifiable}")

FREE_PROD = FREE_CORE + (["material_k"] if material_k_identifiable else [])
print(f"production free parameters: {FREE_PROD}")

theta0_prod = np.array([NOMINAL_PARAMS[k] for k in FREE_PROD])
cost_before = cost(theta0_prod, df0, FREE_PROD, model)

result_prod = minimize(cost, theta0_prod, args=(df0, FREE_PROD, model), method="Nelder-Mead",
                        options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-14, "adaptive": True})
theta_star_prod = result_prod.x
std_prod, corr_prod = bootstrap_uncertainty(theta_star_prod, df0, FREE_PROD, model)

print(f"cost before fit: {cost_before:.4f} mm^2")
print(f"theta*     = {dict(zip(FREE_PROD, theta_star_prod))}")
print(f"RMSE       = {np.sqrt(result_prod.fun):.4f} mm  (before: {np.sqrt(cost_before):.4f} mm)")
print("bootstrap std:", dict(zip(FREE_PROD, std_prod.round(5))))
print(f"\nsanity check vs. spec nominal: tilt_deg near 0 deg? "
      f"{dict(zip(FREE_PROD, theta_star_prod)).get('tilt_deg'):.3f} deg; "
      f"rho near 15.1 grooves/mm? {dict(zip(FREE_PROD, theta_star_prod)).get('rho'):.3f}")

pred_before = predict_centroids(df0["y_nisp"], df0["z_nisp"], df0["wavelength"], theta0_prod, FREE_PROD, model)
pred_after = predict_centroids(df0["y_nisp"], df0["z_nisp"], df0["wavelength"], theta_star_prod, FREE_PROD, model)

fig, axes = plt.subplots(1, 2, figsize=(11, 5), sharex=True, sharey=True)
for ax, pred, label in zip(axes, [pred_before, pred_after], ["before fit (theta0)", "after fit (theta*)"]):
    ax.scatter(df0["cent_y"], df0["cent_z"], s=4, label="observed", alpha=0.6)
    ax.scatter(pred[0], pred[1], s=4, label="predicted", alpha=0.6)
    ax.set_xlabel("cent_y [mm]")
    ax.set_title(label)
    ax.legend(markerscale=3)
axes[0].set_ylabel("cent_z [mm]")
fig.suptitle("bgs000_0: observed vs. predicted centroids")
plt.tight_layout()
plt.show()

def write_bgs_fit_toml(df, free_names, theta, theta_std, cost_before, cost_after, excluded_ids):
    result_lines = "\n".join(f"{name} = {v:.6f}" for name, v in zip(free_names, theta))
    std_lines = "\n".join(f"{name}_bootstrap_std = {s:.6f}" for name, s in zip(free_names, theta_std))
    p = dict(NOMINAL_PARAMS)
    p.update(dict(zip(free_names, theta)))
    fixed_lines = []
    if "A_deg" not in free_names:
        fixed_lines.append(f"A_deg = {p['A_deg']}")
    if "material_n0" not in free_names:
        fixed_lines.append(f"material_n0 = {p['material_n0']}")
    if "material_k" not in free_names:
        fixed_lines.append(f"material_k = {p['material_k']}")
    fixed_block = "\n".join(fixed_lines)

    text = f'''# Stage 5 Phase 4 -- physical-model fit to bgs000_0.
# Generated by notebooks/5.4-BGS_Model.ipynb. Do not hand-edit the fitted
# values; re-run the notebook and regenerate this file instead.

[base_config]
path = "stage1_bgs_instrument.toml"  # all other parameters unchanged from this

[fit]
dataset = "bgs000_0_first.csv"
loader = "median_per_spectrum(load_spectra(...))"  # dispcraft.measurement, default sig_max
excluded_spectra_ids = {excluded_ids}
n_points = {len(df)}
method = "Nelder-Mead"
free_parameters = {free_names}
cost_before_mm2 = {cost_before:.6f}
cost_after_mm2 = {cost_after:.6f}
rmse_mm = {np.sqrt(cost_after):.6f}

[fit.result]
{result_lines}
{std_lines}

[fit.fixed]
# Held at the BGS000 nominal (stage1_bgs_instrument.toml) -- see notebook
# Section 3's identifiability check for why each fixed parameter isn't free.
{fixed_block}
'''
    path = MODELS_DIR / "bgs000_0_fit.toml"
    path.write_text(text)
    print(f"wrote {path}")
    return path


bgs_fit_path = write_bgs_fit_toml(df0, FREE_PROD, theta_star_prod, std_prod, cost_before, result_prod.fun,
                                   EXCLUDED_SPECTRA_IDS)
print(bgs_fit_path.read_text())

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)

pred_final = predict_centroids(df0["y_nisp"], df0["z_nisp"], df0["wavelength"],
                                theta_star_prod, FREE_PROD, model)
d = df0.assign(r_y=pred_final[0] - df0["cent_y"].values, r_z=pred_final[1] - df0["cent_z"].values)

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

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

physical_rmse_y = float(np.sqrt(mean_squared_error(r_y_test, np.zeros_like(r_y_test))))
physical_rmse_z = float(np.sqrt(mean_squared_error(r_z_test, np.zeros_like(r_z_test))))
physical_dist_test = float(np.mean(np.hypot(r_y_test, r_z_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 = r_y_test - mlp_y.predict(X_test_s)
hyb_z_test = r_z_test - mlp_z.predict(X_test_s)

hybrid_rmse_y = float(np.sqrt(mean_squared_error(r_y_test, mlp_y.predict(X_test_s))))
hybrid_rmse_z = float(np.sqrt(mean_squared_error(r_z_test, mlp_z.predict(X_test_s))))
hybrid_dist_test = float(np.mean(np.hypot(hyb_y_test, hyb_z_test)))

bgs_results = pd.DataFrame({
    "physical_rmse_y_mm": [physical_rmse_y], "physical_rmse_z_mm": [physical_rmse_z],
    "hybrid_rmse_y_mm": [hybrid_rmse_y], "hybrid_rmse_z_mm": [hybrid_rmse_z],
    "physical_dist_mm": [physical_dist_test], "hybrid_dist_mm": [hybrid_dist_test],
}, index=["bgs000_0"])
bgs_results["improvement_pct"] = 100 * (1 - bgs_results["hybrid_dist_mm"] / bgs_results["physical_dist_mm"])
bgs_results.round(5)

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

with mlflow.start_run(run_name="bgs000_0_hybrid"):
    mlflow.log_param("dataset", "bgs000_0")
    mlflow.log_param("free_parameters", FREE_PROD)
    for name, v in zip(FREE_PROD, theta_star_prod):
        mlflow.log_param(name, v)
    for axis, p in [("y", mlp_y_params), ("z", mlp_z_params)]:
        for k, v in p.items():
            mlflow.log_param(f"mlp_{axis}_{k}", v)
    mlflow.log_metric("physical_dist_test", physical_dist_test)
    mlflow.log_metric("hybrid_dist_test", hybrid_dist_test)
    mlflow.log_metric("hybrid_rmse_y", hybrid_rmse_y)
    mlflow.log_metric("hybrid_rmse_z", hybrid_rmse_z)
    mlflow.log_metric("improvement_pct", float(bgs_results["improvement_pct"].iloc[0]))

print(f"logged 1 run to the '{EXPERIMENT_NAME}' MLflow experiment")

PX_MM_CORRECT = 0.018  # 18 um pixel pitch (Jahnke+2024, via Euclid-NISP-Specs.md)

paper_in_sample_y_mm = 7e-4
paper_in_sample_z_mm = 1e-3
paper_held_out_px = 0.5
paper_held_out_mm = paper_held_out_px * PX_MM_CORRECT

bgs_summary_rows = [
    {"model": "Physical model (this notebook)", "rmse_y_mm": physical_rmse_y, "rmse_z_mm": physical_rmse_z,
     "note": "held-out test split"},
    {"model": "Hybrid, physical + ML (this notebook)", "rmse_y_mm": hybrid_rmse_y, "rmse_z_mm": hybrid_rmse_z,
     "note": "held-out test split"},
    {"model": "Reference (arXiv 2506.08378), in-sample", "rmse_y_mm": paper_in_sample_y_mm, "rmse_z_mm": paper_in_sample_z_mm,
     "note": "FoV-pooled, in-sample (Fabry-Perot), stated for every grism config incl. BGS000"},
    {"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, pooled across all grisms, not BGS-specific ({paper_held_out_px} px = {paper_held_out_mm:.4f} mm)"},
]
bgs_comparison_table = pd.DataFrame(bgs_summary_rows).set_index("model")
bgs_comparison_table["rmse_y_px"] = bgs_comparison_table["rmse_y_mm"] / PX_MM_CORRECT
bgs_comparison_table["rmse_z_px"] = bgs_comparison_table["rmse_z_mm"] / PX_MM_CORRECT
bgs_comparison_table[["rmse_y_mm", "rmse_z_mm", "rmse_y_px", "rmse_z_px", "note"]].round(5)

fig, ax = plt.subplots(figsize=(8, 4.5))
labels = ["Physical\n(this notebook)", "Hybrid\n(this notebook)", "Reference\n(in-sample)"]
y_vals = [physical_rmse_y, hybrid_rmse_y, paper_in_sample_y_mm]
z_vals = [physical_rmse_z, hybrid_rmse_z, paper_in_sample_z_mm]
x = np.arange(len(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, all grisms): {paper_held_out_mm:.4f} mm")
ax.set_yscale("log")
ax.set_xticks(x)
ax.set_xticklabels(labels)
ax.set_ylabel("RMSE [mm], held-out test split (log scale)")
ax.legend(loc="upper right", fontsize=8)
ax.set_title("BGS000: this notebook's models vs. the reference paper (arXiv 2506.08378)")
plt.tight_layout()
plt.show()

print("=== BGS000 physical fit ===")
print(f"free parameters: {FREE_PROD}")
print(f"fitted values: {dict(zip(FREE_PROD, theta_star_prod.round(5)))}")
print(f"RMSE (full dataset): {np.sqrt(result_prod.fun):.4f} mm  (before fit: {np.sqrt(cost_before):.4f} mm)")
print(f"excluded_spectra_ids: {EXCLUDED_SPECTRA_IDS}")
print()
print("=== held-out test split ===")
print(f"physical: y={physical_rmse_y:.4f} mm  z={physical_rmse_z:.4f} mm")
print(f"hybrid:   y={hybrid_rmse_y:.4f} mm  z={hybrid_rmse_z:.4f} mm  "
      f"({float(bgs_results['improvement_pct'].iloc[0]):.1f}% combined-distance improvement)")
print(f"reference paper, held-out (Argon, combined, all grisms): {paper_held_out_mm:.4f} mm")
print(f"gap: y {hybrid_rmse_y/paper_held_out_mm:.2f}x, z {hybrid_rmse_z/paper_held_out_mm:.2f}x worse")
print()
print(f"material_k identifiable from BGS000's own passband: {material_k_identifiable}")
print()
print("Follow-up not attempted here: fuse this fit into Stage 5 Phase 3's zeroth-order")
print("material_k analysis -- needs dispcraft/zeroth_dispersion.py's hardcoded RGS rank")
print("wavelengths (1206/1892 nm) generalized to BGS's own passband edges (926/1366 nm) first.")