From bccf40c583936b0e8a4098dfd526ed57f86322c9 Mon Sep 17 00:00:00 2001 From: Sandro Campos Date: Fri, 1 Nov 2024 11:07:32 -0400 Subject: [PATCH] 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 <113376043+delucchi-cmu@users.noreply.github.com> --------- Co-authored-by: Melissa DeLucchi <113376043+delucchi-cmu@users.noreply.github.com> --- src/hats/io/file_io/file_io.py | 14 ++++++++------ tests/hats/io/file_io/test_file_io.py | 12 ++++++++++++ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/src/hats/io/file_io/file_io.py b/src/hats/io/file_io/file_io.py index 67d3aa82..be46b105 100644 --- a/src/hats/io/file_io/file_io.py +++ b/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()) diff --git a/tests/hats/io/file_io/test_file_io.py b/tests/hats/io/file_io/test_file_io.py index 746596f2..f4119820 100644 --- a/tests/hats/io/file_io/test_file_io.py +++ b/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)