Coverage for tests/test_models.py: 13%
191 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-19 10:38 +0000
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-19 10:38 +0000
1# This file is part of lsst.scarlet.lite.
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 os
23from functools import partial
24from typing import cast
26import lsst.scarlet.lite.models as models
27import numpy as np
28from lsst.scarlet.lite import Blend, Box, FistaParameter, Image, Observation, Source
29from lsst.scarlet.lite.component import (
30 Component,
31 default_adaprox_parameterization,
32 default_fista_parameterization,
33)
34from lsst.scarlet.lite.initialization import FactorizedChi2Initialization
35from lsst.scarlet.lite.models import (
36 CartesianFrame,
37 EllipseFrame,
38 EllipticalParametricComponent,
39 FittedPsfBlend,
40 FittedPsfObservation,
41 FreeFormComponent,
42 ParametricComponent,
43)
44from lsst.scarlet.lite.operators import Monotonicity
45from lsst.scarlet.lite.parameters import AdaproxParameter, parameter, relative_step
46from lsst.scarlet.lite.utils import integrated_circular_gaussian
47from numpy.testing import assert_array_equal
48from utils import ScarletTestCase
51def parameterize(component: Component):
52 assert isinstance(component, ParametricComponent)
53 component._spectrum = AdaproxParameter(
54 component.spectrum,
55 step=partial(relative_step, factor=1e-2, minimum=1e-16),
56 )
57 component._params = AdaproxParameter(
58 component._params.x,
59 step=1e-2,
60 )
63class TestFreeForm(ScarletTestCase):
64 def setUp(self) -> None:
65 yx0 = (1000, 2000)
66 filename = os.path.join(__file__, "..", "..", "data", "hsc_cosmos_35.npz")
67 filename = os.path.abspath(filename)
68 data = np.load(filename)
69 self.data = data
70 model_psf = integrated_circular_gaussian(sigma=0.8)
71 self.detect = np.sum(data["images"], axis=0)
72 self.centers = np.array([data["catalog"]["y"], data["catalog"]["x"]]).T + np.array(yx0)
73 bands = data["filters"]
74 self.observation = Observation(
75 Image(data["images"], bands=bands, yx0=yx0),
76 Image(data["variance"], bands=bands, yx0=yx0),
77 Image(1 / data["variance"], bands=bands, yx0=yx0),
78 data["psfs"],
79 model_psf[None],
80 bands=bands,
81 )
83 def tearDown(self):
84 del self.data
86 def test_free_form_component(self):
87 images = self.data["images"]
89 # Test with no thresholding (sparsity)
90 sources = []
91 for i in range(5):
92 component = FreeFormComponent(
93 self.observation.bands,
94 np.ones(5),
95 images[i].copy(),
96 self.observation.bbox,
97 )
98 sources.append(Source([component]))
100 blend = Blend(sources, self.observation).fit_spectra()
101 blend.parameterize(default_adaprox_parameterization)
102 blend.fit(12, e_rel=1e-6)
104 # Test with thresholding (sparsity)
105 sources = []
106 for i in range(5):
107 component = FreeFormComponent(
108 self.observation.bands,
109 np.ones(5),
110 images[i].copy(),
111 self.observation.bbox,
112 bg_rms=self.observation.noise_rms,
113 bg_thresh=0.25,
114 min_area=4,
115 )
116 sources.append(Source([component]))
118 blend = Blend(sources, self.observation).fit_spectra()
119 blend.parameterize(default_adaprox_parameterization)
120 blend.fit(12, e_rel=1e-6)
122 # Test with peak centers specified
123 sources = []
124 peaks = list(np.array([self.data["catalog"]["y"], self.data["catalog"]["x"]]).T.astype(int))
125 for i in range(5):
126 component = FreeFormComponent(
127 self.observation.bands,
128 np.ones(5),
129 images[i].copy(),
130 self.observation.bbox,
131 peaks=peaks,
132 )
133 sources.append(Source([component]))
135 blend = Blend(sources, self.observation).fit_spectra()
136 blend.parameterize(default_adaprox_parameterization)
137 blend.fit(12, e_rel=1e-6)
139 # Tests for code blocks that are difficult to reach,
140 # to complete test coverage
141 component = blend.sources[-1].components[0]
142 self.assertFalse(component.resize(self.observation.bbox))
143 component.morph[:] = 0
144 component.prox_morph(component.morph)
147class TestParametric(ScarletTestCase):
148 def setUp(self) -> None:
149 filename = os.path.join(__file__, "..", "..", "data", "hsc_cosmos_35.npz")
150 filename = os.path.abspath(filename)
151 data = np.load(filename)
152 self.data = data
153 self.model_psf = integrated_circular_gaussian(sigma=0.8)
154 self.detect = np.sum(data["images"], axis=0)
155 self.centers = np.array([data["catalog"]["y"], data["catalog"]["x"]]).T
156 bands = data["filters"]
157 self.observation = Observation(
158 Image(data["images"], bands=bands),
159 Image(data["variance"], bands=bands),
160 Image(1 / data["variance"], bands=bands),
161 data["psfs"],
162 self.model_psf[None],
163 bands=bands,
164 )
166 def tearDown(self):
167 del self.data
169 def test_cartesian_frame(self):
170 bbox = Box((31, 60), (1, 2))
171 frame = CartesianFrame(bbox)
172 y = np.linspace(1, 31, 31)
173 x = np.linspace(2, 61, 60)
174 x, y = np.meshgrid(x, y)
176 self.assertBoxEqual(frame.bbox, bbox)
177 self.assertTupleEqual(frame.shape, bbox.shape)
178 assert_array_equal(frame.x_grid, x)
179 assert_array_equal(frame.y_grid, y)
180 self.assertIsNone(frame._r2)
181 self.assertIsNone(frame._r)
183 def test_ellipse_frame(self):
184 y0 = 23
185 x0 = 36
186 major = 3
187 minor = 2
188 theta = 2 * np.pi / 3
189 bbox = Box((50, 39), (2, 7))
190 r_min = 1e-20
191 frame = EllipseFrame(
192 y0,
193 x0,
194 major,
195 minor,
196 theta,
197 bbox,
198 )
200 self.assertEqual(frame.x0, x0)
201 self.assertEqual(frame.y0, y0)
202 self.assertEqual(frame.major, major)
203 self.assertEqual(frame.minor, minor)
204 self.assertEqual(frame.theta, theta)
205 self.assertEqual(frame.bbox, bbox)
207 y = np.linspace(2, 51, 50)
208 x = np.linspace(7, 45, 39)
209 x, y = np.meshgrid(x, y)
210 r2 = frame._xa**2 + frame._yb**2 + r_min**2
211 r = np.sqrt(r2)
213 self.assertEqual(frame._sin, np.sin(theta))
214 self.assertEqual(frame._cos, np.cos(theta))
215 self.assertBoxEqual(frame.bbox, bbox)
216 self.assertTupleEqual(frame.shape, bbox.shape)
217 assert_array_equal(frame.x_grid, x)
218 assert_array_equal(frame.y_grid, y)
219 assert_array_equal(frame.r2_grid, r2)
220 assert_array_equal(frame.r_grid, r)
222 # There doesn't seem to be much utility to testing the
223 # gradients exactly, however we test that they all run
224 # properly.
225 gradients = (
226 frame.grad_x0,
227 frame.grad_y0,
228 frame.grad_major,
229 frame.grad_minor,
230 frame.grad_theta,
231 )
233 input_grad = np.ones(bbox.shape, dtype=float)
234 original_grad = input_grad.copy()
235 for grad in gradients:
236 grad(input_grad, False)
237 grad(input_grad, True)
239 # Make sure that none of the methods changed the input gradient
240 assert_array_equal(input_grad, original_grad)
242 def test_parametric_component(self):
243 observation = self.observation
244 bands = observation.bands
245 spectrum = np.ones((observation.n_bands,), dtype=float)
246 frame = CartesianFrame(observation.bbox)
248 # Test integrated Gaussian PSF and sersic
249 sources = []
250 for idx, center in enumerate(self.centers):
251 # Get the integer center of the source
252 cy, cx = int(np.round(center[0])), int(np.round(center[1]))
253 # For now we use a fixed bounding box that is the size of
254 # the observed PSF image
255 bbox = Box((41, 41), origin=(cy - 20, cx - 20)) & observation.bbox
256 # Keep track of the initial positions
257 yi, xi = cy, cx
258 # Restrict the values of the parameters
259 _proxmin = np.array([yi - 2, xi - 2, 1e-1, 1e-1, -np.pi / 2, 1.1])
260 _proxmax = np.array([yi + 2, xi + 2, frame.shape[-2] / 2, frame.shape[-1] / 2, np.pi / 2, 3])
262 __proxmin = np.array([yi - 2, xi - 2, 0.8])
263 __proxmax = np.array([yi + 2, xi + 2, 1.2])
265 # Initialize a PSF-like component of the source using a
266 # non-pixel integrated gaussian
267 component = ParametricComponent(
268 bands,
269 bbox,
270 spectrum=parameter(spectrum.copy()),
271 morph_params=parameter(np.array([center[0], center[1], 0.8])),
272 morph_func=models.integrated_gaussian,
273 morph_grad=models.grad_integrated_gaussian,
274 morph_prox=partial(models.bounded_prox, proxmin=__proxmin, proxmax=__proxmax),
275 morph_step=np.array([1e-2, 1e-2, 1e-2]),
276 prox_spectrum=lambda x: x,
277 )
278 # Define the component to use ADAPROX as the optimizer
279 components = [component]
281 # Initialize an n=1 sersic component
282 component = EllipticalParametricComponent(
283 bands,
284 bbox,
285 spectrum=parameter(spectrum.copy()),
286 morph_params=parameter(np.array([center[0], center[1], 2 * 1.2**2, 2 * 1.2**2, 0.0, 1])),
287 morph_func=models.sersic,
288 morph_grad=models.grad_sersic,
289 morph_prox=partial(models.bounded_prox, proxmin=_proxmin, proxmax=_proxmax),
290 morph_step=np.array([1e-2, 1e-2, 1e-3, 1e-3, 1e-2, 1e-2]),
291 )
292 # Define the component to use ADAPROX as the optimizer
293 components.append(component)
295 component = EllipticalParametricComponent(
296 bands,
297 bbox,
298 spectrum=parameter(spectrum.copy()),
299 morph_params=parameter(np.array([center[0], center[1], 2 * 1.2**2, 2 * 1.2**2, 0.0, 1])),
300 morph_func=models.sersic,
301 morph_grad=models.grad_sersic,
302 morph_prox=partial(models.bounded_prox, proxmin=_proxmin, proxmax=_proxmax),
303 morph_step=np.array([1e-2, 1e-2, 1e-3, 1e-3, 1e-2, 1e-2]),
304 )
305 # Define the component to use ADAPROX as the optimizer
306 components.append(component)
308 # Create a new source using the two components
309 sources.append(Source(components))
311 # Fit the models
312 blend = Blend(sources, observation)
313 blend.parameterize(parameterize)
314 blend.fit_spectra()
316 # Check properties of a component
317 component = cast(ParametricComponent, blend.components[0])
318 self.assertTupleEqual(component.peak, tuple(self.centers[0]))
319 self.assertEqual(component.y0, component._params.x[0])
320 self.assertEqual(component.x0, component._params.x[1])
321 assert_array_equal(component.morph_step, np.array([1e-2, 1e-2, 1e-2]))
323 component = cast(EllipticalParametricComponent, blend.components[1])
324 self.assertTupleEqual(component.peak, tuple(self.centers[0]))
325 self.assertEqual(component.y0, component._params.x[0])
326 self.assertEqual(component.x0, component._params.x[1])
327 assert_array_equal(component.semi_major, component._params.x[2])
328 assert_array_equal(component.semi_minor, component._params.x[3])
329 assert_array_equal(component.theta, component._params.x[4])
330 assert_array_equal(component.ellipse_params, component._params.x[:5])
331 assert_array_equal(component.radial_params, component._params.x[5:])
332 self.assertIsNotNone(component.morph_prox)
333 self.assertIsNotNone(component.morph_grad)
335 blend.fit(12)
337 # Test elliptical and circular Gaussian models
338 sources = []
339 for idx, center in enumerate(self.centers):
340 # Get the integer center of the source
341 cy, cx = int(np.round(center[0])), int(np.round(center[1]))
342 # For now we use a fixed bounding box that is the size of
343 # the observed PSF image
344 bbox = Box((41, 41), origin=(cy - 20, cx - 20)) & observation.bbox
345 # Keep track of the initial positions
346 yi, xi = cy, cx
347 # Restrict the values of the parameters
348 _proxmin = np.array([yi - 2, xi - 2, 1e-1, 1e-1, -np.pi / 2])
349 _proxmax = np.array([yi + 2, xi + 2, frame.shape[-2] / 2, frame.shape[-1] / 2, np.pi / 2])
351 __proxmin = np.array([yi - 2, xi - 2])
352 __proxmax = np.array([yi + 2, xi + 2])
354 # Initialize a PSF-like component of the source using a
355 # non-pixel integrated gaussian
356 component = ParametricComponent(
357 bands,
358 bbox,
359 spectrum=parameter(spectrum.copy()),
360 morph_params=parameter(np.array([center[0], center[1]])),
361 morph_func=partial(models.circular_gaussian, sigma=0.8),
362 morph_grad=partial(models.grad_circular_gaussian, sigma=0.8),
363 morph_prox=partial(models.bounded_prox, proxmin=__proxmin, proxmax=__proxmax),
364 morph_step=np.array([1e-2, 1e-2, 1e-2]),
365 )
366 # Define the component to use ADAPROX as the optimizer
367 components = [component]
369 # Initialize an n=1 sersic component
370 component = EllipticalParametricComponent(
371 bands,
372 bbox,
373 spectrum=parameter(spectrum.copy()),
374 morph_params=parameter(np.array([center[0], center[1], 2 * 1.2**2, 2 * 1.2**2, 0.0])),
375 morph_func=models.gaussian2d,
376 morph_grad=models.grad_gaussian2,
377 morph_prox=partial(models.bounded_prox, proxmin=_proxmin, proxmax=_proxmax),
378 morph_step=np.array([1e-2, 1e-2, 1e-3, 1e-3, 1e-2, 1e-2]),
379 )
380 # Define the component to use ADAPROX as the optimizer
381 components.append(component)
382 # Create a new source using the two components
383 sources.append(Source(components))
385 # Fit the models
386 blend = Blend(sources, observation)
387 blend.parameterize(parameterize)
388 blend.fit_spectra()
389 blend.fit(12)
391 def test_psf_fitting(self):
392 # Use flat weights for FISTA optimization
393 weights = np.ones(self.data["images"].shape)
395 monotonicity = Monotonicity((101, 101))
397 observation = FittedPsfObservation(
398 self.data["images"],
399 self.data["variance"],
400 weights,
401 self.data["psfs"],
402 self.model_psf[None],
403 bands=self.data["filters"],
404 )
406 def obs_params(cls):
407 if isinstance(cls, FittedPsfObservation):
408 cls._fitted_kernel = FistaParameter(cls._fitted_kernel.x, step=1e-2)
410 init = FactorizedChi2Initialization(observation, self.centers, monotonicity=monotonicity)
411 blend = FittedPsfBlend(init.sources, observation).fit_spectra()
412 blend.parameterize(default_fista_parameterization)
413 cast(FittedPsfObservation, blend.observation).parameterize(obs_params)
414 blend.fit(12, e_rel=1e-4)