Coverage for tests/test_blend.py: 19%
125 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 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/>.
22from typing import Callable, cast
24import numpy as np
25from lsst.scarlet.lite import Blend, Box, Image, Observation, Source
26from lsst.scarlet.lite.component import Component, FactorizedComponent, default_adaprox_parameterization
27from lsst.scarlet.lite.initialization import FactorizedChi2Initialization
28from lsst.scarlet.lite.operators import Monotonicity
29from lsst.scarlet.lite.parameters import Parameter
30from lsst.scarlet.lite.utils import integrated_circular_gaussian
31from numpy.testing import assert_almost_equal, assert_raises
32from scipy.signal import convolve as scipy_convolve
33from utils import ObservationData, ScarletTestCase
36class DummyCubeComponent(Component):
37 def __init__(self, model: Image):
38 super().__init__(model.bands, model.bbox)
39 self._model = Parameter(model.data, {}, 0)
41 @property
42 def data(self) -> np.ndarray:
43 return self._model.x
45 def resize(self, model_box: Box) -> bool:
46 pass
48 def update(self, it: int, input_grad: np.ndarray):
49 pass
51 def get_model(self) -> Image:
52 return Image(self.data, bands=self.bands, yx0=self.bbox.origin)
54 def parameterize(self, parameterization: Callable) -> None:
55 pass
58class TestBlend(ScarletTestCase):
59 def setUp(self):
60 bands = ("g", "r", "i")
61 yx0 = (1000, 2000)
62 # The PSF in each band of the "observation"
63 psfs = np.array([integrated_circular_gaussian(sigma=sigma) for sigma in [1.05, 0.9, 1.2]])
64 # The PSF of the model
65 model_psf = integrated_circular_gaussian(sigma=0.8)
67 # The spectrum of each source
68 spectra = np.array(
69 [
70 [40, 10, 0],
71 [0, 25, 40],
72 [15, 8, 3],
73 [20, 3, 4],
74 [0, 30, 60],
75 ],
76 dtype=float,
77 )
79 # Use a point source for all of the sources
80 morphs = [integrated_circular_gaussian(sigma=sigma) for sigma in [0.8, 2.5, 1.1, 2.1, 1.5]]
81 # Make the second component a disk component
82 morphs[1] = scipy_convolve(morphs[1], model_psf, mode="same")
84 # Give the first two components the same center, and unique centers
85 # for the remaining sources
86 centers = [
87 (1010, 2012),
88 (1010, 2012),
89 (1020, 2023),
90 (1020, 2010),
91 (1025, 2020),
92 ]
94 # Create the simulated image and associated data products
95 test_data = ObservationData(bands, psfs, spectra, morphs, centers, model_psf, yx0=yx0)
97 # Create the Observation
98 variance = np.ones((3, 35, 35), dtype=float) * 1e-2
99 weights = 1 / variance
100 weights = weights / np.max(weights)
101 self.observation = Observation(
102 test_data.convolved,
103 variance,
104 weights,
105 psfs,
106 model_psf[None],
107 bands=bands,
108 bbox=Box(variance.shape[-2:], origin=yx0),
109 )
110 self.data = test_data
111 self.spectra = spectra
112 self.centers = centers
113 self.morphs = morphs
115 components = []
116 for spectrum, center, morph, data_morph in zip(
117 self.spectra, self.centers, self.morphs, self.data.morphs
118 ):
119 components.append(
120 FactorizedComponent(
121 bands=bands,
122 spectrum=spectrum,
123 morph=morph,
124 bbox=data_morph.bbox,
125 peak=center,
126 )
127 )
129 sources = [Source(components[:2])]
130 sources += [Source([component]) for component in components[2:]]
132 self.blend = Blend(sources, self.observation)
134 def test_exact(self):
135 """Test that a blend model initialized with the exact solution
136 builds the model correctly
137 """
138 blend = self.blend
139 self.assertEqual(len(blend.components), 5)
140 self.assertEqual(len(blend.sources), 4)
141 self.assertBoxEqual(blend.bbox, Box(self.data.images.shape[1:], self.observation.bbox.origin))
142 self.assertImageAlmostEqual(blend.get_model(), self.data.images)
143 self.assertImageAlmostEqual(blend.get_model(convolve=True), self.observation.images)
144 self.assertImageAlmostEqual(
145 self.observation.convolve(blend.get_model(), mode="real"),
146 self.observation.images,
147 )
149 # Test that the log likelihood is very small
150 assert_almost_equal([blend.log_likelihood], [0])
152 # Test that grad_log_likelihood updates the loss
153 self.assertListEqual(blend.loss, [])
154 blend._grad_log_likelihood()
155 assert_almost_equal(blend.loss, [0])
157 # Remove one of the sources and calculate the non-zero log_likelihood
158 del blend.sources[-1]
159 # Update the loss function and check that the loss changed
160 blend._grad_log_likelihood()
161 assert_almost_equal(blend.log_likelihood, -60.011720889007485)
162 assert_almost_equal(blend.loss, [0, -60.011720889007485])
164 def test_fit_spectra(self):
165 """Test that fitting the spectra with exact morphologies is
166 identical to the multiband image
167 """
168 np.random.seed(0)
169 blend = self.blend
171 # Change the initial spectra so that they can be fit later
172 for component in blend.components:
173 c = cast(FactorizedComponent, component)
174 c.spectrum[:] = np.random.rand(3) * 10
176 with assert_raises(AssertionError):
177 # Since the spectra have not yet been fit,
178 # the model and images should not be equal
179 self.assertImageEqual(blend.get_model(), self.data.images)
181 # We initialized all of the morphologies exactly,
182 # so fitting the spectra should give a nearly exact solution
183 blend.fit_spectra()
185 self.assertEqual(len(blend.components), 5)
186 self.assertEqual(len(blend.sources), 4)
187 self.assertBoxEqual(blend.bbox, self.observation.bbox)
188 self.assertImageAlmostEqual(blend.get_model(), self.data.images)
189 self.assertImageAlmostEqual(blend.get_model(convolve=True), self.observation.images)
191 def test_fit(self):
192 observation = self.observation
193 np.random.seed(0)
194 images = observation.images.copy()
195 noise = np.random.normal(size=observation.images.shape) * 1e-2
196 observation.images._data += noise
198 monotonicity = Monotonicity((101, 101))
199 init = FactorizedChi2Initialization(observation, self.centers, monotonicity=monotonicity)
201 blend = Blend(init.sources, self.observation).fit_spectra()
202 blend.parameterize(default_adaprox_parameterization)
203 blend.fit(100)
205 self.assertImageAlmostEqual(blend.get_model(convolve=True), images, decimal=1)
207 def test_non_factorized(self):
208 np.random.seed(1)
209 blend = self.blend
210 # Remove the disk component from the first source
211 model = self.spectra[1][:, None, None] * self.morphs[1][None, :, :]
212 yx0 = blend.sources[0].components[1].bbox.origin
213 blend.sources[0].components = blend.sources[0].components[:1]
215 # Change the initial spectra so that they can be fit later
216 for component in blend.components:
217 c = cast(FactorizedComponent, component)
218 c.spectrum[:] = np.random.rand(3) * 10
220 with assert_raises(AssertionError):
221 # Since the spectra have not yet been fit,
222 # the model and images should not be equal
223 self.assertImageEqual(blend.get_model(), self.data.images)
225 # Remove the disk component from the first source
226 blend.sources[0].components = blend.sources[0].components[:1]
227 # Create a new source for the disk with a non-factorized component
228 component = DummyCubeComponent(Image(model, bands=self.blend.observation.bands, yx0=yx0))
229 blend.sources.append(Source([component]))
231 blend.fit_spectra()
233 self.assertEqual(len(blend.components), 5)
234 self.assertEqual(len(blend.sources), 5)
235 self.assertImageAlmostEqual(blend.get_model(), self.data.images)
237 def test_clipping(self):
238 blend = self.blend
240 # Change the initial spectra so that they can be fit later
241 for component in blend.components:
242 c = cast(FactorizedComponent, component)
243 c.spectrum[:] = np.random.rand(3) * 10
245 with assert_raises(AssertionError):
246 # Since the spectra have not yet been fit,
247 # the model and images should not be equal
248 self.assertImageEqual(blend.get_model(), self.data.images)
250 # Add an empty source
251 zero_model = Image.from_box(Box((5, 5), (30, 0)), bands=blend.observation.bands)
252 component = DummyCubeComponent(zero_model)
253 blend.sources.append(Source([component]))
255 blend.fit_spectra(clip=True)
257 self.assertEqual(len(blend.components), 5)
258 self.assertEqual(len(blend.sources), 5)
259 self.assertImageAlmostEqual(blend.get_model(), self.data.images)