Coverage for tests / utils_tests.py: 13%
108 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-17 00:50 -0700
1# This file is part of ap_association.
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/>.
22"""Helper functions for tests of DIA catalogs, including generating mock
23catalogs for simulated APDB access.
24"""
25import astropy.units
26import pandas as pd
27import numpy as np
29from lsst.afw.cameraGeom.testUtils import DetectorWrapper
30import lsst.afw.geom as afwGeom
31import lsst.afw.image as afwImage
32from lsst.daf.base import DateTime, PropertySet
33import lsst.daf.butler as dafButler
34import lsst.geom
35from lsst.pipe.base.utils import RegionTimeInfo
36import lsst.sphgeom
38from lsst.ap.association.utils import getRegion
41def makeDiaObjects(nObjects, exposure, rng, startId=1):
42 """Make a test set of DiaObjects.
44 Parameters
45 ----------
46 nObjects : `int`
47 Number of objects to create.
48 exposure : `lsst.afw.image.Exposure`
49 Exposure to create objects over.
50 rng : `numpy.random.Generator`
51 A NumPy random number generator initialized with a fixed seed for reproducibility.
53 Returns
54 -------
55 diaObjects : `pandas.DataFrame`
56 DiaObjects generated across the exposure.
57 """
58 bbox = lsst.geom.Box2D(exposure.getBBox())
59 rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nObjects)
60 rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nObjects)
62 data = []
63 for idx, (x, y) in enumerate(zip(rand_x, rand_y)):
64 coord = exposure.wcs.pixelToSky(x, y)
65 newObject = {"ra": coord.getRa().asDegrees(),
66 "dec": coord.getDec().asDegrees(),
67 "diaObjectId": idx + startId,
68 "nDiaSources": 5}
69 for f in ["u", "g", "r", "i", "z", "y"]:
70 newObject["%s_psfFluxNdata" % f] = 0
71 data.append(newObject)
73 return pd.DataFrame(data=data).set_index("diaObjectId", drop=False)
76def makeDiaSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False, flagList=None, startId=1):
77 """Make a test set of DiaSources.
79 Parameters
80 ----------
81 nSources : `int`
82 Number of sources to create.
83 diaObjectIds : `numpy.ndarray`
84 Integer Ids of diaobjects to "associate" with the DiaSources.
85 exposure : `lsst.afw.image.Exposure`
86 Exposure to create sources over.
87 rng : `numpy.random.Generator`
88 A NumPy random number generator initialized with a fixed seed for reproducibility.
89 randomizeObjects : `bool`, optional
90 If True, randomly draw from `diaObjectIds` to generate the ids in the
91 output catalog, otherwise just iterate through them, repeating as
92 necessary to get nSources objectIds.
93 flagList : `list` of `str`, optional
94 Optional list of names of flag columns to add to the DiaSource table.
96 Returns
97 -------
98 diaSources : `pandas.DataFrame`
99 DiaSources generated across the exposure.
100 """
101 bbox = lsst.geom.Box2D(exposure.getBBox())
102 rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nSources)
103 rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nSources)
104 if randomizeObjects:
105 objectIds = diaObjectIds[rng.integers(len(diaObjectIds), size=nSources)]
106 else:
107 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nSources)]]
109 midpointMjdTai = exposure.visitInfo.date.get(system=DateTime.MJD)
111 data = []
112 flags = {}
113 if flagList is not None:
114 for flag in flagList:
115 flags[flag] = False
116 for idx, (x, y, objId) in enumerate(zip(rand_x, rand_y, objectIds)):
117 coord = exposure.wcs.pixelToSky(x, y)
118 # Put together the minimum values for the alert.
119 diaSource = {"ra": coord.getRa().asDegrees(),
120 "dec": coord.getDec().asDegrees(),
121 "x": x,
122 "y": y,
123 "visit": exposure.visitInfo.id,
124 "detector": exposure.detector.getId(),
125 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD,
126 scale=DateTime.TAI),
127 "diaObjectId": objId,
128 "ssObjectId": 0,
129 "parentDiaSourceId": 0,
130 "diaSourceId": idx + startId,
131 "midpointMjdTai": midpointMjdTai + 1.0 * idx,
132 "band": exposure.getFilter().bandLabel,
133 "psfNdata": 0,
134 "bboxSize": 23,
135 "trailNdata": 0,
136 "dipoleNdata": 0}
137 data.append(diaSource | flags)
139 return pd.DataFrame(data=data).set_index(["diaObjectId", "band", "diaSourceId"], drop=False)
142def makeSolarSystemSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False, startId=1):
143 """Make a test set of solar system sources.
145 Parameters
146 ----------
147 nSources : `int`
148 Number of sources to create.
149 diaObjectIds : `numpy.ndarray`
150 Integer Ids of diaobjects to "associate" with the DiaSources.
151 exposure : `lsst.afw.image.Exposure`
152 Exposure to create sources over.
153 rng : `numpy.random.Generator`
154 A NumPy random number generator initialized with a fixed seed for reproducibility.
155 randomizeObjects : `bool`, optional
156 If True, randomly draw from `diaObjectIds` to generate the ids in the
157 output catalog, otherwise just iterate through them, repeating as
158 necessary to get nSources objectIds.
160 Returns
161 -------
162 solarSystemSources : `pandas.DataFrame`
163 Solar system sources generated across the exposure.
164 """
165 solarSystemSources = makeDiaSources(nSources, diaObjectIds, exposure, rng,
166 randomizeObjects=False,
167 startId=startId)
168 solarSystemSources["ssObjectId"] = rng.integers(0, high=2**63-1, size=nSources)
169 solarSystemSources["Err(arcsec)"] = rng.uniform(0.2, 0.4, size=nSources)
171 return solarSystemSources
174def makeDiaForcedSources(nForcedSources, diaObjectIds, exposure, rng, randomizeObjects=False, startId=1):
175 """Make a test set of DiaSources.
177 Parameters
178 ----------
179 nForcedSources : `int`
180 Number of sources to create.
181 diaObjectIds : `numpy.ndarray`
182 Integer Ids of diaobjects to "associate" with the DiaSources.
183 exposure : `lsst.afw.image.Exposure`
184 Exposure to create sources over.
185 rng : `numpy.random.Generator`
186 A NumPy random number generator initialized with a fixed seed for reproducibility.
187 randomizeObjects : `bool`, optional
188 If True, randomly draw from `diaObjectIds` to generate the ids in the
189 output catalog, otherwise just iterate through them.
191 Returns
192 -------
193 diaForcedSources : `pandas.DataFrame`
194 DiaForcedSources generated across the exposure.
195 """
196 midpointMjdTai = exposure.visitInfo.date.get(system=DateTime.MJD)
197 visit = exposure.visitInfo.id
198 detector = exposure.detector.getId()
199 if randomizeObjects:
200 objectIds = diaObjectIds[rng.randint(len(diaObjectIds), size=nForcedSources)]
201 else:
202 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nForcedSources)]]
204 data = []
205 bbox = exposure.getBBox()
207 for i, objId in enumerate(objectIds):
208 # Put together the minimum values for the alert.
209 x = rng.uniform(bbox.minX, bbox.maxX)
210 y = rng.uniform(bbox.minY, bbox.maxY)
211 coord = exposure.wcs.pixelToSky(x, y)
212 data.append({"diaForcedSourceId": i + startId,
213 "visit": visit + i,
214 "detector": detector,
215 "diaObjectId": objId,
216 "ra": coord.getRa().asDegrees(),
217 "dec": coord.getDec().asDegrees(),
218 "midpointMjdTai": midpointMjdTai + 1.0 * i,
219 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD,
220 scale=DateTime.TAI),
221 "band": exposure.getFilter().bandLabel})
223 return pd.DataFrame(data=data).set_index(["diaObjectId", "diaForcedSourceId"], drop=False)
226def makeExposure(flipX=False, flipY=False):
227 """Create an exposure and flip the x or y (or both) coordinates.
229 Returns bounding boxes that are right or left handed around the bounding
230 polygon.
232 Parameters
233 ----------
234 flipX : `bool`
235 Flip the x coordinate in the WCS.
236 flipY : `bool`
237 Flip the y coordinate in the WCS.
239 Returns
240 -------
241 exposure : `lsst.afw.image.Exposure`
242 Exposure with a valid bounding box and wcs.
243 """
244 metadata = PropertySet()
246 metadata.set("SIMPLE", "T")
247 metadata.set("BITPIX", -32)
248 metadata.set("NAXIS", 2)
249 metadata.set("NAXIS1", 1024)
250 metadata.set("NAXIS2", 1153)
251 metadata.set("RADECSYS", 'FK5')
252 metadata.set("EQUINOX", 2000.)
254 metadata.setDouble("CRVAL1", 215.604025685476)
255 metadata.setDouble("CRVAL2", 53.1595451514076)
256 metadata.setDouble("CRPIX1", 1109.99981456774)
257 metadata.setDouble("CRPIX2", 560.018167811613)
258 metadata.set("CTYPE1", 'RA---SIN')
259 metadata.set("CTYPE2", 'DEC--SIN')
261 xFlip = 1
262 if flipX:
263 xFlip = -1
264 yFlip = 1
265 if flipY:
266 yFlip = -1
267 metadata.setDouble("CD1_1", xFlip * 5.10808596133527E-05)
268 metadata.setDouble("CD1_2", yFlip * 1.85579539217196E-07)
269 metadata.setDouble("CD2_2", yFlip * -5.10281493481982E-05)
270 metadata.setDouble("CD2_1", xFlip * -8.27440751733828E-07)
272 wcs = afwGeom.makeSkyWcs(metadata)
273 exposure = afwImage.makeExposure(
274 afwImage.makeMaskedImageFromArrays(np.ones((1024, 1153))), wcs)
275 detector = DetectorWrapper(id=23, bbox=exposure.getBBox()).detector
276 visit = afwImage.VisitInfo(
277 exposureTime=200.,
278 date=DateTime("2014-05-13T17:00:00.000000000",
279 DateTime.Timescale.TAI))
280 exposure.info.id = 1234
281 exposure.setDetector(detector)
282 exposure.info.setVisitInfo(visit)
283 exposure.setFilter(afwImage.FilterLabel(band='g'))
284 exposure.setPhotoCalib(afwImage.PhotoCalib(1., 0., exposure.getBBox()))
286 return exposure
289def makeRegionTime(exposure=None, begin=None, end=None):
290 """Make a `RegionTimeInfo` for testing
292 Parameters
293 ----------
294 exposure : `lsst.afw.image.Exposure`, optional
295 Exposure to construct a ``RegionTimeInfo`` for.
296 If None, a default exposure will be created.
297 begin : `astropy.time.Time`, optional
298 The start time of the interval.
299 If `None`, calculate the start time from the midpoint of the exposure
300 and the exposure length.
301 end : `astropy.time.Time`, optional
302 The end time of the interval.
303 If `None`, calculate the end time from the midpoint of the exposure
304 and the exposure length.
306 Returns
307 -------
308 regionTime : `lsst.pipe.base.utils.RegionTimeInfo`
309 Object containing the spatial region and temporal timespan for an exposure.
310 """
311 if exposure is None:
312 exposure = makeExposure()
313 region = getRegion(exposure)
314 expTime = exposure.visitInfo.exposureTime*astropy.units.second
315 # visitInfo time is the midpoint of the exposure.
316 if begin is None:
317 begin = exposure.visitInfo.date.toAstropy() - expTime/2
318 if end is None:
319 end = exposure.visitInfo.date.toAstropy() + expTime/2
320 timespan = dafButler.Timespan(begin=begin, end=end)
321 return RegionTimeInfo(region=region, timespan=timespan)