Make trailedAssociatorTask

bsmartradio · bsmartradio · commit 62f1e4b24c13 · 2023-06-29T11:38:17.000-07:00
Make trailedAssociatorTask which filters out trails whose lengths are above 0.416 arcseconds/second in length.
diff --git a/python/lsst/ap/association/__init__.py b/python/lsst/ap/association/__init__.py
@@ -20,6 +20,7 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 from .version import *
+from .trailedSourceFilter import *
 from .association import *
 from .diaForcedSource import *
 from .loadDiaCatalogs import *
diff --git a/python/lsst/ap/association/association.py b/python/lsst/ap/association/association.py
@@ -32,6 +32,7 @@
 import lsst.pex.config as pexConfig
 import lsst.pipe.base as pipeBase
 from lsst.utils.timer import timeMethod
+from .trailedSourceFilter import TrailedSourceFilterTask
 
 # Enforce an error for unsafe column/array value setting in pandas.
 pd.options.mode.chained_assignment = 'raise'
@@ -47,10 +48,22 @@ class AssociationConfig(pexConfig.Config):
         default=1.0,
     )
 
+    trailedSourceFilter = pexConfig.ConfigurableField(
+        target=TrailedSourceFilterTask,
+        doc="Docs here"
+
+    )
+
+    doLongTrailFilter = pexConfig.Field(
+        doc="Filter artifact candidates based on morphological criteria, i.g. those that appear to "
+            "be streaks.",
+        dtype=bool,
+        default=True
+    )
+
 
 class AssociationTask(pipeBase.Task):
     """Associate DIAOSources into existing DIAObjects.
-
     This task performs the association of detected DIASources in a visit
     with the previous DIAObjects detected over time. It also creates new
     DIAObjects out of DIASources that cannot be associated with previously
@@ -60,10 +73,16 @@ class AssociationTask(pipeBase.Task):
     ConfigClass = AssociationConfig
     _DefaultName = "association"
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.config.doLongTrailFilter:
+            self.makeSubtask("trailedSourceFilter")
+
     @timeMethod
     def run(self,
             diaSources,
-            diaObjects):
+            diaObjects,
+            exposure):
         """Associate the new DiaSources with existing DiaObjects.
 
         Parameters
@@ -72,6 +91,9 @@ def run(self,
             New DIASources to be associated with existing DIAObjects.
         diaObjects : `pandas.DataFrame`
             Existing diaObjects from the Apdb.
+        exposure : `pandas.DataFrame`
+            Calibrated exposure differenced with a template image during
+            image differencing.
 
         Returns
         -------
@@ -98,7 +120,18 @@ def run(self,
                 nUpdatedDiaObjects=0,
                 nUnassociatedDiaObjects=0)
 
-        matchResult = self.associate_sources(diaObjects, diaSources)
+        if self.config.doLongTrailFilter:
+            diaTrailedResult = self.trailedSourceFilter.run(diaSources, exposure)
+
+            if len(diaTrailedResult.trailedDiaSources) > 0:
+                print("Trailed sources cleaned.")
+            else:
+                print("No trailed sources to clean.")
+
+            matchResult = self.associate_sources(diaObjects, diaTrailedResult.diaSources)
+
+        else:
+            matchResult = self.associate_sources(diaObjects, diaSources)
 
         mask = matchResult.diaSources["diaObjectId"] != 0
 
diff --git a/python/lsst/ap/association/diaPipe.py b/python/lsst/ap/association/diaPipe.py
@@ -28,6 +28,10 @@
 Currently loads directly from the Apdb rather than pre-loading.
 """
 
+__all__ = ("DiaPipelineConfig",
+           "DiaPipelineTask",
+           "DiaPipelineConnections")
+
 import pandas as pd
 
 import lsst.dax.apdb as daxApdb
@@ -41,13 +45,10 @@
     AssociationTask,
     DiaForcedSourceTask,
     LoadDiaCatalogsTask,
-    PackageAlertsTask)
+    PackageAlertsTask,
+    TrailedSourceFilterTask)
 from lsst.ap.association.ssoAssociation import SolarSystemAssociationTask
 
-__all__ = ("DiaPipelineConfig",
-           "DiaPipelineTask",
-           "DiaPipelineConnections")
-
 
 class DiaPipelineConnections(
         pipeBase.PipelineTaskConnections,
@@ -221,6 +222,10 @@ class DiaPipelineConfig(pipeBase.PipelineTaskConfig,
         target=AssociationTask,
         doc="Task used to associate DiaSources with DiaObjects.",
     )
+    trailedFilter = pexConfig.ConfigurableField(
+        target=TrailedSourceFilterTask,
+        doc="Task used to find trailed DiaSources.",
+    )
     doSolarSystemAssociation = pexConfig.Field(
         dtype=bool,
         default=False,
@@ -368,7 +373,7 @@ def run(self,
 
         # Associate new DiaSources with existing DiaObjects.
         assocResults = self.associator.run(diaSourceTable,
-                                           loaderResult.diaObjects)
+                                           loaderResult.diaObjects, exposure)
         if self.config.doSolarSystemAssociation:
             ssoAssocResult = self.solarSystemAssociator.run(
                 assocResults.unAssocDiaSources,
diff --git a/python/lsst/ap/association/trailedSourceFilter.py b/python/lsst/ap/association/trailedSourceFilter.py
@@ -0,0 +1,110 @@
+# This file is part of ap_association.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+"""A simple implementation of source association task for ap_verify.
+"""
+
+__all__ = ["TrailedSourceFilterTask", "TrailedSourceFilterConfig"]
+
+import lsst.pex.config as pexConfig
+import lsst.pipe.base as pipeBase
+from lsst.utils.timer import timeMethod
+
+# import numpy as np
+import pandas as pd
+
+# Enforce an error for unsafe column/array value setting in pandas.
+pd.options.mode.chained_assignment = 'raise'
+
+
+class TrailedSourceFilterConfig(pexConfig.Config):
+    """Config class for TrailedSourceFilterTask.
+    """
+    maxTrailLength = pexConfig.Field(
+        dtype=float,
+        doc='Maximum trail length permitted is less than 10 degrees/day. This is a rate '
+            'of 0.416 arcseconds per second.As trail length is measured in '
+            'arcseconds, it is dependant on the length of the exposure.',
+        default=0.416,  # HSC Default. Should Decam and LSST defaults be passed? Does it change possibly?
+    )
+
+
+class TrailedSourceFilterTask(pipeBase.Task):
+    """Find trailed sources in DIAObjects.
+
+    """
+
+    ConfigClass = TrailedSourceFilterConfig
+    _DefaultName = "trailedAssociation"
+
+    @timeMethod
+    def run(self,
+            dia_sources, exposure):
+        """Find trailed sources which have not been filtered out and will
+        not be included in the diaSource catalog.
+
+        Parameters
+        ----------
+        dia_sources : `pandas.DataFrame`
+            New DIASources to be checked for trailed sources.
+        exposure : `pandas.DataFrame`
+            Calibrated exposure differenced with a template image during
+            image differencing.
+
+        Returns
+        -------
+        result : `lsst.pipe.base.Struct`
+            Results struct with components.
+
+            - ``"dia_sources"`` : DiaSource table that is free from unwanted
+             trailed sources (`pandas.DataFrame`)
+
+            - ``"trailed_dia_sources"`` : DiaSources that have trailed more than
+             0.416 arcseconds/second*exposure_time(`pandas.DataFrame`)
+        """
+        trail_mask = self.check_dia_source_trail(dia_sources, exposure)
+
+        return pipeBase.Struct(
+            diaSources=dia_sources[~trail_mask].reset_index(drop=True),
+            trailedDiaSources=dia_sources[trail_mask].reset_index(drop=True))
+
+    def check_dia_source_trail(self, dia_sources, exposure):
+        """Check that all DiaSources have trails.
+
+        Creates a mask for sources with lengths greater than 0.416 arcseconds/second*exposure time.
+
+        Parameters
+        ----------
+        dia_sources : `pandas.DataFrame`
+            Input DiaSources to check for trail lengths.
+        exposure : `pandas.DataFrame`
+            Calibrated exposure differenced with a template image during
+            image differencing.
+
+        Returns
+        -------
+        trail_mask : `pandas.DataFrame`
+            Boolean mask for dia_sources which are greater than the cuttoff length.
+        """
+        exposure_time = exposure.getInfo().getVisitInfo().getExposureTime()
+        trail_mask = (dia_sources.loc[:, "trailLength"].values[:] >= self.config.maxTrailLength*exposure_time)
+
+        return trail_mask
diff --git a/tests/test_association_task.py b/tests/test_association_task.py
@@ -22,6 +22,8 @@
 import numpy as np
 import pandas as pd
 import unittest
+import lsst.afw.image as afwImage
+from lsst.afw.coord import Weather
 
 import lsst.geom as geom
 import lsst.utils.tests
@@ -46,20 +48,38 @@ def setUp(self):
         self.diaSources = pd.DataFrame(data=[
             {"ra": 0.04*idx + scatter*rng.uniform(-1, 1),
              "dec": 0.04*idx + scatter*rng.uniform(-1, 1),
-             "diaSourceId": idx + 1 + self.nObjects, "diaObjectId": 0}
+             "diaSourceId": idx + 1 + self.nObjects, "diaObjectId": 0, "trailLength": 5.5*idx}
             for idx in range(self.nSources)])
         self.diaSourceZeroScatter = pd.DataFrame(data=[
             {"ra": 0.04*idx,
              "dec": 0.04*idx,
-             "diaSourceId": idx + 1 + self.nObjects, "diaObjectId": 0}
+             "diaSourceId": idx + 1 + self.nObjects, "diaObjectId": 0, "trailLength": 5.5*idx}
             for idx in range(self.nSources)])
+        exposureId = 5
+        exposureTime = 30
+        boresightRotAngle = 45.6 * lsst.geom.degrees
+        weather = Weather(1.1, 2.2, 0.3)
+        visitInfo = afwImage.VisitInfo(
+            exposureId=exposureId,
+            exposureTime=exposureTime,
+            boresightRotAngle=boresightRotAngle,
+            weather=weather,
+        )
+        exposureInfo = afwImage.ExposureInfo()
+        exposureInfo.setVisitInfo(visitInfo)
+        maskedImage = afwImage.MaskedImageF(lsst.geom.Extent2I(64, 64))
+        self.exposure = afwImage.ExposureF(maskedImage, exposureInfo)
 
     def test_run(self):
         """Test the full task by associating a set of diaSources to
         existing diaObjects.
-        """
-        assocTask = AssociationTask()
-        results = assocTask.run(self.diaSources, self.diaObjects)
+        # """
+
+        config = AssociationTask.ConfigClass()
+        config.doLongTrailFilter = False
+        assocTask = AssociationTask(config=config)
+
+        results = assocTask.run(self.diaSources, self.diaObjects, self.exposure)
 
         self.assertEqual(results.nUpdatedDiaObjects, len(self.diaObjects) - 1)
         self.assertEqual(results.nUnassociatedDiaObjects, 1)
@@ -75,13 +95,40 @@ def test_run(self):
                 [0]):
             self.assertEqual(test_obj_id, expected_obj_id)
 
+    def test_run_trailed_sources(self):
+        """Test the full task by associating a set of diaSources to
+        existing diaObjects when trailed sources are filtered
+        # """
+
+        config = AssociationTask.ConfigClass()
+        config.doLongTrailFilter = True
+        assocTask = AssociationTask(config=config)
+
+        results = assocTask.run(self.diaSources, self.diaObjects, self.exposure)
+
+        print(results.nUpdatedDiaObjects, "results nupdated")
+        print(len(self.diaObjects) - 3, " length of dia objects")
+        self.assertEqual(results.nUpdatedDiaObjects, len(self.diaObjects) - 3)
+        self.assertEqual(results.nUnassociatedDiaObjects, 3)
+        self.assertEqual(len(results.matchedDiaSources),
+                         len(self.diaObjects) - 3)
+        self.assertEqual(len(results.unAssocDiaSources), 1)
+        for test_obj_id, expected_obj_id in zip(
+                results.matchedDiaSources["diaObjectId"].to_numpy(),
+                [1, 2, 3, 4]):
+            self.assertEqual(test_obj_id, expected_obj_id)
+        for test_obj_id, expected_obj_id in zip(
+                results.unAssocDiaSources["diaObjectId"].to_numpy(),
+                [0]):
+            self.assertEqual(test_obj_id, expected_obj_id)
+
     def test_run_no_existing_objects(self):
         """Test the run method with a completely empty database.
         """
         assocTask = AssociationTask()
         results = assocTask.run(
             self.diaSources,
-            pd.DataFrame(columns=["ra", "dec", "diaObjectId"]))
+            pd.DataFrame(columns=["ra", "dec", "diaObjectId", "trailLength"]), self.exposure)
         self.assertEqual(results.nUpdatedDiaObjects, 0)
         self.assertEqual(results.nUnassociatedDiaObjects, 0)
         self.assertEqual(len(results.matchedDiaSources), 0)
@@ -100,6 +147,9 @@ def test_associate_sources(self):
                 [0, 1, 2, 3, 4]):
             self.assertEqual(test_obj_id, expected_obj_id)
 
+    def test_trailed_source(self):
+        """Test the removal of trailed sources"""
+
     def test_score_and_match(self):
         """Test association between a set of sources and an existing
         DIAObjectCollection.
diff --git a/tests/test_diaPipe.py b/tests/test_diaPipe.py
@@ -121,7 +121,7 @@ def solarSystemAssociator_run(self, unAssocDiaSources, solarSystemObjectTable, d
                                          unAssocDiaSources=MagicMock(spec=pd.DataFrame()))
 
         @lsst.utils.timer.timeMethod
-        def associator_run(self, table, diaObjects):
+        def associator_run(self, table, diaObjects, exposure):
             return lsst.pipe.base.Struct(nUpdatedDiaObjects=2, nUnassociatedDiaObjects=3,
                                          matchedDiaSources=MagicMock(spec=pd.DataFrame()),
                                          unAssocDiaSources=MagicMock(spec=pd.DataFrame()))
