Coverage for tests/assemble_coadd_test_utils.py: 96%
189 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-04 09:13 +0000
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-04 09:13 +0000
1# This file is part of drp_tasks.
2#
3# LSST Data Management System
4# This product includes software developed by the
5# LSST Project (http://www.lsst.org/).
6# See COPYRIGHT file at the top of the source tree.
7#
8# This program is free software: you can redistribute it and/or modify
9# it under the terms of the GNU General Public License as published by
10# the Free Software Foundation, either version 3 of the License, or
11# (at your option) any later version.
12#
13# This program is distributed in the hope that it will be useful,
14# but WITHOUT ANY WARRANTY; without even the implied warranty of
15# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
16# GNU General Public License for more details.
17#
18# You should have received a copy of the LSST License Statement and
19# the GNU General Public License along with this program. If not,
20# see <https://www.lsstcorp.org/LegalNotices/>.
21#
22"""Set up simulated test data and simplified APIs for AssembleCoaddTask
23and its derived classes.
25This is not intended to test accessing data with the Butler and instead uses
26mock Butler data references to pass in the simulated data.
27"""
29import numpy as np
30from astro_metadata_translator import makeObservationInfo
31from astropy import units as u
32from astropy.coordinates import Angle, EarthLocation, SkyCoord
33from astropy.time import Time
35import lsst.afw.geom as afwGeom
36import lsst.afw.image as afwImage
37import lsst.afw.math as afwMath
38import lsst.afw.table as afwTable
39import lsst.geom as geom
40import lsst.pipe.base as pipeBase
41from lsst.afw.cameraGeom.testUtils import DetectorWrapper
42from lsst.cell_coadds.test_utils import generate_data_id
43from lsst.geom import arcseconds, degrees
44from lsst.meas.algorithms.testUtils import plantSources
45from lsst.obs.base import MakeRawVisitInfoViaObsInfo
46from lsst.pipe.base import InMemoryDatasetHandle
47from lsst.pipe.tasks.coaddBase import growValidPolygons
48from lsst.pipe.tasks.coaddInputRecorder import CoaddInputRecorderConfig, CoaddInputRecorderTask
49from lsst.skymap import Index2D, PatchInfo
51__all__ = ["makeMockSkyInfo", "MockCoaddTestData"]
54class _MockTractInfo:
56 def __init__(self, patch_info: PatchInfo):
57 self._patch_info = patch_info
59 def __getitem__(self, _) -> PatchInfo:
60 return self._patch_info
62 def getBBox(self) -> geom.Box2I:
63 return self._patch_info.getOuterBBox()
66def makeMockSkyInfo(bbox, wcs, patch):
67 """Construct a `Struct` containing the geometry of the patch to be coadded.
69 Parameters
70 ----------
71 bbox : `lsst.geom.Box`
72 Bounding box of the patch to be coadded.
73 wcs : `lsst.afw.geom.SkyWcs`
74 Coordinate system definition (wcs) for the exposure.
76 Returns
77 -------
78 skyInfo : `lsst.pipe.base.Struct`
79 Patch geometry information.
80 """
81 patchInfo = PatchInfo(
82 index=Index2D(0, 0),
83 sequentialIndex=patch,
84 innerBBox=bbox,
85 outerBBox=bbox,
86 tractWcs=wcs,
87 numCellsPerPatchInner=1,
88 cellInnerDimensions=(bbox.width, bbox.height),
89 )
90 skyInfo = pipeBase.Struct(bbox=bbox, wcs=wcs, tractInfo=_MockTractInfo(patchInfo), patchInfo=patchInfo)
91 return skyInfo
94class MockCoaddTestData:
95 """Generate repeatable simulated exposures with consistent metadata that
96 are realistic enough to test the image coaddition algorithms.
98 Notes
99 -----
100 The simple GaussianPsf used by lsst.meas.algorithms.testUtils.plantSources
101 will always return an average position of (0, 0).
102 The bounding box of the exposures MUST include (0, 0), or else the PSF will
103 not be valid and `AssembleCoaddTask` will fail with the error
104 'Could not find a valid average position for CoaddPsf'.
106 Parameters
107 ----------
108 shape : `lsst.geom.Extent2I`, optional
109 Size of the bounding box of the exposures to be simulated, in pixels.
110 offset : `lsst.geom.Point2I`, optional
111 Pixel coordinate of the lower left corner of the bounding box.
112 backgroundLevel : `float`, optional
113 Background value added to all pixels in the simulated images.
114 seed : `int`, optional
115 Seed value to initialize the random number generator.
116 nSrc : `int`, optional
117 Number of sources to simulate.
118 fluxRange : `float`, optional
119 Range in flux amplitude of the simulated sources.
120 noiseLevel : `float`, optional
121 Standard deviation of the noise to add to each pixel.
122 sourceSigma : `float`, optional
123 Average amplitude of the simulated sources,
124 relative to ``noiseLevel``
125 minPsfSize : `float`, optional
126 The smallest PSF width (sigma) to use, in pixels.
127 maxPsfSize : `float`, optional
128 The largest PSF width (sigma) to use, in pixels.
129 pixelScale : `lsst.geom.Angle`, optional
130 The plate scale of the simulated images.
131 ra : `lsst.geom.Angle`, optional
132 Right Ascension of the boresight of the camera for the observation.
133 dec : `lsst.geom.Angle`, optional
134 Declination of the boresight of the camera for the observation.
135 ccd : `int`, optional
136 CCD number to put in the metadata of the exposure.
137 patch : `int`, optional
138 Unique identifier for a subdivision of a tract.
139 tract : `int`, optional
140 Unique identifier for a tract of a skyMap.
142 Raises
143 ------
144 ValueError
145 If the bounding box does not contain the pixel coordinate (0, 0).
146 This is due to `GaussianPsf` that is used by
147 `~lsst.meas.algorithms.testUtils.plantSources`
148 lacking the option to specify the pixel origin.
149 """
151 rotAngle = 0.0 * degrees
152 """Rotation of the pixel grid on the sky, East from North
153 (`lsst.geom.Angle`).
154 """
155 filterLabel = None
156 """The filter definition, usually set in the current instruments' obs
157 package. For these tests, a simple filter is defined without using an obs
158 package (`lsst.afw.image.FilterLabel`).
159 """
160 rngData = None
161 """Pre-initialized random number generator for constructing the test images
162 repeatably (`numpy.random.Generator`).
163 """
164 rngMods = None
165 """Pre-initialized random number generator for applying modifications to
166 the test images for only some test cases (`numpy.random.Generator`).
167 """
168 kernelSize = None
169 "Width of the kernel used for simulating sources, in pixels."
170 exposures = {}
171 """The simulated test data, with variable PSF size
172 (`dict` of `lsst.afw.image.Exposure`)
173 """
174 matchedExposures = {}
175 """The simulated exposures, all with PSF width set to ``maxPsfSize``
176 (`dict` of `lsst.afw.image.Exposure`).
177 """
178 photoCalib = afwImage.makePhotoCalibFromCalibZeroPoint(27, 10)
179 """The photometric zero point to use for converting counts to flux units
180 (`lsst.afw.image.PhotoCalib`).
181 """
182 badMaskPlanes = ["NO_DATA", "BAD"]
183 """Mask planes that, if set, the associated pixel should not be included in
184 the coaddTempExp.
185 """
186 detector = None
187 "Properties of the CCD for the exposure (`lsst.afw.cameraGeom.Detector`)."
189 def __init__(
190 self,
191 shape=geom.Extent2I(201, 301),
192 offset=geom.Point2I(-123, -45),
193 backgroundLevel=314.592,
194 seed=42,
195 nSrc=37,
196 fluxRange=2.0,
197 noiseLevel=5,
198 sourceSigma=200.0,
199 minPsfSize=1.5,
200 maxPsfSize=3.0,
201 pixelScale=0.2 * arcseconds,
202 ra=209.0 * degrees,
203 dec=-20.25 * degrees,
204 ccd=37,
205 patch=42,
206 tract=0,
207 ):
208 self.ra = ra
209 self.dec = dec
210 self.pixelScale = pixelScale
211 self.patch = patch
212 self.tract = tract
213 self.filterLabel = afwImage.FilterLabel(band="gTest", physical="gTest")
214 self.rngData = np.random.default_rng(seed)
215 self.rngMods = np.random.default_rng(seed + 1)
216 self.bbox = geom.Box2I(offset, shape)
217 if not self.bbox.contains(0, 0): 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true
218 raise ValueError(f"The bounding box must contain the coordinate (0, 0). {repr(self.bbox)}")
219 self.wcs = self.makeDummyWcs()
221 # Set up properties of the simulations
222 nSigmaForKernel = 5
223 self.kernelSize = (int(maxPsfSize * nSigmaForKernel + 0.5) // 2) * 2 + 1 # make sure it is odd
225 bufferSize = self.kernelSize // 2
226 x0, y0 = self.bbox.getBegin()
227 xSize, ySize = self.bbox.getDimensions()
228 # Set the pixel coordinates and fluxes of the simulated sources.
229 self.xLoc = self.rngData.random(nSrc) * (xSize - 2 * bufferSize) + bufferSize + x0
230 self.yLoc = self.rngData.random(nSrc) * (ySize - 2 * bufferSize) + bufferSize + y0
231 self.flux = (self.rngData.random(nSrc) * (fluxRange - 1.0) + 1.0) * sourceSigma * noiseLevel
233 self.backgroundLevel = backgroundLevel
234 self.noiseLevel = noiseLevel
235 self.minPsfSize = minPsfSize
236 self.maxPsfSize = maxPsfSize
237 self.detector = DetectorWrapper(name=f"detector {ccd}", id=ccd).detector
239 def setDummyCoaddInputs(self, exposure, expId):
240 """Generate an `ExposureCatalog` as though the exposures had been
241 processed using `make_direct_warp`.
243 Parameters
244 ----------
245 exposure : `lsst.afw.image.Exposure`
246 The exposure to construct a `CoaddInputs` `ExposureCatalog` for.
247 expId : `int`
248 A unique identifier for the visit.
249 """
250 badPixelMask = afwImage.Mask.getPlaneBitMask(self.badMaskPlanes)
251 nGoodPix = np.sum(exposure.getMask().getArray() & badPixelMask == 0)
253 config = CoaddInputRecorderConfig()
254 inputRecorder = CoaddInputRecorderTask(config=config, name="inputRecorder")
255 tempExpInputRecorder = inputRecorder.makeCoaddTempExpRecorder(expId, num=1)
256 tempExpInputRecorder.addCalExp(exposure, expId, nGoodPix)
257 tempExpInputRecorder.finish(exposure, nGoodPix=nGoodPix)
258 growValidPolygons(exposure.getInfo().getCoaddInputs(), growBy=0)
260 def makeCoaddTempExp(self, rawExposure, visitInfo, expId, apCorrMap=None):
261 """Add the metadata required by `AssembleCoaddTask` to an exposure.
263 Parameters
264 ----------
265 rawExposure : `lsst.afw.image.Exposure`
266 The simulated exposure.
267 visitInfo : `lsst.afw.image.VisitInfo`
268 VisitInfo containing metadata for the exposure.
269 expId : `int`
270 A unique identifier for the visit.
272 Returns
273 -------
274 tempExp : `lsst.afw.image.Exposure`
275 The exposure, with all of the metadata needed for coaddition.
276 """
277 tempExp = rawExposure.clone()
278 tempExp.setWcs(self.wcs)
280 tempExp.setFilter(self.filterLabel)
281 tempExp.setPhotoCalib(self.photoCalib)
282 tempExp.setApCorrMap(apCorrMap)
283 tempExp.getInfo().setVisitInfo(visitInfo)
284 tempExp.getInfo().setDetector(self.detector)
285 self.setDummyCoaddInputs(tempExp, expId)
286 return tempExp
288 def makeDummyWcs(self, rotAngle=None, pixelScale=None, crval=None, flipX=True):
289 """Make a World Coordinate System object for testing.
291 Parameters
292 ----------
293 rotAngle : `lsst.geom.Angle`
294 Rotation of the CD matrix, East from North
295 pixelScale : `lsst.geom.Angle`
296 Pixel scale of the projection.
297 crval : `lsst.afw.geom.SpherePoint`
298 Coordinates of the reference pixel of the wcs.
299 flipX : `bool`, optional
300 Flip the direction of increasing Right Ascension.
302 Returns
303 -------
304 wcs : `lsst.afw.geom.skyWcs.SkyWcs`
305 A wcs that matches the inputs.
306 """
307 if rotAngle is None: 307 ↛ 309line 307 didn't jump to line 309 because the condition on line 307 was always true
308 rotAngle = self.rotAngle
309 if pixelScale is None: 309 ↛ 311line 309 didn't jump to line 311 because the condition on line 309 was always true
310 pixelScale = self.pixelScale
311 if crval is None: 311 ↛ 313line 311 didn't jump to line 313 because the condition on line 311 was always true
312 crval = geom.SpherePoint(self.ra, self.dec)
313 crpix = geom.Box2D(self.bbox).getCenter()
314 cdMatrix = afwGeom.makeCdMatrix(scale=pixelScale, orientation=rotAngle, flipX=flipX)
315 wcs = afwGeom.makeSkyWcs(crpix=crpix, crval=crval, cdMatrix=cdMatrix)
316 return wcs
318 def makeDummyVisitInfo(self, exposureId, randomizeTime=False):
319 """Make a self-consistent visitInfo object for testing.
321 Parameters
322 ----------
323 exposureId : `int`, optional
324 Unique integer identifier for this observation.
325 randomizeTime : `bool`, optional
326 Add a random offset within a 6 hour window to the observation time.
328 Returns
329 -------
330 visitInfo : `lsst.afw.image.VisitInfo`
331 VisitInfo for the exposure.
332 """
333 lsstLat = -30.244639 * u.degree
334 lsstLon = -70.749417 * u.degree
335 lsstAlt = 2663.0 * u.m
336 lsstTemperature = 20.0 * u.Celsius
337 lsstHumidity = 40.0 # in percent
338 lsstPressure = 73892.0 * u.pascal
339 loc = EarthLocation(lat=lsstLat, lon=lsstLon, height=lsstAlt)
341 time = Time(2000.0, format="jyear", scale="tt")
342 if randomizeTime: 342 ↛ 345line 342 didn't jump to line 345 because the condition on line 342 was always true
343 # Pick a random time within a 6 hour window
344 time += 6 * u.hour * (self.rngMods.random() - 0.5)
345 radec = SkyCoord(
346 dec=self.dec.asDegrees(),
347 ra=self.ra.asDegrees(),
348 unit="deg",
349 obstime=time,
350 frame="icrs",
351 location=loc,
352 )
353 airmass = float(1.0 / np.sin(radec.altaz.alt))
354 obsInfo = makeObservationInfo(
355 location=loc,
356 detector_exposure_id=exposureId,
357 datetime_begin=time,
358 datetime_end=time,
359 boresight_airmass=airmass,
360 boresight_rotation_angle=Angle(0.0 * u.degree),
361 boresight_rotation_coord="sky",
362 temperature=lsstTemperature,
363 pressure=lsstPressure,
364 relative_humidity=lsstHumidity,
365 tracking_radec=radec,
366 altaz_begin=radec.altaz,
367 observation_type="science",
368 )
369 visitInfo = MakeRawVisitInfoViaObsInfo.observationInfo2visitInfo(obsInfo)
370 return visitInfo
372 def makeDummyApCorrMap(self):
373 """Make a dummy aperture correction map for testing.
375 This method returns an `~lsst.afw.image.ApCorrMap` instance with
376 random values for the "algo1_instFlux", "algo1_instFluxErr",
377 "algo2_instFlux", and "algo2_instFluxErr" fields that are spatially
378 constant.
380 Returns
381 -------
382 apCorrMap : `lsst.afw.image.ApCorrMap`
383 Aperture correction map for the exposure.
384 """
385 apCorrMap = afwImage.ApCorrMap()
386 apCorrMap["algo1_instFlux"] = afwMath.ChebyshevBoundedField(
387 self.bbox,
388 np.array([[self.rngMods.random()]]),
389 )
390 apCorrMap["algo1_instFluxErr"] = afwMath.ChebyshevBoundedField(
391 self.bbox,
392 np.array([[self.rngMods.random()]]),
393 )
394 apCorrMap["algo2_instFlux"] = afwMath.ChebyshevBoundedField(
395 self.bbox,
396 np.array([[self.rngMods.random()]]),
397 )
398 apCorrMap["algo2_instFluxErr"] = afwMath.ChebyshevBoundedField(
399 self.bbox,
400 np.array([[self.rngMods.random()]]),
401 )
403 return apCorrMap
405 def makeTestImage(
406 self,
407 expId,
408 noiseLevel=None,
409 psfSize=None,
410 backgroundLevel=None,
411 detectionSigma=5.0,
412 badRegionBox=None,
413 ):
414 """Make a reproduceable PSF-convolved masked image for testing.
416 Parameters
417 ----------
418 expId : `int`
419 A unique identifier to use to refer to the visit.
420 noiseLevel : `float`, optional
421 Standard deviation of the noise to add to each pixel.
422 psfSize : `float`, optional
423 Width of the PSF of the simulated sources, in pixels.
424 backgroundLevel : `float`, optional
425 Background value added to all pixels in the simulated images.
426 detectionSigma : `float`, optional
427 Threshold amplitude of the image to set the "DETECTED" mask.
428 badRegionBox : `lsst.geom.Box2I`, optional
429 Add a bad region bounding box (set to "BAD").
430 """
431 if backgroundLevel is None: 431 ↛ 433line 431 didn't jump to line 433 because the condition on line 431 was always true
432 backgroundLevel = self.backgroundLevel
433 if noiseLevel is None: 433 ↛ 435line 433 didn't jump to line 435 because the condition on line 433 was always true
434 noiseLevel = 5.0
435 visitInfo = self.makeDummyVisitInfo(expId, randomizeTime=True)
437 if psfSize is None: 437 ↛ 439line 437 didn't jump to line 439 because the condition on line 437 was always true
438 psfSize = self.rngMods.random() * (self.maxPsfSize - self.minPsfSize) + self.minPsfSize
439 nSrc = len(self.flux)
440 sigmas = [psfSize for src in range(nSrc)]
441 sigmasPsfMatched = [self.maxPsfSize for src in range(nSrc)]
442 coordList = list(zip(self.xLoc, self.yLoc, self.flux, sigmas))
443 coordListPsfMatched = list(zip(self.xLoc, self.yLoc, self.flux, sigmasPsfMatched))
444 xSize, ySize = self.bbox.getDimensions()
445 model = plantSources(
446 self.bbox, self.kernelSize, self.backgroundLevel, coordList, addPoissonNoise=False
447 )
448 modelPsfMatched = plantSources(
449 self.bbox, self.kernelSize, self.backgroundLevel, coordListPsfMatched, addPoissonNoise=False
450 )
451 model.variance.array = np.abs(model.image.array) + noiseLevel
452 modelPsfMatched.variance.array = np.abs(modelPsfMatched.image.array) + noiseLevel
453 noise = self.rngData.random((ySize, xSize)) * noiseLevel
454 noise -= np.median(noise)
455 model.image.array += noise
456 modelPsfMatched.image.array += noise
457 detectedMask = afwImage.Mask.getPlaneBitMask("DETECTED")
458 detectionThreshold = self.backgroundLevel + detectionSigma * noiseLevel
459 model.mask.array[model.image.array > detectionThreshold] += detectedMask
461 if badRegionBox is not None:
462 model.mask[badRegionBox] = afwImage.Mask.getPlaneBitMask("BAD")
464 apCorrMap = self.makeDummyApCorrMap()
465 exposure = self.makeCoaddTempExp(model, visitInfo, expId, apCorrMap)
466 matchedExposure = self.makeCoaddTempExp(modelPsfMatched, visitInfo, expId, apCorrMap)
467 exposure.metadata["BUNIT"] = "nJy"
468 matchedExposure.metadata["BUNIT"] = "nJy"
469 return exposure, matchedExposure
471 def makeVisitSummaryTableHandle(self, warpHandle):
472 schema = afwTable.ExposureTable.makeMinimalSchema()
473 schema.addField("visit", type=np.int64, doc="Visit ID")
474 schema.addField("ccd", type=np.int64, doc="CCD ID")
475 schema.addField("meanVar", type=float, units="nJy**2", doc="Mean variance")
476 visitSummaryTable = afwTable.ExposureCatalog(schema)
478 warp = warpHandle.get()
479 for ccd in warp.getInfo().getCoaddInputs().ccds:
480 record = visitSummaryTable.addNew()
481 record.set("id", ccd["ccd"])
482 record.set("ccd", ccd["ccd"])
483 record.set("visit", ccd["visit"])
484 record.set("meanVar", 0.3 + self.rngMods.random())
485 record.setPhotoCalib(afwImage.PhotoCalib(calibrationMean=10.0))
487 handle = InMemoryDatasetHandle(visitSummaryTable, dataId=warpHandle.dataId)
488 return handle
490 @staticmethod
491 def makeDataRefList(exposures, tract=0, patch=42):
492 """Make data references from the simulated exposures that can be
493 retrieved using the Gen 3 Butler API.
495 Parameters
496 ----------
497 exposures : `Mapping` [`Any`, `~lsst.afw.image.ExposureF`]
498 A mapping of exposure IDs to ExposureF objects.
499 tract : `int`, optional
500 Unique identifier for a tract of a skyMap.
501 patch : `int`, optional
502 Unique identifier for a subdivision of a tract.
504 Returns
505 -------
506 dataRefList : `list` [`~lsst.pipe.base.InMemoryDatasetHandle`]
507 The data references.
509 Raises
510 ------
511 ValueError
512 If an unknown `warpType` is supplied.
513 """
514 dataRefList = []
515 for expId in exposures:
516 exposure = exposures[expId]
517 dataRef = pipeBase.InMemoryDatasetHandle(
518 exposure,
519 storageClass="ExposureF",
520 copy=True,
521 dataId=generate_data_id(
522 tract=tract,
523 patch=patch,
524 visit_id=expId,
525 ),
526 )
527 dataRefList.append(dataRef)
528 return dataRefList