Read grid acq probability model data from in chandra models repo

chandra_aca/drift.py

@@ -17,28 +17,12 @@
 from astropy.table import Table
 from astropy.utils.data import download_file
 from Chandra.Time import DateTime
+from ska_helpers.paths import aca_drift_model_path
 from ska_helpers.utils import LazyDict
 
-# Name of environment variable to override default root for data files
-DATA_ROOT_ENV_VAR = "THERMAL_MODELS_DIR_FOR_MATLAB_TOOLS_SW"
-
-
-# Define path to best fit model parameters for ACA drift model.
-# See notebooks fit_aimpoint_drift_*.ipynb for fit details.
-def DRIFT_MODEL_PATH():
-    default_root = Path(
-        os.environ["SKA"],
-        "data",
-        "chandra_models",
-    )
-    root = os.environ.get(DATA_ROOT_ENV_VAR, default_root)
-    path = Path(root) / "chandra_models" / "aca_drift" / "aca_drift_model.json"
-
-    return path
-
 
 def load_drift_pars():
-    pars = json.loads(DRIFT_MODEL_PATH().read_text())
+    pars = json.loads(aca_drift_model_path().read_text())
     return pars
 
 
@@ -127,8 +111,8 @@ def __init__(self, scale, offset, trend, jumps, year0):
 
     def calc(self, times, t_ccd):
         """
-        Calculate the drift model for aspect solution SIM DY/DZ values for input ``times``
-        and ``t_ccd``.  The two arrays are broadcasted to match.
+        Calculate the drift model for aspect solution SIM DY/DZ values for input
+        ``times`` and ``t_ccd``.  The two arrays are broadcasted to match.
 
         The returned drifts are in arcsec and provide the expected aspect solution
         SIM DY or DZ values in mm.  This can be converted to a drift in arcsec via

chandra_aca/star_probs.py

@@ -10,33 +10,30 @@
 SSAWG review: 2020-01-29
 """
 
-import os
+import functools
 import warnings
 
 import numpy as np
 import scipy.stats
 from Chandra.Time import DateTime
 from numba import jit
 from scipy.optimize import bisect, brentq
+from ska_helpers.paths import aca_acq_prob_models_path
 
 from chandra_aca.transform import (
     broadcast_arrays,
     broadcast_arrays_flatten,
     snr_mag_for_t_ccd,
 )
 
-STAR_PROBS_DATA_DIR = os.path.join(os.path.dirname(__file__), "data", "star_probs")
-
 # Default acquisition probability model
 DEFAULT_MODEL = "grid-floor-2020-02"
 
 # Cache of cubic spline functions.  Eval'd only on the first time.
 SPLINE_FUNCS = {}
 
-# Cache of grid model interpolation functions
-GRID_FUNCS = {}
-
-WARM_THRESHOLD = 100  # Value (N100) used for fitting
+# Value (N100) used for fitting
+WARM_THRESHOLD = 100
 
 
 # Min and max star acquisition probabilities, regardless of model predictions
@@ -116,7 +113,7 @@ def t_ccd_warm_limit(
     :param warm_t_ccd: warmest CCD temperature to consider (default=-5 C)
     :param model: probability model (see acq_success_prob() for allowed values, default)
 
-    :returns: (t_ccd, n_acq | prob_n_or_fewer) tuple with CCD temperature upper limit and:
+    :returns: (t_ccd, n_acq | prob_n_or_fewer) tuple with CCD temperature upper limit:
               - number of expected ACQ stars at that temperature (scalar min_n_acq)
               - probability of acquiring ``n`` or fewer stars (tuple min_n_acq)
     """
@@ -165,9 +162,9 @@ def prob_n_or_fewer_below_max(t_ccd):
         t_ccd = warm_t_ccd
 
     elif merit_func(cold_t_ccd) <= 0:
-        # If there are not enough ACQ stars at the coldest CCD temperature then stop there
-        # as well.  The ACA thermal model will never predict a temperature below this
-        # value so this catalog will fail thermal check.
+        # If there are not enough ACQ stars at the coldest CCD temperature then stop
+        # there as well.  The ACA thermal model will never predict a temperature below
+        # this value so this catalog will fail thermal check.
         t_ccd = cold_t_ccd
 
     else:
@@ -232,8 +229,8 @@ def acq_success_prob(
     model=None,
 ):
     """
-    Return probability of acquisition success for given date, temperature, star properties
-    and search box size.
+    Return probability of acquisition success for given date, temperature, star
+    properties and search box size.
 
     Any of the inputs can be scalars or arrays, with the output being the result of
     the broadcasted dimension of the inputs.
@@ -330,6 +327,83 @@ def clip_and_warn(name, val, val_lo, val_hi, model):
     return val
 
 
+def get_grid_axis_values(hdr, axis):
+    """Get grid model axis values from FITS header.
+
+    This is an irregularly-spaced grid if ``hdr`` has ``{axis}_0`` .. ``{axis}_<N-1>``.
+    Otherwise it is a regularly-spaced grid:
+        linspace(hdr[f"{axis}_lo"], hdr[f"{axis}_hi"], n_vals)
+
+    :param hdr: FITS header (dict-like)
+    :param axis: Axis name (e.g. "mag")
+    """
+    n_vals = hdr[f"{axis}_n"]
+    if all(f"{axis}_{ii}" in hdr for ii in range(n_vals)):
+        # New style grid model file
+        vals = np.array([hdr[f"{axis}_{ii}"] for ii in range(n_vals)])
+    else:
+        # Old style grid model file
+        vals = np.linspace(hdr[f"{axis}_lo"], hdr[f"{axis}_hi"], n_vals)
+
+    return vals
+
+
+@functools.lru_cache
+def get_grid_func_model(model):
+    from astropy.io import fits
+    from scipy.interpolate import RegularGridInterpolator
+
+    # Read the model file and put into local vars
+    filepath = aca_acq_prob_models_path() / (model + ".fits.gz")
+    if not filepath.exists():
+        raise IOError(f"model file {filepath} does not exist")
+
+    hdus = fits.open(filepath)
+    hdu0 = hdus[0]
+    probit_p_fail_no_1p5 = hdus[1].data
+    probit_p_fail_1p5 = hdus[2].data
+
+    hdr = hdu0.header
+    grid_mags = get_grid_axis_values(hdr, "mag")
+    grid_t_ccds = get_grid_axis_values(hdr, "t_ccd")
+    grid_halfws = get_grid_axis_values(hdr, "halfw")
+
+    # Sanity checks on model data
+    assert probit_p_fail_no_1p5.shape == (
+        len(grid_mags),
+        len(grid_t_ccds),
+        len(grid_halfws),
+    )
+    assert probit_p_fail_1p5.shape == probit_p_fail_no_1p5.shape
+
+    # Generate the 3-d linear interpolation functions
+    func_no_1p5 = RegularGridInterpolator(
+        points=(grid_mags, grid_t_ccds, grid_halfws), values=probit_p_fail_no_1p5
+    )
+    func_1p5 = RegularGridInterpolator(
+        points=(grid_mags, grid_t_ccds, grid_halfws), values=probit_p_fail_1p5
+    )
+    mag_lo = hdr["mag_lo"]
+    mag_hi = hdr["mag_hi"]
+    t_ccd_lo = hdr["t_ccd_lo"]
+    t_ccd_hi = hdr["t_ccd_hi"]
+    halfw_lo = hdr["halfw_lo"]
+    halfw_hi = hdr["halfw_hi"]
+
+    out = {
+        "filename": filepath,
+        "func_no_1p5": func_no_1p5,
+        "func_1p5": func_1p5,
+        "mag_lo": mag_lo,
+        "mag_hi": mag_hi,
+        "t_ccd_lo": t_ccd_lo,
+        "t_ccd_hi": t_ccd_hi,
+        "halfw_lo": halfw_lo,
+        "halfw_hi": halfw_hi,
+    }
+    return out
+
+
 def grid_model_acq_prob(
     mag=10.0, t_ccd=-12.0, color=0.6, halfwidth=120, probit=False, model=None
 ):
@@ -349,69 +423,18 @@ def grid_model_acq_prob(
     :returns: Acquisition success probability(s)
 
     """
-    # Info about model is cached in GRID_FUNCS
-    if model not in GRID_FUNCS:
-        from astropy.io import fits
-        from scipy.interpolate import RegularGridInterpolator
-
-        # Read the model file and put into local vars
-        filename = os.path.join(STAR_PROBS_DATA_DIR, model) + ".fits.gz"
-        if not os.path.exists(filename):
-            raise IOError("model file {} does not exist".format(filename))
-
-        hdus = fits.open(filename)
-        hdu0 = hdus[0]
-        probit_p_fail_no_1p5 = hdus[1].data
-        probit_p_fail_1p5 = hdus[2].data
-
-        hdr = hdu0.header
-        grid_mags = np.linspace(hdr["mag_lo"], hdr["mag_hi"], hdr["mag_n"])
-        grid_t_ccds = np.linspace(hdr["t_ccd_lo"], hdr["t_ccd_hi"], hdr["t_ccd_n"])
-        grid_halfws = np.linspace(hdr["halfw_lo"], hdr["halfw_hi"], hdr["halfw_n"])
-
-        # Sanity checks on model data
-        assert probit_p_fail_no_1p5.shape == (
-            len(grid_mags),
-            len(grid_t_ccds),
-            len(grid_halfws),
-        )
-        assert probit_p_fail_1p5.shape == probit_p_fail_no_1p5.shape
-
-        # Generate the 3-d linear interpolation functions
-        func_no_1p5 = RegularGridInterpolator(
-            points=(grid_mags, grid_t_ccds, grid_halfws), values=probit_p_fail_no_1p5
-        )
-        func_1p5 = RegularGridInterpolator(
-            points=(grid_mags, grid_t_ccds, grid_halfws), values=probit_p_fail_1p5
-        )
-        mag_lo = hdr["mag_lo"]
-        mag_hi = hdr["mag_hi"]
-        t_ccd_lo = hdr["t_ccd_lo"]
-        t_ccd_hi = hdr["t_ccd_hi"]
-        halfw_lo = hdr["halfw_lo"]
-        halfw_hi = hdr["halfw_hi"]
-
-        GRID_FUNCS[model] = {
-            "filename": filename,
-            "func_no_1p5": func_no_1p5,
-            "func_1p5": func_1p5,
-            "mag_lo": mag_lo,
-            "mag_hi": mag_hi,
-            "t_ccd_lo": t_ccd_lo,
-            "t_ccd_hi": t_ccd_hi,
-            "halfw_lo": halfw_lo,
-            "halfw_hi": halfw_hi,
-        }
-    else:
-        gfm = GRID_FUNCS[model]
-        func_no_1p5 = gfm["func_no_1p5"]
-        func_1p5 = gfm["func_1p5"]
-        mag_lo = gfm["mag_lo"]
-        mag_hi = gfm["mag_hi"]
-        t_ccd_lo = gfm["t_ccd_lo"]
-        t_ccd_hi = gfm["t_ccd_hi"]
-        halfw_lo = gfm["halfw_lo"]
-        halfw_hi = gfm["halfw_hi"]
+    # Get the grid model function and model parameters from a FITS file. This function
+    # call is cached.
+    gfm = get_grid_func_model(model)
+
+    func_no_1p5 = gfm["func_no_1p5"]
+    func_1p5 = gfm["func_1p5"]
+    mag_lo = gfm["mag_lo"]
+    mag_hi = gfm["mag_hi"]
+    t_ccd_lo = gfm["t_ccd_lo"]
+    t_ccd_hi = gfm["t_ccd_hi"]
+    halfw_lo = gfm["halfw_lo"]
+    halfw_hi = gfm["halfw_hi"]
 
     # Make sure inputs are within range of gridded model
     mag = clip_and_warn("mag", mag, mag_lo, mag_hi, model)
@@ -473,7 +496,7 @@ def spline_model_acq_prob(
     :param probit: if True then return Probit(p_success). Default=False
 
     :returns: Acquisition success probability(s)
-    """
+    """  # noqa: E501
     try:
         from scipy.interpolate import CubicSpline
     except ImportError:
@@ -562,11 +585,11 @@ def spline_model_acq_prob(
         tcm = tc12[mask]
         boxm = box_deltas[mask]
 
-        # The model is only calibrated betweeen 8.5 and 10.7.  First, clip mags going into
-        # the spline to be larger than 8.5.  (Extrapolating slightly above 10.7 is OK).
-        # Second, subtract a linearly varying term from probit_p_fail from stars brighter
-        # than 8.5 mag ==> from 0.0 for mag=8.5 to 0.25 for mag=6.0.  This is to allow
-        # star selection to favor a 6.0 mag star over 8.5 mag.
+        # The model is only calibrated betweeen 8.5 and 10.7.  First, clip mags going
+        # into the spline to be larger than 8.5.  (Extrapolating slightly above 10.7 is
+        # OK). Second, subtract a linearly varying term from probit_p_fail from stars
+        # brighter than 8.5 mag ==> from 0.0 for mag=8.5 to 0.25 for mag=6.0.  This is
+        # to allow star selection to favor a 6.0 mag star over 8.5 mag.
         magmc = magm.clip(8.5, None)
         bright = (8.5 - magm.clip(None, 8.5)) / 10.0
 
@@ -598,9 +621,11 @@ def model_acq_success_prob(mag, warm_frac, color=0, halfwidth=120):
 
 def sota_model_acq_prob(mag, warm_frac, color=0, halfwidth=120):
     """
-    Calculate raw SOTA model probability of acquisition success for a star with ``mag``
-    magnitude and a CCD warm fraction ``warm_frac``.  This is not typically used directly
-    since it does not account for star properties like spoiled or color=0.7.
+    Calculate raw SOTA model probability of acquisition success.
+
+    This is for a star with ``mag`` magnitude and a CCD warm fraction ``warm_frac``.
+    This is not typically used directly since it does not account for star properties
+    like spoiled or color=0.7.
 
     Uses the empirical relation::
 
@@ -609,9 +634,9 @@ def sota_model_acq_prob(mag, warm_frac, color=0, halfwidth=120):
        P_acq_success = 1 - P_acq_fail
 
     This is based on the dark model and acquisition success model presented in the State
-    of the ACA 2013, and subsequently updated to use a Probit transform and separately fit
-    B-V=1.5 stars.  It was updated in 2017 to include a fitted dependence on the search
-    box halfwidth.  See:
+    of the ACA 2013, and subsequently updated to use a Probit transform and separately
+    fit B-V=1.5 stars.  It was updated in 2017 to include a fitted dependence on the
+    search box halfwidth.  See:
 
     https://github.com/sot/skanb/blob/master/pea-test-set/fit_box_size_acq_prob.ipynb
     https://github.com/sot/aca_stats/blob/master/fit_acq_prob_model-2017-07-sota.ipynb

chandra_aca/tests/test_drift.py

@@ -4,6 +4,7 @@
 import astropy.table as tbl
 import numpy as np
 import pytest
+import ska_helpers.paths
 
 from chandra_aca import drift
 
@@ -58,7 +59,7 @@ def test_get_aca_offsets(kwargs, env_override, monkeypatch):
     """Regression test that ACA offsets match the original flight values to expected
     precision."""
     if env_override:
-        monkeypatch.setenv(drift.DATA_ROOT_ENV_VAR, env_override)
+        monkeypatch.setenv(ska_helpers.paths.CHANDRA_MODELS_ROOT_ENV_VAR, env_override)
     kwargs = kwargs.copy()
     aca_offset_y = kwargs.pop("aca_offset_y")
     aca_offset_z = kwargs.pop("aca_offset_z")

chandra_aca/tests/test_star_probs.py

@@ -2,15 +2,14 @@
 import hashlib
 import itertools
 import os
-from pathlib import Path
 
 import numpy as np
 import pytest
 from astropy.table import Table
+from ska_helpers.paths import aca_acq_prob_models_path
 
 from chandra_aca.star_probs import (
     DEFAULT_MODEL,
-    STAR_PROBS_DATA_DIR,
     acq_success_prob,
     binom_ppf,
     grid_model_acq_prob,
@@ -591,7 +590,7 @@ def test_md5_2020_02():
     """Test that model data file has expected MD5 sum.  See cell 29 of
     aca_stats/fit_acq_model-2020-02-binned-poly-binom-floor.ipynb
     """
-    filename = Path(STAR_PROBS_DATA_DIR, "grid-floor-2020-02.fits.gz")
+    filename = aca_acq_prob_models_path() / "grid-floor-2020-02.fits.gz"
     md5 = hashlib.md5(open(filename, "rb").read()).hexdigest()
     assert md5 == "39960b6254acc4a7500397aac9908412"