Stage 3 — Fitting, Residuals, and Hybrid Models (3-Intro_ML)¶

# Load + clean rgs000_0 with the Stage 2 functions, one row per (spectrum, line)
df0_raw = median_per_spectrum(load_spectra(DATA_DIR / "rgs000_0_first.csv"))
print(f"rgs000_0: {len(df0_raw)} rows, {df0_raw['spectra_id'].nunique()} spectra")

FREE_OFFSET_ONLY = ["offset_y_mm", "offset_z_mm"]
theta0_offset = np.array([NOMINAL_PARAMS[k] for k in FREE_OFFSET_ONLY])
print(f"theta0 (nominal offset, mm) = {theta0_offset}")
print(f"cost(theta0) = {cost(theta0_offset, df0_raw, FREE_OFFSET_ONLY, model):.4f} mm^2")

result_raw = minimize(cost, theta0_offset, args=(df0_raw, FREE_OFFSET_ONLY, model), method="Nelder-Mead",
                       options={"maxiter": 5000, "xatol": 1e-8, "fatol": 1e-12})
print("Nelder-Mead result on the full rgs000_0 dataset:")
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)")
print(f"  before-fit cost was {cost(theta0_offset, df0_raw, FREE_OFFSET_ONLY, model):.4f} mm^2 -- barely moved.")
print("  This is suspicious for a 2-parameter linear-ish fit; investigated in the next cell.")

# Per-spectrum mean |r_y| at theta0, sorted -- find what's dominating the (quadratic) MSE cost
pred0 = predict_centroids(df0_raw["y_nisp"], df0_raw["z_nisp"], df0_raw["wavelength"], theta0_offset, FREE_OFFSET_ONLY, model)
df0_raw = df0_raw.assign(r_y=pred0[0] - df0_raw["cent_y"].values, r_z=pred0[1] - df0_raw["cent_z"].values)
per_spectrum = df0_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()),
).sort_values("mean_abs_r_y", ascending=False)
per_spectrum.head(5)

OUTLIER_SPECTRA_ID = int(per_spectrum.index[0])
df0 = df0_raw[df0_raw["spectra_id"] != OUTLIER_SPECTRA_ID].drop(columns=["r_y", "r_z"]).reset_index(drop=True)
print(f"excluded spectra_id={OUTLIER_SPECTRA_ID} (y_nisp=z_nisp=0); "
      f"{len(df0_raw) - len(df0)} rows dropped, {len(df0)} remain")

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

FREE_WIDE = ["offset_y_mm", "offset_z_mm", "tilt_deg", "A_deg", "rho"]
theta0_wide = np.array([NOMINAL_PARAMS[k] for k in FREE_WIDE])

result_wide = minimize(cost, theta0_wide, args=(df0, FREE_WIDE, model), method="Nelder-Mead",
                        options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-14, "adaptive": True})
print("theta* (wide) =", dict(zip(FREE_WIDE, result_wide.x)))
print(f"RMSE = {np.sqrt(result_wide.fun):.4f} mm")

std_wide, corr_wide = bootstrap_uncertainty(result_wide.x, df0, FREE_WIDE, model)
print("\nbootstrap std:")
for name, s in zip(FREE_WIDE, std_wide):
    print(f"  {name:12s} {s:.4f}")
print("\nbootstrap correlation matrix:")
corr_wide.round(2)

FREE_EX1 = ["offset_y_mm", "offset_z_mm", "tilt_deg", "rho"]  # A_deg, material fixed at nominal
theta0_ex1 = np.array([NOMINAL_PARAMS[k] for k in FREE_EX1])

result = minimize(cost, theta0_ex1, args=(df0, FREE_EX1, model), method="Nelder-Mead",
                   options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-14, "adaptive": True})
theta_star = result.x
print(f"cost before fit: {cost(theta0_ex1, df0, FREE_EX1, model):.4f} mm^2")
print(f"theta*     = {dict(zip(FREE_EX1, theta_star))}")
print(f"result.fun = {result.fun:.4f} mm^2  (RMSE = {np.sqrt(result.fun):.4f} mm)")
print(f"(2-param offset-only RMSE was {np.sqrt(result_offset.fun):.4f} mm)")

theta_star_std, corr_ex1 = bootstrap_uncertainty(theta_star, df0, FREE_EX1, model)
print("\nbootstrap std (50 resamples):")
for name, s in zip(FREE_EX1, theta_star_std):
    print(f"  {name:12s} {s:.4f}")
print("\nbootstrap correlation matrix:")
corr_ex1.round(2)

# Compare predictions before (theta0) and after (theta_star) optimization
pred_before = predict_centroids(df0["y_nisp"], df0["z_nisp"], df0["wavelength"], theta0_ex1, FREE_EX1, model)
pred_after = predict_centroids(df0["y_nisp"], df0["z_nisp"], df0["wavelength"], theta_star, FREE_EX1, 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("rgs000_0: observed vs. predicted centroids")
plt.tight_layout()
plt.show()

def write_fit_toml(cfg, df, free_names, theta, theta_std, cost_before, cost_after, excluded_ids=None):
    """Write a Stage 3 per-dataset fit to `models/stage3_fit_<cfg>.toml`.

    Reused for all 6 datasets (Exercise 1 + Exercise 2) so the schema stays
    identical across configs. `A_deg`/material are always fixed at the
    Stage 1 nominal (see the identifiability check above) -- recorded here
    too, so the file is self-contained.
    """
    excluded_ids = excluded_ids or []
    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)))

    text = f'''# Stage 3 -- physical-model fit to {cfg}.
# Generated by notebooks/3-Intro_ML.ipynb. Do not hand-edit the fitted
# values; re-run the notebook and regenerate this file instead.

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

[fit]
dataset = "{cfg}_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 Stage 1 nominal -- not identifiable from a single narrow-band
# ground-test dataset (see notebook Section 2, identifiability check).
A_deg = {p["A_deg"]}
material_n0 = {base_config["material"]["n0"]}
material_k = {base_config["material"]["k"]}

# Fitted parameters, in the same units/frame as stage1_instrument.toml --
# apply these in place of [detector].offset and [grism].tilt_deg / [grating].rho.
[detector]
offset = [{p["offset_y_mm"] / 1000.0:.9f}, {p["offset_z_mm"] / 1000.0:.9f}]

[grism]
tilt_deg = {p["tilt_deg"]:.6f}

[grating]
rho = {p["rho"]:.6f}
'''
    path = MODELS_DIR / f"stage3_fit_{cfg}.toml"
    path.write_text(text)
    print(f"wrote {path}")
    return path


fit_config_path = write_fit_toml("rgs000_0", df0, FREE_EX1, theta_star, theta_star_std,
                                  cost(theta0_ex1, df0, FREE_EX1, model), result.fun,
                                  excluded_ids=[OUTLIER_SPECTRA_ID])
print(fit_config_path.read_text())

OTHER_CONFIGS = [c for c in CONFIGS if c != "rgs000_0"]


def worst_spectrum(cfg, theta0=None, free_names=FREE_OFFSET_ONLY):
    """Per-spectrum mean |r| at a nominal-ish theta -- same diagnostic as Exercise 1's outlier check."""
    theta0 = theta0 if theta0 is not None else np.array([NOMINAL_PARAMS[k] for k in free_names])
    d = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
    pred = predict_centroids(d["y_nisp"], d["z_nisp"], d["wavelength"], theta0, free_names, model)
    r = np.hypot(pred[0] - d["cent_y"].values, pred[1] - d["cent_z"].values)
    d = d.assign(r=r)
    per_spec = d.groupby("spectra_id")["r"].mean().sort_values(ascending=False)
    return d, per_spec


outlier_rows = []
for cfg in OTHER_CONFIGS:
    d, per_spec = worst_spectrum(cfg)
    outlier_rows.append({"config": cfg, "n_spectra": d["spectra_id"].nunique(),
                          "worst_mean_r": per_spec.iloc[0], "2nd_worst_mean_r": per_spec.iloc[1],
                          "ratio": per_spec.iloc[0] / per_spec.iloc[1]})
pd.DataFrame(outlier_rows).set_index("config").round(3)

d180_0 = median_per_spectrum(load_spectra(DATA_DIR / "rgs180_0_first.csv"))

result_180_naive = minimize(cost, theta0_ex1, args=(d180_0, FREE_EX1, model), method="Nelder-Mead",
                             options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-14, "adaptive": True})
print("naive (nominal-start) fit on rgs180_0:")
print(f"  theta* = {dict(zip(FREE_EX1, result_180_naive.x))}")
print(f"  RMSE   = {np.sqrt(result_180_naive.fun):.4f} mm")

# Alternate start: same tilt_deg + 180 deg and same rho as rgs000_0's own fit
theta0_180_alt = theta_star.copy()
theta0_180_alt[FREE_EX1.index("tilt_deg")] += 180.0
result_180_alt = minimize(cost, theta0_180_alt, args=(d180_0, FREE_EX1, model), method="Nelder-Mead",
                           options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-14, "adaptive": True})
print("\nalternate (tilt+180-start) fit on rgs180_0:")
print(f"  theta* = {dict(zip(FREE_EX1, result_180_alt.x))}")
print(f"  RMSE   = {np.sqrt(result_180_alt.fun):.4f} mm")

TILT_IDX = FREE_EX1.index("tilt_deg")


def fit_config(cfg, theta0_seed):
    d = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
    cost_before = cost(theta0_seed, d, FREE_EX1, model)
    r = minimize(cost, theta0_seed, args=(d, FREE_EX1, model), method="Nelder-Mead",
                 options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-14, "adaptive": True})
    std, _ = bootstrap_uncertainty(r.x, d, FREE_EX1, model)
    return {"df": d, "theta": r.x, "std": std, "cost_before": cost_before, "cost_after": r.fun}


fits = {"rgs000_0": {"df": df0, "theta": theta_star, "std": theta_star_std,
                      "cost_before": cost(theta0_ex1, df0, FREE_EX1, model), "cost_after": result.fun}}

# rgs000_m4/p4: same physical grism as rgs000_0, no degeneracy -- nominal start is fine
for cfg in ["rgs000_m4", "rgs000_p4"]:
    fits[cfg] = fit_config(cfg, theta0_ex1)
    print(f"{cfg}: RMSE {np.sqrt(fits[cfg]['cost_before']):.4f} -> {np.sqrt(fits[cfg]['cost_after']):.4f} mm  "
          f"theta*={dict(zip(FREE_EX1, fits[cfg]['theta']))}")

# rgs180_*: seed from the matching rgs000_* fit's tilt+180, same rho (see degeneracy discussion above)
for cfg180, cfg000 in [("rgs180_0", "rgs000_0"), ("rgs180_m4", "rgs000_m4"), ("rgs180_p4", "rgs000_p4")]:
    theta0_seed = fits[cfg000]["theta"].copy()
    theta0_seed[TILT_IDX] += 180.0
    fits[cfg180] = fit_config(cfg180, theta0_seed)
    print(f"{cfg180}: RMSE {np.sqrt(fits[cfg180]['cost_before']):.4f} -> {np.sqrt(fits[cfg180]['cost_after']):.4f} mm  "
          f"theta*={dict(zip(FREE_EX1, fits[cfg180]['theta']))}")

# Compare fitted parameters across the 6 configs -- does rgs180 track rgs000 + 180deg tilt / close rho?
param_rows = []
for cfg in CONFIGS:
    p = dict(zip(FREE_EX1, fits[cfg]["theta"]))
    param_rows.append({"config": cfg, "offset_y_mm": p["offset_y_mm"], "offset_z_mm": p["offset_z_mm"],
                        "tilt_deg_mod360": p["tilt_deg"] % 360, "rho": p["rho"],
                        "RMSE_mm": np.sqrt(fits[cfg]["cost_after"])})
param_table = pd.DataFrame(param_rows).set_index("config")
print(param_table.round(4))

print("\nrgs180 vs. matching rgs000 (tilt_deg + 180, rho):")
for cfg180, cfg000 in [("rgs180_0", "rgs000_0"), ("rgs180_m4", "rgs000_m4"), ("rgs180_p4", "rgs000_p4")]:
    t000 = fits[cfg000]["theta"][TILT_IDX] % 360
    t180 = fits[cfg180]["theta"][TILT_IDX] % 360
    rho000 = dict(zip(FREE_EX1, fits[cfg000]["theta"]))["rho"]
    rho180 = dict(zip(FREE_EX1, fits[cfg180]["theta"]))["rho"]
    print(f"  {cfg180}: tilt={t180:.3f} vs. {cfg000}+180={((t000 + 180) % 360):.3f}  "
          f"(diff {(t180 - t000 - 180 + 180) % 360 - 180:+.3f} deg)   "
          f"rho={rho180:.4f} vs. {rho000:.4f} (diff {100 * (rho180 - rho000) / rho000:+.2f}%)")

for cfg in OTHER_CONFIGS:
    write_fit_toml(cfg, fits[cfg]["df"], FREE_EX1, fits[cfg]["theta"], fits[cfg]["std"],
                    fits[cfg]["cost_before"], fits[cfg]["cost_after"])

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


# Each dataset evaluated at its OWN Exercise-2 fit (fits[cfg]["theta"]), not a shared theta_star
datasets = {cfg: add_residuals(fits[cfg]["df"], fits[cfg]["theta"], FREE_EX1) for cfg in CONFIGS}

summary_rows = []
for cfg, d in datasets.items():
    summary_rows.append({
        "config": cfg,
        "N": len(d),
        "RMSE_y": np.sqrt(np.mean(d["r_y"] ** 2)),
        "RMSE_z": np.sqrt(np.mean(d["r_z"] ** 2)),
        "MAE_y": np.mean(np.abs(d["r_y"])),
        "MAE_z": np.mean(np.abs(d["r_z"])),
        "max_y": np.max(np.abs(d["r_y"])),
        "max_z": np.max(np.abs(d["r_z"])),
    })
summary = pd.DataFrame(summary_rows).set_index("config")
summary.round(3)

# Histograms of r_y, r_z for all 6 datasets
fig, axes = plt.subplots(2, 6, figsize=(20, 6), sharex="row")
for col, (cfg, d) in enumerate(datasets.items()):
    axes[0, col].hist(d["r_y"], bins=40, color="tab:blue")
    axes[0, col].set_title(cfg)
    axes[1, col].hist(d["r_z"], bins=40, color="tab:orange")
axes[0, 0].set_ylabel("r_y")
axes[1, 0].set_ylabel("r_z")
fig.suptitle("Residual histograms, all 6 datasets (each at its own Exercise 2 fit)")
plt.tight_layout()
plt.show()

# r_y vs y_nisp (trend with field position?) for all 6 datasets
fig, axes = plt.subplots(1, 6, figsize=(20, 3.2), sharey=True)
for ax, (cfg, d) in zip(axes, datasets.items()):
    ax.scatter(d["y_nisp"], d["r_y"], s=3, alpha=0.4)
    ax.axhline(0, color="k", lw=0.5)
    ax.set_title(cfg)
    ax.set_xlabel("y_nisp [mm]")
axes[0].set_ylabel("r_y [mm]")
fig.suptitle("r_y vs. field position y_nisp")
plt.tight_layout()
plt.show()

# 2D residual maps: (cent_y, cent_z) colored by |r_y|, for all 6 datasets
fig, axes = plt.subplots(1, 6, figsize=(22, 3.6))
for ax, (cfg, d) in zip(axes, datasets.items()):
    sc = ax.scatter(d["cent_y"], d["cent_z"], c=d["r_y"].abs(), s=4, cmap="viridis")
    ax.set_title(cfg)
    ax.set_xlabel("cent_y [mm]")
    fig.colorbar(sc, ax=ax, label="|r_y| [mm]", fraction=0.046)
axes[0].set_ylabel("cent_z [mm]")
fig.suptitle("2D residual maps, |r_y|")
plt.tight_layout()
plt.show()

d0 = datasets["rgs000_0"]
X = d0[["y_nisp", "z_nisp", "wavelength"]].values
groups = d0["spectra_id"].values

gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=RNG_SEED)
train_idx, test_idx = next(gss.split(X, groups=groups))
print(f"train: {len(train_idx)} rows, {d0['spectra_id'].iloc[train_idx].nunique()} spectra")
print(f"test:  {len(test_idx)} rows, {d0['spectra_id'].iloc[test_idx].nunique()} spectra")

scaler = StandardScaler().fit(X[train_idx])  # fit on train only
X_train_s = scaler.transform(X[train_idx])
X_test_s = scaler.transform(X[test_idx])

ml_models = {}  # axis -> {"linear": ..., "rf": ...}
ml_rows = []
for axis, col in [("y", "r_y"), ("z", "r_z")]:
    y = d0[col].values
    y_train, y_test = y[train_idx], y[test_idx]

    linear = LinearRegression().fit(X_train_s, y_train)
    rf = RandomForestRegressor(n_estimators=200, max_depth=8, random_state=RNG_SEED).fit(X_train_s, y_train)
    ml_models[axis] = {"linear": linear, "rf": rf}

    phys_rmse = np.sqrt(np.mean(y_test**2))  # physics-only = residual itself, no correction
    lin_rmse = np.sqrt(mean_squared_error(y_test, linear.predict(X_test_s)))
    rf_rmse = np.sqrt(mean_squared_error(y_test, rf.predict(X_test_s)))
    cv_rf = cross_val_score(RandomForestRegressor(n_estimators=200, max_depth=8, random_state=RNG_SEED),
                             X_train_s, y_train, cv=5, scoring="neg_root_mean_squared_error")

    ml_rows.append({"axis": axis, "physics_only_RMSE": phys_rmse, "linear_RMSE": lin_rmse,
                     "RF_RMSE": rf_rmse, "RF_cv_RMSE_mean": -cv_rf.mean(), "RF_cv_RMSE_std": cv_rf.std()})

ml_summary = pd.DataFrame(ml_rows).set_index("axis")
ml_summary.round(4)

# Predicted vs. actual residual, test set, both axes
fig, axes = plt.subplots(1, 2, figsize=(10, 5))
for ax, axis, col in zip(axes, ["y", "z"], ["r_y", "r_z"]):
    y_test = d0[col].values[test_idx]
    lin_pred = ml_models[axis]["linear"].predict(X_test_s)
    rf_pred = ml_models[axis]["rf"].predict(X_test_s)
    lims = [y_test.min(), y_test.max()]
    ax.plot(lims, lims, "k--", lw=1, label="perfect")
    ax.scatter(y_test, lin_pred, s=6, alpha=0.5, label="Linear")
    ax.scatter(y_test, rf_pred, s=6, alpha=0.5, label="RandomForest")
    ax.set_xlabel(f"actual {col} [mm]")
    ax.set_ylabel(f"predicted {col} [mm]")
    ax.legend()
fig.suptitle("rgs000_0 test set: predicted vs. actual residual")
plt.tight_layout()
plt.show()

# Hybrid vs. physics-only RMSE on the rgs000_0 test set (RandomForest as the ML component)
hybrid_rows = []
for axis, col in [("y", "r_y"), ("z", "r_z")]:
    r_test = d0[col].values[test_idx]
    ml_pred = ml_models[axis]["rf"].predict(X_test_s)
    hybrid_resid = r_test - ml_pred  # what's left after the ML correction
    hybrid_rows.append({
        "axis": axis,
        "physics_only_RMSE": np.sqrt(np.mean(r_test**2)),
        "hybrid_RMSE": np.sqrt(np.mean(hybrid_resid**2)),
    })
hybrid_summary = pd.DataFrame(hybrid_rows).set_index("axis")
hybrid_summary.round(4)

# Overlaid residual histograms + side-by-side 2D maps, physics-only vs hybrid (y-axis, test set)
r_y_test = d0["r_y"].values[test_idx]
hybrid_r_y_test = r_y_test - ml_models["y"]["rf"].predict(X_test_s)
cent_y_test = d0["cent_y"].values[test_idx]
cent_z_test = d0["cent_z"].values[test_idx]

fig, axes = plt.subplots(1, 3, figsize=(16, 4.5))
axes[0].hist(r_y_test, bins=30, alpha=0.6, label="physics-only")
axes[0].hist(hybrid_r_y_test, bins=30, alpha=0.6, label="hybrid")
axes[0].set_xlabel("r_y [mm]")
axes[0].legend()
axes[0].set_title("Residual histogram (overlaid)")

for ax, resid, title in zip(axes[1:], [r_y_test, hybrid_r_y_test], ["physics-only", "hybrid"]):
    sc = ax.scatter(cent_y_test, cent_z_test, c=np.abs(resid), s=8, cmap="viridis",
                     vmin=0, vmax=np.abs(r_y_test).max())
    ax.set_xlabel("cent_y [mm]")
    ax.set_title(f"2D map, |r_y| -- {title}")
    fig.colorbar(sc, ax=ax, label="|r_y| [mm]")
axes[1].set_ylabel("cent_z [mm]")
fig.suptitle("rgs000_0 test set: physics-only vs. hybrid")
plt.tight_layout()
plt.show()

X_all_s = scaler.transform(X)  # all of rgs000_0, same scaler (fit on the train split above)
d180 = datasets["rgs180_0"]  # residuals from rgs180_0's OWN Exercise-2 fit (fits["rgs180_0"]), not rgs000_0's
X180_s = scaler.transform(d180[["y_nisp", "z_nisp", "wavelength"]].values)

gen_rows = []
for axis, col in [("y", "r_y"), ("z", "r_z")]:
    rf_full = RandomForestRegressor(n_estimators=200, max_depth=8, random_state=RNG_SEED)
    rf_full.fit(X_all_s, d0[col].values)
    r180 = d180[col].values
    ml_pred180 = rf_full.predict(X180_s)
    hybrid180 = r180 - ml_pred180
    gen_rows.append({
        "axis": axis,
        "physics_only_RMSE": np.sqrt(np.mean(r180**2)),
        "hybrid_RMSE": np.sqrt(np.mean(hybrid180**2)),
        "generalizes": np.sqrt(np.mean(hybrid180**2)) < 0.9 * np.sqrt(np.mean(r180**2)),
    })
gen_summary = pd.DataFrame(gen_rows).set_index("axis")
gen_summary.round(4)

from sklearn.ensemble import HistGradientBoostingRegressor

GEN_THRESHOLD = 0.9  # hybrid RMSE must beat physics-only by >10% on rgs180_0 to count as "generalizes"

# Physical model, before vs. after the Exercise-1/2 fit, on the rgs000_0 test rows
pred_test_before = predict_centroids(d0["y_nisp"].values[test_idx], d0["z_nisp"].values[test_idx],
                                      d0["wavelength"].values[test_idx], theta0_ex1, FREE_EX1, model)
r_y_test_before = pred_test_before[0] - d0["cent_y"].values[test_idx]
r_z_test_before = pred_test_before[1] - d0["cent_z"].values[test_idx]

rows = [
    {"model": "Physical (before fit)", "RMSE_y": np.sqrt(np.mean(r_y_test_before**2)),
     "RMSE_z": np.sqrt(np.mean(r_z_test_before**2)), "Generalizes?": "--"},
    {"model": "Physical (after fit)", "RMSE_y": np.sqrt(np.mean(d0["r_y"].values[test_idx] ** 2)),
     "RMSE_z": np.sqrt(np.mean(d0["r_z"].values[test_idx] ** 2)), "Generalizes?": "Yes (by construction)"},
]

ml_estimators = {
    "Physical + Linear": lambda: LinearRegression(),
    "Physical + Random Forest": lambda: RandomForestRegressor(n_estimators=200, max_depth=8, random_state=RNG_SEED),
    "Physical + Gradient Boosting": lambda: HistGradientBoostingRegressor(max_depth=6, random_state=RNG_SEED),
}

for name, make_estimator in ml_estimators.items():
    row = {"model": name}
    gen_flags = []
    for axis, col in [("y", "r_y"), ("z", "r_z")]:
        est_test = make_estimator().fit(X_train_s, d0[col].values[train_idx])
        hybrid_test = d0[col].values[test_idx] - est_test.predict(X_test_s)
        row[f"RMSE_{axis}"] = np.sqrt(np.mean(hybrid_test**2))

        est_full = make_estimator().fit(X_all_s, d0[col].values)
        hybrid180 = d180[col].values - est_full.predict(X180_s)
        phys180_rmse = np.sqrt(np.mean(d180[col].values ** 2))
        gen_flags.append(np.sqrt(np.mean(hybrid180**2)) < GEN_THRESHOLD * phys180_rmse)
    row["Generalizes?"] = "y: Yes, z: Yes" if all(gen_flags) else (
        "y: Yes, z: No" if gen_flags[0] and not gen_flags[1] else
        "y: No, z: Yes" if gen_flags[1] and not gen_flags[0] else "No")
    rows.append(row)

comparison_table = pd.DataFrame(rows).set_index("model")
comparison_table.round(4)

final_rows = []

# rgs000_0: own-trained RF, test-set accuracy (Exercise 4/5)
r_y_test = d0["r_y"].values[test_idx]
r_z_test = d0["r_z"].values[test_idx]
hybrid_y_test = r_y_test - ml_models["y"]["rf"].predict(X_test_s)
hybrid_z_test = r_z_test - ml_models["z"]["rf"].predict(X_test_s)
final_rows.append({
    "config": "rgs000_0",
    "physics_RMSE_y": ml_summary.loc["y", "physics_only_RMSE"], "physics_RMSE_z": ml_summary.loc["z", "physics_only_RMSE"],
    "hybrid_RMSE_y": ml_summary.loc["y", "RF_RMSE"], "hybrid_RMSE_z": ml_summary.loc["z", "RF_RMSE"],
    "physics_mean_dist_mm": np.mean(np.hypot(r_y_test, r_z_test)),
    "hybrid_mean_dist_mm": np.mean(np.hypot(hybrid_y_test, hybrid_z_test)),
    "used_ML_y": True, "used_ML_z": True, "ML_source": "own-trained (Exercise 4 test set)",
})

# rgs000_0-trained RF, applied to each other config's own (full-dataset) Exercise-2 residual --
# adopt the correction per-axis only where it actually beats that config's own physics-only RMSE
rf_full_by_axis = {axis: RandomForestRegressor(n_estimators=200, max_depth=8, random_state=RNG_SEED).fit(X_all_s, d0[col].values)
                   for axis, col in [("y", "r_y"), ("z", "r_z")]}

for cfg in OTHER_CONFIGS:
    d = datasets[cfg]
    Xc_s = scaler.transform(d[["y_nisp", "z_nisp", "wavelength"]].values)
    row = {"config": cfg, "ML_source": "rgs000_0-trained (cross-instance)"}
    hybrid_resid = {}
    for axis, col in [("y", "r_y"), ("z", "r_z")]:
        r = d[col].values
        phys_rmse = np.sqrt(np.mean(r ** 2))
        ml_resid = r - rf_full_by_axis[axis].predict(Xc_s)
        hybrid_rmse = np.sqrt(np.mean(ml_resid ** 2))
        used_ml = hybrid_rmse < phys_rmse
        row[f"physics_RMSE_{axis}"] = phys_rmse
        row[f"hybrid_RMSE_{axis}"] = hybrid_rmse if used_ml else phys_rmse  # only adopt if it helps
        row[f"used_ML_{axis}"] = used_ml
        hybrid_resid[axis] = ml_resid if used_ml else r  # per-axis best residual, for the joint distance below
    row["physics_mean_dist_mm"] = np.mean(np.hypot(d["r_y"].values, d["r_z"].values))
    row["hybrid_mean_dist_mm"] = np.mean(np.hypot(hybrid_resid["y"], hybrid_resid["z"]))
    final_rows.append(row)

final_table = pd.DataFrame(final_rows).set_index("config")
final_table = final_table[["physics_RMSE_y", "hybrid_RMSE_y", "used_ML_y",
                            "physics_RMSE_z", "hybrid_RMSE_z", "used_ML_z",
                            "physics_mean_dist_mm", "hybrid_mean_dist_mm", "ML_source"]]
final_table.round(4)