Coverage for tests/utils_tests.py: 13%
112 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-14 02:23 -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
39def makeDiaObjects(nObjects, exposure, rng):
40 """Make a test set of DiaObjects.
42 Parameters
43 ----------
44 nObjects : `int`
45 Number of objects to create.
46 exposure : `lsst.afw.image.Exposure`
47 Exposure to create objects over.
48 rng : `numpy.random.Generator`
49 A NumPy random number generator initialized with a fixed seed for reproducibility.
51 Returns
52 -------
53 diaObjects : `pandas.DataFrame`
54 DiaObjects generated across the exposure.
55 """
56 bbox = lsst.geom.Box2D(exposure.getBBox())
57 rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nObjects)
58 rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nObjects)
60 data = []
61 for idx, (x, y) in enumerate(zip(rand_x, rand_y)):
62 coord = exposure.wcs.pixelToSky(x, y)
63 newObject = {"ra": coord.getRa().asDegrees(),
64 "dec": coord.getDec().asDegrees(),
65 "diaObjectId": idx + 1,
66 "nDiaSources": 5}
67 for f in ["u", "g", "r", "i", "z", "y"]:
68 newObject["%s_psfFluxNdata" % f] = 0
69 data.append(newObject)
71 return pd.DataFrame(data=data).set_index("diaObjectId", drop=False)
74def makeDiaSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False, flagList=None):
75 """Make a test set of DiaSources.
77 Parameters
78 ----------
79 nSources : `int`
80 Number of sources to create.
81 diaObjectIds : `numpy.ndarray`
82 Integer Ids of diaobjects to "associate" with the DiaSources.
83 exposure : `lsst.afw.image.Exposure`
84 Exposure to create sources over.
85 rng : `numpy.random.Generator`
86 A NumPy random number generator initialized with a fixed seed for reproducibility.
87 randomizeObjects : `bool`, optional
88 If True, randomly draw from `diaObjectIds` to generate the ids in the
89 output catalog, otherwise just iterate through them, repeating as
90 necessary to get nSources objectIds.
91 flagList : `list` of `str`, optional
92 Optional list of names of flag columns to add to the DiaSource table.
94 Returns
95 -------
96 diaSources : `pandas.DataFrame`
97 DiaSources generated across the exposure.
98 """
99 bbox = lsst.geom.Box2D(exposure.getBBox())
100 rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nSources)
101 rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nSources)
102 if randomizeObjects:
103 objectIds = diaObjectIds[rng.integers(len(diaObjectIds), size=nSources)]
104 else:
105 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nSources)]]
107 midpointMjdTai = exposure.visitInfo.date.get(system=DateTime.MJD)
109 data = []
110 flags = {}
111 if flagList is not None:
112 for flag in flagList:
113 flags[flag] = False
114 for idx, (x, y, objId) in enumerate(zip(rand_x, rand_y, objectIds)):
115 coord = exposure.wcs.pixelToSky(x, y)
116 # Put together the minimum values for the alert.
117 diaSource = {"ra": coord.getRa().asDegrees(),
118 "dec": coord.getDec().asDegrees(),
119 "x": x,
120 "y": y,
121 "visit": exposure.visitInfo.id,
122 "detector": exposure.detector.getId(),
123 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD,
124 scale=DateTime.TAI),
125 "diaObjectId": objId,
126 "ssObjectId": 0,
127 "parentDiaSourceId": 0,
128 "diaSourceId": idx + 1,
129 "midpointMjdTai": midpointMjdTai + 1.0 * idx,
130 "band": exposure.getFilter().bandLabel,
131 "psfNdata": 0,
132 "trailNdata": 0,
133 "dipoleNdata": 0}
134 data.append(diaSource | flags)
136 return pd.DataFrame(data=data).set_index(["diaObjectId", "band", "diaSourceId"], drop=False)
139def makeSolarSystemSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False):
140 """Make a test set of solar system sources.
142 Parameters
143 ----------
144 nSources : `int`
145 Number of sources to create.
146 diaObjectIds : `numpy.ndarray`
147 Integer Ids of diaobjects to "associate" with the DiaSources.
148 exposure : `lsst.afw.image.Exposure`
149 Exposure to create sources over.
150 rng : `numpy.random.Generator`
151 A NumPy random number generator initialized with a fixed seed for reproducibility.
152 randomizeObjects : `bool`, optional
153 If True, randomly draw from `diaObjectIds` to generate the ids in the
154 output catalog, otherwise just iterate through them, repeating as
155 necessary to get nSources objectIds.
157 Returns
158 -------
159 solarSystemSources : `pandas.DataFrame`
160 Solar system sources generated across the exposure.
161 """
162 solarSystemSources = makeDiaSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False)
163 solarSystemSources["ssObjectId"] = rng.integers(0, high=2**63-1, size=nSources)
164 solarSystemSources["Err(arcsec)"] = rng.uniform(0.2, 0.4, size=nSources)
166 return solarSystemSources
169def makeDiaForcedSources(nForcedSources, diaObjectIds, exposure, rng, randomizeObjects=False):
170 """Make a test set of DiaSources.
172 Parameters
173 ----------
174 nForcedSources : `int`
175 Number of sources to create.
176 diaObjectIds : `numpy.ndarray`
177 Integer Ids of diaobjects to "associate" with the DiaSources.
178 exposure : `lsst.afw.image.Exposure`
179 Exposure to create sources over.
180 rng : `numpy.random.Generator`
181 A NumPy random number generator initialized with a fixed seed for reproducibility.
182 randomizeObjects : `bool`, optional
183 If True, randomly draw from `diaObjectIds` to generate the ids in the
184 output catalog, otherwise just iterate through them.
186 Returns
187 -------
188 diaForcedSources : `pandas.DataFrame`
189 DiaForcedSources generated across the exposure.
190 """
191 midpointMjdTai = exposure.visitInfo.date.get(system=DateTime.MJD)
192 visit = exposure.visitInfo.id
193 detector = exposure.detector.getId()
194 if randomizeObjects:
195 objectIds = diaObjectIds[rng.randint(len(diaObjectIds), size=nForcedSources)]
196 else:
197 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nForcedSources)]]
199 data = []
200 bbox = exposure.getBBox()
202 for i, objId in enumerate(objectIds):
203 # Put together the minimum values for the alert.
204 x = rng.uniform(bbox.minX, bbox.maxX)
205 y = rng.uniform(bbox.minY, bbox.maxY)
206 coord = exposure.wcs.pixelToSky(x, y)
207 data.append({"diaForcedSourceId": i + 1,
208 "visit": visit + i,
209 "detector": detector,
210 "diaObjectId": objId,
211 "ra": coord.getRa().asDegrees(),
212 "dec": coord.getDec().asDegrees(),
213 "midpointMjdTai": midpointMjdTai + 1.0 * i,
214 "timeProcessedMjdTai": DateTime.now().get(system=DateTime.MJD,
215 scale=DateTime.TAI),
216 "band": exposure.getFilter().bandLabel})
218 return pd.DataFrame(data=data).set_index(["diaObjectId", "diaForcedSourceId"], drop=False)
221def makeExposure(flipX=False, flipY=False):
222 """Create an exposure and flip the x or y (or both) coordinates.
224 Returns bounding boxes that are right or left handed around the bounding
225 polygon.
227 Parameters
228 ----------
229 flipX : `bool`
230 Flip the x coordinate in the WCS.
231 flipY : `bool`
232 Flip the y coordinate in the WCS.
234 Returns
235 -------
236 exposure : `lsst.afw.image.Exposure`
237 Exposure with a valid bounding box and wcs.
238 """
239 metadata = PropertySet()
241 metadata.set("SIMPLE", "T")
242 metadata.set("BITPIX", -32)
243 metadata.set("NAXIS", 2)
244 metadata.set("NAXIS1", 1024)
245 metadata.set("NAXIS2", 1153)
246 metadata.set("RADECSYS", 'FK5')
247 metadata.set("EQUINOX", 2000.)
249 metadata.setDouble("CRVAL1", 215.604025685476)
250 metadata.setDouble("CRVAL2", 53.1595451514076)
251 metadata.setDouble("CRPIX1", 1109.99981456774)
252 metadata.setDouble("CRPIX2", 560.018167811613)
253 metadata.set("CTYPE1", 'RA---SIN')
254 metadata.set("CTYPE2", 'DEC--SIN')
256 xFlip = 1
257 if flipX:
258 xFlip = -1
259 yFlip = 1
260 if flipY:
261 yFlip = -1
262 metadata.setDouble("CD1_1", xFlip * 5.10808596133527E-05)
263 metadata.setDouble("CD1_2", yFlip * 1.85579539217196E-07)
264 metadata.setDouble("CD2_2", yFlip * -5.10281493481982E-05)
265 metadata.setDouble("CD2_1", xFlip * -8.27440751733828E-07)
267 wcs = afwGeom.makeSkyWcs(metadata)
268 exposure = afwImage.makeExposure(
269 afwImage.makeMaskedImageFromArrays(np.ones((1024, 1153))), wcs)
270 detector = DetectorWrapper(id=23, bbox=exposure.getBBox()).detector
271 visit = afwImage.VisitInfo(
272 exposureTime=200.,
273 date=DateTime("2014-05-13T17:00:00.000000000",
274 DateTime.Timescale.TAI))
275 exposure.info.id = 1234
276 exposure.setDetector(detector)
277 exposure.info.setVisitInfo(visit)
278 exposure.setFilter(afwImage.FilterLabel(band='g'))
279 exposure.setPhotoCalib(afwImage.PhotoCalib(1., 0., exposure.getBBox()))
281 return exposure
284def getRegion(exposure):
285 """Calculate an enveloping region for an exposure.
287 Parameters
288 ----------
289 exposure : `lsst.afw.image.Exposure`
290 Exposure object with calibrated WCS.
292 Returns
293 -------
294 region : `lsst.sphgeom.Region`
295 Region enveloping an exposure.
296 """
297 # Bounding box needs to be a `Box2D` not a `Box2I` for `wcs.pixelToSky()`
298 bbox = lsst.geom.Box2D(exposure.getBBox())
299 wcs = exposure.getWcs()
301 region = lsst.sphgeom.ConvexPolygon([pp.getVector() for pp in wcs.pixelToSky(bbox.getCorners())])
303 return region
306def makeRegionTime(exposure=None, begin=None, end=None):
307 """Make a `RegionTimeInfo` for testing
309 Parameters
310 ----------
311 exposure : `lsst.afw.image.Exposure`, optional
312 Exposure to construct a ``RegionTimeInfo`` for.
313 If None, a default exposure will be created.
314 begin : `astropy.time.Time`, optional
315 The start time of the interval.
316 If `None`, calculate the start time from the midpoint of the exposure
317 and the exposure length.
318 end : `astropy.time.Time`, optional
319 The end time of the interval.
320 If `None`, calculate the end time from the midpoint of the exposure
321 and the exposure length.
323 Returns
324 -------
325 regionTime : `lsst.pipe.base.utils.RegionTimeInfo`
326 Object containing the spatial region and temporal timespan for an exposure.
327 """
328 if exposure is None:
329 exposure = makeExposure()
330 region = getRegion(exposure)
331 expTime = exposure.visitInfo.exposureTime*astropy.units.second
332 # visitInfo time is the midpoint of the exposure.
333 if begin is None:
334 begin = exposure.visitInfo.date.toAstropy() - expTime/2
335 if end is None:
336 end = exposure.visitInfo.date.toAstropy() + expTime/2
337 timespan = dafButler.Timespan(begin=begin, end=end)
338 return RegionTimeInfo(region=region, timespan=timespan)