Coverage for tests/test_plotDiaSourceLightcurve.py: 97%
141 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 15:23 +0000
1# This file is part of analysis_ap.
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 unittest
24import lsst.afw.image
25import lsst.geom
26import lsst.meas.base.tests
27import lsst.utils.tests
28import numpy as np
29import pandas as pd
30import PIL
32from lsst.analysis.ap import plotDiaSourceLightcurve
34# Pull in the DATA fixture from the cutouts test so we get the full set of
35# flag columns expected by _annotate_image.
36from test_plotImageSubtractionCutouts import DATA, make_test_images, skyCenter
39DIA_OBJECT_ID = 999999999999000001
42# Add the extra columns required by the lightcurve task.
43def _augment_with_object_columns(data):
44 data = data.copy()
45 data["diaObjectId"] = [DIA_OBJECT_ID, DIA_OBJECT_ID]
46 data["midpointMjdTai"] = [60100.5, 60110.5]
47 return data
50class _StubApdbQuery:
51 """Minimal DbQuery stub returning canned sources/forced DataFrames.
53 Tracks the number of calls so tests can verify the per-diaObject cache.
54 """
56 def __init__(self, sources, forced):
57 self._sources = sources
58 self._forced = forced
59 self.source_calls = 0
60 self.forced_calls = 0
62 def load_sources_for_object(self, dia_object_id, exclude_flagged=False, limit=100000):
63 self.source_calls += 1
64 self.last_source_kwargs = {"exclude_flagged": exclude_flagged, "limit": limit}
65 return self._sources.copy()
67 def load_forced_sources_for_object(self, dia_object_id, exclude_flagged=False, limit=100000):
68 self.forced_calls += 1
69 self.last_forced_kwargs = {"exclude_flagged": exclude_flagged, "limit": limit}
70 return self._forced.copy()
73def _make_sources_frame(visits, bands, mjds, fluxes, errs, with_err=True):
74 frame = pd.DataFrame({
75 "visit": visits,
76 "band": bands,
77 "midpointMjdTai": mjds,
78 "psfFlux": fluxes,
79 })
80 if with_err:
81 frame["psfFluxErr"] = errs
82 return frame
85class TestPlotDiaSourceLightcurve(lsst.utils.tests.TestCase):
86 """Tests for PlotDiaSourceLightcurveTask."""
88 def setUp(self):
89 self.scale = 0.3
90 _, self.science, self.template, self.difference = make_test_images()
91 self.data = _augment_with_object_columns(DATA)
93 def _make_task(self, sources=None, forced=None, **config_kwargs):
94 config = plotDiaSourceLightcurve.PlotDiaSourceLightcurveConfig()
95 for key, value in config_kwargs.items():
96 setattr(config, key, value)
97 query = None
98 if sources is not None or forced is not None:
99 query = _StubApdbQuery(
100 sources if sources is not None else pd.DataFrame(),
101 forced if forced is not None else pd.DataFrame(),
102 )
103 task = plotDiaSourceLightcurve.PlotDiaSourceLightcurveTask(
104 config=config, output_path="", apdb_query=query)
105 return task, query
107 def _record(self, row):
108 """Return a single-row numpy.record matching DataFrame ``iloc``."""
109 return self.data.iloc[[row]].to_records(index=False)[0]
111 def test_generate_image_no_apdb(self):
112 """Without an APDB handle, the lightcurve panel renders a placeholder
113 but the figure still produces a valid PNG.
114 """
115 task, _ = self._make_task()
116 cutout = task.generate_image(self.science, self.template, self.difference,
117 skyCenter, self.scale,
118 source=self._record(0))
119 with PIL.Image.open(cutout) as im:
120 self.assertGreater(im.width, 0)
121 self.assertGreater(im.height, 0)
123 def test_generate_image_with_lightcurve(self):
124 """With matching sources and forced rows, both groups render."""
125 sources = _make_sources_frame(
126 visits=[1234, 5678, 9999],
127 bands=["r", "g", "r"],
128 mjds=[60100.5, 60110.5, 60120.5],
129 fluxes=[1234.5, 2345.6, 3456.7],
130 errs=[123.5, 234.5, 345.6],
131 )
132 # Forced has a visit (8888) that does NOT appear in sources — that
133 # one should be drawn with the forced-only marker. Visit 1234 IS in
134 # sources, so the forced entry for it should be suppressed.
135 forced = _make_sources_frame(
136 visits=[1234, 8888],
137 bands=["r", "g"],
138 mjds=[60100.5, 60115.5],
139 fluxes=[1200.0, 800.0],
140 errs=[120.0, 80.0],
141 )
142 task, query = self._make_task(sources=sources, forced=forced)
143 cutout = task.generate_image(self.science, self.template, self.difference,
144 skyCenter, self.scale,
145 source=self._record(0))
146 self.assertEqual(query.source_calls, 1)
147 self.assertEqual(query.forced_calls, 1)
148 with PIL.Image.open(cutout) as im:
149 self.assertGreater(im.width, 0)
150 self.assertGreater(im.height, 0)
152 def test_lightcurve_cache_reuses_query(self):
153 """Adjacent diaSources on the same diaObject should hit the cache."""
154 sources = _make_sources_frame(
155 visits=[1234, 5678],
156 bands=["r", "g"],
157 mjds=[60100.5, 60110.5],
158 fluxes=[1234.5, 2345.6],
159 errs=[123.5, 234.5],
160 )
161 forced = pd.DataFrame(columns=["visit", "band", "midpointMjdTai",
162 "psfFlux", "psfFluxErr"])
163 task, query = self._make_task(sources=sources, forced=forced)
164 for i in range(2):
165 task.generate_image(self.science, self.template, self.difference,
166 skyCenter, self.scale,
167 source=self._record(i))
168 # Both sources share the same diaObjectId; the second call must come
169 # from the cache.
170 self.assertEqual(query.source_calls, 1)
171 self.assertEqual(query.forced_calls, 1)
173 def test_forced_only_dedup_by_visit(self):
174 """Forced rows whose ``visit`` matches a diaSource are suppressed."""
175 sources = _make_sources_frame(
176 visits=[1234, 5678],
177 bands=["r", "g"],
178 mjds=[60100.5, 60110.5],
179 fluxes=[1234.5, 2345.6],
180 errs=[123.5, 234.5],
181 )
182 forced = _make_sources_frame(
183 visits=[1234, 5678, 7777, 8888],
184 bands=["r", "g", "r", "g"],
185 mjds=[60100.5, 60110.5, 60112.5, 60115.5],
186 fluxes=[1200.0, 2400.0, 500.0, 800.0],
187 errs=[120.0, 240.0, 50.0, 80.0],
188 )
189 task, _ = self._make_task(sources=sources, forced=forced)
190 forced_only = forced[~forced["visit"].isin(sources["visit"])]
191 self.assertEqual(set(forced_only["visit"]), {7777, 8888})
192 # Also exercise the rendering path to confirm it does not raise.
193 task.generate_image(self.science, self.template, self.difference,
194 skyCenter, self.scale, source=self._record(0))
196 def test_no_psfFluxErr(self):
197 """Missing or NaN psfFluxErr should fall back to no error bars."""
198 sources = _make_sources_frame(
199 visits=[1234, 5678],
200 bands=["r", "g"],
201 mjds=[60100.5, 60110.5],
202 fluxes=[1234.5, 2345.6],
203 errs=None, with_err=False,
204 )
205 forced = _make_sources_frame(
206 visits=[7777],
207 bands=["r"],
208 mjds=[60112.5],
209 fluxes=[500.0],
210 errs=[np.nan], with_err=True,
211 )
212 task, _ = self._make_task(sources=sources, forced=forced)
213 cutout = task.generate_image(self.science, self.template, self.difference,
214 skyCenter, self.scale,
215 source=self._record(0))
216 with PIL.Image.open(cutout) as im:
217 self.assertGreater(im.width, 0)
219 def test_forced_query_skips_exclude_flagged(self):
220 """DiaForcedSource lacks the diaSource flag columns; the forced
221 query must always be called with exclude_flagged=False, even when
222 the config asks to exclude flagged diaSources.
223 """
224 sources = _make_sources_frame(
225 visits=[1234], bands=["r"], mjds=[60100.5],
226 fluxes=[1234.5], errs=[123.5])
227 forced = pd.DataFrame(columns=["visit", "band", "midpointMjdTai",
228 "psfFlux", "psfFluxErr"])
229 task, query = self._make_task(sources=sources, forced=forced,
230 lightcurve_exclude_flagged=True)
231 task.generate_image(self.science, self.template, self.difference,
232 skyCenter, self.scale,
233 source=self._record(0))
234 self.assertTrue(query.last_source_kwargs["exclude_flagged"])
235 self.assertFalse(query.last_forced_kwargs["exclude_flagged"])
237 def test_nan_dia_object_id(self):
238 """A NaN diaObjectId (unassociated diaSource) must not crash; the
239 lightcurve panel renders its empty placeholder.
240 """
241 data = self.data.copy()
242 data["diaObjectId"] = data["diaObjectId"].astype(float)
243 data.loc[0, "diaObjectId"] = np.nan
244 sources = _make_sources_frame(
245 visits=[1234], bands=["r"], mjds=[60100.5],
246 fluxes=[1234.5], errs=[123.5])
247 forced = pd.DataFrame(columns=["visit", "band", "midpointMjdTai",
248 "psfFlux", "psfFluxErr"])
249 task, query = self._make_task(sources=sources, forced=forced)
250 record = data.iloc[[0]].to_records(index=False)[0]
251 cutout = task.generate_image(self.science, self.template, self.difference,
252 skyCenter, self.scale, source=record)
253 # APDB was never queried — diaObjectId is NaN.
254 self.assertEqual(query.source_calls, 0)
255 self.assertEqual(query.forced_calls, 0)
256 with PIL.Image.open(cutout) as im:
257 self.assertGreater(im.width, 0)
259 def test_forced_legend_single_black_entry(self):
260 """Forced points are colored per-band, but contribute a single
261 black legend entry regardless of band count.
262 """
263 sources = _make_sources_frame(
264 visits=[1234, 5678],
265 bands=["r", "g"],
266 mjds=[60100.5, 60110.5],
267 fluxes=[1234.5, 2345.6],
268 errs=[123.5, 234.5],
269 )
270 # Forced-only visits span two bands: should still produce one entry.
271 forced = _make_sources_frame(
272 visits=[7777, 8888],
273 bands=["r", "g"],
274 mjds=[60112.5, 60115.5],
275 fluxes=[500.0, 800.0],
276 errs=[50.0, 80.0],
277 )
278 task, _ = self._make_task(sources=sources, forced=forced)
279 import matplotlib.pyplot as plt
280 fig, ax = plt.subplots()
281 try:
282 task._draw_lightcurve(ax, sources, forced,
283 current_source=self._record(0))
284 handles, labels = ax.get_legend_handles_labels()
285 forced_labels = [lbl for lbl in labels if "forced" in lbl]
286 self.assertEqual(len(forced_labels), 1)
287 self.assertIn("(n=2)", forced_labels[0])
288 # The single forced legend handle should be drawn in black.
289 forced_handle = handles[labels.index(forced_labels[0])]
290 self.assertEqual(forced_handle.get_color(), "black")
291 finally:
292 plt.close(fig)
294 def test_njobs_downgraded(self):
295 """Requesting multiprocessing should be downgraded with a warning."""
296 sources = _make_sources_frame(
297 visits=[1234], bands=["r"], mjds=[60100.5],
298 fluxes=[1234.5], errs=[123.5])
299 forced = pd.DataFrame(columns=["visit", "band", "midpointMjdTai",
300 "psfFlux", "psfFluxErr"])
301 task, _ = self._make_task(sources=sources, forced=forced)
302 # write_images would normally attempt multiprocessing; the override
303 # should silently drop njobs to 0 and not raise.
304 with self.assertLogs(task.log.name, level="WARNING") as ctx:
305 # Use just one row so we don't need real files on disk; the
306 # cutouts task will try to look them up via butler_cache and
307 # fail, but the warning should fire first.
308 try:
309 task.write_images(self.data.head(0), butler=None, njobs=4)
310 except Exception:
311 pass
312 self.assertTrue(any("njobs" in msg for msg in ctx.output))
315class MemoryTester(lsst.utils.tests.MemoryTestCase):
316 pass
319def setup_module(module):
320 lsst.utils.tests.init()
323if __name__ == "__main__": 323 ↛ 324line 323 didn't jump to line 324 because the condition on line 323 was never true
324 lsst.utils.tests.init()
325 unittest.main()