Added round-tripping test via Dask class and fix a number of issues that came up

Author: astrofrog

reproject/hips/_dask_array.py

@@ -0,0 +1,133 @@
+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.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
+
+__all__ = ['hips_as_dask_and_wcs']
+
+
+class HiPSArray:
+
+    def __init__(self, directory_or_url, level=None):
+
+        self._directory_or_url = directory_or_url
+
+        self._is_url = is_url(directory_or_url)
+
+        self._properties = load_properties(directory_or_url)
+
+        self._tile_width = int(self._properties["hips_tile_width"])
+        self._order = int(self._properties["hips_order"])
+        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"]
+        self._frame = VALID_COORD_SYSTEM[self._frame_str]
+
+        self._hp = HEALPix(nside=level_to_nside(self._level), frame=self._frame, order="nested")
+
+        self._header = map_header(level=self._level, frame=self._frame, tile_size=self._tile_width)
+
+        self.wcs = WCS(self._header)
+        self.shape = self.wcs.array_shape
+
+        self.dtype = float
+        self.ndim = 2
+
+        self.chunksize = (self._tile_width, self._tile_width)
+
+        self._nan = np.nan * np.ones(self.chunksize, dtype=self.dtype)
+
+        self._blank = np.broadcast_to(np.nan, self.shape)
+
+    def __getitem__(self, item):
+
+        if item[0].start == item[0].stop or item[1].start == item[1].stop:
+            return self._blank[item]
+
+        # For now assume item is a list of slices. Find
+
+        # imid = (item[0].start + item[0].stop) // 2
+        # jmid = (item[1].start + item[1].stop) // 2
+
+        istart = item[0].start
+        irange = item[0].stop - item[0].start
+        imid = np.array([istart + 0.25 * irange, istart + 0.75 * irange])
+
+        jstart = item[1].start
+        jrange = item[1].stop - item[1].start
+        jmid = np.array([jstart + 0.25 * jrange, jstart + 0.75 * jrange])
+
+        # Convert pixel coordinates to HEALPix indices
+
+        coord = self.wcs.pixel_to_world(jmid, imid)
+
+        if self._frame_str == "equatorial":
+            lon, lat = coord.ra.deg, coord.dec.deg
+        elif self._frame_str == "galactic":
+            lon, lat = coord.l.deg, coord.b.deg
+        else:
+            raise NotImplementedError()
+
+        if np.all(np.isnan(lon) | np.isnan(lat)):
+            return self._nan
+
+        index = self._hp.skycoord_to_healpix(coord)
+
+        if np.all(index == -1):
+            return self._nan
+
+        index = np.max(index)
+
+        return self._get_tile(level=self._level, index=index)
+
+    @functools.lru_cache(maxsize=128)
+    def _get_tile(self, *, level, index):
+
+        filename_or_url = tile_filename(
+            level=self._level,
+            index=index,
+            output_directory=self._directory_or_url,
+            extension="fits",
+        )
+
+        if self._is_url:
+            try:
+                filename = download_file(filename_or_url, cache=True)
+            except urllib.error.HTTPError:
+                return self._nan
+        elif not os.path.exists(filename_or_url):
+            return self._nan
+        else:
+            filename = filename_or_url
+
+        with fits.open(filename) as hdulist:
+            hdu = hdulist[0]
+            # data = hdu.data[::-1]
+            data = hdu.data
+
+            # return np.ones(hdu.data.shape) * index
+
+        return data
+
+
+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

reproject/hips/high_level.py

@@ -10,7 +10,6 @@
 import numpy as np
 from astropy.coordinates import ICRS, BarycentricTrueEcliptic, Galactic
 from astropy.io import fits
-from astropy.nddata import block_reduce
 from astropy_healpix import (
     HEALPix,
     level_to_nside,
@@ -204,6 +203,8 @@ def reproject_to_hips(
     # Determine center of image and radius to furthest corner, to determine
     # which HiPS tiles need to be generated
 
+    # TODO: this will fail for e.g. allsky maps
+
     ny, nx = array_in.shape[-2:]
 
     cen_x, cen_y = (nx - 1) / 2, (ny - 1) / 2
@@ -216,9 +217,6 @@ def reproject_to_hips(
 
     radius = cor_world.separation(cen_world).max() * 2
 
-    print(cen_world)
-    print(cor_world)
-
     # TODO: in future if astropy-healpix implements polygon searches, we could
     # use that instead
 
@@ -227,9 +225,9 @@ def reproject_to_hips(
     nside = level_to_nside(level)
     hp = HEALPix(nside=nside, order="nested", frame=frame)
 
-    indices = hp.cone_search_skycoord(cen_world, radius=radius)
+    # indices = hp.cone_search_skycoord(cen_world, radius=radius)
 
-    print(indices)
+    indices = np.arange(hp.npix)
 
     logger.info(f"Found {len(indices)} tiles (at most) to generate at level {level}")
 
@@ -241,12 +239,28 @@ def reproject_to_hips(
 
     # Iterate over the tiles and generate them
     def process(index):
-        header = tile_header(level=level, index=index, frame=frame, tile_size=tile_size)
         if hasattr(wcs_in, "deepcopy"):
             wcs_in_copy = wcs_in.deepcopy()
         else:
             wcs_in_copy = deepcopy(wcs_in)
-        array_out, footprint = reproject_function((array_in, wcs_in_copy), header, **kwargs)
+
+        header = tile_header(level=level, index=index, frame=frame, tile_size=tile_size)
+
+        if isinstance(header, tuple):
+            array_out1, footprint1 = reproject_function(
+                (array_in, wcs_in_copy), header[0], **kwargs
+            )
+            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
+            header = header[0]
+        else:
+            array_out, footprint = reproject_function((array_in, wcs_in_copy), header, **kwargs)
+
         if tile_format != "png":
             array_out[np.isnan(array_out)] = 0.0
         if np.all(footprint == 0):
@@ -260,6 +274,7 @@ def process(index):
                     extension=EXTENSION[tile_format],
                 ),
                 array_out,
+                header,
             )
         else:
             if tile_format == "png":
@@ -294,76 +309,76 @@ def process(index):
 
     indices = np.array(generated_indices)
 
-    # Iterate over higher levels and compute lower resolution tiles
-    for ilevel in range(level - 1, -1, -1):
-
-        # Find index of tiles to produce at lower-resolution levels
-        indices = np.sort(np.unique(indices // 4))
-
-        make_tile_folders(level=ilevel, indices=indices, output_directory=output_directory)
-
-        for index in indices:
-
-            header = tile_header(level=ilevel, index=index, frame=frame, tile_size=tile_size)
-
-            if tile_format == "fits":
-                array = np.zeros((tile_size, tile_size))
-            elif tile_format == "png":
-                array = np.zeros((tile_size, tile_size, 4), dtype=np.uint8)
-            else:
-                array = np.zeros((tile_size, tile_size, 3), dtype=np.uint8)
-
-            for subindex in range(4):
-
-                current_index = 4 * index + subindex
-                subtile_filename = tile_filename(
-                    level=ilevel + 1,
-                    index=current_index,
-                    output_directory=output_directory,
-                    extension=EXTENSION[tile_format],
-                )
-
-                if os.path.exists(subtile_filename):
-
-                    if tile_format == "fits":
-                        data = block_reduce(fits.getdata(subtile_filename), 2, func=np.mean)
-                    else:
-                        data = block_reduce(
-                            np.array(Image.open(subtile_filename))[::-1], (2, 2, 1), func=np.mean
-                        )
-
-                    if subindex == 0:
-                        array[256:, :256] = data
-                    elif subindex == 2:
-                        array[256:, 256:] = data
-                    elif subindex == 1:
-                        array[:256, :256] = data
-                    elif subindex == 3:
-                        array[:256, 256:] = data
-
-            if tile_format == "fits":
-                fits.writeto(
-                    tile_filename(
-                        level=ilevel,
-                        index=index,
-                        output_directory=output_directory,
-                        extension=EXTENSION[tile_format],
-                    ),
-                    array,
-                    header,
-                )
-            else:
-                image = as_transparent_rgb(array.transpose(2, 0, 1))
-                if tile_format == "jpeg":
-                    image = image.convert("RGB")
-                image.save(
-                    tile_filename(
-                        level=ilevel,
-                        index=index,
-                        output_directory=output_directory,
-                        extension=EXTENSION[tile_format],
-                    )
-                )
+    # # Iterate over higher levels and compute lower resolution tiles
+    # for ilevel in range(level - 1, -1, -1):
+
+    #     # Find index of tiles to produce at lower-resolution levels
+    #     indices = np.sort(np.unique(indices // 4))
+
+    #     make_tile_folders(level=ilevel, indices=indices, output_directory=output_directory)
+
+    #     for index in indices:
+
+    #         header = tile_header(level=ilevel, index=index, frame=frame, tile_size=tile_size)
+
+    #         if tile_format == "fits":
+    #             array = np.zeros((tile_size, tile_size))
+    #         elif tile_format == "png":
+    #             array = np.zeros((tile_size, tile_size, 4), dtype=np.uint8)
+    #         else:
+    #             array = np.zeros((tile_size, tile_size, 3), dtype=np.uint8)
+
+    #         for subindex in range(4):
+
+    #             current_index = 4 * index + subindex
+    #             subtile_filename = tile_filename(
+    #                 level=ilevel + 1,
+    #                 index=current_index,
+    #                 output_directory=output_directory,
+    #                 extension=EXTENSION[tile_format],
+    #             )
+
+    #             if os.path.exists(subtile_filename):
+
+    #                 if tile_format == "fits":
+    #                     data = block_reduce(fits.getdata(subtile_filename), 2, func=np.mean)
+    #                 else:
+    #                     data = block_reduce(
+    #                         np.array(Image.open(subtile_filename))[::-1], (2, 2, 1), func=np.mean
+    #                     )
+
+    #                 if subindex == 0:
+    #                     array[256:, :256] = data
+    #                 elif subindex == 2:
+    #                     array[256:, 256:] = data
+    #                 elif subindex == 1:
+    #                     array[:256, :256] = data
+    #                 elif subindex == 3:
+    #                     array[:256, 256:] = data
+
+    #         if tile_format == "fits":
+    #             fits.writeto(
+    #                 tile_filename(
+    #                     level=ilevel,
+    #                     index=index,
+    #                     output_directory=output_directory,
+    #                     extension=EXTENSION[tile_format],
+    #                 ),
+    #                 array,
+    #                 header,
+    #             )
+    #         else:
+    #             image = as_transparent_rgb(array.transpose(2, 0, 1))
+    #             if tile_format == "jpeg":
+    #                 image = image.convert("RGB")
+    #             image.save(
+    #                 tile_filename(
+    #                     level=ilevel,
+    #                     index=index,
+    #                     output_directory=output_directory,
+    #                     extension=EXTENSION[tile_format],
+    #                 )
+    #             )
 
     # Generate properties file