"""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),
}