dispcraft.chebyshev_residual

chebyshev_residual

Chebyshev-polynomial residual model: physics (frozen joint fit) plus a Chebyshev-in-wavelength, Chebyshev-in-field-position residual correction.

Stage 4's production hybrid model (hybrid_joint_ml, docs/Status_Report .md Section 5) corrects the joint 3-tier physical fit's residual with a generic sklearn MLPRegressor on (y_nisp, z_nisp, wavelength). The reference paper (arXiv:2506.08378, docs/2506.08378v2-nisp_grism_trace_model .md) instead models its entire trace with an explicit functional form: Eq. (2) expands the trace as a Chebyshev polynomial in normalized wavelength lambda', whose own coefficients are themselves a 2D Chebyshev polynomial in normalized field position (y0', z0') (Eq. 3). This module applies that same nested functional form to this project's residual (not the raw trace, which the physical model already captures most of), as an interpretable, deterministic alternative to the MLP residual -- directly swappable with it (Stage 8, notebooks/8-Chebyshev_Residual.ipynb).

Because the model is linear in its coefficients (a Chebyshev basis times a Chebyshev basis), fitting is ordinary least squares (np.linalg.lstsq) -- no random seed, no early stopping, no architecture search, unlike dispcraft/ml.py/dispcraft/field_calibration.py. wave_order/ spatial_order play the role of a hyperparameter (a term count, (wave_order+1) * (spatial_order+1)**2 per axis), compared directly rather than searched.

The joint physical fit's tilt_deg/offset_y_mm/offset_z_mm are used exactly as already frozen (models/joint_specific_fit_<cfg>.toml) -- this model does not re-fit them (unlike field_calibration.py's alternating tilt_deg loop, which replaces the offsets with its own NN correction and is a different, non-default tier). This is a direct, drop-in alternative to the generic MLP residual step (notebooks/4.3-Multi_Dataset_Fitting.ipynb Section 4), so it matches that step's residual sign convention exactly, NOT field_calibration.py's: r_y/r_z = pred_phys - data (dispcraft.calibration.cost's own convention), the model is fit to predict r directly (not -r), and the corrected prediction is pred_phys - predicted_r (subtracted, not added -- the opposite sign from field_calibration.predict_centroids_field's pred + correction). Getting this backwards would double the residual instead of correcting it, the same class of bug Stage 5 Phase 2 Step 3 caught for the field-dependent NN tier -- guarded here by an end-to-end synthetic-recovery test (tests/test_chebyshev_residual.py).

ChebyshevResidualModel dataclass 

ChebyshevResidualModel(coef_y, coef_z, wave_order, spatial_order, lam_min_nm, lam_max_nm, field_scale_mm=FIELD_SCALE_MM)

Fitted Chebyshev residual model: one coefficient vector per axis.

predict_chebyshev_residual evaluates Phi @ coef_y/Phi @ coef_z -- predicting the residual r = pred_phys - data directly (module docstring's sign convention), not a correction to add.

chebyshev_design_matrix 

chebyshev_design_matrix(y0_mm, z0_mm, wavelength_nm, wave_order, spatial_order, lam_min_nm, lam_max_nm, field_scale_mm=FIELD_SCALE_MM)

Design matrix Phi with one column per (i, k, l) term:

Phi[:, i*(spatial_order+1)**2 + k*(spatial_order+1) + l]
    = T_i(lambda') * T_k(z0') * T_l(y0')

i = 0..wave_order (wavelength order, Eq. 2); k, l = 0..spatial_order (z0, y0 spatial order, Eq. 3). Mirrors the reference paper's own nested Chebyshev structure exactly: residual_kappa = sum_i C_kappa_i(y0,z0) * T_i(lambda'), C_kappa_i(y0,z0) = sum_kl a_kl * T_k(z0') * T_l(y0').

Returns:
  • (N, (wave_order+1)*(spatial_order+1)**2) array –
Source code in dispcraft/chebyshev_residual.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def chebyshev_design_matrix(y0_mm, z0_mm, wavelength_nm, wave_order, spatial_order,
                             lam_min_nm, lam_max_nm, field_scale_mm=FIELD_SCALE_MM):
    """Design matrix `Phi` with one column per `(i, k, l)` term:

        Phi[:, i*(spatial_order+1)**2 + k*(spatial_order+1) + l]
            = T_i(lambda') * T_k(z0') * T_l(y0')

    `i` = 0..wave_order (wavelength order, Eq. 2); `k, l` = 0..spatial_order
    (z0, y0 spatial order, Eq. 3). Mirrors the reference paper's own nested
    Chebyshev structure exactly: `residual_kappa = sum_i C_kappa_i(y0,z0) *
    T_i(lambda')`, `C_kappa_i(y0,z0) = sum_kl a_kl * T_k(z0') * T_l(y0')`.

    Returns
    -------
    (N, (wave_order+1)*(spatial_order+1)**2) array
    """
    y0_mm = np.asarray(y0_mm, dtype=float)
    z0_mm = np.asarray(z0_mm, dtype=float)
    lam_prime = standardize_wavelength(wavelength_nm, lam_min_nm, lam_max_nm)
    y0_prime = y0_mm / field_scale_mm
    z0_prime = z0_mm / field_scale_mm

    v_lam = chebvander(lam_prime, wave_order)  # (N, wave_order+1)
    v_z = chebvander(z0_prime, spatial_order)  # (N, spatial_order+1)
    v_y = chebvander(y0_prime, spatial_order)  # (N, spatial_order+1)

    phi = v_lam[:, :, None, None] * v_z[:, None, :, None] * v_y[:, None, None, :]
    n = phi.shape[0]
    return phi.reshape(n, -1)

fit_chebyshev_residual 

fit_chebyshev_residual(df, wave_order, spatial_order, lam_min_nm, lam_max_nm, field_scale_mm=FIELD_SCALE_MM)

Fit a ChebyshevResidualModel to df's residual columns via OLS.

Parameters:
  • df (DataFrame with `y_nisp`, `z_nisp`, `wavelength`, `r_y`, `r_z`) –

    columns -- r_y/r_z must already be the physical-fit residual, pred_phys - data (dispcraft.calibration.cost's convention, e.g. computed by the caller from predict_centroids with the frozen joint-fit fixed dict). Same convention Stage 4's generic MLP residual step is trained on -- this model is a drop-in alternative to it, not to field_calibration's differently-signed correction (see module docstring).

  • wave_order (int -- polynomial orders (see) –

    chebyshev_design_matrix)

  • spatial_order (int -- polynomial orders (see) –

    chebyshev_design_matrix)

  • lam_min_nm (float -- this grism's passband, for `lambda'`'s) –

    normalization (Eq. 2)

  • lam_max_nm (float -- this grism's passband, for `lambda'`'s) –

    normalization (Eq. 2)
Returns:
Source code in dispcraft/chebyshev_residual.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def fit_chebyshev_residual(df, wave_order, spatial_order, lam_min_nm, lam_max_nm,
                            field_scale_mm=FIELD_SCALE_MM):
    """Fit a `ChebyshevResidualModel` to `df`'s residual columns via OLS.

    Parameters
    ----------
    df : DataFrame with `y_nisp`, `z_nisp`, `wavelength`, `r_y`, `r_z`
        columns -- `r_y`/`r_z` must already be the physical-fit residual,
        `pred_phys - data` (`dispcraft.calibration.cost`'s convention, e.g.
        computed by the caller from `predict_centroids` with the frozen
        joint-fit `fixed` dict). Same convention Stage 4's generic MLP
        residual step is trained on -- this model is a drop-in alternative
        to it, not to `field_calibration`'s differently-signed correction
        (see module docstring).
    wave_order, spatial_order : int -- polynomial orders (see
        `chebyshev_design_matrix`)
    lam_min_nm, lam_max_nm : float -- this grism's passband, for `lambda'`'s
        normalization (Eq. 2)

    Returns
    -------
    ChebyshevResidualModel
    """
    phi = chebyshev_design_matrix(df["y_nisp"].values, df["z_nisp"].values, df["wavelength"].values,
                                   wave_order, spatial_order, lam_min_nm, lam_max_nm, field_scale_mm)
    coef_y, *_ = np.linalg.lstsq(phi, df["r_y"].values, rcond=None)
    coef_z, *_ = np.linalg.lstsq(phi, df["r_z"].values, rcond=None)
    return ChebyshevResidualModel(coef_y=coef_y, coef_z=coef_z, wave_order=wave_order, spatial_order=spatial_order,
                                   lam_min_nm=lam_min_nm, lam_max_nm=lam_max_nm, field_scale_mm=field_scale_mm)

n_chebyshev_terms 

n_chebyshev_terms(wave_order, spatial_order)

Number of coefficients per axis at a given (wave_order, spatial_order).

Source code in dispcraft/chebyshev_residual.py 
64
65
66
def n_chebyshev_terms(wave_order, spatial_order):
    """Number of coefficients per axis at a given (wave_order, spatial_order)."""
    return (wave_order + 1) * (spatial_order + 1) ** 2

predict_centroids_chebyshev 

predict_centroids_chebyshev(y_nisp_mm, z_nisp_mm, wavelength_nm, model, fixed, cheb_model)

Predict centroids from the frozen physical fit (fixed, e.g. a full Tier 1-3 models/joint_specific_fit_<cfg>.toml) minus the Chebyshev model's predicted residual: corrected_pred = pred_phys - r_pred, the same sign convention as Stage 4's generic MLP hybrid step (module docstring) -- the opposite sign from field_calibration .predict_centroids_field's pred + correction.

Returns:
  • (2, N) array -- [cent_y, cent_z], in mm (same shape as `predict_centroids`) –
Source code in dispcraft/chebyshev_residual.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def predict_centroids_chebyshev(y_nisp_mm, z_nisp_mm, wavelength_nm, model, fixed, cheb_model):
    """Predict centroids from the frozen physical fit (`fixed`, e.g. a full
    Tier 1-3 `models/joint_specific_fit_<cfg>.toml`) minus the Chebyshev
    model's predicted residual: `corrected_pred = pred_phys - r_pred`, the
    same sign convention as Stage 4's generic MLP hybrid step (module
    docstring) -- the opposite sign from `field_calibration
    .predict_centroids_field`'s `pred + correction`.

    Returns
    -------
    (2, N) array -- [cent_y, cent_z], in mm (same shape as `predict_centroids`)
    """
    pred = predict_centroids(y_nisp_mm, z_nisp_mm, wavelength_nm, theta=[], free_names=[], model=model, fixed=fixed)
    r_y_pred, r_z_pred = predict_chebyshev_residual(cheb_model, np.asarray(y_nisp_mm), np.asarray(z_nisp_mm),
                                                      wavelength_nm)
    return np.stack([pred[0] - r_y_pred, pred[1] - r_z_pred])

predict_chebyshev_residual 

predict_chebyshev_residual(cheb_model, y0_mm, z0_mm, wavelength_nm)

Predict the physical-fit residual (r_y, r_z), pred_phys - data convention -- NOT a correction to add (see module docstring).

Source code in dispcraft/chebyshev_residual.py 
149
150
151
152
153
154
def predict_chebyshev_residual(cheb_model, y0_mm, z0_mm, wavelength_nm):
    """Predict the physical-fit residual `(r_y, r_z)`, `pred_phys - data`
    convention -- NOT a correction to add (see module docstring)."""
    phi = chebyshev_design_matrix(y0_mm, z0_mm, wavelength_nm, cheb_model.wave_order, cheb_model.spatial_order,
                                   cheb_model.lam_min_nm, cheb_model.lam_max_nm, cheb_model.field_scale_mm)
    return phi @ cheb_model.coef_y, phi @ cheb_model.coef_z

standardize_wavelength 

standardize_wavelength(wavelength_nm, lam_min_nm, lam_max_nm)

wavelength [nm] -> lambda' in [-1,1], Eq. (2)'s own normalization: (lambda - 0.5(lmax+lmin)) / (0.5(lmax-lmin)).

Source code in dispcraft/chebyshev_residual.py 
55
56
57
58
59
60
61
def standardize_wavelength(wavelength_nm, lam_min_nm, lam_max_nm):
    """wavelength [nm] -> lambda' in [-1,1], Eq. (2)'s own normalization:
    `(lambda - 0.5(lmax+lmin)) / (0.5(lmax-lmin))`."""
    wavelength_nm = np.asarray(wavelength_nm, dtype=float)
    mid = 0.5 * (lam_max_nm + lam_min_nm)
    half_range = 0.5 * (lam_max_nm - lam_min_nm)
    return (wavelength_nm - mid) / half_range