Simplify tweakreg unit tests

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "crds >=13.0.2",
     "drizzle >= 2.1.1,<2.2",
     "gwcs >=0.25.2,<0.27.0",
-    "stcal>=1.15.1,<1.16.0",
+    "stcal>=1.15.2,<1.16.0",
     "stpipe >=0.11.0,<0.12.0",
     "spherical-geometry >=1.3.3,<1.4",
     "stsci.imagestats >=1.8.3,<1.9",  # remove when stcal fixes this missing dep

romancal/tweakreg/tests/conftest.py

@@ -0,0 +1,128 @@
+import numpy as np
+import pytest
+from astropy import coordinates as coord
+from astropy import units as u
+from astropy.modeling.models import RotationSequence3D, Scale, Shift
+from astropy.table import Table
+from astropy.time import Time
+from gwcs import coordinate_frames as cf
+from gwcs import wcs
+from gwcs.geometry import CartesianToSpherical, SphericalToCartesian
+from roman_datamodels import datamodels as rdm
+from stcal.tweakreg.astrometric_utils import get_catalog
+
+from romancal.tweakreg.tweakreg_step import DEFAULT_ABS_REFCAT
+
+
+@pytest.fixture(scope="session")
+def gaia_coords():
+    gaia_cat = get_catalog(
+        right_ascension=270,
+        declination=66,
+        search_radius=100 / 3600.0,
+        catalog=DEFAULT_ABS_REFCAT,
+        timeout=120,
+    )
+    return [(ra, dec) for ra, dec in zip(gaia_cat["ra"], gaia_cat["dec"], strict=False)]
+
+
+@pytest.fixture
+def tweakreg_image(tmp_path, gaia_coords):
+    """
+    Create a base image with a realistic WCSInfo and a WCS.
+
+    Notes
+    -----
+    The size of the image needs to be relatively large in order for
+    the source catalog step to find a reasonable number of sources in the image.
+
+    shift_1 and shift_2 (units in pixel) are used to shift the WCS projection plane.
+    """
+
+    def _tweakreg_image(shift_1=0, shift_2=0, catalog_filename="catalog.ecsv"):
+        l2 = rdm.ImageModel.create_fake_data(shape=(2000, 2000))
+        l2.meta.filename = "none"
+        l2.meta.cal_step = {}
+        for step_name in l2.schema_info("required")["roman"]["meta"]["cal_step"][
+            "required"
+        ].info:
+            l2.meta.cal_step[step_name] = "INCOMPLETE"
+        l2.meta.cal_logs = []
+
+        # to match GAIADR3 epoch
+        l2.meta.exposure.start_time = Time("2016-01-01T00:00:00")
+
+        # update wcsinfo
+        l2.meta.wcsinfo.v2_ref = 0.42955128282521254
+        l2.meta.wcsinfo.v3_ref = -0.2479976768255853
+        l2.meta.wcsinfo.vparity = -1
+        l2.meta.wcsinfo.v3yangle = -999999
+        l2.meta.wcsinfo.ra_ref = 270.0
+        l2.meta.wcsinfo.dec_ref = 66.0
+        l2.meta.wcsinfo.roll_ref = 60
+        l2.meta.wcsinfo.s_region = (
+            "POLYGON ICRS "
+            "269.3318903230621 65.56866666048172 "
+            "269.32578768154605 65.69246311613287 "
+            "269.02457173246125 65.69201346248587 "
+            "269.0333096074621 65.56870823657276 "
+        )
+
+        # create necessary transformations
+        distortion = Shift(-shift_1) & Shift(-shift_2)
+        v2_ref = l2.meta.wcsinfo.v2_ref / 3600
+        v3_ref = l2.meta.wcsinfo.v3_ref / 3600
+        roll_ref = l2.meta.wcsinfo.roll_ref
+        ra_ref = l2.meta.wcsinfo.ra_ref
+        dec_ref = l2.meta.wcsinfo.dec_ref
+
+        angles = np.array([v2_ref, -v3_ref, roll_ref, dec_ref, -ra_ref])
+        axes = "zyxyz"
+        rot = RotationSequence3D(angles, axes_order=axes)
+
+        # The sky rotation expects values in deg.
+        tel2sky = (
+            (Scale(0.1 / 3600) & Scale(0.1 / 3600))
+            | SphericalToCartesian(wrap_lon_at=180)
+            | rot
+            | CartesianToSpherical(wrap_lon_at=360)
+        )
+        tel2sky.name = "v23tosky"
+
+        # create required frames
+        detector = cf.Frame2D(name="detector", axes_order=(0, 1), unit=(u.pix, u.pix))
+        v2v3 = cf.Frame2D(
+            name="v2v3",
+            axes_order=(0, 1),
+            axes_names=("v2", "v3"),
+            unit=(u.arcsec, u.arcsec),
+        )
+        world = cf.CelestialFrame(reference_frame=coord.ICRS(), name="world")
+
+        # create pipeline
+        pipeline = [
+            wcs.Step(detector, distortion),
+            wcs.Step(v2v3, tel2sky),
+            wcs.Step(world, None),
+        ]
+
+        wcs_obj = wcs.WCS(pipeline)
+        wcs_obj.bounding_box = (
+            (-0.5, l2.data.shape[-2] + 0.5),
+            (-0.5, l2.data.shape[-1] + 0.5),
+        )
+
+        l2.meta.wcs = wcs_obj
+
+        catalog_path = tmp_path / catalog_filename
+        catalog_data = np.array(
+            [wcs_obj.world_to_pixel_values(ra, dec) for ra, dec in gaia_coords]
+        )
+        Table(catalog_data, names=("x", "y")).write(catalog_path, format="ascii.ecsv")
+        l2.meta["source_catalog"] = {
+            "tweakreg_catalog_name": str(catalog_path),
+        }
+
+        return l2
+
+    return _tweakreg_image

romancal/tweakreg/tests/test_parse_catfile.py

@@ -0,0 +1,56 @@
+import pytest
+
+from romancal.tweakreg.tweakreg_step import _parse_catfile
+
+
+def test_parse_catfile_raises_error_on_invalid_content(tmp_path):
+    """
+    Test that _parse_catfile raises an error if catfile contains more than
+    two columns (i.e. image name and corresponding catalog path).
+    """
+    # create custom catalog file and input datamodels
+    catfile = str(tmp_path / "catfile.txt")
+    with open(catfile, "w") as f:
+        f.write("img1.asdf column1 column2 column3")
+
+    with pytest.raises(ValueError):
+        _parse_catfile(catfile)
+
+
+def test_parse_catfile_valid_catalog(tmp_path):
+    """
+    Test that _parse_catfile can parse a custom catalog with valid format.
+    """
+    catfile = str(tmp_path / "catfile.txt")
+    with open(catfile, mode="w") as f:
+        f.write("img1.asdf cat1\nimg2.asdf cat2")
+    catdict = _parse_catfile(catfile)
+
+    assert catdict == {
+        "img1.asdf": str(tmp_path / "cat1"),
+        "img2.asdf": str(tmp_path / "cat2"),
+    }
+
+
+@pytest.mark.parametrize("catfile", (None, ""))
+def test_parse_catfile_returns_none(catfile):
+    """
+    Test that _parse_catfile returns None when catfile = None or catfile = "".
+    """
+    catdict = _parse_catfile(catfile=catfile)
+
+    assert catdict is None
+
+
+def test_parse_catfile_returns_none_on_invalid_content(tmp_path):
+    """
+    Test that _parse_catfile returns a dict where all the values are None
+    if only filename is present in catfile (i.e. no associated catalog).
+    """
+    catfile = str(tmp_path / "catfile.txt")
+    with open(catfile, mode="w") as f:
+        f.write("img1.asdf\nimg2.asdf\nimg3.asdf")
+
+    catdict = _parse_catfile(catfile)
+
+    assert not all(catdict.values())

romancal/tweakreg/tests/test_prepare_input.py

@@ -0,0 +1,57 @@
+import pytest
+
+from romancal.datamodels import ModelLibrary
+from romancal.tweakreg.tweakreg_step import TweakRegStep
+
+
+@pytest.fixture()
+def datamodel(tweakreg_image):
+    model = tweakreg_image()
+    model.meta.filename = "test.asdf"
+    return model
+
+
+@pytest.fixture()
+def datamodel_filename(datamodel, tmp_path):
+    fn = tmp_path / "test.asdf"
+    datamodel.save(fn)
+    return fn
+
+
+@pytest.fixture()
+def list_of_models(datamodel):
+    return [datamodel]
+
+
+@pytest.fixture()
+def list_of_filenames(datamodel_filename):
+    return [datamodel_filename]
+
+
+@pytest.fixture()
+def library(list_of_models):
+    return ModelLibrary(list_of_models)
+
+
+@pytest.fixture()
+def association(library, tmp_path):
+    fn = tmp_path / "asn.json"
+    library._save(fn.parent)
+    return fn
+
+
+@pytest.mark.parametrize(
+    "init",
+    [
+        "datamodel",
+        "datamodel_filename",
+        "list_of_models",
+        "list_of_filenames",
+        "library",
+        "association",
+    ],
+)
+def test_tweakreg_handle_input(init, request):
+    assert isinstance(
+        TweakRegStep()._prepare_input(request.getfixturevalue(init)), ModelLibrary
+    )