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