Coverage for tests/utils_tests.py: 93%
115 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-29 22:47 +0000
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 data.append(diaSource | flags)
142 return pd.DataFrame(data=data).set_index(["diaObjectId", "band", "diaSourceId"], drop=False)
145def makeSolarSystemSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False, startId=1):
146 """Make a test set of solar system sources.
148 Parameters
149 ----------
150 nSources : `int`
151 Number of sources to create.
152 diaObjectIds : `numpy.ndarray`
153 Integer Ids of diaobjects to "associate" with the DiaSources.
154 exposure : `lsst.afw.image.Exposure`
155 Exposure to create sources over.
156 rng : `numpy.random.Generator`
157 A NumPy random number generator initialized with a fixed seed for reproducibility.
158 randomizeObjects : `bool`, optional
159 If True, randomly draw from `diaObjectIds` to generate the ids in the
160 output catalog, otherwise just iterate through them, repeating as
161 necessary to get nSources objectIds.
163 Returns
164 -------
165 solarSystemSources : `pandas.DataFrame`
166 Solar system sources generated across the exposure.
167 """
168 solarSystemSources = makeDiaSources(nSources, diaObjectIds, exposure, rng,
169 randomizeObjects=False,
170 startId=startId)
171 solarSystemSources["ssObjectId"] = rng.integers(0, high=2**63-1, size=nSources)
172 solarSystemSources["Err(arcsec)"] = rng.uniform(0.2, 0.4, size=nSources)
174 return solarSystemSources
177def makeDiaForcedSources(nForcedSources, diaObjectIds, exposure, rng, randomizeObjects=False, startId=1):
178 """Make a test set of DiaSources.
180 Parameters
181 ----------
182 nForcedSources : `int`
183 Number of sources to create.
184 diaObjectIds : `numpy.ndarray`
185 Integer Ids of diaobjects to "associate" with the DiaSources.
186 exposure : `lsst.afw.image.Exposure`
187 Exposure to create sources over.
188 rng : `numpy.random.Generator`
189 A NumPy random number generator initialized with a fixed seed for reproducibility.
190 randomizeObjects : `bool`, optional
191 If True, randomly draw from `diaObjectIds` to generate the ids in the
192 output catalog, otherwise just iterate through them.
194 Returns
195 -------
196 diaForcedSources : `pandas.DataFrame`
197 DiaForcedSources generated across the exposure.
198 """
199 midpointMjdTai = exposure.visitInfo.date.get(system=DateTime.MJD)
200 visit = exposure.visitInfo.id
201 detector = exposure.detector.getId()
202 if randomizeObjects: 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true
203 objectIds = diaObjectIds[rng.randint(len(diaObjectIds), size=nForcedSources)]
204 else:
205 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nForcedSources)]]
207 data = []
208 bbox = exposure.getBBox()
210 for i, objId in enumerate(objectIds):
211 # Put together the minimum values for the alert.
212 x = rng.uniform(bbox.minX, bbox.maxX)
213 y = rng.uniform(bbox.minY, bbox.maxY)
214 coord = exposure.wcs.pixelToSky(x, y)
215 data.append({"diaForcedSourceId": i + startId,
216 "visit": visit + i,
217 "detector": detector,
218 "diaObjectId": objId,
219 "ra": coord.getRa().asDegrees(),
220 "dec": coord.getDec().asDegrees(),
221 "midpointMjdTai": midpointMjdTai + 1.0 * i,
222 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD,
223 scale=DateTime.TAI),
224 "band": exposure.getFilter().bandLabel})
226 return pd.DataFrame(data=data).set_index(["diaObjectId", "diaForcedSourceId"], drop=False)
229def makeExposure(flipX=False, flipY=False):
230 """Create an exposure and flip the x or y (or both) coordinates.
232 Returns bounding boxes that are right or left handed around the bounding
233 polygon.
235 Parameters
236 ----------
237 flipX : `bool`
238 Flip the x coordinate in the WCS.
239 flipY : `bool`
240 Flip the y coordinate in the WCS.
242 Returns
243 -------
244 exposure : `lsst.afw.image.Exposure`
245 Exposure with a valid bounding box and wcs.
246 """
247 metadata = PropertySet()
249 metadata.set("SIMPLE", "T")
250 metadata.set("BITPIX", -32)
251 metadata.set("NAXIS", 2)
252 metadata.set("NAXIS1", 1024)
253 metadata.set("NAXIS2", 1153)
254 metadata.set("RADECSYS", 'FK5')
255 metadata.set("EQUINOX", 2000.)
256 ra = 215.604025685476
257 dec = 53.1595451514076
259 metadata.setDouble("CRVAL1", ra)
260 metadata.setDouble("CRVAL2", dec)
261 metadata.setDouble("CRPIX1", 1109.99981456774)
262 metadata.setDouble("CRPIX2", 560.018167811613)
263 metadata.set("CTYPE1", 'RA---SIN')
264 metadata.set("CTYPE2", 'DEC--SIN')
266 xFlip = 1
267 if flipX: 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true
268 xFlip = -1
269 yFlip = 1
270 if flipY: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 yFlip = -1
272 metadata.setDouble("CD1_1", xFlip * 5.10808596133527E-05)
273 metadata.setDouble("CD1_2", yFlip * 1.85579539217196E-07)
274 metadata.setDouble("CD2_2", yFlip * -5.10281493481982E-05)
275 metadata.setDouble("CD2_1", xFlip * -8.27440751733828E-07)
277 wcs = afwGeom.makeSkyWcs(metadata)
278 exposure = afwImage.makeExposure(
279 afwImage.makeMaskedImageFromArrays(np.ones((1024, 1153))), wcs)
280 detector = DetectorWrapper(id=23, bbox=exposure.getBBox()).detector
281 visit = afwImage.VisitInfo(
282 exposureTime=200.,
283 date=DateTime("2014-05-13T17:00:00.000000000",
284 DateTime.Timescale.TAI),
285 boresightRaDec=lsst.geom.SpherePoint(ra, dec, lsst.geom.degrees),
286 boresightRotAngle=73.2*lsst.geom.degrees,
287 rotType=afwImage.RotType.SKY,
288 observatory=Observatory(11.1*lsst.geom.degrees, 22.2*lsst.geom.degrees, 0.333),
289 )
290 kernel = measAlg.DoubleGaussianPsf(7, 7, 2.0).getKernel()
291 exposure.setPsf(measAlg.KernelPsf(kernel, exposure.getBBox().getCenter()))
292 exposure.info.id = 1234
293 exposure.setDetector(detector)
294 exposure.info.setVisitInfo(visit)
295 exposure.setFilter(afwImage.FilterLabel(band='g'))
296 exposure.setPhotoCalib(afwImage.PhotoCalib(1., 0., exposure.getBBox()))
298 return exposure
301def makeRegionTime(exposure=None, begin=None, end=None):
302 """Make a `RegionTimeInfo` for testing
304 Parameters
305 ----------
306 exposure : `lsst.afw.image.Exposure`, optional
307 Exposure to construct a ``RegionTimeInfo`` for.
308 If None, a default exposure will be created.
309 begin : `astropy.time.Time`, optional
310 The start time of the interval.
311 If `None`, calculate the start time from the midpoint of the exposure
312 and the exposure length.
313 end : `astropy.time.Time`, optional
314 The end time of the interval.
315 If `None`, calculate the end time from the midpoint of the exposure
316 and the exposure length.
318 Returns
319 -------
320 regionTime : `lsst.pipe.base.utils.RegionTimeInfo`
321 Object containing the spatial region and temporal timespan for an exposure.
322 """
323 if exposure is None: 323 ↛ 324line 323 didn't jump to line 324 because the condition on line 323 was never true
324 exposure = makeExposure()
325 region = getRegion(exposure)
326 expTime = exposure.visitInfo.exposureTime*astropy.units.second
327 # visitInfo time is the midpoint of the exposure.
328 if begin is None: 328 ↛ 330line 328 didn't jump to line 330 because the condition on line 328 was always true
329 begin = exposure.visitInfo.date.toAstropy() - expTime/2
330 if end is None: 330 ↛ 332line 330 didn't jump to line 332 because the condition on line 330 was always true
331 end = exposure.visitInfo.date.toAstropy() + expTime/2
332 timespan = dafButler.Timespan(begin=begin, end=end)
333 return RegionTimeInfo(region=region, timespan=timespan)