Coverage for tests/test_extraction.py: 99%
265 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-13 09:57 +0000
1#!/usr/bin/env python
2# This file is part of atmospec.
3#
4# Developed for the LSST Data Management System.
5# This product includes software developed by the LSST Project
6# (https://www.lsst.org).
7# See the COPYRIGHT file at the top-level directory of this distribution
8# for details of code ownership.
9#
10# This program is free software: you can redistribute it and/or modify
11# it under the terms of the GNU General Public License as published by
12# the Free Software Foundation, either version 3 of the License, or
13# (at your option) any later version.
14#
15# This program is distributed in the hope that it will be useful,
16# but WITHOUT ANY WARRANTY; without even the implied warranty of
17# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
18# GNU General Public License for more details.
19#
20# You should have received a copy of the GNU General Public License
21# along with this program. If not, see <https://www.gnu.org/licenses/>.
22"""Test cases for lsst.atmospec.extraction."""
24import types
25import unittest
26import unittest.mock as mock
27import numpy as np
29import lsst.utils
30import lsst.utils.tests
31import lsst.geom as geom
32import lsst.afw.image as afwImage
33from lsst.atmospec.dispersion import DispersionRelation
34from lsst.atmospec.extraction import (
35 SpectralExtractionTask,
36 SpectralExtractionTaskConfig,
37 gausMoffatModel,
38 moffatModel,
39)
42def makeGaussTrace(width, height, amp=50.0, center=None, sigma=3.0, background=5.0):
43 """Build a synthetic 2D array with a Gaussian spatial profile per row.
45 Each row (constant y) contains a Gaussian in x, emulating a spectral
46 trace whose spatial profile runs across the columns.
47 """
48 if center is None:
49 center = width / 2.0
50 xs = np.arange(width)
51 profile = amp * np.exp(-(xs - center) ** 2 / (2.0 * sigma ** 2)) + background
52 array = np.tile(profile, (height, 1)).astype(np.float32)
53 return array
56def makeExposure(width, height, **kwargs):
57 """Make an ExposureF containing a synthetic Gaussian trace."""
58 bbox = geom.Box2I(geom.Point2I(0, 0), geom.Extent2I(width, height))
59 exp = afwImage.ExposureF(bbox)
60 exp.image.array[:] = makeGaussTrace(width, height, **kwargs)
61 exp.variance.array[:] = 1.0
62 return exp, bbox
65class ModelFunctionTestCase(lsst.utils.tests.TestCase):
66 """Tests for the module-level pure model/fit helpers."""
68 def test_gauss1D(self):
69 x = np.linspace(-10, 10, 101)
70 amp, mu, sigma = 5.0, 0.0, 2.0
71 y = SpectralExtractionTask.gauss1D(x, amp, mu, sigma)
72 # Peak at mu, equal to amp.
73 self.assertAlmostEqual(np.max(y), amp, places=6)
74 self.assertAlmostEqual(x[np.argmax(y)], mu, places=6)
75 # Symmetric about the mean.
76 self.assertFloatsAlmostEqual(y, y[::-1], atol=1e-6)
78 def test_moffatModel(self):
79 model = moffatModel((10.0, 5.0, 2.0))
80 self.assertAlmostEqual(model.amplitude.value, 10.0)
81 self.assertAlmostEqual(model.x_0.value, 5.0)
82 self.assertAlmostEqual(model.gamma.value, 2.0)
83 # Peak of a Moffat is at x_0 and equals the amplitude.
84 self.assertAlmostEqual(model(5.0), 10.0, places=6)
86 def test_gausMoffatModel(self):
87 pars = (10.0, 15.0, 1.5, 1.5, 500.0, 15.0, 3.0)
88 model = gausMoffatModel(pars)
89 # Compound model (Gaussian + Moffat) evaluates to a finite array.
90 xs = np.arange(30)
91 vals = model(xs)
92 self.assertEqual(vals.shape, xs.shape)
93 self.assertTrue(np.all(np.isfinite(vals)))
94 # Sub-model parameters are accessible on the compound model.
95 self.assertAlmostEqual(model.amplitude_0.value, 500.0)
96 self.assertAlmostEqual(model.amplitude_1.value, 10.0)
98 def test_moffatFit(self):
99 pixels = np.arange(40)
100 truth = moffatModel((1000.0, 20.0, 3.0))
101 footprint = np.asarray(truth(pixels), dtype=float)
102 integral, x0, gamma, alpha = SpectralExtractionTask.moffatFit(
103 pixels, footprint, amp=900.0, mu=19.0, sigma=2.5)
104 self.assertTrue(np.isfinite(integral))
105 self.assertGreater(integral, 0.0)
106 # The recovered centroid should be close to the injected one.
107 self.assertAlmostEqual(x0, 20.0, delta=1.0)
108 self.assertTrue(np.isfinite(gamma))
109 self.assertTrue(np.isfinite(alpha))
111 def test_gaussMoffatFit(self):
112 pixels = np.arange(40)
113 truthPars = (200.0, 20.0, 1.5, 1.5, 800.0, 20.0, 3.0)
114 truth = gausMoffatModel(truthPars)
115 footprint = np.asarray(truth(pixels), dtype=float)
116 initialPars = (150.0, 20.0, 1.5, 1.5, 700.0, 20.0, 3.0)
117 result = SpectralExtractionTask.gaussMoffatFit(pixels, footprint, initialPars)
118 self.assertEqual(len(result), 9)
119 intGM, intG, fitGausAmp, fitX0, fitGam, fitAlpha, fitMofAmp, fitMofMean, fitMofWid = result
120 for val in result:
121 self.assertTrue(np.isfinite(val))
123 def test_subtractBkgd(self):
124 width = 30
125 bkgd_size = 5
126 slc = np.zeros(width)
127 slc[:] = 7.0 # constant background of 7
128 slc[12:18] += 100.0 # a bump in the middle (inside the kept region)
129 result = SpectralExtractionTask.subtractBkgd(slc, width, bkgd_size)
130 # Kept region has length width - 2*bkgd_size.
131 self.assertEqual(len(result), width - 2 * bkgd_size)
132 # Constant background should be subtracted to ~zero at the edges of
133 # the kept region.
134 self.assertAlmostEqual(result[0], 0.0, places=6)
137class ConfigTestCase(lsst.utils.tests.TestCase):
138 """Tests for the task configuration."""
140 def test_configDefaults(self):
141 config = SpectralExtractionTaskConfig()
142 self.assertFalse(config.perRowBackground)
143 self.assertEqual(config.perRowBackgroundSize, 10)
144 self.assertFalse(config.writeResiduals)
145 self.assertTrue(config.doSmoothBackround)
146 self.assertFalse(config.doSigmaClipBackground)
147 config.validate()
149 def test_configConstruction(self):
150 task = SpectralExtractionTask()
151 self.assertIsNotNone(task.config)
152 task2 = SpectralExtractionTask(config=SpectralExtractionTaskConfig())
153 self.assertIsNotNone(task2.config)
156class BackgroundTestCase(lsst.utils.tests.TestCase):
157 """Additional edge cases for _calculateBackground."""
159 def test_noSmooth(self):
160 task = SpectralExtractionTask()
161 mi = afwImage.MaskedImageF(5, 20)
162 mi.image.array[:] = 3.0
163 bgImg = task._calculateBackground(mi, 5, smooth=False)
164 self.assertEqual(np.shape(mi.image.array), np.shape(bgImg.array))
165 self.assertFloatsAlmostEqual(np.max(bgImg.array), 3.0, atol=1e-5)
167 def test_sigmaClip(self):
168 config = SpectralExtractionTaskConfig()
169 config.doSigmaClipBackground = True
170 task = SpectralExtractionTask(config=config)
171 mi = afwImage.MaskedImageF(5, 20)
172 mi.image.array[:] = 2.0
173 bgImg = task._calculateBackground(mi, 5)
174 self.assertEqual(np.shape(mi.image.array), np.shape(bgImg.array))
176 def test_reduceNbinsWarning(self):
177 task = SpectralExtractionTask()
178 mi = afwImage.MaskedImageF(5, 6)
179 mi.image.array[:] = 1.0
180 # nbins larger than height-1 triggers the reduction branch.
181 bgImg = task._calculateBackground(mi, 50, smooth=False)
182 self.assertEqual(np.shape(mi.image.array), np.shape(bgImg.array))
185class InitialiseAndFluxTestCase(lsst.utils.tests.TestCase):
186 """Tests for initialise() and getFluxBasic()."""
188 def setUp(self):
189 self.width = 30
190 self.height = 25
191 self.dispersion = DispersionRelation([100, 200], [400, 800])
193 def _makeTask(self, config=None):
194 return SpectralExtractionTask(config=config)
196 def test_initialise(self):
197 exp, bbox = makeExposure(self.width, self.height)
198 task = self._makeTask()
199 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
200 task.initialise(exp, centroid, bbox, self.dispersion)
202 self.assertEqual(task.spectrumWidth, self.width)
203 self.assertEqual(task.spectrumHeight, self.height)
204 self.assertEqual(len(task.apertureFlux), self.height)
205 self.assertEqual(len(task.psfFitPars), self.height)
206 self.assertEqual(len(task.moffatFitPars), self.height)
207 self.assertEqual(len(task.gausMoffatFitPars), self.height)
208 # Background-subtracted image should exist and match the footprint.
209 self.assertEqual(task.bgSubMi.getDimensions(), task.footprintMi.getDimensions())
210 self.assertIs(task.dispersionRelation, self.dispersion)
212 def test_initialise_noSmoothBackground(self):
213 config = SpectralExtractionTaskConfig()
214 config.doSmoothBackround = False
215 exp, bbox = makeExposure(self.width, self.height)
216 task = self._makeTask(config=config)
217 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
218 task.initialise(exp, centroid, bbox, self.dispersion)
219 self.assertEqual(task.spectrumHeight, self.height)
221 def test_getFluxBasic(self):
222 exp, bbox = makeExposure(self.width, self.height)
223 task = self._makeTask()
224 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
225 task.initialise(exp, centroid, bbox, self.dispersion)
226 result = task.getFluxBasic()
227 self.assertIs(result, task)
229 # Aperture flux and row-wise max must be populated and finite.
230 self.assertEqual(len(task.apertureFlux), self.height)
231 self.assertTrue(np.all(np.isfinite(task.apertureFlux)))
232 self.assertTrue(np.all(np.isfinite(task.rowWiseMax)))
234 # Basic Gaussian fit should have succeeded on at least some rows,
235 # recovering a centroid near the trace center.
236 fittedRows = [p for p in task.psfFitPars if p is not None]
237 self.assertGreater(len(fittedRows), 0)
238 for (psfAmp, psfMu, psfSigma, psfFlux) in fittedRows:
239 self.assertGreater(psfAmp, 0)
240 self.assertGreater(psfFlux, 0)
242 def test_getFluxBasic_perRowBackground(self):
243 config = SpectralExtractionTaskConfig()
244 config.perRowBackground = True
245 config.perRowBackgroundSize = 10
246 exp, bbox = makeExposure(self.width, self.height)
247 task = self._makeTask(config=config)
248 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
249 task.initialise(exp, centroid, bbox, self.dispersion)
250 result = task.getFluxBasic()
251 self.assertIs(result, task)
252 self.assertEqual(len(task.apertureFlux), self.height)
254 def test_getFluxBasic_offCenterWarning(self):
255 # Trace centered near an edge so |mu - width/2| >= 10, triggering
256 # the "initial mu more than 10 pixels from footprint center" warning.
257 exp, bbox = makeExposure(self.width, self.height, center=3.0, sigma=2.0)
258 task = self._makeTask()
259 centroid = geom.Point2D(3.0, self.height / 2.0)
260 task.initialise(exp, centroid, bbox, self.dispersion)
261 result = task.getFluxBasic()
262 self.assertIs(result, task)
264 def test_getFluxBasic_writeResiduals(self):
265 config = SpectralExtractionTaskConfig()
266 config.writeResiduals = True
267 exp, bbox = makeExposure(self.width, self.height)
268 task = self._makeTask(config=config)
269 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
270 task.initialise(exp, centroid, bbox, self.dispersion)
271 result = task.getFluxBasic()
272 self.assertIs(result, task)
274 def test_getFluxBasic_preventRunaway(self):
275 # Enable the PREVENT_RUNAWAY module flag and use an off-centre trace
276 # with a large per-row value spread so that both the psfMu and the
277 # psfSigma "runaway" reset branches fire.
278 exp, bbox = makeExposure(self.width, self.height, amp=100.0,
279 center=3.0, sigma=3.0)
280 task = self._makeTask()
281 centroid = geom.Point2D(3.0, self.height / 2.0)
282 task.initialise(exp, centroid, bbox, self.dispersion)
283 with mock.patch("lsst.atmospec.extraction.PREVENT_RUNAWAY", True):
284 result = task.getFluxBasic()
285 self.assertIs(result, task)
287 def test_getFluxBasic_failingRows(self):
288 # Row 0 has a clean Gaussian (so the bootstrap succeeds), but some
289 # later rows are driven below background so their max is negative.
290 # This makes the curve_fit bounds invalid (upper <= lower), raising a
291 # ValueError caught by the basic-Gauss handler, leaving psfFitPars
292 # None for those rows and exercising the "is None" gaussMoffat branch.
293 exp, bbox = makeExposure(self.width, self.height, amp=50.0)
294 # Drive a few rows well below the background level.
295 exp.image.array[10, :] = -1000.0
296 exp.image.array[15, :] = -1000.0
297 task = self._makeTask()
298 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
299 task.initialise(exp, centroid, bbox, self.dispersion)
300 result = task.getFluxBasic()
301 self.assertIs(result, task)
302 # The bad rows should have failed the basic Gaussian fit.
303 self.assertIsNone(task.psfFitPars[10])
304 # Row 0 should still have succeeded.
305 self.assertIsNotNone(task.psfFitPars[0])
307 def test_getFluxBasic_moffatFitRaises(self):
308 # Force the Moffat fitter to raise on every row so the
309 # RuntimeError/ValueError handler around it is exercised. (The
310 # GaussMoffat handler cannot be driven on every row because later
311 # rows reference the previous iteration's fitValsGM, so we leave the
312 # GaussMoffat fit working here.)
313 exp, bbox = makeExposure(self.width, self.height, amp=50.0)
314 task = self._makeTask()
315 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
316 task.initialise(exp, centroid, bbox, self.dispersion)
317 with mock.patch.object(task, "moffatFit",
318 side_effect=RuntimeError("boom")):
319 result = task.getFluxBasic()
320 self.assertIs(result, task)
321 # All Moffat fits failed, so that parameter list is unpopulated.
322 self.assertTrue(all(p is None for p in task.moffatFitPars))
324 def test_getFluxBasic_gaussMoffatRaisesLastRow(self):
325 # Make the GaussMoffat fit raise only on the final row so the
326 # RuntimeError/ValueError handler around it is exercised without
327 # leaving a later row referencing an unbound fitValsGM.
328 exp, bbox = makeExposure(self.width, self.height, amp=50.0)
329 task = self._makeTask()
330 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
331 task.initialise(exp, centroid, bbox, self.dispersion)
333 realFit = task.gaussMoffatFit
334 lastRow = self.height - 1
335 state = {"calls": 0}
337 def sometimesRaise(pixels, footprint, initialPars):
338 state["calls"] += 1
339 if state["calls"] > lastRow: # final row only
340 raise RuntimeError("boom")
341 return realFit(pixels, footprint, initialPars)
343 with mock.patch.object(task, "gaussMoffatFit",
344 side_effect=sometimesRaise):
345 result = task.getFluxBasic()
346 self.assertIs(result, task)
347 self.assertIsNone(task.gausMoffatFitPars[lastRow])
349 def test_getFluxBasic_debugPlot(self):
350 # Exercise the debug.plot block. Note this block references several
351 # attributes that the production code never actually sets
352 # (self.psf_gauss_flux, self.spectrum.object_name, ...), so we must
353 # inject them here for the code path to run at all.
354 exp, bbox = makeExposure(self.width, self.height, amp=50.0)
355 task = self._makeTask()
356 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
357 task.initialise(exp, centroid, bbox, self.dispersion)
358 task.debug = types.SimpleNamespace(display=False, plot="all")
359 task.psf_gauss_flux = np.zeros(self.height)
360 task.psf_gauss_psfSigma = np.zeros(self.height)
361 task.psf_gauss_psfMu = np.zeros(self.height)
362 task.spectrum = types.SimpleNamespace(object_name="testStar")
363 with mock.patch("lsst.atmospec.extraction.pl.show"):
364 result = task.getFluxBasic()
365 self.assertIs(result, task)
367 def test_initialise_debugDisplayFailure(self):
368 # A debug display with an invalid backend must be caught and disable
369 # the display without raising (covers the except branch).
370 exp, bbox = makeExposure(self.width, self.height)
371 task = self._makeTask()
372 task.debug = types.SimpleNamespace(
373 display=True, displayBackend="not-a-real-backend",
374 displayItems=[], plot=False)
375 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
376 task.initialise(exp, centroid, bbox, self.dispersion)
377 # The failed display should have been disabled.
378 self.assertFalse(task.debug.display)
380 def test_initialise_debugDisplaySuccess(self):
381 # Patch the real afw.display module so the display path succeeds and
382 # the spectrumBgSub display branch is exercised without a real viewer.
383 import lsst.afw.display as afwDisp
384 fakeDisplayInstance = mock.MagicMock()
385 exp, bbox = makeExposure(self.width, self.height)
386 task = self._makeTask()
387 task.debug = types.SimpleNamespace(
388 display=True, displayBackend="fake",
389 displayItems=["spectrumBgSub"], plot=False)
390 centroid = geom.Point2D(self.width / 2.0, self.height / 2.0)
391 with mock.patch.object(afwDisp, "setDefaultBackend"), \
392 mock.patch.object(afwDisp, "Display",
393 return_value=fakeDisplayInstance):
394 task.initialise(exp, centroid, bbox, self.dispersion)
395 # Display stayed enabled and mtv was invoked (init + spectrumBgSub).
396 self.assertTrue(task.debug.display)
397 self.assertTrue(fakeDisplayInstance.mtv.called)
400class TestMemory(lsst.utils.tests.MemoryTestCase):
401 pass
404def setup_module(module):
405 lsst.utils.tests.init()
408if __name__ == "__main__": 408 ↛ 409line 408 didn't jump to line 409 because the condition on line 408 was never true
409 lsst.utils.tests.init()
410 unittest.main()