RCAL-905: Add ePSF, ABVegaOffset, and ApCorr Datamodels (#393)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.github/pull_request_template.md

@@ -10,13 +10,13 @@ This PR addresses ...
 
 <!-- if you can't perform these tasks due to permissions, please ask a maintainer to do them -->
 ## Tasks
-- [ ] update or add relevant tests
-- [ ] update relevant docstrings and / or `docs/` page
-- [ ] Does this PR change any API used downstream? (if not, label with `no-changelog-entry-needed`)
-  - [ ] write news fragment(s) in `changes/`: `echo "changed something" > changes/<PR#>.<changetype>.rst` (see below for change types)
-  - [ ] [start a `romancal` regression test](https://github.com/spacetelescope/RegressionTests/actions/workflows/romancal.yml) with this branch installed (`"git+https://github.com/<fork>/roman_datamodels@<branch>"`)
+- [ ] Update or add relevant `roman_datamodels` tests.
+- [ ] Update relevant docstrings and / or `docs/` page.
+- [ ] Does this PR change any API used downstream? (If not, label with `no-changelog-entry-needed`.)
+  - [ ] Write news fragment(s) in `changes/`: `echo "changed something" > changes/<PR#>.<changetype>.rst` (see below for change types).
+  - [ ] Start a `romancal` regression test (https://github.com/spacetelescope/RegressionTests/actions/workflows/romancal.yml) with this branch installed (`"git+https://github.com/<fork>/rad@<branch>"`).
 
-<details><summary>news fragment change types...</summary>
+<details><summary>News fragment change types:</summary>
 
 - ``changes/<PR#>.feature.rst``: new feature
 - ``changes/<PR#>.bugfix.rst``: fixes an issue

changes/393.feature.rst

@@ -0,0 +1 @@
+Added datamodels and tests for ePSF, ABVegaOffset, and ApCorr reference files.

pyproject.toml

@@ -18,8 +18,8 @@ dependencies = [
     "gwcs >=0.19.0",
     "numpy >=1.22",
     "astropy >=5.3.0",
-    "rad >= 0.21.0",
-    # "rad @ git+https://github.com/spacetelescope/rad.git",
+    # "rad >= 0.21.0",
+    "rad @ git+https://github.com/spacetelescope/rad.git",
     "asdf-standard >=1.1.0",
 ]
 dynamic = [

src/roman_datamodels/datamodels/_datamodels.py

@@ -227,6 +227,14 @@ class FlatRefModel(_DataModel):
     _node_type = stnode.FlatRef
 
 
+class AbvegaoffsetRefModel(_DataModel):
+    _node_type = stnode.AbvegaoffsetRef
+
+
+class ApcorrRefModel(_DataModel):
+    _node_type = stnode.ApcorrRef
+
+
 class DarkRefModel(_DataModel):
     _node_type = stnode.DarkRef
 
@@ -235,6 +243,10 @@ class DistortionRefModel(_DataModel):
     _node_type = stnode.DistortionRef
 
 
+class EpsfRefModel(_DataModel):
+    _node_type = stnode.EpsfRef
+
+
 class GainRefModel(_DataModel):
     _node_type = stnode.GainRef
 

src/roman_datamodels/maker_utils/_common_meta.py

@@ -660,6 +660,24 @@ def mk_ref_dark_meta(**kwargs):
     return meta
 
 
+def mk_ref_epsf_meta(**kwargs):
+    """
+    Create dummy metadata for ePSF reference file instances.
+
+    Returns
+    -------
+    dict (follows reference_file/ref_common-1.0.0 schema + ePSF reference file metadata)
+    """
+    meta = mk_ref_common("EPSF", **kwargs)
+    meta["oversample"] = kwargs.get("oversample", NONUM)
+    meta["effective_temperature"] = kwargs.get("effective_temperature", np.arange(1, 10).tolist())
+    meta["defocus"] = kwargs.get("defocus", np.arange(1, 10).tolist())
+    meta["pixel_x"] = kwargs.get("pixel_x", np.arange(1, 10, dtype=np.float32).tolist())
+    meta["pixel_y"] = kwargs.get("pixel_y", np.arange(1, 10, dtype=np.float32).tolist())
+
+    return meta
+
+
 def mk_ref_distoriton_meta(**kwargs):
     """
     Create dummy metadata for distortion reference file instances.

src/roman_datamodels/maker_utils/_ref_files.py

@@ -6,20 +6,24 @@
 
 from roman_datamodels import stnode
 
-from ._base import MESSAGE, save_node
+from ._base import MESSAGE, NONUM, save_node
 from ._common_meta import (
     mk_ref_common,
     mk_ref_dark_meta,
     mk_ref_distoriton_meta,
+    mk_ref_epsf_meta,
     mk_ref_pixelarea_meta,
     mk_ref_readnoise_meta,
     mk_ref_units_dn_meta,
 )
 
 __all__ = [
+    "mk_abvegaoffset",
+    "mk_apcorr",
     "mk_flat",
     "mk_dark",
     "mk_distortion",
+    "mk_epsf",
     "mk_gain",
     "mk_ipc",
     "mk_linearity",
@@ -33,6 +37,108 @@
     "mk_refpix",
 ]
 
+OPT_ELEM = ("F062", "F087", "F106", "F129", "F146", "F158", "F184", "F213", "GRISM", "PRISM", "DARK")
+
+
+def mk_ref_abvegaoffset_data(**kwargs):
+    """
+    Create dummy data for AB Vega Offset reference file instances.
+
+    Returns
+    -------
+    dict
+    """
+    data = {}
+    for optical_elem in OPT_ELEM:
+        data[optical_elem] = {}
+        data[optical_elem]["abvega_offset"] = kwargs.get(optical_elem, {}).get("abvega_offset", float(NONUM))
+
+    return data
+
+
+def mk_abvegaoffset(*, filepath=None, **kwargs):
+    """
+    Create a dummy AB Vega Offset instance (or file) with valid values
+    for attributes required by the schema.
+
+    Parameters
+    ----------
+    filepath
+        (optional, keyword-only) File name and path to write model to.
+
+    Returns
+    -------
+    roman_datamodels.stnode.AbvegaoffsetRef
+    """
+    abvegaref = stnode.AbvegaoffsetRef()
+    abvegaref["meta"] = mk_ref_common("ABVEGAOFFSET", **kwargs.get("meta", {}))
+    abvegaref["data"] = mk_ref_abvegaoffset_data(**kwargs.get("data", {}))
+
+    return save_node(abvegaref, filepath=filepath)
+
+
+def mk_ref_apcorr_data(shape=(10,), **kwargs):
+    """
+    Create dummy data for Aperture Correction reference file instances.
+
+    Parameters
+    ----------
+    shape
+        (optional, keyword-only) Shape of arrays in the model.
+        If shape is greater than 1D, the first dimension is used.
+
+    Returns
+    -------
+    dict
+    """
+    if len(shape) > 1:
+        shape = shape[:1]
+
+        warnings.warn(f"{MESSAGE} assuming the first entry. The remaining are thrown out!", UserWarning)
+
+    data = {}
+    for optical_elem in OPT_ELEM:
+        data[optical_elem] = {}
+        data[optical_elem]["ap_corrections"] = kwargs.get(optical_elem, {}).get(
+            "ap_corrections", np.zeros(shape, dtype=np.float64)
+        )
+        data[optical_elem]["ee_fractions"] = kwargs.get(optical_elem, {}).get("ee_fractions", np.zeros(shape, dtype=np.float64))
+        data[optical_elem]["ee_radii"] = kwargs.get(optical_elem, {}).get("ee_radii", np.zeros(shape, dtype=np.float64))
+        data[optical_elem]["sky_background_rin"] = kwargs.get(optical_elem, {}).get("sky_background_rin", float(NONUM))
+        data[optical_elem]["sky_background_rout"] = kwargs.get(optical_elem, {}).get("sky_background_rout", float(NONUM))
+
+    return data
+
+
+def mk_apcorr(*, shape=(10,), filepath=None, **kwargs):
+    """
+    Create a dummy Aperture Correction instance (or file) with arrays and valid values
+    for attributes required by the schema.
+
+    Parameters
+    ----------
+    shape
+        (optional, keyword-only) Shape of arrays in the model.
+        If shape is greater than 1D, the first dimension is used.
+
+    filepath
+        (optional, keyword-only) File name and path to write model to.
+
+    Returns
+    -------
+    roman_datamodels.stnode.ApcorrRef
+    """
+    if len(shape) > 1:
+        shape = shape[:1]
+
+        warnings.warn(f"{MESSAGE} assuming the first entry. The remaining is thrown out!", UserWarning)
+
+    apcorrref = stnode.ApcorrRef()
+    apcorrref["meta"] = mk_ref_common("APCORR", **kwargs.get("meta", {}))
+    apcorrref["data"] = mk_ref_apcorr_data(shape, **kwargs.get("data", {}))
+
+    return save_node(apcorrref, filepath=filepath)
+
 
 def mk_flat(*, shape=(4096, 4096), filepath=None, **kwargs):
     """
@@ -128,6 +234,37 @@ def mk_distortion(*, filepath=None, **kwargs):
     return save_node(distortionref, filepath=filepath)
 
 
+def mk_epsf(*, shape=(3, 9, 361, 361), filepath=None, **kwargs):
+    """
+    Create a dummy ePSF instance (or file) with arrays and valid values
+    for attributes required by the schema.
+
+    Parameters
+    ----------
+    shape
+        (optional, keyword-only) Shape of arrays in the model.
+        If shape is greater than 1D, the first dimension is used.
+
+    filepath
+        (optional, keyword-only) File name and path to write model to.
+
+    Returns
+    -------
+    roman_datamodels.stnode.EpsfRef
+    """
+    if len(shape) != 4:
+        shape = (3, 9, 361, 361)
+        warnings.warn("Input shape must be 4D. Defaulting to (3, 9, 361, 361)")
+
+    epsfref = stnode.EpsfRef()
+    epsfref["meta"] = mk_ref_epsf_meta(**kwargs.get("meta", {}))
+
+    epsfref["psf"] = kwargs.get("psf", np.zeros(shape, dtype=np.float32))
+    epsfref["extended_psf"] = kwargs.get("extended_psf", np.zeros(shape[2:], dtype=np.float32))
+
+    return save_node(epsfref, filepath=filepath)
+
+
 def mk_gain(*, shape=(4096, 4096), filepath=None, **kwargs):
     """
     Create a dummy Gain instance (or file) with arrays and valid values for
@@ -340,9 +477,7 @@ def _mk_phot_table(**kwargs):
     """
     Create the phot_table for the photom reference file.
     """
-    entries = ("F062", "F087", "F106", "F129", "F146", "F158", "F184", "F213", "GRISM", "PRISM", "DARK")
-
-    return {entry: _mk_phot_table_entry(entry, **kwargs.get(entry, {})) for entry in entries}
+    return {entry: _mk_phot_table_entry(entry, **kwargs.get(entry, {})) for entry in OPT_ELEM}
 
 
 def mk_wfi_img_photom(*, filepath=None, **kwargs):

tests/test_maker_utils.py

@@ -26,6 +26,7 @@ def test_maker_utility_implemented(node_class):
 
 @pytest.mark.parametrize("node_class", stnode.NODE_CLASSES)
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_instance_valid(node_class):
     """
@@ -39,6 +40,7 @@ def test_instance_valid(node_class):
 
 @pytest.mark.parametrize("node_class", [c for c in stnode.NODE_CLASSES if issubclass(c, stnode.TaggedObjectNode)])
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_no_extra_fields(node_class, manifest):
     instance = maker_utils.mk_node(node_class, shape=(8, 8, 8))
@@ -100,6 +102,7 @@ def test_deprecated():
 
 @pytest.mark.parametrize("model_class", [mdl for mdl in maker_utils.NODE_REGISTRY])
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_datamodel_maker(model_class):
     """
@@ -120,6 +123,7 @@ def test_datamodel_maker(model_class):
 
 @pytest.mark.parametrize("node_class", [node for node in datamodels.MODEL_REGISTRY])
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_override_data(node_class):
     """

tests/test_models.py

@@ -53,6 +53,7 @@ def test_datamodel_exists(name):
 
 @pytest.mark.parametrize("model", datamodels.MODEL_REGISTRY.values())
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_node_type_matches_model(model):
     """
@@ -349,6 +350,30 @@ def test_reference_file_model_base(tmp_path):
             raise ValueError("Reference schema does not include ref_common")  # pragma: no cover
 
 
+# AB Vega Offset Correction tests
+def test_make_abvegaoffset():
+    abvegaoffset = utils.mk_abvegaoffset()
+    assert abvegaoffset.meta.reftype == "ABVEGAOFFSET"
+    assert isinstance(abvegaoffset.data.GRISM["abvega_offset"], float)
+
+    # Test validation
+    abvegaoffset_model = datamodels.AbvegaoffsetRefModel(abvegaoffset)
+    assert abvegaoffset_model.validate() is None
+
+
+# Aperture Correction tests
+def test_make_apcorr():
+    apcorr = utils.mk_apcorr()
+    assert apcorr.meta.reftype == "APCORR"
+    assert isinstance(apcorr.data.DARK["sky_background_rin"], float)
+    assert isinstance(apcorr.data.DARK["ap_corrections"], np.ndarray)
+    assert isinstance(apcorr.data.DARK["ap_corrections"][0], float)
+
+    # Test validation
+    apcorr_model = datamodels.ApcorrRefModel(apcorr)
+    assert apcorr_model.validate() is None
+
+
 # Flat tests
 def test_make_flat():
     flat = utils.mk_flat(shape=(8, 8))
@@ -421,6 +446,21 @@ def test_make_distortion():
     assert distortion_model.validate() is None
 
 
+# ePSF tests
+def test_make_epsf():
+    epsf = utils.mk_epsf(shape=(2, 4, 8, 8))
+    assert epsf.meta.reftype == "EPSF"
+    assert isinstance(epsf.meta["pixel_x"], list)
+    assert isinstance(epsf.meta["pixel_x"][0], float)
+    assert epsf["psf"].shape == (2, 4, 8, 8)
+    print(f"XXX type(epsf['psf'][0,0,0,0]) = {type(epsf['psf'][0,0,0,0])}")
+    assert isinstance(epsf["psf"][0, 0, 0, 0], (float, np.float32))
+
+    # Test validation
+    epsf_model = datamodels.EpsfRefModel(epsf)
+    assert epsf_model.validate() is None
+
+
 # Gain tests
 def test_make_gain():
     gain = utils.mk_gain(shape=(8, 8))
@@ -903,6 +943,7 @@ def test_model_validate_without_save():
 @pytest.mark.parametrize("node", datamodels.MODEL_REGISTRY.keys())
 @pytest.mark.parametrize("correct, model", datamodels.MODEL_REGISTRY.items())
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_model_only_init_with_correct_node(node, correct, model):
     """
@@ -945,6 +986,7 @@ def test_ramp_from_science_raw():
 
 @pytest.mark.parametrize("model", datamodels.MODEL_REGISTRY.values())
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_datamodel_construct_like_from_like(model):
     """

tests/test_open.py

@@ -183,6 +183,7 @@ def test_no_memmap(tmp_path, kwargs):
 
 @pytest.mark.parametrize("node_class", [node for node in datamodels.MODEL_REGISTRY])
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_node_round_trip(tmp_path, node_class):
     file_path = tmp_path / "test.asdf"
@@ -195,6 +196,7 @@ def test_node_round_trip(tmp_path, node_class):
 
 @pytest.mark.parametrize("node_class", [node for node in datamodels.MODEL_REGISTRY])
 @pytest.mark.filterwarnings("ignore:This function assumes shape is 2D")
+@pytest.mark.filterwarnings("ignore:Input shape must be 4D")
 @pytest.mark.filterwarnings("ignore:Input shape must be 5D")
 def test_opening_model(tmp_path, node_class):
     file_path = tmp_path / "test.asdf"