Coverage for tests/test_diffractionSpikeMask.py: 17%
119 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-08 01:33 -0800
1# This file is part of pipe_tasks.
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/>.
22import os
23import glob
25import numpy as np
27import lsst.afw.image as afwImage
28import lsst.geom as geom
29from lsst.meas.algorithms.testUtils import MockReferenceObjectLoaderFromFiles
30import lsst.meas.base.tests
31from lsst.pipe.tasks.diffractionSpikeMask import DiffractionSpikeMaskTask, DiffractionSpikeMaskConfig
32from lsst.pipe.tasks.colorterms import Colorterm, ColortermDict, ColortermLibrary
33import lsst.utils.tests
35from utils import makeTestVisitInfo
38TESTDIR = os.path.abspath(os.path.dirname(__file__))
39RefCatDir = os.path.join(TESTDIR, "data", "sdssrefcat")
41testColorterms = ColortermLibrary(data={
42 "test*": ColortermDict(data={
43 "test-g": Colorterm(primary="g", secondary="r", c0=0.00, c1=0.00),
44 "test-r": Colorterm(primary="r", secondary="i", c0=0.00, c1=0.00, c2=0.00),
45 "test-i": Colorterm(primary="i", secondary="z", c0=1.00, c1=0.00, c2=0.00),
46 "test-z": Colorterm(primary="z", secondary="i", c0=0.00, c1=0.00, c2=0.00),
47 })
48})
51def setup_module(module):
52 lsst.utils.tests.init()
55class DiffractionSpikeMaskTest(lsst.utils.tests.TestCase):
57 def setUp(self):
59 # Load sample input from disk
60 testDir = os.path.dirname(__file__)
62 # The .xy.fits file has sources in the range ~ [0,2000],[0,4500]
63 # which is bigger than the exposure
64 self.bbox = geom.Box2I(geom.Point2I(0, 0), geom.Extent2I(2048, 4612))
66 smallExposure = afwImage.ExposureF(os.path.join(testDir, "data", "v695833-e0-c000-a00.sci.fits"))
67 self.exposure = afwImage.ExposureF(self.bbox)
68 self.exposure.setWcs(smallExposure.getWcs())
69 self.exposure.setFilter(afwImage.FilterLabel(band="i", physical="test-i"))
70 self.exposure.info.setVisitInfo(makeTestVisitInfo())
72 # Make a reference loader
73 filenames = sorted(glob.glob(os.path.join(RefCatDir, 'ref_cats', 'cal_ref_cat', '??????.fits')))
74 self.refObjLoader = MockReferenceObjectLoaderFromFiles(filenames, htmLevel=8)
76 def tearDown(self):
77 del self.exposure
78 del self.refObjLoader
80 def test_raiseWithoutLoader(self):
81 """The task should raise an error if no reference catalog loader is
82 configured.
83 """
84 config = DiffractionSpikeMaskConfig()
85 task = DiffractionSpikeMaskTask(config=config)
86 with self.assertRaises(RuntimeError):
87 task.run(self.exposure)
89 def test_loadAndMaskStars(self):
90 """Run the bright star mask with a selection of reference sources."""
92 config = DiffractionSpikeMaskConfig(magnitudeThreshold=16)
93 task = DiffractionSpikeMaskTask(self.refObjLoader, config=config)
94 exposure = self.exposure.clone()
95 # Set the saturated mask plane in half of the image
96 saturatedMaskBit = exposure.mask.getPlaneBitMask(config.saturatedMaskPlane)
97 bbox = exposure.getBBox()
98 bbox.grow(-lsst.geom.Extent2I(0, bbox.height//4))
99 exposure[bbox].mask.array |= saturatedMaskBit
100 brightCat = task.run(exposure=exposure)
101 self.assertGreater(len(brightCat), 0)
102 # Verify that the new mask plane has been added
103 spikeMaskBit = exposure.mask.getPlaneBitMask(config.spikeMask)
104 # The images should not be modified
105 self.assertImagesEqual(self.exposure.image, exposure.image)
106 self.assertImagesEqual(self.exposure.variance, exposure.variance)
107 # Ensure that the mask has changed
108 self.assertFloatsNotEqual(self.exposure.mask.array, exposure.mask.array)
110 # Check that the mask is set for the bright sources inside the image
111 # Note that the catalog will include bright sources *off* the image that
112 # have long enough diffraction spikes to overlap the edge of the image
113 xvals, yvals = exposure.wcs.skyToPixelArray(brightCat[config.raKey], brightCat[config.decKey])
114 bbox = exposure.getBBox()
115 # Shrink the bounding box so that the bbox.contains check below is
116 # sufficient to avoid errors.
117 bbox.grow(-2)
118 points = [geom.Point2D(xv, yv) for xv, yv in zip(xvals, yvals)]
119 inside = 0
120 outside = 0
121 for pt in points:
122 if bbox.contains(int(pt.getX()), int(pt.getY())):
123 ptBox = geom.Box2I.makeCenteredBox(pt, geom.Extent2I(3, 3))
124 maskSet = exposure[ptBox].mask.array & spikeMaskBit > 0
125 self.assertTrue(np.all(maskSet))
126 inside += 1
127 else:
128 outside += 1
129 self.assertGreater(inside, 0)
130 self.assertGreater(outside, 0)
132 def test_noBrightStars(self):
133 """Run the bright star mask with no bright stars."""
135 # Set a very high magnitude limit so that no stars are selected
136 config = DiffractionSpikeMaskConfig(magnitudeThreshold=0)
137 task = DiffractionSpikeMaskTask(self.refObjLoader, config=config)
138 exposure = self.exposure.clone()
139 brightCat = task.run(exposure=exposure)
140 self.assertEqual(len(brightCat), 0)
141 # Verify that the new mask plane has been added
142 exposure.mask.getPlaneBitMask(config.spikeMask)
143 # The images should not be modified
144 self.assertImagesEqual(self.exposure.image, exposure.image)
146 def test_maskSources(self):
147 """Verify that sources on and off the image are masked correctly."""
148 task = DiffractionSpikeMaskTask(self.refObjLoader, config=DiffractionSpikeMaskConfig())
149 task.set_diffraction_angle(self.exposure)
150 self.exposure.mask.addMaskPlane(task.config.spikeMask)
152 xSize, ySize = self.bbox.getDimensions()
153 x0, y0 = self.bbox.getBegin()
154 x1, y1 = self.bbox.getEnd()
156 nBright = 50
157 rng = np.random.RandomState(3)
158 xLoc = np.arange(x0 - xSize/4, x1 + xSize/4)
159 rng.shuffle(xLoc)
160 xLoc = xLoc[:nBright]
161 yLoc = np.arange(y0 - ySize/4, y1 + ySize/4)
162 rng.shuffle(yLoc)
163 yLoc = yLoc[:nBright]
164 spikeRadii = np.arange(10, 200)
165 rng.shuffle(spikeRadii)
166 spikeRadii = spikeRadii[:nBright]
167 saturatedBox = geom.Box2I(self.bbox.getBegin(), geom.Extent2I(xSize, ySize//2))
168 baseMask = self.exposure.mask.clone()
169 baseMask[saturatedBox].array |= baseMask.getPlaneBitMask(task.config.saturatedMaskPlane)
170 # There are four classes of sources:
171 # 1. Bright sources on the image with saturated cores - masked
172 # 2. Sources on the image without saturated cores - not masked
173 # 3. Bright sources off the image with predicted diffraction spikes that
174 # overlap the image - masked
175 # 4. Bright sources off the image that are far away enough that any
176 # diffraction spikes would not overlap the image - not masked
177 nClass1 = 0
178 nClass2 = 0
179 nClass3 = 0
180 nClass4 = 0
181 selectedSources = task.selectSources(xLoc, yLoc, spikeRadii, baseMask)
182 for x, y, r, selected in zip(xLoc, yLoc, spikeRadii, selectedSources):
183 mask = baseMask.clone()
184 isInImage = self.bbox.contains(geom.Point2I(x, y))
185 # Bright sources on the image with saturated cores.
186 if isInImage and selected:
187 nClass1 += 1
188 task.maskSources([x], [y], [r], mask)
189 self.assertGreater(np.sum(mask.array & mask.getPlaneBitMask(task.config.spikeMask) > 0), 0)
191 # Bright sources on the image without saturated cores.
192 if isInImage and not selected:
193 nClass2 += 1
194 # Do *not* run task.maskSources in this case, since these are
195 # skipped.
197 # Bright sources off the image that we predict should overlap the
198 # image, should set the SPIKE mask for some pixels.
199 if not isInImage and selected:
200 nClass3 += 1
201 task.maskSources([x], [y], [r], mask)
202 self.assertGreater(np.sum(mask.array & mask.getPlaneBitMask(task.config.spikeMask) > 0), 0)
204 # Sources off the image that are skipped in source selection should
205 # not change the mask even if we do calculate their SPIKE mask.
206 if not isInImage and not selected:
207 nClass4 += 1
208 task.maskSources([x], [y], [r], mask)
209 self.assertMasksEqual(mask, baseMask)
210 # Verify that the test points were sufficient to exercise all classes.
211 self.assertGreater(nClass1, 0)
212 self.assertGreater(nClass2, 0)
213 self.assertGreater(nClass3, 0)
214 self.assertGreater(nClass4, 0)