Coverage for tests/test_modeller.py: 92%
288 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-23 09:54 +0000
1# This file is part of multiprofit.
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 math
23import time
25import numpy as np
26import pytest
28import lsst.gauss2d as g2
29import lsst.gauss2d.fit as g2f
30from lsst.multiprofit.componentconfig import (
31 CentroidConfig,
32 FluxFractionParameterConfig,
33 FluxParameterConfig,
34 GaussianComponentConfig,
35 ParameterConfig,
36 SersicComponentConfig,
37 SersicIndexParameterConfig,
38)
39from lsst.multiprofit.model_utils import make_image_gaussians, make_psf_model_null
40from lsst.multiprofit.modelconfig import ModelConfig
41from lsst.multiprofit.modeller import FitInputs, LinearGaussians, Modeller, fit_methods_linear
42from lsst.multiprofit.observationconfig import CoordinateSystemConfig, ObservationConfig
43from lsst.multiprofit.sourceconfig import ComponentGroupConfig, SourceConfig
44from lsst.multiprofit.utils import get_params_uniq
46sigma_inv = 1e4
49@pytest.fixture(scope="module")
50def channels() -> dict[str, g2f.Channel]:
51 """Return dict of generic RGB channels."""
52 return {band: g2f.Channel.get(band) for band in ("R", "G", "B")}
55@pytest.fixture(scope="module")
56def data(channels) -> g2f.DataD:
57 """Return initialized data in all bands."""
58 n_rows, n_cols = 25, 27
59 x_min, y_min = 0, 0
61 dn_rows, dn_cols = 2, -3
62 dx_min, dy_min = -1, 1
64 observations = []
65 for idx, band in enumerate(channels):
66 config = ObservationConfig(
67 band=band,
68 coordsys=CoordinateSystemConfig(
69 x_min=x_min + idx * dx_min,
70 y_min=y_min + idx * dy_min,
71 ),
72 n_rows=n_rows + idx * dn_rows,
73 n_cols=n_cols + idx * dn_cols,
74 )
75 observation = config.make_observation()
76 observation.image.fill(0)
77 observation.sigma_inv.fill(sigma_inv)
78 observation.mask_inv.fill(1e4)
79 observations.append(observation)
80 return g2f.DataD(observations)
83@pytest.fixture(scope="module")
84def psf_models(channels) -> list[g2f.PsfModel]:
85 """Return a double Gaussian PSF model for each band."""
86 rho, size_x, size_y = 0.12, 1.6, 1.2
87 drho, dsize_x, dsize_y = -0.3, 1.1, 1.9
88 drho_chan, dsize_x_chan, dsize_y_chan = 0.03, 0.12, 0.14
89 frac, dfrac = 0.62, -0.08
91 n_components = 2
92 psf_models = []
94 for idx_chan, channel in enumerate(channels.values()):
95 frac_chan = frac + idx_chan * dfrac
96 config = SourceConfig(
97 component_groups={
98 "psf": ComponentGroupConfig(
99 components_gauss={
100 str(idx): GaussianComponentConfig(
101 rho=ParameterConfig(value_initial=rho + idx * drho + idx_chan * drho_chan),
102 size_x=ParameterConfig(
103 value_initial=size_x + idx * dsize_x + idx_chan * dsize_x_chan
104 ),
105 size_y=ParameterConfig(
106 value_initial=size_y + idx * dsize_y + idx_chan * dsize_y_chan
107 ),
108 **(
109 {
110 "flux": FluxParameterConfig(value_initial=1.0, fixed=True),
111 "fluxfrac": FluxFractionParameterConfig(
112 value_initial=frac_chan, fixed=False
113 ),
114 }
115 if (idx == 0)
116 else {}
117 ),
118 )
119 for idx in range(n_components)
120 },
121 is_fractional=True,
122 )
123 },
124 )
125 config.validate()
126 psf_model, priors = config.make_psf_model(
127 [
128 component_group.get_fluxes_default(
129 channels=(g2f.Channel.NONE,),
130 component_configs=component_group.get_component_configs(),
131 is_fractional=component_group.is_fractional,
132 )
133 for component_group in config.component_groups.values()
134 ]
135 )
136 psf_models.append(psf_model)
137 return psf_models
140@pytest.fixture(scope="module")
141def model(channels, data, psf_models) -> g2f.ModelD:
142 """Return the configured model."""
143 rho, size_x, size_y, sersicn, flux = 0.4, 1.5, 1.9, 1.0, 4.7
144 drho, dsize_x, dsize_y, dsersicn, dflux = -0.9, 2.5, 5.4, 3.0, 13.9
146 components_sersic = {}
147 fluxes_group = []
149 # Linear interpolators fail to compute accurate likelihoods at knot values
150 is_linear_interp = (
151 g2f.SersicMixComponentIndexParameterD(
152 interpolator=SersicComponentConfig().get_interpolator(4)
153 ).interptype
154 == g2f.InterpType.linear
155 )
157 for idx, name in enumerate(("exp", "dev")):
158 components_sersic[name] = SersicComponentConfig(
159 rho=ParameterConfig(value_initial=rho + idx * drho),
160 size_x=ParameterConfig(value_initial=size_x + idx * dsize_x),
161 size_y=ParameterConfig(value_initial=size_y + idx * dsize_y),
162 sersic_index=SersicIndexParameterConfig(
163 # Add a small offset since 1.0 and 4.0 are bound to be knots
164 value_initial=sersicn + idx * dsersicn + 1e-4 * is_linear_interp,
165 fixed=idx == 0,
166 prior_mean=None,
167 ),
168 )
169 fluxes_comp = {
170 channel: flux + idx_channel * dflux * idx for idx_channel, channel in enumerate(channels.values())
171 }
172 fluxes_group.append(fluxes_comp)
174 modelconfig = ModelConfig(
175 sources={
176 "src": SourceConfig(
177 component_groups={
178 "": ComponentGroupConfig(
179 components_sersic=components_sersic,
180 centroids={
181 "default": CentroidConfig(
182 x=ParameterConfig(value_initial=12.14, fixed=True),
183 y=ParameterConfig(value_initial=13.78, fixed=True),
184 )
185 },
186 ),
187 }
188 ),
189 },
190 )
191 model = modelconfig.make_model([[fluxes_group]], data=data, psf_models=psf_models)
192 model.setup_evaluators(g2f.EvaluatorMode.loglike_image)
193 model.evaluate()
195 rng = np.random.default_rng(2)
197 n_obs = len(model.data)
198 for idx_obs in range(n_obs):
199 observation = model.data[idx_obs]
200 output = model.outputs[idx_obs]
201 observation.image.data.flat = (
202 output.data.flat + rng.standard_normal(output.data.size) / observation.sigma_inv.data.flat
203 )
205 return model
208@pytest.fixture
209def model_func_scope(model) -> g2f.ModelD:
210 """Return a shallow copy of the configured model."""
211 model_func_scope = g2f.ModelD(data=model.data, psfmodels=model.psfmodels, sources=model.sources)
212 return model_func_scope
215@pytest.fixture(scope="module")
216def psf_observations(psf_models) -> list[g2f.ObservationD]:
217 """Return the PSF model observations for each band."""
218 config = ObservationConfig(n_rows=17, n_cols=19)
219 rng = np.random.default_rng(1)
221 observations = []
222 for psf_model in psf_models:
223 observation = config.make_observation()
224 # Have to make a duplicate image here because one can only call
225 # make_image_gaussians with an owning pointer, whereas
226 # observation.image is a reference
227 image = g2.ImageD(observation.image.data)
228 # Make the kernel centered
229 gaussians_source = psf_model.gaussians(g2f.Channel.NONE)
230 for idx in range(len(gaussians_source)):
231 gaussian_idx = gaussians_source.at(idx)
232 gaussian_idx.centroid.x = image.n_cols / 2.0
233 gaussian_idx.centroid.y = image.n_rows / 2.0
234 gaussians_kernel = g2.Gaussians([g2.Gaussian()])
235 make_image_gaussians(
236 gaussians_source=gaussians_source,
237 gaussians_kernel=gaussians_kernel,
238 output=image,
239 )
240 image.data.flat += 1e-4 * rng.standard_normal(image.data.size)
241 observation.mask_inv.fill(1)
242 observation.sigma_inv.fill(1e3)
243 observations.append(observation)
244 return observations
247@pytest.fixture(scope="module")
248def psf_fit_models(psf_models, psf_observations):
249 """Return initialized models for each band's PSF."""
250 psf_null = [make_psf_model_null()]
251 return [
252 g2f.ModelD(g2f.DataD([observation]), psf_null, [g2f.Source(psf_model.components)])
253 for psf_model, observation in zip(psf_models, psf_observations)
254 ]
257def test_model_evaluation(channels, model, model_func_scope):
258 """Test that each kind of model evaluation works correctly."""
259 with pytest.raises(RuntimeError):
260 model_func_scope.evaluate()
262 printout = False
263 # Freeze the PSF params - they can't be fit anyway
264 for m in (model, model_func_scope):
265 for psf_model in m.psfmodels:
266 params = psf_model.parameters()
267 for param in params:
268 param.fixed = True
270 model.setup_evaluators(print=printout, force=True)
271 model.evaluate()
273 n_priors = 0
274 n_obs = len(model.data)
275 n_rows = np.zeros(n_obs, dtype=int)
276 n_cols = np.zeros(n_obs, dtype=int)
277 datasizes = np.zeros(n_obs, dtype=int)
278 ranges_params = [None] * n_obs
279 params_free = tuple(get_params_uniq(model_func_scope, fixed=False))
281 # There's one extra validation array
282 n_params_jac = len(params_free) + 1
283 assert n_params_jac > 1
285 for idx_obs in range(n_obs):
286 observation = model.data[idx_obs]
287 n_rows[idx_obs] = observation.image.n_rows
288 n_cols[idx_obs] = observation.image.n_cols
289 datasizes[idx_obs] = n_rows[idx_obs] * n_cols[idx_obs]
290 params = tuple(get_params_uniq(model, fixed=False, channel=observation.channel))
291 n_params_obs = len(params)
292 ranges_params_obs = [0] * (n_params_obs + 1)
293 for idx_param in range(n_params_obs):
294 ranges_params_obs[idx_param + 1] = params_free.index(params[idx_param]) + 1
295 ranges_params[idx_obs] = ranges_params_obs
297 n_free_first = len(ranges_params[0])
298 assert all([len(rp) == n_free_first for rp in ranges_params[1:]])
300 jacobians = [None] * n_obs
301 residuals = [None] * n_obs
302 datasize = np.sum(datasizes) + n_priors
303 jacobian = np.zeros((datasize, n_params_jac))
304 residual = np.zeros(datasize)
306 offset = 0
307 for idx_obs in range(n_obs):
308 size_obs = datasizes[idx_obs]
309 end = offset + size_obs
310 shape = (n_rows[idx_obs], n_cols[idx_obs])
311 jacobians_obs = [None] * n_params_jac
312 for idx_jac in range(n_params_jac):
313 jacobians_obs[idx_jac] = g2.ImageD(jacobian[offset:end, idx_jac].view().reshape(shape))
314 jacobians[idx_obs] = jacobians_obs
315 residuals[idx_obs] = g2.ImageD(residual[offset:end].view().reshape(shape))
316 offset = end
318 model.setup_evaluators(evaluatormode=g2f.EvaluatorMode.loglike)
319 loglike_init = model.evaluate()
321 model_func_scope.setup_evaluators(
322 evaluatormode=g2f.EvaluatorMode.jacobian,
323 outputs=jacobians,
324 residuals=residuals,
325 print=printout,
326 )
327 model_func_scope.verify_jacobian()
328 loglike_jac = model_func_scope.evaluate()
330 assert all(np.isclose(loglike_init, loglike_jac))
333@pytest.fixture(scope="module")
334def psf_models_linear_gaussians(channels, psf_models):
335 """Return individual Gaussians for each PSF model."""
336 gaussians = [None] * len(psf_models)
337 for idx, psf_model in enumerate(psf_models):
338 params = psf_model.parameters(paramfilter=g2f.ParamFilter(nonlinear=False, channel=g2f.Channel.NONE))
339 params[0].fixed = False
340 gaussians[idx] = LinearGaussians.make(psf_model, is_psf=True)
341 # Return the param to its original state
342 params[0].fixed = True
343 return gaussians
346def test_make_psf_source_linear(psf_models, psf_models_linear_gaussians):
347 """Test that the list of PSF Gaussians matches the model."""
348 for psf_model, linear_gaussians in zip(psf_models, psf_models_linear_gaussians):
349 gaussians = psf_model.gaussians(g2f.Channel.NONE)
350 assert len(gaussians) == (
351 len(linear_gaussians.gaussians_free) + len(linear_gaussians.gaussians_fixed)
352 )
355def test_modeller(model):
356 """Test that Modellers can fit models and return sensible values."""
357 # For debugging purposes
358 printout = False
359 # Force to ensure test-order independence
360 model.setup_evaluators(evaluatormode=g2f.EvaluatorMode.loglike_image, force=True)
361 # Get the model images
362 model.evaluate()
363 rng = np.random.default_rng(3)
365 for idx_obs, observation in enumerate(model.data):
366 output = model.outputs[idx_obs]
367 observation.image.data.flat = (
368 output.data.flat + rng.standard_normal(output.data.size) / observation.sigma_inv.data.flat
369 )
371 # Freeze the PSF params - they can't be fit anyway
372 for psf_model in model.psfmodels:
373 for param in psf_model.parameters():
374 param.fixed = True
376 params_free = tuple(get_params_uniq(model, fixed=False))
377 values_true = tuple(param.value for param in params_free)
379 modeller = Modeller()
381 dloglike = model.compute_loglike_grad(verify=True, findiff_frac=1e-8, findiff_add=1e-8)
382 assert all(np.isfinite(dloglike))
384 time_init = time.process_time()
385 kwargs_fit = dict(ftol=1e-6, xtol=1e-6)
387 for delta_param in (0, 0.2):
388 model = g2f.ModelD(data=model.data, psfmodels=model.psfmodels, sources=model.sources)
389 values_init = values_true
390 if delta_param != 0:
391 for param, value_init in zip(params_free, values_init):
392 param.value = value_init
393 try:
394 param.value_transformed += delta_param
395 except RuntimeError:
396 param.value_transformed -= delta_param
398 model.setup_evaluators(evaluatormode=g2f.EvaluatorMode.loglike)
399 loglike_init = np.array(model.evaluate())
400 results = modeller.fit_model(model, **kwargs_fit)
401 params_best = results.params_best
403 for param, value in zip(params_free, params_best):
404 param.value_transformed = value
406 loglike_noprior = model.evaluate()
407 assert np.sum(loglike_noprior) > np.sum(loglike_init)
409 errors = modeller.compute_variances(model)
410 # TODO: This should check >0, and < (some reasonable value)
411 # However, but the scipy least squares does not do a great job
412 # optimizing and the loglike_grad isn't even negative...
413 assert np.all(np.isfinite(errors))
415 if printout: 415 ↛ 416line 415 didn't jump to line 416 because the condition on line 415 was never true
416 print(
417 f"got loglike={loglike_noprior} (init={sum(loglike_noprior)})"
418 f" from modeller.fit_model in t={time.process_time() - time_init:.3e}, x={params_best},"
419 f" results: \n{results}"
420 )
422 loglike_noprior_sum = sum(loglike_noprior)
423 for offset in (0, 1e-6):
424 for param, value in zip(params_free, params_best):
425 param.value_transformed = value
426 priors = tuple(
427 g2f.GaussianPrior(param, param.value_transformed + offset, 1.0, transformed=True)
428 for param in params_free
429 )
430 if offset == 0:
431 for p in priors:
432 assert p.evaluate().loglike == 0
433 assert p.loglike_const_terms[0] == -math.log(math.sqrt(2 * math.pi))
434 model_new = g2f.ModelD(
435 data=model.data, psfmodels=model.psfmodels, sources=model.sources, priors=priors
436 )
437 model_new.setup_evaluators(evaluatormode=g2f.EvaluatorMode.loglike)
438 loglike_init = sum(loglike_eval for loglike_eval in model_new.evaluate())
439 if offset == 0:
440 assert np.isclose(loglike_init, loglike_noprior_sum, rtol=1e-10, atol=1e-10)
441 else:
442 assert loglike_init < loglike_noprior_sum
444 time_init = time.process_time()
445 results = modeller.fit_model(model_new, **kwargs_fit)
446 time_init = time.process_time() - time_init
447 loglike_new = -results.result.cost
448 for param, value in zip(params_free, results.params_best):
449 param.value_transformed = value
451 model_new.setup_evaluators(evaluatormode=g2f.EvaluatorMode.loglike)
452 loglike_model = sum(loglike_eval for loglike_eval in model_new.evaluate())
453 assert np.isclose(loglike_new, loglike_model, rtol=1e-10, atol=1e-10)
454 # This should be > 0. TODO: Determine why it isn't always
455 assert (loglike_new - loglike_init) > -1e-3
457 if printout: 457 ↛ 458line 457 didn't jump to line 458 because the condition on line 457 was never true
458 print(
459 f"got loglike={loglike_new} (first={loglike_noprior})"
460 f" from modeller.fit_model in t={time_init:.3e}, x={results.params_best},"
461 f" results: \n{results}"
462 )
463 # Adding a suitably-scaled prior far from the truth should always
464 # worsen loglikel, but doesn't - why? noise bias? bad convergence?
465 # assert (loglike_new >= loglike_noprior) == (offset == 0)
467 # Return parameters to original values since they are shared between
468 # model instances, which is a useful persistence test.
469 for param, value in zip(params_free, values_true):
470 param.value = value
473def test_psf_model_fit(psf_fit_models):
474 """Test that the PSF models evaluate correctly."""
475 for model in psf_fit_models:
476 params = get_params_uniq(model.sources[0])
477 params_freed = set()
478 for param in params:
479 # Fitting the total flux won't work in a fractional model (yet)
480 if isinstance(param, g2f.IntegralParameterD):
481 assert param.fixed
482 else:
483 params_freed.add(param)
484 param.fixed = False
485 # Necessary whenever parameters are freed/fixed
486 model.setup_evaluators(g2f.EvaluatorMode.jacobian, force=True)
487 errors = model.verify_jacobian(rtol=5e-4, atol=5e-4, findiff_add=1e-6, findiff_frac=1e-6)
488 if errors: 488 ↛ 489line 488 didn't jump to line 489 because the condition on line 488 was never true
489 import matplotlib.pyplot as plt
491 print(model.parameters())
493 fitinputs = FitInputs.from_model(model)
494 model.setup_evaluators(
495 evaluatormode=g2f.EvaluatorMode.jacobian,
496 outputs=fitinputs.jacobians,
497 residuals=fitinputs.residuals,
498 print=True,
499 force=True,
500 )
501 model.evaluate(print=True)
502 assert (fitinputs.jacobians[0][0].data == 0).all()
503 assert np.sum(np.abs(fitinputs.jacobians[0][1].data)) > 0
504 model.setup_evaluators(evaluatormode=g2f.EvaluatorMode.loglike_image)
505 model.evaluate()
506 outputs = model.outputs
507 diffs = [g2.ImageD(img.data.copy()) for img in outputs]
508 delta = 1e-5
509 param.value -= delta
510 model.evaluate()
511 for diff, output in zip(diffs, outputs):
512 diff = (output.data - diff.data) / delta
513 jacobian = fitinputs.jacobians[0][1].data
514 fig, ax = plt.subplots(1, 2)
515 ax[0].imshow(diff)
516 ax[1].imshow(jacobian)
517 plt.show()
518 assert len(errors) == 0
519 # Return the params to their original settings
520 for param in params_freed:
521 param.fixed = True
524def test_psf_models_linear_gaussians(data, psf_models_linear_gaussians, psf_observations):
525 """Test that PSF model linear Gaussians can be used for least squares
526 fitting.
527 """
528 results = [None] * len(psf_observations)
529 for idx, (gaussians_linear, observation_psf) in enumerate(
530 zip(psf_models_linear_gaussians, psf_observations)
531 ):
532 results[idx] = Modeller.fit_gaussians_linear(
533 gaussians_linear=gaussians_linear,
534 observation=observation_psf,
535 fit_methods=fit_methods_linear,
536 plot=False,
537 )
538 assert len(results[idx]) > 0
541def test_modeller_fit_linear(model):
542 """Test that a Modeller can do linear fitting."""
543 modeller = Modeller()
544 params = get_params_uniq(model)
545 params_linear_free = {}
546 params_other = {}
547 for param in params:
548 (params_linear_free if (param.free and param.linear) else params_other)[param] = param.value
549 results = modeller.fit_model_linear(model, validate=True)
550 assert results is not None
552 for param, value in params_linear_free.items():
553 assert value != param.value
554 param.value = value
555 assert all([param.value == value for param, value in params_other.items()])