Skip to content

Commit

Permalink
Review changes
Browse files Browse the repository at this point in the history
  • Loading branch information
@bsmartradio
bsmartradio committed Oct 2, 2023
1 parent c55879a commit 918b188
Show file tree
Hide file tree
Showing 2 changed files with 153 additions and 59 deletions.
20 changes: 12 additions & 8 deletions python/lsst/meas/extensions/trailedSources/NaivePlugin.py
View file
Expand Up @@ -27,7 +27,7 @@
from scipy.special import erf
from math import sqrt

from lsst.geom import Point2D
from lsst.geom import Point2D, Point2I
from lsst.meas.base.pluginRegistry import register
from lsst.meas.base import SingleFramePlugin, SingleFramePluginConfig
from lsst.meas.base import FlagHandler, FlagDefinitionList, SafeCentroidExtractor
Expand Down Expand Up @@ -124,7 +124,7 @@ def __init__(self, config, name, schema, metadata, logName=None):
self.NO_CONVERGE = flagDefs.add("flag_noConverge", "The root finder did not converge")
self.NO_SIGMA = flagDefs.add("flag_noSigma", "No PSF width (sigma)")
self.SAFE_CENTROID = flagDefs.add("flag_safeCentroid", "Fell back to safe centroid extractor")
self.EDGE = flagDefs.add("flag_edge", "trail contains edge pixels or extends off chip")
self.EDGE = flagDefs.add("flag_edge", "Trail contains edge pixels or extends off chip")
self.flagHandler = FlagHandler.addFields(schema, name, flagDefs)

self.centriodExtractor = SafeCentroidExtractor(schema, name)
Expand Down Expand Up @@ -192,16 +192,20 @@ def measure(self, measRecord, exposure):
x1 = xc + dydtheta
y1 = yc + dxdtheta

if not (0 <= x0 <= exposure.getBBox().endX and 0 <= x1 <= exposure.getBBox().endX
and 0 <= y0 <= exposure.getBBox().endY and 0 <= y1 <= exposure.getBBox().endY):
if not (exposure.getBBox().beginX <= x0 <= exposure.getBBox().endX
and exposure.getBBox().beginX <= x1 <= exposure.getBBox().endX
and exposure.getBBox().beginY <= y0 <= exposure.getBBox().endY
and exposure.getBBox().beginY <= y1 <= exposure.getBBox().endY):

self.flagHandler.setValue(measRecord, self.EDGE.number, True)

else:

if exposure.mask[np.round(x0), np.round(y0)] and exposure.mask[np.round(x1), np.round(y1)]:
if ((exposure.mask[np.round(x0), np.round(y0)] & exposure.mask.getPlaneBitMask('EDGE') != 0)
or (exposure.mask[np.round(x1), np.round(y1)]
# Check whether the beginning or end point of the trail has the
# edge flag set. The end points are not whole pixel values, so
# the pixel value must be rounded.
if exposure.mask[Point2I(int(x0), int(y0))] and exposure.mask[Point2I(int(x1), int(y1))]:
if ((exposure.mask[Point2I(int(x0), int(y0))] & exposure.mask.getPlaneBitMask('EDGE') != 0)
or (exposure.mask[Point2I(int(x1), int(y1))]
& exposure.mask.getPlaneBitMask('EDGE') != 0)):

self.flagHandler.setValue(measRecord, self.EDGE.number, True)
Expand Down
192 changes: 141 additions & 51 deletions tests/test_trailedEdgeSources.py
View file
Expand Up @@ -27,12 +27,14 @@
import lsst.meas.extensions.trailedSources
from lsst.meas.base.tests import AlgorithmTestCase
from lsst.utils.tests import classParameters
from lsst.geom import Point2I, Point2D, Box2I, Extent2I

# Trailed-source length, angle, and centroid, one on the edge, one not.
Ls = np.array([5, 5])
thetas = np.array([100, 0])
xcs = np.array([100, 20])
ycs = np.array([100, 20])

# Trailed-source length, angle, and centroid.
length = np.array([5, 5, 10])
angle = np.array([100, 0, 5])
x_coord = np.array([100, 20, -20])
y_coord = np.array([100, 20, -30])


class TrailedEdgeSource:
Expand All @@ -43,7 +45,7 @@ def __init__(self, instFlux, length, angle, xc, yc):
self.instFlux = instFlux
self.length = length
self.angle = angle
self.center = lsst.geom.Point2D(xc, yc)
self.center = Point2D(xc, yc)
self.x0 = xc - length / 2 * np.cos(angle)
self.y0 = yc - length / 2 * np.sin(angle)
self.x1 = xc + length / 2 * np.cos(angle)
Expand All @@ -61,13 +63,14 @@ def __init__(self, bbox, threshold=10.0, exposure=None, **kwds):

super().__init__(bbox, threshold, exposure, **kwds)

def addTrailedSource(self, trail):
def addTrailedSource(self, trail, edge=True):
"""Add a trailed source to the simulation.
'Re-implemented' version of
Re-implemented version of
`lsst.meas.base.tests.TestDataset.addSource`. Numerically integrates a
Gaussian PSF over a line to obtain am image of a trailed source.
Gaussian PSF over a line to obtain an image of a trailed source and
adds edge flags to the image.
"""

record = self.catalog.addNew()
record.set(self.keys["centroid"], trail.center)
rng = np.random.default_rng(32)
Expand All @@ -82,21 +85,23 @@ def addTrailedSource(self, trail):
xp = np.linspace(trail.x0, trail.x1, num=numIter)
yp = np.linspace(trail.y0, trail.y1, num=numIter)
for (x, y) in zip(xp, yp):
pt = lsst.geom.Point2D(x, y)
pt = Point2D(x, y)
im = self.drawGaussian(self.exposure.getBBox(), trail.instFlux,
lsst.afw.geom.Ellipse(self.psfShape, pt))
self.exposure.getMaskedImage().getImage().getArray()[:, :] += im.getArray()

planes = self.exposure.mask.getMaskPlaneDict()
dim = self.exposure.getBBox().getDimensions()

for y in range(20):
self.exposure.mask.setMaskPlaneValues(planes['EDGE'], 0, dim[0] - 1, y)
self.exposure.mask.setMaskPlaneValues(planes['EDGE'], 0, dim[0] - 1, y + dim[1] - 20)
# Add edge flags to the first and last 20 columns and rows.
if edge:
for y in range(20):
self.exposure.mask.setMaskPlaneValues(planes['EDGE'], 0, dim[0] - 1, y)
self.exposure.mask.setMaskPlaneValues(planes['EDGE'], 0, dim[0] - 1, y + dim[1] - 20)

for y in range(dim[1]):
self.exposure.mask.setMaskPlaneValues(planes['EDGE'], 0, 20, y)
self.exposure.mask.setMaskPlaneValues(planes['EDGE'], dim[0] - 20, dim[0] - 1, y)
for y in range(dim[1]):
self.exposure.mask.setMaskPlaneValues(planes['EDGE'], 0, 20, y)
self.exposure.mask.setMaskPlaneValues(planes['EDGE'], dim[0] - 20, dim[0] - 1, y)

totFlux = self.exposure.image.array.sum()
self.exposure.image.array /= totFlux
Expand All @@ -109,13 +114,19 @@ def addTrailedSource(self, trail):


# Following from test_trailedSources
@classParameters(length=Ls, theta=thetas, xc=xcs, yc=ycs)
@classParameters(length=length, theta=angle, xc=x_coord, yc=y_coord)
class TrailedEdgeSourcesTestCase(AlgorithmTestCase, lsst.utils.tests.TestCase):
""" Test if ext_trailedSources_Naive_flag_edge is set correctly.
Given a `TrailedSource`, test if the edge flag is set correctly in the
source catalog after the
'lsst.meas.extensions.trailedSources.Naive.Plugin.makeTrailedSourceMeasurementTask'
has been run on the source catalog.
"""

def setUp(self):
self.center = lsst.geom.Point2D(50.1, 49.8)
self.bbox = lsst.geom.Box2I(lsst.geom.Point2I(-20, -30),
lsst.geom.Extent2I(140, 160))
self.center = Point2D(50.1, 49.8)
self.bbox = Box2I(lsst.geom.Point2I(-20, -30), Extent2I(140, 160))
self.dataset = TrailedTestDataset(self.bbox)

# Trail which extends into edge pixels
Expand All @@ -133,7 +144,6 @@ def makeTrailedSourceMeasurementTask(self, plugin=None, dependencies=(),
algMetadata=None):
"""Set up a measurement task for a trailed source plugin.
"""

config = self.makeSingleFrameMeasurementConfig(plugin=plugin,
dependencies=dependencies)

Expand All @@ -146,65 +156,145 @@ def makeTrailedSourceMeasurementTask(self, plugin=None, dependencies=(),
algMetadata=algMetadata)

def testEdgeFlag(self):
"""Test the NaivePlugin with one trailed source crossing edge pixels
and one source which does not.
"""Test if edge flags are correctly set in NaivePlugin.py
Given a `TrailedTestDataset`, run the NaivePlugin measurement and
check that the trailed sources have the edge flag set.
check that the trailed sources have the edge flag set. [100,100] does
not contain any edge pixels and should not have a flag set, [20,20]
crosses into the edge region on onle one side and should have the edge
flag set, and [-20,-30] extends off the chip and should have the edge
flag set.
"""

# Set up and run Naive measurement.
task = self.makeTrailedSourceMeasurementTask(
plugin="ext_trailedSources_Naive",
dependencies=("base_SdssCentroid", "base_SdssShape")
)
exposure, catalog = self.dataset.realize(5.0, task.schema,
randomSeed=0)
exposure, catalog = self.dataset.realize(5.0, task.schema, randomSeed=0)
task.run(catalog, exposure)
record = catalog[0]

# Check that x0, y0 or x1, y1 is flagged as an edge pixel
x1 = np.round(record['ext_trailedSources_Naive_x1'])
y1 = np.round(record['ext_trailedSources_Naive_y1'])
x0 = np.round(record['ext_trailedSources_Naive_x0'])
y0 = np.round(record['ext_trailedSources_Naive_y0'])

edge_pixel_set = np.array([(exposure.mask[x0, y0] & exposure.mask.getPlaneBitMask(
'EDGE') != 0), (exposure.mask[x1, y1] & exposure.mask.getPlaneBitMask('EDGE') != 0)])
x1 = int(record['ext_trailedSources_Naive_x1'])
y1 = int(record['ext_trailedSources_Naive_y1'])
x0 = int(record['ext_trailedSources_Naive_x0'])
y0 = int(record['ext_trailedSources_Naive_y0'])

# Checks trailed which crosses into edge pixels
# Test Case with no edge pixels
if record['truth_x'] == 100:

self.assertTrue(edge_pixel_set.any())
# These are used to ensure the mask pixels the trailed sources are
# compared with have the correct flags set
begin_edge_pixel_set = (exposure.mask[Point2I(x0, y0)] & exposure.mask.getPlaneBitMask(
'EDGE') != 0)
end_edge_pixel_set = (exposure.mask[Point2I(x1, y1)] & exposure.mask.getPlaneBitMask(
'EDGE') != 0)

self.assertFalse(begin_edge_pixel_set)
self.assertTrue(end_edge_pixel_set)

# Make sure measurement edge flag is set, but Naive_flag is not.
# A failed trailed source measurement can still have the edge flag
# set but we are looking for only successful measurements with the
# flag set.
# A failed trailed source measurement with the edge flag
# set means the edge flag was set despite the measurement
# failing.
self.assertTrue(record.get("ext_trailedSources_Naive_flag_edge"))
self.assertFalse(record.get("ext_trailedSources_Naive_flag"))

x1 = np.round(record['ext_trailedSources_Naive_x1'])
y1 = np.round(record['ext_trailedSources_Naive_y1'])
x1 = int(record['ext_trailedSources_Naive_x1'])
y1 = int(record['ext_trailedSources_Naive_y1'])

self.assertFalse(exposure.mask[x0, y0] & exposure.mask.getPlaneBitMask('EDGE') != 0)
self.assertTrue(exposure.mask[x1, y1] & exposure.mask.getPlaneBitMask('EDGE') != 0)
self.assertFalse(exposure.mask[Point2I(x0, y0)] & exposure.mask.getPlaneBitMask('EDGE') != 0)
self.assertTrue(exposure.mask[Point2I(x1, y1)] & exposure.mask.getPlaneBitMask('EDGE') != 0)

# Checks trailed source which does not contain edge pixels
else:
# Test case with one end of trail containing edge pixels
elif record['truth_x'] == 20:

self.assertFalse(edge_pixel_set.any())
begin_edge_pixel_set = (exposure.mask[Point2I(x0, y0)] & exposure.mask.getPlaneBitMask(
'EDGE') != 0)
end_edge_pixel_set = (exposure.mask[Point2I(x1, y1)] & exposure.mask.getPlaneBitMask(
'EDGE') != 0)

self.assertFalse(begin_edge_pixel_set)
self.assertFalse(end_edge_pixel_set)

# Make sure measurement Naive_flag_edge and Naive_flag not set
self.assertFalse(record.get("ext_trailedSources_Naive_flag_edge"))
self.assertFalse(record.get("ext_trailedSources_Naive_flag"))

x1 = np.round(record['ext_trailedSources_Naive_x1'])
y1 = np.round(record['ext_trailedSources_Naive_y1'])
x1 = int(record['ext_trailedSources_Naive_x1'])
y1 = int(record['ext_trailedSources_Naive_y1'])

self.assertFalse(exposure.mask[Point2I(x0, y0)] & exposure.mask.getPlaneBitMask('EDGE') != 0)
self.assertFalse(exposure.mask[Point2I(x1, y1)] & exposure.mask.getPlaneBitMask('EDGE') != 0)

# Test case with trailed source extending off chip.
else:

self.assertTrue(record.get("ext_trailedSources_Naive_flag_edge"))
self.assertFalse(record.get("ext_trailedSources_Naive_flag"))


@classParameters(length=[10], theta=[5], xc=[-20], yc=[-30])
class TrailedEdgeSourcesOffImageTest(AlgorithmTestCase, lsst.utils.tests.TestCase):
""" Test if ext_trailedSources_Naive_flag_edge is set correctly.
Given a `TrailedSource`, test if the edge flag is set correctly in the
source catalog after the
'lsst.meas.extensions.trailedSources.Naive.Plugin.makeTrailedSourceMeasurementTask'
has been run on the source catalog.
"""

def setUp(self):
self.center = Point2D(50.1, 49.8)
self.bbox = Box2I(lsst.geom.Point2I(-20, -30), Extent2I(140, 160))
self.dataset = TrailedTestDataset(self.bbox)

# Trail which extends into edge pixels
self.trail = TrailedEdgeSource(100000.0, self.length, self.theta,
self.xc, self.yc)
self.dataset.addTrailedSource(self.trail, edge=False)

def tearDown(self):
del self.center
del self.bbox
del self.trail
del self.dataset

def makeTrailedSourceMeasurementTask(self, plugin=None, dependencies=(),
config=None, schema=None,
algMetadata=None):
"""Set up a measurement task for a trailed source plugin.
"""
config = self.makeSingleFrameMeasurementConfig(plugin=plugin,
dependencies=dependencies)

self.assertFalse(exposure.mask[x0, y0] & exposure.mask.getPlaneBitMask('EDGE') != 0)
# Make sure the shape slot is base_SdssShape
config.slots.shape = "base_SdssShape"
return self.makeSingleFrameMeasurementTask(plugin=plugin,
dependencies=dependencies,
config=config,
schema=schema,
algMetadata=algMetadata)

def testOffImageEdgeFlag(self):
"""Test if edge flags are correctly set in NaivePlugin.py when source
extends off the the image.
Given a `TrailedTestDataset`, run the NaivePlugin measurement and
check that the edge flag set when a source extends off the chip.
Edge pixels are not set in this test.
"""
# Set up and run Naive measurement.
task = self.makeTrailedSourceMeasurementTask(
plugin="ext_trailedSources_Naive",
dependencies=("base_SdssCentroid", "base_SdssShape")
)
exposure, catalog = self.dataset.realize(5.0, task.schema, randomSeed=0)
task.run(catalog, exposure)
record = catalog[0]

self.assertFalse(exposure.mask[x1, y1] & exposure.mask.getPlaneBitMask('EDGE') != 0)
self.assertTrue(record.get("ext_trailedSources_Naive_flag_edge"))
self.assertFalse(record.get("ext_trailedSources_Naive_flag"))


class TestMemory(lsst.utils.tests.MemoryTestCase):
Expand Down

0 comments on commit 918b188

Please sign in to comment.