Coverage for tests/utils_tests.py: 13%
108 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-04 02:27 -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
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):
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 + 1,
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):
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 + 1,
131 "midpointMjdTai": midpointMjdTai + 1.0 * idx,
132 "band": exposure.getFilter().bandLabel,
133 "psfNdata": 0,
134 "trailNdata": 0,
135 "dipoleNdata": 0}
136 data.append(diaSource | flags)
138 return pd.DataFrame(data=data).set_index(["diaObjectId", "band", "diaSourceId"], drop=False)
141def makeSolarSystemSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False):
142 """Make a test set of solar system sources.
144 Parameters
145 ----------
146 nSources : `int`
147 Number of sources to create.
148 diaObjectIds : `numpy.ndarray`
149 Integer Ids of diaobjects to "associate" with the DiaSources.
150 exposure : `lsst.afw.image.Exposure`
151 Exposure to create sources over.
152 rng : `numpy.random.Generator`
153 A NumPy random number generator initialized with a fixed seed for reproducibility.
154 randomizeObjects : `bool`, optional
155 If True, randomly draw from `diaObjectIds` to generate the ids in the
156 output catalog, otherwise just iterate through them, repeating as
157 necessary to get nSources objectIds.
159 Returns
160 -------
161 solarSystemSources : `pandas.DataFrame`
162 Solar system sources generated across the exposure.
163 """
164 solarSystemSources = makeDiaSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False)
165 solarSystemSources["ssObjectId"] = rng.integers(0, high=2**63-1, size=nSources)
166 solarSystemSources["Err(arcsec)"] = rng.uniform(0.2, 0.4, size=nSources)
168 return solarSystemSources
171def makeDiaForcedSources(nForcedSources, diaObjectIds, exposure, rng, randomizeObjects=False):
172 """Make a test set of DiaSources.
174 Parameters
175 ----------
176 nForcedSources : `int`
177 Number of sources to create.
178 diaObjectIds : `numpy.ndarray`
179 Integer Ids of diaobjects to "associate" with the DiaSources.
180 exposure : `lsst.afw.image.Exposure`
181 Exposure to create sources over.
182 rng : `numpy.random.Generator`
183 A NumPy random number generator initialized with a fixed seed for reproducibility.
184 randomizeObjects : `bool`, optional
185 If True, randomly draw from `diaObjectIds` to generate the ids in the
186 output catalog, otherwise just iterate through them.
188 Returns
189 -------
190 diaForcedSources : `pandas.DataFrame`
191 DiaForcedSources generated across the exposure.
192 """
193 midpointMjdTai = exposure.visitInfo.date.get(system=DateTime.MJD)
194 visit = exposure.visitInfo.id
195 detector = exposure.detector.getId()
196 if randomizeObjects:
197 objectIds = diaObjectIds[rng.randint(len(diaObjectIds), size=nForcedSources)]
198 else:
199 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nForcedSources)]]
201 data = []
202 bbox = exposure.getBBox()
204 for i, objId in enumerate(objectIds):
205 # Put together the minimum values for the alert.
206 x = rng.uniform(bbox.minX, bbox.maxX)
207 y = rng.uniform(bbox.minY, bbox.maxY)
208 coord = exposure.wcs.pixelToSky(x, y)
209 data.append({"diaForcedSourceId": i + 1,
210 "visit": visit + i,
211 "detector": detector,
212 "diaObjectId": objId,
213 "ra": coord.getRa().asDegrees(),
214 "dec": coord.getDec().asDegrees(),
215 "midpointMjdTai": midpointMjdTai + 1.0 * i,
216 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD,
217 scale=DateTime.TAI),
218 "band": exposure.getFilter().bandLabel})
220 return pd.DataFrame(data=data).set_index(["diaObjectId", "diaForcedSourceId"], drop=False)
223def makeExposure(flipX=False, flipY=False):
224 """Create an exposure and flip the x or y (or both) coordinates.
226 Returns bounding boxes that are right or left handed around the bounding
227 polygon.
229 Parameters
230 ----------
231 flipX : `bool`
232 Flip the x coordinate in the WCS.
233 flipY : `bool`
234 Flip the y coordinate in the WCS.
236 Returns
237 -------
238 exposure : `lsst.afw.image.Exposure`
239 Exposure with a valid bounding box and wcs.
240 """
241 metadata = PropertySet()
243 metadata.set("SIMPLE", "T")
244 metadata.set("BITPIX", -32)
245 metadata.set("NAXIS", 2)
246 metadata.set("NAXIS1", 1024)
247 metadata.set("NAXIS2", 1153)
248 metadata.set("RADECSYS", 'FK5')
249 metadata.set("EQUINOX", 2000.)
251 metadata.setDouble("CRVAL1", 215.604025685476)
252 metadata.setDouble("CRVAL2", 53.1595451514076)
253 metadata.setDouble("CRPIX1", 1109.99981456774)
254 metadata.setDouble("CRPIX2", 560.018167811613)
255 metadata.set("CTYPE1", 'RA---SIN')
256 metadata.set("CTYPE2", 'DEC--SIN')
258 xFlip = 1
259 if flipX:
260 xFlip = -1
261 yFlip = 1
262 if flipY:
263 yFlip = -1
264 metadata.setDouble("CD1_1", xFlip * 5.10808596133527E-05)
265 metadata.setDouble("CD1_2", yFlip * 1.85579539217196E-07)
266 metadata.setDouble("CD2_2", yFlip * -5.10281493481982E-05)
267 metadata.setDouble("CD2_1", xFlip * -8.27440751733828E-07)
269 wcs = afwGeom.makeSkyWcs(metadata)
270 exposure = afwImage.makeExposure(
271 afwImage.makeMaskedImageFromArrays(np.ones((1024, 1153))), wcs)
272 detector = DetectorWrapper(id=23, bbox=exposure.getBBox()).detector
273 visit = afwImage.VisitInfo(
274 exposureTime=200.,
275 date=DateTime("2014-05-13T17:00:00.000000000",
276 DateTime.Timescale.TAI))
277 exposure.info.id = 1234
278 exposure.setDetector(detector)
279 exposure.info.setVisitInfo(visit)
280 exposure.setFilter(afwImage.FilterLabel(band='g'))
281 exposure.setPhotoCalib(afwImage.PhotoCalib(1., 0., exposure.getBBox()))
283 return exposure
286def makeRegionTime(exposure=None, begin=None, end=None):
287 """Make a `RegionTimeInfo` for testing
289 Parameters
290 ----------
291 exposure : `lsst.afw.image.Exposure`, optional
292 Exposure to construct a ``RegionTimeInfo`` for.
293 If None, a default exposure will be created.
294 begin : `astropy.time.Time`, optional
295 The start time of the interval.
296 If `None`, calculate the start time from the midpoint of the exposure
297 and the exposure length.
298 end : `astropy.time.Time`, optional
299 The end time of the interval.
300 If `None`, calculate the end time from the midpoint of the exposure
301 and the exposure length.
303 Returns
304 -------
305 regionTime : `lsst.pipe.base.utils.RegionTimeInfo`
306 Object containing the spatial region and temporal timespan for an exposure.
307 """
308 if exposure is None:
309 exposure = makeExposure()
310 region = getRegion(exposure)
311 expTime = exposure.visitInfo.exposureTime*astropy.units.second
312 # visitInfo time is the midpoint of the exposure.
313 if begin is None:
314 begin = exposure.visitInfo.date.toAstropy() - expTime/2
315 if end is None:
316 end = exposure.visitInfo.date.toAstropy() + expTime/2
317 timespan = dafButler.Timespan(begin=begin, end=end)
318 return RegionTimeInfo(region=region, timespan=timespan)