From 3bfd0ee79b354055bf43e8cae392400ee5dcddb4 Mon Sep 17 00:00:00 2001 From: Lee Kelvin Date: Tue, 27 Jan 2026 06:00:08 -0800 Subject: [PATCH 1/4] Add brightStarCutout.py --- .../tasks/brightStarSubtraction/__init__.py | 1 + .../brightStarSubtraction/brightStarCutout.py | 516 ++++++++++++++++++ 2 files changed, 517 insertions(+) create mode 100644 python/lsst/pipe/tasks/brightStarSubtraction/__init__.py create mode 100644 python/lsst/pipe/tasks/brightStarSubtraction/brightStarCutout.py diff --git a/python/lsst/pipe/tasks/brightStarSubtraction/__init__.py b/python/lsst/pipe/tasks/brightStarSubtraction/__init__.py new file mode 100644 index 000000000..fb8e3320d --- /dev/null +++ b/python/lsst/pipe/tasks/brightStarSubtraction/__init__.py @@ -0,0 +1 @@ +from .brightStarCutout import * diff --git a/python/lsst/pipe/tasks/brightStarSubtraction/brightStarCutout.py b/python/lsst/pipe/tasks/brightStarSubtraction/brightStarCutout.py new file mode 100644 index 000000000..d182e09ac --- /dev/null +++ b/python/lsst/pipe/tasks/brightStarSubtraction/brightStarCutout.py @@ -0,0 +1,516 @@ +# 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 . + +__all__ = ["BrightStarCutoutConnections", "BrightStarCutoutConfig", "BrightStarCutoutTask"] + +from typing import Any + +import astropy.units as u +import numpy as np +from astropy.coordinates import SkyCoord +from astropy.table import Table + +from lsst.afw.cameraGeom import FIELD_ANGLE, FOCAL_PLANE, PIXELS +from lsst.afw.detection import footprintsToNumpy +from lsst.afw.geom import makeModifiedWcs +from lsst.afw.geom.transformFactory import makeTransform +from lsst.afw.image import ExposureF, MaskedImageF +from lsst.afw.math import BackgroundList, FixedKernel, WarpingControl, warpImage +from lsst.afw.table import SourceCatalog +from lsst.daf.base import PropertyList +from lsst.daf.butler import DataCoordinate +from lsst.geom import ( + AffineTransform, + Angle, + Box2I, + Extent2D, + Extent2I, + Point2D, + Point2I, + SpherePoint, + arcseconds, + floor, + radians, +) +from lsst.meas.algorithms import ( + BrightStarStamp, + BrightStarStamps, + KernelPsf, + LoadReferenceObjectsConfig, + ReferenceObjectLoader, + WarpedPsf, +) +from lsst.pex.config import ChoiceField, ConfigField, Field, ListField +from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct +from lsst.pipe.base.connectionTypes import Input, Output, PrerequisiteInput +from lsst.utils.timer import timeMethod + +NEIGHBOR_MASK_PLANE = "NEIGHBOR" + + +class BrightStarCutoutConnections( + PipelineTaskConnections, + dimensions=("instrument", "visit", "detector"), +): + """Connections for BrightStarCutoutTask.""" + + ref_cat = PrerequisiteInput( + name="gaia_dr3_20230707", + storageClass="SimpleCatalog", + doc="Reference catalog that contains bright star positions.", + dimensions=("skypix",), + multiple=True, + deferLoad=True, + ) + input_exposure = Input( + name="visit_image", + storageClass="ExposureF", + doc="Background-subtracted input exposure from which to extract bright star stamp cutouts.", + dimensions=("visit", "detector"), + ) + input_background = Input( + name="visit_image_background", + storageClass="Background", + doc="Background model for the input exposure, to be added back on during processing.", + dimensions=("visit", "detector"), + ) + input_source_catalog = Input( + name="source_footprints", + storageClass="SourceCatalog", + doc="Source catalog containing footprints on the input exposure, used to mask neighboring sources.", + dimensions=("visit", "detector"), + ) + bright_star_stamps = Output( + name="bright_star_stamps", + storageClass="BrightStarStamps", + doc="Set of preprocessed postage stamp cutouts, each centered on a single bright star.", + dimensions=("visit", "detector"), + ) + + +class BrightStarCutoutConfig( + PipelineTaskConfig, + pipelineConnections=BrightStarCutoutConnections, +): + """Configuration parameters for BrightStarCutoutTask.""" + + # Star selection + mag_range = ListField[float]( + doc="Magnitude range in Gaia G. Cutouts will be made for all stars in this range.", + default=[10, 18], + ) + exclude_arcsec_radius = Field[float]( + doc="Stars with a star in the range ``exclude_mag_range`` mag in ``exclude_arcsec_radius`` are not " + "used.", + default=5, + ) + exclude_mag_range = ListField[float]( + doc="Stars with a star in the range ``exclude_mag_range`` mag in ``exclude_arcsec_radius`` are not " + "used.", + default=[0, 20], + ) + min_area_fraction = Field[float]( + doc="Minimum fraction of the stamp area, post-masking, that must remain for a cutout to be retained.", + default=0.1, + ) + bad_mask_planes = ListField[str]( + doc="Mask planes that identify excluded pixels for the calculation of ``min_area_fraction``.", + default=[ + "BAD", + "CR", + "CROSSTALK", + "EDGE", + "NO_DATA", + "SAT", + "SUSPECT", + "UNMASKEDNAN", + NEIGHBOR_MASK_PLANE, + ], + ) + min_focal_plane_radius = Field[float]( + doc="Minimum distance to the center of the focal plane, in mm. " + "Stars with a focal plane radius smaller than this will be omitted.", + default=0.0, + ) + max_focal_plane_radius = Field[float]( + doc="Maximum distance to the center of the focal plane, in mm. " + "Stars with a focal plane radius larger than this will be omitted.", + default=np.inf, + ) + + # Stamp geometry + stamp_size = ListField[int]( + doc="Size of the stamps to be extracted, in pixels.", + default=[251, 251], + ) + warping_kernel_name = ChoiceField[str]( + doc="Warping kernel for image data warping.", + default="lanczos5", + allowed={ + "bilinear": "bilinear interpolation", + "lanczos3": "Lanczos kernel of order 3", + "lanczos4": "Lanczos kernel of order 4", + "lanczos5": "Lanczos kernel of order 5", + }, + ) + mask_warping_kernel_name = ChoiceField[str]( + doc="Warping kernel for mask warping. Typically a more conservative kernel (e.g. with less ringing) " + "is desirable for warping masks than for warping image data.", + default="bilinear", + allowed={ + "bilinear": "bilinear interpolation", + "lanczos3": "Lanczos kernel of order 3", + "lanczos4": "Lanczos kernel of order 4", + "lanczos5": "Lanczos kernel of order 5", + }, + ) + + # Misc + load_reference_objects_config = ConfigField[LoadReferenceObjectsConfig]( + doc="Reference object loader for astrometric calibration.", + ) + + +class BrightStarCutoutTask(PipelineTask): + """Extract bright star cutouts, and warp to the same pixel grid. + + The BrightStarCutoutTask is used to extract, process, and store small image + cutouts (or "postage stamps") around bright stars. + This task essentially consists of two principal steps. + First, it identifies bright stars within an exposure using a reference + catalog and extracts a stamp around each. + Second, it shifts and warps each stamp to remove optical distortions and + sample all stars on the same pixel grid. + """ + + ConfigClass = BrightStarCutoutConfig + _DefaultName = "brightStarCutout" + config: BrightStarCutoutConfig + + def runQuantum(self, butlerQC, inputRefs, outputRefs): + inputs = butlerQC.get(inputRefs) + inputs["data_id"] = butlerQC.quantum.dataId + ref_obj_loader = ReferenceObjectLoader( + dataIds=[ref.datasetRef.dataId for ref in inputRefs.ref_cat], + refCats=inputs.pop("ref_cat"), + name=self.config.connections.ref_cat, + config=self.config.load_reference_objects_config, + ) + output = self.run(**inputs, ref_obj_loader=ref_obj_loader) + # Only ingest Stamp if it exists; prevents ingesting an empty FITS file + if output: + butlerQC.put(output, outputRefs) + + @timeMethod + def run( + self, + input_exposure: ExposureF, + input_background: BackgroundList, + input_source_catalog: SourceCatalog, + ref_obj_loader: ReferenceObjectLoader, + data_id: dict[str, Any] | DataCoordinate, + ): + """Identify bright stars within an exposure using a reference catalog, + extract stamps around each and warp/shift stamps onto a common frame. + + Parameters + ---------- + input_exposure : `~lsst.afw.image.ExposureF` + The background-subtracted image to extract bright star stamps. + input_background : `~lsst.afw.math.BackgroundList` + The background model associated with the input exposure. + input_source_catalog : `~lsst.afw.table.SourceCatalog` + The source catalog containing footprints on the input exposure. + ref_obj_loader : `~lsst.meas.algorithms.ReferenceObjectLoader` + Loader to find objects within a reference catalog. + data_id : `dict` or `~lsst.daf.butler.DataCoordinate` + The data_id of the exposure that bright stars are extracted from. + Both 'visit' and 'detector' will be persisted in the output data. + + Returns + ------- + bright_star_stamps : `~lsst.meas.algorithms.BrightStarStamps` + A set of postage stamp cutouts, each centered on a bright star. + """ + bright_stars = self._get_bright_stars(ref_obj_loader, input_exposure) + + bright_star_stamps = self._get_bright_star_stamps( + input_exposure, + input_background, + input_source_catalog, + bright_stars, + visit=int(data_id["visit"]), + ) + + return Struct(bright_star_stamps=bright_star_stamps) + + def _get_bright_stars( + self, + ref_obj_loader: ReferenceObjectLoader, + input_exposure: ExposureF, + ) -> Table: + """Get a table of bright stars from the reference catalog. + + Trim the reference catalog to only those objects within the exposure + bounding box. + Then, select bright stars based on the specified magnitude range, + isolation criteria, and optionally focal plane radius criteria. + Finally, add columns with pixel coordinates and focal plane coordinates + for each bright star. + + Parameters + ---------- + ref_obj_loader : `~lsst.meas.algorithms.ReferenceObjectLoader` + Loader to find objects within a reference catalog. + input_exposure : `~lsst.afw.image.ExposureF` + The exposure for which bright stars are being selected. + + Returns + ------- + bright_stars : `~astropy.table.Table` + Table of bright stars within the exposure. + """ + bbox = input_exposure.getBBox() + wcs = input_exposure.getWcs() + detector = input_exposure.detector + + within_region = ref_obj_loader.loadPixelBox(bbox, wcs, "phot_g_mean") + ref_cat_full = within_region.refCat + flux_field: str = within_region.fluxField + exclude_arcsec_radius = self.config.exclude_arcsec_radius * u.arcsec + + flux_range_neighbor = sorted(((self.config.exclude_mag_range * u.ABmag).to(u.nJy)).to_value()) + flux_range_candidate = sorted(((self.config.mag_range * u.ABmag).to(u.nJy)).to_value()) + + flux_min = np.min((flux_range_neighbor[0], flux_range_candidate[0])) + flux_max = np.max((flux_range_neighbor[1], flux_range_candidate[1])) + stars_subset = (ref_cat_full[flux_field] >= flux_min) & (ref_cat_full[flux_field] <= flux_max) + ref_cat_subset_columns = ("id", "coord_ra", "coord_dec", flux_field) + ref_cat_subset = Table(ref_cat_full.extract(*ref_cat_subset_columns, where=stars_subset)) + flux_subset = ref_cat_subset[flux_field] + + is_neighbor = (flux_subset >= flux_range_neighbor[0]) & (flux_subset <= flux_range_neighbor[1]) + is_candidate = (flux_subset >= flux_range_candidate[0]) & (flux_subset <= flux_range_candidate[1]) + + coords = SkyCoord(ref_cat_subset["coord_ra"], ref_cat_subset["coord_dec"], unit="rad") + coords_neighbor = coords[is_neighbor] + coords_candidate = coords[is_candidate] + + is_candidate_isolated = np.ones(len(coords_candidate), dtype=bool) + if len(coords_neighbor) > 0: + index_candidate, _, separation2d, _ = coords_candidate.search_around_sky( + coords_neighbor, exclude_arcsec_radius + ) + index_candidate = index_candidate[separation2d > 0 * u.arcsec] # Exclude self-matches + is_candidate_isolated[index_candidate] = False + + bright_stars = ref_cat_subset[is_candidate][is_candidate_isolated] + + flux_nanojansky = bright_stars[flux_field][:] * u.nJy + bright_stars["mag"] = flux_nanojansky.to(u.ABmag).to_value() # AB magnitudes + + zip_ra_dec = zip(bright_stars["coord_ra"] * radians, bright_stars["coord_dec"] * radians) + sphere_points = [SpherePoint(ra, dec) for ra, dec in zip_ra_dec] + pixel_coords = wcs.skyToPixel(sphere_points) + bright_stars["pixel_x"] = [pixel_coord.x for pixel_coord in pixel_coords] + bright_stars["pixel_y"] = [pixel_coord.y for pixel_coord in pixel_coords] + + mm_coords = detector.transform(pixel_coords, PIXELS, FOCAL_PLANE) + mm_coords_x = np.array([mm_coord.x for mm_coord in mm_coords]) + mm_coords_y = np.array([mm_coord.y for mm_coord in mm_coords]) + radius_mm = np.sqrt(mm_coords_x**2 + mm_coords_y**2) + theta_radians = np.arctan2(mm_coords_y, mm_coords_x) + bright_stars["radius_mm"] = radius_mm + bright_stars["theta_radians"] = theta_radians + + within_bbox = bright_stars["pixel_x"] >= bbox.getMinX() + within_bbox &= bright_stars["pixel_x"] <= bbox.getMaxX() + within_bbox &= bright_stars["pixel_y"] >= bbox.getMinY() + within_bbox &= bright_stars["pixel_y"] <= bbox.getMaxY() + + within_radii = bright_stars["radius_mm"] >= self.config.min_focal_plane_radius + within_radii &= bright_stars["radius_mm"] <= self.config.max_focal_plane_radius + + bright_stars = bright_stars[within_bbox & within_radii] + + self.log.info( + "Identified %i of %i reference catalog star%s that are in the field of view, are in the " + "range %s mag, and that have no neighboring stars in the range %s mag within %s arcsec.", + len(bright_stars), + len(ref_cat_full), + "" if len(ref_cat_full) == 1 else "s", + self.config.mag_range, + self.config.exclude_mag_range, + self.config.exclude_arcsec_radius, + ) + + return bright_stars + + def _get_bright_star_stamps( + self, + input_exposure: ExposureF, + input_background: BackgroundList | None, + footprints: SourceCatalog | np.ndarray, + bright_stars: Table, + visit: int | None = None, + ) -> BrightStarStamps | None: + """Extract and warp bright star stamps. + + For each bright star, extract a stamp from the input exposure centered + on the star's pixel coordinates. + Then, shift and warp the stamp to recenter on the star and align each + to the same orientation. + Finally, check the fraction of the stamp area that is masked + (e.g. due to neighboring sources or bad pixels), and only retain stamps + with sufficient unmasked area. + + Parameters + ---------- + input_exposure : `~lsst.afw.image.ExposureF` + The science image to extract bright star stamps. + input_background : `~lsst.afw.math.BackgroundList` | None + The background model associated with the input exposure. + If provided, this will be added back on to the input image. + footprints : `~lsst.afw.table.SourceCatalog` | `numpy.ndarray` + The source catalog containing footprints on the input exposure, or + a 2D numpy array with the same dimensions as the input exposure + where each pixel value corresponds to the source footprint ID. + bright_stars : `~astropy.table.Table` + Table of bright stars for which to extract stamps. + visit : `int`, optional + The visit number stored in the output `BrightStarStamp`metadata. + + Returns + ------- + bright_star_stamps : `~lsst.meas.algorithms.BrightStarStamps` | None + A set of postage stamp cutouts, each centered on a bright star. + If no bright star stamps are retained post-masking, returns `None`. + """ + warp_control = WarpingControl(self.config.warping_kernel_name, self.config.mask_warping_kernel_name) + bbox = input_exposure.getBBox() + + # Prepare masked image for warping + input_MI = input_exposure.getMaskedImage() + if input_background is not None: + input_MI += input_background.getImage() + input_MI = input_exposure.photoCalib.calibrateImage(input_MI) # to nJy + + # Generate unique footprint IDs for NEIGHBOR masking + input_MI.mask.addMaskPlane(NEIGHBOR_MASK_PLANE) + if isinstance(footprints, SourceCatalog): + footprints = footprintsToNumpy(footprints, bbox, asBool=False) + + # Establish pixel-to-boresight-pseudopixel transform + pixel_scale = input_exposure.wcs.getPixelScale(bbox.getCenter()).asArcseconds() * arcseconds + pixels_to_boresight_pseudopixels = input_exposure.detector.getTransform(PIXELS, FIELD_ANGLE).then( + makeTransform(AffineTransform.makeScaling(1 / pixel_scale.asRadians())) + ) + + # Stamp bounding boxes + stamp_radius = floor(Extent2D(*self.config.stamp_size) / 2) + stamp_bbox = Box2I(Point2I(0, 0), Extent2I(1, 1)).dilatedBy(stamp_radius) + stamp_radius_padded = floor((Extent2D(*self.config.stamp_size) * 1.42) / 2) + stamp_bbox_padded = Box2I(Point2I(0, 0), Extent2I(1, 1)).dilatedBy(stamp_radius_padded) + + stamps = [] + for bright_star in bright_stars: + pix_coord = Point2D(bright_star["pixel_x"], bright_star["pixel_y"]) + + # Set NEIGHBOR mask plane for all sources except the current one + neighbor_bit_mask = input_MI.mask.getPlaneBitMask(NEIGHBOR_MASK_PLANE) + input_MI.mask.clearMaskPlane(input_MI.mask.getMaskPlaneDict()[NEIGHBOR_MASK_PLANE]) + bright_star_id = footprints[int(pix_coord.y), int(pix_coord.x)] + neighbor_mask = (footprints != 0) & (footprints != bright_star_id) + input_MI.mask.array[neighbor_mask] |= neighbor_bit_mask + + # Define linear shifting and rotation to recenter and align stamps + boresight_pseudopixel_coord = pixels_to_boresight_pseudopixels.applyForward(pix_coord) + shift = makeTransform(AffineTransform(Point2D(0, 0) - boresight_pseudopixel_coord)) + rotation = makeTransform(AffineTransform.makeRotation(-bright_star["theta_radians"] * radians)) + pixels_to_stamp_frame = pixels_to_boresight_pseudopixels.then(shift).then(rotation) + + # Warp the image and mask to the stamp frame + stamp_MI = MaskedImageF(stamp_bbox_padded) + warpImage(stamp_MI, input_MI, pixels_to_stamp_frame, warp_control) + stamp_MI = stamp_MI[stamp_bbox] + + # Check mask coverage, update metadata + bad_bit_mask = stamp_MI.mask.getPlaneBitMask(self.config.bad_mask_planes) + good = (stamp_MI.mask.array & bad_bit_mask) == 0 + good_frac = np.sum(good) / stamp_MI.mask.array.size + if good_frac < self.config.min_area_fraction: + continue + + warped_psf = WarpedPsf(input_exposure.getPsf(), pixels_to_stamp_frame, warp_control) + stamp_psf = KernelPsf(FixedKernel(warped_psf.computeKernelImage(Point2D(0, 0)))) + stamp_wcs = makeModifiedWcs(pixels_to_stamp_frame, input_exposure.wcs, False) + + stamp = BrightStarStamp( + stamp_im=stamp_MI, + psf=stamp_psf, + wcs=stamp_wcs, + visit=visit, + detector=input_exposure.detector.getId(), + ref_id=bright_star["id"], + ref_mag=bright_star["mag"], + position=pix_coord, + focal_plane_radius=bright_star["radius_mm"], + focal_plane_angle=Angle(bright_star["theta_radians"], radians), + scale=None, + scale_err=None, + pedestal=None, + pedestal_err=None, + pedestal_scale_cov=None, + gradient_x=None, + gradient_y=None, + curvature_x=None, + curvature_y=None, + curvature_xy=None, + global_reduced_chi_squared=None, + global_degrees_of_freedom=None, + psf_reduced_chi_squared=None, + psf_degrees_of_freedom=None, + psf_masked_flux_fraction=None, + ) + stamps.append(stamp) + + num_stars = len(bright_stars) + num_excluded = num_stars - len(stamps) + percent_excluded = 100.0 * num_excluded / num_stars if num_stars > 0 else 0.0 + self.log.info( + "Extracted %i bright star stamp%s. " + "Excluded %i star%s (%.1f%%) due to a masked area fraction < %s.", + len(stamps), + "" if len(stamps) == 1 else "s", + num_excluded, + "" if num_excluded == 1 else "s", + percent_excluded, + self.config.min_area_fraction, + ) + + if not stamps: + return None + + focal_plane_radii = [stamp.focal_plane_radius for stamp in stamps] + metadata = PropertyList() + metadata["FOCAL_PLANE_RADIUS_MIN"] = np.min(focal_plane_radii) + metadata["FOCAL_PLANE_RADIUS_MAX"] = np.max(focal_plane_radii) + return BrightStarStamps(stamps, metadata=metadata) From a7ab52866c083d4616600c61841f3cc667a60498 Mon Sep 17 00:00:00 2001 From: Lee Kelvin Date: Tue, 27 Jan 2026 06:00:30 -0800 Subject: [PATCH 2/4] Add unit test for brightStarCutout.py --- tests/test_brightStarSubtraction.py | 189 ++++++++++++++++++++++++++++ 1 file changed, 189 insertions(+) create mode 100644 tests/test_brightStarSubtraction.py diff --git a/tests/test_brightStarSubtraction.py b/tests/test_brightStarSubtraction.py new file mode 100644 index 000000000..50f1035d4 --- /dev/null +++ b/tests/test_brightStarSubtraction.py @@ -0,0 +1,189 @@ +# 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 . + +import unittest + +import lsst.utils.tests +import numpy as np +from astropy.table import Table +from lsst.afw.cameraGeom import FOCAL_PLANE, PIXELS +from lsst.afw.cameraGeom.testUtils import CameraWrapper +from lsst.afw.geom import makeCdMatrix, makeSkyWcs +from lsst.afw.image import ExposureF, ImageD, ImageF, MaskedImageF, makePhotoCalibFromCalibZeroPoint +from lsst.afw.math import FixedKernel +from lsst.geom import Box2I, Extent2I, Point2D, Point2I, SpherePoint, arcseconds, degrees +from lsst.meas.algorithms import KernelPsf +from lsst.pipe.tasks.brightStarSubtraction import BrightStarCutoutConfig, BrightStarCutoutTask + + +class BrightStarSubtractionTestCase(lsst.utils.tests.TestCase): + def setUp(self): + rng = np.random.default_rng(42) + + # Background coefficients + sigma = 60.0 # noise + pedestal = 3210.1 + coef_x = 1e-2 + coef_y = 2e-2 + coef_x2 = 1e-5 + coef_xy = 2e-5 + coef_y2 = 3e-5 + + # Make an input exposure + wcs = makeSkyWcs( + crpix=Point2D(0, 0), + crval=SpherePoint(0, 0, degrees), + cdMatrix=makeCdMatrix(scale=0.2 * arcseconds, flipX=True), + ) + self.exposure = ExposureF(1001, 1001, wcs) + self.exposure.setPhotoCalib(makePhotoCalibFromCalibZeroPoint(10 ** (0.4 * 30), 1.0)) + ny, nx = self.exposure.image.array.shape + grid_y, grid_x = np.mgrid[(-ny + 1) // 2 : ny // 2 + 1, (-nx + 1) // 2 : nx // 2 + 1] + self.exposure.image.array[:] += rng.normal(scale=sigma, size=self.exposure.image.array.shape) + self.exposure.image.array += pedestal + self.exposure.image.array += coef_x * grid_x + self.exposure.image.array += coef_y * grid_y + self.exposure.image.array += coef_x2 * grid_x * grid_x + self.exposure.image.array += coef_xy * grid_x * grid_y + self.exposure.image.array += coef_y2 * grid_y * grid_y + camera = CameraWrapper().camera + detector = camera[10] + self.exposure.setDetector(detector) + for mask_plane in [ + "BAD", + "CR", + "CROSSTALK", + "EDGE", + "NO_DATA", + "SAT", + "SUSPECT", + "UNMASKEDNAN", + "NEIGHBOR", + ]: + _ = self.exposure.mask.addMaskPlane(mask_plane) + self.exposure.variance.array.fill(1.0) + + # Make a bright stars table + corners = self.exposure.wcs.pixelToSky([Point2D(x) for x in self.exposure.getBBox().getCorners()]) + corner_ras = [corner.getRa().asDegrees() for corner in corners] + corner_decs = [corner.getDec().asDegrees() for corner in corners] + num_stars = 10 + ras = rng.uniform(np.min(corner_ras), np.max(corner_ras), num_stars) + decs = rng.uniform(np.min(corner_decs), np.max(corner_decs), num_stars) + sky_coords = [SpherePoint(ra, dec, degrees) for ra, dec in zip(ras, decs)] + pixel_coords = self.exposure.wcs.skyToPixel(sky_coords) + pixel_x = [coord.getX() for coord in pixel_coords] + pixel_y = [coord.getY() for coord in pixel_coords] + mags = rng.uniform(10, 20, num_stars) + fluxes = [self.exposure.photoCalib.magnitudeToInstFlux(mag) for mag in mags] + mm_coords = detector.transform(pixel_coords, PIXELS, FOCAL_PLANE) + mm_coords_x = np.array([mm_coord.x for mm_coord in mm_coords]) + mm_coords_y = np.array([mm_coord.y for mm_coord in mm_coords]) + radius_mm = np.sqrt(mm_coords_x**2 + mm_coords_y**2) + theta_radians = np.arctan2(mm_coords_y, mm_coords_x) + self.bright_stars = Table( + { + "id": np.arange(num_stars), + "coord_ra": ras, + "coord_dec": decs, + "phot_g_mean_flux": fluxes, + "mag": mags, + "pixel_x": pixel_x, + "pixel_y": pixel_y, + "radius_mm": radius_mm, + "theta_radians": theta_radians, + } + ) + + # Make a synthetic star + stamp_radius = 25 + grid_y, grid_x = np.mgrid[-stamp_radius : stamp_radius + 1, -stamp_radius : stamp_radius + 1] + dist_from_center = np.sqrt(grid_x**2 + grid_y**2) + sigma = 1.5 + psf_array = np.exp(-(dist_from_center**2) / (2 * sigma**2)) + psf_array /= np.sum(psf_array) + fixed_kernel = FixedKernel(ImageD(psf_array)) + kernel_psf = KernelPsf(fixed_kernel) + self.exposure.setPsf(kernel_psf) + psf = kernel_psf.computeKernelImage(kernel_psf.getAveragePosition()) + + # Add synthetic stars to the exposure + footprints = ImageF(self.exposure.getBBox()) + for bright_star_id, bright_star in enumerate(self.bright_stars): + bbox_star = Box2I( + Point2I(bright_star["pixel_x"], bright_star["pixel_y"]), Extent2I(1, 1) + ).dilatedBy(stamp_radius) + bbox_star_clipped = bbox_star.clippedTo(self.exposure.getBBox()) + bright_star_im = MaskedImageF(bbox_star) + bright_star_im.image.array = bright_star["phot_g_mean_flux"] * psf.getArray() + bright_star_im = bright_star_im[bbox_star_clipped] + detection_threshold = self.exposure.getPhotoCalib().magnitudeToInstFlux(25) + detected = bright_star_im.image.array > detection_threshold + footprints[bbox_star_clipped].array = detected * (bright_star_id + 1) + _ = bright_star_im.mask.addMaskPlane("DETECTED") + bright_star_im.mask.array[detected] |= bright_star_im.mask.getPlaneBitMask("DETECTED") + bright_star_im.variance.array.fill(1.0) + self.exposure.maskedImage[bbox_star_clipped] += bright_star_im + self.footprints = footprints.array + + # Run the bright star cutout task + brightStarCutoutConfig = BrightStarCutoutConfig() + brightStarCutoutTask = BrightStarCutoutTask(config=brightStarCutoutConfig) + self.bright_star_stamps = brightStarCutoutTask._get_bright_star_stamps( + input_exposure=self.exposure, + input_background=None, + footprints=self.footprints, + bright_stars=self.bright_stars, + visit=12345, + ) + + def test_brightStarCutout(self): + """Test BrightStarCutoutTask.""" + assert self.bright_star_stamps is not None + self.assertAlmostEqual(self.bright_star_stamps.metadata["FOCAL_PLANE_RADIUS_MIN"], 5.22408977, 7) + self.assertAlmostEqual(self.bright_star_stamps.metadata["FOCAL_PLANE_RADIUS_MAX"], 14.6045832, 7) + self.assertEqual(len(self.bright_star_stamps), len(self.bright_stars)) + self.assertEqual(self.bright_star_stamps[0].visit, 12345) + self.assertEqual(self.bright_star_stamps[0].detector, 10) + + for bright_star_stamp, bright_star_row in zip(self.bright_star_stamps, self.bright_stars): + self.assertEqual(bright_star_stamp.ref_id, bright_star_row["id"]) + self.assertEqual(bright_star_stamp.ref_mag, bright_star_row["mag"]) + assert bright_star_stamp.position is not None + self.assertEqual(bright_star_stamp.position.getX(), bright_star_row["pixel_x"]) + self.assertEqual(bright_star_stamp.position.getY(), bright_star_row["pixel_y"]) + self.assertEqual(bright_star_stamp.focal_plane_radius, bright_star_row["radius_mm"]) + assert bright_star_stamp.focal_plane_angle is not None + focal_plane_angle = bright_star_stamp.focal_plane_angle.asRadians() + self.assertEqual(focal_plane_angle, bright_star_row["theta_radians"]) + + +def setup_module(module): + lsst.utils.tests.init() + + +class MemoryTestCase(lsst.utils.tests.MemoryTestCase): + pass + + +if __name__ == "__main__": + lsst.utils.tests.init() + unittest.main() From e2a2cabe7a765db1353ca1ee63a3ccc215a178fc Mon Sep 17 00:00:00 2001 From: Lee Kelvin Date: Sun, 1 Feb 2026 08:34:30 -0800 Subject: [PATCH 3/4] Add brightStarStack.py --- .../tasks/brightStarSubtraction/__init__.py | 1 + .../brightStarSubtraction/brightStarStack.py | 603 ++++++++++++++++++ 2 files changed, 604 insertions(+) create mode 100644 python/lsst/pipe/tasks/brightStarSubtraction/brightStarStack.py diff --git a/python/lsst/pipe/tasks/brightStarSubtraction/__init__.py b/python/lsst/pipe/tasks/brightStarSubtraction/__init__.py index fb8e3320d..fe5088369 100644 --- a/python/lsst/pipe/tasks/brightStarSubtraction/__init__.py +++ b/python/lsst/pipe/tasks/brightStarSubtraction/__init__.py @@ -1 +1,2 @@ from .brightStarCutout import * +from .brightStarStack import * diff --git a/python/lsst/pipe/tasks/brightStarSubtraction/brightStarStack.py b/python/lsst/pipe/tasks/brightStarSubtraction/brightStarStack.py new file mode 100644 index 000000000..ead645324 --- /dev/null +++ b/python/lsst/pipe/tasks/brightStarSubtraction/brightStarStack.py @@ -0,0 +1,603 @@ +# 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 . + +"""Stack bright star postage stamp cutouts to produce an extended PSF model.""" + +__all__ = ["BrightStarStackConnections", "BrightStarStackConfig", "BrightStarStackTask"] + +from collections.abc import Sequence + +import numpy as np +from astropy.modeling.fitting import LMLSQFitter +from astropy.modeling.models import Moffat2D + +from lsst.afw.geom import SpanSet +from lsst.afw.image import ImageF, Mask, MaskedImageF +from lsst.afw.math import StatisticsControl, statisticsStack, stringToStatisticsProperty +from lsst.daf.base import PropertyList +from lsst.meas.algorithms import BrightStarStamp, BrightStarStamps, ImagePsf +from lsst.pex.config import Field, ListField +from lsst.pipe.base import PipelineTask, PipelineTaskConfig, PipelineTaskConnections, Struct +from lsst.pipe.base.connectionTypes import Input, Output +from lsst.utils.timer import timeMethod + +from .brightStarCutout import NEIGHBOR_MASK_PLANE + + +class BrightStarStackConnections( + PipelineTaskConnections, + dimensions=("instrument", "visit", "detector"), +): + """Connections for BrightStarStackTask.""" + + bright_star_stamps = Input( + name="bright_star_stamps", + storageClass="BrightStarStamps", + doc="Set of preprocessed postage stamp cutouts, each centered on a single bright star.", + dimensions=("visit", "detector"), + ) + extended_psf = Output( + name="extended_psf", + storageClass="MaskedImageF", + doc="Extended PSF model, built from stacking bright star cutouts.", + dimensions=("visit", "detector"), + ) + extended_psf_moffat_fit = Output( + name="extended_psf_moffat_fit", + storageClass="PropertyList", + doc="Fitted Moffat2D parameters and fit statistics for the extended PSF model.", + dimensions=("visit", "detector"), + ) + + +class BrightStarStackConfig( + PipelineTaskConfig, + pipelineConnections=BrightStarStackConnections, +): + """Configuration parameters for BrightStarStackTask.""" + + # Star selection + bad_mask_planes = ListField[str]( + doc="Mask planes that identify excluded (masked) pixels.", + default=[ + "BAD", + "CR", + "CROSSTALK", + "EDGE", + "NO_DATA", + "SAT", + "SUSPECT", + "UNMASKEDNAN", + NEIGHBOR_MASK_PLANE, + ], + ) + psf_masked_flux_frac_threshold = Field[float]( + doc="Maximum allowed fraction of masked PSF flux for fitting to occur.", + default=0.97, + ) + min_focal_plane_radius = Field[float]( + doc="Minimum distance to the center of the focal plane, in mm. " + "Stars with a focal plane radius smaller than this will be omitted.", + default=0.0, + ) + max_focal_plane_radius = Field[float]( + doc="Maximum distance to the center of the focal plane, in mm. " + "Stars with a focal plane radius larger than this will be omitted.", + default=np.inf, + ) + + # Stacking control + stack_type = Field[str]( + default="MEANCLIP", + doc="Statistic name to use for stacking (from `~lsst.afw.math.Property`)", + ) + stack_num_sigma_clip = Field[float]( + doc="Number of sigma to use for clipping when stacking.", + default=3.0, + ) + stack_num_iter = Field[int]( + doc="Number of iterations to use for clipping when stacking.", + default=5, + ) + + # Fitting + use_median_variance = Field[bool]( + doc="Use the median of the variance plane for fitting.", + default=False, + ) + fit_iterations = Field[int]( + doc="Number of iterations over pedestal-gradient and scaling fit.", + default=5, + ) + bg_order = Field[int]( + doc="Order of polynomial to fit for background in bright star stamps.", + default=2, + ) + + +class BrightStarStackTask(PipelineTask): + """Stack bright star postage stamps to produce an extended PSF model.""" + + ConfigClass = BrightStarStackConfig + _DefaultName = "brightStarStack" + config: BrightStarStackConfig + + def __init__(self, initInputs=None, *args, **kwargs): + super().__init__(*args, **kwargs) + + order = self.config.bg_order + self._bg_powers = [(i, j) for i in range(order + 1) for j in range(order + 1) if i + j <= order] + + def runQuantum(self, butlerQC, inputRefs, outputRefs): + inputs = butlerQC.get(inputRefs) + output = self.run(**inputs) + # Guard against empty outputs, resulting from failed fitting runs + if output: + butlerQC.put(output, outputRefs) + + @timeMethod + def run( + self, + bright_star_stamps: BrightStarStamps, + ): + """Run the BrightStarStackTask to produce an extended PSF model from + bright star stamps. + + Parameters + ---------- + bright_star_stamps : `~lsst.meas.algorithms.BrightStarStamps + Set of preprocessed postage stamp cutouts, each centered on a + single bright star. + + Returns + ------- + extended_psf : `~lsst.afw.image.MaskedImageF` or `None` + The extended PSF model produced by stacking bright star stamps, or + `None` if the stacking was unsuccessful (e.g., no valid stamps). + """ + stamps = self._get_stamps(bright_star_stamps) + if len(stamps) == 0: + return None + extended_psf = self._make_extended_psf(stamps) + moffat_results = self._fit_moffat(extended_psf) + + return Struct(extended_psf=extended_psf, moffat_results=moffat_results) + + def _get_stamps(self, bright_star_stamps: BrightStarStamps) -> list[BrightStarStamp]: + """Get bright star stamps that are within the specified fp radius. + + Parameters + ---------- + bright_star_stamps : `~lsst.meas.algorithms.BrightStarStamps + Set of preprocessed postage stamp cutouts, each centered on a + single bright star. + + Returns + ------- + stamps : `list`[`~lsst.meas.algorithms.BrightStarStamp] + List of bright star stamps that are within the specified focal + plane radii. + """ + stamps = [ + stamp + for stamp in bright_star_stamps + if stamp.focal_plane_radius is not None + and self.config.min_focal_plane_radius + <= stamp.focal_plane_radius + <= self.config.max_focal_plane_radius + ] + if len(stamps) == 0: + self.log.warning( + "No bright star stamps found within the focal plane radius range [%s, %s] mm.", + self.config.min_focal_plane_radius, + self.config.max_focal_plane_radius, + ) + elif len(stamps) < len(bright_star_stamps): + self.log.info( + "Only %d of %d bright star stamps are within the focal plane radius range [%s, %s] mm.", + len(stamps), + len(bright_star_stamps), + self.config.min_focal_plane_radius, + self.config.max_focal_plane_radius, + ) + self.log.info( + "Constructing an extended PSF from %d bright star stamp%s.", + len(stamps), + "s" if len(stamps) > 1 else "", + ) + return stamps + + def _make_extended_psf(self, stamps: list[BrightStarStamp]) -> MaskedImageF: + """Stack bright star stamps to produce an extended PSF model. + + Parameters + ---------- + stamps : `list`[`~lsst.meas.algorithms.BrightStarStamp] + List of bright star stamps to stack. + + Returns + ------- + extended_psf : `~lsst.afw.image.MaskedImageF` + The extended PSF model produced by stacking bright star stamps. + """ + stack_type_property = stringToStatisticsProperty(self.config.stack_type) + statistics_control = StatisticsControl( + numSigmaClip=self.config.stack_num_sigma_clip, numIter=self.config.stack_num_iter + ) + bad_mask_bit_mask = stamps[0].stamp_im.mask.getPlaneBitMask(self.config.bad_mask_planes) + statistics_control.setAndMask(bad_mask_bit_mask) + + psf_amplitudes: dict[int, float | None] = {stamp_id: None for stamp_id in range(len(stamps))} + + extended_psf = None + for iteration in range(self.config.fit_iterations): + stamp_ims_normalized = [] + for stamp_id, stamp in enumerate(stamps): + masked_image_normalized, psf_amplitude = self._fit_and_normalize( + stamp.stamp_im, + stamp.psf if extended_psf is None else extended_psf, + psf_amplitudes[stamp_id], + ) + if masked_image_normalized is not None and psf_amplitude is not None: + stamp_ims_normalized.append(masked_image_normalized) + psf_amplitudes[stamp_id] = psf_amplitude + + if len(stamp_ims_normalized) == 0: + self.log.warning( + "Iteration %d: No stamps were successfully fit and normalized. " + "No further iterations will be attempted.", + iteration + 1, + ) + break + else: + self.log.info( + "Iteration %d: Successfully fit and normalized %d out of %d stamp%s using the %s model.", + iteration + 1, + len(stamp_ims_normalized), + len(stamps), + "s" if len(stamps) > 1 else "", + "baseline PSF" if extended_psf is None else "extended PSF", + ) + + # Stack normalized stamps from an iteration to make an extended PSF + extended_psf = MaskedImageF(stamp_ims_normalized[0].getBBox()) + statisticsStack(extended_psf, stamp_ims_normalized, stack_type_property, statistics_control) + + return extended_psf + + def _fit_and_normalize( + self, + masked_image: MaskedImageF, + psf: ImagePsf | MaskedImageF, + psf_amplitude: float | None, + ) -> tuple[MaskedImageF | None, float | None]: + """Fit the PSF model to a single bright star stamp. + + Parameters + ---------- + masked_image : `~lsst.afw.image.MaskedImageF` + The masked image of the bright star stamp to fit. + psf : `~lsst.meas.algorithms.ImagePsf` | `~lsst.afw.image.MaskedImageF` + The PSF model to fit to the data. This can be either an `ImagePsf` + or a `MaskedImageF` (from, for example, a prior iteration). + psf_amplitude : `float` | `None` + The amplitude to scale the PSF image by before subtraction in the + background fitting step. If `None`, no PSF subtraction is performed + during background fitting. + + Returns + ------- + masked_image_normalized : `~lsst.afw.image.MaskedImageF` | `None` + The masked image of the bright star stamp after fitting and + normalization, or `None` if the fit was unsuccessful. + psf_amplitude : `float` | `None` + The fitted amplitude of the PSF model, or `None` if the fit was + unsuccessful. + """ + if isinstance(psf, ImagePsf): + psf_image = psf.computeKernelImage(psf.getAveragePosition()) + else: + psf_image = psf.image # Assumed to be warped, center at [0,0] + + bit_mask_bad = masked_image.mask.getPlaneBitMask(self.config.bad_mask_planes) + + # Calculate the fraction of the PSF image flux masked by bad pixels + psf_mask = ImageF(psf_image.getBBox()) + psf_mask.array[:, :] = (masked_image.mask[psf_image.getBBox()].array & bit_mask_bad).astype(bool) + psf_masked_flux_frac = ( + np.dot(psf_image.array.flat, psf_mask.array.flat).astype(np.float64) / psf_image.array.sum() + ) + if psf_masked_flux_frac > self.config.psf_masked_flux_frac_threshold: + return None, None # Handle cases where the PSF image is mostly masked + + # Fit the background + bg_image = self._fit_bg( + masked_image, list(self.config.bad_mask_planes) + ["DETECTED"], psf_image, psf_amplitude + ) + if bg_image is None: + return None, None + masked_image_fit = masked_image.clone() + masked_image_fit.image -= bg_image + + # Fit the PSF amplitude + psf_image_padded = ImageF(masked_image.getBBox()) + psf_image_padded[psf_image.getBBox()] = psf_image.convertF() + psf_amplitude = self._fit_psf_amplitude( + masked_image_fit, psf_image_padded, self.config.bad_mask_planes + ) + if psf_amplitude is None or psf_amplitude <= 0: + return None, None + masked_image_fit /= psf_amplitude + + return masked_image_fit, psf_amplitude + + def _fit_bg( + self, + masked_image: MaskedImageF, + mask_planes: Sequence[str], + psf_image: ImageF, + psf_amplitude: float | None, + ) -> ImageF | None: + """Fit a polynomial background to a bright star stamp, optionally + subtracting a scaled PSF model. + + Parameters + ---------- + masked_image : `~lsst.afw.image.MaskedImageF` + The masked image of the bright star stamp to fit. + mask_planes : `Sequence`[`str`] + Sequence of mask planes to use for identifying bad pixels. + psf_image : `~lsst.afw.image.ImageF` + The PSF image to optionally subtract from the data before fitting. + psf_amplitude : `float` | `None` + The amplitude to scale the PSF image by before subtraction. + If `None`, no PSF subtraction is performed. + + Returns + ------- + bg_image : `~lsst.afw.image.ImageF` | `None` + The fitted background image, or `None` if the fit was unsuccessful. + """ + if psf_amplitude is not None: + psf_image_scaled = psf_image.clone() + psf_image_scaled.array *= psf_amplitude + + spans, image, sigma = self._get_span_data( + masked_image, mask_planes, subtract_image=psf_image_scaled if psf_amplitude is not None else None + ) + spans_yx = spans.indices() + spans_x = spans_yx[1, :] + spans_y = spans_yx[0, :] + + design_matrix = np.ones((len(image), len(self._bg_powers)), dtype=float) + for k, (i, j) in enumerate(self._bg_powers): + if i == j == 0: + design_matrix[:, k] /= sigma # constant term + else: + design_matrix[:, k] = (spans_x**i * spans_y**j) / sigma + + solutions, _ = self._solve_design_matrix(design_matrix, image) + if solutions is None: + return None + + bbox = masked_image.getBBox() + grid_x, grid_y = np.meshgrid(bbox.getX().arange(), bbox.getY().arange()) + bg_array = sum(s * (grid_x**i * grid_y**j) for s, (i, j) in zip(solutions, self._bg_powers)) + + return ImageF(bg_array.astype(np.float32), xy0=masked_image.getXY0()) + + def _fit_psf_amplitude( + self, + masked_image: MaskedImageF, + psf_image: ImageF, + mask_planes: Sequence[str], + mask_zeros: bool = True, + ) -> float | None: + """Fit the amplitude of the PSF model to a bright star stamp. + + Parameters + ---------- + masked_image : `~lsst.afw.image.MaskedImageF` + The masked image of the bright star stamp to fit. + psf_image : `~lsst.afw.image.ImageF` + The PSF image to fit to the data. + mask_planes : `Sequence`[`str`] + Sequence of mask planes to use for identifying bad pixels. + mask_zeros : `bool`, optional + Whether to mask pixels where the PSF image is zero, in addition to + the bad pixels identified by the mask planes. + + Returns + ------- + psf_amplitude : `float` | `None` + The fitted amplitude of the PSF model, or `None` if the fit was + unsuccessful. + + Notes + ----- + NaN and negative values in the PSF image are always masked, in addition + to the specified mask planes and optionally zero-valued pixels. + """ + psf_mask_array = np.isnan(psf_image.array) | (psf_image.array < 0) + if mask_zeros: + psf_mask_array |= psf_image.array == 0 + psf_mask = Mask(masked_image.getBBox()) # type: ignore + psf_mask.array = psf_mask_array + + spans, image, sigma = self._get_span_data(masked_image, mask_planes, psf_mask) + psf = spans.flatten(psf_image.array, psf_image.getXY0()) + psf /= sigma + design_matrix = psf[:, None] # A single column for amplitude fit alone + + solutions, _ = self._solve_design_matrix(design_matrix, image) + if solutions is None: + return None + + return float(solutions[0]) + + def _get_span_data( + self, + masked_image: MaskedImageF, + mask_planes: Sequence[str], + additional_mask: Mask | None = None, + subtract_image: ImageF | None = None, + ) -> tuple[SpanSet, np.ndarray, np.ndarray | float]: + """Get the data and corresponding spans for fitting. + + Parameters + ---------- + masked_image : `~lsst.afw.image.MaskedImageF` + The masked image from which to extract data for fitting. + mask_planes : `Sequence`[`str`] + Sequence of mask planes to use for identifying bad pixels. + additional_mask : `~lsst.afw.image.Mask` or `None`, optional + An additional mask to combine with the masked_image's mask. + subtract_image : `~lsst.afw.image.ImageF` or `None`, optional + An image to subtract from the masked_image data before fitting + (e.g., a PSF model). If `None`, no subtraction is performed. + + Returns + ------- + spans : `~lsst.afw.geom.SpanSet` + The spans corresponding to the good pixels to be used for fitting. + image : `numpy.ndarray` + The pixel values from the masked_image data (after optional + subtraction) corresponding to the derived span sets. + sigma : `numpy.ndarray` | `float` + The pixel-wise uncertainties from the masked_image's variance + corresponding to the derived span sets. + This can be a single value if `use_median_variance` is `True`. + """ + bit_mask = masked_image.mask.getPlaneBitMask(mask_planes) + bad_spans = SpanSet.fromMask(masked_image.mask, bit_mask) + if additional_mask is not None: + additional_bit_mask = additional_mask.getPlaneBitMask(mask_planes) + additional_bad_spans = SpanSet.fromMask(additional_mask, additional_bit_mask) + bad_spans = bad_spans.union(additional_bad_spans) + spans = SpanSet(masked_image.getBBox()).intersectNot(bad_spans) + image_data = masked_image.image.array - (subtract_image.array if subtract_image is not None else 0) + image = spans.flatten(image_data, masked_image.getXY0()) + variance = spans.flatten(masked_image.variance.array, masked_image.getXY0()) + if self.config.use_median_variance: + variance = np.median(variance) + sigma = np.sqrt(variance) + sigma[sigma == 0] = np.inf # Guard against zero variance pixels + return spans, image / sigma, sigma + + def _solve_design_matrix( + self, + design_matrix: np.ndarray, + data: np.ndarray, + calculate_covariance: bool = False, + ) -> tuple[np.ndarray | None, np.ndarray | None]: + """Solve the linear system for the given design matrix and data. + + Parameters + ---------- + design_matrix : `numpy.ndarray` + The design matrix for the linear system. Each column corresponds to + a basis function evaluated at the data points. The number of rows + should match the length of the data vector. + data : `numpy.ndarray` + The data vector for the linear system. The length should match the + number of rows in the design matrix. + calculate_covariance : `bool`, optional + Whether to calculate and return the covariance matrix of the + solutions. + + Returns + ------- + solutions : `numpy.ndarray` or `None` + The solutions to the linear system, or `None` if the system could + not be solved. + covariance_matrix : `numpy.ndarray` or `None` + The covariance matrix of the solutions, or `None` if the system + could not be solved. Only returned if `calculate_covariance` is + `True`, otherwise `None`. + """ + covariance_matrix = None + try: + solutions, sum_squared_residuals, *_ = np.linalg.lstsq(design_matrix, data, rcond=None) + if calculate_covariance: + covariance_matrix = np.linalg.pinv(design_matrix.T @ design_matrix) + except np.linalg.LinAlgError: + return None, None # Handle singular matrix errors + if sum_squared_residuals.size == 0: + return None, None # Handle cases where sum of the squared residuals are empty + return solutions, covariance_matrix + + def _fit_moffat(self, extended_psf: MaskedImageF, fix_center: bool = False) -> PropertyList: + """Fit a Moffat2D model to the extended PSF image. + + Parameters + ---------- + extended_psf : `~lsst.afw.image.MaskedImageF` + The extended PSF image to fit. + fix_center : `bool`, optional + Whether to fix the center of the Moffat2D model to the center of + the extended PSF image. If `False`, the center will be allowed to + vary in the fit. + + Returns + ------- + moffat_results : `~lsst.daf.base.PropertyList` + A PropertyList containing the fitted Moffat2D parameters and fit + statistics (chi-squared and reduced chi-squared). + """ + bbox = extended_psf.getBBox() + grid_x, grid_y = np.meshgrid(bbox.getX().arange(), bbox.getY().arange()) + + extended_psf_image = extended_psf.image.array + extended_psf_sigma = np.sqrt(extended_psf.variance.array) + + fitter = LMLSQFitter() + moffat_init = Moffat2D( + amplitude=extended_psf_image.max(), + x_0=bbox.getCenter().x, + y_0=bbox.getCenter().y, + gamma=3.0, + alpha=2.5, + fixed={"x_0": fix_center, "y_0": fix_center}, + ) + weights = 1.0 / np.clip(extended_psf_sigma, 1e-12, None) + moffat_fit = fitter(moffat_init, grid_x, grid_y, extended_psf_image, weights=weights) + + residuals = extended_psf_image - moffat_fit(grid_x, grid_y) + chi2 = np.sum((residuals / extended_psf_sigma) ** 2) + dof = extended_psf_image.size - len(moffat_fit.parameters) + reduced_chi2 = chi2 / dof + + moffat_results = PropertyList() + for name, value in zip(moffat_fit.param_names, moffat_fit.parameters): + moffat_results.setDouble(f"MOFFAT_{name.upper()}", float(value)) + moffat_results.setDouble("MOFFAT_CHI2", float(chi2)) + moffat_results.setDouble("MOFFAT_REDUCED_CHI2", float(reduced_chi2)) + + self.log.info( + "Extended PSF Moffat fit results: x_0=%.2f, y_0=%.2f, gamma=%.2f, alpha=%.2f, reduced_chi2=%.2f", + moffat_fit.x_0.value, + moffat_fit.y_0.value, + moffat_fit.gamma.value, + moffat_fit.alpha.value, + reduced_chi2, + ) + + return moffat_results From 1d8c6f66b49e7a0e626462dc2e0bc1009b9bf986 Mon Sep 17 00:00:00 2001 From: Lee Kelvin Date: Sun, 1 Feb 2026 08:34:43 -0800 Subject: [PATCH 4/4] Add unit test for brightStarStack.py --- tests/test_brightStarSubtraction.py | 30 ++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/test_brightStarSubtraction.py b/tests/test_brightStarSubtraction.py index 50f1035d4..78413dc38 100644 --- a/tests/test_brightStarSubtraction.py +++ b/tests/test_brightStarSubtraction.py @@ -31,7 +31,12 @@ from lsst.afw.math import FixedKernel from lsst.geom import Box2I, Extent2I, Point2D, Point2I, SpherePoint, arcseconds, degrees from lsst.meas.algorithms import KernelPsf -from lsst.pipe.tasks.brightStarSubtraction import BrightStarCutoutConfig, BrightStarCutoutTask +from lsst.pipe.tasks.brightStarSubtraction import ( + BrightStarCutoutConfig, + BrightStarCutoutTask, + BrightStarStackConfig, + BrightStarStackTask, +) class BrightStarSubtractionTestCase(lsst.utils.tests.TestCase): @@ -154,6 +159,15 @@ def setUp(self): bright_stars=self.bright_stars, visit=12345, ) + assert self.bright_star_stamps is not None + + # Run the bright star stack task + brightStarStackConfig = BrightStarStackConfig() + brightStarStackTask = BrightStarStackTask(config=brightStarStackConfig) + bss_results = brightStarStackTask.run(bright_star_stamps=self.bright_star_stamps) + assert bss_results is not None + self.extended_psf = bss_results.extended_psf + self.moffat_results = bss_results.moffat_results def test_brightStarCutout(self): """Test BrightStarCutoutTask.""" @@ -175,6 +189,20 @@ def test_brightStarCutout(self): focal_plane_angle = bright_star_stamp.focal_plane_angle.asRadians() self.assertEqual(focal_plane_angle, bright_star_row["theta_radians"]) + def test_brightStarStack(self): + """Test BrightStarStackTask.""" + assert self.extended_psf is not None + assert self.moffat_results is not None + self.assertAlmostEqual(np.sum(self.extended_psf.image.array), 0.8233416, places=5) + self.assertAlmostEqual(np.sum(self.extended_psf.variance.array), 0.007561891, places=5) + self.assertAlmostEqual(self.moffat_results["MOFFAT_AMPLITUDE"], 0.078900464260488, places=5) + self.assertAlmostEqual(self.moffat_results["MOFFAT_X_0"], -0.68834523633912, places=5) + self.assertAlmostEqual(self.moffat_results["MOFFAT_Y_0"], -0.069005412739451, places=5) + self.assertAlmostEqual(self.moffat_results["MOFFAT_GAMMA"], 8.0966823485900, places=5) + self.assertAlmostEqual(self.moffat_results["MOFFAT_ALPHA"], 16.048683662812, places=5) + self.assertAlmostEqual(self.moffat_results["MOFFAT_CHI2"], 107652.97393353, places=5) + self.assertAlmostEqual(self.moffat_results["MOFFAT_REDUCED_CHI2"], 1.7088858647141, places=5) + def setup_module(module): lsst.utils.tests.init()