Coverage for tests/test_inject_engine.py: 97%
151 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-11 09:07 +0000
1# This file is part of source_injection.
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 logging
23import unittest
24from types import GeneratorType
26import astropy.units as u
27import galsim
28import numpy as np
29import pytest
30from galsim import BoundsI, GSObject
32import lsst.utils.tests
33from lsst.geom import Point2D, SpherePoint, degrees
34from lsst.images import YX
35from lsst.images.cells import CellIJ
36from lsst.source.injection.inject_engine import (
37 generate_galsim_objects,
38 get_gain_map,
39 infer_gain_from_image,
40 inject_galsim_objects_into_exposure,
41 make_galsim_object,
42)
43from lsst.source.injection.utils.test_utils import (
44 make_test_cell_coadd,
45 make_test_exposure,
46 make_test_injection_catalog,
47)
48from lsst.utils.tests import TestCase
51class InjectEngineTestCase(TestCase):
52 """Test the inject_engine.py module."""
54 def setUp(self):
55 """Set up synthetic source injection inputs.
57 This method sets up a noisy synthetic image with Gaussian PSFs injected
58 into the frame, an example source injection catalog, and a generator of
59 GalSim objects intended for injection.
60 """
61 self.exposure = make_test_exposure()
62 # Mark a cell as missing to test that it's ignored, although the data
63 # in it may still actually be fine
64 self.cell_bad = CellIJ(3, 3)
65 self.cell_bad_inverse = {
66 CellIJ(i, j)
67 for i in range(90)
68 for j in range(90)
69 if not (i == self.cell_bad.i and j == self.cell_bad.j)
70 }
71 self.cell_coadd, self.cell_coadd_inverse = (
72 make_test_cell_coadd(
73 exposure=self.exposure.clone(),
74 cell_shape=YX(x=35, y=35),
75 psf_shape=(33, 33),
76 missing=missing,
77 band="r",
78 )
79 for missing in ({self.cell_bad}, self.cell_bad_inverse)
80 )
81 bbox_exp = self.exposure.getBBox()
82 bbox_cell_bad = self.cell_coadd.grid.bbox_of(self.cell_bad)
83 self.slice_x_cell_bad = slice(
84 bbox_cell_bad.start.x - bbox_exp.beginX,
85 bbox_cell_bad.stop.x - bbox_exp.beginX,
86 )
87 self.slice_y_cell_bad = slice(
88 bbox_cell_bad.start.y - bbox_exp.beginY,
89 bbox_cell_bad.stop.y - bbox_exp.beginY,
90 )
91 self.bbox_cell_bad = bbox_cell_bad
92 self.bbox_cell_bad_legacy = bbox_cell_bad.to_legacy()
93 cen_cell_bad = self.exposure.wcs.pixelToSky(self.bbox_cell_bad_legacy.getCenter())
94 injection_catalog = make_test_injection_catalog(
95 self.exposure.getWcs(),
96 bbox_exp,
97 )
98 row_last = injection_catalog[-1]
99 injection_catalog.add_row(
100 {
101 "ra": cen_cell_bad.getRa().asDegrees(),
102 "dec": cen_cell_bad.getDec().asDegrees(),
103 "mag": row_last["mag"],
104 "source_type": row_last["source_type"],
105 }
106 )
107 # Add one galaxy centered in the "bad" cell
108 self.cell_galaxy = self.cell_bad
109 cen_cell_galaxy = self.exposure.wcs.pixelToSky(
110 self.cell_coadd.grid.bbox_of(self.cell_galaxy).to_legacy().getCenter()
111 )
112 n_rows = len(injection_catalog)
113 columns_sersic = ("n", "half_light_radius", "q", "beta")
114 injection_catalog.add_columns(
115 cols=tuple(
116 np.ma.masked_array(data=np.zeros(n_rows, dtype=float), mask=np.ones(n_rows, dtype=bool))
117 for _ in range(len(columns_sersic))
118 ),
119 names=columns_sersic,
120 )
121 injection_catalog.add_row(
122 {
123 "ra": cen_cell_galaxy.getRa().asDegrees(),
124 "dec": cen_cell_galaxy.getDec().asDegrees(),
125 "mag": injection_catalog["mag"][0],
126 "source_type": "Sersic",
127 "n": 2.0,
128 "half_light_radius": 10 * 0.036,
129 "q": 0.8,
130 "beta": 15.0,
131 }
132 )
133 injection_catalog["half_light_radius"].unit = u.arcsec
134 self.injection_catalog = injection_catalog
136 self.galsim_objects = generate_galsim_objects(
137 injection_catalog=self.injection_catalog,
138 photo_calib=self.exposure.photoCalib,
139 wcs=self.exposure.wcs,
140 fits_alignment="wcs",
141 stamp_prefix="",
142 )
143 self.photoCalib = self.exposure.getPhotoCalib()
144 self.inst_fluxes = [
145 float(self.photoCalib.magnitudeToInstFlux(mag)) for mag in self.injection_catalog["mag"]
146 ]
148 def tearDown(self):
149 del self.exposure
150 del self.cell_coadd
151 del self.injection_catalog
152 del self.galsim_objects
153 del self.photoCalib
154 del self.inst_fluxes
156 def test_make_galsim_object(self):
157 source_data = self.injection_catalog[0]
158 sky_coords = SpherePoint(float(source_data["ra"]), float(source_data["dec"]), degrees)
159 pixel_coords = self.exposure.wcs.skyToPixel(sky_coords)
160 inst_flux = self.exposure.photoCalib.magnitudeToInstFlux(source_data["mag"], pixel_coords)
161 object = make_galsim_object(
162 source_data=source_data,
163 source_type=source_data["source_type"],
164 inst_flux=inst_flux,
165 )
166 self.assertIsInstance(object, GSObject)
167 self.assertIsInstance(object, getattr(galsim, source_data["source_type"]))
169 def test_generate_galsim_objects(self):
170 self.assertTrue(isinstance(self.galsim_objects, GeneratorType))
171 for galsim_object in self.galsim_objects:
172 self.assertIsInstance(galsim_object, tuple)
173 self.assertEqual(len(galsim_object), 4)
174 self.assertIsInstance(galsim_object[0], SpherePoint) # RA/Dec
175 self.assertIsInstance(galsim_object[1], Point2D) # x/y
176 self.assertIsInstance(galsim_object[2], int) # draw size
177 self.assertIsInstance(galsim_object[3], GSObject) # GSObject
179 def test_infer_gain_nonexistent_mask(self):
180 """Test that infer_gain_from_image returns a gain value and logs a
181 warning when provided with nonexistent mask plane names.
182 """
183 logger = logging.getLogger(__name__)
184 with self.assertLogs(logger, level="WARNING") as cm:
185 gain = infer_gain_from_image(
186 self.exposure,
187 bad_mask_names=["NONEXISTENT_MASK1", "NONEXISTENT_MASK2"],
188 logger=logger,
189 )
190 self.assertTrue(any("NONEXISTENT_MASK1" in msg for msg in cm.output))
191 self.assertIsInstance(gain, float)
192 self.assertTrue(np.isfinite(gain))
193 gain_no_mask = infer_gain_from_image(self.exposure, bad_mask_names=[])
194 self.assertAlmostEqual(gain, gain_no_mask)
196 def test_infer_gain_no_valid_pixels(self):
197 """Test that infer_gain_from_image returns NaN (rather than raising)
198 when a region has no valid pixels to fit.
199 """
200 logger = logging.getLogger(__name__)
201 # Every pixel flagged with a bad mask plane.
202 all_bad = make_test_exposure()
203 all_bad.mask.addMaskPlane("BAD")
204 all_bad.mask.array[:] = all_bad.mask.getPlaneBitMask("BAD")
205 with self.assertLogs(logger, level="WARNING"):
206 gain = infer_gain_from_image(all_bad, bad_mask_names=["BAD"], logger=logger)
207 self.assertTrue(np.isnan(gain))
209 # Every variance pixel non-finite.
210 all_nan = make_test_exposure()
211 all_nan.variance.array[:] = np.nan
212 self.assertTrue(np.isnan(infer_gain_from_image(all_nan, bad_mask_names=[])))
214 def test_get_gain_map_no_valid_pixels(self):
215 """Test that get_gain_map produces a finite, positive map even when no
216 region can be fit, falling back to unit gain.
217 """
218 all_bad = make_test_exposure()
219 all_bad.mask.addMaskPlane("BAD")
220 all_bad.mask.array[:] = all_bad.mask.getPlaneBitMask("BAD")
221 gain_map = get_gain_map(all_bad, bad_mask_names=["BAD"])
222 self.assertTrue(np.all(np.isfinite(gain_map.array)))
223 self.assertTrue(np.all(gain_map.array > 0))
225 def test_inject_galsim_objects_into_cell_coadd(self):
226 self._test_inject_galsim_objects_into_exposure(
227 self.cell_coadd,
228 exposure_inverse=self.cell_coadd_inverse,
229 is_cell=True,
230 )
232 def test_inject_galsim_objects_into_exposure(self):
233 self._test_inject_galsim_objects_into_exposure(self.exposure, is_cell=False)
235 def _test_inject_galsim_objects_into_exposure(
236 self,
237 exposure,
238 exposure_inverse=None,
239 is_cell: bool = True,
240 ):
241 array_orig = exposure.image.array.copy()
242 flux0 = np.sum(array_orig)
243 cell_bad_orig = array_orig[self.slice_y_cell_bad, self.slice_x_cell_bad]
244 for injection_core_size in (None, 0, 1.5):
245 with pytest.raises(ValueError):
246 inject_galsim_objects_into_exposure(
247 exposure=exposure,
248 objects=(),
249 injection_core_size=injection_core_size,
250 )
251 galsim_objects = self.galsim_objects
252 if is_cell:
253 # At least the last object will be needed later
254 galsim_objects = list(galsim_objects)
256 injected_outputs, injected_outputs_inverse = (
257 (
258 inject_galsim_objects_into_exposure(
259 exposure=to_inject,
260 objects=galsim_objects,
261 mask_plane_name="INJECTED",
262 calib_flux_radius=12.0,
263 draw_size_max=1000,
264 add_noise=False,
265 injection_core_size=5,
266 )
267 if to_inject is not None
268 else None
269 )
270 for to_inject in (exposure, exposure_inverse)
271 )
273 cell_bad = exposure.image.array[self.slice_y_cell_bad, self.slice_x_cell_bad]
274 mask_cell_bad = exposure.mask[self.bbox_cell_bad if is_cell else self.bbox_cell_bad_legacy]
276 check_total_flux = True
277 if is_cell:
278 # Test that nothing was injected in the bad cell
279 np.testing.assert_array_equal(cell_bad_orig, cell_bad)
280 np.testing.assert_array_equal(mask_cell_bad.get("INJECTED"), False)
281 np.testing.assert_array_equal(mask_cell_bad.get("INJECTED_CORE"), False)
282 if exposure_inverse is not None: 282 ↛ 293line 282 didn't jump to line 293 because the condition on line 282 was always true
283 cell_bad_inverse = exposure_inverse.image.array[self.slice_y_cell_bad, self.slice_x_cell_bad]
284 mask_cell_bad_inverse = exposure_inverse.mask[self.bbox_cell_bad]
285 np.testing.assert_array_less(cell_bad_orig, cell_bad_inverse)
286 np.testing.assert_array_equal(mask_cell_bad_inverse.get("INJECTED"), True)
287 # Default core size is 3x3 and there should be at least one
288 # injected core from the galaxy. Others may exist depending on
289 # the random seed used
290 self.assertGreaterEqual(np.sum(mask_cell_bad_inverse.get("INJECTED_CORE")), 9)
291 exposure.image[self.bbox_cell_bad].array = exposure_inverse.image[self.bbox_cell_bad].array
292 else:
293 check_total_flux = False
294 else:
295 # Test that the entire good cell is filled
296 # Note this requires that the injected galaxy be large enough, so
297 # if setUp is modified to inject a smaller object it will fail
298 np.testing.assert_array_less(cell_bad_orig, cell_bad)
299 injected_bitmask = mask_cell_bad.getPlaneBitMask("INJECTED")
300 np.testing.assert_array_equal(mask_cell_bad.array & injected_bitmask, injected_bitmask)
302 if check_total_flux: 302 ↛ 309line 302 didn't jump to line 309 because the condition on line 302 was always true
303 self.assertAlmostEqual(
304 np.sum(exposure.image.array) - flux0,
305 np.sum(self.inst_fluxes),
306 delta=0.00015 * np.sum(self.inst_fluxes),
307 )
309 for injected_exposure in (exposure, exposure_inverse):
310 if injected_exposure is None:
311 continue
312 mask_dict = exposure.mask.schema if is_cell else exposure.mask.getMaskPlaneDict()
313 assert "INJECTED" in mask_dict
314 assert "INJECTED_CORE" in mask_dict
316 for injected_outputs in (injected_outputs, injected_outputs_inverse):
317 if injected_outputs is None:
318 continue
319 draw_sizes, common_bounds, fft_size_errors, psf_compute_errors = injected_outputs
320 self.assertEqual(len(draw_sizes), len(self.injection_catalog["ra"]))
321 self.assertTrue(all(isinstance(injected_output, list) for injected_output in injected_outputs))
322 self.assertTrue(all(isinstance(item, int) and (item > 0) for item in draw_sizes))
323 self.assertTrue(all(isinstance(item, BoundsI) for item in common_bounds)) # common bounds
324 self.assertTrue(all(item is False for item in fft_size_errors)) # FFT size errors
325 self.assertTrue(all(isinstance(item, bool) for item in psf_compute_errors)) # PSF compute errors
328class MemoryTestCase(lsst.utils.tests.MemoryTestCase):
329 """Test memory usage of functions in this script."""
331 pass
334def setup_module(module):
335 """Configure pytest."""
336 lsst.utils.tests.init()
339if __name__ == "__main__": 339 ↛ 340line 339 didn't jump to line 340 because the condition on line 339 was never true
340 lsst.utils.tests.init()
341 unittest.main()