dispcraft.calibration¶

def bootstrap_uncertainty_joint(theta_star, dfs, free_names, fixed_by_cfg, model: GroundTestModel,
                                 n_boot=30, seed=None):
    """Bootstrap std + correlation matrix for a `cost_joint` fit.

    Resamples each dataset in `dfs` independently (own size, with
    replacement), recombines into `cost_joint`, and re-minimizes from
    `theta_star` for each replicate. Generic over how many datasets `dfs`
    holds -- a single-entry `dfs` reduces this to single-dataset bootstrap
    uncertainty.

    Returns
    -------
    std  : (len(free_names),) array
    corr : DataFrame, index/columns = free_names
    """
    rng = np.random.default_rng(seed)
    boots = []
    for _ in range(n_boot):
        resampled = {cfg: dfs[cfg].iloc[rng.integers(0, len(dfs[cfg]), len(dfs[cfg]))] for cfg in dfs}
        r = minimize(cost_joint, theta_star, args=(resampled, free_names, fixed_by_cfg, model),
                     method="Nelder-Mead", options=_JOINT_MINIMIZE_OPTIONS)
        boots.append(r.x)
    boots = np.array(boots)
    std = boots.std(axis=0)
    corr = pd.DataFrame(np.corrcoef(boots.T), index=free_names, columns=free_names)
    return std, corr

def bootstrap_uncertainty_wide(theta_star, dfs, shared_names, per_dataset_name, fixed_by_cfg, model: GroundTestModel,
                                n_boot=30, seed=None):
    """Bootstrap std + correlation matrix for a `cost_joint_wide` fit.

    Same resampling scheme as `bootstrap_uncertainty_joint`, applied to
    `cost_joint_wide`'s shared+per-dataset parameter layout.

    Returns
    -------
    std  : (len(shared_names) + len(dfs),) array
    corr : DataFrame, index/columns = shared_names + one
        "{per_dataset_name}__{cfg}" per dataset in dfs
    """
    rng = np.random.default_rng(seed)
    names = shared_names + [f"{per_dataset_name}__{cfg}" for cfg in dfs]
    boots = []
    for _ in range(n_boot):
        resampled = {cfg: dfs[cfg].iloc[rng.integers(0, len(dfs[cfg]), len(dfs[cfg]))] for cfg in dfs}
        r = minimize(cost_joint_wide, theta_star, args=(resampled, shared_names, per_dataset_name, fixed_by_cfg, model),
                     method="Nelder-Mead", options=_BOOTSTRAP_MINIMIZE_OPTIONS)
        boots.append(r.x)
    boots = np.array(boots)
    std = boots.std(axis=0)
    corr = pd.DataFrame(np.corrcoef(boots.T), index=names, columns=names)
    return std, corr

def cost(theta, df, free_names, model: GroundTestModel, fixed=None):
    """MSE cost over (r_y, r_z), vectorized over all rows of df.

    `df` must have `y_nisp`, `z_nisp`, `wavelength`, `cent_y`, `cent_z`
    columns (e.g. from `dispcraft.measurement.median_per_spectrum`).
    """
    pred = predict_centroids(df["y_nisp"], df["z_nisp"], df["wavelength"], theta, free_names, model, fixed)
    r_y = pred[0] - df["cent_y"].values
    r_z = pred[1] - df["cent_z"].values
    return np.mean(r_y**2 + r_z**2)

def cost_joint(theta, dfs, free_names, fixed_by_cfg, model: GroundTestModel):
    """Unweighted sum of `cost(...)` over every dataset in `dfs`.

    The core of a joint multi-dataset fit: `free_names`/`theta` are shared
    across all datasets in `dfs`, while `fixed_by_cfg[cfg]` supplies
    whatever each dataset holds at its own value (e.g. an already-resolved
    tilt/offset from a previous fit tier). A single-entry `dfs` reduces this
    to an ordinary single-dataset fit via `cost`.

    No `n_points` weighting -- datasets with more rows get proportionally
    more influence over the shared parameters, same as summing their MSEs
    directly would. Be aware of this if dataset sizes differ substantially.

    `dfs`, `fixed_by_cfg` : dicts keyed by dataset id, same keys.
    """
    return sum(cost(theta, dfs[cfg], free_names, model, fixed=fixed_by_cfg[cfg]) for cfg in dfs)

def cost_joint_wide(theta, dfs, shared_names, per_dataset_name, fixed_by_cfg, model: GroundTestModel):
    """Like `cost_joint`, but `per_dataset_name` is free PER DATASET instead
    of shared, while `shared_names` stay one shared value across `dfs`.

    `theta` layout: `[*one value per shared_names, *one value of
    per_dataset_name per dataset in dfs' key order]`.

    Used to check whether a shared parameter is actually identifiable
    independently of a suspected degenerate partner: free that partner
    per-dataset instead of assuming it's safe to hold fixed, and see how
    much the shared parameter's bootstrap correlation with it (via
    `bootstrap_uncertainty_wide`) actually drops once datasets with
    different geometries are fit jointly.
    """
    n_shared = len(shared_names)
    shared_theta = theta[:n_shared]
    total = 0.0
    for i, cfg in enumerate(dfs):
        theta_d = np.concatenate([shared_theta, [theta[n_shared + i]]])
        total += cost(theta_d, dfs[cfg], shared_names + [per_dataset_name], model, fixed=fixed_by_cfg[cfg])
    return total

def ground_test_model_from_config(base_config: dict) -> GroundTestModel:
    """Build a GroundTestModel from a config shaped like `models/stage1_instrument.toml`."""
    _tel, coll, _grism_nominal, cam, _det = build_instrument_from_config(base_config)
    material = Material(n0=base_config["material"]["n0"], k=base_config["material"]["k"])
    nominal_params = {
        "offset_y_mm": 0.0,
        "offset_z_mm": 0.0,
        "tilt_deg": base_config["grism"]["tilt_deg"],
        "A_deg": base_config["prism"]["A_deg"],
        "rho": base_config["grating"]["rho"],
        "coll_f": coll.f,
        "cam_f": cam.f,
        "material_n0": material.n0,
        "material_k": material.k,
    }
    return GroundTestModel(coll=coll, cam=cam, material=material,
                            m_order=base_config["grating"]["m"], nominal_params=nominal_params)

def per_dataset_rmse(theta, free_names, fixed_by_cfg, dfs, model: GroundTestModel):
    """RMSE per dataset at a shared `theta`, each evaluated against its own
    `fixed_by_cfg[cfg]` (e.g. that dataset's own tilt/offsets).

    `dfs`, `fixed_by_cfg` : dicts keyed by dataset id, same keys.
    """
    return {cfg: float(np.sqrt(cost(theta, dfs[cfg], free_names, model, fixed=fixed_by_cfg[cfg]))) for cfg in dfs}

def predict_centroids(y_nisp_mm, z_nisp_mm, wavelength_nm, theta, free_names, model: GroundTestModel, fixed=None):
    """Forward model: (y_nisp, z_nisp, wavelength) -> (cent_y, cent_z), in mm.

    Parameters
    ----------
    y_nisp_mm, z_nisp_mm, wavelength_nm : array-like, same length
    theta      : free-parameter values, ordered like `free_names`
    free_names : subset of `model.nominal_params`' keys -- which parameters
        `theta` supplies. Everything else is held at `model.nominal_params`
        (or at `fixed`, to override a subset of those constants).
    model      : GroundTestModel -- nominal optics/material/parameter values
    fixed      : dict, optional -- override a subset of nominal_params
        without making it a free parameter (e.g. a fit-specific A_deg)

    Returns
    -------
    (2, N) array -- [cent_y, cent_z], in mm
    """
    p = dict(model.nominal_params)
    if fixed:
        p.update(fixed)
    p.update(dict(zip(free_names, theta)))

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

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