Merge pull request #18 from lincc-frameworks/issue/15/todos

olivialynn · web-flow · commit f99287fd97d6 · 2025-08-15T00:58:23.000-07:00
Address various todos from comments; add unit tests for skymap writer
diff --git a/src/skymap_convert/skymap_readers.py b/src/skymap_convert/skymap_readers.py
@@ -31,7 +31,8 @@ class ConvertedSkymapReader:
     file_path : Path
         Path to the skymap directory.
     safe_loading : bool
-        Whether to verify polygon non-degeneracy when loading vertices.
+        Whether to verify polygon non-degeneracy when loading vertices. Additionally, if true,
+        checks that the metadata matches the array shapes.
     """
 
     def __init__(self, file_path: str | Path = None, safe_loading: bool = False, preset: str = None):
@@ -47,6 +48,15 @@ def __init__(self, file_path: str | Path = None, safe_loading: bool = False, pre
         preset : str, optional
             Name of a built-in skymap preset to load. If specified, file_path is ignored.
             Available presets can be listed with skymap_convert.presets.list_available_presets().
+
+        Raises
+        ------
+        ValueError
+            If neither file_path nor preset is provided, or if the specified preset does not exist.
+        FileNotFoundError
+            If the specified file_path does not exist or is not a directory.
+        AssertionError
+            If safe_loading is True and the metadata does not match the array shapes.
         """
         # Use preset if provided, otherwise check for file_path
         if preset is not None:
@@ -87,7 +97,12 @@ def __init__(self, file_path: str | Path = None, safe_loading: bool = False, pre
         self.tracts = np.load(self.tracts_path, mmap_mode="r")
         self.patches = np.load(self.patches_path, mmap_mode="r")
 
-        # TODO could be nice to check if the metadata matches the arrays here.
+        # Check if the metadata matches the arrays
+        if self.safe_loading:
+            assert self.n_tracts == self.tracts.shape[0], "Metadata n_tracts does not match array shape"
+            assert (
+                self.n_patches_per_tract == self.patches.shape[1]
+            ), "Metadata n_patches_per_tract does not match array shape"
 
     def _decompress_patches_gz(self) -> Path:
         """Decompress patches.npy.gz to a temporary file if not already done.
diff --git a/src/skymap_convert/skymap_writers.py b/src/skymap_convert/skymap_writers.py
@@ -46,8 +46,9 @@ def write(self, skymap, output_path: str | Path, skymap_name: str = "converted_s
 
         Notes
         -----
-        The method assumes 100 patches per tract, which is standard for LSST skymaps.
-        All polygon vertices are stored as [RA, Dec] coordinates in degrees.
+        - The method assumes all tracts have the same number of patches, which is standard for LSST
+          skymaps.
+        - All polygon vertices are stored as [RA, Dec] coordinates in degrees.
 
         Examples
         --------
@@ -59,8 +60,24 @@ def write(self, skymap, output_path: str | Path, skymap_name: str = "converted_s
         output_path = Path(output_path)
         self._ensure_output_directory(output_path)
 
-        n_tracts = len(skymap)  # TODO - is this accounting for poles?
-        n_patches = 100  # fixed per tract # TODO - would be nice to make this dynamic, though we expect 100
+        n_tracts = skymap._numTracts
+
+        # Get number of patches from the first tract
+        first_tract = skymap[0]
+        n_patches_0 = (
+            first_tract._tractBuilder._numPatches.x * first_tract._tractBuilder._numPatches.y
+        )  #  Number of patches per tract, typically 100.
+        # Get the number of patches from the second tract
+        second_tract = skymap[1]
+        n_patches_1 = second_tract._tractBuilder._numPatches.x * second_tract._tractBuilder._numPatches.y
+        # Check that they are the same
+        if n_patches_0 != n_patches_1:
+            raise ValueError(
+                f"Number of patches per tract mismatch: {n_patches_0} != {n_patches_1}. "
+                "Ensure all tracts have the same number of patches."
+            )
+        n_patches = n_patches_0  # Use the first tract's patch count
+
         tract_array = np.zeros((n_tracts, 4, 2), dtype=np.float64)
         patch_array = np.zeros((n_tracts, n_patches, 4, 2), dtype=np.float64)
 
diff --git a/src/skymap_convert/utils.py b/src/skymap_convert/utils.py
@@ -118,66 +118,6 @@ def radians_to_degrees(radians):
     return radians * (180.0 / math.pi)
 
 
-class IterateTractAndRing:
-    """An iterator to traverse through tract IDs and ring IDs in a skymap.
-
-    This iterator yields tuples of (tract_index, ring_index) for each tract in the skymap.
-    The first tract is the south pole (ring -1), and the last tract is the north pole.
-
-    Parameters
-    ----------
-    ring_nums : list of int
-        A list where each element represents the number of tracts in each ring.
-        For example, [5, 10, 15] represents 5 tracts in ring 0, 10 in ring 1,
-        and 15 in ring 2.
-    add_poles : bool, optional
-        If True, include the south pole (tract 0) and north pole (last tract)
-        in the iteration. If False, only iterate through the rings (default: True).
-
-    Examples
-    --------
-    >>> iterator = IterateTractAndRing([5, 10, 15], add_poles=True)
-    >>> for tract_id, ring_id in iterator:
-    ...     print(f"Tract {tract_id} is in ring {ring_id}")
-    Tract 0 is in ring -1
-    Tract 1 is in ring 0
-    ...
-
-    TODO : this is more or less deprecated, I belive. Consider removing.
-    """
-
-    def __init__(self, ring_nums, add_poles=True):
-        self.ring_nums = ring_nums
-        if add_poles:
-            self.total_tracts = sum(ring_nums) + 2
-            self.current_tract = 0
-            self.current_ring = -1
-        else:
-            self.total_tracts = sum(ring_nums)
-            self.current_tract = 1
-            self.current_ring = 0
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        # End iteration if we have processed all tracts.
-        if self.current_tract >= self.total_tracts:
-            raise StopIteration
-        tract_and_ring = (self.current_tract, self.current_ring)
-
-        # Increase tract.
-        self.current_tract += 1
-
-        # Check if we need to move to the next ring.clear
-        if self.current_ring == -1:
-            self.current_ring += 1
-        elif self.current_tract > sum(self.ring_nums[: self.current_ring + 1]):
-            self.current_ring += 1
-
-        return tract_and_ring
-
-
 def get_poly_from_tract_id(skymap, tract_id, inner=False):
     """Get the ConvexPolygon for a tract by its ID.
 
diff --git a/tests/skymap_convert/conftest.py b/tests/skymap_convert/conftest.py
@@ -52,9 +52,15 @@ def converted_skymap_reader():
     Note
     ----
     This uses the pre-converted LSST skymap for performance reasons.
+    The reader will be automatically cleaned up at the end of the test session.
     """
+    reader = ConvertedSkymapReader(preset="lsst_skymap")
 
-    return ConvertedSkymapReader(preset="lsst_skymap")
+    # Yield the reader for use in tests
+    yield reader
+
+    # Cleanup after all tests are done
+    reader.cleanup()
 
 
 @pytest.fixture(scope="session")
@@ -74,3 +80,59 @@ def lsst_skymap():
     """
     pytest.importorskip("lsst.skymap")
     return load_pickle_skymap(RAW_SKYMAP_DIR / "skyMap_lsst_cells_v1_skymaps.pickle")
+
+
+@pytest.fixture(scope="session")
+def written_skymap_data(lsst_skymap, tmp_path_factory, request):
+    """Session-scoped fixture that writes skymap data once and provides paths for testing.
+
+    Parameters
+    ----------
+    lsst_skymap : lsst.skymap.SkyMap
+        The original LSST skymap object to write.
+    tmp_path_factory : pytest.TempPathFactory
+        Pytest factory for creating temporary directories.
+    request : pytest.FixtureRequest
+        Pytest request object to access command line options.
+
+    Returns
+    -------
+    dict
+        Dictionary containing:
+        - 'output_dir': Path to the directory containing written skymap files
+        - 'skymap_name': Name used for the skymap
+        - 'reader': ConvertedSkymapReader instance for the written data
+
+    Notes
+    -----
+    This fixture only runs if --longrun is specified, otherwise it skips.
+    The expensive skymap writing operation happens only once per test session.
+    The reader will be automatically cleaned up at the end of the test session.
+    """
+    # Skip if not running longrun tests
+    if not request.config.option.longrun:
+        pytest.skip("Skipping written_skymap_data fixture - requires --longrun")
+
+    pytest.importorskip("lsst.skymap")
+    pytest.importorskip("lsst.sphgeom")
+
+    from skymap_convert.skymap_writers import ConvertedSkymapWriter
+
+    # Create session-scoped temporary directory
+    output_dir = tmp_path_factory.mktemp("session_skymap_data")
+    skymap_name = "test_session_skymap"
+
+    # Write the skymap once
+    writer = ConvertedSkymapWriter()
+    writer.write(lsst_skymap, output_dir, skymap_name)
+
+    # Create reader for the written data
+    reader = ConvertedSkymapReader(output_dir)
+
+    data = {"output_dir": output_dir, "skymap_name": skymap_name, "reader": reader}
+
+    # Yield the data for use in tests
+    yield data
+
+    # Cleanup after all tests are done
+    reader.cleanup()
diff --git a/tests/skymap_convert/test_skymap_reader.py b/tests/skymap_convert/test_skymap_reader.py
diff --git a/tests/skymap_convert/test_skymap_writer.py b/tests/skymap_convert/test_skymap_writer.py
@@ -0,0 +1,149 @@
+import pytest
+import yaml
+from skymap_convert.utils import (
+    get_quad_from_patch_id,
+    get_quad_from_tract_id,
+    quads_are_equiv,
+)
+from tqdm import tqdm
+
+TRACT_SAMPLES = [1, 250, 1899]
+PATCH_SAMPLES = [0, 42, 99]
+
+
+@pytest.mark.longrun
+def test_writer_creates_expected_files(lsst_skymap, written_skymap_data):
+    """Test that ConvertedSkymapWriter creates all expected output files.
+
+    Parameters
+    ----------
+    lsst_skymap : lsst.skymap.SkyMap
+        The original LSST skymap object to write.
+    written_skymap_data : dict
+        Fixture providing written skymap data with output_dir, skymap_name, and reader.
+
+    Notes
+    -----
+    Verifies that:
+    - All expected files (metadata.yaml, tracts.npy, patches.npy.gz) are created
+    - No temporary files are left behind
+    - Files have non-zero size
+    """
+    pytest.importorskip("lsst.skymap")
+    pytest.importorskip("lsst.sphgeom")
+
+    output_dir = written_skymap_data["output_dir"]
+    skymap_name = written_skymap_data["skymap_name"]
+
+    # Check expected files exist
+    assert (output_dir / "metadata.yaml").exists()
+    assert (output_dir / "tracts.npy").exists()
+    assert (output_dir / "patches.npy.gz").exists()
+
+    # Check temporary files are cleaned up
+    assert not (output_dir / "patches.npy").exists()
+
+    # Check files have content
+    assert (output_dir / "metadata.yaml").stat().st_size > 0
+    assert (output_dir / "tracts.npy").stat().st_size > 0
+    assert (output_dir / "patches.npy.gz").stat().st_size > 0
+
+    # Load and check metadata
+    with open(output_dir / "metadata.yaml", "r") as f:
+        metadata = yaml.safe_load(f)
+
+    assert metadata["name"] == skymap_name
+    assert metadata["n_tracts"] == len(lsst_skymap)
+    assert metadata["n_patches_per_tract"] == 100
+    assert metadata["format_version"] == 1
+    assert "generated" in metadata
+    assert metadata["generated"].endswith("Z")  # ISO format with UTC
+
+
+@pytest.mark.longrun
+def test_writer_sample_tracts(lsst_skymap, written_skymap_data):
+    """Test complete round-trip: write with ConvertedSkymapWriter, read with ConvertedSkymapReader.
+
+    Parameters
+    ----------
+    lsst_skymap : lsst.skymap.SkyMap
+        The original LSST skymap object.
+    written_skymap_data : dict
+        Fixture providing written skymap data with output_dir, skymap_name, and reader.
+
+    Notes
+    -----
+    Verifies that data written by ConvertedSkymapWriter can be correctly
+    read by ConvertedSkymapReader and produces identical results.
+    """
+    pytest.importorskip("lsst.skymap")
+    pytest.importorskip("lsst.sphgeom")
+
+    reader = written_skymap_data["reader"]
+    skymap_name = written_skymap_data["skymap_name"]
+
+    # Verify metadata matches
+    assert reader.metadata["name"] == skymap_name
+    assert reader.n_tracts == len(lsst_skymap)
+    assert reader.n_patches_per_tract == 100
+
+    # Verify a few sample geometries match original
+    for tract_id in TRACT_SAMPLES:
+        original_tract = get_quad_from_tract_id(lsst_skymap, tract_id, inner=True)
+        round_trip_tract = reader.get_tract_vertices(tract_id)
+
+        assert quads_are_equiv(original_tract, round_trip_tract), f"Round-trip failed for tract {tract_id}"
+
+        for patch_id in PATCH_SAMPLES:
+            original_patch = get_quad_from_patch_id(lsst_skymap, tract_id, patch_id)
+            round_trip_patch = reader.get_patch_vertices(tract_id, patch_id)
+
+            assert quads_are_equiv(
+                original_patch, round_trip_patch
+            ), f"Round-trip failed for patch {patch_id} in tract {tract_id}"
+
+
+@pytest.mark.longrun
+def test_writer_full_skymap_integrity(lsst_skymap, written_skymap_data):
+    """Comprehensive test that writer preserves all tract and patch geometries.
+
+    Parameters
+    ----------
+    lsst_skymap : lsst.skymap.SkyMap
+        The original LSST skymap object.
+    written_skymap_data : dict
+        Fixture providing written skymap data with output_dir, skymap_name, and reader.
+
+    Notes
+    -----
+    This is a comprehensive test that verifies every tract and patch geometry
+    is preserved correctly using the session-scoped written skymap data.
+    """
+    if True:
+        pytest.skip(
+            "Skipping full integrity test as it takes an exceptionally long time to run. "
+            "Recommend running periodically, especially after major changes to the writer. "
+            "To run, manually remove the branching logic in the code of the test."
+        )
+    else:
+        pytest.importorskip("lsst.skymap")
+        pytest.importorskip("lsst.sphgeom")
+
+        reader = written_skymap_data["reader"]
+
+        tract_ids = range(len(lsst_skymap))
+        for tract_id in tqdm(tract_ids, desc="Verifying written tracts", leave=False):
+            # Verify tract geometry
+            truth_quad = get_quad_from_tract_id(lsst_skymap, tract_id, inner=True)
+            written_quad = reader.get_tract_vertices(tract_id)
+
+            assert quads_are_equiv(truth_quad, written_quad), f"Tract {tract_id} geometry not preserved"
+
+            # Verify all patch geometries
+            for patch_id in range(100):
+                truth_patch = get_quad_from_patch_id(lsst_skymap, tract_id, patch_id)
+                written_patch = reader.get_patch_vertices(tract_id, patch_id)
+
+                assert quads_are_equiv(
+                    truth_patch, written_patch
+                ), f"Patch {patch_id} in tract {tract_id} geometry not preserved"
