diff --git a/python/lsst/pipe/tasks/fits2hips/__init__.py b/python/lsst/pipe/tasks/fits2hips/__init__.py new file mode 100644 index 000000000..d325a5e64 --- /dev/null +++ b/python/lsst/pipe/tasks/fits2hips/__init__.py @@ -0,0 +1,3 @@ +from ._all_sky_hips import * +from ._high_order_hips import * +from ._low_order_hips import * diff --git a/python/lsst/pipe/tasks/fits2hips/_all_sky_hips.py b/python/lsst/pipe/tasks/fits2hips/_all_sky_hips.py new file mode 100644 index 000000000..e501f337f --- /dev/null +++ b/python/lsst/pipe/tasks/fits2hips/_all_sky_hips.py @@ -0,0 +1,423 @@ +# This file is part of pipe_tasks. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +from __future__ import annotations + +__all__ = ("AllSkyFitsHipsTaskConnections", "AllSkyFitsHipsTaskConfig", "AllSkyFitsHipsTask") + +import math +import re +import hpgeom as hpg +import numpy as np +import healsparse as hsp +from dataclasses import replace +from datetime import datetime +from collections.abc import Iterable + + +from lsst.resources import ResourcePath +from lsst.pex.config import Field, ConfigField +from lsst.sphgeom import RangeSet +from lsst.images import Image + +from lsst.pipe.base import ( + PipelineTask, + PipelineTaskConfig, + PipelineTaskConnections, + Struct, + QuantumContext, + InputQuantizedConnection, + OutputQuantizedConnection, + TaskMetadata, +) +from lsst.pipe.base.connectionTypes import Input + + +from ..rgb2hips._properties import HipsPropertiesConfig, _write_property +from ..healSparseMapping import _is_power_of_two + + +class AllSkyFitsHipsTaskConnections( + PipelineTaskConnections, + dimensions=("band",), + defaultTemplates={"input_task_label": "generateLowOrderFitsHips"}, +): + low_order_metadata = Input( + doc="Metadata produced by the LowOrderHipsTask", + name="{input_task_label}_metadata", + storageClass="TaskMetadata", + multiple=True, + deferLoad=True, + dimensions=tuple(), + ) + input_hips = Input( + doc="Hips pixels at level 8 used to build higher orders", + name="fits_picture_hips8", + storageClass="NumpyArray", + multiple=True, + deferLoad=True, + dimensions=("healpix8", "band"), + ) + + def __init__(self, *, config: AllSkyFitsHipsTaskConfig): + # Set the input dimensions to whatever the minimum order healpix + # to produce is. + self.low_order_metadata = replace( + self.low_order_metadata, dimensions=set((f"healpix{config.min_order}", "band")) + ) + + +class AllSkyFitsHipsTaskConfig(PipelineTaskConfig, pipelineConnections=AllSkyFitsHipsTaskConnections): + hips_base_uri = Field[str]( + doc="URI to HiPS base for output.", + optional=False, + ) + properties = ConfigField[HipsPropertiesConfig]( + doc="Configuration for properties file.", + ) + allsky_tilesize = Field[int]( + dtype=int, + doc="Allsky tile size; must be power of 2. HiPS standard recommends 64x64 tiles.", + default=64, + check=_is_power_of_two, + ) + max_order = Field[int](doc="The maximum order hips that was produced", default=11) + shift_order = Field[int]( + doc="Shift order of hips, right now must be 9 configuration for future options", default=9 + ) + min_order = Field[int]( + doc="Minimum healpix order for HiPS tree.", + default=3, + ) + + def validate(self): + if self.shift_order != 9: + raise ValueError("Shift order must be 9.") + return super().validate() + + +class AllSkyFitsHipsTask(PipelineTask): + """Pipeline task for generating all-sky HealPix (HiPS) tiles and associated metadata.""" + + _DefaultName = "allSkyHipsTask" + ConfigClass = AllSkyFitsHipsTaskConfig + config: ConfigClass + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.hips_base_path = ResourcePath(self.config.hips_base_uri, forceDirectory=True) + + def runQuantum( + self, + butlerQC: QuantumContext, + inputRefs: InputQuantizedConnection, + outputRefs: OutputQuantizedConnection, + ) -> None: + band_name = butlerQC.quantum.dataId["band"] + + self.hips_base_path = self.hips_base_path.join(f"band_{band_name}", forceDirectory=True) + + # Extract the healpix8 pixel ids. + hpx8_pixels = [] + for ref in inputRefs.input_hips: + hpx8_pixels.append((ref.dataId["healpix8"])) + + # Scale pixel IDS to higher order (hpx11) that were already produced + hpx8_rangeset = RangeSet(hpx8_pixels) + hpx11_rangeset = hpx8_rangeset.scaled(4**3) + hpx11_pixels = set() + for begin, end in hpx11_rangeset: + hpx11_pixels.update(range(begin, end)) + hpx11_pixels = np.array([s for s in hpx11_pixels]) + + low_order_metadata = butlerQC.get(inputRefs.low_order_metadata) + + outputs = self.run(low_order_metadata, hpx11_pixels, band_name) + butlerQC.put(outputs, outputRefs) + + def run(self, low_order_metadata: Iterable[TaskMetadata], hpx11_pixels, band_name) -> Struct: + """Generate all-sky HealPix tiles and metadata. + + Parameters + ---------- + low_order_metadata : Iterable[TaskMetadata] + Low-order metadata from previous processing steps. + hpx11_pixels : array-like + Array of HPX11 pixel IDs to process. + band_name : `str` + The band in which this data was collected + + Returns + ------- + Struct + This task produces no outputs so an empty struct is returned + """ + self._write_properties_and_moc( + self.config.max_order, hpx11_pixels, self.config.shift_order, band_name + ) + self._write_allsky_file(self.config.min_order) + return Struct() + + def _write_properties_and_moc(self, max_order, pixels, shift_order, band): + """Write HiPS properties file and MOC. + + Parameters + ---------- + max_order : `int` + Maximum HEALPix order. + pixels : `np.ndarray` (N,) + Array of pixels used. + shift_order : `int` + HPX shift order. + band : `str` + Band (or color). + """ + area = hpg.nside_to_pixel_area(2**max_order, degrees=True) * len(pixels) + + initial_ra = self.config.properties.initial_ra + initial_dec = self.config.properties.initial_dec + initial_fov = self.config.properties.initial_fov + + if initial_ra is None or initial_dec is None or initial_fov is None: + # We want to point to an arbitrary pixel in the footprint. + # Just take the median pixel value for simplicity. + temp_pixels = pixels.copy() + if temp_pixels.size % 2 == 0: + temp_pixels = np.append(temp_pixels, [temp_pixels[0]]) + medpix = int(np.median(temp_pixels)) + _initial_ra, _initial_dec = hpg.pixel_to_angle(2**max_order, medpix) + _initial_fov = hpg.nside_to_resolution(2**max_order, units="arcminutes") / 60.0 + + if initial_ra is None or initial_dec is None: + initial_ra = _initial_ra + initial_dec = _initial_dec + if initial_fov is None: + initial_fov = _initial_fov + + self._write_hips_properties_file( + self.config.properties, + band, + max_order, + shift_order, + area, + initial_ra, + initial_dec, + initial_fov, + ) + + # Write the MOC coverage + self._write_hips_moc_file( + max_order, + pixels, + ) + + def _write_hips_properties_file( + self, + properties_config, + band, + max_order, + shift_order, + area, + initial_ra, + initial_dec, + initial_fov, + ): + """Write HiPS properties file. + + Parameters + ---------- + properties_config : `lsst.pipe.tasks.hips.HipsPropertiesConfig` + Configuration for properties values. + band : `str` + Name of band(s) for HiPS tree. + max_order : `int` + Maximum HEALPix order. + shift_order : `int` + HPX shift order. + area : `float` + Coverage area in square degrees. + initial_ra : `float` + Initial HiPS RA position (degrees). + initial_dec : `float` + Initial HiPS Dec position (degrees). + initial_fov : `float` + Initial HiPS display size (degrees). + """ + + bitpix = 32 + hbitpix = 32 + + date_iso8601 = datetime.utcnow().isoformat(timespec="seconds") + "Z" + pixel_scale = hpg.nside_to_resolution(2 ** (max_order + shift_order), units="degrees") + + uri = self.hips_base_path.join("properties") + with ResourcePath.temporary_uri(suffix=uri.getExtension()) as temporary_uri: + with open(temporary_uri.ospath, "w") as fh: + _write_property( + fh, + "creator_did", + properties_config.creator_did_template.format(band=band), + ) + if properties_config.obs_collection is not None: + _write_property(fh, "obs_collection", properties_config.obs_collection) + _write_property( + fh, + "obs_title", + properties_config.obs_title_template.format(band=band), + ) + if properties_config.obs_description_template is not None: + _write_property( + fh, + "obs_description", + properties_config.obs_description_template.format(band=band), + ) + if len(properties_config.prov_progenitor) > 0: + for prov_progenitor in properties_config.prov_progenitor: + _write_property(fh, "prov_progenitor", prov_progenitor) + if properties_config.obs_ack is not None: + _write_property(fh, "obs_ack", properties_config.obs_ack) + _write_property(fh, "obs_regime", "Optical") + _write_property(fh, "data_pixel_bitpix", str(bitpix)) + _write_property(fh, "dataproduct_type", "image") + _write_property(fh, "moc_sky_fraction", str(area / 41253.0)) + _write_property(fh, "data_ucd", "phot.flux") + _write_property(fh, "hips_creation_date", date_iso8601) + _write_property(fh, "hips_builder", "lsst.pipe.tasks.hips.GenerateHipsTask") + _write_property(fh, "hips_creator", "Vera C. Rubin Observatory") + _write_property(fh, "hips_version", "1.4") + _write_property(fh, "hips_release_date", date_iso8601) + _write_property(fh, "hips_frame", "equatorial") + _write_property(fh, "hips_order", str(max_order)) + _write_property(fh, "hips_tile_width", str(2**shift_order)) + _write_property(fh, "hips_status", "private master clonableOnce") + _write_property(fh, "hips_tile_format", "fits") + _write_property(fh, "dataproduct_subtype", "color") + _write_property(fh, "hips_pixel_bitpix", str(hbitpix)) + _write_property(fh, "hips_pixel_scale", str(pixel_scale)) + _write_property(fh, "hips_initial_ra", str(initial_ra)) + _write_property(fh, "hips_initial_dec", str(initial_dec)) + _write_property(fh, "hips_initial_fov", str(initial_fov)) + if band in properties_config.spectral_ranges: + em_min = properties_config.spectral_ranges[band].lambda_min / 1e9 + else: + self.log.warning("blue band %s not in self.config.spectral_ranges.", band) + em_min = 3e-7 + if band in properties_config.spectral_ranges: + em_max = properties_config.spectral_ranges[band].lambda_max / 1e9 + else: + self.log.warning("red band %s not in self.config.spectral_ranges.", band) + em_max = 1e-6 + _write_property(fh, "em_min", str(em_min)) + _write_property(fh, "em_max", str(em_max)) + if properties_config.t_min is not None: + _write_property(fh, "t_min", properties_config.t_min) + if properties_config.t_max is not None: + _write_property(fh, "t_max", properties_config.t_max) + + uri.transfer_from(temporary_uri, transfer="copy", overwrite=True) + + def _write_hips_moc_file(self, max_order, pixels, min_uniq_order=1): + """Write HiPS MOC file. + + Parameters + ---------- + max_order : `int` + Maximum HEALPix order. + pixels : `np.ndarray` + Array of pixels covered. + min_uniq_order : `int`, optional + Minimum HEALPix order for looking for fully covered pixels. + """ + # WARNING: In general PipelineTasks are not allowed to do any outputs + # outside of the butler. This task has been given (temporary) + # Special Dispensation because of the nature of HiPS outputs until + # a more controlled solution can be found. + + # Make a healsparse map which provides easy degrade/comparisons. + hspmap = hsp.HealSparseMap.make_empty(2**min_uniq_order, 2**max_order, dtype=np.int8) + hspmap[pixels] = 1 + + uri = self.hips_base_path.join("Moc.fits") + with ResourcePath.temporary_uri(suffix=uri.getExtension()) as temporary_uri: + hspmap.write_moc(temporary_uri.ospath) + uri.transfer_from(temporary_uri, transfer="copy", overwrite=True) + + def _write_allsky_file(self, allsky_order): + """Write an Allsky.png file. + + Parameters + ---------- + allsky_order : `int` + HEALPix order of the minimum order to make allsky file. + """ + tile_size = self.config.allsky_tilesize + + # The Allsky file format is described in + # https://www.ivoa.net/documents/HiPS/20170519/REC-HIPS-1.0-20170519.pdf + # From S4.3.2: + # The Allsky file is built as an array of tiles, stored side by side in + # the left-to-right order. The width of this array must be the square + # root of the number of the tiles of the order. For instance, the width + # of this array at order 3 is 27 ( (int)sqrt(768) ). To avoid having a + # too large Allsky file, the resolution of each tile may be reduced but + # must stay a power of two (typically 64x64 pixels rather than 512x512). + + n_tiles = hpg.nside_to_npixel(hpg.order_to_nside(allsky_order)) + n_tiles_wide = int(np.floor(np.sqrt(n_tiles))) + n_tiles_high = int(np.ceil(n_tiles / n_tiles_wide)) + + allsky_image = None + + allsky_order_uri = self.hips_base_path.join(f"Norder{allsky_order}", forceDirectory=True) + pixel_regex = re.compile(r"Npix([0-9]+)\.fits$") + + image_uris = list( + ResourcePath.findFileResources( + candidates=[allsky_order_uri], + file_filter=pixel_regex, + ) + ) + + for image_uri in image_uris: + matches = re.match(pixel_regex, image_uri.basename()) + pix_num = int(matches.group(1)) + tile_image = Image.read(image_uri).array + row = math.floor(pix_num // n_tiles_wide) + column = pix_num % n_tiles_wide + left, upper, right, lower = ( + column * tile_size, + row * tile_size, + (column + 1) * tile_size, + (row + 1) * tile_size, + ) + tile_image_shrunk = tile_image.resize((tile_size, tile_size)) + + if allsky_image is None: + allsky_image = np.zeros( + (n_tiles_wide * tile_size, n_tiles_high * tile_size), + ) + allsky_image[upper:lower, left:right] = tile_image_shrunk + + uri = allsky_order_uri.join("Allsky.fits") + + with ResourcePath.temporary_uri(suffix=uri.getExtension()) as temporary_uri: + Image(allsky_image).write(temporary_uri.ospath) + + uri.transfer_from(temporary_uri, transfer="copy", overwrite=True) diff --git a/python/lsst/pipe/tasks/fits2hips/_high_order_hips.py b/python/lsst/pipe/tasks/fits2hips/_high_order_hips.py new file mode 100644 index 000000000..c31bf2dd8 --- /dev/null +++ b/python/lsst/pipe/tasks/fits2hips/_high_order_hips.py @@ -0,0 +1,338 @@ +# This file is part of pipe_tasks. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +from __future__ import annotations + +__all__ = ("HighOrderFitsHipsTaskConnections", "HighOrderFitsHipsTaskConfig", "HighOrderFitsHipsTask") + +import numpy as np +from numpy.typing import NDArray + +from lsst.afw.geom import makeHpxWcs +from lsst.pipe.base import ( + PipelineTask, + PipelineTaskConfig, + PipelineTaskConnections, + Struct, + QuantumContext, + InputQuantizedConnection, + OutputQuantizedConnection, +) +from lsst.pex.config import ConfigField, Field +from lsst.pipe.base.connectionTypes import Input, Output +from lsst.skymap import BaseSkyMap +from lsst.afw.geom import SkyWcs +from lsst.geom import Box2I, Point2I, Extent2I +from lsst.afw.math import Warper +from lsst.daf.butler import DeferredDatasetHandle +from lsst.afw.image import ImageF +from lsst.resources import ResourcePath + +from collections.abc import Iterable +from lsst.sphgeom import RangeSet + +import cv2 + +from ..rgb2hips._utils import _write_hips_image +from ..prettyPictureMaker import FeatheredMosaicCreator + + +class HighOrderFitsHipsTaskConnections(PipelineTaskConnections, dimensions=("healpix8", "band")): + input_images = Input( + doc="Fits images which are to be turned into hips tiles", + name="deep_coadd", + storageClass="ExposureF", + dimensions=("tract", "patch", "skymap", "band"), + multiple=True, + deferLoad=True, + ) + skymap = Input( + doc="The skymap which the data has been mapped onto", + storageClass="SkyMap", + name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, + dimensions=("skymap",), + ) + output_hpx = Output( + doc="Healpix tiles at order 8, but binned to 256x256", + name="fits_picture_hips8", + storageClass="NumpyArray", + dimensions=("healpix8", "band"), + ) + + +class HighOrderFitsHipsTaskConfig(PipelineTaskConfig, pipelineConnections=HighOrderFitsHipsTaskConnections): + """Configuration class for the HighOrderHipsTask pipeline task.""" + + hips_order = 8 + """HealPix order to generate tiles for.""" + warp = ConfigField[Warper.ConfigClass]( + doc="Warper configuration", + ) + hips_base_uri = Field[str]( + doc="URI to HiPS base for output.", + optional=False, + ) + + def setDefaults(self): + self.warp.warpingKernelName = "lanczos5" + + +class HighOrderFitsHipsTask(PipelineTask): + """Pipeline task that generates high-order HealPix tiles from Fits images. + + Of Note; This task has special dispensation to write "out-of-tree" to a + location not within the butler. DO NOT model other tasks on this one. + + This task takes in Fits images generated on a tract patch grid. It assembles + them into a 4096 x 4096 image aligned with the wcs coordinates of hips + order 8 pixels. This is then divided up into an 8x8 grid to produce 512x512 + images at hips order 11. The images is then resampled using lanczos order 4 + such that the image is half the size. The original image is then divided + into a 4x4 grid to produce hips images at order 10. The process is repeated + to produce hips images at order 9, and finally the image is resampled down + to 512x512 and saved out at hips order 8. + + The order 8 image is resampled one more time to 256x256 and presisted by + the butler for later consumption in the `LowOrderHipsTask`. + + The difference at producding wcs at order 8 and working up to 11, is tested + to be less than 6 decimal places when converting ra dec to pixel coordinates, + and even that is likely to be due to differences in warping kernels, + and not an intrinsic error. Doing processing like this allows hips generation + to be more effectively split across compute nodes. + """ + + _DefaultName = "highOrderFitsHipsTask" + ConfigClass = HighOrderFitsHipsTaskConfig + + config: ConfigClass + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.warper = Warper.fromConfig(self.config.warp) + + # Set the base resource path that will be used for all outputs + self.hips_base_path = ResourcePath(self.config.hips_base_uri, forceDirectory=True) + + def run(self, input_images: Iterable[tuple[NDArray, SkyWcs, Box2I]], healpix_id) -> Struct: + """Main execution method for generating HealPix tiles. + + Parameters + ---------- + input_images : Iterable[tuple[NDArray, SkyWcs, Box2I]] + Iterable of tuples containing image data, WCS, and bounding box information. + healpix_id : int + The HealPix order 8 ID to process. + + Returns + ------- + Struct + Output structure containing the processed HealPix order 8 tile. + This has been downsampled to 256x256 corresponding to a quarter of a healpix + order 7 image. + """ + # Make the WCS for the transform, intentionally over-sampled to shift order 12. + # This creates as 4096 x 4096 image that can be broken apart to form the higher + # orders, binning each as needed + target_wcs = makeHpxWcs(8, healpix_id, 12, False) + + # construct a bounding box that holds the warping results for each channel + exp_bbox = Box2I(corner=Point2I(0, 0), dimensions=Extent2I(2**12, 2**12)) + + output_array_hpx = np.zeros((4096, 4096), dtype=np.float32) + output_array_hpx[:, :] = np.nan + + self.log.info("Warping input exposures and populating hpx8 super tile.") + # Need to loop over input arrays + # Warp and combine input images into the HealPix tile + for input_image, in_wcs, in_box in input_images: + tmp_image = ImageF(in_box) + in_image: NDArray = input_image + + # construct an Exposure object from one channel in the array + tmp_image.array[:, :] = in_image + + # Warp the image to the target WCS + warpped = self.warper.warpImage(target_wcs, tmp_image, in_wcs, maxBBox=exp_bbox) + warpped_box_slices = warpped.getBBox().slices + + # Update the output array with valid (non-NaN) values + are_warpped = np.isfinite(warpped.array) + output_array_hpx[warpped_box_slices][are_warpped] = warpped.array[are_warpped] + + # Replace any remaining NaN values with zeros + output_array_hpx[np.isnan(output_array_hpx)] = 0 + + # Flip the y-axis to match HealPix indexing + output_array_hpx = output_array_hpx[::-1, :] + + # Generate tiles for different HealPix orders using Lanczos resampling instead of binning. + # This handles how intensities should change as the hips level changes. + # + # what this does is take a single 4096 x 4096 image and resamples it in a courser grain such + # that the output pixels correspond to a 4x4 grid of hips pixels at an increasingly lower scale. + # This works because hips is a hierarchy of tiles all contained in the same area of the sky. + # This allows us to generate all the output images by resampling the inputs and saves the time + # required to generate whole new images at each scale. + # + # The loop variables are the resampling factor, the hips order, and the number of sub-divisions + # a pixel has gone through (used to determine quadrant). + for zoom, hips_level, factor in zip((0, 2, 4, 8), (11, 10, 9, 8), (3, 2, 1, 0)): + self.log.info("Generating tiles for hxp level %d", hips_level) + if zoom: + size = 4096 // zoom + binned_array = cv2.resize(output_array_hpx, (size, size), interpolation=cv2.INTER_AREA) + else: + binned_array = output_array_hpx + # always create blocks of 512x512 as that is native shift order 9 size + # + # Figure out the hips pixel ids at this hips order. This is complicated because each hipx pixel + # turns into 4 at a higher level, but must be in a specific order to correspond to how the data + # is layed out in an y,x grid. So if a hips order 8 pixel A turns into four pixels b,c,d,e, they + # are layed out like [[b,d], [c,e]]. This is true for every pixel as you go up in order. So + # if you start at order 8 with one pixel, you need to do order 9 and calculate the layout. Then + # for each order 9 pixel, do the same to get the layout in order 10, etc. This leaves a grid + # of pixels that are the ids of the corresponding 512,512 sub grid pixel in the input image. + tmp_pixels = np.array([[healpix_id]]) + for _ in range(factor): + tmp_array = np.zeros(np.array(tmp_pixels.shape) * 2) + for ii in range(tmp_pixels.shape[0]): + for jj in range(tmp_pixels.shape[1]): + tmp_array_view = tmp_array[ii * 2 : ii * 2 + 2, jj * 2 : jj * 2 + 2] + tmp_range_set = RangeSet(int(tmp_pixels[ii, jj])) + tmp_array_view[:, :] = ( + np.array([x for x in range(*tmp_range_set.scaled(4)[0])], dtype=int)[[0, 2, 1, 3]] + ).reshape(2, 2) + tmp_pixels = tmp_array + + # now for each 512x512 sub pixel region write the hips image with the corresponding healpix id + hpx_id_array = tmp_pixels + for i in range(binned_array.shape[0] // 512): + for j in range(binned_array.shape[1] // 512): + pixel_id = int(hpx_id_array[i, j]) + sub_pixel = binned_array[i * 512 : i * 512 + 512, j * 512 : j * 512 + 512] + self.log.info(f"writing sub_pixel {pixel_id}") + tile_wcs = makeHpxWcs(hips_level, pixel_id, 9) + _write_hips_image( + sub_pixel, + pixel_id, + hips_level, + self.hips_base_path, + "fits", + "float", + tile_wcs=tile_wcs, + ) + + # Finally, bin the level 8 hpx to 256x256 (1/4 order 7) to save to the butler. + # This makes smaller arrays to load, and saves the binning operation in the joint phase. + zoomed = cv2.resize(output_array_hpx, (256, 256), interpolation=cv2.INTER_LANCZOS4) + + return Struct(output_hpx=zoomed) + + def _assemble_sub_region( + self, tract_patch: dict[int, Iterable[tuple[DeferredDatasetHandle, SkyWcs, Box2I]]], patch_grow: int + ) -> list[tuple[NDArray, SkyWcs, Box2I]]: + """Assemble all the patches in each tract into images. + + This function takes in an input keyed by tract, with values + corresponding the patches in that tract that overlap the quatum's + healpix value. It assembles each of these into a single image such + that the return values is a list of images (and metadata) one element + for each input tract. + + Parameters + ---------- + tract_patch : `dict` of `int` to `iterable` of `tuple` of + `DeferredDatasetHandle`, `SkyWcs` and `Box2I` + Input images and metadata organized into corresponding tracts. + patch_grow : `int` + Amount to grow patches by + + Returns + ------- + output_list : `list` of `tuple` of `NDArray` `SkyWcs` and `Box2I` + List of assembled images and metadata, one element for each tract + + """ + + boxes = [] + for _, iterable in tract_patch.items(): + mosaic_maker = FeatheredMosaicCreator(patch_grow) + new_box = Box2I() + for _, _, bbox in iterable: + new_box.include(bbox) + # allocate tmp array + new_array = np.zeros((new_box.getHeight(), new_box.getWidth()), dtype=np.float32) + for handle, skyWcs, box in iterable: + # Make a new box of the same size, but with the origin centered + # on the lowest corner were there is data. + localOrigin = box.getBegin() - new_box.getBegin() + localOrigin = Point2I( + x=int(np.floor(localOrigin.x)), + y=int(np.floor(localOrigin.y)), + ) + + localExtent = Extent2I( + x=int(np.floor(box.getWidth())), + y=int(np.floor(box.getHeight())), + ) + tmpBox = Box2I(localOrigin, localExtent) + tmp_new_box = Box2I(Point2I(x=0, y=0), Extent2I(x=new_box.getWidth(), y=new_box.getHeight())) + + image = handle.get() + mosaic_maker.add_to_image(new_array, image.image.array, tmp_new_box, tmpBox, reverse=False) + boxes.append((new_array, skyWcs, new_box)) + return boxes + + def runQuantum( + self, + butlerQC: QuantumContext, + inputRefs: InputQuantizedConnection, + outputRefs: OutputQuantizedConnection, + ) -> None: + # First get what healpix pixel this task is working on + healpix_id = butlerQC.quantum.dataId["healpix8"] + + # grab the skymap + skymap: BaseSkyMap = butlerQC.get(inputRefs.skymap) + + # Iterate over the input image refs, to get the corresponding bbox + # and assemble into container for run + inputs_by_tract = {} + for input_image_ref in inputRefs.input_images: + tract = input_image_ref.dataId["tract"] + patch = input_image_ref.dataId["patch"] + # All boxes in a given skymap will have the same inner dimensions + # for x and y and will be the same for all patches + imageWcs = skymap[tract][patch].getWcs() + box = skymap[tract][patch].getOuterBBox() + patch_grow = skymap[tract][patch].getCellInnerDimensions().getX() + imageHandle = butlerQC.get(input_image_ref) + container = inputs_by_tract.setdefault(tract, list()) + container.append((imageHandle, imageWcs, box)) + + band_name = butlerQC.quantum.dataId["band"] + + self.hips_base_path = self.hips_base_path.join(f"band_{band_name}", forceDirectory=True) + + input_images = self._assemble_sub_region(inputs_by_tract, patch_grow) + + outputs = self.run(input_images, healpix_id) + butlerQC.put(outputs, outputRefs) diff --git a/python/lsst/pipe/tasks/fits2hips/_low_order_hips.py b/python/lsst/pipe/tasks/fits2hips/_low_order_hips.py new file mode 100644 index 000000000..cbb2f809c --- /dev/null +++ b/python/lsst/pipe/tasks/fits2hips/_low_order_hips.py @@ -0,0 +1,211 @@ +# This file is part of pipe_tasks. +# +# Developed for the LSST Data Management System. +# This product includes software developed by the LSST Project +# (https://www.lsst.org). +# See the COPYRIGHT file at the top-level directory of this distribution +# for details of code ownership. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with this program. If not, see . +from __future__ import annotations + +__all__ = ("LowOrderFitsHipsTaskConnections", "LowOrderFitsHipsTaskConfig", "LowOrderFitsHipsTask") + +import numpy as np +import cv2 + +from lsst.afw.geom import makeHpxWcs + +from lsst.daf.butler import DeferredDatasetHandle +from lsst.pipe.base import ( + PipelineTask, + PipelineTaskConfig, + PipelineTaskConnections, + Struct, + QuantumContext, + InputQuantizedConnection, + OutputQuantizedConnection, +) + +from lsst.pex.config import Field +from lsst.pipe.base.connectionTypes import Input +from lsst.resources import ResourcePath + +from collections.abc import Iterable + +from numpy.typing import NDArray + +from ..rgb2hips._utils import _write_hips_image + + +class LowOrderFitsHipsTaskConnections( + PipelineTaskConnections, + dimensions=tuple(), +): + input_hips = Input( + doc="Hips pixels at level 8 used to build higher orders", + name="fits_picture_hips8", + storageClass="NumpyArray", + multiple=True, + deferLoad=True, + dimensions=("healpix8", "band"), + ) + + def __init__(self, *, config: LowOrderFitsHipsTaskConfig): + # Set the quantum dimensions to whatever the minimum order healpix + # to produce is. + self.dimensions = set( + (f"healpix{config.min_order}", "band"), + ) + + +class LowOrderFitsHipsTaskConfig(PipelineTaskConfig, pipelineConnections=LowOrderFitsHipsTaskConnections): + min_order = Field[int]( + doc="Minimum healpix order for HiPS tree.", + default=3, + ) + hips_base_uri = Field[str]( + doc="URI to HiPS base for output.", + optional=False, + ) + + def validate(self): + if self.min_order >= 8: + raise ValueError("The minimum order must be less than 8.") + + +class LowOrderFitsHipsTask(PipelineTask): + """`PipelineTask` to create low order hips tiles. + + This task reads in healpix 8 tiles, which have already been down sampled, + and assembles them into progressively lower hips order tiles. + + This task has special permission to write to locations outside the butler. + Don't emulate this in other tasks. + """ + + _DefaultName = "lowOrderHipsTask" + ConfigClass = LowOrderFitsHipsTaskConfig + + config: ConfigClass + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.hips_base_path = ResourcePath(self.config.hips_base_uri, forceDirectory=True) + + def run(self, hpx_container: Iterable[tuple[DeferredDatasetHandle, int]]) -> Struct: + """Produce Hips images with hips order 8 inputs to the configured min_order. + + Parameters + ---------- + hpx_container : `Iterable` of `tuple` of `DeferredDatasetHanle`, `int` + This is an iterable of handles to already down-sampled hpx order 8 + arrays and their corresponding order 8 pixel id. + + Returns + ------- + result : `Struct` + This tasks does not produce an output, so will return an empty `Struct` + + """ + # loop over each order, assembling the previous order tiles into + # an array, and writing the image. Resample each image smaller, + # and continue downward in order. + # This must 7 here based on the outputs of HighOrderHipsTask being + # healpix order 8 pixels. + for order in range(7, self.config.min_order - 1, -1): + self.log.info("Processing order %d", order) + # sort the previous order's pixels into a mapping with keys of + # this order's pixel to the corresponding previous orders pixels + # that are contained within that key. + hpx_next_mapping = self._create_sorted_container(hpx_container) + + hpx_next_container = [] + npix = 512 + size_thresh = len(hpx_next_mapping) // 10 + size_counter = 0 + percent_counter = 0 + for hpx_next_id, hpx_next_items in hpx_next_mapping.items(): + # Print out a log message every so often for a liveness + # check + if size_counter > size_thresh: + percent_counter += 10 + self.log.info("Done %d percent", percent_counter) + size_counter = 0 + # allocate a container for the pixel being assembled + hpx_next_array = np.zeros((npix, npix), dtype=np.float32) + for img_prev, hpx_prev_id in hpx_next_items: + if order == 7: + # These are saved out in float32 from the previous task + img_prev: NDArray = img_prev.get() + # determine which sub pixel quadrant this belongs to in the next orders + # pixel and assign. + sub_index = hpx_prev_id - np.left_shift(hpx_next_id, 2) + match sub_index: + case 0: + hpx_next_array[0 : npix // 2 :, 0 : npix // 2] = img_prev + case 1: + hpx_next_array[npix // 2 :, 0 : npix // 2] = img_prev + case 2: + hpx_next_array[0 : npix // 2, npix // 2 :] = img_prev + case 3: + hpx_next_array[npix // 2 :, npix // 2 :] = img_prev + # Write out the hips image + tile_wcs = makeHpxWcs(order, hpx_next_id, 9) + + _write_hips_image( + hpx_next_array, + hpx_next_id, + order, + self.hips_base_path, + "fits", + "float", + tile_wcs=tile_wcs, + ) + size_counter += 1 + + # resample the image to a smaller grid and store it for the next order + zoomed = cv2.resize(hpx_next_array, (256, 256), interpolation=cv2.INTER_AREA) + + hpx_next_container.append((zoomed, hpx_next_id)) + hpx_container = hpx_next_container + return Struct() + + def _create_sorted_container( + self, + hpx_container: Iterable[tuple[DeferredDatasetHandle, int]], + ) -> dict[int, Iterable[tuple[DeferredDatasetHandle, int]]]: + """Sort a list of [images (or handels), hpx_id] into corresponding pixels at a higher order.""" + hpx_output_mapping = {} + for pair in hpx_container: + hpx_output_id = np.right_shift(pair[1], 2) + hpx_output_container = hpx_output_mapping.setdefault(hpx_output_id, []) + hpx_output_container.append(pair) + return hpx_output_mapping + + def runQuantum( + self, + butlerQC: QuantumContext, + inputRefs: InputQuantizedConnection, + outputRefs: OutputQuantizedConnection, + ) -> None: + band_name = butlerQC.quantum.dataId["band"] + self.hips_base_path = self.hips_base_path.join(f"band_{band_name}", forceDirectory=True) + # get the hips handles and their pixel + hpx_container = [] + for ref in inputRefs.input_hips: + hpx_container.append((butlerQC.get(ref), ref.dataId["healpix8"])) + + outputs = self.run(hpx_container) + butlerQC.put(outputs, outputRefs) diff --git a/python/lsst/pipe/tasks/prettyPictureMaker/_utils.py b/python/lsst/pipe/tasks/prettyPictureMaker/_utils.py index f9bd3fb3b..628c44004 100644 --- a/python/lsst/pipe/tasks/prettyPictureMaker/_utils.py +++ b/python/lsst/pipe/tasks/prettyPictureMaker/_utils.py @@ -142,4 +142,4 @@ def add_to_image( patch = mixer * patch - image[*box.slices] += patch[::-1, :, :] if reverse else patch + image[*box.slices] += patch[::-1, ...] if reverse else patch diff --git a/python/lsst/pipe/tasks/rgb2hips/_utils.py b/python/lsst/pipe/tasks/rgb2hips/_utils.py index f8e9de3e9..836d6039c 100644 --- a/python/lsst/pipe/tasks/rgb2hips/_utils.py +++ b/python/lsst/pipe/tasks/rgb2hips/_utils.py @@ -26,9 +26,13 @@ from PIL import Image import numpy as np from numpy.typing import NDArray +from astropy import units from lsst.resources import ResourcePath +from lsst.afw.geom import SkyWcs +from lsst.images import Image as LsstImage +from lsst.images import Projection, GeneralFrame # allow PIL to work with really large images @@ -58,6 +62,7 @@ def _write_hips_image( hips_base_path: ResourcePath, file_extension: str, output_type: str, + tile_wcs: SkyWcs | None = None, ) -> None: """Write a processed image to disk in the HealPix tile format. @@ -76,15 +81,23 @@ def _write_hips_image( hips_base_path : `ResourcePath` Base directory path where the HealPix tiles will be stored. file_extension : `str` - File extension (format) for saving the image ('png' or 'webp'). + File extension (format) for saving the image ('png' or 'webp' or 'fits'). output_type : `str` Data type of the output array, which can be: - "uint8": 8-bit unsigned integers (0-255) - "uint16": 16-bit unsigned integers (0-65535) - "half": 16-bit floating-point numbers - "float": 32-bit floating-point numbers + tile_wcs : `~lsst.afw.geom.SkyWcs` or `None` + If supplied, this wcs is written to fits images. + Raises + ------ + ValueError : + Raised if the fits file extension is used with a type other than float """ + if file_extension == "fits" and output_type != "float": + raise ValueError("Fits files must be written with data type float") # clip in case any of the warping caused values over 1 image_data = np.clip(image_data, 0, 1) # remap the image_data to the chosen output_type @@ -107,6 +120,19 @@ def _write_hips_image( # Create the file URI for saving uri = hips_dir.join(f"Npix{pixel_id}.{file_extension}") + if file_extension == "fits": + im = LsstImage( + image_data, + projection=Projection.from_legacy(tile_wcs, pixel_frame=GeneralFrame(unit=units.pix)) + if tile_wcs + else None, + ) + # Save the image to a temporary file and transfer to final location + with ResourcePath.temporary_uri(suffix=uri.getExtension()) as temporary_uri: + im.write(temporary_uri.ospath) + uri.transfer_from(temporary_uri, transfer="copy", overwrite=True) + return + # Convert numpy array to PIL Image and save with appropriate arguments im = Image.fromarray(image_data, mode="RGB")