diff --git a/python/lsst/source/injection/inject_base.py b/python/lsst/source/injection/inject_base.py index 3beb7fb..597bf21 100644 --- a/python/lsst/source/injection/inject_base.py +++ b/python/lsst/source/injection/inject_base.py @@ -117,6 +117,16 @@ class BaseInjectConfig(PipelineTaskConfig, pipelineConnections=BaseInjectConnect doc="String to prefix to the entries in the *col_stamp* column, for example, a directory path.", default="", ) + add_noise = Field[bool]( + doc="Whether to randomly vary the injected flux in each pixel by an amount consistent with " + "the injected variance.", + default=True, + ) + noise_seed = Field[int]( + doc="Initial seed for random noise generation. This value increments by 1 for each injected " + "object, so each object has an independent noise realization.", + default=0, + ) # Custom column names. col_ra = Field[str]( @@ -163,7 +173,7 @@ class BaseInjectTask(PipelineTask): _DefaultName = "baseInjectTask" ConfigClass = BaseInjectConfig - def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs): + def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs, variance_scale=0.0): """Inject sources into an image. Parameters @@ -179,6 +189,9 @@ def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs): Photometric calibration used to calibrate injected sources. wcs : `lsst.afw.geom.SkyWcs` WCS used to calibrate injected sources. + variance_scale : `float` + Scale by which to multiply injected image flux to determine the + amount of variance to add. Returns ------- @@ -267,6 +280,9 @@ def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs): mask_plane_name=self.config.mask_plane_name, calib_flux_radius=self.config.calib_flux_radius, draw_size_max=10000, # TODO: replace draw_size logic with GS logic. + variance_scale=self.config.variance_scale, + add_noise=self.config.add_noise, + noise_seed=self.config.noise_seed, logger=self.log, ) # Add inject_galsim_objects_into_exposure outputs into output cat. diff --git a/python/lsst/source/injection/inject_coadd.py b/python/lsst/source/injection/inject_coadd.py index 942fa6d..d12cfa5 100644 --- a/python/lsst/source/injection/inject_coadd.py +++ b/python/lsst/source/injection/inject_coadd.py @@ -23,7 +23,12 @@ __all__ = ["CoaddInjectConnections", "CoaddInjectConfig", "CoaddInjectTask"] +import numpy as np +from lsst.pex.config import Field from lsst.pipe.base.connectionTypes import Input, Output +from sklearn.cluster import KMeans +from sklearn.linear_model import LinearRegression, RANSACRegressor +from sklearn.metrics import mean_squared_error from .inject_base import BaseInjectConfig, BaseInjectConnections, BaseInjectTask @@ -63,7 +68,58 @@ class CoaddInjectConfig( # type: ignore [call-arg] ): """Coadd-level configuration for source injection tasks.""" - pass + n_fits_1 = Field[int]( + doc="Perform this many RANSAC fits in the first round, to get a sample " + "of different slopes based on the different random samples of points used " + "in the fit.", + default=20, + ) + n_fits_2 = Field[int]( + doc="Perform this many RANSAC fits in the second round, to get a sample " + "of different slopes based on the different random samples of points used " + "in the fit.", + default=20, + ) + threshold_scale_1 = Field[float]( + doc="An outlier in the first RANSAC fit is farther from the fit line, " + "in terms of squared error, than this multiple of the initial linear MSE.", + default=0.1, + ) + threshold_scale_2 = Field[float]( + doc="An outlier in the second RANSAC fit is farther from the fit line, " + "in terms of squared error, than this multiple of the initial linear MSE.", + default=0.1, + ) + max_trials_1 = Field[int]( + doc="Maximum number of trials the first RANSAC fit is allowed to run.", + default=1000, + ) + max_trials_2 = Field[int]( + doc="Maximum number of trials the second RANSAC fit is allowed to run.", + default=1000, + ) + variance_fit_seed_1 = Field[int](doc="Seed for first RANSAC fit of flux vs. variance.", default=0) + variance_fit_seed_2 = Field[int](doc="Seed for second RANSAC fit of flux vs. variance.", default=0) + n_clusters_1 = Field[int]( + doc="K-means cluster the first set of RANSAC fits using this many clusters, " + "in order to find the most stable slope (biggest cluster).", + default=4, + ) + n_clusters_2 = Field[int]( + doc="K-means cluster the second set of RANSAC fits using this many clusters, " + "in order to find the most stable slope (biggest cluster).", + default=3, + ) + kmeans_seed_1 = Field[int](doc="Seed for first round of k-means clustering.", default=0) + kmeans_seed_2 = Field[int](doc="Seed for second round of k-means clustering.", default=0) + kmeans_n_init_1 = Field[int]( + doc="Number of times the first k-means clustering is run with different initial centroids.", + default=10, + ) + kmeans_n_init_2 = Field[int]( + doc="Number of times the second k-means clustering is run with different initial centroids.", + default=10, + ) class CoaddInjectTask(BaseInjectTask): @@ -72,6 +128,13 @@ class CoaddInjectTask(BaseInjectTask): _DefaultName = "coaddInjectTask" ConfigClass = CoaddInjectConfig + def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs): + self.log.info("Fitting flux vs. variance in each pixel.") + variance_scale = self.get_variance_scale(input_exposure) + self.log.info("Variance scale factor: %.6f", variance_scale) + + return super().run(injection_catalogs, input_exposure, psf, photo_calib, wcs, variance_scale) + def runQuantum(self, butler_quantum_context, input_refs, output_refs): inputs = butler_quantum_context.get(input_refs) @@ -82,3 +145,175 @@ def runQuantum(self, butler_quantum_context, input_refs, output_refs): input_keys = ["injection_catalogs", "input_exposure", "sky_map", "psf", "photo_calib", "wcs"] outputs = self.run(**{key: value for (key, value) in inputs.items() if key in input_keys}) butler_quantum_context.put(outputs, output_refs) + + """ + Establish the variance scale by a linear fit of variance vs. flux. + In practice we see that most pixels in a coadd obey a consistent, simple + linear relationship between variance and flux, but a small sample do not. + + To identify and hence ignore these odd pixels, we perform two rounds of + RANSAC fits. RANSAC iteratively finds a least-squares straight line fit on + random subsamples of points, while identifying outliers. The inlier pixels + from a first RANSAC fit tend to be regions of empty space and the extreme + outer edges of galaxies, while the outliers tend to be the inner regions of + galaxies. + + We can pick up some more pixels in the galaxies by running a second RANSCAC + fit on just the outliers of the first fit. In the second round, the inliers + tend to be the bulk of the galaxy, while the outliers are the innermost + cores of galaxies. In practice, the inliers from this second round tend to + have a qualitatively similar variance-vs-flux relationship to the outliers + from the first fit and do not strongly alter the final variance scale we + get; while the outliers from the second round are truly wild and should + clearly be omitted from a simple linear fit. + + In both rounds, random variation of the points sampled by the RANSAC fit + causes the fitted slope and intercept to vary, and sometimes the fit can + settle on a pathological sample of points as its inliers. We try to avoid + such pathologies by running each fit multiple times with different seeds, + and using K-Means clustering on the resulting slopes to identify the most + stable value. + """ + + def get_variance_scale(self, exposure): + flux = exposure.image.array.ravel() + var = exposure.variance.array.ravel() + + # Ignore pixels with nan or infinite values + good_pixels = np.isfinite(flux) & np.isfinite(var) + flux = flux[good_pixels] + var = var[good_pixels] + + # Simple linear regression to establish MSE. + linear = LinearRegression() + linear.fit(flux.reshape(-1, 1), var) + linear_mse = mean_squared_error(var, linear.predict(flux.reshape(-1, 1))) + + # First RANSAC fit + fit_results = [] + for seed in range( + self.config.variance_fit_seed_1, self.config.variance_fit_seed_1 + self.config.n_fits_1 + ): + ransac = RANSACRegressor( + loss="squared_error", + residual_threshold=self.config.threshold_scale_1 * linear_mse, + max_trials=self.config.max_trials_1, + random_state=seed, + ) + ransac.fit(flux.reshape(-1, 1), var) + # Remember fit results + slope = ransac.estimator_.coef_[0] + fit_results.append((slope, seed)) + + # K-means cluster the first round of fits, + # to find the most stable results. + kmeans = KMeans( + n_clusters=self.config.n_clusters_1, + random_state=self.config.kmeans_seed_1, + n_init=self.config.kmeans_n_init_1, + ) + kmeans.fit(np.log(np.array([f[0] for f in fit_results if f[0] > 0])).reshape(-1, 1)) + label_counts = [np.sum(kmeans.labels_ == idx) for idx in range(self.config.n_clusters_1)] + + # Recall one of the fits, chosen arbitrarily from those which are both + # stable, according to the first k-means fit, and positive-slope. + stable_fit_seeds = np.array([f[1] for f in fit_results if f[0] > 0])[ + kmeans.labels_ == np.argmax(label_counts) + ] + if len(stable_fit_seeds == 0): + # No positive-slope fit found. + # Allow the fitted slope to be negative but throw a warning. + self.log.warning( + "No positive-slope result in the first round of " + "RANSAC fits. Proceeding with a negative-slope fit." + ) + stable_fit_seeds = np.array([f[1] for f in fit_results])[ + kmeans.labels_ == np.argmax(label_counts) + ] + seed = stable_fit_seeds[0] + ransac = RANSACRegressor( + loss="squared_error", + residual_threshold=self.config.threshold_scale_1 * linear_mse, + max_trials=self.config.max_trials_1, + random_state=seed, + ) + ransac.fit(flux.reshape(-1, 1), var) + + # Label the pixels with a "good" variance vs. flux relationship + # (the inliers), together with the ones that are further from a simple + # straight line (the outliers). + inlier_mask_1 = ransac.inlier_mask_ + outlier_mask_1 = ~inlier_mask_1 + + # Second RANSAC fit, + # on just the outliers of the 1st fit. + fit_results = [] + for seed in range( + self.config.variance_fit_seed_2, self.config.variance_fit_seed_2 + self.config.n_fits_2 + ): + ransac = RANSACRegressor( + loss="squared_error", + residual_threshold=self.config.threshold_scale_2 * linear_mse, + max_trials=self.config.max_trials_2, + random_state=seed, + ) + ransac.fit(flux[outlier_mask_1].reshape(-1, 1), var[outlier_mask_1]) + # Remember fit results + slope = ransac.estimator_.coef_[0] + fit_results.append((slope, seed)) + + # K-Means cluster the second round of fits, + # to find the most stable result + kmeans = KMeans( + n_clusters=self.config.n_clusters_2, + random_state=self.config.kmeans_seed_2, + n_init=self.config.kmeans_n_init_2, + ) + kmeans.fit(np.log(np.array([f[0] for f in fit_results if f[0] > 0])).reshape(-1, 1)) + label_counts = [np.sum(kmeans.labels_ == idx) for idx in range(self.config.n_clusters_2)] + + # Recall one of the stable fits + stable_fit_seeds = np.array([f[1] for f in fit_results if f[0] > 0])[ + kmeans.labels_ == np.argmax(label_counts) + ] + if len(stable_fit_seeds == 0): + # No positive-slope fit found. + # Allow the fitted slope to be negative but throw a warning. + self.log.warning( + "No positive-slope result in the second round of " + "RANSAC fits. Proceeding with a negative-slope fit." + ) + stable_fit_seeds = np.array([f[1] for f in fit_results])[ + kmeans.labels_ == np.argmax(label_counts) + ] + seed = stable_fit_seeds[0] + ransac = RANSACRegressor( + loss="squared_error", + residual_threshold=self.config.threshold_scale_2 * linear_mse, + max_trials=self.config.max_trials_2, + random_state=seed, + ) + ransac.fit(flux[outlier_mask_1].reshape(-1, 1), var[outlier_mask_1]) + + # Pixels with a "good" variance vs. flux relationship: + # Union of the inliers from the first fit + # together with the inliers from the second fit. + flux_good = np.concatenate( + (flux[inlier_mask_1], flux[outlier_mask_1][ransac.inlier_mask_]), + axis=None, + ) + var_good = np.concatenate( + (var[inlier_mask_1], var[outlier_mask_1][ransac.inlier_mask_]), + axis=None, + ) + + # Fit all the good pixels with a simple least squares regression. + linear = LinearRegression() + linear.fit(flux_good.reshape(-1, 1), var_good) + variance_scale = float(linear.coef_[0]) + + if variance_scale < 0: + self.log.warning("Slope of final variance vs. flux fit is negative.") + + # Return the slope of the final fit. + return variance_scale diff --git a/python/lsst/source/injection/inject_engine.py b/python/lsst/source/injection/inject_engine.py index aceb2e0..488411f 100644 --- a/python/lsst/source/injection/inject_engine.py +++ b/python/lsst/source/injection/inject_engine.py @@ -389,6 +389,9 @@ def inject_galsim_objects_into_exposure( mask_plane_name: str = "INJECTED", calib_flux_radius: float = 12.0, draw_size_max: int = 1000, + variance_scale: float = 0.0, + add_noise: bool = True, + noise_seed: int = 0, logger: Any | None = None, ) -> tuple[list[int], list[galsim.BoundsI], list[bool], list[bool]]: """Inject sources into given exposure using GalSim. @@ -411,6 +414,14 @@ def inject_galsim_objects_into_exposure( draw_size_max : `int` Maximum allowed size of the drawn object. If the object is larger than this, the draw size will be clipped to this size. + variance_scale : `float` + Factor by which to multiply injected image flux to obtain the injected + variance level. + add_noise : `bool` + Whether to randomly vary the amount of injected image flux by an amount + consistent with the amount of injected variance. + noise_seed : `int` + Initial seed for random noise generation. logger : `lsst.utils.logging.LsstLogAdapter`, optional Logger to use for logging messages. @@ -439,6 +450,7 @@ def inject_galsim_objects_into_exposure( bbox = exposure.getBBox() full_bounds = galsim.BoundsI(bbox.minX, bbox.maxX, bbox.minY, bbox.maxY) galsim_image = galsim.Image(exposure.image.array, bounds=full_bounds) + galsim_variance = galsim.Image(exposure.variance.array, bounds=full_bounds) pixel_scale = wcs.getPixelScale(bbox.getCenter()).asArcseconds() draw_sizes: list[int] = [] @@ -525,13 +537,17 @@ def inject_galsim_objects_into_exposure( # Inject the source if there is any overlap. if object_common_bounds.area() > 0: common_image = galsim_image[object_common_bounds] + common_variance = galsim_variance[object_common_bounds] offset = posd - object_common_bounds.true_center + # Attempt to draw a smooth version of the image, + # representing an expected light profile. # Note, for calexp injection, pixel is already part of the PSF and # for coadd injection, it's incorrect to include the output pixel. # So for both cases, we draw using method='no_pixel'. + image_template = common_image.copy() try: - conv.drawImage( - common_image, add_to_image=True, offset=offset, wcs=galsim_wcs, method="no_pixel" + image_template = conv.drawImage( + image_template, add_to_image=False, offset=offset, wcs=galsim_wcs, method="no_pixel" ) except GalSimFFTSizeError as err: fft_size_errors[i] = True @@ -542,6 +558,42 @@ def inject_galsim_objects_into_exposure( err, ) continue + + # Set a variance level in each pixel + # corresponding to the drawn light profile. + variance_template = image_template.copy() + variance_template *= variance_scale + + if add_noise: + # For generating noise, + # variance must be meaningful. + if np.any(variance_template.array < 0): + if logger: + logger.debug("Setting negative-variance pixels to 0 for noise generation.") + variance_template.array[variance_template.array < 0] = 0 + if np.any(~np.isfinite(variance_template.array)): + if logger: + logger.debug("Setting non-finite-variance pixels to 0 for noise generation.") + variance_template.array[~np.isfinite(variance_template.array)] = 0 + + # Randomly vary the injected flux in each pixel, + # consistent with the true variance level. + rng = galsim.BaseDeviate(noise_seed) + variable_noise = galsim.VariableGaussianNoise(rng, variance_template) + image_template.addNoise(variable_noise) + + # Set an "estimated" variance level in each pixel, + # corresponding to the randomly varied image. + variance_template = image_template.copy() + variance_template *= variance_scale + + # Add the randomly varied synthetic image to the original + # image. + common_image += image_template + # Add the estimated variance of the injection to the original + # variance. + common_variance += variance_template + common_box = Box2I( Point2I(object_common_bounds.xmin, object_common_bounds.ymin), Point2I(object_common_bounds.xmax, object_common_bounds.ymax), @@ -564,4 +616,8 @@ def inject_galsim_objects_into_exposure( if logger: logger.debug("No area overlap for object at %s; flagging and skipping.", sky_coords) + # Increment the seed so different noise is generated for different + # objects. + noise_seed += 1 + return draw_sizes, common_bounds, fft_size_errors, psf_compute_errors diff --git a/python/lsst/source/injection/inject_visit.py b/python/lsst/source/injection/inject_visit.py index ab97d7e..447449e 100644 --- a/python/lsst/source/injection/inject_visit.py +++ b/python/lsst/source/injection/inject_visit.py @@ -25,8 +25,11 @@ from typing import cast +import numpy as np from lsst.pex.config import Field from lsst.pipe.base.connectionTypes import Input, Output +from sklearn.linear_model import LinearRegression, RANSACRegressor +from sklearn.metrics import mean_squared_error from .inject_base import BaseInjectConfig, BaseInjectConnections, BaseInjectTask @@ -96,6 +99,12 @@ class VisitInjectConfig( # type: ignore [call-arg] dtype=bool, default=False, ) + brightness_cutoff = Field[float]( + doc="Ignore pixels with image flux below this level when fitting flux vs. variance.", + dtype=float, + default=500, + ) + variance_fit_seed = Field[int](doc="Seed for RANSAC fit of flux vs. variance.", default=0) class VisitInjectTask(BaseInjectTask): @@ -104,6 +113,13 @@ class VisitInjectTask(BaseInjectTask): _DefaultName = "visitInjectTask" ConfigClass = VisitInjectConfig + def run(self, injection_catalogs, input_exposure, psf, photo_calib, wcs): + self.log.info("Fitting flux vs. variance in each pixel.") + variance_scale = self.get_variance_scale(input_exposure) + self.log.info("Variance scale factor: %.3f", variance_scale) + + return super().run(injection_catalogs, input_exposure, psf, photo_calib, wcs, variance_scale) + def runQuantum(self, butler_quantum_context, input_refs, output_refs): inputs = butler_quantum_context.get(input_refs) detector_id = inputs["input_exposure"].getDetector().getId() @@ -128,3 +144,50 @@ def runQuantum(self, butler_quantum_context, input_refs, output_refs): input_keys = ["injection_catalogs", "input_exposure", "sky_map", "psf", "photo_calib", "wcs"] outputs = self.run(**{key: value for (key, value) in inputs.items() if key in input_keys}) butler_quantum_context.put(outputs, output_refs) + + """ + Establish the variance scale by a linear fit of variance vs. flux. + In practice we see that most pixels in a coadd obey a consistent, simple + linear relationship between variance and flux, but a small sample skews + somewhat away from linearity and results in a fit that seems less than + ideal. + + To identify and hence ignore these odd pixels, we perform a RANSAC fit. + RANSAC iteratively finds a least-squares straight line fit on random + subsamples of points, while identifying outliers. The final results tend to + be much more stable and outlier-robust than a simple linear fit. + + Pixels below a flux level specified by brightness_cutoff are more likely to + have wild outliers, so for simplicity's sake we simply ignore them. + """ + + def get_variance_scale(self, exposure): + flux = exposure.image.array.ravel() + var = exposure.variance.array.ravel() + + # Ignore pixels with nan or infinite values + good_pixels = np.isfinite(flux) & np.isfinite(var) + flux = flux[good_pixels] + var = var[good_pixels] + + # Only fit pixels with at least this much inst flux + bright_pixels = flux > self.config.brightness_cutoff + + # Simple linear regression to establish MSE + linear = LinearRegression() + linear.fit(flux[bright_pixels].reshape(-1, 1), var[bright_pixels]) + linear_mse = mean_squared_error(var, linear.predict(flux.reshape(-1, 1))) + + # RANSAC regression + ransac = RANSACRegressor( + loss="squared_error", + residual_threshold=0.1 * linear_mse, + random_state=self.config.variance_fit_seed, + ) + ransac.fit(flux[bright_pixels].reshape(-1, 1), var[bright_pixels]) + variance_scale = float(ransac.estimator_.coef_[0]) + + if variance_scale < 0: + self.log.warning("Slope of final variance vs. flux fit is negative.") + + return variance_scale