Coverage for python/lsst/meas/algorithms/gp_interpolation.py: 71%
115 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-01 08:58 +0000
1# This file is part of meas_algorithms.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (https://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
22import numpy as np
23from lsst.meas.algorithms import CloughTocher2DInterpolatorUtils as ctUtils
24from lsst.geom import Box2I, Point2I
25from lsst.afw.geom import SpanSet
26import copy
27import treecorr
28import treegp
30import logging
32# We need to explicitly turn off multiprocessing in treecorr which is used
33# by treegp.
34treecorr.set_max_omp_threads(1)
36__all__ = [
37 "InterpolateOverDefectGaussianProcess",
38 "GaussianProcessTreegp",
39]
42def updateMaskFromArray(mask, bad_pixel, interpBit):
43 """
44 Update the mask array with the given bad pixels.
46 Parameters
47 ----------
48 mask : `lsst.afw.image.MaskedImage`
49 The mask image to update.
50 bad_pixel : `np.array`
51 An array-like object containing the coordinates of the bad pixels.
52 Each row should contain the x and y coordinates of a bad pixel.
53 interpBit : `int`
54 The bit value to set for the bad pixels in the mask.
55 """
56 x0 = mask.getX0()
57 y0 = mask.getY0()
58 for row in bad_pixel:
59 x = int(row[0] - x0)
60 y = int(row[1] - y0)
61 mask.array[y, x] |= interpBit
62 # TO DO --> might be better: mask.array[int(bad_pixel[:,1]-y0), int(bad_pixel[:,0]-x)] |= interpBit
65def median_with_mad_clipping(data, mad_multiplier=2.0):
66 """
67 Calculate the median of the input data after applying Median Absolute Deviation (MAD) clipping.
69 The MAD clipping method is used to remove outliers from the data. The median of the data is calculated,
70 and then the MAD is calculated as the median absolute deviation from the median. The data is then clipped
71 by removing values that are outside the range of median +/- mad_multiplier * MAD. Finally, the median of
72 the clipped data is returned.
74 Parameters:
75 -----------
76 data : `np.array`
77 Input data array.
78 mad_multiplier : `float`, optional
79 Multiplier for the MAD value used for clipping. Default is 2.0.
81 Returns:
82 --------
83 median_clipped : `float`
84 Median value of the clipped data.
86 Examples:
87 ---------
88 >>> data = [1, 2, 3, 4, 5, 100]
89 >>> median_with_mad_clipping(data)
90 3.5
91 """
92 median = np.median(data)
93 mad = np.median(np.abs(data - median))
94 clipping_range = mad_multiplier * mad
95 clipped_data = np.clip(data, median - clipping_range, median + clipping_range)
96 median_clipped = np.median(clipped_data)
97 return median_clipped
100class GaussianProcessTreegp:
101 """
102 Gaussian Process Treegp class for Gaussian Process interpolation.
104 The basic GP regression, which uses Cholesky decomposition.
106 Parameters:
107 -----------
108 std : `float`, optional
109 Standard deviation of the Gaussian Process kernel. Default is 1.0.
110 correlation_length : `float`, optional
111 Correlation length of the Gaussian Process kernel. Default is 1.0.
112 white_noise : `float`, optional
113 White noise level of the Gaussian Process. Default is 0.0.
114 mean : `float`, optional
115 Mean value of the Gaussian Process. Default is 0.0.
116 """
118 def __init__(self, std=1.0, correlation_length=1.0, white_noise=0.0, mean=0.0):
119 self.std = std
120 self.correlation_length = correlation_length
121 self.white_noise = white_noise
122 self.mean = mean
124 # Looks like weird to do that, but this is justified.
125 # in GP if no noise is provided, even if matrix
126 # can be inverted, it wont invert because of numerical
127 # issue (det(K)~0). Add a little bit of noise allow
128 # to compute a numerical solution in the case of no
129 # external noise is added. Wont happened on real
130 # image but help for unit test.
131 if self.white_noise == 0.0: 131 ↛ 132line 131 didn't jump to line 132 because the condition on line 131 was never true
132 self.white_noise = 1e-5
134 def fit(self, x_train, y_train):
135 """
136 Fit the Gaussian Process to the given training data.
138 Parameters:
139 -----------
140 x_train : `np.array`
141 Input features for the training data.
142 y_train : `np.array`
143 Target values for the training data.
144 """
145 kernel = f"{self.std}**2 * RBF({self.correlation_length})"
146 self.gp = treegp.GPInterpolation(
147 kernel=kernel,
148 optimizer="none",
149 normalize=False,
150 white_noise=self.white_noise,
151 )
152 self.gp.initialize(x_train, y_train - self.mean)
153 self.gp.solve()
155 def predict(self, x_predict):
156 """
157 Predict the target values for the given input features.
159 Parameters:
160 -----------
161 x_predict : `np.array`
162 Input features for the prediction.
164 Returns:
165 --------
166 y_pred : `np.array`
167 Predicted target values.
168 """
169 y_pred = self.gp.predict(x_predict)
170 return y_pred + self.mean
173class InterpolateOverDefectGaussianProcess:
174 """
175 InterpolateOverDefectGaussianProcess class performs Gaussian Process
176 (GP) interpolation over defects in an image.
178 Parameters:
179 -----------
180 masked_image : `lsst.afw.image.MaskedImage`
181 The masked image containing the defects to be interpolated.
182 defects : `list`[`str`], optional
183 The types of defects to be interpolated. Default is ["SAT"].
184 fwhm : `float`, optional
185 The full width at half maximum (FWHM) of the PSF. Default is 5.
186 bin_spacing : `int`, optional
187 The spacing between bins for good pixel binning. Default is 10.
188 threshold_dynamic_binning : `int`, optional
189 The threshold for dynamic binning. Default is 1000.
190 threshold_subdivide : `int`, optional
191 The threshold for sub-dividing the bad pixel array to avoid memory error. Default is 20000.
192 correlation_length_cut : `int`, optional
193 The factor by which to dilate the bounding box around defects. Default is 5.
194 log : `lsst.log.Log`, `logging.Logger` or `None`, optional
195 Logger object used to write out messages. If `None` a default
196 logger will be used.
197 """
199 def __init__(
200 self,
201 masked_image,
202 defects=["SAT"],
203 fwhm=5,
204 bin_image=True,
205 bin_spacing=10,
206 threshold_dynamic_binning=1000,
207 threshold_subdivide=20000,
208 correlation_length_cut=5,
209 log=None,
210 ):
212 self.log = log or logging.getLogger(__name__)
214 self.bin_image = bin_image
215 self.bin_spacing = bin_spacing
216 self.threshold_subdivide = threshold_subdivide
217 self.threshold_dynamic_binning = threshold_dynamic_binning
219 self.masked_image = masked_image
220 self.defects = defects
221 self.correlation_length = fwhm
222 self.correlation_length_cut = correlation_length_cut
224 self.interpBit = self.masked_image.mask.getPlaneBitMask("INTRP")
226 def run(self):
227 """
228 Interpolate over the defects in the image.
230 Change self.masked_image .
231 """
232 if self.defects == [] or self.defects is None: 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true
233 self.log.info("No defects found. No interpolation performed.")
234 else:
235 mask = self.masked_image.getMask()
236 bad_pixel_mask = mask.getPlaneBitMask(self.defects)
237 bad_mask_span_set = SpanSet.fromMask(mask, bad_pixel_mask).split()
239 bbox = self.masked_image.getBBox()
240 global_xmin, global_xmax = bbox.minX, bbox.maxX
241 global_ymin, global_ymax = bbox.minY, bbox.maxY
243 for spanset in bad_mask_span_set:
244 bbox = spanset.getBBox()
245 # Dilate the bbox to make sure we have enough good pixels around the defect
246 # For now, we dilate by 5 times the correlation length
247 # For GP with the isotropic kernel, points at the default value of
248 # correlation_length_cut=5 have negligible effect on the prediction.
249 bbox = bbox.dilatedBy(
250 int(self.correlation_length * self.correlation_length_cut)
251 ) # need integer as input.
252 xmin, xmax = max([global_xmin, bbox.minX]), min(global_xmax, bbox.maxX)
253 ymin, ymax = max([global_ymin, bbox.minY]), min(global_ymax, bbox.maxY)
254 localBox = Box2I(Point2I(xmin, ymin), Point2I(xmax - xmin, ymax - ymin))
255 masked_sub_image = self.masked_image[localBox]
257 masked_sub_image = self.interpolate_masked_sub_image(masked_sub_image)
258 self.masked_image[localBox] = masked_sub_image
260 def _good_pixel_binning(self, pixels):
261 """
262 Performs pixel binning using treegp.meanify
264 Parameters:
265 -----------
266 pixels : `np.array`
267 The array of pixels.
269 Returns:
270 --------
271 `np.array`
272 The binned array of pixels.
273 """
275 n_pixels = len(pixels[:, 0])
276 dynamic_binning = int(np.sqrt(n_pixels / self.threshold_dynamic_binning))
277 if n_pixels / self.bin_spacing**2 < n_pixels / dynamic_binning**2:
278 bin_spacing = self.bin_spacing
279 else:
280 bin_spacing = dynamic_binning
281 binning = treegp.meanify(bin_spacing=bin_spacing, statistics="mean")
282 binning.add_field(
283 pixels[:, :2],
284 pixels[:, 2:].T,
285 )
286 binning.meanify()
287 return np.array(
288 [binning.coords0[:, 0], binning.coords0[:, 1], binning.params0]
289 ).T
291 def interpolate_masked_sub_image(self, masked_sub_image):
292 """
293 Interpolate the masked sub-image.
295 Parameters:
296 -----------
297 masked_sub_image : `lsst.afw.image.MaskedImage`
298 The sub-masked image to be interpolated.
300 Returns:
301 --------
302 `lsst.afw.image.MaskedImage`
303 The interpolated sub-masked image.
304 """
306 cut = int(
307 self.correlation_length * self.correlation_length_cut
308 ) # need integer as input.
309 bad_pixel, good_pixel = ctUtils.findGoodPixelsAroundBadPixels(
310 masked_sub_image, self.defects, buffer=cut
311 )
312 # Do nothing if bad pixel is None.
313 if bad_pixel.size == 0 or good_pixel.size == 0: 313 ↛ 314line 313 didn't jump to line 314 because the condition on line 313 was never true
314 self.log.info("No bad or good pixels found. No interpolation performed.")
315 return masked_sub_image
316 # Do GP interpolation if bad pixel found.
317 else:
318 # gp interpolation
319 sub_image_array = masked_sub_image.getVariance().array
320 white_noise = np.sqrt(
321 np.mean(sub_image_array[np.isfinite(sub_image_array)])
322 )
323 kernel_amplitude = np.max(good_pixel[:, 2:])
324 if not np.isfinite(kernel_amplitude): 324 ↛ 325line 324 didn't jump to line 325 because the condition on line 324 was never true
325 filter_finite = np.isfinite(good_pixel[:, 2:]).T[0]
326 good_pixel = good_pixel[filter_finite]
327 if good_pixel.size == 0:
328 self.log.info(
329 "No bad or good pixels found. No interpolation performed."
330 )
331 return masked_sub_image
332 # kernel amplitude might be better described by maximum value of good pixel given
333 # the data and not really a random gaussian field.
334 kernel_amplitude = np.max(good_pixel[:, 2:])
336 if self.bin_image: 336 ↛ 337line 336 didn't jump to line 337 because the condition on line 336 was never true
337 try:
338 good_pixel = self._good_pixel_binning(copy.deepcopy(good_pixel))
339 except Exception:
340 self.log.info(
341 "Binning failed, use original good pixel array in interpolation."
342 )
344 # put this after binning as computing median is O(n*log(n))
345 clipped_median = median_with_mad_clipping(good_pixel[:, 2:])
347 gp = GaussianProcessTreegp(
348 std=np.sqrt(kernel_amplitude),
349 correlation_length=self.correlation_length,
350 white_noise=white_noise,
351 mean=clipped_median,
352 )
353 gp.fit(good_pixel[:, :2], np.squeeze(good_pixel[:, 2:]))
354 if bad_pixel.size < self.threshold_subdivide: 354 ↛ 358line 354 didn't jump to line 358 because the condition on line 354 was always true
355 gp_predict = gp.predict(bad_pixel[:, :2])
356 bad_pixel[:, 2:] = gp_predict.reshape(np.shape(bad_pixel[:, 2:]))
357 else:
358 self.log.info("sub-divide bad pixel array to avoid memory error.")
359 for i in range(0, len(bad_pixel), self.threshold_subdivide):
360 end = min(i + self.threshold_subdivide, len(bad_pixel))
361 gp_predict = gp.predict(bad_pixel[i:end, :2])
362 bad_pixel[i:end, 2:] = gp_predict.reshape(
363 np.shape(bad_pixel[i:end, 2:])
364 )
366 # Update values
367 ctUtils.updateImageFromArray(masked_sub_image.image, bad_pixel)
368 updateMaskFromArray(masked_sub_image.mask, bad_pixel, self.interpBit)
369 return masked_sub_image