Source code for vmpt.wavelengths

"""Per-shutter NIRSpec dispersion model for MSA shutters.

All supported (disperser, filter) combinations use a *per-shutter*
lookup table derived from spacetelescope/msaviz's numerical
integration of the pipeline dispersion polynomials. The table lives
in `data/dispersion_cutoffs.npz` and is regenerated by
`scripts/precompute_dispersion_cutoffs.py` (re-run when the
underlying msaviz reference files change).

For each (disperser, filter) the table stores four (4, 171, 365)
float32 arrays — one slice per quadrant — under keys

    {DISPERSER}_{FILTER}_blue_edge
    {DISPERSER}_{FILTER}_gap_lo
    {DISPERSER}_{FILTER}_gap_hi
    {DISPERSER}_{FILTER}_red_edge

NaN means the shutter's spectrum doesn't reach the corresponding
detector. A linear V2-shift fallback is retained for the case where
the table is missing on disk.
"""

from pathlib import Path

import numpy as np

from .coords import MSA_V2_REF

# Per-grating fiducial wavelength ranges (microns), sourced from
# spacetelescope/msaviz's `ranges.json`, which is derived from the
# pipeline reference files. These are the wavelength bounds the
# detector ACTUALLY sees through each (disperser, filter) — wider
# than the JDox "useful range" in a few places (e.g. G140M/F100LP
# JDox quotes 0.97-1.84 but the pipeline reference is 0.97-1.89).
# We use the pipeline values so the per-shutter table lookups and
# this fallback agree.
GRATING_RANGES: dict[str, dict[str, tuple[float, float]]] = {
    "PRISM": {"CLEAR": (0.60, 5.30)},
    "G140M": {"F070LP": (0.70, 1.27), "F100LP": (0.97, 1.89)},
    "G235M": {"F170LP": (1.66, 3.17)},
    "G395M": {"F290LP": (2.87, 5.27)},
    "G140H": {"F070LP": (0.70, 1.27), "F100LP": (0.97, 1.89)},
    "G235H": {"F170LP": (1.66, 3.17)},
    "G395H": {"F290LP": (2.87, 5.27)},
}

# Filter blue cutoffs (microns) — hard lower bounds.
FILTER_BLUE_CUTOFF: dict[str, float] = {
    "CLEAR": 0.60,
    "F070LP": 0.70,
    "F100LP": 0.97,
    "F170LP": 1.66,
    "F290LP": 2.87,
}

# Projected MSA-to-detector dispersion span along V2 (arcsec). Placeholder per PLAN.md.
V2_DISP_EXTENT: float = 180.0

LAMBDA_PER_ARCSEC: float = (5.30 - 0.60) / V2_DISP_EXTENT  # ~0.026 μm/arcsec


# Half-width of the on-detector spectrum, in V2 arcsec, per disperser.
#
# Source: eMPT (Bonaventura et al. 2023, A&A 672, A40),
# `reference_files/prism_sep.dat`, which tabulates per-shutter +/- column
# separations for the PRISM spectrum. At a central shutter the values are
# sep_p ≈ +176, sep_m ≈ -177 columns. Each MSA column is ~0.20" along V2,
# so the projected V2 half-extent of PRISM's spectrum is ~0.20 × 176 ≈ 35".
# Within ~10 % across the MSA.
#
# For the grating modes, eMPT applies no column cutoff at all — any pair
# of shutters at the same detector-y row collide (the spectrum spans the
# entire detector). For an H-grating the spectrum reaches beyond the MSA
# in V2, so cross-quadrant pairs collide as well. We approximate both
# behaviours with a single (large) V2 half-extent below.
#
# IMPORTANT — what the visualisation actually shows. We compute the
# orange spectral-conflict shutters as "same q AND same s in the MSA grid,
# within `v2_overlap_distance` of the open shutter in V2." So:
#
#   - PRISM at a center shutter (d ~ 200): ~94 % of the row in the same
#     quadrant. (Yes, "most of the row" — physical: PRISM spectrum on
#     detector is ~70" wide in V2, the row is 73" wide.)
#   - PRISM at an edge shutter (d ~ 10): only ~51 % of the row. (The
#     spectrum runs off the end.)
#   - M / H gratings: ~100 % of the row in same quadrant.
#
# Not yet modelled: cross-quadrant collisions for grating modes. Two
# shutters in different quadrants but at the same detector y can collide;
# we'd need the eMPT shval tables to do this exactly. Reasonable
# enhancement for a future release.
SPECTRUM_V2_HALFEXTENT: dict[str, float] = {
    "PRISM": 35.0,
    "G140M": 200.0,
    "G235M": 200.0,
    "G395M": 200.0,
    "G140H": 500.0,
    "G235H": 500.0,
    "G395H": 500.0,
}


[docs] def v2_overlap_distance(disperser: str, filt: str) -> float: """V2 half-extent (arcsec) of the spectrum on the detector. Two same-row shutters whose V2 separation is less than this value collide on the detector. See module-level comment for the eMPT reference.""" return SPECTRUM_V2_HALFEXTENT.get(disperser.upper(), 180.0)
[docs] def disperser_range(disperser: str, filt: str) -> tuple[float, float] | None: """Nominal ``(lam_min, lam_max)`` in μm for a (disperser, filter) combination. Returns the wavelength range the chosen mode actually delivers to the detector (the values in :data:`GRATING_RANGES`) — the same numbers shown in the JDox "useful range" docs, except where msaviz / the pipeline reference disagree (we follow the pipeline). Returns ``None`` when the combination isn't supported (e.g. ``G140H`` + ``F290LP``); the optimizer treats that as "no constraint can pass" and drops the affected sources. """ if disperser is None or filt is None: return None d = disperser.upper() f = filt.upper() bands = GRATING_RANGES.get(d) if bands is None: return None return bands.get(f)
[docs] def disperser_min_lambda(disperser: str, filt: str) -> float | None: """Bluest wavelength the (disperser, filter) combo can deliver. ``None`` when the combination isn't recognised.""" r = disperser_range(disperser, filt) return None if r is None else r[0]
[docs] def disperser_max_lambda(disperser: str, filt: str) -> float | None: """Reddest wavelength the (disperser, filter) combo can deliver. ``None`` when the combination isn't recognised.""" r = disperser_range(disperser, filt) return None if r is None else r[1]
[docs] def interval_covered( lo: float, hi: float, blue: float, gap_lo: float, gap_hi: float, red: float, ) -> bool: """Does the spectrum ``[blue, red]`` (with the detector gap ``[gap_lo, gap_hi]`` excluded) fully cover the requested ``[lo, hi]`` range? Used by the per-target "required wavelength range" constraint. Parameters ---------- lo, hi : float Requested interval, in μm. Must satisfy ``lo <= hi``. blue, red : float Bluest / reddest λ the centre shutter delivers to the detector. Both NaN when the source doesn't reach the detector at all (cutoffs returns NaN for off-grid shutters). gap_lo, gap_hi : float NRS1/NRS2 detector-gap wavelength bounds for this shutter. Both NaN ⇒ no gap (e.g. M-grating modes within their nominal range). Returns ------- bool True iff every wavelength in ``[lo, hi]`` lands somewhere on the detector (i.e. inside ``[blue, gap_lo] ∪ [gap_hi, red]``, where the gap is skipped iff it has finite bounds). """ if not (np.isfinite(lo) and np.isfinite(hi)): return False if not (np.isfinite(blue) and np.isfinite(red)): return False if lo > hi: lo, hi = hi, lo if lo < blue or hi > red: return False has_gap = np.isfinite(gap_lo) and np.isfinite(gap_hi) if not has_gap: return True # Gap is INSIDE [blue, red]; reject if [lo, hi] dips into it. # Strict: any wavelength in [lo, hi] that falls in (gap_lo, gap_hi) # is missing. return not (lo < gap_hi and hi > gap_lo)
# Fallback fractional gap parameters used when neither the per- # shutter PRISM table NOR a per-disperser fiducial value is # available. Width tightened from the previous 10 % so the gratings # show a more believably narrow detector gap. GAP_CENTER_REL: float = 0.50 GAP_WIDTH_REL: float = 0.04 # Per-disperser fiducial gap fallback (used only when the per-shutter # PRISM table can't be located — kept as a safety net so vMPT still # loads if `data/prism_cutoffs.npz` goes missing). DETECTOR_GAP_FIDUCIAL: dict[str, tuple[float, float]] = { "PRISM": (1.87, 3.93), } # Per-shutter dispersion bounds for every supported (disperser, # filter) combo, generated by scripts/precompute_dispersion_cutoffs.py # and shipped at vmpt/data/dispersion_cutoffs.npz. Loaded lazily on # first lookup so import stays cheap. _DISPERSION_TABLE_PATH = ( Path(__file__).resolve().parent / "data" / "dispersion_cutoffs.npz" ) _DISPERSION_TABLE: dict | None = None _DISPERSION_TABLE_LOAD_ATTEMPTED: bool = False def _load_dispersion_table() -> dict | None: """Lazy-load the per-shutter dispersion bounds table. Returns None if the npz isn't present (fresh checkout, fall back to the linear-shift model).""" global _DISPERSION_TABLE, _DISPERSION_TABLE_LOAD_ATTEMPTED if _DISPERSION_TABLE_LOAD_ATTEMPTED: return _DISPERSION_TABLE _DISPERSION_TABLE_LOAD_ATTEMPTED = True if not _DISPERSION_TABLE_PATH.exists(): return None try: # Eager-copy the relevant arrays so the npz handle can close. with np.load(_DISPERSION_TABLE_PATH) as data: _DISPERSION_TABLE = {key: data[key] for key in data.files} except (OSError, KeyError, ValueError): _DISPERSION_TABLE = None return _DISPERSION_TABLE def _table_lookup( disperser: str, filt: str, q: int, s: int, d: int, ) -> tuple[float, float, float, float] | None: """Read the precomputed wavelength bounds for shutter (q, s, d) under the given (disperser, filter). Returns None if the table or the specific combo isn't loadable. Per-value NaN means the spectrum doesn't reach that detector.""" tbl = _load_dispersion_table() if tbl is None: return None qi, si, di = q - 1, s - 1, d - 1 if not (0 <= qi < 4 and 0 <= si < 171 and 0 <= di < 365): return None key_blue = f"{disperser}_{filt}_blue_edge" key_glo = f"{disperser}_{filt}_gap_lo" key_ghi = f"{disperser}_{filt}_gap_hi" key_red = f"{disperser}_{filt}_red_edge" if key_blue not in tbl: return None return ( float(tbl[key_blue][qi, si, di]), float(tbl[key_glo][qi, si, di]), float(tbl[key_ghi][qi, si, di]), float(tbl[key_red][qi, si, di]), )
[docs] def cutoffs( v2_arcsec: float, v3_arcsec: float, disperser: str, filt: str, *, q: int | None = None, s: int | None = None, d: int | None = None, ) -> dict: """Wavelength endpoints of the dispersed spectrum on the detector for a shutter at (V2, V3). If `q, s, d` are supplied AND the precomputed dispersion table is available for the requested (disperser, filter), returns the per-shutter values from `data/dispersion_cutoffs.npz` (derived from msaviz's integration of the pipeline dispersion models). This is the accurate path; the gap location and spectrum edges vary substantially across the MSA for every disperser, but especially for PRISM (non-linear dispersion). Without shutter indices, or if the table is missing, the function falls back to a linear V2-shift model. The fallback exists so vMPT still loads on a fresh checkout that hasn't run the precompute, and so the existing test suite (which calls `cutoffs(V2, V3, ...)` without indices) keeps working. """ disperser = disperser.upper() filt = filt.upper() if disperser not in GRATING_RANGES or filt not in GRATING_RANGES[disperser]: raise ValueError(f"Unsupported (disperser, filter) = ({disperser}, {filt})") lam_min, lam_max = GRATING_RANGES[disperser][filt] blue_cut = FILTER_BLUE_CUTOFF.get(filt, 0.0) def _maybe(lam, is_blue_edge: bool = False, lam_red_ref: float | None = None): """Apply filter blue-cut and convert NaN/None → None.""" if lam is None: return None if isinstance(lam, float) and (lam != lam): # NaN check w/o numpy return None if lam < blue_cut: if is_blue_edge and lam_red_ref is not None and lam_red_ref > blue_cut: return float(blue_cut) return None return float(lam) # Per-shutter lookup — preferred when indices + table are available # for the given (disperser, filter). if q is not None and s is not None and d is not None: hit = _table_lookup(disperser, filt, int(q), int(s), int(d)) if hit is not None: blue, glo, ghi, red = hit # If both bounds are NaN the shutter has no spectrum at all # on either detector for this combo. if blue != blue and red != red: return {"lam_blue": None, "lam_gap_lo": None, "lam_gap_hi": None, "lam_red": None} return { "lam_blue": _maybe(blue, is_blue_edge=True, lam_red_ref=red), "lam_gap_lo": _maybe(glo), "lam_gap_hi": _maybe(ghi), "lam_red": _maybe(red), } # Fall through to fiducial / linear path if the lookup missed. # Fallback / grating path: linear V2 shift, fiducial or fractional gap. span = lam_max - lam_min dlam_dv2 = span / V2_DISP_EXTENT shift = (v2_arcsec - MSA_V2_REF) * dlam_dv2 if disperser == "PRISM": # PRISM dispersion isn't linear in V2; without the table we # hold endpoints at the disperser's full range (per msaviz # the spread across shutters is ~0.01 μm anyway). lam_blue_raw = lam_min lam_red_raw = lam_max else: lam_blue_raw = lam_min + shift lam_red_raw = lam_max + shift lam_blue = max(lam_min, min(lam_max, lam_blue_raw)) lam_red = max(lam_min, min(lam_max, lam_red_raw)) fixed_gap = DETECTOR_GAP_FIDUCIAL.get(disperser) if fixed_gap is not None: lam_gap_lo_raw, lam_gap_hi_raw = fixed_gap else: gap_half = 0.5 * GAP_WIDTH_REL * span lam_gap_center_raw = lam_min + GAP_CENTER_REL * span + shift lam_gap_lo_raw = lam_gap_center_raw - gap_half lam_gap_hi_raw = lam_gap_center_raw + gap_half if lam_gap_hi_raw < lam_blue or lam_gap_lo_raw > lam_red: lam_gap_lo = None lam_gap_hi = None else: lam_gap_lo = max(lam_blue, min(lam_red, lam_gap_lo_raw)) lam_gap_hi = max(lam_blue, min(lam_red, lam_gap_hi_raw)) return { "lam_blue": _maybe(lam_blue, is_blue_edge=True, lam_red_ref=lam_red), "lam_gap_lo": _maybe(lam_gap_lo), "lam_gap_hi": _maybe(lam_gap_hi), "lam_red": _maybe(lam_red), }