dispcraft
  • Home

User Guide

  • Getting Started
  • Using the Calibrated Models
  • Notebooks
    • Overview
    • 1. Subject & Instrument Model
    • 2. Data Introduction
    • 3. ML Introduction
    • 4.1 ML Model Comparison
    • 4.2 Per-Dataset Pipeline
    • 4.3 Multi-Dataset Joint Fitting
    • 4.5 Comparison With Published Results
    • 5.1 PyTorch Migration
    • 5.2 Field-Dependent Parameters
    • 5.3 Zeroth-Order Dispersion
      • 1. Loading 0th-Order Data — a Real Bug in
      • 2. Outlier Investigation — Why Some Spectra Give a ~30 mm "Separation"
      • 3. Per-Config Baseline and Field-Position Dependence
      • 4. Findings
      • 5. Step 2 Setup — One Point per Field Position
      • 6. Architecture Grid — Input Encoding × Output Structure ()
      • 7. Sharing Granularity — Per-Dataset vs. Per-Instance Pooled Network
      • 8. Findings
      • 9. Step 3 — Fitting the Field-Dependent Model per Config
      • 10. Residuals After the Field-Dependent Fit
      • 11. Findings
      • 12. Step 4 — Joining 0th and 1st Order via a Shared
      • 13. Fitting a Shared
      • 14. Residual After the Shared-Physics Baseline
      • 15. Retraining the Field-Dependent NN on the Smaller Residual
      • 16. Findings
      • 17. Step 5 — Testing the Collimator/Camera Chromatic-Aberration Hypothesis Directly
      • 18. Fitting a Shared Chromatic Coefficient
      • 19. Residual After the Chromatic Term
      • 20. Isotropic Is Not Enough — Testing an Anisotropic Variant
      • 21. Findings
    • 5.4 BGS Model
    • 6. Status Report Assembly
    • 8. Chebyshev Residual Model
  • Status Report
  • Beginner Introduction (Slides)
  • Interactive Model (Webapp)

Reference

  • Euclid NISP Specs
  • Reference Paper Summary

API Reference

  • Overview
  • optics
  • measurement
  • calibration
  • field_calibration
  • zeroth_dispersion
  • chebyshev_residual
  • ml
  • model_registry
  • prediction
dispcraft
  • User Guide
  • Notebooks
  • 5.3 Zeroth-Order Dispersion

Stage 5, Phase 3 — 0th-Order Residual Dispersion (5.3-Zeroth_Order_Dispersion)¶

Per CLAUDE.md's Stage 5 Phase 3 plan. The physical model's m=0 grating term predicts zero wavelength dependence for the 0th order, but the actual ground-test 0th-order data is measurably wavelength-dependent: zeroth_order_centers (dispcraft/measurement.py) reduces each spectrum's 0th order to a two-blob model -- a rank-1 (blue, 1206 nm) blob and a rank-2 (red, 1892 nm) blob, the RGS passband's transmission edges (doi:10.1051/0004-6361/202555859) -- whose separation is real dispersion, not noise. This is produced by an optical effect the physical model's hypotheses exclude by construction (chromatic behaviour of the grism substrate/prism surviving the undiffracted beam, or chromatic aberration in collimator/camera -- root cause not identified yet).

Step 1 (this notebook, so far): residual-structure diagnostic. Before modeling anything, characterize the two-blob separation itself -- across configs and field position -- using *_zeroth.csv data that has never been fed into any physical fit (Stage 2's compare_orders was exploratory only).

In [1]:
Copied!
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression

from dispcraft.measurement import load_spectra

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

CONFIGS = ["rgs000_0", "rgs000_m4", "rgs000_p4", "rgs180_0", "rgs180_m4", "rgs180_p4", "bgs000_0"]

SIG_MAX = 0.05  # mm, same threshold as first-order (measurement.load_spectra default)
from pathlib import Path import numpy as np import pandas as pd import matplotlib.pyplot as plt from sklearn.linear_model import LinearRegression from dispcraft.measurement import load_spectra REPO_ROOT = Path("..") DATA_DIR = REPO_ROOT / "data" CONFIGS = ["rgs000_0", "rgs000_m4", "rgs000_p4", "rgs180_0", "rgs180_m4", "rgs180_p4", "bgs000_0"] SIG_MAX = 0.05 # mm, same threshold as first-order (measurement.load_spectra default)

1. Loading 0th-Order Data — a Real Bug in zeroth_order_centers¶

load_spectra's original quality filter required both sig_y < sig_max and sig_z < sig_max. In every *_zeroth.csv, sig_z (e_PSF_centroid[1]) is 100% NaN — it is never measured for the 0th order, only sig_y is. A NaN < sig_max comparison is always False, so the combined filter dropped every single row: zeroth_order_centers silently returned an empty DataFrame on all 7 real configs. This went unnoticed because its only caller so far was Stage 2's exploratory compare_orders (never fed into a fit), and its unit test used synthetic sig_z values that happened to be valid.

Fixed in dispcraft/measurement.py: load_spectra now takes a sig_columns parameter (defaults to both axes, for first-order data); zeroth_order_centers passes sig_columns=("sig_y",). Covered by two new tests in tests/test_measurement.py (all-NaN sig_z, and restricting sig_columns directly).

In [2]:
Copied!
def load_clean_zeroth(cfg):
    """Load one config's *_zeroth.csv, sig_y-filtered, one row per (spectra_id, rank).

    ~50% of spectra have 3-4 repeated exposures per rank (same physical
    blob measured multiple times) -- median-collapse them, same treatment
    `median_per_spectrum` gives first-order data, before computing the
    rank2-rank1 separation.
    """
    df = load_spectra(DATA_DIR / f"{cfg}_zeroth.csv", sig_max=SIG_MAX, sig_columns=("sig_y",))
    return df.groupby(["spectra_id", "rank"]).median(numeric_only=True).reset_index()

medians = {cfg: load_clean_zeroth(cfg) for cfg in CONFIGS}

for cfg, df in medians.items():
    n_spec = df["spectra_id"].nunique()
    n_both = (df.groupby("spectra_id")["rank"].nunique() == 2).sum()
    print(f"{cfg:12s} rows={len(df):5d}  spectra={n_spec:4d}  with_both_ranks={n_both:4d}")
def load_clean_zeroth(cfg): """Load one config's *_zeroth.csv, sig_y-filtered, one row per (spectra_id, rank). ~50% of spectra have 3-4 repeated exposures per rank (same physical blob measured multiple times) -- median-collapse them, same treatment `median_per_spectrum` gives first-order data, before computing the rank2-rank1 separation. """ df = load_spectra(DATA_DIR / f"{cfg}_zeroth.csv", sig_max=SIG_MAX, sig_columns=("sig_y",)) return df.groupby(["spectra_id", "rank"]).median(numeric_only=True).reset_index() medians = {cfg: load_clean_zeroth(cfg) for cfg in CONFIGS} for cfg, df in medians.items(): n_spec = df["spectra_id"].nunique() n_both = (df.groupby("spectra_id")["rank"].nunique() == 2).sum() print(f"{cfg:12s} rows={len(df):5d} spectra={n_spec:4d} with_both_ranks={n_both:4d}")
rgs000_0     rows=  687  spectra= 344  with_both_ranks= 343
rgs000_m4    rows=  596  spectra= 298  with_both_ranks= 298
rgs000_p4    rows=  576  spectra= 288  with_both_ranks= 288
rgs180_0     rows=  802  spectra= 402  with_both_ranks= 400
rgs180_m4    rows=  576  spectra= 288  with_both_ranks= 288
rgs180_p4    rows=  575  spectra= 288  with_both_ranks= 287
bgs000_0     rows=  603  spectra= 302  with_both_ranks= 301

2. Outlier Investigation — Why Some Spectra Give a ~30 mm "Separation"¶

Computing dz = cent_z(rank2) - cent_z(rank1) per spectrum, most configs have a tight, physically sensible distribution (sub-mm), but 5 of 7 show a handful of spectra with |dz| of 20-30 mm — an order of magnitude larger than the detector itself would make plausible for a chromatic effect. Median-collapsing repeated exposures (Section 1) does not remove these: they turn out to come from spectra where one rank has only a single, unreplicated measurement, and that lone measurement is itself a bad crossmatch (mismatched to a different, nearby PSF), e.g.:

spectra_id 433994101 (rgs000_0):
  rank=1: 4 rows, three at cent_z≈6.10, one at cent_z=36.71  -> median≈6.10 (protected)
  rank=2: 1 row  at cent_z=37.08                              -> no redundancy to protect it
  dz = 37.08 - 6.10 = 30.98 mm

This is a data-quality issue in the underlying nearest-match crossmatch (upstream of this project), not measurement noise — and not something a larger sig_max cut would catch, since the bad row's own sig_y is small (it's a confident measurement of the wrong source). A robust MAD-based clip on the separation, per config, isolates exactly these cases: the clean population has MAD*1.4826 of ~0.01-0.02 mm, so a >5*MAD cut is a wide margin against the true signal while catching >20 mm outliers cleanly.

In [3]:
Copied!
def separations(cfg):
    """One row per spectrum with both ranks present: dy, dz, and rank-1 field position."""
    df = medians[cfg]
    piv_y = df.pivot_table(index="spectra_id", columns="rank", values="cent_y")
    piv_z = df.pivot_table(index="spectra_id", columns="rank", values="cent_z")
    pos = df[df["rank"] == 1].set_index("spectra_id")[["y_nisp", "z_nisp"]]
    out = pd.DataFrame({
        "dy": piv_y[2] - piv_y[1],
        "dz": piv_z[2] - piv_z[1],
    }).dropna().join(pos)
    return out

def mad_clip(sub, n_mad=5.0):
    mad = (sub - sub.median()).abs().median() * 1.4826
    keep = (sub - sub.median()).abs() <= n_mad * mad
    return keep, mad

raw_sep = {cfg: separations(cfg) for cfg in CONFIGS}

print(f"{'config':12s} {'n':>5s} {'n_outlier':>10s} {'mad_dz':>8s} {'max|dz|_before':>15s} {'max|dz|_after':>14s}")
clean_sep = {}
for cfg in CONFIGS:
    sub = raw_sep[cfg]
    keep, mad = mad_clip(sub["dz"])
    clean_sep[cfg] = sub[keep]
    print(f"{cfg:12s} {len(sub):5d} {(~keep).sum():10d} {mad:8.4f} {sub['dz'].abs().max():15.3f} {clean_sep[cfg]['dz'].abs().max():14.3f}")
def separations(cfg): """One row per spectrum with both ranks present: dy, dz, and rank-1 field position.""" df = medians[cfg] piv_y = df.pivot_table(index="spectra_id", columns="rank", values="cent_y") piv_z = df.pivot_table(index="spectra_id", columns="rank", values="cent_z") pos = df[df["rank"] == 1].set_index("spectra_id")[["y_nisp", "z_nisp"]] out = pd.DataFrame({ "dy": piv_y[2] - piv_y[1], "dz": piv_z[2] - piv_z[1], }).dropna().join(pos) return out def mad_clip(sub, n_mad=5.0): mad = (sub - sub.median()).abs().median() * 1.4826 keep = (sub - sub.median()).abs() <= n_mad * mad return keep, mad raw_sep = {cfg: separations(cfg) for cfg in CONFIGS} print(f"{'config':12s} {'n':>5s} {'n_outlier':>10s} {'mad_dz':>8s} {'max|dz|_before':>15s} {'max|dz|_after':>14s}") clean_sep = {} for cfg in CONFIGS: sub = raw_sep[cfg] keep, mad = mad_clip(sub["dz"]) clean_sep[cfg] = sub[keep] print(f"{cfg:12s} {len(sub):5d} {(~keep).sum():10d} {mad:8.4f} {sub['dz'].abs().max():15.3f} {clean_sep[cfg]['dz'].abs().max():14.3f}")
config           n  n_outlier   mad_dz  max|dz|_before  max|dz|_after
rgs000_0       343          4   0.0167          30.983          0.244
rgs000_m4      298          1   0.0115           0.243          0.243
rgs000_p4      288         13   0.0132          15.237          0.252
rgs180_0       400          6   0.0126          30.707          0.241
rgs180_m4      288          0   0.0128           0.251          0.251
rgs180_p4      287          1   0.0103          29.973          0.237
bgs000_0       301         52   0.0059          24.755          0.148

3. Per-Config Baseline and Field-Position Dependence¶

With outliers removed, check two things per config/axis: the near-constant baseline separation (does it split by grism identity, as offset_z_mm did in Stage 4 Phase 2?), and how much of the remaining scatter a plain linear separation ~ y0 + z0 term explains — the same quantitative lower bound Stage 5 Phase 2 Step 1 used before considering any NN.

In [4]:
Copied!
rows = []
for cfg in CONFIGS:
    sub = clean_sep[cfg]
    X = sub[["y_nisp", "z_nisp"]].values
    for axis in ["dy", "dz"]:
        y = sub[axis].values
        lr = LinearRegression().fit(X, y)
        resid = y - lr.predict(X)
        r2 = lr.score(X, y)
        rows.append({
            "config": cfg,
            "axis": axis,
            "n": len(sub),
            "median_mm": float(np.median(y)),
            "rms_before_mm": float(np.std(y)),
            "rms_after_linear_mm": float(np.std(resid)),
            "r2": float(r2),
        })

summary = pd.DataFrame(rows)
pd.set_option("display.float_format", lambda v: f"{v:.4f}")
summary
rows = [] for cfg in CONFIGS: sub = clean_sep[cfg] X = sub[["y_nisp", "z_nisp"]].values for axis in ["dy", "dz"]: y = sub[axis].values lr = LinearRegression().fit(X, y) resid = y - lr.predict(X) r2 = lr.score(X, y) rows.append({ "config": cfg, "axis": axis, "n": len(sub), "median_mm": float(np.median(y)), "rms_before_mm": float(np.std(y)), "rms_after_linear_mm": float(np.std(resid)), "r2": float(r2), }) summary = pd.DataFrame(rows) pd.set_option("display.float_format", lambda v: f"{v:.4f}") summary
Out[4]:
config axis n median_mm rms_before_mm rms_after_linear_mm r2
0 rgs000_0 dy 339 0.0005 0.0025 0.0015 0.6458
1 rgs000_0 dz 339 0.2185 0.0135 0.0071 0.7236
2 rgs000_m4 dy 297 0.0162 0.0024 0.0016 0.5501
3 rgs000_m4 dz 297 0.2138 0.0115 0.0078 0.5363
4 rgs000_p4 dy 275 -0.0143 0.0032 0.0015 0.7719
5 rgs000_p4 dz 275 0.2280 0.0146 0.0079 0.7065
6 rgs180_0 dy 394 -0.0004 0.0038 0.0035 0.1502
7 rgs180_0 dz 394 -0.2177 0.0109 0.0072 0.5672
8 rgs180_m4 dy 288 -0.0155 0.0038 0.0031 0.3177
9 rgs180_m4 dz 288 -0.2240 0.0126 0.0068 0.7056
10 rgs180_p4 dy 286 0.0147 0.0021 0.0015 0.4498
11 rgs180_p4 dz 286 -0.2124 0.0095 0.0075 0.3727
12 bgs000_0 dy 249 0.0005 0.0028 0.0018 0.6132
13 bgs000_0 dz 249 0.1256 0.0101 0.0082 0.3351
In [5]:
Copied!
vmax_dz = max(clean_sep[cfg]["dz"].abs().max() for cfg in CONFIGS)

fig, axes = plt.subplots(len(CONFIGS), 1, figsize=(4.5, 3 * len(CONFIGS)), constrained_layout=True)
for ax, cfg in zip(axes, CONFIGS):
    sub = clean_sep[cfg]
    sc = ax.scatter(sub["y_nisp"], sub["z_nisp"], c=sub["dz"], cmap="RdBu_r", vmin=-vmax_dz, vmax=vmax_dz, s=12)
    ax.set_title(cfg)
    ax.set_xlabel("y_nisp [mm]")
    ax.set_ylabel("z_nisp [mm]")
    fig.colorbar(sc, ax=ax, label="dz = cent_z(rank2) - cent_z(rank1) [mm]")
plt.show()
vmax_dz = max(clean_sep[cfg]["dz"].abs().max() for cfg in CONFIGS) fig, axes = plt.subplots(len(CONFIGS), 1, figsize=(4.5, 3 * len(CONFIGS)), constrained_layout=True) for ax, cfg in zip(axes, CONFIGS): sub = clean_sep[cfg] sc = ax.scatter(sub["y_nisp"], sub["z_nisp"], c=sub["dz"], cmap="RdBu_r", vmin=-vmax_dz, vmax=vmax_dz, s=12) ax.set_title(cfg) ax.set_xlabel("y_nisp [mm]") ax.set_ylabel("z_nisp [mm]") fig.colorbar(sc, ax=ax, label="dz = cent_z(rank2) - cent_z(rank1) [mm]") plt.show()
No description has been provided for this image

4. Findings¶

The two-blob separation is real, structured, and dominated by z, not y. dz (rank2−rank1, i.e. red−blue) has a near-constant baseline 2-3 orders of magnitude larger than dy:

  • dz splits cleanly by grism identity, the same tilt+180°-mounting pattern Stage 4 Phase 2 traced offset_z_mm's sign to: rgs000_* ≈ +0.21 to +0.23 mm across all three tilts (0/m4/p4 agree to ~0.015 mm), rgs180_* ≈ -0.21 to -0.22 mm. bgs000_0 (a different device — broad-band grism, not part of the 6-config RGS set) sits at +0.126 mm, roughly half the RGS magnitude — consistent with it being a materially different optical path, not a data-quality artifact.
  • dy's baseline is near-zero for the two _0 (no-tilt) configs (~0.0005 mm) but ~0.014-0.016 mm for every ±4°-tilt config, with sign depending on both tilt direction and grism instance (rgs000_m4=+0.016, rgs000_p4=-0.014, rgs180_m4=-0.016, rgs180_p4=+0.015) — a tilt-grism interaction, not a pure grism-identity split like dz.
  • Field position explains a substantial share of what's left: a plain linear separation ~ y0 + z0 term gives R²=0.34-0.77 across all config/axis combinations (only rgs180_0's dy is weak, R²=0.15 — its baseline is already ~0, so there's less linear signal to explain to begin with) and cuts RMS 25-53%. This is stronger field dependence than Stage 5 Phase 2 Step 1 found for the first-order residual (R²=29-60%) — the 0th order's field structure is, if anything, easier to model, not harder.
  • A real, upstream data-quality issue was found and isolated, not papered over: 0-52 spectra per config (rgs180_m4: 0, bgs000_0: 52) have a >20mm "separation" caused by a single unreplicated, mismatched-crossmatch measurement — traced to a concrete example (Section 2), removed via a per-config 5×MAD clip (clean population MAD≈0.006-0.017mm, so the margin against the true ~0.2mm signal is wide). bgs000_0's unusually high 52/301 (17%) outlier rate is itself notable — this device may have a systematically worse crossmatch than the RGS configs, worth keeping in mind if bgs000_0 is carried into modeling.

Implication for a Phase 3 model: like Phase 2, the natural structure is tiered — a per-grism-instance (or per-config, given the tilt-dy interaction) near-constant baseline, plus a field-dependent term on top, with bgs000_0 likely needing separate treatment as a different device rather than folding into the 6-config RGS set. Nothing promoted to dispcraft/ yet — this step is diagnostic only, matching Phase 2 Step 1's own precedent. Next step (not started): pick and compare candidate model families for the field-dependent term, the same explore-before-commit approach Phase 2 Step 2 used.

5. Step 2 Setup — One Point per Field Position¶

Multiple spectra share (nearly) the same exact (y_nisp, z_nisp) — e.g. rgs000_0 has 339 clean spectra but only 167 distinct field positions. Group by exact position (rounded to 3 d.p., same convention dispcraft. field_calibration.aggregate_by_position uses) and average dy/dz within each — the same treatment Phase 2 Step 2 gave its residuals, and for the same reason: repeated measurements at one field position should be averaged down for noise, not treated as independent samples.

bgs000_0 is excluded from here on — Step 1 flagged it as a different device (broad-band grism) with a much higher outlier rate (17% vs 0-4.5% for the RGS configs), so it doesn't belong pooled or compared alongside the 6-config RGS set without separate treatment.

In [6]:
Copied!
CONFIGS_RGS = [c for c in CONFIGS if c != "bgs000_0"]


def to_field_df(cfg):
    """Collapse clean_sep[cfg] to one row per exact field position, averaging
    dy/dz -- plus the within/between-position std used to justify doing so."""
    d = clean_sep[cfg].copy()
    d["pos_key"] = list(zip(d["y_nisp"].round(3), d["z_nisp"].round(3)))
    within = d.groupby("pos_key")[["dy", "dz"]].std().mean()
    field = d.groupby("pos_key").agg(
        y_nisp=("y_nisp", "mean"), z_nisp=("z_nisp", "mean"),
        dy=("dy", "mean"), dz=("dz", "mean"),
    ).reset_index(drop=True)
    between = field[["dy", "dz"]].std()
    return field, within, between


field_dfs = {}
agg_rows = []
for cfg in CONFIGS_RGS:
    field, within, between = to_field_df(cfg)
    field_dfs[cfg] = field
    agg_rows.append({
        "config": cfg, "n_positions": len(field),
        "within_position_std_dy_mm": within["dy"], "between_position_std_dy_mm": between["dy"],
        "within_position_std_dz_mm": within["dz"], "between_position_std_dz_mm": between["dz"],
    })
agg_summary = pd.DataFrame(agg_rows)
agg_summary
CONFIGS_RGS = [c for c in CONFIGS if c != "bgs000_0"] def to_field_df(cfg): """Collapse clean_sep[cfg] to one row per exact field position, averaging dy/dz -- plus the within/between-position std used to justify doing so.""" d = clean_sep[cfg].copy() d["pos_key"] = list(zip(d["y_nisp"].round(3), d["z_nisp"].round(3))) within = d.groupby("pos_key")[["dy", "dz"]].std().mean() field = d.groupby("pos_key").agg( y_nisp=("y_nisp", "mean"), z_nisp=("z_nisp", "mean"), dy=("dy", "mean"), dz=("dz", "mean"), ).reset_index(drop=True) between = field[["dy", "dz"]].std() return field, within, between field_dfs = {} agg_rows = [] for cfg in CONFIGS_RGS: field, within, between = to_field_df(cfg) field_dfs[cfg] = field agg_rows.append({ "config": cfg, "n_positions": len(field), "within_position_std_dy_mm": within["dy"], "between_position_std_dy_mm": between["dy"], "within_position_std_dz_mm": within["dz"], "between_position_std_dz_mm": between["dz"], }) agg_summary = pd.DataFrame(agg_rows) agg_summary
Out[6]:
config n_positions within_position_std_dy_mm between_position_std_dy_mm within_position_std_dz_mm between_position_std_dz_mm
0 rgs000_0 167 0.0010 0.0022 0.0055 0.0128
1 rgs000_m4 146 0.0011 0.0021 0.0073 0.0100
2 rgs000_p4 139 0.0007 0.0032 0.0055 0.0145
3 rgs180_0 168 0.0011 0.0031 0.0058 0.0101
4 rgs180_m4 144 0.0012 0.0031 0.0058 0.0116
5 rgs180_p4 144 0.0011 0.0018 0.0073 0.0074

6. Architecture Grid — Input Encoding × Output Structure (rgs000_0)¶

Same protocol as Phase 2 Step 2: 5-fold KFold over field positions (each row already one position, no grouping needed), rgs000_0 as the representative single dataset, every candidate logged to a new zeroth_field_dependent_arch MLflow experiment. Two encoding × output-mode axes plus a hidden_layer_sizes/alpha grid:

  • n_bands: 0 = plain standardized (y0,z0); 4 = plus 4 Fourier bands per coordinate.
  • output_mode: independent (one MLP per axis) vs joint (one two-output MLP).

Baselines: baseline_constant (predicts the training fold's mean dy/dz — the "no field dependence, just Step 1's per-config baseline" assumption) and baseline_linear (Step 1's separation ~ y0 + z0 term). baseline_constant replaces Phase 2's baseline_zero, since here the target is the raw separation (with a real, large, non-zero baseline), not an already-baseline-subtracted residual.

In [7]:
Copied!
import logging
import warnings

import mlflow
from sklearn.model_selection import KFold

import dispcraft.ml as ml
from dispcraft.field_calibration import FIELD_SCALE_MM

logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
warnings.filterwarnings("ignore", message=".*does not have many workers.*")

RNG_SEED = 42


def encode_field(y0, z0, n_bands):
    u, v = y0 / FIELD_SCALE_MM, z0 / FIELD_SCALE_MM
    feats = [u[:, None], v[:, None]]
    for k in range(1, n_bands + 1):
        feats += [np.sin(k * np.pi * u)[:, None], np.cos(k * np.pi * u)[:, None],
                  np.sin(k * np.pi * v)[:, None], np.cos(k * np.pi * v)[:, None]]
    return np.concatenate(feats, axis=1)


def cv_rmse(X, y, fit_predict, kf):
    train_rmses, test_rmses = [], []
    for train_idx, test_idx in kf.split(X):
        pred_train, pred_test = fit_predict(X[train_idx], y[train_idx], X[train_idx], X[test_idx])
        train_rmses.append(ml.rmse(y[train_idx], pred_train))
        test_rmses.append(ml.rmse(y[test_idx], pred_test))
    return float(np.mean(train_rmses)), float(np.mean(test_rmses))


field0 = field_dfs["rgs000_0"]
y0_arr, z0_arr = field0["y_nisp"].values, field0["z_nisp"].values
targets0 = field0[["dy", "dz"]].values
kf = KFold(n_splits=5, shuffle=True, random_state=RNG_SEED)

mlflow.set_tracking_uri(f"sqlite:///{(REPO_ROOT / 'mlflow.db').resolve()}")
EXPERIMENT_NAME = "zeroth_field_dependent_arch"
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)

arch_rows = []


def _constant_fit_predict(X_train, y_train, X_train2, X_test):
    mean = y_train.mean(axis=0, keepdims=True)
    return np.repeat(mean, len(X_train2), axis=0), np.repeat(mean, len(X_test), axis=0)


def _linear_fit_predict(X_train, y_train, X_train2, X_test):
    preds_train, preds_test = [], []
    for j in range(y_train.shape[1]):
        lr = LinearRegression().fit(X_train[:, :2], y_train[:, j])
        preds_train.append(lr.predict(X_train2[:, :2]))
        preds_test.append(lr.predict(X_test[:, :2]))
    return np.stack(preds_train, axis=1), np.stack(preds_test, axis=1)


X0_raw = encode_field(y0_arr, z0_arr, 0)
for name, fit_predict in [("baseline_constant", _constant_fit_predict), ("baseline_linear", _linear_fit_predict)]:
    train_rmse, test_rmse = cv_rmse(X0_raw, targets0, fit_predict, kf)
    arch_rows.append({"n_bands": None, "hidden_layer_sizes": None, "alpha": None, "output_mode": name,
                       "cv_train_rmse_mm": train_rmse, "cv_test_rmse_mm": test_rmse})

# NN candidates
for n_bands in [0, 4]:
    X0 = encode_field(y0_arr, z0_arr, n_bands)
    for hidden in [(16,), (32, 16)]:
        for alpha in [1e-3, 1e-1]:
            mlp_params = {"hidden_layer_sizes": hidden, "activation": "relu", "alpha": alpha,
                          "random_state": RNG_SEED, "early_stopping": True}
            for output_mode in ["independent", "joint"]:
                if output_mode == "joint":
                    def fit_predict(X_train, y_train, X_train2, X_test, p=mlp_params):
                        m = ml.train_residual_mlp(X_train, y_train, p, n_outputs=2)
                        return ml.predict(m, X_train2), ml.predict(m, X_test)
                else:
                    def fit_predict(X_train, y_train, X_train2, X_test, p=mlp_params):
                        preds_train, preds_test = [], []
                        for j in range(2):
                            m = ml.train_residual_mlp(X_train, y_train[:, j:j + 1], p, n_outputs=1)
                            preds_train.append(ml.predict(m, X_train2)[:, 0])
                            preds_test.append(ml.predict(m, X_test)[:, 0])
                        return np.stack(preds_train, axis=1), np.stack(preds_test, axis=1)

                train_rmse, test_rmse = cv_rmse(X0, targets0, fit_predict, kf)
                with mlflow.start_run(run_name=f"bands{n_bands}_{hidden}_{alpha}_{output_mode}"):
                    mlflow.log_params({"n_bands": n_bands, "hidden_layer_sizes": str(hidden),
                                        "alpha": alpha, "output_mode": output_mode})
                    mlflow.log_metrics({"cv_train_rmse": train_rmse, "cv_test_rmse": test_rmse})
                arch_rows.append({"n_bands": n_bands, "hidden_layer_sizes": hidden, "alpha": alpha,
                                   "output_mode": output_mode, "cv_train_rmse_mm": train_rmse,
                                   "cv_test_rmse_mm": test_rmse})

arch_results = pd.DataFrame(arch_rows).sort_values("cv_test_rmse_mm").reset_index(drop=True)
arch_results
import logging import warnings import mlflow from sklearn.model_selection import KFold import dispcraft.ml as ml from dispcraft.field_calibration import FIELD_SCALE_MM logging.getLogger("pytorch_lightning").setLevel(logging.ERROR) warnings.filterwarnings("ignore", message=".*does not have many workers.*") RNG_SEED = 42 def encode_field(y0, z0, n_bands): u, v = y0 / FIELD_SCALE_MM, z0 / FIELD_SCALE_MM feats = [u[:, None], v[:, None]] for k in range(1, n_bands + 1): feats += [np.sin(k * np.pi * u)[:, None], np.cos(k * np.pi * u)[:, None], np.sin(k * np.pi * v)[:, None], np.cos(k * np.pi * v)[:, None]] return np.concatenate(feats, axis=1) def cv_rmse(X, y, fit_predict, kf): train_rmses, test_rmses = [], [] for train_idx, test_idx in kf.split(X): pred_train, pred_test = fit_predict(X[train_idx], y[train_idx], X[train_idx], X[test_idx]) train_rmses.append(ml.rmse(y[train_idx], pred_train)) test_rmses.append(ml.rmse(y[test_idx], pred_test)) return float(np.mean(train_rmses)), float(np.mean(test_rmses)) field0 = field_dfs["rgs000_0"] y0_arr, z0_arr = field0["y_nisp"].values, field0["z_nisp"].values targets0 = field0[["dy", "dz"]].values kf = KFold(n_splits=5, shuffle=True, random_state=RNG_SEED) mlflow.set_tracking_uri(f"sqlite:///{(REPO_ROOT / 'mlflow.db').resolve()}") EXPERIMENT_NAME = "zeroth_field_dependent_arch" 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) arch_rows = [] def _constant_fit_predict(X_train, y_train, X_train2, X_test): mean = y_train.mean(axis=0, keepdims=True) return np.repeat(mean, len(X_train2), axis=0), np.repeat(mean, len(X_test), axis=0) def _linear_fit_predict(X_train, y_train, X_train2, X_test): preds_train, preds_test = [], [] for j in range(y_train.shape[1]): lr = LinearRegression().fit(X_train[:, :2], y_train[:, j]) preds_train.append(lr.predict(X_train2[:, :2])) preds_test.append(lr.predict(X_test[:, :2])) return np.stack(preds_train, axis=1), np.stack(preds_test, axis=1) X0_raw = encode_field(y0_arr, z0_arr, 0) for name, fit_predict in [("baseline_constant", _constant_fit_predict), ("baseline_linear", _linear_fit_predict)]: train_rmse, test_rmse = cv_rmse(X0_raw, targets0, fit_predict, kf) arch_rows.append({"n_bands": None, "hidden_layer_sizes": None, "alpha": None, "output_mode": name, "cv_train_rmse_mm": train_rmse, "cv_test_rmse_mm": test_rmse}) # NN candidates for n_bands in [0, 4]: X0 = encode_field(y0_arr, z0_arr, n_bands) for hidden in [(16,), (32, 16)]: for alpha in [1e-3, 1e-1]: mlp_params = {"hidden_layer_sizes": hidden, "activation": "relu", "alpha": alpha, "random_state": RNG_SEED, "early_stopping": True} for output_mode in ["independent", "joint"]: if output_mode == "joint": def fit_predict(X_train, y_train, X_train2, X_test, p=mlp_params): m = ml.train_residual_mlp(X_train, y_train, p, n_outputs=2) return ml.predict(m, X_train2), ml.predict(m, X_test) else: def fit_predict(X_train, y_train, X_train2, X_test, p=mlp_params): preds_train, preds_test = [], [] for j in range(2): m = ml.train_residual_mlp(X_train, y_train[:, j:j + 1], p, n_outputs=1) preds_train.append(ml.predict(m, X_train2)[:, 0]) preds_test.append(ml.predict(m, X_test)[:, 0]) return np.stack(preds_train, axis=1), np.stack(preds_test, axis=1) train_rmse, test_rmse = cv_rmse(X0, targets0, fit_predict, kf) with mlflow.start_run(run_name=f"bands{n_bands}_{hidden}_{alpha}_{output_mode}"): mlflow.log_params({"n_bands": n_bands, "hidden_layer_sizes": str(hidden), "alpha": alpha, "output_mode": output_mode}) mlflow.log_metrics({"cv_train_rmse": train_rmse, "cv_test_rmse": test_rmse}) arch_rows.append({"n_bands": n_bands, "hidden_layer_sizes": hidden, "alpha": alpha, "output_mode": output_mode, "cv_train_rmse_mm": train_rmse, "cv_test_rmse_mm": test_rmse}) arch_results = pd.DataFrame(arch_rows).sort_values("cv_test_rmse_mm").reset_index(drop=True) arch_results
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
Out[7]:
n_bands hidden_layer_sizes alpha output_mode cv_train_rmse_mm cv_test_rmse_mm
0 4.0000 (32, 16) 0.1000 independent 0.0030 0.0034
1 0.0000 (32, 16) 0.0010 joint 0.0032 0.0037
2 0.0000 (16,) 0.1000 independent 0.0038 0.0038
3 0.0000 (32, 16) 0.0010 independent 0.0034 0.0038
4 4.0000 (16,) 0.1000 joint 0.0030 0.0039
5 NaN None NaN baseline_linear 0.0040 0.0040
6 0.0000 (16,) 0.0010 independent 0.0037 0.0042
7 0.0000 (32, 16) 0.1000 independent 0.0041 0.0043
8 4.0000 (32, 16) 0.1000 joint 0.0040 0.0046
9 0.0000 (32, 16) 0.1000 joint 0.0045 0.0046
10 0.0000 (16,) 0.1000 joint 0.0051 0.0055
11 0.0000 (16,) 0.0010 joint 0.0056 0.0063
12 NaN None NaN baseline_constant 0.0091 0.0091
13 4.0000 (32, 16) 0.0010 independent 0.0044 0.0129
14 4.0000 (16,) 0.0010 joint 0.0090 0.0196
15 4.0000 (16,) 0.1000 independent 0.0105 0.0211
16 4.0000 (32, 16) 0.0010 joint 0.0080 0.0221
17 4.0000 (16,) 0.0010 independent 0.0182 0.0402

7. Sharing Granularity — Per-Dataset vs. Per-Instance Pooled Network¶

Same question Phase 2 Step 2 asked: fit the field-dependent network per-dataset, or shared per grism instance (pooling {rgs000_0, rgs000_m4, rgs000_p4} / {rgs180_0, rgs180_m4, rgs180_p4})? Step 1 already found a concrete reason to expect per-dataset to win here even more clearly than in Phase 2: dy's baseline sign depends on tilt direction, not just grism identity (rgs000_m4=+0.016mm vs rgs000_p4=-0.014mm) — pooling across tilts would average away that real, tilt-specific signal, not just smooth out noise.

Tested directly: for each instrument, the position-level intersection across its 3 configs, 5-fold KFold over positions, comparing a per-dataset network (trained on that config's own training-fold positions) against a pooled network (trained on all 3 configs' training-fold positions/targets together) — both evaluated on each config's own held-out positions. Uses the winning plain-coordinate joint MLP from Section 6 (Fourier encoding's one nominal win there is flagged as fragile — see Section 8 — so it's not carried forward here).

In [8]:
Copied!
BEST_MLP_PARAMS = {"hidden_layer_sizes": (32, 16), "activation": "relu", "alpha": 1e-3,
                    "random_state": RNG_SEED, "early_stopping": True}


def train_predict_joint(X_train, y_train, X_test):
    m = ml.train_residual_mlp(X_train, y_train, BEST_MLP_PARAMS, n_outputs=2)
    return ml.predict(m, X_test)


def shared_positions(cfgs):
    grids = [set(zip(field_dfs[c]["y_nisp"].round(3), field_dfs[c]["z_nisp"].round(3))) for c in cfgs]
    return sorted(set.intersection(*grids))


def aligned_targets(cfg, shared):
    f = field_dfs[cfg].copy()
    f["key"] = list(zip(f["y_nisp"].round(3), f["z_nisp"].round(3)))
    return f.set_index("key").loc[list(shared), ["dy", "dz"]].values


INSTANCES = {"rgs000": ["rgs000_0", "rgs000_m4", "rgs000_p4"], "rgs180": ["rgs180_0", "rgs180_m4", "rgs180_p4"]}

sharing_rows = []
for instance, cfgs in INSTANCES.items():
    shared = shared_positions(cfgs)
    y0_s = np.array([p[0] for p in shared])
    z0_s = np.array([p[1] for p in shared])
    X_shared = encode_field(y0_s, z0_s, 0)
    targets = {cfg: aligned_targets(cfg, shared) for cfg in cfgs}
    print(f"{instance}: {len(shared)} shared positions across {cfgs}")

    kf_pos = KFold(n_splits=5, shuffle=True, random_state=RNG_SEED)
    for train_idx, test_idx in kf_pos.split(X_shared):
        X_train, X_test = X_shared[train_idx], X_shared[test_idx]

        X_train_pooled = np.concatenate([X_train] * len(cfgs), axis=0)
        y_train_pooled = np.concatenate([targets[cfg][train_idx] for cfg in cfgs], axis=0)
        pred_pooled = train_predict_joint(X_train_pooled, y_train_pooled, X_test)

        for cfg in cfgs:
            pred_own = train_predict_joint(X_train, targets[cfg][train_idx], X_test)
            sharing_rows.append({
                "instance": instance, "config": cfg,
                "rmse_per_dataset": ml.rmse(targets[cfg][test_idx], pred_own),
                "rmse_pooled": ml.rmse(targets[cfg][test_idx], pred_pooled),
            })

sharing_results = pd.DataFrame(sharing_rows)
sharing_summary = sharing_results.groupby("config")[["rmse_per_dataset", "rmse_pooled"]].mean()
sharing_summary["winner"] = np.where(
    sharing_summary["rmse_per_dataset"] < sharing_summary["rmse_pooled"], "per_dataset", "pooled")
sharing_summary
BEST_MLP_PARAMS = {"hidden_layer_sizes": (32, 16), "activation": "relu", "alpha": 1e-3, "random_state": RNG_SEED, "early_stopping": True} def train_predict_joint(X_train, y_train, X_test): m = ml.train_residual_mlp(X_train, y_train, BEST_MLP_PARAMS, n_outputs=2) return ml.predict(m, X_test) def shared_positions(cfgs): grids = [set(zip(field_dfs[c]["y_nisp"].round(3), field_dfs[c]["z_nisp"].round(3))) for c in cfgs] return sorted(set.intersection(*grids)) def aligned_targets(cfg, shared): f = field_dfs[cfg].copy() f["key"] = list(zip(f["y_nisp"].round(3), f["z_nisp"].round(3))) return f.set_index("key").loc[list(shared), ["dy", "dz"]].values INSTANCES = {"rgs000": ["rgs000_0", "rgs000_m4", "rgs000_p4"], "rgs180": ["rgs180_0", "rgs180_m4", "rgs180_p4"]} sharing_rows = [] for instance, cfgs in INSTANCES.items(): shared = shared_positions(cfgs) y0_s = np.array([p[0] for p in shared]) z0_s = np.array([p[1] for p in shared]) X_shared = encode_field(y0_s, z0_s, 0) targets = {cfg: aligned_targets(cfg, shared) for cfg in cfgs} print(f"{instance}: {len(shared)} shared positions across {cfgs}") kf_pos = KFold(n_splits=5, shuffle=True, random_state=RNG_SEED) for train_idx, test_idx in kf_pos.split(X_shared): X_train, X_test = X_shared[train_idx], X_shared[test_idx] X_train_pooled = np.concatenate([X_train] * len(cfgs), axis=0) y_train_pooled = np.concatenate([targets[cfg][train_idx] for cfg in cfgs], axis=0) pred_pooled = train_predict_joint(X_train_pooled, y_train_pooled, X_test) for cfg in cfgs: pred_own = train_predict_joint(X_train, targets[cfg][train_idx], X_test) sharing_rows.append({ "instance": instance, "config": cfg, "rmse_per_dataset": ml.rmse(targets[cfg][test_idx], pred_own), "rmse_pooled": ml.rmse(targets[cfg][test_idx], pred_pooled), }) sharing_results = pd.DataFrame(sharing_rows) sharing_summary = sharing_results.groupby("config")[["rmse_per_dataset", "rmse_pooled"]].mean() sharing_summary["winner"] = np.where( sharing_summary["rmse_per_dataset"] < sharing_summary["rmse_pooled"], "per_dataset", "pooled") sharing_summary
rgs000: 134 shared positions across ['rgs000_0', 'rgs000_m4', 'rgs000_p4']
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
rgs180: 143 shared positions across ['rgs180_0', 'rgs180_m4', 'rgs180_p4']
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
Out[8]:
rmse_per_dataset rmse_pooled winner
config
rgs000_0 0.0039 0.0053 per_dataset
rgs000_m4 0.0041 0.0131 per_dataset
rgs000_p4 0.0052 0.0127 per_dataset
rgs180_0 0.0035 0.0038 per_dataset
rgs180_m4 0.0046 0.0121 per_dataset
rgs180_p4 0.0033 0.0121 per_dataset

8. Findings¶

Aggregation (Section 5) helps, more modestly than Phase 2. Within- position std is 0.0007-0.0012 mm (dy) / 0.0055-0.0073 mm (dz); between-position std is 0.0018-0.0032 mm (dy) / 0.0074-0.0145 mm (dz) — a ~2-2.3x gap, not Phase 2's 20-30x. Makes sense: unlike Phase 2's residual (pure field-position signal on top of near-zero noise), this separation's "between" spread already contains its own baseline-driven variation, so the gap here specifically isolates the repeated-exposure noise reduction, which is real but smaller.

Architecture grid (Section 6). Every plain-coordinate (n_bands=0) NN candidate lands in a tight 0.0037-0.0063 mm test-RMSE band, beating baseline_constant (0.0091 mm, ~-59%) clearly but baseline_linear (0.0040 mm) only modestly (best plain-coordinate candidate: 0.0037 mm, ~6-14% better) — a much smaller NN-over-linear margin than Phase 2's ~75%, consistent with Step 1 already finding a strong linear component here (R²=0.34-0.77 vs. Phase 2's 29-60%): there's simply less non-linear structure left for an NN to earn its complexity on.

  • Output structure is a wash here, unlike Phase 2's decisive joint win: best plain-coordinate candidate is joint (hidden_layer_sizes= (32,16), alpha=1e-3, 0.00374 mm) but the best independent candidates are statistically tied with it (0.00377 mm, <1% apart) — closer to Stage 4's original wash than Phase 2's clear joint win. Plausible reason: Phase 2's r_y/r_z shared one corner feature with one likely physical cause; here dz's baseline splits by grism identity while dy's splits by tilt and grism identity — different physical drivers per axis, so less cross-axis structure for a joint model to exploit.
  • Fourier encoding is not a reliable win, despite one nominal top result. The single best row overall is n_bands=4, hidden=(32,16), alpha=0.1, independent (0.00344 mm) — but every other n_bands=4 configuration with lighter regularization (alpha=1e-3) catastrophically overfits (test RMSE 0.013-0.040 mm, 3-10x worse than any plain-coordinate candidate) on this ~134-position training set. A single win surrounded by that much instability reads as a fold-specific fluke, not a real architecture edge — plain coordinates are the robust choice, same conclusion Phase 2 reached, now for a second, independent residual target.
  • Winning architecture: standardized (y0,z0), joint MLP, hidden_layer_sizes=(32,16), alpha=1e-3 — identical hyperparameters to Phase 2's winner, arrived at independently on a different target (0th-order two-blob separation vs. first-order MLP residual).

Sharing granularity (Section 7): per-dataset wins decisively, 6/6 configs (vs. Phase 2's 5/6), pooled RMSE 1.4-3.7x worse — sharpest for the tilt configs (rgs000_m4/rgs000_p4/rgs180_m4/rgs180_p4: pooled 2.6-3.7x worse), consistent with Step 1's finding that dy's baseline sign flips with tilt direction, not just grism identity — pooling across the 3 tilts of one instrument averages away that real tilt-specific signal, more severely than Phase 2's parameter pooling did.

Nothing promoted to dispcraft/ yet — architecture settled (matching Phase 2 Step 2's own precedent, no promotion until the tier is actually fit). Next (not started): fit the field-dependent NN into the two-blob separation per config (Phase 2 Step 3's pattern), then decide how this 0th-order correction composes with the existing physical/hybrid model — open question, not yet resolved, since the 0th order isn't part of dispcraft.calibration's current predict_centroids at all.

9. Step 3 — Fitting the Field-Dependent Model per Config¶

Unlike Phase 2 Step 3, there's no existing physical parameter this competes with or replaces: the physical model's m=0 grating term already predicts exactly zero 0th-order separation, and nothing in dispcraft.calibration is fit against 0th-order data at all. So the NN here is a standalone model of the separation itself (baseline + field dependence together, both learned in one fit), not a residual against a physical prediction — no alternating fit or "replace, not add" degeneracy concern (Phase 2's tilt_deg/offset coupling) applies.

Promoted to dispcraft/, the two-phase workflow's real logic move:

  • dispcraft.measurement.zeroth_order_separation — the per-spectrum dy/dz computation from Sections 1-2 (median-collapsing repeats per rank first). Also fixed the same latent bug in zeroth_order_centers while touching this code: it averaged raw rows over spectra_id without median-collapsing per-rank repeats first, so it was just as exposed to Section 2's single-bad-exposure crossmatch outliers as the separation was — median-collapsing now happens before combining ranks in both functions. Covered by new tests in tests/test_measurement.py.
  • dispcraft.zeroth_dispersion (new module) — mad_outlier_mask, aggregate_by_position, fit_zeroth_dispersion_model, predict_zeroth_dispersion, composing measurement's data with dispcraft.ml's NN machinery, mirroring how field_calibration composes measurement with calibration. Tests in tests/test_zeroth_dispersion.py, including a synthetic-recovery check (known field-dependent + constant-baseline ground truth, recovered from generated data).

Fits one model per config (all 7, including bgs000_0 — Step 1/2's "needs separate treatment" means not pooled with the RGS instances, which a per-config fit already respects), using Step 2's winning architecture.

In [9]:
Copied!
from dispcraft.measurement import zeroth_order_separation
from dispcraft.zeroth_dispersion import (
    fit_zeroth_dispersion_model,
    mad_outlier_mask,
    predict_zeroth_dispersion,
)

tier_rows = []
final_sep = {}
nn_models = {}
for cfg in CONFIGS:
    sep = zeroth_order_separation(DATA_DIR / f"{cfg}_zeroth.csv")
    sep = sep[mad_outlier_mask(sep["dz"].values)].reset_index(drop=True)
    final_sep[cfg] = sep

    nn_model = fit_zeroth_dispersion_model(sep)
    nn_models[cfg] = nn_model
    pred = predict_zeroth_dispersion(sep["y_nisp"].values, sep["z_nisp"].values, nn_model)
    resid_dy = pred[:, 0] - sep["dy"].values
    resid_dz = pred[:, 1] - sep["dz"].values
    rmse_nn = float(np.sqrt(np.mean(resid_dy**2 + resid_dz**2)))

    baseline_pred = sep[["dy", "dz"]].mean().values
    rmse_baseline = float(np.sqrt(np.mean((sep["dy"].values - baseline_pred[0])**2
                                            + (sep["dz"].values - baseline_pred[1])**2)))

    nn_mean = pred.mean(axis=0)  # mean prediction over the training field positions

    tier_rows.append({
        "config": cfg, "n": len(sep),
        "rmse_baseline_constant_mm": rmse_baseline,
        "rmse_field_dependent_mm": rmse_nn,
        "rmse_reduction_pct": 100 * (1 - rmse_nn / rmse_baseline),
        "nn_mean_dy_mm": float(nn_mean[0]),
        "nn_mean_dz_mm": float(nn_mean[1]),
        "median_dy_mm": float(sep["dy"].median()),
        "median_dz_mm": float(sep["dz"].median()),
    })

tier_results = pd.DataFrame(tier_rows)
tier_results
from dispcraft.measurement import zeroth_order_separation from dispcraft.zeroth_dispersion import ( fit_zeroth_dispersion_model, mad_outlier_mask, predict_zeroth_dispersion, ) tier_rows = [] final_sep = {} nn_models = {} for cfg in CONFIGS: sep = zeroth_order_separation(DATA_DIR / f"{cfg}_zeroth.csv") sep = sep[mad_outlier_mask(sep["dz"].values)].reset_index(drop=True) final_sep[cfg] = sep nn_model = fit_zeroth_dispersion_model(sep) nn_models[cfg] = nn_model pred = predict_zeroth_dispersion(sep["y_nisp"].values, sep["z_nisp"].values, nn_model) resid_dy = pred[:, 0] - sep["dy"].values resid_dz = pred[:, 1] - sep["dz"].values rmse_nn = float(np.sqrt(np.mean(resid_dy**2 + resid_dz**2))) baseline_pred = sep[["dy", "dz"]].mean().values rmse_baseline = float(np.sqrt(np.mean((sep["dy"].values - baseline_pred[0])**2 + (sep["dz"].values - baseline_pred[1])**2))) nn_mean = pred.mean(axis=0) # mean prediction over the training field positions tier_rows.append({ "config": cfg, "n": len(sep), "rmse_baseline_constant_mm": rmse_baseline, "rmse_field_dependent_mm": rmse_nn, "rmse_reduction_pct": 100 * (1 - rmse_nn / rmse_baseline), "nn_mean_dy_mm": float(nn_mean[0]), "nn_mean_dz_mm": float(nn_mean[1]), "median_dy_mm": float(sep["dy"].median()), "median_dz_mm": float(sep["dz"].median()), }) tier_results = pd.DataFrame(tier_rows) tier_results
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
Out[9]:
config n rmse_baseline_constant_mm rmse_field_dependent_mm rmse_reduction_pct nn_mean_dy_mm nn_mean_dz_mm median_dy_mm median_dz_mm
0 rgs000_0 339 0.0137 0.0063 53.9163 0.0003 0.2170 0.0005 0.2185
1 rgs000_m4 297 0.0118 0.0070 40.5778 0.0161 0.2135 0.0162 0.2138
2 rgs000_p4 275 0.0149 0.0088 40.8299 -0.0144 0.2239 -0.0143 0.2280
3 rgs180_0 394 0.0116 0.0073 36.8632 -0.0001 -0.2161 -0.0004 -0.2177
4 rgs180_m4 288 0.0131 0.0073 44.3872 -0.0158 -0.2221 -0.0155 -0.2240
5 rgs180_p4 286 0.0097 0.0070 27.8132 0.0147 -0.2132 0.0147 -0.2124
6 bgs000_0 249 0.0105 0.0080 24.0171 0.0002 0.1224 0.0005 0.1256

10. Residuals After the Field-Dependent Fit¶

Same diagnostic as Section 4's plot, now for data - NN(y0,z0). If the field-dependent model genuinely absorbed the structure Steps 1-2 found, this should look like unstructured noise at a much smaller color scale.

In [10]:
Copied!
final_residuals = {}
for cfg in CONFIGS:
    sep = final_sep[cfg]
    pred = predict_zeroth_dispersion(sep["y_nisp"].values, sep["z_nisp"].values, nn_models[cfg])
    final_residuals[cfg] = pd.DataFrame({
        "y_nisp": sep["y_nisp"].values, "z_nisp": sep["z_nisp"].values,
        "r_dy": pred[:, 0] - sep["dy"].values, "r_dz": pred[:, 1] - sep["dz"].values,
    })

vmax_dz2 = max(final_residuals[cfg]["r_dz"].abs().max() for cfg in CONFIGS)

fig, axes = plt.subplots(len(CONFIGS), 1, figsize=(4.5, 3 * len(CONFIGS)), constrained_layout=True)
for ax, cfg in zip(axes, CONFIGS):
    r = final_residuals[cfg]
    sc = ax.scatter(r["y_nisp"], r["z_nisp"], c=r["r_dz"], cmap="RdBu_r", vmin=-vmax_dz2, vmax=vmax_dz2, s=12)
    ax.set_title(f"{cfg} (|max r_dz|={r['r_dz'].abs().max():.4f}, vs. Section 4's {vmax_dz:.3f})")
    ax.set_xlabel("y_nisp [mm]")
    ax.set_ylabel("z_nisp [mm]")
    fig.colorbar(sc, ax=ax, label="r_dz = NN(y0,z0) - dz [mm]")
plt.show()
final_residuals = {} for cfg in CONFIGS: sep = final_sep[cfg] pred = predict_zeroth_dispersion(sep["y_nisp"].values, sep["z_nisp"].values, nn_models[cfg]) final_residuals[cfg] = pd.DataFrame({ "y_nisp": sep["y_nisp"].values, "z_nisp": sep["z_nisp"].values, "r_dy": pred[:, 0] - sep["dy"].values, "r_dz": pred[:, 1] - sep["dz"].values, }) vmax_dz2 = max(final_residuals[cfg]["r_dz"].abs().max() for cfg in CONFIGS) fig, axes = plt.subplots(len(CONFIGS), 1, figsize=(4.5, 3 * len(CONFIGS)), constrained_layout=True) for ax, cfg in zip(axes, CONFIGS): r = final_residuals[cfg] sc = ax.scatter(r["y_nisp"], r["z_nisp"], c=r["r_dz"], cmap="RdBu_r", vmin=-vmax_dz2, vmax=vmax_dz2, s=12) ax.set_title(f"{cfg} (|max r_dz|={r['r_dz'].abs().max():.4f}, vs. Section 4's {vmax_dz:.3f})") ax.set_xlabel("y_nisp [mm]") ax.set_ylabel("z_nisp [mm]") fig.colorbar(sc, ax=ax, label="r_dz = NN(y0,z0) - dz [mm]") plt.show()
No description has been provided for this image

11. Findings¶

A clean, decisive win across all 7 configs, no bugs or surprises this time (unlike Phase 2 Step 3's sign bug) — likely because there's no alternating fit here to get a sign backwards in: RMSE dropped 24-54% relative to the per-config constant baseline (rgs000_0: 0.0137→ 0.0063mm; weakest reduction is bgs000_0 at 24%, consistent with it being the noisiest config throughout Steps 1-2).

The NN's mean prediction recovers the constant baseline almost exactly (e.g. rgs000_0: nn_mean_dz=0.2170 vs. median_dz=0.2185; rgs180_m4: nn_mean_dy=-0.0158 vs. median_dy=-0.0155) — confirming the single combined fit correctly learns "baseline + field shape" together without needing a separate scalar term, across every config including the sign-flipped rgs180_* ones and the tilt-dependent dy baselines.

Residuals after the fit (Section 10) lose the field-position structure Steps 1-2 found — no visible gradient or corner feature remains at a much smaller color scale than Section 4's pre-fit plot, for every config.

Nothing wired into the physical/hybrid model yet — this is a standalone predictive model of the 0th-order separation, not a correction inside predict_centroids. The physical model doesn't represent the 0th order at all currently (only dispcraft.measurement.zeroth_order_centers/ zeroth_order_separation extract it from data); composing this NN with the rest of dispcraft.calibration — e.g. using it to predict where the 0th-order center itself should land, or combining it with the first-order hybrid model in a shared training/reporting pipeline — is an open design question for a later step, not resolved here.

12. Step 4 — Joining 0th and 1st Order via a Shared material_k¶

Steps 1-3 built a standalone model of the two-blob separation, with no connection to the physical model already fit against first-order data. CLAUDE.md's Phase 3 plan names two candidate physical causes for the "unmodeled" separation: chromatic behavior of the prism material, and chromatic aberration in the collimator/camera. Neither had been tested.

dispcraft.calibration.GroundTestModel's m_order is a plain dataclass field, so evaluating the existing, unmodified predict_centroids at m_order=0 (dataclasses.replace) with each config's already-frozen first-order parameters (models/joint_specific_fit_<cfg>.toml) requires no new physics code. The Grism's deviation is prism.deviation(λ) - grating.deviation(λ); at m=0 the grating term is zero, but the prism's Cauchy chromatic dispersion is not — so the model already predicts a non-zero rank2-rank1 separation, from parameters fit purely on first-order data. Below: how much of the observed separation that alone explains.

In [11]:
Copied!
import dataclasses
import tomllib

from dispcraft.calibration import ground_test_model_from_config
from dispcraft.zeroth_dispersion import fit_shared_material_k, predict_zeroth_order_physical

MODELS_DIR = REPO_ROOT / "models"
CONFIGS_RGS = [c for c in CONFIGS if c != "bgs000_0"]  # bgs000_0 has no frozen first-order physical fit

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


def load_frozen_fit(cfg):
    with open(MODELS_DIR / f"joint_specific_fit_{cfg}.toml", "rb") as f:
        spec = tomllib.load(f)
    fit = spec["fit"]
    return {**fit["fixed"], **{k: fit["result"][k] for k in fit["free_parameters"]}}


frozen_params = {cfg: load_frozen_fit(cfg) for cfg in CONFIGS_RGS}

baseline_rows = []
for cfg in CONFIGS_RGS:
    sep = final_sep[cfg]
    params = frozen_params[cfg]
    pred1 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values,
                                           np.full(len(sep), 1206.0), model, params)
    pred2 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values,
                                           np.full(len(sep), 1892.0), model, params)
    dy_phys = (pred2[0] - pred1[0])[0]  # field-position-independent (prism deviation depends on wavelength only)
    dz_phys = (pred2[1] - pred1[1])[0]
    baseline_rows.append({
        "config": cfg, "dy_phys_mm": dy_phys, "dy_obs_median_mm": sep["dy"].median(),
        "dy_frac_pct": 100 * dy_phys / sep["dy"].median(),
        "dz_phys_mm": dz_phys, "dz_obs_median_mm": sep["dz"].median(),
        "dz_frac_pct": 100 * dz_phys / sep["dz"].median(),
    })

baseline_summary = pd.DataFrame(baseline_rows)
baseline_summary
import dataclasses import tomllib from dispcraft.calibration import ground_test_model_from_config from dispcraft.zeroth_dispersion import fit_shared_material_k, predict_zeroth_order_physical MODELS_DIR = REPO_ROOT / "models" CONFIGS_RGS = [c for c in CONFIGS if c != "bgs000_0"] # bgs000_0 has no frozen first-order physical fit with open(MODELS_DIR / "stage1_instrument.toml", "rb") as f: base_config = tomllib.load(f) model = ground_test_model_from_config(base_config) def load_frozen_fit(cfg): with open(MODELS_DIR / f"joint_specific_fit_{cfg}.toml", "rb") as f: spec = tomllib.load(f) fit = spec["fit"] return {**fit["fixed"], **{k: fit["result"][k] for k in fit["free_parameters"]}} frozen_params = {cfg: load_frozen_fit(cfg) for cfg in CONFIGS_RGS} baseline_rows = [] for cfg in CONFIGS_RGS: sep = final_sep[cfg] params = frozen_params[cfg] pred1 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values, np.full(len(sep), 1206.0), model, params) pred2 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values, np.full(len(sep), 1892.0), model, params) dy_phys = (pred2[0] - pred1[0])[0] # field-position-independent (prism deviation depends on wavelength only) dz_phys = (pred2[1] - pred1[1])[0] baseline_rows.append({ "config": cfg, "dy_phys_mm": dy_phys, "dy_obs_median_mm": sep["dy"].median(), "dy_frac_pct": 100 * dy_phys / sep["dy"].median(), "dz_phys_mm": dz_phys, "dz_obs_median_mm": sep["dz"].median(), "dz_frac_pct": 100 * dz_phys / sep["dz"].median(), }) baseline_summary = pd.DataFrame(baseline_rows) baseline_summary
Out[11]:
config dy_phys_mm dy_obs_median_mm dy_frac_pct dz_phys_mm dz_obs_median_mm dz_frac_pct
0 rgs000_0 0.0001 0.0005 29.4354 0.0611 0.2185 27.9653
1 rgs000_m4 0.0044 0.0162 27.3959 0.0609 0.2138 28.5045
2 rgs000_p4 -0.0040 -0.0143 28.0347 0.0610 0.2280 26.7398
3 rgs180_0 -0.0003 -0.0004 69.7881 -0.0611 -0.2177 28.0668
4 rgs180_m4 -0.0045 -0.0155 28.8527 -0.0609 -0.2240 27.2155
5 rgs180_p4 0.0040 0.0147 27.4947 -0.0610 -0.2124 28.7062

13. Fitting a Shared material_k¶

Stage 3 found material_k "not identifiable from centroid-position data alone" and fixed it at nominal (0.004) — but that used only first-order, single-narrow-wavelength-range data. 0th order's two widely-separated wavelengths (1206nm/1892nm) supply exactly the baseline first-order data lacks. Fit one shared material_k across all 6 RGS configs at once (dispcraft.zeroth_dispersion.fit_shared_material_k), holding every other parameter at its already-frozen first-order value, m_order=0. Also fit material_k independently per config as a consistency check (mirrors Stage 4 Phase 3's shared-vs-specific methodology) — if a shared material property is really being identified, independent per-config fits should land close to the joint one and to each other.

Fit against the separation, not absolute position. The first version of this fit reused dispcraft.calibration.cost_joint directly on absolute 0th-order centroids — a real bug, caught before trusting the result: it landed on a physically meaningless material_k≈-0.017 with several-mm RMSE (rgs000_0 alone: ~7mm), because the 0th order's absolute position doesn't match the shared-physics prediction at all (a much bigger, unrelated registration problem Stage 2 explicitly deferred and this step doesn't reopen). Fixed by fitting the separation instead (dispcraft.zeroth_dispersion.predict_zeroth_order_separation_physical), which cancels the absolute-registration mismatch out and isolates the actual chromatic effect — recovering the material_k≈0.0143 found while planning this step.

In [12]:
Copied!
fixed_by_cfg = {}
for cfg in CONFIGS_RGS:
    fixed_by_cfg[cfg] = {k: v for k, v in frozen_params[cfg].items() if k != "material_k"}

dfs_sep = {cfg: final_sep[cfg] for cfg in CONFIGS_RGS}
k_fit_joint = fit_shared_material_k(dfs_sep, fixed_by_cfg, model)
print(f"joint material_k = {k_fit_joint:.5f}  (nominal 0.004)")

fit_k_rows = []
for cfg in CONFIGS_RGS:
    k_fit_solo = fit_shared_material_k({cfg: dfs_sep[cfg]}, {cfg: fixed_by_cfg[cfg]}, model)

    sep = dfs_sep[cfg]
    rmse_by_k = {}
    for label, k in [("nominal", 0.004), ("joint_fit", k_fit_joint)]:
        params = {**fixed_by_cfg[cfg], "material_k": k}
        pred1 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values,
                                               np.full(len(sep), 1206.0), model, params)
        pred2 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values,
                                               np.full(len(sep), 1892.0), model, params)
        r_y = (pred2[0] - pred1[0]) - sep["dy"].values
        r_z = (pred2[1] - pred1[1]) - sep["dz"].values
        rmse_by_k[label] = float(np.sqrt(np.mean(r_y**2 + r_z**2)))

    fit_k_rows.append({
        "config": cfg, "material_k_per_config": k_fit_solo,
        "rmse_nominal_k_mm": rmse_by_k["nominal"], "rmse_joint_fit_k_mm": rmse_by_k["joint_fit"],
        "rmse_reduction_pct": 100 * (1 - rmse_by_k["joint_fit"] / rmse_by_k["nominal"]),
    })

fit_k_summary = pd.DataFrame(fit_k_rows)
fit_k_summary
fixed_by_cfg = {} for cfg in CONFIGS_RGS: fixed_by_cfg[cfg] = {k: v for k, v in frozen_params[cfg].items() if k != "material_k"} dfs_sep = {cfg: final_sep[cfg] for cfg in CONFIGS_RGS} k_fit_joint = fit_shared_material_k(dfs_sep, fixed_by_cfg, model) print(f"joint material_k = {k_fit_joint:.5f} (nominal 0.004)") fit_k_rows = [] for cfg in CONFIGS_RGS: k_fit_solo = fit_shared_material_k({cfg: dfs_sep[cfg]}, {cfg: fixed_by_cfg[cfg]}, model) sep = dfs_sep[cfg] rmse_by_k = {} for label, k in [("nominal", 0.004), ("joint_fit", k_fit_joint)]: params = {**fixed_by_cfg[cfg], "material_k": k} pred1 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values, np.full(len(sep), 1206.0), model, params) pred2 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values, np.full(len(sep), 1892.0), model, params) r_y = (pred2[0] - pred1[0]) - sep["dy"].values r_z = (pred2[1] - pred1[1]) - sep["dz"].values rmse_by_k[label] = float(np.sqrt(np.mean(r_y**2 + r_z**2))) fit_k_rows.append({ "config": cfg, "material_k_per_config": k_fit_solo, "rmse_nominal_k_mm": rmse_by_k["nominal"], "rmse_joint_fit_k_mm": rmse_by_k["joint_fit"], "rmse_reduction_pct": 100 * (1 - rmse_by_k["joint_fit"] / rmse_by_k["nominal"]), }) fit_k_summary = pd.DataFrame(fit_k_rows) fit_k_summary
joint material_k = 0.01430  (nominal 0.004)
Out[12]:
config material_k_per_config rmse_nominal_k_mm rmse_joint_fit_k_mm rmse_reduction_pct
0 rgs000_0 0.0142 0.1564 0.0138 91.1620
1 rgs000_m4 0.0140 0.1534 0.0126 91.7845
2 rgs000_p4 0.0147 0.1641 0.0161 90.1645
3 rgs180_0 0.0142 0.1559 0.0118 92.4617
4 rgs180_m4 0.0146 0.1620 0.0138 91.5112
5 rgs180_p4 0.0140 0.1526 0.0110 92.8007

14. Residual After the Shared-Physics Baseline¶

With material_k fit, recompute the residual (actual separation - physics prediction with the fit material_k) and rerun Section 3's plain linear residual ~ y0 + z0 diagnostic on it. If what's left is now relatively more field-position-dependent than before (Section 3's R²), that's evidence for the second candidate cause (chromatic aberration in the collimator/camera, which would vary across the field) rather than the prism material (which doesn't).

In [13]:
Copied!
residual_rows = []
residual_sep = {}
for cfg in CONFIGS_RGS:
    sep = final_sep[cfg].copy()
    params = {**fixed_by_cfg[cfg], "material_k": k_fit_joint}
    pred1 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values,
                                           np.full(len(sep), 1206.0), model, params)
    pred2 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values,
                                           np.full(len(sep), 1892.0), model, params)
    sep["dy_resid"] = sep["dy"].values - (pred2[0] - pred1[0])
    sep["dz_resid"] = sep["dz"].values - (pred2[1] - pred1[1])
    residual_sep[cfg] = sep

    X = sep[["y_nisp", "z_nisp"]].values
    for axis, resid_col in [("dy", "dy_resid"), ("dz", "dz_resid")]:
        y = sep[resid_col].values
        lr = LinearRegression().fit(X, y)
        resid_after = y - lr.predict(X)
        pre_fit_row = summary[(summary["config"] == cfg) & (summary["axis"] == axis)].iloc[0]
        residual_rows.append({
            "config": cfg, "axis": axis,
            "r2_pre_fit_k": pre_fit_row["r2"], "r2_post_fit_k": float(lr.score(X, y)),
            "rms_before_mm": float(np.sqrt(np.mean(y**2))),
            "rms_after_linear_mm": float(np.sqrt(np.mean(resid_after**2))),
        })

residual_summary = pd.DataFrame(residual_rows)
residual_summary
residual_rows = [] residual_sep = {} for cfg in CONFIGS_RGS: sep = final_sep[cfg].copy() params = {**fixed_by_cfg[cfg], "material_k": k_fit_joint} pred1 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values, np.full(len(sep), 1206.0), model, params) pred2 = predict_zeroth_order_physical(sep["y_nisp"].values, sep["z_nisp"].values, np.full(len(sep), 1892.0), model, params) sep["dy_resid"] = sep["dy"].values - (pred2[0] - pred1[0]) sep["dz_resid"] = sep["dz"].values - (pred2[1] - pred1[1]) residual_sep[cfg] = sep X = sep[["y_nisp", "z_nisp"]].values for axis, resid_col in [("dy", "dy_resid"), ("dz", "dz_resid")]: y = sep[resid_col].values lr = LinearRegression().fit(X, y) resid_after = y - lr.predict(X) pre_fit_row = summary[(summary["config"] == cfg) & (summary["axis"] == axis)].iloc[0] residual_rows.append({ "config": cfg, "axis": axis, "r2_pre_fit_k": pre_fit_row["r2"], "r2_post_fit_k": float(lr.score(X, y)), "rms_before_mm": float(np.sqrt(np.mean(y**2))), "rms_after_linear_mm": float(np.sqrt(np.mean(resid_after**2))), }) residual_summary = pd.DataFrame(residual_rows) residual_summary
Out[13]:
config axis r2_pre_fit_k r2_post_fit_k rms_before_mm rms_after_linear_mm
0 rgs000_0 dy 0.6458 0.6458 0.0025 0.0015
1 rgs000_0 dz 0.7236 0.7236 0.0136 0.0071
2 rgs000_m4 dy 0.5501 0.5501 0.0024 0.0016
3 rgs000_m4 dz 0.5363 0.5363 0.0124 0.0078
4 rgs000_p4 dy 0.7719 0.7719 0.0032 0.0015
5 rgs000_p4 dz 0.7065 0.7065 0.0158 0.0079
6 rgs180_0 dy 0.1502 0.1502 0.0038 0.0035
7 rgs180_0 dz 0.5672 0.5672 0.0111 0.0072
8 rgs180_m4 dy 0.3177 0.3177 0.0038 0.0031
9 rgs180_m4 dz 0.7056 0.7056 0.0132 0.0068
10 rgs180_p4 dy 0.4498 0.4498 0.0021 0.0015
11 rgs180_p4 dz 0.3727 0.3727 0.0108 0.0075

15. Retraining the Field-Dependent NN on the Smaller Residual¶

Retrain Step 3's field-dependent NN (dispcraft.zeroth_dispersion.fit_zeroth_dispersion_model, architecture unchanged) on this smaller residual instead of the raw separation, and compare final RMSE to Step 3's own per-config numbers (Section 9's tier_results).

In [14]:
Copied!
step4_rows = []
for cfg in CONFIGS_RGS:
    sep = residual_sep[cfg][["y_nisp", "z_nisp", "dy_resid", "dz_resid"]].rename(
        columns={"dy_resid": "dy", "dz_resid": "dz"})
    nn_model = fit_zeroth_dispersion_model(sep)
    pred = predict_zeroth_dispersion(sep["y_nisp"].values, sep["z_nisp"].values, nn_model)
    resid_dy = pred[:, 0] - sep["dy"].values
    resid_dz = pred[:, 1] - sep["dz"].values
    rmse_step4 = float(np.sqrt(np.mean(resid_dy**2 + resid_dz**2)))

    step3_row = tier_results[tier_results["config"] == cfg].iloc[0]
    step4_rows.append({
        "config": cfg,
        "rmse_step3_raw_target_mm": step3_row["rmse_field_dependent_mm"],
        "rmse_step4_physics_plus_nn_mm": rmse_step4,
        "change_pct": 100 * (rmse_step4 / step3_row["rmse_field_dependent_mm"] - 1),
    })

step4_summary = pd.DataFrame(step4_rows)
step4_summary
step4_rows = [] for cfg in CONFIGS_RGS: sep = residual_sep[cfg][["y_nisp", "z_nisp", "dy_resid", "dz_resid"]].rename( columns={"dy_resid": "dy", "dz_resid": "dz"}) nn_model = fit_zeroth_dispersion_model(sep) pred = predict_zeroth_dispersion(sep["y_nisp"].values, sep["z_nisp"].values, nn_model) resid_dy = pred[:, 0] - sep["dy"].values resid_dz = pred[:, 1] - sep["dz"].values rmse_step4 = float(np.sqrt(np.mean(resid_dy**2 + resid_dz**2))) step3_row = tier_results[tier_results["config"] == cfg].iloc[0] step4_rows.append({ "config": cfg, "rmse_step3_raw_target_mm": step3_row["rmse_field_dependent_mm"], "rmse_step4_physics_plus_nn_mm": rmse_step4, "change_pct": 100 * (rmse_step4 / step3_row["rmse_field_dependent_mm"] - 1), }) step4_summary = pd.DataFrame(step4_rows) step4_summary
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
/home/zoubian/Workspace/dispers/dispcraft/.pixi/envs/default/lib/python3.12/site-packages/pytorch_lightning/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
Out[14]:
config rmse_step3_raw_target_mm rmse_step4_physics_plus_nn_mm change_pct
0 rgs000_0 0.0063 0.0062 -1.4900
1 rgs000_m4 0.0070 0.0070 -0.2291
2 rgs000_p4 0.0088 0.0068 -23.2757
3 rgs180_0 0.0073 0.0073 0.3818
4 rgs180_m4 0.0073 0.0072 -1.5802
5 rgs180_p4 0.0070 0.0069 -2.2692

16. Findings¶

A single shared material_k explains the baseline, decisively. Fitting material_k jointly across all 6 RGS configs against the two-blob separation (not absolute position — see the caught-bug note in Section 13) lands at 0.0143 (vs. nominal 0.004), consistent to 3 significant figures across independent per-config fits (0.0140-0.0147) — as clean a confirmation as this project has produced that a physical parameter is identifiable and shared. Separation-prediction RMSE drops ~90-93% (0.15-0.16mm → 0.011-0.016mm) for every config. The nominal-material_k baseline (Section 12) explains a consistent ~27-29% of the observed separation across configs (rgs180_0's 70% dy figure is a division artifact — its dy baseline is ~0, so the ratio is meaningless there, not a real outlier).

The remaining field-dependence is untouched, and that's informative, not a gap. R² of the linear residual ~ y0+z0 diagnostic (Section 14) is identical before and after fitting material_k, to 4 decimal places, for every config/axis. This isn't a bug — it's mathematically guaranteed: material_k's effect is a per-config constant (Section 12 confirmed dy_phys/dz_phys don't vary with field position at all), and subtracting a constant shifts a linear regression's intercept without changing its R² (variance explained relative to the residual's own variance is invariant under a constant shift). So this cleanly separates the two candidate causes CLAUDE.md named: the prism material accounts for the field-independent baseline and nothing else; whatever field-dependence exists (R²=0.15-0.77, Phase 3 Step 1's finding) must come from a different, additive mechanism — consistent with the second candidate, collimator/camera chromatic aberration (a field-position-dependent effect by nature), though this step doesn't test that hypothesis directly.

Retraining the field-dependent NN on the smaller residual doesn't meaningfully change its own accuracy (Section 15): RMSE is within ±2% of Step 3's raw-target number for 5/6 configs, and better by 23% for rgs000_p4 (0.0088mm → 0.0068mm). This makes sense once R²'s invariance above is understood: the NN could already represent a per-config constant perfectly via its own bias term (Step 3 already confirmed this), so removing that constant via physics doesn't make the NN's job easier in absolute terms. The value Step 4 adds is interpretability, not accuracy: the ~27-29% baseline is now attributed to a real, shared, newly-identified material property (material_k, shared with the first-order model) instead of being folded invisibly into an opaque per-config learned constant — directly answering one of CLAUDE.md's two open candidate causes for the "unmodeled" 0th-order dispersion, with the other (collimator/camera chromatic aberration) now isolated as the remaining, not-yet-explained field-dependent share.

Promoted to dispcraft/: calibration.py's material_k is now genuinely fittable (nominal_params, no longer hardcoded to model.material.k); zeroth_dispersion.py gained predict_zeroth_order_physical (the shared-physics baseline, m_order=0), predict_zeroth_order_separation_physical (the two-blob separation this step actually fits against), and fit_shared_material_k. A real bug was caught and fixed before trusting the result (Section 13) — the module docstring documents it, and the promoted fit_shared_material_k structurally can't repeat it (its dfs are separation-shaped, no absolute cent_y/ cent_z columns to accidentally fit against). bgs000_0 still excluded (no frozen first-order fit) — stays on Step 3's standalone empirical model, an open gap. fit_zeroth_dispersion_model/predict_zeroth_dispersion (Steps 1-3) are unchanged code, now documented as expecting a baseline-subtracted target from Step 4 onward.

Still open: whether collimator/camera chromatic aberration is really the remaining mechanism is not tested here (would need independent evidence -- e.g. a wavelength-and-field-position model of the first-order residual that Phase 2 didn't consider, or optical modeling of lateral color); no models/*.toml recording of the fit material_k (this step establishes it computationally, not as a frozen, versioned result yet); and the "joined model" is still two separate function calls (predict_zeroth_order_physical + a residual NN), not a single unified prediction entry point -- left for whenever this composes with the rest of dispcraft.calibration.

17. Step 5 — Testing the Collimator/Camera Chromatic-Aberration Hypothesis Directly¶

Section 16 left one candidate cause untested: the prism material accounts for the field-independent baseline and nothing else, so whatever field-dependence remains (R²=0.15-0.77, Section 4) must come from a different, additive mechanism -- CLAUDE.md's second candidate, collimator/camera chromatic aberration.

Deriving a testable, not guessed, functional form. At m_order=0, Grism.forward reduces to: rotate the incoming angle, add the prism's (field-independent) deviation, rotate back. The only field-dependence in the whole chain comes from Collimator.forward (theta = -x/f_coll) and Camera.forward (x = -f_cam*theta), which are currently achromatic -- dispcraft.optics.elements.Collimator/Camera take one scalar f, no wavelength dependence. Composing the two (and using R.T@R=I) gives:

pos_cam(y0,z0,λ) = (cam_f(λ)/coll_f(λ)) * pos_foc(y0,z0)  -  cam_f(λ) * R.T@[0, prism.deviation(λ)]

The second term is field-independent (already Section 12-13's material_k baseline). The first term is a magnification ratio M(λ) = cam_f(λ)/coll_f(λ) multiplying the field position directly -- exactly the textbook signature of lateral color / transverse chromatic aberration: a field-position-proportional, wavelength-dependent image shift, with zero intercept (unlike the general y0+z0 linear fit Section 4 used, which allows one).

Stage 4 Phase 3 already found coll_f/cam_f individually non-identifiable (only their ratio matters, 99.8% bootstrap-correlated) -- so parametrizing two separate chromatic coefficients would just reintroduce that degeneracy. Instead: one shared chromaticity coefficient mu on the ratio, in the same Cauchy form already used for the prism material (n(λ) = n0 + k/λ²), applied to cam_f alone (equivalent to putting it on the ratio, since coll_f stays fixed):

cam_f_eff(λ) = cam_f * (1 + mu / λ_um²)

mu is fit shared across all 6 RGS configs (same bench collimator/camera, same reasoning Stage 4 Phase 3 gave coll_f/cam_f global Tier-1 status), against the Section 14 residual (dy_resid/dz_resid, already material_k-baseline-subtracted) -- diagnostic/notebook-local for now, no new dispcraft/ code, matching every other Step in this phase before its architecture was validated.

In [15]:
Copied!
from scipy.optimize import minimize as scipy_minimize

from dispcraft.optics.elements import Camera, Collimator, Grating, Grism, Prism
from dispcraft.optics.materials import Material
from dispcraft.zeroth_dispersion import _RANK_WAVELENGTHS_NM, predict_zeroth_order_separation_physical


def predict_zeroth_order_separation_chromatic(y0_mm, z0_mm, model, params, mu):
    """Physics-predicted rank2-rank1 separation with a chromatic
    collimator/camera magnification (`cam_f_eff(λ) = cam_f*(1+mu/λ_um²)`)
    added on top of `predict_zeroth_order_separation_physical`'s achromatic
    chain. `mu=0` reproduces it exactly (checked in the next cell)."""
    y0_mm = np.asarray(y0_mm, dtype=float)
    z0_mm = np.asarray(z0_mm, dtype=float)
    n = len(y0_mm)

    material = Material(n0=params["material_n0"], k=params["material_k"])
    prism = Prism(material=material, A=np.radians(params["A_deg"]))
    grating = Grating(m=0, rho=params["rho"])
    grism = Grism(prism, grating, tilt=np.radians(params["tilt_deg"]))
    coll = Collimator(f=params["coll_f"])
    pos_foc = np.stack([y0_mm, z0_mm]) / 1000.0  # mm -> m
    angle_col = coll.forward(pos_foc)

    def cam_pos(wavelength_nm):
        wavelength_um = np.full(n, wavelength_nm / 1000.0)
        angle_gr = grism.forward(angle_col, wavelength_um)
        cam_f_eff = params["cam_f"] * (1.0 + mu / wavelength_um**2)
        return -cam_f_eff * angle_gr  # (2, N), m -- Camera.forward with a per-row f

    pos1, pos2 = cam_pos(_RANK_WAVELENGTHS_NM[0]), cam_pos(_RANK_WAVELENGTHS_NM[1])
    return np.stack([(pos2[0] - pos1[0]) * 1000.0, (pos2[1] - pos1[1]) * 1000.0], axis=1)  # mm


# Sanity check: mu=0 must reproduce Section 12's achromatic prediction exactly.
_cfg_check = "rgs000_0"
_sep_check = final_sep[_cfg_check]
_params_check = {**fixed_by_cfg[_cfg_check], "material_k": k_fit_joint}
_chrom0 = predict_zeroth_order_separation_chromatic(
    _sep_check["y_nisp"].values, _sep_check["z_nisp"].values, model, _params_check, mu=0.0)
_achrom = predict_zeroth_order_separation_physical(
    _sep_check["y_nisp"].values, _sep_check["z_nisp"].values, model, _params_check)
print("max |mu=0 - achromatic| (mm):", np.max(np.abs(_chrom0 - _achrom)))
from scipy.optimize import minimize as scipy_minimize from dispcraft.optics.elements import Camera, Collimator, Grating, Grism, Prism from dispcraft.optics.materials import Material from dispcraft.zeroth_dispersion import _RANK_WAVELENGTHS_NM, predict_zeroth_order_separation_physical def predict_zeroth_order_separation_chromatic(y0_mm, z0_mm, model, params, mu): """Physics-predicted rank2-rank1 separation with a chromatic collimator/camera magnification (`cam_f_eff(λ) = cam_f*(1+mu/λ_um²)`) added on top of `predict_zeroth_order_separation_physical`'s achromatic chain. `mu=0` reproduces it exactly (checked in the next cell).""" y0_mm = np.asarray(y0_mm, dtype=float) z0_mm = np.asarray(z0_mm, dtype=float) n = len(y0_mm) material = Material(n0=params["material_n0"], k=params["material_k"]) prism = Prism(material=material, A=np.radians(params["A_deg"])) grating = Grating(m=0, rho=params["rho"]) grism = Grism(prism, grating, tilt=np.radians(params["tilt_deg"])) coll = Collimator(f=params["coll_f"]) pos_foc = np.stack([y0_mm, z0_mm]) / 1000.0 # mm -> m angle_col = coll.forward(pos_foc) def cam_pos(wavelength_nm): wavelength_um = np.full(n, wavelength_nm / 1000.0) angle_gr = grism.forward(angle_col, wavelength_um) cam_f_eff = params["cam_f"] * (1.0 + mu / wavelength_um**2) return -cam_f_eff * angle_gr # (2, N), m -- Camera.forward with a per-row f pos1, pos2 = cam_pos(_RANK_WAVELENGTHS_NM[0]), cam_pos(_RANK_WAVELENGTHS_NM[1]) return np.stack([(pos2[0] - pos1[0]) * 1000.0, (pos2[1] - pos1[1]) * 1000.0], axis=1) # mm # Sanity check: mu=0 must reproduce Section 12's achromatic prediction exactly. _cfg_check = "rgs000_0" _sep_check = final_sep[_cfg_check] _params_check = {**fixed_by_cfg[_cfg_check], "material_k": k_fit_joint} _chrom0 = predict_zeroth_order_separation_chromatic( _sep_check["y_nisp"].values, _sep_check["z_nisp"].values, model, _params_check, mu=0.0) _achrom = predict_zeroth_order_separation_physical( _sep_check["y_nisp"].values, _sep_check["z_nisp"].values, model, _params_check) print("max |mu=0 - achromatic| (mm):", np.max(np.abs(_chrom0 - _achrom)))
max |mu=0 - achromatic| (mm): 1.099120794378905e-14

18. Fitting a Shared Chromatic Coefficient mu¶

Two fits, mirroring Section 13's shared-vs-per-config check for material_k:

  1. mu alone, holding material_k at Section 13's joint value (k_fit_joint) -- fit against the Section 14 residual (dy_resid/dz_resid), the direct test of "does this specific functional form explain what's left."
  2. material_k and mu refit jointly from the raw separation -- an identifiability check: if the two are genuinely orthogonal (one a pure field-independent constant, the other a pure field-proportional term through the origin), material_k should barely move from Section 13's k_fit_joint.
In [16]:
Copied!
def _residual_chromatic_cost(mu, dfs_resid, fixed_by_cfg, model):
    mu = mu[0]
    total, n = 0.0, 0
    for cfg, df in dfs_resid.items():
        pred = predict_zeroth_order_separation_chromatic(
            df["y_nisp"].values, df["z_nisp"].values, model,
            {**fixed_by_cfg[cfg], "material_k": k_fit_joint}, mu)
        pred0 = predict_zeroth_order_separation_chromatic(
            df["y_nisp"].values, df["z_nisp"].values, model,
            {**fixed_by_cfg[cfg], "material_k": k_fit_joint}, 0.0)
        chrom_only = pred - pred0  # isolate the mu-only contribution (residual is already material_k-subtracted)
        r = chrom_only - df[["dy_resid", "dz_resid"]].values
        total += np.sum(r**2)
        n += len(df)
    return total / n


dfs_resid = {cfg: residual_sep[cfg] for cfg in CONFIGS_RGS}
mu_fit = scipy_minimize(_residual_chromatic_cost, [0.0], args=(dfs_resid, fixed_by_cfg, model),
                         method="Nelder-Mead").x[0]
print(f"joint mu = {mu_fit:.5f}  (material_k for scale: {k_fit_joint:.5f})")

mu_rows = []
for cfg in CONFIGS_RGS:
    mu_solo = scipy_minimize(_residual_chromatic_cost, [0.0], args=({cfg: dfs_resid[cfg]}, fixed_by_cfg, model),
                              method="Nelder-Mead").x[0]
    mu_rows.append({"config": cfg, "mu_per_config": mu_solo})
mu_summary = pd.DataFrame(mu_rows)
mu_summary


def _joint_k_mu_cost(theta, dfs_sep, fixed_by_cfg, model):
    k, mu = theta
    total, n = 0.0, 0
    for cfg, df in dfs_sep.items():
        params = {**fixed_by_cfg[cfg], "material_k": k}
        pred = predict_zeroth_order_separation_chromatic(df["y_nisp"].values, df["z_nisp"].values, model, params, mu)
        r = pred - df[["dy", "dz"]].values
        total += np.sum(r**2)
        n += len(df)
    return total / n


joint_fit = scipy_minimize(_joint_k_mu_cost, [k_fit_joint, mu_fit], args=(dfs_sep, fixed_by_cfg, model),
                            method="Nelder-Mead")
k_joint_refit, mu_joint_refit = joint_fit.x
print(f"joint refit: material_k={k_joint_refit:.5f} (vs. Section 13's {k_fit_joint:.5f}), "
      f"mu={mu_joint_refit:.5f} (vs. mu-alone fit {mu_fit:.5f})")
def _residual_chromatic_cost(mu, dfs_resid, fixed_by_cfg, model): mu = mu[0] total, n = 0.0, 0 for cfg, df in dfs_resid.items(): pred = predict_zeroth_order_separation_chromatic( df["y_nisp"].values, df["z_nisp"].values, model, {**fixed_by_cfg[cfg], "material_k": k_fit_joint}, mu) pred0 = predict_zeroth_order_separation_chromatic( df["y_nisp"].values, df["z_nisp"].values, model, {**fixed_by_cfg[cfg], "material_k": k_fit_joint}, 0.0) chrom_only = pred - pred0 # isolate the mu-only contribution (residual is already material_k-subtracted) r = chrom_only - df[["dy_resid", "dz_resid"]].values total += np.sum(r**2) n += len(df) return total / n dfs_resid = {cfg: residual_sep[cfg] for cfg in CONFIGS_RGS} mu_fit = scipy_minimize(_residual_chromatic_cost, [0.0], args=(dfs_resid, fixed_by_cfg, model), method="Nelder-Mead").x[0] print(f"joint mu = {mu_fit:.5f} (material_k for scale: {k_fit_joint:.5f})") mu_rows = [] for cfg in CONFIGS_RGS: mu_solo = scipy_minimize(_residual_chromatic_cost, [0.0], args=({cfg: dfs_resid[cfg]}, fixed_by_cfg, model), method="Nelder-Mead").x[0] mu_rows.append({"config": cfg, "mu_per_config": mu_solo}) mu_summary = pd.DataFrame(mu_rows) mu_summary def _joint_k_mu_cost(theta, dfs_sep, fixed_by_cfg, model): k, mu = theta total, n = 0.0, 0 for cfg, df in dfs_sep.items(): params = {**fixed_by_cfg[cfg], "material_k": k} pred = predict_zeroth_order_separation_chromatic(df["y_nisp"].values, df["z_nisp"].values, model, params, mu) r = pred - df[["dy", "dz"]].values total += np.sum(r**2) n += len(df) return total / n joint_fit = scipy_minimize(_joint_k_mu_cost, [k_fit_joint, mu_fit], args=(dfs_sep, fixed_by_cfg, model), method="Nelder-Mead") k_joint_refit, mu_joint_refit = joint_fit.x print(f"joint refit: material_k={k_joint_refit:.5f} (vs. Section 13's {k_fit_joint:.5f}), " f"mu={mu_joint_refit:.5f} (vs. mu-alone fit {mu_fit:.5f})")
joint mu = -0.00025  (material_k for scale: 0.01430)
joint refit: material_k=0.01439 (vs. Section 13's 0.01430), mu=-0.00026 (vs. mu-alone fit -0.00025)

19. Residual After the Chromatic Term¶

Same diagnostic as Section 14: recompute the residual after subtracting the fitted chromatic prediction, and rerun the linear residual ~ y0+z0 probe. Unlike material_k (a pure constant, mathematically guaranteed to leave R² unchanged, Section 16), the chromatic term is field-proportional -- if it's really the missing mechanism, R² should drop substantially toward 0.

In [17]:
Copied!
chromatic_rows = []
for cfg in CONFIGS_RGS:
    sep = residual_sep[cfg]
    params = {**fixed_by_cfg[cfg], "material_k": k_fit_joint}
    pred_chrom = predict_zeroth_order_separation_chromatic(
        sep["y_nisp"].values, sep["z_nisp"].values, model, params, mu_fit)
    pred0 = predict_zeroth_order_separation_chromatic(
        sep["y_nisp"].values, sep["z_nisp"].values, model, params, 0.0)
    chrom_only = pred_chrom - pred0

    X = sep[["y_nisp", "z_nisp"]].values
    for i, (axis, resid_col) in enumerate([("dy", "dy_resid"), ("dz", "dz_resid")]):
        y = sep[resid_col].values
        y_after_chrom = y - chrom_only[:, i]

        lr_before = LinearRegression().fit(X, y)
        lr_after = LinearRegression().fit(X, y_after_chrom)

        pre_row = residual_summary[(residual_summary["config"] == cfg) & (residual_summary["axis"] == axis)].iloc[0]
        chromatic_rows.append({
            "config": cfg, "axis": axis,
            "r2_before_chromatic": pre_row["r2_post_fit_k"],
            "r2_after_chromatic": float(lr_after.score(X, y_after_chrom)),
            "rms_before_mm": float(np.sqrt(np.mean(y**2))),
            "rms_after_chromatic_mm": float(np.sqrt(np.mean(y_after_chrom**2))),
            "rms_reduction_pct": 100 * (1 - np.sqrt(np.mean(y_after_chrom**2)) / np.sqrt(np.mean(y**2))),
        })

chromatic_summary = pd.DataFrame(chromatic_rows)
chromatic_summary
chromatic_rows = [] for cfg in CONFIGS_RGS: sep = residual_sep[cfg] params = {**fixed_by_cfg[cfg], "material_k": k_fit_joint} pred_chrom = predict_zeroth_order_separation_chromatic( sep["y_nisp"].values, sep["z_nisp"].values, model, params, mu_fit) pred0 = predict_zeroth_order_separation_chromatic( sep["y_nisp"].values, sep["z_nisp"].values, model, params, 0.0) chrom_only = pred_chrom - pred0 X = sep[["y_nisp", "z_nisp"]].values for i, (axis, resid_col) in enumerate([("dy", "dy_resid"), ("dz", "dz_resid")]): y = sep[resid_col].values y_after_chrom = y - chrom_only[:, i] lr_before = LinearRegression().fit(X, y) lr_after = LinearRegression().fit(X, y_after_chrom) pre_row = residual_summary[(residual_summary["config"] == cfg) & (residual_summary["axis"] == axis)].iloc[0] chromatic_rows.append({ "config": cfg, "axis": axis, "r2_before_chromatic": pre_row["r2_post_fit_k"], "r2_after_chromatic": float(lr_after.score(X, y_after_chrom)), "rms_before_mm": float(np.sqrt(np.mean(y**2))), "rms_after_chromatic_mm": float(np.sqrt(np.mean(y_after_chrom**2))), "rms_reduction_pct": 100 * (1 - np.sqrt(np.mean(y_after_chrom**2)) / np.sqrt(np.mean(y**2))), }) chromatic_summary = pd.DataFrame(chromatic_rows) chromatic_summary
Out[17]:
config axis r2_before_chromatic r2_after_chromatic rms_before_mm rms_after_chromatic_mm rms_reduction_pct
0 rgs000_0 dy 0.6458 0.7915 0.0025 0.0032 -30.1455
1 rgs000_0 dz 0.7236 0.4763 0.0136 0.0099 27.2356
2 rgs000_m4 dy 0.5501 0.8201 0.0024 0.0038 -57.2721
3 rgs000_m4 dz 0.5363 0.2075 0.0124 0.0092 25.9398
4 rgs000_p4 dy 0.7719 0.6191 0.0032 0.0025 22.3819
5 rgs000_p4 dz 0.7065 0.4766 0.0158 0.0133 15.8270
6 rgs180_0 dy 0.1502 0.4751 0.0038 0.0049 -27.5844
7 rgs180_0 dz 0.5672 0.2213 0.0111 0.0082 26.4996
8 rgs180_m4 dy 0.3177 0.3979 0.0038 0.0040 -6.3466
9 rgs180_m4 dz 0.7056 0.4260 0.0132 0.0108 18.6210
10 rgs180_p4 dy 0.4498 0.8515 0.0021 0.0040 -92.3369
11 rgs180_p4 dz 0.3727 0.1008 0.0108 0.0086 20.3008

20. Isotropic mu Is Not Enough — Testing an Anisotropic Variant¶

Section 19's single, shared, isotropic mu (one magnification-ratio chromaticity applied equally to y and z) gives a genuinely mixed result, not a clean confirmation: for dz, R² drops substantially (e.g. rgs000_0: 0.72→0.48) and RMS improves 16-27% in all 6 configs -- real explanatory power. But for dy, RMS gets worse in 5 of 6 configs (-6% to -92%) and R² mostly goes up -- the same scalar mu that helps z overcorrects y in the wrong shape entirely.

An isotropic magnification chromaticity is the simplest version of the hypothesis, not the only one: real lateral color can be anisotropic (chromatic difference of astigmatism, or the two axes' effective optical paths differing after the grism's tilt). Testing a version with independent mu_y, mu_z isolates whether the mechanism is "collimator/camera chromatic aberration, just not isotropic" vs. "not this mechanism at all for y."

In [18]:
Copied!
def predict_zeroth_order_separation_chromatic_aniso(y0_mm, z0_mm, model, params, mu_y, mu_z):
    """Like `predict_zeroth_order_separation_chromatic`, but with independent
    chromaticity coefficients per axis (`cam_f_eff` becomes a 2-vector,
    applied component-wise to `angle_gr`) -- tests an anisotropic chromatic
    aberration instead of assuming the isotropic `mu` form."""
    y0_mm = np.asarray(y0_mm, dtype=float)
    z0_mm = np.asarray(z0_mm, dtype=float)
    n = len(y0_mm)

    material = Material(n0=params["material_n0"], k=params["material_k"])
    prism = Prism(material=material, A=np.radians(params["A_deg"]))
    grating = Grating(m=0, rho=params["rho"])
    grism = Grism(prism, grating, tilt=np.radians(params["tilt_deg"]))
    coll = Collimator(f=params["coll_f"])
    pos_foc = np.stack([y0_mm, z0_mm]) / 1000.0
    angle_col = coll.forward(pos_foc)

    def cam_pos(wavelength_nm):
        wavelength_um = np.full(n, wavelength_nm / 1000.0)
        angle_gr = grism.forward(angle_col, wavelength_um)
        cam_f_eff = params["cam_f"] * np.stack([1.0 + mu_y / wavelength_um**2, 1.0 + mu_z / wavelength_um**2])
        return -cam_f_eff * angle_gr

    pos1, pos2 = cam_pos(_RANK_WAVELENGTHS_NM[0]), cam_pos(_RANK_WAVELENGTHS_NM[1])
    return np.stack([(pos2[0] - pos1[0]) * 1000.0, (pos2[1] - pos1[1]) * 1000.0], axis=1)


def _residual_chromatic_aniso_cost(theta, dfs_resid, fixed_by_cfg, model):
    mu_y, mu_z = theta
    total, n = 0.0, 0
    for cfg, df in dfs_resid.items():
        params = {**fixed_by_cfg[cfg], "material_k": k_fit_joint}
        pred = predict_zeroth_order_separation_chromatic_aniso(
            df["y_nisp"].values, df["z_nisp"].values, model, params, mu_y, mu_z)
        pred0 = predict_zeroth_order_separation_chromatic_aniso(
            df["y_nisp"].values, df["z_nisp"].values, model, params, 0.0, 0.0)
        chrom_only = pred - pred0
        r = chrom_only - df[["dy_resid", "dz_resid"]].values
        total += np.sum(r**2)
        n += len(df)
    return total / n


aniso_fit = scipy_minimize(_residual_chromatic_aniso_cost, [mu_fit, mu_fit], args=(dfs_resid, fixed_by_cfg, model),
                            method="Nelder-Mead")
mu_y_fit, mu_z_fit = aniso_fit.x
print(f"joint mu_y={mu_y_fit:.5f}, mu_z={mu_z_fit:.5f}  (isotropic mu was {mu_fit:.5f})")

aniso_rows = []
for cfg in CONFIGS_RGS:
    sep = residual_sep[cfg]
    params = {**fixed_by_cfg[cfg], "material_k": k_fit_joint}
    pred_aniso = predict_zeroth_order_separation_chromatic_aniso(
        sep["y_nisp"].values, sep["z_nisp"].values, model, params, mu_y_fit, mu_z_fit)
    pred0 = predict_zeroth_order_separation_chromatic_aniso(
        sep["y_nisp"].values, sep["z_nisp"].values, model, params, 0.0, 0.0)
    chrom_only = pred_aniso - pred0

    X = sep[["y_nisp", "z_nisp"]].values
    for i, (axis, resid_col) in enumerate([("dy", "dy_resid"), ("dz", "dz_resid")]):
        y = sep[resid_col].values
        y_after = y - chrom_only[:, i]
        lr_after = LinearRegression().fit(X, y_after)
        iso_row = chromatic_summary[(chromatic_summary["config"] == cfg) & (chromatic_summary["axis"] == axis)].iloc[0]
        aniso_rows.append({
            "config": cfg, "axis": axis,
            "r2_isotropic": iso_row["r2_after_chromatic"],
            "r2_anisotropic": float(lr_after.score(X, y_after)),
            "rms_reduction_pct_isotropic": iso_row["rms_reduction_pct"],
            "rms_reduction_pct_anisotropic": 100 * (1 - np.sqrt(np.mean(y_after**2)) / np.sqrt(np.mean(y**2))),
        })

aniso_summary = pd.DataFrame(aniso_rows)
aniso_summary
def predict_zeroth_order_separation_chromatic_aniso(y0_mm, z0_mm, model, params, mu_y, mu_z): """Like `predict_zeroth_order_separation_chromatic`, but with independent chromaticity coefficients per axis (`cam_f_eff` becomes a 2-vector, applied component-wise to `angle_gr`) -- tests an anisotropic chromatic aberration instead of assuming the isotropic `mu` form.""" y0_mm = np.asarray(y0_mm, dtype=float) z0_mm = np.asarray(z0_mm, dtype=float) n = len(y0_mm) material = Material(n0=params["material_n0"], k=params["material_k"]) prism = Prism(material=material, A=np.radians(params["A_deg"])) grating = Grating(m=0, rho=params["rho"]) grism = Grism(prism, grating, tilt=np.radians(params["tilt_deg"])) coll = Collimator(f=params["coll_f"]) pos_foc = np.stack([y0_mm, z0_mm]) / 1000.0 angle_col = coll.forward(pos_foc) def cam_pos(wavelength_nm): wavelength_um = np.full(n, wavelength_nm / 1000.0) angle_gr = grism.forward(angle_col, wavelength_um) cam_f_eff = params["cam_f"] * np.stack([1.0 + mu_y / wavelength_um**2, 1.0 + mu_z / wavelength_um**2]) return -cam_f_eff * angle_gr pos1, pos2 = cam_pos(_RANK_WAVELENGTHS_NM[0]), cam_pos(_RANK_WAVELENGTHS_NM[1]) return np.stack([(pos2[0] - pos1[0]) * 1000.0, (pos2[1] - pos1[1]) * 1000.0], axis=1) def _residual_chromatic_aniso_cost(theta, dfs_resid, fixed_by_cfg, model): mu_y, mu_z = theta total, n = 0.0, 0 for cfg, df in dfs_resid.items(): params = {**fixed_by_cfg[cfg], "material_k": k_fit_joint} pred = predict_zeroth_order_separation_chromatic_aniso( df["y_nisp"].values, df["z_nisp"].values, model, params, mu_y, mu_z) pred0 = predict_zeroth_order_separation_chromatic_aniso( df["y_nisp"].values, df["z_nisp"].values, model, params, 0.0, 0.0) chrom_only = pred - pred0 r = chrom_only - df[["dy_resid", "dz_resid"]].values total += np.sum(r**2) n += len(df) return total / n aniso_fit = scipy_minimize(_residual_chromatic_aniso_cost, [mu_fit, mu_fit], args=(dfs_resid, fixed_by_cfg, model), method="Nelder-Mead") mu_y_fit, mu_z_fit = aniso_fit.x print(f"joint mu_y={mu_y_fit:.5f}, mu_z={mu_z_fit:.5f} (isotropic mu was {mu_fit:.5f})") aniso_rows = [] for cfg in CONFIGS_RGS: sep = residual_sep[cfg] params = {**fixed_by_cfg[cfg], "material_k": k_fit_joint} pred_aniso = predict_zeroth_order_separation_chromatic_aniso( sep["y_nisp"].values, sep["z_nisp"].values, model, params, mu_y_fit, mu_z_fit) pred0 = predict_zeroth_order_separation_chromatic_aniso( sep["y_nisp"].values, sep["z_nisp"].values, model, params, 0.0, 0.0) chrom_only = pred_aniso - pred0 X = sep[["y_nisp", "z_nisp"]].values for i, (axis, resid_col) in enumerate([("dy", "dy_resid"), ("dz", "dz_resid")]): y = sep[resid_col].values y_after = y - chrom_only[:, i] lr_after = LinearRegression().fit(X, y_after) iso_row = chromatic_summary[(chromatic_summary["config"] == cfg) & (chromatic_summary["axis"] == axis)].iloc[0] aniso_rows.append({ "config": cfg, "axis": axis, "r2_isotropic": iso_row["r2_after_chromatic"], "r2_anisotropic": float(lr_after.score(X, y_after)), "rms_reduction_pct_isotropic": iso_row["rms_reduction_pct"], "rms_reduction_pct_anisotropic": 100 * (1 - np.sqrt(np.mean(y_after**2)) / np.sqrt(np.mean(y**2))), }) aniso_summary = pd.DataFrame(aniso_rows) aniso_summary
joint mu_y=-0.00025, mu_z=-0.00026  (isotropic mu was -0.00025)
Out[18]:
config axis r2_isotropic r2_anisotropic rms_reduction_pct_isotropic rms_reduction_pct_anisotropic
0 rgs000_0 dy 0.7915 0.7915 -30.1455 -30.1455
1 rgs000_0 dz 0.4763 0.4594 27.2356 28.2133
2 rgs000_m4 dy 0.8201 0.8201 -57.2721 -57.2721
3 rgs000_m4 dz 0.2075 0.1921 25.9398 26.8141
4 rgs000_p4 dy 0.6191 0.6191 22.3819 22.3819
5 rgs000_p4 dz 0.4766 0.4609 15.8270 16.3832
6 rgs180_0 dy 0.4751 0.4751 -27.5844 -27.5844
7 rgs180_0 dz 0.2213 0.2075 26.4996 27.0942
8 rgs180_m4 dy 0.3979 0.3979 -6.3466 -6.3466
9 rgs180_m4 dz 0.4260 0.4070 18.6210 19.1653
10 rgs180_p4 dy 0.8515 0.8515 -92.3369 -92.3369
11 rgs180_p4 dz 0.1008 0.0990 20.3008 20.6852

21. Findings¶

The collimator/camera chromatic-aberration hypothesis is confirmed for z, not for y -- a mixed, axis-asymmetric result, not a clean yes/no. The functional form was derived, not guessed: dispcraft.optics.elements .Collimator/Camera are currently achromatic (one scalar f each, no λ dependence); giving their ratio a chromatic magnification M(λ)=cam_f(λ)/coll_f(λ)=M0*(1+mu/λ_um²) (mirroring the Cauchy form already used for the prism material) predicts a field-proportional, zero-intercept separation term from the existing, unmodified Grism.forward -- the textbook lateral-color signature, and structurally orthogonal to material_k's field-independent baseline by construction (confirmed: jointly refitting material_k and mu together moves material_k by only 0.6%, from 0.01430 to 0.01439, Section 18).

dz: real, if partial, explanatory power, consistent across all 6 configs. A single shared mu=-0.00025 (57x smaller than material_k's 0.0143, for scale) fit against the Section 14 residual cuts dz's RMS 15-28% and drops the linear residual~y0+z0 R² substantially in every RGS config -- rgs000_0: 0.72→0.48; rgs180_p4: 0.37→0.10; rgs000_m4: 0.54→0.21 (Section 19). dz is the axis that carries almost all of the observed 0th-order baseline (Section 4) and dominates the remaining field structure Phase 3 Step 1 found -- exactly the part of the puzzle this hypothesis most needed to explain, and it does, partially.

dy: ruled out, not just unconfirmed. The same isotropic mu makes dy's RMS worse in 5 of 6 configs (-6% to -92%, Section 19). Allowing independent mu_y/mu_z (Section 20, an anisotropic variant, in case real lateral color differs by axis) doesn't rescue it: the joint optimizer converges to mu_y=-0.00025, mu_z=-0.00026 -- indistinguishable from the isotropic value -- and dy's R²/RMS come out identical to 4 decimal places between the isotropic and anisotropic fits for every config. That's a stronger conclusion than "wrong coefficient": the optimizer had a free parameter available specifically to fix dy and found no value of it that helps at all, meaning dy's remaining field structure isn't just mis-scaled by this mechanism, it has the wrong shape for it (not proportional to field position through the origin the way this hypothesis predicts).

Interpretation, honestly bounded. CLAUDE.md's open question -- "is collimator/camera chromatic aberration the remaining mechanism?" -- now has a concrete, tested answer: partially, for z only. dy's remaining field dependence (R²=0.15-0.85 after this step, Section 19-20) still has no identified physical cause; candidates not tested here include a chromatic effect that doesn't reduce to a single magnification ratio (e.g. acting on coll_f/cam_f individually rather than only their ratio -- deliberately not tried, since Stage 4 Phase 3 found them jointly non-identifiable that way), a higher-order or non-proportional field dependence, or a mechanism unrelated to chromatic aberration entirely.

Nothing promoted to dispcraft/. This step is diagnostic/exploratory only, matching every architecture-comparison step in this phase before a model was validated well enough to commit -- and here the result itself (partial, axis-split) doesn't clear that bar. No models/*.toml recording. predict_zeroth_order_separation_chromatic/_aniso stay notebook-local.

Previous Next

Built with MkDocs using a theme provided by Read the Docs.
dispers/dispcraft « Previous Next »