Generate Allsky files for HiPS

reproject/hips/_allsky.py

@@ -0,0 +1,95 @@
+import math
+import os
+import warnings
+
+import numpy as np
+from astropy.io import fits
+from astropy.nddata import block_reduce
+from PIL import Image
+
+from ._utils import tile_filename
+
+__all__ = ["save_allsky"]
+
+
+# Per the HiPS standard, Allsky preview files are only generated for the low
+# orders (0 to 3), and each tile is downsampled to a small size that must stay a
+# power of two (typically 64 pixels).
+ALLSKY_MAX_ORDER = 3
+ALLSKY_TILE_SIZE = 64
+
+
+def _reduction_factor(tile_size):
+    # Largest power-of-two factor that divides tile_size and keeps the
+    # downsampled tile at least ALLSKY_TILE_SIZE pixels across.
+    factor = 1
+    while tile_size // (factor * 2) >= ALLSKY_TILE_SIZE and tile_size % (factor * 2) == 0:
+        factor *= 2
+    return factor
+
+
+def save_allsky(*, output_directory, tile_format, extension, tile_size, spatial_level):
+    """
+    Write the ``Allsky`` preview files for a 2-d image HiPS.
+
+    For each order between 0 and 3 (inclusive, and no deeper than the deepest
+    order of the HiPS), all the tiles of that order are packed, downsampled, into
+    a single ``Norder{order}/Allsky.{ext}`` mosaic. The tiles are laid out side
+    by side in HEALPix index order, with a width of ``floor(sqrt(n_tiles))``, as
+    described by the HiPS standard.
+    """
+
+    factor = _reduction_factor(tile_size)
+    allsky_tile = tile_size // factor
+
+    for order in range(min(ALLSKY_MAX_ORDER, spatial_level) + 1):
+
+        n_tiles = 12 * 4**order
+        width = math.isqrt(n_tiles)  # floor(sqrt(n_tiles))
+        height = math.ceil(n_tiles / width)
+
+        if tile_format == "fits":
+            mosaic = np.full((height * allsky_tile, width * allsky_tile), np.nan, dtype=np.float32)
+        else:
+            channels = 4 if tile_format == "png" else 3
+            mosaic = np.zeros((height * allsky_tile, width * allsky_tile, channels), dtype=np.uint8)
+
+        found = False
+        for index in range(n_tiles):
+
+            filename = tile_filename(
+                level=order, index=index, output_directory=output_directory, extension=extension
+            )
+            if not os.path.exists(filename):
+                continue
+            found = True
+
+            row, col = divmod(index, width)
+            ys = slice(row * allsky_tile, (row + 1) * allsky_tile)
+            xs = slice(col * allsky_tile, (col + 1) * allsky_tile)
+
+            if tile_format == "fits":
+                data = fits.getdata(filename)
+                if factor > 1:
+                    # block_reduce with nanmean can warn on all-NaN tiles
+                    with warnings.catch_warnings():
+                        warnings.simplefilter("ignore")
+                        data = block_reduce(data, factor, func=np.nanmean)
+                mosaic[ys, xs] = data
+            else:
+                image = Image.open(filename).convert("RGBA" if tile_format == "png" else "RGB")
+                if factor > 1:
+                    image = image.reduce(factor)
+                mosaic[ys, xs] = np.asarray(image)
+
+        if not found:
+            continue
+
+        allsky_filename = os.path.join(output_directory, f"Norder{order}", f"Allsky.{extension}")
+        if tile_format == "fits":
+            fits.writeto(allsky_filename, mosaic, overwrite=True)
+        else:
+            image = Image.fromarray(mosaic)
+            if tile_format == "jpeg":
+                image = image.convert("RGB")
+            image.save(allsky_filename)

reproject/hips/_high_level.py

@@ -29,6 +29,7 @@
 from .._array_utils import sample_array_edges
 from .._wcs_utils import has_celestial, has_spectral, pixel_scale
 from ..utils import as_transparent_rgb, is_jpeg, is_png, parse_input_data
+from ._allsky import save_allsky
 from ._moc import save_moc
 from ._trim_utils import fits_getdata_untrimmed, fits_writeto_withtrim
 from ._utils import (
@@ -106,6 +107,7 @@ def reproject_to_hips(
     threads=False,
     properties=None,
     generate_moc=True,
+    allsky=True,
     **kwargs,
 ):
     """
@@ -165,6 +167,10 @@ def reproject_to_hips(
         Whether to write a ``Moc.fits`` coverage file at the root of the HiPS
         dataset (a spatial MOC for 2-d data and a space-frequency MOC for 3-d
         data), as recommended by the HiPS standard. Defaults to `True`.
+    allsky : bool, optional
+        Whether to write ``Allsky`` preview files for the low orders (0 to 3) of
+        a 2-d image HiPS, as described by the HiPS standard. Ignored for 3-d
+        data. Defaults to `True`.
     **kwargs
         Keyword arguments to be passed to the reprojection function.
 
@@ -505,6 +511,15 @@ def process(index):
         level_depth=level_depth,
     )
 
+    if allsky and ndim == 2:
+        save_allsky(
+            output_directory=output_directory,
+            tile_format=tile_format,
+            extension=EXTENSION[tile_format],
+            tile_size=tile_size,
+            spatial_level=spatial_level,
+        )
+
 
 def find_indices(*, output_directory, ndim, spatial_level, level_depth):
 

reproject/hips/tests/test_high_level.py

@@ -1,5 +1,6 @@
 import os
 import re
+import warnings
 from math import prod
 
 import numpy as np
@@ -18,9 +19,13 @@
 
 EXPECTED_FILES = [
     "Moc.fits",
+    "Norder0/Allsky.fits",
     "Norder0/Dir0/Npix0.fits",
+    "Norder1/Allsky.fits",
     "Norder1/Dir0/Npix2.fits",
+    "Norder2/Allsky.fits",
     "Norder2/Dir0/Npix9.fits",
+    "Norder3/Allsky.fits",
     "Norder3/Dir0/Npix38.fits",
     "Norder4/Dir0/Npix152.fits",
     "Norder4/Dir0/Npix153.fits",
@@ -83,10 +88,14 @@ def test_reproject_to_hips(tmp_path, valid_celestial_input_data):
 
 EXPECTED_FILES_GALACTIC = [
     "Moc.fits",
+    "Norder0/Allsky.fits",
     "Norder0/Dir0/Npix9.fits",
+    "Norder1/Allsky.fits",
     "Norder1/Dir0/Npix39.fits",
+    "Norder2/Allsky.fits",
     "Norder2/Dir0/Npix156.fits",
     "Norder2/Dir0/Npix157.fits",
+    "Norder3/Allsky.fits",
     "Norder3/Dir0/Npix627.fits",
     "Norder3/Dir0/Npix630.fits",
     "Norder4/Dir0/Npix2511.fits",
@@ -156,8 +165,11 @@ def test_reproject_to_hips_invalid_parameters(tmp_path, simple_celestial_fits_wc
 
 EXPECTED_FILES_AUTO_1 = [
     "Moc.fits",
+    "Norder0/Allsky.fits",
     "Norder0/Dir0/Npix0.fits",
+    "Norder1/Allsky.fits",
     "Norder1/Dir0/Npix2.fits",
+    "Norder2/Allsky.fits",
     "Norder2/Dir0/Npix9.fits",
     "index.html",
     "properties",
@@ -166,9 +178,13 @@ def test_reproject_to_hips_invalid_parameters(tmp_path, simple_celestial_fits_wc
 
 EXPECTED_FILES_AUTO_2 = [
     "Moc.fits",
+    "Norder0/Allsky.fits",
     "Norder0/Dir0/Npix0.fits",
+    "Norder1/Allsky.fits",
     "Norder1/Dir0/Npix2.fits",
+    "Norder2/Allsky.fits",
     "Norder2/Dir0/Npix9.fits",
+    "Norder3/Allsky.fits",
     "Norder3/Dir0/Npix38.fits",
     "Norder4/Dir0/Npix153.fits",
     "Norder5/Dir0/Npix612.fits",
@@ -614,5 +630,101 @@ def test_moc_disabled(tmp_path, simple_celestial_fits_wcs):
     assert not (output_directory / "Moc.fits").exists()
 
 
+def test_allsky_2d(tmp_path, simple_celestial_fits_wcs):
+    # An Allsky preview is written for each low order (0-3), packing all the
+    # tiles of that order into a single downsampled mosaic.
+    from astropy.nddata import block_reduce
+
+    output_directory = tmp_path / "output"
+    reproject_to_hips(
+        (np.arange(30 * 40).reshape(30, 40).astype(float), simple_celestial_fits_wcs),
+        coord_system_out="equatorial",
+        level=5,
+        tile_size=128,
+        reproject_function=reproject_interp,
+        output_directory=output_directory,
+    )
+
+    # Allsky files exist for orders 0-3 only (not for the deeper orders 4-5)
+    for order in range(4):
+        assert (output_directory / f"Norder{order}" / "Allsky.fits").exists()
+    for order in (4, 5):
+        assert not (output_directory / f"Norder{order}" / "Allsky.fits").exists()
+
+    # The order-3 mosaic has the standard layout: width = floor(sqrt(768)) = 27
+    # columns of 64-pixel tiles (128 downsampled by 2)
+    allsky = fits.getdata(output_directory / "Norder3" / "Allsky.fits")
+    assert allsky.shape == (29 * 64, 27 * 64)
+
+    # Each generated tile appears, downsampled, in its mosaic cell
+    tile_path = next((output_directory / "Norder3").glob("Dir*/Npix*.fits"))
+    index = int(tile_path.name[:-5].replace("Npix", ""))
+    row, col = divmod(index, 27)
+    cell = allsky[row * 64 : (row + 1) * 64, col * 64 : (col + 1) * 64]
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", RuntimeWarning)  # all-NaN blocks in nanmean
+        expected = block_reduce(fits.getdata(tile_path), 2, func=np.nanmean)
+    mask = np.isfinite(expected) & np.isfinite(cell)
+    assert mask.any()
+    np.testing.assert_allclose(cell[mask], expected[mask])
+
+
+@pytest.mark.parametrize(("tile_format", "extension"), [("png", "png"), ("jpeg", "jpg")])
+def test_allsky_image_formats(tmp_path, simple_celestial_fits_wcs, tile_format, extension):
+    # Allsky previews are also written for PNG/JPEG image HiPS, with the same
+    # mosaic layout as the FITS case.
+    layer = (np.arange(120 * 120).reshape(120, 120) % 256).astype(np.uint8)
+    original = tmp_path / f"original.{extension}"
+    tagged = tmp_path / f"tagged.{extension}"
+    Image.fromarray(np.dstack([layer, layer, layer])).save(original)
+    AVM.from_wcs(simple_celestial_fits_wcs, shape=(120, 120)).embed(original, tagged)
+
+    output_directory = tmp_path / "output"
+    reproject_to_hips(
+        tagged,
+        coord_system_out="equatorial",
+        level=3,
+        tile_size=128,
+        reproject_function=reproject_interp,
+        output_directory=output_directory,
+    )
+
+    for order in range(4):
+        assert (output_directory / f"Norder{order}" / f"Allsky.{extension}").exists()
+
+    channels = 4 if tile_format == "png" else 3
+    allsky = np.asarray(Image.open(output_directory / "Norder3" / f"Allsky.{extension}"))
+    assert allsky.shape == (29 * 64, 27 * 64, channels)
+
+    # A generated tile appears, downsampled, in its mosaic cell
+    tile_path = next((output_directory / "Norder3").glob(f"Dir*/Npix*.{extension}"))
+    index = int(tile_path.name.replace("Npix", "").split(".")[0])
+    row, col = divmod(index, 27)
+    cell = allsky[row * 64 : (row + 1) * 64, col * 64 : (col + 1) * 64]
+    tile = np.asarray(
+        Image.open(tile_path).convert("RGBA" if tile_format == "png" else "RGB").reduce(2)
+    )
+    if tile_format == "png":
+        np.testing.assert_array_equal(cell, tile)  # lossless
+    else:
+        # JPEG is lossy (and the Allsky re-compresses), so the cell should
+        # broadly match the downsampled tile rather than be pixel-exact
+        assert np.abs(cell.astype(int) - tile.astype(int)).mean() < 10
+
+
+def test_allsky_disabled(tmp_path, simple_celestial_fits_wcs):
+    # No Allsky files are written when allsky=False.
+    output_directory = tmp_path / "output"
+    reproject_to_hips(
+        (np.ones((30, 40)), simple_celestial_fits_wcs),
+        coord_system_out="equatorial",
+        level=4,
+        reproject_function=reproject_interp,
+        output_directory=output_directory,
+        allsky=False,
+    )
+    assert not list(output_directory.glob("Norder*/Allsky.*"))
+
+
 # TODO: Add tests of different spectral frames
 # TODO: Add test of auto level determination for all sky maps in 2D and 3D