Coverage for tests/utils_tests.py: 93%
115 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-10 01:47 -0800
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
30from lsst.afw.coord import Observatory
31import lsst.afw.geom as afwGeom
32import lsst.afw.image as afwImage
33from lsst.daf.base import DateTime, PropertySet
34import lsst.daf.butler as dafButler
35import lsst.geom
36import lsst.meas.algorithms as measAlg
37from lsst.pipe.base.utils import RegionTimeInfo
38import lsst.sphgeom
40from lsst.ap.association.utils import getRegion
43def makeDiaObjects(nObjects, exposure, rng, startId=1):
44 """Make a test set of DiaObjects.
46 Parameters
47 ----------
48 nObjects : `int`
49 Number of objects to create.
50 exposure : `lsst.afw.image.Exposure`
51 Exposure to create objects over.
52 rng : `numpy.random.Generator`
53 A NumPy random number generator initialized with a fixed seed for reproducibility.
55 Returns
56 -------
57 diaObjects : `pandas.DataFrame`
58 DiaObjects generated across the exposure.
59 """
60 bbox = lsst.geom.Box2D(exposure.getBBox())
61 rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nObjects)
62 rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nObjects)
64 data = []
65 for idx, (x, y) in enumerate(zip(rand_x, rand_y)):
66 coord = exposure.wcs.pixelToSky(x, y)
67 newObject = {"ra": coord.getRa().asDegrees(),
68 "dec": coord.getDec().asDegrees(),
69 "diaObjectId": idx + startId,
70 "nDiaSources": 5}
71 for f in ["u", "g", "r", "i", "z", "y"]:
72 newObject["%s_psfFluxNdata" % f] = 0
73 data.append(newObject)
75 return pd.DataFrame(data=data).set_index("diaObjectId", drop=False)
78def makeDiaSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False, flagList=None, startId=1):
79 """Make a test set of DiaSources.
81 Parameters
82 ----------
83 nSources : `int`
84 Number of sources to create.
85 diaObjectIds : `numpy.ndarray`
86 Integer Ids of diaobjects to "associate" with the DiaSources.
87 exposure : `lsst.afw.image.Exposure`
88 Exposure to create sources over.
89 rng : `numpy.random.Generator`
90 A NumPy random number generator initialized with a fixed seed for reproducibility.
91 randomizeObjects : `bool`, optional
92 If True, randomly draw from `diaObjectIds` to generate the ids in the
93 output catalog, otherwise just iterate through them, repeating as
94 necessary to get nSources objectIds.
95 flagList : `list` of `str`, optional
96 Optional list of names of flag columns to add to the DiaSource table.
98 Returns
99 -------
100 diaSources : `pandas.DataFrame`
101 DiaSources generated across the exposure.
102 """
103 bbox = lsst.geom.Box2D(exposure.getBBox())
104 rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nSources)
105 rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nSources)
106 diaObjectIds = diaObjectIds.astype(np.int64)
107 if randomizeObjects:
108 objectIds = diaObjectIds[rng.integers(len(diaObjectIds), size=nSources)]
109 else:
110 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nSources)]]
112 midpointMjdTai = exposure.visitInfo.date.get(system=DateTime.MJD)
114 data = []
115 flags = {}
116 if flagList is not None:
117 for flag in flagList:
118 flags[flag] = False
119 for idx, (x, y, objId) in enumerate(zip(rand_x, rand_y, objectIds)):
120 coord = exposure.wcs.pixelToSky(x, y)
121 # Put together the minimum values for the alert.
122 diaSource = {"ra": coord.getRa().asDegrees(),
123 "dec": coord.getDec().asDegrees(),
124 "x": x,
125 "y": y,
126 "visit": exposure.visitInfo.id,
127 "detector": exposure.detector.getId(),
128 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD,
129 scale=DateTime.TAI),
130 "diaObjectId": objId,
131 "ssObjectId": 0,
132 "parentDiaSourceId": 0,
133 "diaSourceId": idx + startId,
134 "midpointMjdTai": midpointMjdTai + 1.0 * idx,
135 "band": exposure.getFilter().bandLabel,
136 "psfNdata": 0,
137 "bboxSize": 23,
138 "trailNdata": 0,
139 "dipoleNdata": 0,
140 "templateInstrument": "TESTCam",
141 "templateVersion": 0,
142 "mixedInstruments": False}
143 data.append(diaSource | flags)
145 return pd.DataFrame(data=data).set_index(["diaObjectId", "band", "diaSourceId"], drop=False)
148def makeSolarSystemSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False, startId=1):
149 """Make a test set of solar system sources.
151 Parameters
152 ----------
153 nSources : `int`
154 Number of sources to create.
155 diaObjectIds : `numpy.ndarray`
156 Integer Ids of diaobjects to "associate" with the DiaSources.
157 exposure : `lsst.afw.image.Exposure`
158 Exposure to create sources over.
159 rng : `numpy.random.Generator`
160 A NumPy random number generator initialized with a fixed seed for reproducibility.
161 randomizeObjects : `bool`, optional
162 If True, randomly draw from `diaObjectIds` to generate the ids in the
163 output catalog, otherwise just iterate through them, repeating as
164 necessary to get nSources objectIds.
166 Returns
167 -------
168 solarSystemSources : `pandas.DataFrame`
169 Solar system sources generated across the exposure.
170 """
171 solarSystemSources = makeDiaSources(nSources, diaObjectIds, exposure, rng,
172 randomizeObjects=False,
173 startId=startId)
174 solarSystemSources["ssObjectId"] = rng.integers(0, high=2**63-1, size=nSources)
175 solarSystemSources["Err(arcsec)"] = rng.uniform(0.2, 0.4, size=nSources)
177 return solarSystemSources
180def makeDiaForcedSources(nForcedSources, diaObjectIds, exposure, rng, randomizeObjects=False, startId=1):
181 """Make a test set of DiaSources.
183 Parameters
184 ----------
185 nForcedSources : `int`
186 Number of sources to create.
187 diaObjectIds : `numpy.ndarray`
188 Integer Ids of diaobjects to "associate" with the DiaSources.
189 exposure : `lsst.afw.image.Exposure`
190 Exposure to create sources over.
191 rng : `numpy.random.Generator`
192 A NumPy random number generator initialized with a fixed seed for reproducibility.
193 randomizeObjects : `bool`, optional
194 If True, randomly draw from `diaObjectIds` to generate the ids in the
195 output catalog, otherwise just iterate through them.
197 Returns
198 -------
199 diaForcedSources : `pandas.DataFrame`
200 DiaForcedSources generated across the exposure.
201 """
202 midpointMjdTai = exposure.visitInfo.date.get(system=DateTime.MJD)
203 visit = exposure.visitInfo.id
204 detector = exposure.detector.getId()
205 if randomizeObjects: 205 ↛ 206line 205 didn't jump to line 206 because the condition on line 205 was never true
206 objectIds = diaObjectIds[rng.randint(len(diaObjectIds), size=nForcedSources)]
207 else:
208 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nForcedSources)]]
210 data = []
211 bbox = exposure.getBBox()
213 for i, objId in enumerate(objectIds):
214 # Put together the minimum values for the alert.
215 x = rng.uniform(bbox.minX, bbox.maxX)
216 y = rng.uniform(bbox.minY, bbox.maxY)
217 coord = exposure.wcs.pixelToSky(x, y)
218 data.append({"diaForcedSourceId": i + startId,
219 "visit": visit + i,
220 "detector": detector,
221 "diaObjectId": objId,
222 "ra": coord.getRa().asDegrees(),
223 "dec": coord.getDec().asDegrees(),
224 "midpointMjdTai": midpointMjdTai + 1.0 * i,
225 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD,
226 scale=DateTime.TAI),
227 "band": exposure.getFilter().bandLabel})
229 return pd.DataFrame(data=data).set_index(["diaObjectId", "diaForcedSourceId"], drop=False)
232def makeExposure(flipX=False, flipY=False):
233 """Create an exposure and flip the x or y (or both) coordinates.
235 Returns bounding boxes that are right or left handed around the bounding
236 polygon.
238 Parameters
239 ----------
240 flipX : `bool`
241 Flip the x coordinate in the WCS.
242 flipY : `bool`
243 Flip the y coordinate in the WCS.
245 Returns
246 -------
247 exposure : `lsst.afw.image.Exposure`
248 Exposure with a valid bounding box and wcs.
249 """
250 metadata = PropertySet()
252 metadata.set("SIMPLE", "T")
253 metadata.set("BITPIX", -32)
254 metadata.set("NAXIS", 2)
255 metadata.set("NAXIS1", 1024)
256 metadata.set("NAXIS2", 1153)
257 metadata.set("RADECSYS", 'FK5')
258 metadata.set("EQUINOX", 2000.)
259 ra = 215.604025685476
260 dec = 53.1595451514076
262 metadata.setDouble("CRVAL1", ra)
263 metadata.setDouble("CRVAL2", dec)
264 metadata.setDouble("CRPIX1", 1109.99981456774)
265 metadata.setDouble("CRPIX2", 560.018167811613)
266 metadata.set("CTYPE1", 'RA---SIN')
267 metadata.set("CTYPE2", 'DEC--SIN')
269 xFlip = 1
270 if flipX: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 xFlip = -1
272 yFlip = 1
273 if flipY: 273 ↛ 274line 273 didn't jump to line 274 because the condition on line 273 was never true
274 yFlip = -1
275 metadata.setDouble("CD1_1", xFlip * 5.10808596133527E-05)
276 metadata.setDouble("CD1_2", yFlip * 1.85579539217196E-07)
277 metadata.setDouble("CD2_2", yFlip * -5.10281493481982E-05)
278 metadata.setDouble("CD2_1", xFlip * -8.27440751733828E-07)
280 wcs = afwGeom.makeSkyWcs(metadata)
281 exposure = afwImage.makeExposure(
282 afwImage.makeMaskedImageFromArrays(np.ones((1024, 1153))), wcs)
283 detector = DetectorWrapper(id=23, bbox=exposure.getBBox()).detector
284 visit = afwImage.VisitInfo(
285 exposureTime=200.,
286 date=DateTime("2014-05-13T17:00:00.000000000",
287 DateTime.Timescale.TAI),
288 boresightRaDec=lsst.geom.SpherePoint(ra, dec, lsst.geom.degrees),
289 boresightRotAngle=73.2*lsst.geom.degrees,
290 rotType=afwImage.RotType.SKY,
291 observatory=Observatory(11.1*lsst.geom.degrees, 22.2*lsst.geom.degrees, 0.333),
292 )
293 kernel = measAlg.DoubleGaussianPsf(7, 7, 2.0).getKernel()
294 exposure.setPsf(measAlg.KernelPsf(kernel, exposure.getBBox().getCenter()))
295 exposure.info.id = 1234
296 exposure.setDetector(detector)
297 exposure.info.setVisitInfo(visit)
298 exposure.setFilter(afwImage.FilterLabel(band='g'))
299 exposure.setPhotoCalib(afwImage.PhotoCalib(1., 0., exposure.getBBox()))
301 return exposure
304def makeRegionTime(exposure=None, begin=None, end=None):
305 """Make a `RegionTimeInfo` for testing
307 Parameters
308 ----------
309 exposure : `lsst.afw.image.Exposure`, optional
310 Exposure to construct a ``RegionTimeInfo`` for.
311 If None, a default exposure will be created.
312 begin : `astropy.time.Time`, optional
313 The start time of the interval.
314 If `None`, calculate the start time from the midpoint of the exposure
315 and the exposure length.
316 end : `astropy.time.Time`, optional
317 The end time of the interval.
318 If `None`, calculate the end time from the midpoint of the exposure
319 and the exposure length.
321 Returns
322 -------
323 regionTime : `lsst.pipe.base.utils.RegionTimeInfo`
324 Object containing the spatial region and temporal timespan for an exposure.
325 """
326 if exposure is None: 326 ↛ 327line 326 didn't jump to line 327 because the condition on line 326 was never true
327 exposure = makeExposure()
328 region = getRegion(exposure)
329 expTime = exposure.visitInfo.exposureTime*astropy.units.second
330 # visitInfo time is the midpoint of the exposure.
331 if begin is None: 331 ↛ 333line 331 didn't jump to line 333 because the condition on line 331 was always true
332 begin = exposure.visitInfo.date.toAstropy() - expTime/2
333 if end is None: 333 ↛ 335line 333 didn't jump to line 335 because the condition on line 333 was always true
334 end = exposure.visitInfo.date.toAstropy() + expTime/2
335 timespan = dafButler.Timespan(begin=begin, end=end)
336 return RegionTimeInfo(region=region, timespan=timespan)