Write point map with cdshealpix skymap (#409)

* Write point map with cdshealpix skymap

* Add missing mypy types

* Add temp file to handle cloud catalogs

* Update src/hats/io/file_io/file_io.py

Co-authored-by: Melissa DeLucchi <[email protected]>

---------

Co-authored-by: Melissa DeLucchi <[email protected]>

src/hats/io/file_io/file_io.py

@@ -10,10 +10,10 @@
 import pyarrow.dataset as pds
 import pyarrow.parquet as pq
 import yaml
+from cdshealpix.skymap.skymap import Skymap
 from pyarrow.dataset import Dataset
 from upath import UPath
 
-import hats.pixel_math.healpix_shim as hp
 from hats.io.file_io.file_pointer import get_upath
 
 
@@ -213,11 +213,12 @@ def write_parquet_metadata(
     )
 
 
-def read_fits_image(map_file_pointer: str | Path | UPath):
+def read_fits_image(map_file_pointer: str | Path | UPath) -> np.ndarray:
     """Read the object spatial distribution information from a healpix FITS file.
 
     Args:
-        file_pointer: location of file to be written
+        map_file_pointer (path-like): location of file to be read
+
     Returns:
         one-dimensional numpy array of integers where the
         value at each index corresponds to the number of objects found at the healpix pixel.
@@ -227,7 +228,7 @@ def read_fits_image(map_file_pointer: str | Path | UPath):
         with map_file_pointer.open("rb") as _map_file:
             map_data = _map_file.read()
             _tmp_file.write(map_data)
-            return hp.read_map(_tmp_file.name, nest=True)
+            return Skymap.from_fits(_tmp_file.name).values
 
 
 def write_fits_image(histogram: np.ndarray, map_file_pointer: str | Path | UPath):
@@ -236,12 +237,13 @@ def write_fits_image(histogram: np.ndarray, map_file_pointer: str | Path | UPath
     Args:
         histogram (:obj:`np.ndarray`): one-dimensional numpy array of long integers where the
             value at each index corresponds to the number of objects found at the healpix pixel.
-        file_pointer: location of file to be written
+        map_file_pointer (path-like): location of file to be written
     """
     map_file_pointer = get_upath(map_file_pointer)
     with tempfile.NamedTemporaryFile() as _tmp_file:
         with map_file_pointer.open("wb") as _map_file:
-            hp.write_map(_tmp_file.name, histogram, overwrite=True, dtype=np.int32, nest=True, coord="CEL")
+            skymap = Skymap.from_array(histogram)
+            skymap.to_fits(_tmp_file.name)
             _map_file.write(_tmp_file.read())
 
 

tests/hats/io/file_io/test_file_io.py

@@ -2,15 +2,18 @@
 import pandas as pd
 import pytest
 
+from hats.io import paths
 from hats.io.file_io import (
     delete_file,
     load_csv_to_pandas,
     load_csv_to_pandas_generator,
     make_directory,
+    read_fits_image,
     read_parquet_dataset,
     read_parquet_file_to_pandas,
     remove_directory,
     write_dataframe_to_csv,
+    write_fits_image,
     write_string_to_file,
 )
 from hats.io.file_io.file_pointer import does_file_or_directory_exist
@@ -123,3 +126,12 @@ def test_read_parquet_dataset(small_sky_dir, small_sky_order1_dir):
     )
 
     assert ds.count_rows() == 131
+
+
+def test_write_point_map_roundtrip(small_sky_order1_dir, tmp_path):
+    """Test the reading/writing of a catalog point map"""
+    expected_counts_skymap = read_fits_image(paths.get_point_map_file_pointer(small_sky_order1_dir))
+    output_map_pointer = paths.get_point_map_file_pointer(tmp_path)
+    write_fits_image(expected_counts_skymap, output_map_pointer)
+    counts_skymap = read_fits_image(output_map_pointer)
+    np.testing.assert_array_equal(counts_skymap, expected_counts_skymap)