dispcraft.prediction

prediction

Single entry point for both 1st- and 0th-order grism predictions.

Closes Status_Report.md §9/§10's "two separate function calls" gap: 1st-order position prediction (dispcraft.calibration.predict_centroids + an optional per-config hybrid ML residual, docs/User_Guide.md §1-2) and 0th-order two-blob separation prediction (dispcraft.zeroth_dispersion's shared-material_k physics baseline + a residual NN, docs/User_Guide.md §3) were each wired together by hand, once per mode, directly in the guide -- this module promotes that wiring into dispcraft/, so both orders share one object and one calling convention instead of bespoke glue code per mode.

CalibratedGrism bundles the physics shared by both orders (one GroundTestModel + one frozen per-config parameter set -- the same fit, evaluated at m_order=1 for position and at m_order=0 for separation, per dispcraft.zeroth_dispersion's module docstring) with each order's optional ML residual. It does not change either model: predict_first_order is exactly docs/User_Guide.md §1's predict_hybrid, and predict_zeroth_order is exactly §3's predict_zeroth_order_full, moved here unchanged so a caller writes them once instead of copying the recipe from documentation. Any residual slot left unset means "physics only" for 1st order (degrades gracefully); 0th order has no such fallback -- m_order=0's grating term is exactly zero, so a residual NN is required to predict anything but a flat zero (see dispcraft.zeroth_dispersion's module docstring).

load_calibrated_grism builds one from models/*.toml + the registered residual models (dispcraft.model_registry, needs MLFLOW_TRACKING_TOKEN); load_hybrid=False/load_zeroth=False skip the network calls entirely (e.g. to build a physics-only CalibratedGrism for tests, or a config that has no registered residual model, like a hand-rolled sensitivity study).

CalibratedGrism dataclass 

CalibratedGrism(model, params, hybrid_mlp_y=None, hybrid_mlp_z=None, hybrid_scaler_mean=None, hybrid_scaler_scale=None, zeroth_material_k=None, zeroth_residual_nn=None)

One config's calibrated model, both orders.

Attributes:
  • model, params (GroundTestModel, dict -- the physics shared by both) –

    orders: nominal optics/material + this config's frozen fit (Tier 1+2+3), e.g. models/joint_specific_fit_<config>.toml

  • hybrid_mlp_y, hybrid_mlp_z ((LitResidualRegressor, optional)) –

    -- 1st-order per-axis residual MLPs; None for a physics-only model

  • hybrid_scaler_mean, hybrid_scaler_scale (array-like, optional --) –

    standardization for (y, z, wavelength) the hybrid MLPs were trained on; required together with hybrid_mlp_y/hybrid_mlp_z

  • zeroth_material_k (float, optional -- physics-baseline `material_k`) –

    override for 0th order; None for configs with no frozen first-order fit to evaluate at m=0 (currently only bgs000_0, Stage 5 Phase 3), in which case predict_zeroth_order is NN-only

  • zeroth_residual_nn (dispcraft.ml.LitResidualRegressor, optional --) –

    0th-order residual NN; required for predict_zeroth_order (see module docstring -- unlike 1st order, there's no physics-only fallback)

predict_first_order 

predict_first_order(y_nisp_mm, z_nisp_mm, wavelength_nm)

(y, z, wavelength) [mm, mm, nm] -> (2, N) [cent_y, cent_z] mm.

Physical prediction alone if no hybrid MLPs were loaded; physical + ML residual otherwise (docs/User_Guide.md §1's predict_hybrid).

Source code in dispcraft/prediction.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def predict_first_order(self, y_nisp_mm, z_nisp_mm, wavelength_nm):
    """(y, z, wavelength) [mm, mm, nm] -> (2, N) [cent_y, cent_z] mm.

    Physical prediction alone if no hybrid MLPs were loaded; physical +
    ML residual otherwise (`docs/User_Guide.md` §1's `predict_hybrid`).
    """
    pred = predict_centroids(y_nisp_mm, z_nisp_mm, wavelength_nm,
                              theta=[], free_names=[], model=self.model, fixed=self.params)
    if self.hybrid_mlp_y is None or self.hybrid_mlp_z is None:
        return pred
    X = (np.column_stack([y_nisp_mm, z_nisp_mm, wavelength_nm])
         - self.hybrid_scaler_mean) / self.hybrid_scaler_scale
    corr_y = ml_predict(self.hybrid_mlp_y, X).ravel()
    corr_z = ml_predict(self.hybrid_mlp_z, X).ravel()
    return np.array([pred[0] - corr_y, pred[1] - corr_z])

predict_zeroth_order 

predict_zeroth_order(y_nisp_mm, z_nisp_mm)

(y, z) [mm] -> (N, 2) [dy, dz] mm, the rank2-rank1 two-blob separation (docs/User_Guide.md §3's predict_zeroth_order_full).

Adds the shared-physics baseline only if zeroth_material_k is set; NN-only otherwise (bgs000_0, see class docstring). Raises if no residual NN was loaded at all -- there's no physics-only fallback for this order (class docstring).

Source code in dispcraft/prediction.py 
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def predict_zeroth_order(self, y_nisp_mm, z_nisp_mm):
    """(y, z) [mm] -> (N, 2) [dy, dz] mm, the rank2-rank1 two-blob
    separation (`docs/User_Guide.md` §3's `predict_zeroth_order_full`).

    Adds the shared-physics baseline only if `zeroth_material_k` is
    set; NN-only otherwise (`bgs000_0`, see class docstring). Raises if
    no residual NN was loaded at all -- there's no physics-only
    fallback for this order (class docstring).
    """
    if self.zeroth_residual_nn is None:
        raise ValueError("no zeroth-order residual NN loaded for this config -- the physical model alone "
                          "predicts exactly zero separation at m_order=0, see dispcraft.zeroth_dispersion")
    resid = predict_zeroth_dispersion(y_nisp_mm, z_nisp_mm, self.zeroth_residual_nn)
    if self.zeroth_material_k is None:
        return resid
    params_zeroth = {**self.params, "material_k": self.zeroth_material_k}
    phys = predict_zeroth_order_separation_physical(y_nisp_mm, z_nisp_mm, self.model, params_zeroth)
    return phys + resid

frozen_params 

frozen_params(fit_toml)

Merge a models/*_fit.toml's [fit.fixed] + [fit.result] into one dict usable as predict_centroids's fixed=, dropping the *_bootstrap_std uncertainty entries recorded alongside the fitted values.

Source code in dispcraft/prediction.py 
51
52
53
54
55
56
57
58
def frozen_params(fit_toml):
    """Merge a `models/*_fit.toml`'s `[fit.fixed]` + `[fit.result]` into one
    dict usable as `predict_centroids`'s `fixed=`, dropping the
    `*_bootstrap_std` uncertainty entries recorded alongside the fitted
    values."""
    fixed = dict(fit_toml["fit"]["fixed"])
    result = {k: v for k, v in fit_toml["fit"]["result"].items() if not k.endswith("_bootstrap_std")}
    return {**fixed, **result}

load_calibrated_grism 

load_calibrated_grism(config, models_dir='models', instrument_toml='stage1_instrument.toml', fit_toml=None, load_hybrid=True, load_zeroth=True, zeroth_material_k=DEFAULT_ZEROTH_MATERIAL_K, registry_version=None)

Build a CalibratedGrism for config from a frozen models/*.toml fit plus the registered residual models (dispcraft.model_registry) -- one call replacing docs/User_Guide.md §1-3's hand-written per-mode setup code.

Parameters:
  • config (str -- e.g. "rgs000_0", "bgs000_0"; also the registered) –

    residual models' name prefix ("<config>-hybrid-mlp-y", etc.)

  • models_dir (Path or str, default: "models" ) –
  • instrument_toml (str -- nominal instrument config filename, relative to, default: 'stage1_instrument.toml' ) –

    models_dir; default is the RGS one -- pass "stage1_bgs_instrument.toml" for bgs000_0

  • fit_toml (str, optional -- frozen fit filename, relative to, default: None ) –

    models_dir; default f"joint_specific_fit_{config}.toml" (the RGS convention) -- pass "bgs000_0_fit.toml" for bgs000_0

  • load_hybrid (bool -- fetch that order's registered, default: True ) –

    residual model(s) from the registry (network call, needs MLFLOW_TRACKING_TOKEN, see dispcraft.model_registry); set False to build a physics-only CalibratedGrism for that order without network access (e.g. in tests)

  • load_zeroth (bool -- fetch that order's registered, default: True ) –

    residual model(s) from the registry (network call, needs MLFLOW_TRACKING_TOKEN, see dispcraft.model_registry); set False to build a physics-only CalibratedGrism for that order without network access (e.g. in tests)

  • zeroth_material_k (float or None, default: `DEFAULT_ZEROTH_MATERIAL_K` ) –

    -- physics-baseline coefficient for 0th order; pass None for a config not fused with this shared-material_k baseline (currently only bgs000_0 -- it has its own frozen first-order fit, but fusing it needs dispcraft.zeroth_dispersion's hardcoded RGS rank wavelengths generalized first, see that module's docstring)

  • registry_version (int, optional -- pin one registered version for, default: None ) –

    every model loaded here (default: latest of each)
Returns:
Source code in dispcraft/prediction.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def load_calibrated_grism(config, models_dir="models", instrument_toml="stage1_instrument.toml",
                           fit_toml=None, load_hybrid=True, load_zeroth=True,
                           zeroth_material_k=DEFAULT_ZEROTH_MATERIAL_K, registry_version=None):
    """Build a `CalibratedGrism` for `config` from a frozen `models/*.toml`
    fit plus the registered residual models (`dispcraft.model_registry`) --
    one call replacing `docs/User_Guide.md` §1-3's hand-written per-mode
    setup code.

    Parameters
    ----------
    config : str -- e.g. "rgs000_0", "bgs000_0"; also the registered
        residual models' name prefix (`"<config>-hybrid-mlp-y"`, etc.)
    models_dir : Path or str, default "models"
    instrument_toml : str -- nominal instrument config filename, relative to
        `models_dir`; default is the RGS one -- pass
        `"stage1_bgs_instrument.toml"` for `bgs000_0`
    fit_toml : str, optional -- frozen fit filename, relative to
        `models_dir`; default `f"joint_specific_fit_{config}.toml"` (the
        RGS convention) -- pass `"bgs000_0_fit.toml"` for `bgs000_0`
    load_hybrid, load_zeroth : bool -- fetch that order's registered
        residual model(s) from the registry (network call, needs
        `MLFLOW_TRACKING_TOKEN`, see `dispcraft.model_registry`); set
        `False` to build a physics-only `CalibratedGrism` for that order
        without network access (e.g. in tests)
    zeroth_material_k : float or None, default `DEFAULT_ZEROTH_MATERIAL_K`
        -- physics-baseline coefficient for 0th order; pass `None` for a
        config not fused with this shared-`material_k` baseline (currently
        only `bgs000_0` -- it has its own frozen first-order fit, but
        fusing it needs `dispcraft.zeroth_dispersion`'s hardcoded RGS rank
        wavelengths generalized first, see that module's docstring)
    registry_version : int, optional -- pin one registered version for
        every model loaded here (default: latest of each)

    Returns
    -------
    CalibratedGrism
    """
    models_dir = Path(models_dir)
    fit_toml = fit_toml or f"joint_specific_fit_{config}.toml"

    with open(models_dir / instrument_toml, "rb") as f:
        model = ground_test_model_from_config(tomllib.load(f))
    with open(models_dir / fit_toml, "rb") as f:
        params = frozen_params(tomllib.load(f))

    hybrid_mlp_y = hybrid_mlp_z = hybrid_scaler_mean = hybrid_scaler_scale = None
    if load_hybrid:
        hybrid_mlp_y, extra = load_model(f"{config}-hybrid-mlp-y", version=registry_version)
        hybrid_mlp_z, _ = load_model(f"{config}-hybrid-mlp-z", version=registry_version)
        hybrid_scaler_mean, hybrid_scaler_scale = extra["scaler_mean"], extra["scaler_scale"]

    zeroth_residual_nn = None
    if load_zeroth:
        zeroth_residual_nn, _ = load_model(f"{config}-zeroth-residual-nn", version=registry_version)

    return CalibratedGrism(
        model=model, params=params,
        hybrid_mlp_y=hybrid_mlp_y, hybrid_mlp_z=hybrid_mlp_z,
        hybrid_scaler_mean=hybrid_scaler_mean, hybrid_scaler_scale=hybrid_scaler_scale,
        zeroth_material_k=zeroth_material_k, zeroth_residual_nn=zeroth_residual_nn,
    )