dispcraft.zeroth_dispersion¶

def aggregate_by_position(df, value_cols=("dy", "dz"), decimals=3):
    """Collapse spectra sharing (nearly) the same field position to one row.

    Many spectra share the same `(y_nisp, z_nisp)` (repeated exposures at
    one field angle) -- Phase 3 Step 2 found averaging them first reduces
    noise (within-position std 2-2.3x smaller than between-position std)
    before training, same rationale as `dispcraft.field_calibration
    .aggregate_by_position` for the first-order residual, keyed on exact
    rounded position here rather than `spectra_id` (0th order has no
    repeated-wavelength rows to collapse first).
    """
    d = df.copy()
    d["_y_key"] = d["y_nisp"].round(decimals)
    d["_z_key"] = d["z_nisp"].round(decimals)
    agg = {"y_nisp": "mean", "z_nisp": "mean", **{c: "mean" for c in value_cols}}
    return d.groupby(["_y_key", "_z_key"]).agg(agg).reset_index(drop=True)

def fit_shared_material_k(dfs, fixed_by_cfg, model, k_init=0.004):
    """Fit one `material_k`, shared across every config in `dfs`, against
    the 0th-order two-blob *separation* (not absolute position, see
    `predict_zeroth_order_separation_physical`'s docstring).

    Parameters
    ----------
    dfs : dict, config -> DataFrame with `y_nisp`, `z_nisp`, `dy`, `dz`
        columns, one row per spectrum -- e.g.
        `dispcraft.measurement.zeroth_order_separation`'s output with
        `mad_outlier_mask` already applied
    fixed_by_cfg : dict, same keys as `dfs` -- each config's frozen
        physical fit (Tier 1+2+3), `material_k` excluded (it's the one
        free parameter here)
    model : dispcraft.calibration.GroundTestModel (any `m_order` -- swapped
        to 0 internally)
    k_init : float -- optimizer starting point, default the project's
        nominal `material_k` (`models/stage1_instrument.toml`)

    Returns
    -------
    float -- the fit, shared `material_k`
    """
    result = minimize(_separation_cost, [k_init], args=(dfs, fixed_by_cfg, model), method="Nelder-Mead")
    return float(result.x[0])

def fit_zeroth_dispersion_model(df, mlp_params=None):
    """Fit the field-dependent NN predicting the 0th-order two-blob
    separation left over after the shared-physics baseline (module
    docstring): Step 4 onwards, `df["dy"]`/`df["dz"]` should already have
    `predict_zeroth_order_physical`'s prediction subtracted (Steps 1-3
    fit this directly against the raw separation instead, since Step 4's
    physical baseline didn't exist yet).

    Parameters
    ----------
    df : DataFrame with `y_nisp`, `z_nisp`, `dy`, `dz` columns, one row per
        spectrum -- e.g. `dispcraft.measurement.zeroth_order_separation`'s
        output (or that minus the physical baseline), with
        `mad_outlier_mask` already applied.
    mlp_params : dict, default `DEFAULT_ZEROTH_MLP_PARAMS`

    Returns
    -------
    dispcraft.ml.LitResidualRegressor -- predicts (dy,dz) from
    `standardize_field(y0,z0)`.
    """
    mlp_params = mlp_params or DEFAULT_ZEROTH_MLP_PARAMS
    field = aggregate_by_position(df)
    X = standardize_field(field["y_nisp"].values, field["z_nisp"].values)
    y = field[["dy", "dz"]].values
    return train_residual_mlp(X, y, mlp_params, n_outputs=2)

def mad_outlier_mask(x, n_mad=5.0):
    """Boolean mask (True = keep), median-absolute-deviation clip.

    Stage 5 Phase 3 Step 1 found some 0th-order spectra have a single
    unreplicated, mismatched-crossmatch measurement that survives
    `measurement.zeroth_order_separation`'s median-collapsing (it only
    protects ranks with more than one exposure) -- producing `dz`/`dy`
    of 20-30mm next to a clean population with MAD*1.4826 of ~0.01mm. A
    wide `n_mad` clip isolates exactly these without touching the real
    signal.
    """
    x = np.asarray(x, dtype=float)
    mad = np.median(np.abs(x - np.median(x))) * 1.4826
    if mad == 0:
        return np.ones_like(x, dtype=bool)
    return np.abs(x - np.median(x)) <= n_mad * mad

def predict_zeroth_dispersion(y_nisp_mm, z_nisp_mm, nn_model):
    """(y0,z0) in mm -> (N,2) [dy,dz] predicted two-blob separation
    residual, mm (add `predict_zeroth_order_physical`'s output for the full
    Step-4-onwards prediction; the full raw separation itself for Steps 1-3's
    `nn_model`s, which were fit with no physical baseline)."""
    X = standardize_field(np.asarray(y_nisp_mm), np.asarray(z_nisp_mm))
    return ml_predict(nn_model, X)

def predict_zeroth_order_physical(y_nisp_mm, z_nisp_mm, wavelength_nm, model, params):
    """Shared-physics 0th-order prediction: `predict_centroids` at
    `m_order=0` (the prism's own chromatic dispersion, unaffected by `m`),
    using the same parameters already fit against first-order data.

    Parameters
    ----------
    y_nisp_mm, z_nisp_mm, wavelength_nm : array-like, same length
    model : dispcraft.calibration.GroundTestModel (any `m_order` -- swapped
        to 0 internally)
    params : dict -- a config's frozen physical fit (Tier 1+2+3), e.g.
        `models/joint_specific_fit_<cfg>.toml`'s merged `fit.fixed`/
        `fit.result`, optionally with `material_k` overridden (see
        `fit_shared_material_k`)

    Returns
    -------
    (2, N) array -- [cent_y, cent_z], in mm (same shape as `predict_centroids`)
    """
    model_m0 = dataclasses.replace(model, m_order=0)
    return predict_centroids(y_nisp_mm, z_nisp_mm, wavelength_nm, theta=[], free_names=[], model=model_m0, fixed=params)

def predict_zeroth_order_separation_physical(y_nisp_mm, z_nisp_mm, model, params):
    """Physics-predicted rank2-rank1 (red-blue) separation at the RGS
    passband's two edge wavelengths (`_RANK_WAVELENGTHS_NM`) -- the
    quantity `fit_shared_material_k` actually fits against.

    Deliberately *not* absolute position: `predict_zeroth_order_physical`'s
    absolute prediction (same offset_y_mm/offset_z_mm as first-order data)
    does not match the 0th order's actual absolute centroids -- checked
    while developing this step and found off by several mm, a much bigger,
    unrelated registration problem Stage 2 explicitly deferred. The
    *separation* between the two blobs of the same spectrum cancels that
    out and isolates the actual chromatic effect this step targets.

    Returns
    -------
    (N, 2) array -- [dy, dz], in mm
    """
    y_nisp_mm = np.asarray(y_nisp_mm)
    z_nisp_mm = np.asarray(z_nisp_mm)
    n = len(y_nisp_mm)
    pred1 = predict_zeroth_order_physical(y_nisp_mm, z_nisp_mm, np.full(n, _RANK_WAVELENGTHS_NM[0]), model, params)
    pred2 = predict_zeroth_order_physical(y_nisp_mm, z_nisp_mm, np.full(n, _RANK_WAVELENGTHS_NM[1]), model, params)
    return np.stack([pred2[0] - pred1[0], pred2[1] - pred1[1]], axis=1)