Stage 4, Phase 3 — Multi-Dataset Fitting (4.3-Multi_Dataset_Fitting)¶

GLOBAL_PARAMS = ["coll_f", "cam_f"]
INSTANCE_PARAMS = ["A_deg", "material_n0", "rho"]
SPECIFIC_PARAMS = ["tilt_deg", "offset_y_mm", "offset_z_mm"]
ALL_PARAMS = GLOBAL_PARAMS + INSTANCE_PARAMS + SPECIFIC_PARAMS

print("Tier 1 (global, all 6 datasets):       ", GLOBAL_PARAMS)
print("Tier 2 (per grism instance, 3 each):    ", INSTANCE_PARAMS)
print("Tier 3 (dataset-specific, each of 6):   ", SPECIFIC_PARAMS)

FULL_NOMINAL = {
    **model.nominal_params,
    "coll_f": model.coll.f,
    "cam_f": model.cam.f,
    "material_n0": model.material.n0,
}


def predict_centroids_full(y_nisp_mm, z_nisp_mm, wavelength_nm, theta, free_names, base_model, fixed=None):
    """Like calibration.predict_centroids, but also allows coll_f, cam_f, and
    material_n0 to vary. Notebook-local for Phase 3's 3-tier joint fit --
    not promoted to dispcraft/calibration.py pending validation."""
    p = dict(FULL_NOMINAL)
    if fixed:
        p.update(fixed)
    p.update(dict(zip(free_names, theta)))

    material = Material(n0=p["material_n0"], k=base_model.material.k)
    prism = Prism(material=material, A=np.radians(p["A_deg"]))
    grating = Grating(m=base_model.m_order, rho=p["rho"])
    grism = Grism(prism, grating, tilt=np.radians(p["tilt_deg"]))
    coll, cam = Collimator(f=p["coll_f"]), Camera(f=p["cam_f"])

    pos_foc = np.stack([np.asarray(y_nisp_mm), np.asarray(z_nisp_mm)]) / 1000.0  # mm -> m
    wavelength_um = np.asarray(wavelength_nm) / 1000.0  # nm -> um
    angle_col = coll.forward(pos_foc)
    angle_gr = grism.forward(angle_col, wavelength_um)
    pos_cam = cam.forward(angle_gr)  # m
    offset_mm = np.array([p["offset_y_mm"], p["offset_z_mm"]])
    return pos_cam * 1000.0 + offset_mm[:, None]  # mm


def cost_full(theta, df, free_names, base_model, fixed=None):
    """MSE cost over (r_y, r_z), vectorized over all rows of df. Same shape as
    calibration.cost, built on predict_centroids_full."""
    pred = predict_centroids_full(df["y_nisp"], df["z_nisp"], df["wavelength"], theta, free_names, base_model, fixed)
    r_y = pred[0] - df["cent_y"].values
    r_z = pred[1] - df["cent_z"].values
    return np.mean(r_y**2 + r_z**2)


# Sanity check: with coll_f/cam_f/material_n0 fixed at nominal and the same
# free_names/theta calibration.predict_centroids supports, the two must agree.
from dispcraft.calibration import cost as cost_orig

_df_check = median_per_spectrum(load_spectra(DATA_DIR / "rgs000_0_first.csv"))
_free_check = ["offset_y_mm", "offset_z_mm", "tilt_deg", "rho"]
_theta_check = np.array([FULL_NOMINAL[k] for k in _free_check])
_c_orig = cost_orig(_theta_check, _df_check, _free_check, model)
_c_full = cost_full(_theta_check, _df_check, _free_check, model)
assert np.isclose(_c_orig, _c_full), (_c_orig, _c_full)
print(f"predict_centroids_full/cost_full agree with calibration.cost at nominal: {_c_orig:.6f} mm^2")

def load_fit(cfg):
    with open(MODELS_DIR / f"stage3_fit_{cfg}.toml", "rb") as f:
        fit = tomllib.load(f)
    free_names = fit["fit"]["free_parameters"]
    theta = np.array([fit["fit"]["result"][k] for k in free_names])
    excluded_ids = fit["fit"]["excluded_spectra_ids"]
    return fit, free_names, theta, excluded_ids


dfs = {}
stage1_fits = {}
for cfg in CONFIGS:
    fit, free_names, theta, excluded_ids = load_fit(cfg)
    df = median_per_spectrum(load_spectra(DATA_DIR / f"{cfg}_first.csv"))
    df = df[~df["spectra_id"].isin(excluded_ids)].reset_index(drop=True)
    assert len(df) == fit["fit"]["n_points"], f"{cfg}: row count must match the TOML's recorded fit"
    dfs[cfg] = df
    stage1_fits[cfg] = {"free_names": free_names, "theta": theta, "params": dict(zip(free_names, theta)),
                         "rmse_mm": fit["fit"]["rmse_mm"], "excluded_ids": excluded_ids}

print(f"loaded {len(dfs)} datasets, row counts: {[len(dfs[c]) for c in CONFIGS]}")

stage1_table = pd.DataFrame({
    cfg: {**stage1_fits[cfg]["params"], "tilt_deg_mod360": stage1_fits[cfg]["params"]["tilt_deg"] % 360,
          "rmse_mm": stage1_fits[cfg]["rmse_mm"]}
    for cfg in CONFIGS
}).T[["offset_y_mm", "offset_z_mm", "tilt_deg_mod360", "rho", "rmse_mm"]]
stage1_table.round(4)

def fixed_for_tier1(cfg, stage1_fits):
    """All non-global params held at Stage-1 values (A_deg/material_n0 at
    nominal, matching Stage 1's own fixed treatment of them)."""
    p = dict(stage1_fits[cfg]["params"])  # offset_y_mm, offset_z_mm, tilt_deg, rho
    p["A_deg"] = FULL_NOMINAL["A_deg"]
    p["material_n0"] = FULL_NOMINAL["material_n0"]
    return p


def cost_joint(theta, dfs, free_names, fixed_by_cfg, model):
    """Unweighted sum of cost_full(...) over all datasets in `dfs`, matching
    the deck's own pseudocode (total += mse(...), no n_points weighting)."""
    return sum(cost_full(theta, dfs[cfg], free_names, model, fixed=fixed_by_cfg[cfg]) for cfg in dfs)


def per_dataset_rmse(theta, free_names, fixed_by_cfg, dfs, model):
    return {cfg: float(np.sqrt(cost_full(theta, dfs[cfg], free_names, model, fixed=fixed_by_cfg[cfg])))
            for cfg in dfs}


fixed_tier1 = {cfg: fixed_for_tier1(cfg, stage1_fits) for cfg in CONFIGS}
theta0_tier1 = np.array([FULL_NOMINAL[k] for k in GLOBAL_PARAMS])
print(f"theta0 (nominal coll_f, cam_f) = {theta0_tier1}")
print(f"cost_joint(theta0) = {cost_joint(theta0_tier1, dfs, GLOBAL_PARAMS, fixed_tier1, model):.4f} mm^2 (summed over 6 datasets)")

result_tier1 = minimize(cost_joint, theta0_tier1, args=(dfs, GLOBAL_PARAMS, fixed_tier1, model),
                         method="Nelder-Mead",
                         options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-13, "adaptive": True})
theta_star_tier1 = result_tier1.x
print(f"theta* (coll_f, cam_f) = {dict(zip(GLOBAL_PARAMS, theta_star_tier1))}")
print(f"cost_joint: {cost_joint(theta0_tier1, dfs, GLOBAL_PARAMS, fixed_tier1, model):.4f} -> {result_tier1.fun:.4f} mm^2 (sum over 6 datasets)")

# Multi-start validation: only 2 free dims, no degenerate branches reachable
# since Tier 2/3 params are frozen at already-resolved Stage-1 values -- a
# cheap check, not a hard requirement.
starts = [theta0_tier1, theta0_tier1 + [0.05, 0], theta0_tier1 - [0.05, 0],
          theta0_tier1 + [0, 0.05], theta0_tier1 - [0, 0.05]]
multi_start_thetas = []
for s in starts:
    r = minimize(cost_joint, s, args=(dfs, GLOBAL_PARAMS, fixed_tier1, model), method="Nelder-Mead",
                 options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-13, "adaptive": True})
    multi_start_thetas.append(r.x)
multi_start_thetas = np.array(multi_start_thetas)
print(f"\nmulti-start thetas:\n{multi_start_thetas}")
print(f"multi-start spread (std): {multi_start_thetas.std(axis=0)}  -- should be << bootstrap std below")

def bootstrap_uncertainty_joint(theta_star, dfs, free_names, fixed_by_cfg, model, n_boot=30, seed=RNG_SEED):
    """Bootstrap std + correlation matrix for a joint multi-dataset fit.

    Resamples each dataset independently (own size, with replacement),
    recombines into cost_joint, re-minimizes from theta_star each replicate.
    No existing function covers this multi-dataset case (bootstrap_uncertainty
    in 3-Intro_ML.ipynb only resamples one dataset)."""
    rng = np.random.default_rng(seed)
    boots = []
    for _ in range(n_boot):
        resampled = {cfg: dfs[cfg].iloc[rng.integers(0, len(dfs[cfg]), len(dfs[cfg]))] for cfg in dfs}
        r = minimize(cost_joint, theta_star, args=(resampled, free_names, fixed_by_cfg, model),
                     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


std_tier1, corr_tier1 = bootstrap_uncertainty_joint(theta_star_tier1, dfs, GLOBAL_PARAMS, fixed_tier1, model)
print("bootstrap std (30 resamples):")
for name, s in zip(GLOBAL_PARAMS, std_tier1):
    print(f"  {name:10s} {s:.6f}")
print("\nbootstrap correlation matrix:")
corr_tier1.round(3)

# Per-dataset RMSE: Stage 1 (independent fit) vs. Tier-1-forced (coll_f/cam_f
# jointly fit, everything else still at Stage-1 values). Note cost_joint's
# converged value is a SUM of 6 MSEs, not a pooled quantity given unequal
# n_points -- always report the per-dataset breakdown, never sqrt(cost/6).
rmse_tier1 = per_dataset_rmse(theta_star_tier1, GLOBAL_PARAMS, fixed_tier1, dfs, model)
rmse_before_tier1 = per_dataset_rmse(theta0_tier1, GLOBAL_PARAMS, fixed_tier1, dfs, model)

comparison_tier1 = pd.DataFrame({
    "stage1_rmse_mm": {cfg: stage1_fits[cfg]["rmse_mm"] for cfg in CONFIGS},
    "tier1_before_rmse_mm": rmse_before_tier1,
    "tier1_after_rmse_mm": rmse_tier1,
}).loc[CONFIGS]
comparison_tier1.round(4)

# Residual (data - model) histograms at the Tier-1-forced parameters, all 6
# datasets -- inspected for structure, not just aggregate RMSE.
fig, axes = plt.subplots(2, 6, figsize=(20, 6), sharex="row")
for col, cfg in enumerate(CONFIGS):
    pred = predict_centroids_full(dfs[cfg]["y_nisp"], dfs[cfg]["z_nisp"], dfs[cfg]["wavelength"],
                                   theta_star_tier1, GLOBAL_PARAMS, model, fixed=fixed_tier1[cfg])
    r_y = pred[0] - dfs[cfg]["cent_y"].values
    r_z = pred[1] - dfs[cfg]["cent_z"].values
    axes[0, col].hist(r_y, bins=40, color="tab:blue")
    axes[0, col].set_title(cfg)
    axes[1, col].hist(r_z, bins=40, color="tab:orange")
axes[0, 0].set_ylabel("r_y [mm]")
axes[1, 0].set_ylabel("r_z [mm]")
fig.suptitle("Residuals at Tier-1-forced parameters (coll_f/cam_f jointly fit, Tier 2/3 at Stage-1 values)")
plt.tight_layout()
plt.show()

def write_tier1_toml(theta, std, cost_before, cost_after, rmse_by_cfg):
    result_lines = "\n".join(f"{name} = {v:.6f}" for name, v in zip(GLOBAL_PARAMS, theta))
    std_lines = "\n".join(f"{name}_bootstrap_std = {s:.6f}" for name, s in zip(GLOBAL_PARAMS, std))
    rmse_lines = "\n".join(f"{cfg} = {rmse_by_cfg[cfg]:.6f}" for cfg in CONFIGS)

    text = f'''# Stage 4 Phase 3 -- Tier 1: global joint fit of coll_f/cam_f across all 6
# datasets. Generated by notebooks/4.3-Multi_Dataset_Fitting.ipynb. Do not
# hand-edit; re-run the notebook and regenerate this file instead.

[base_config]
path = "stage1_instrument.toml"

[fit]
tier = "global"
datasets = {CONFIGS}
method = "Nelder-Mead"
free_parameters = {GLOBAL_PARAMS}
cost_joint_before_mm2 = {cost_before:.6f}
cost_joint_after_mm2 = {cost_after:.6f}

[fit.result]
{result_lines}
{std_lines}

# Per-dataset RMSE at this Tier-1 result, with Tier 2/3 params still at
# their Stage-1 (independent) values -- not yet re-optimized.
[fit.per_dataset_rmse_mm]
{rmse_lines}
'''
    path = MODELS_DIR / "joint_global_shared.toml"
    path.write_text(text)
    print(f"wrote {path}")
    return path


cost_before_tier1 = cost_joint(theta0_tier1, dfs, GLOBAL_PARAMS, fixed_tier1, model)
tier1_toml_path = write_tier1_toml(theta_star_tier1, std_tier1, cost_before_tier1, result_tier1.fun, rmse_tier1)
print(tier1_toml_path.read_text())

mlflow.set_tracking_uri(f"sqlite:///{(REPO_ROOT / 'mlflow.db').resolve()}")
EXPERIMENT_NAME_TIER1 = "stage2"  # deck-mandated name (Task 3 step 2); collides in spelling only
# with this project's own project-Stage-3 -- disambiguated via TOML filenames instead.
if mlflow.get_experiment_by_name(EXPERIMENT_NAME_TIER1) is None:
    mlflow.create_experiment(EXPERIMENT_NAME_TIER1, artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(EXPERIMENT_NAME_TIER1)

with mlflow.start_run(run_name="global_shared_fit"):
    mlflow.log_param("tier", "global")
    mlflow.log_param("free_parameters", GLOBAL_PARAMS)
    mlflow.log_param("n_datasets", len(CONFIGS))
    for name, v, s in zip(GLOBAL_PARAMS, theta_star_tier1, std_tier1):
        mlflow.log_param(name, v)
        mlflow.log_metric(f"{name}_bootstrap_std", s)
    mlflow.log_metric("cost_joint_before_mm2", cost_before_tier1)
    mlflow.log_metric("cost_joint_after_mm2", result_tier1.fun)
    for cfg in CONFIGS:
        mlflow.log_metric(f"rmse_{cfg}", rmse_tier1[cfg])

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

def bootstrap_uncertainty_wide(theta_star, dfs_group, shared_names, cfgs, fixed_by_cfg, model,
                                n_boot=30, seed=RNG_SEED):
    rng = np.random.default_rng(seed)
    names = shared_names + [f"offset_z_mm__{cfg}" for cfg in cfgs]
    boots = []
    for _ in range(n_boot):
        resampled = {cfg: dfs_group[cfg].iloc[rng.integers(0, len(dfs_group[cfg]), len(dfs_group[cfg]))] for cfg in cfgs}
        r = minimize(cost_joint_wide, theta_star, args=(resampled, shared_names, cfgs, fixed_by_cfg, model),
                     method="Nelder-Mead",
                     options={"maxiter": 12000, "maxfev": 12000, "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=names, columns=names)
    return std, corr


wide_results = {}
for instance, cfgs in GRISM_INSTANCES.items():
    fixed_wide = {cfg: fixed_for_tier2(cfg, stage1_fits, theta_star_tier1) for cfg in cfgs}
    rho0 = np.mean([stage1_fits[cfg]["params"]["rho"] for cfg in cfgs])
    theta0_wide = np.concatenate([
        [FULL_NOMINAL["A_deg"], FULL_NOMINAL["material_n0"], rho0],
        [stage1_fits[cfg]["params"]["offset_z_mm"] for cfg in cfgs],
    ])
    dfs_group = {cfg: dfs[cfg] for cfg in cfgs}

    r_wide = minimize(cost_joint_wide, theta0_wide, args=(dfs_group, INSTANCE_PARAMS, cfgs, fixed_wide, model),
                       method="Nelder-Mead",
                       options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-13, "adaptive": True})
    std_wide, corr_wide = bootstrap_uncertainty_wide(r_wide.x, dfs_group, INSTANCE_PARAMS, cfgs, fixed_wide, model)
    wide_results[instance] = {"theta": r_wide.x, "std": std_wide, "corr": corr_wide, "cost_after": r_wide.fun}

    names = INSTANCE_PARAMS + [f"offset_z_mm__{cfg}" for cfg in cfgs]
    print(f"=== {instance} wide identifiability check ===")
    print(f"theta* = {dict(zip(names, r_wide.x))}")
    print("bootstrap std:", dict(zip(names, std_wide.round(5))))
    print()

for instance in GRISM_INSTANCES:
    print(f"=== {instance}: bootstrap correlation matrix ===")
    corr = wide_results[instance]["corr"]
    print(corr.round(3))
    a_deg_corrs = corr.loc["A_deg"].drop("A_deg")
    print(f"\nmax |corr(A_deg, other)| = {a_deg_corrs.abs().max():.3f} "
          f"(with {a_deg_corrs.abs().idxmax()})  -- vs. ~0.99 single-dataset (3-Intro_ML.ipynb)")
    print()

INSTANCE_FREE_PARAMS = ["A_deg", "rho"]  # material_n0 fixed at nominal -- see identifiability check above

tier2_results = {}
for instance, cfgs in GRISM_INSTANCES.items():
    fixed_tier2 = {cfg: fixed_for_tier2(cfg, stage1_fits, theta_star_tier1) for cfg in cfgs}
    rho0 = np.mean([stage1_fits[cfg]["params"]["rho"] for cfg in cfgs])
    theta0 = np.array([FULL_NOMINAL["A_deg"], rho0])
    dfs_group = {cfg: dfs[cfg] for cfg in cfgs}

    cost_before = cost_joint(theta0, dfs_group, INSTANCE_FREE_PARAMS, fixed_tier2, model)
    r = minimize(cost_joint, theta0, args=(dfs_group, INSTANCE_FREE_PARAMS, fixed_tier2, model),
                 method="Nelder-Mead",
                 options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-13, "adaptive": True})
    std, corr = bootstrap_uncertainty_joint(r.x, dfs_group, INSTANCE_FREE_PARAMS, fixed_tier2, model)

    tier2_results[instance] = {"theta": r.x, "std": std, "corr": corr,
                                "cost_before": cost_before, "cost_after": r.fun, "fixed_by_cfg": fixed_tier2}
    print(f"=== {instance} (production Tier 2: A_deg + rho) ===")
    print(f"theta* = {dict(zip(INSTANCE_FREE_PARAMS, r.x))}")
    print(f"bootstrap std = {dict(zip(INSTANCE_FREE_PARAMS, std.round(5)))}")
    print(f"cost_joint: {cost_before:.4f} -> {r.fun:.4f} mm^2 (sum over {len(cfgs)} datasets)")
    print("bootstrap correlation:")
    print(corr.round(3))
    print()

def full_fixed_for_tier12(cfg, stage1_fits, theta_star_tier1, tier2_results):
    """Complete parameter dict for `cfg` at the current Tier-1+2-forced state:
    coll_f/cam_f (Tier 1), A_deg/rho (Tier 2, this cfg's grism instance,
    material_n0 at nominal), tilt_deg/offset_y_mm/offset_z_mm (Stage 1,
    Tier 3 not yet re-optimized)."""
    p = fixed_for_tier2(cfg, stage1_fits, theta_star_tier1)
    inst = instance_of(cfg)
    p.update(dict(zip(INSTANCE_FREE_PARAMS, tier2_results[inst]["theta"])))
    return p


fixed_tier12 = {cfg: full_fixed_for_tier12(cfg, stage1_fits, theta_star_tier1, tier2_results) for cfg in CONFIGS}
rmse_tier12 = {cfg: float(np.sqrt(cost_full(np.array([]), dfs[cfg], [], model, fixed=fixed_tier12[cfg])))
               for cfg in CONFIGS}

comparison_tier12 = pd.DataFrame({
    "stage1_rmse_mm": {cfg: stage1_fits[cfg]["rmse_mm"] for cfg in CONFIGS},
    "tier1_forced_rmse_mm": rmse_tier1,
    "tier1+2_forced_rmse_mm": rmse_tier12,
}).loc[CONFIGS]
comparison_tier12.round(4)

# rho: production Tier-2 (instance-shared) vs. Stage-1 per-dataset values
rho_compare_rows = []
for instance, cfgs in GRISM_INSTANCES.items():
    tier2_rho = tier2_results[instance]["theta"][INSTANCE_FREE_PARAMS.index("rho")]
    tier2_A = tier2_results[instance]["theta"][INSTANCE_FREE_PARAMS.index("A_deg")]
    for cfg in cfgs:
        rho_compare_rows.append({"dataset": cfg, "instance": instance,
                                  "stage1_rho": stage1_fits[cfg]["params"]["rho"],
                                  "tier2_rho_shared": tier2_rho, "tier2_A_deg_shared": tier2_A})
pd.DataFrame(rho_compare_rows).set_index("dataset").round(4)

# Residual histograms at Tier-1+2-forced parameters, all 6 datasets
fig, axes = plt.subplots(2, 6, figsize=(20, 6), sharex="row")
for col, cfg in enumerate(CONFIGS):
    pred = predict_centroids_full(dfs[cfg]["y_nisp"], dfs[cfg]["z_nisp"], dfs[cfg]["wavelength"],
                                   np.array([]), [], model, fixed=fixed_tier12[cfg])
    r_y = pred[0] - dfs[cfg]["cent_y"].values
    r_z = pred[1] - dfs[cfg]["cent_z"].values
    axes[0, col].hist(r_y, bins=40, color="tab:blue")
    axes[0, col].set_title(cfg)
    axes[1, col].hist(r_z, bins=40, color="tab:orange")
axes[0, 0].set_ylabel("r_y [mm]")
axes[1, 0].set_ylabel("r_z [mm]")
fig.suptitle("Residuals at Tier-1+2-forced parameters (coll_f/cam_f + A_deg/rho per instance jointly fit, Tier 3 at Stage-1 values)")
plt.tight_layout()
plt.show()

def write_tier2_toml(instance, cfgs, result):
    theta, std = result["theta"], result["std"]
    result_lines = "\n".join(f"{name} = {v:.6f}" for name, v in zip(INSTANCE_FREE_PARAMS, theta))
    std_lines = "\n".join(f"{name}_bootstrap_std = {s:.6f}" for name, s in zip(INSTANCE_FREE_PARAMS, std))
    rmse_lines = "\n".join(f"{cfg} = {rmse_tier12[cfg]:.6f}" for cfg in cfgs)

    text = f'''# Stage 4 Phase 3 -- Tier 2: per-grism-instance joint fit of A_deg/rho for
# the {instance} instance (3 datasets: {cfgs}). Generated by
# notebooks/4.3-Multi_Dataset_Fitting.ipynb. Do not hand-edit; re-run the
# notebook and regenerate this file instead.

[base_config]
path = "stage1_instrument.toml"

[tier1]
path = "joint_global_shared.toml"  # coll_f/cam_f, applied fixed here

[fit]
tier = "instance"
instance = "{instance}"
datasets = {cfgs}
method = "Nelder-Mead"
free_parameters = {INSTANCE_FREE_PARAMS}
cost_joint_before_mm2 = {result["cost_before"]:.6f}
cost_joint_after_mm2 = {result["cost_after"]:.6f}

[fit.result]
{result_lines}
{std_lines}

[fit.fixed]
# material_n0 fixed at nominal -- found ~95% bootstrap-correlated with
# A_deg even when jointly fit across 3 tilt geometries (product
# degeneracy in Prism.deviation = A*(n-1), not resolved by tilt diversity).
material_n0 = {FULL_NOMINAL["material_n0"]}
material_k = {base_config["material"]["k"]}

# Per-dataset RMSE at the current Tier-1+2-forced state (Tier 3 still at
# Stage-1 values, not yet re-optimized).
[fit.per_dataset_rmse_mm]
{rmse_lines}
'''
    path = MODELS_DIR / f"joint_instance_shared_{instance}.toml"
    path.write_text(text)
    print(f"wrote {path}")
    return path


for instance, cfgs in GRISM_INSTANCES.items():
    write_tier2_toml(instance, cfgs, tier2_results[instance])

mlflow.set_experiment("stage2")  # same deck-mandated experiment as Tier 1's global run

for instance, cfgs in GRISM_INSTANCES.items():
    result = tier2_results[instance]
    with mlflow.start_run(run_name=f"instance_shared_fit_{instance}"):
        mlflow.log_param("tier", "instance")
        mlflow.log_param("instance", instance)
        mlflow.log_param("datasets", cfgs)
        mlflow.log_param("free_parameters", INSTANCE_FREE_PARAMS)
        mlflow.log_param("material_n0_fixed_at", FULL_NOMINAL["material_n0"])
        for name, v, s in zip(INSTANCE_FREE_PARAMS, result["theta"], result["std"]):
            mlflow.log_param(name, v)
            mlflow.log_metric(f"{name}_bootstrap_std", s)
        mlflow.log_metric("cost_joint_before_mm2", result["cost_before"])
        mlflow.log_metric("cost_joint_after_mm2", result["cost_after"])
        for cfg in cfgs:
            mlflow.log_metric(f"rmse_{cfg}", rmse_tier12[cfg])

print("logged 2 runs ('instance_shared_fit_rgs000', 'instance_shared_fit_rgs180') to the 'stage2' MLflow experiment")

tier3_results = {}
for cfg in CONFIGS:
    theta0 = np.array([stage1_fits[cfg]["params"][k] for k in SPECIFIC_PARAMS])
    dfs_single = {cfg: dfs[cfg]}
    fixed_single = {cfg: fixed_tier12[cfg]}  # coll_f/cam_f (tier1), A_deg/rho (tier2 instance), material_n0 nominal

    cost_before = cost_joint(theta0, dfs_single, SPECIFIC_PARAMS, fixed_single, model)
    r = minimize(cost_joint, theta0, args=(dfs_single, SPECIFIC_PARAMS, fixed_single, model), method="Nelder-Mead",
                 options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-14, "adaptive": True})
    std, _ = bootstrap_uncertainty_joint(r.x, dfs_single, SPECIFIC_PARAMS, fixed_single, model)

    tier3_results[cfg] = {"theta": r.x, "std": std, "cost_before": cost_before, "cost_after": r.fun}
    print(f"{cfg}: RMSE {np.sqrt(cost_before):.4f} -> {np.sqrt(r.fun):.4f} mm  "
          f"theta*={dict(zip(SPECIFIC_PARAMS, r.x.round(4)))}")

# Tier-3 param table: Stage 1 vs. this section, per dataset, with bootstrap std
param_rows = []
for cfg in CONFIGS:
    s1 = stage1_fits[cfg]["params"]
    t3 = dict(zip(SPECIFIC_PARAMS, tier3_results[cfg]["theta"]))
    t3_std = dict(zip(SPECIFIC_PARAMS, tier3_results[cfg]["std"]))
    row = {"dataset": cfg}
    for k in SPECIFIC_PARAMS:
        row[f"stage1_{k}"] = s1[k]
        row[f"tier3_{k}"] = t3[k]
        row[f"tier3_{k}_std"] = t3_std[k]
    param_rows.append(row)
tier3_param_table = pd.DataFrame(param_rows).set_index("dataset")
tier3_param_table.round(4)

# RMSE across every stage of the alternating strategy, per dataset
rmse_tier3 = {cfg: float(np.sqrt(tier3_results[cfg]["cost_after"])) for cfg in CONFIGS}

rmse_all_stages = pd.DataFrame({
    "stage1_independent": {cfg: stage1_fits[cfg]["rmse_mm"] for cfg in CONFIGS},
    "tier1_forced": rmse_tier1,
    "tier1+2_forced": rmse_tier12,
    "tier3_joint_final": rmse_tier3,
}).loc[CONFIGS]
rmse_all_stages["delta_stage1_to_tier3"] = rmse_all_stages["tier3_joint_final"] - rmse_all_stages["stage1_independent"]
rmse_all_stages.round(4)

# Residual histograms at the final Tier-1+2+3 joint parameters, all 6 datasets
fig, axes = plt.subplots(2, 6, figsize=(20, 6), sharex="row")
for col, cfg in enumerate(CONFIGS):
    fixed_final = dict(fixed_tier12[cfg])
    fixed_final.update(dict(zip(SPECIFIC_PARAMS, tier3_results[cfg]["theta"])))
    pred = predict_centroids_full(dfs[cfg]["y_nisp"], dfs[cfg]["z_nisp"], dfs[cfg]["wavelength"],
                                   np.array([]), [], model, fixed=fixed_final)
    r_y = pred[0] - dfs[cfg]["cent_y"].values
    r_z = pred[1] - dfs[cfg]["cent_z"].values
    axes[0, col].hist(r_y, bins=40, color="tab:blue")
    axes[0, col].set_title(cfg)
    axes[1, col].hist(r_z, bins=40, color="tab:orange")
axes[0, 0].set_ylabel("r_y [mm]")
axes[1, 0].set_ylabel("r_z [mm]")
fig.suptitle("Residuals at the final joint fit (Tier 1 + Tier 2 + Tier 3), all 6 datasets")
plt.tight_layout()
plt.show()

fixed_tier1_iter2 = {}
for cfg in CONFIGS:
    p = dict(fixed_tier12[cfg])  # A_deg/rho (tier2 instance), material_n0 nominal
    p.update(dict(zip(SPECIFIC_PARAMS, tier3_results[cfg]["theta"])))  # tilt/offsets now Tier-3 refit
    fixed_tier1_iter2[cfg] = p

result_tier1_iter2 = minimize(cost_joint, theta_star_tier1, args=(dfs, GLOBAL_PARAMS, fixed_tier1_iter2, model),
                               method="Nelder-Mead",
                               options={"maxiter": 20000, "maxfev": 20000, "xatol": 1e-9, "fatol": 1e-13, "adaptive": True})
shift = result_tier1_iter2.x - theta_star_tier1
shift_in_std = shift / std_tier1

print(f"Tier 1 re-fit with Tier-2/3-updated baseline: {dict(zip(GLOBAL_PARAMS, result_tier1_iter2.x))}")
print(f"original Section-1 result:                     {dict(zip(GLOBAL_PARAMS, theta_star_tier1))}")
print(f"shift: {dict(zip(GLOBAL_PARAMS, shift))}")
print(f"shift in bootstrap-std units: {dict(zip(GLOBAL_PARAMS, shift_in_std.round(3)))}")
print(f"\n{'Converged: shift << 1 bootstrap std, one alternation pass is sufficient.' if np.all(np.abs(shift_in_std) < 0.5) else 'Shift is non-negligible relative to bootstrap std -- would warrant another alternation pass.'}")

def write_tier3_toml(cfg, result, fixed):
    theta, std = result["theta"], result["std"]
    result_lines = "\n".join(f"{name} = {v:.6f}" for name, v in zip(SPECIFIC_PARAMS, theta))
    std_lines = "\n".join(f"{name}_bootstrap_std = {s:.6f}" for name, s in zip(SPECIFIC_PARAMS, std))

    text = f'''# Stage 4 Phase 3 -- Tier 3: per-dataset re-fit of tilt_deg/offset_y_mm/
# offset_z_mm for {cfg}, with Tier 1 (coll_f/cam_f) and Tier 2 (A_deg/rho,
# this dataset's grism instance) held fixed. Generated by
# notebooks/4.3-Multi_Dataset_Fitting.ipynb. Do not hand-edit; re-run the
# notebook and regenerate this file instead.

[base_config]
path = "stage1_instrument.toml"

[tier1]
path = "joint_global_shared.toml"

[tier2]
path = "joint_instance_shared_{instance_of(cfg)}.toml"

[fit]
dataset = "{cfg}"
method = "Nelder-Mead"
free_parameters = {SPECIFIC_PARAMS}
cost_before_mm2 = {result["cost_before"]:.6f}
cost_after_mm2 = {result["cost_after"]:.6f}
rmse_mm = {np.sqrt(result["cost_after"]):.6f}

[fit.result]
{result_lines}
{std_lines}

[fit.fixed]
coll_f = {fixed["coll_f"]:.6f}
cam_f = {fixed["cam_f"]:.6f}
A_deg = {fixed["A_deg"]:.6f}
rho = {fixed["rho"]:.6f}
material_n0 = {fixed["material_n0"]:.6f}
material_k = {base_config["material"]["k"]}
'''
    path = MODELS_DIR / f"joint_specific_fit_{cfg}.toml"
    path.write_text(text)
    print(f"wrote {path}")
    return path


for cfg in CONFIGS:
    write_tier3_toml(cfg, tier3_results[cfg], fixed_tier12[cfg])

EXPERIMENT_NAME_TIER3 = "stage3"  # deck-mandated name (Task 3 step 3)
if mlflow.get_experiment_by_name(EXPERIMENT_NAME_TIER3) is None:
    mlflow.create_experiment(EXPERIMENT_NAME_TIER3, artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(EXPERIMENT_NAME_TIER3)

for cfg in CONFIGS:
    result = tier3_results[cfg]
    fixed = fixed_tier12[cfg]
    with mlflow.start_run(run_name=f"{cfg}_specific_refit"):
        mlflow.log_param("dataset", cfg)
        mlflow.log_param("instance", instance_of(cfg))
        mlflow.log_param("free_parameters", SPECIFIC_PARAMS)
        for k in ["coll_f", "cam_f", "A_deg", "rho", "material_n0"]:
            mlflow.log_param(k, fixed[k])
        for name, v, s in zip(SPECIFIC_PARAMS, result["theta"], result["std"]):
            mlflow.log_param(name, v)
            mlflow.log_metric(f"{name}_bootstrap_std", s)
        mlflow.log_metric("rmse_stage1", stage1_fits[cfg]["rmse_mm"])
        mlflow.log_metric("rmse_tier3", rmse_tier3[cfg])
        mlflow.log_metric("delta_rmse", rmse_tier3[cfg] - stage1_fits[cfg]["rmse_mm"])

print(f"logged {len(CONFIGS)} runs to the '{EXPERIMENT_NAME_TIER3}' MLflow experiment")

import ast

from sklearn.model_selection import GroupShuffleSplit
from sklearn.preprocessing import StandardScaler
from sklearn.neural_network import MLPRegressor
from sklearn.metrics import mean_squared_error

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)

def run_stage4_pipeline(cfg):
    df = dfs[cfg]
    fixed = fixed_final_by_cfg[cfg]
    pred = predict_centroids_full(df["y_nisp"], df["z_nisp"], df["wavelength"], np.array([]), [], model, fixed=fixed)
    d = df.assign(r_y=pred[0] - df["cent_y"].values, r_z=pred[1] - df["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 = np.sqrt(mean_squared_error(r_y_test, np.zeros_like(r_y_test)))
    physical_rmse_z = np.sqrt(mean_squared_error(r_z_test, np.zeros_like(r_z_test)))
    physical_dist_test = 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 = r_y_test - mlp_y.predict(X_test_s)
    hyb_z = r_z_test - mlp_z.predict(X_test_s)

    return {
        "cfg": cfg, "physical_rmse_y": physical_rmse_y, "physical_rmse_z": physical_rmse_z,
        "physical_dist_test": physical_dist_test,
        "hybrid_rmse_y": np.sqrt(mean_squared_error(r_y_test, mlp_y.predict(X_test_s))),
        "hybrid_rmse_z": np.sqrt(mean_squared_error(r_z_test, mlp_z.predict(X_test_s))),
        "hybrid_dist_test": np.mean(np.hypot(hyb_y, hyb_z)),
    }


stage4_results = {cfg: run_stage4_pipeline(cfg) for cfg in CONFIGS}
stage4_df = pd.DataFrame(stage4_results.values()).set_index("cfg").loc[CONFIGS]
stage4_df["improvement_pct"] = 100 * (1 - stage4_df["hybrid_dist_test"] / stage4_df["physical_dist_test"])
stage4_df.round(4)

EXPERIMENT_NAME_TIER4 = "stage4"  # deck-mandated name (Task 3 step 4)
if mlflow.get_experiment_by_name(EXPERIMENT_NAME_TIER4) is None:
    mlflow.create_experiment(EXPERIMENT_NAME_TIER4, artifact_location=f"file:{(REPO_ROOT / 'mlruns').resolve()}")
mlflow.set_experiment(EXPERIMENT_NAME_TIER4)

for cfg in CONFIGS:
    res = stage4_results[cfg]
    with mlflow.start_run(run_name=f"{cfg}_hybrid"):
        mlflow.log_param("dataset", cfg)
        mlflow.log_param("instance", instance_of(cfg))
        for k, v in fixed_final_by_cfg[cfg].items():
            mlflow.log_param(k, 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", res["physical_dist_test"])
        mlflow.log_metric("hybrid_dist_test", res["hybrid_dist_test"])
        mlflow.log_metric("hybrid_rmse_y", res["hybrid_rmse_y"])
        mlflow.log_metric("hybrid_rmse_z", res["hybrid_rmse_z"])
        mlflow.log_metric("improvement_pct", 100 * (1 - res["hybrid_dist_test"] / res["physical_dist_test"]))

print(f"logged {len(CONFIGS)} runs to the '{EXPERIMENT_NAME_TIER4}' MLflow experiment")

fig, ax = plt.subplots(figsize=(9, 4.5))
x = np.arange(len(CONFIGS))
width = 0.35
ax.bar(x - width/2, stage4_df["physical_dist_test"], width, label="physical-only (Tier 1+2+3 joint fit)")
ax.bar(x + width/2, stage4_df["hybrid_dist_test"], width, label="hybrid, MLP(y)+MLP(z)")
ax.set_xticks(x)
ax.set_xticklabels(CONFIGS, rotation=30, ha="right")
ax.set_ylabel("combined point distance [mm]")
ax.legend()
ax.set_title("Physical-only (final joint fit) vs. hybrid, per dataset")
plt.tight_layout()
plt.show()