Added comments and added a test for level validation

Author: astrofrog

reproject/hips/_dask_array.py

@@ -1,22 +1,19 @@
+import functools
 import os
-import struct
 import urllib
 import uuid
-import functools
 
 import numpy as np
-from astropy import units as u
 from astropy.io import fits
+from astropy.utils.data import download_file
 from astropy.wcs import WCS
 from astropy_healpix import HEALPix, level_to_nside
 from dask import array as da
-from astropy.utils.data import download_file
-from astropy.wcs.utils import celestial_frame_to_wcs
 
-from .utils import is_url, load_properties, tile_filename, tile_header, map_header
 from .high_level import VALID_COORD_SYSTEM
+from .utils import is_url, load_properties, map_header, tile_filename
 
-__all__ = ['hips_as_dask_and_wcs']
+__all__ = ["hips_as_dask_and_wcs"]
 
 
 class HiPSArray:
@@ -31,6 +28,17 @@ def __init__(self, directory_or_url, level=None):
 
         self._tile_width = int(self._properties["hips_tile_width"])
         self._order = int(self._properties["hips_order"])
+        if level is None:
+            self._level = self._order
+        else:
+            if level > self._order:
+                raise ValueError(
+                    f"HiPS dataset at {directory_or_url} does not contain level {level} data"
+                )
+            elif level < 0:
+                raise ValueError("level should be positive")
+            else:
+                self._level = int(level)
         self._level = self._order if level is None else level
         self._tile_format = self._properties["hips_tile_format"]
         self._frame_str = self._properties["hips_frame"]
@@ -81,8 +89,12 @@ def __getitem__(self, item):
         else:
             raise NotImplementedError()
 
-        if np.all(np.isnan(lon) | np.isnan(lat)):
+        invalid = np.isnan(lon) | np.isnan(lat)
+
+        if np.all(invalid):
             return self._nan
+        elif np.any(invalid):
+            coord = coord[~invalid]
 
         index = self._hp.skycoord_to_healpix(coord)
 
@@ -125,9 +137,12 @@ def _get_tile(self, *, level, index):
 
 def hips_as_dask_and_wcs(directory_or_url, *, level=None):
     array_wrapper = HiPSArray(directory_or_url, level=level)
-    return da.from_array(
-        array_wrapper,
-        chunks=array_wrapper.chunksize,
-        name=str(uuid.uuid4()),
-        meta=np.array([], dtype=float)
-    ), array_wrapper.wcs
+    return (
+        da.from_array(
+            array_wrapper,
+            chunks=array_wrapper.chunksize,
+            name=str(uuid.uuid4()),
+            meta=np.array([], dtype=float),
+        ),
+        array_wrapper.wcs,
+    )

reproject/hips/high_level.py

@@ -253,10 +253,11 @@ def process(index):
             array_out2, footprint2 = reproject_function(
                 (array_in, wcs_in_copy), header[1], **kwargs
             )
-            array_out = (
-                np.nan_to_num(array_out1) * footprint1 + np.nan_to_num(array_out2) * footprint2
-            ) / (footprint1 + footprint2)
-            footprint = (footprint1 + footprint2) / 2
+            with np.errstate(invalid="ignore"):
+                array_out = (
+                    np.nan_to_num(array_out1) * footprint1 + np.nan_to_num(array_out2) * footprint2
+                ) / (footprint1 + footprint2)
+                footprint = (footprint1 + footprint2) / 2
             header = header[0]
         else:
             array_out, footprint = reproject_function((array_in, wcs_in_copy), header, **kwargs)

reproject/hips/tests/test_dask_array.py

@@ -1,40 +1,50 @@
-import pytest
 import numpy as np
-
-from astropy.wcs import WCS
+import pytest
 from astropy.io import fits
+from astropy.utils.data import get_pkg_data_filename
+from astropy.wcs import WCS
 
 from reproject import reproject_interp
 from reproject.hips import reproject_to_hips
 from reproject.hips._dask_array import hips_as_dask_and_wcs
-from astropy.utils.data import get_pkg_data_filename
+
 
 class TestHIPSDaskArray:
 
     def setup_method(self):
-
-        hdu = fits.open(get_pkg_data_filename('allsky/allsky_rosat.fits'))[0]
-        self.original_header = hdu.header
+        # We use an all-sky WCS image as input since this will test all parts
+        # of the HiPS projection (some issues happen around boundaries for instance)
+        hdu = fits.open(get_pkg_data_filename("allsky/allsky_rosat.fits"))[0]
         self.original_wcs = WCS(hdu.header)
         self.original_array = hdu.data.size + np.arange(hdu.data.size).reshape(hdu.data.shape)
 
-    @pytest.mark.parametrize('frame', ('galactic', 'equatorial'))
-    @pytest.mark.parametrize('level',  (0, 1))
+    @pytest.mark.parametrize("frame", ("galactic", "equatorial"))
+    @pytest.mark.parametrize("level", (0, 1))
     def test_roundtrip(self, tmp_path, frame, level):
 
-        self.output_directory = tmp_path / 'roundtrip'
+        output_directory = tmp_path / "roundtrip"
 
+        # Note that we always use level=1 to generate, but use a variable level
+        # to construct the dask array - this is deliberate and ensure that the
+        # dask array has a proper separation of maximum and current level.
         reproject_to_hips(
             (self.original_array, self.original_wcs),
             coord_system_out=frame,
-            level=level,
+            level=1,
             reproject_function=reproject_interp,
-            output_directory=self.output_directory,
+            output_directory=output_directory,
+            tile_size=256,
         )
 
-        dask_array, wcs = hips_as_dask_and_wcs(self.output_directory, level=level)
+        # Represent the HiPS as a dask array
+        dask_array, wcs = hips_as_dask_and_wcs(output_directory, level=level)
 
-        final_array, footprint = reproject_interp((dask_array, wcs), self.original_wcs, shape_out=self.original_array.shape)
+        # Reproject back to the original WCS
+        final_array, footprint = reproject_interp(
+            (dask_array, wcs),
+            self.original_wcs,
+            shape_out=self.original_array.shape,
+        )
 
         # FIXME: Due to boundary effects and the fact there are NaN values in
         # the whole-map dask array, there are a few NaN pixels in the image in
@@ -47,10 +57,33 @@ def test_roundtrip(self, tmp_path, frame, level):
         # values.
 
         valid = ~np.isnan(final_array)
-
         assert np.sum(valid) > 90400
-
         np.testing.assert_allclose(final_array[valid], self.original_array[valid], rtol=0.01)
 
+    def test_level_validation(self, tmp_path):
+
+        output_directory = tmp_path / "levels"
+
+        reproject_to_hips(
+            (self.original_array, self.original_wcs),
+            coord_system_out="equatorial",
+            level=1,
+            reproject_function=reproject_interp,
+            output_directory=output_directory,
+            tile_size=32,
+        )
+
+        dask_array, wcs = hips_as_dask_and_wcs(output_directory, level=0)
+        assert dask_array.shape == (160, 160)
+
+        dask_array, wcs = hips_as_dask_and_wcs(output_directory, level=1)
+        assert dask_array.shape == (320, 320)
+
+        dask_array, wcs = hips_as_dask_and_wcs(output_directory)
+        assert dask_array.shape == (320, 320)
+
+        with pytest.raises(Exception, match=r"does not contain level 2 data"):
+            hips_as_dask_and_wcs(output_directory, level=2)
 
-    # VALIDATE LEVEL
+        with pytest.raises(Exception, match=r"should be positive"):
+            hips_as_dask_and_wcs(output_directory, level=-1)