From 6a82ca476d1f3ac3a91830e0505d301eac6409fc Mon Sep 17 00:00:00 2001
From: Brianna Smart <drbsmart@uw.edu>
Date: Wed, 28 Feb 2024 16:22:48 -0800
Subject: [PATCH] Update flagmap and filter

Update trail source filter and unit tests
---
 data/association-flag-map.yaml                | 11 +++++++++-
 .../ap/association/trailedSourceFilter.py     | 21 ++++++++++++++-----
 tests/test_trailedSourceFilter.py             |  3 ++-
 3 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/data/association-flag-map.yaml b/data/association-flag-map.yaml
index 1dda5015..bb25a025 100644
--- a/data/association-flag-map.yaml
+++ b/data/association-flag-map.yaml
@@ -90,7 +90,7 @@ columns:
       doc: parent source, ignored; only valid for HsmShape
     - name: ext_trailedSources_Naive_flag_edge
       bit: 27
-      doc: source is trailed and extends off chip
+      doc: Trail beginning or end falls on edge pixels or extends off image
     - name: base_PixelFlags_flag_streak
       bit: 28
       doc: Streak in the Source footprint
@@ -109,3 +109,12 @@ columns:
     - name: base_PixelFlags_flag_injected_templateCenter
       bit: 33
       doc: Fake source template injection center in source footprint
+    - name: ext_trailedSources_Naive_flag_off_image
+      bit: 34
+      doc: Trail extends off image
+    - name: ext_trailedSources_Naive_flag_nan
+      bit: 35
+      doc: Trail start or end pixels contain nan
+    - name: ext_trailedSources_Naive_flag_suspect_long_trail
+      bit: 36
+      doc: Trail longer than three psf radii
\ No newline at end of file
diff --git a/python/lsst/ap/association/trailedSourceFilter.py b/python/lsst/ap/association/trailedSourceFilter.py
index d7721928..13c48b6b 100644
--- a/python/lsst/ap/association/trailedSourceFilter.py
+++ b/python/lsst/ap/association/trailedSourceFilter.py
@@ -96,11 +96,15 @@ def run(self, dia_sources, exposure_time):
             longTrailedDiaSources=dia_sources[trail_mask].reset_index(drop=True))
 
     def _check_dia_source_trail(self, dia_sources, exposure_time, flags):
-        """Find DiaSources that have long trails.
+        """Find DiaSources that have long trails or trails with indeterminant
+        end points.
 
         Return a mask of sources with lengths greater than
-        ``config.max_trail_length``  multiplied by the exposure time in seconds
-        or have ext_trailedSources_Naive_flag_edge set.
+        ``config.max_trail_length``  multiplied by the exposure time in
+        seconds. Additionally, set mask if
+        ``ext_trailedSources_Naive_flag_off_image`` is set or if
+        ``ext_trailedSources_Naive_flag_suspect_long_trail`` and
+        ``ext_trailedSources_Naive_flag_edge`` are both set.
 
         Parameters
         ----------
@@ -115,11 +119,18 @@ def _check_dia_source_trail(self, dia_sources, exposure_time, flags):
         -------
         trail_mask : `pandas.DataFrame`
             Boolean mask for DIASources which are greater than the
-            cutoff length and have the edge flag set.
+            cutoff length or have off_image or suspect_long_trail
+            flag set.
         """
         trail_mask = (dia_sources.loc[:, "trailLength"].values[:]
                       >= (self.config.max_trail_length*exposure_time))
 
-        trail_mask[np.where(flags['ext_trailedSources_Naive_flag_edge'])] = True
+        long_flags = flags['ext_trailedSources_Naive_flag_suspect_long_trail']
+        edge_flags = flags['ext_trailedSources_Naive_flag_edge']
+
+        trail_mask[np.where(flags['ext_trailedSources_Naive_flag_off_image'])] = True
+        for index, value in enumerate(long_flags):
+            if value and edge_flags[index]:
+                trail_mask[index] = True
 
         return trail_mask
diff --git a/tests/test_trailedSourceFilter.py b/tests/test_trailedSourceFilter.py
index e80516dd..6b2c02d6 100644
--- a/tests/test_trailedSourceFilter.py
+++ b/tests/test_trailedSourceFilter.py
@@ -57,7 +57,7 @@ def setUp(self):
              "flags": 0}
             for idx in range(self.nSources)])
 
-        self.edgeDiaSources.loc[[1, 4], 'flags'] = np.power(2, 27)
+        self.edgeDiaSources.loc[[1, 4], 'flags'] = np.power(2, 27) + np.power(2, 36)
 
     def test_run(self):
         """Run trailedSourceFilterTask with the default max distance.
@@ -116,6 +116,7 @@ def test_run_edge(self):
         """
         trailedSourceFilterTask = TrailedSourceFilterTask()
 
+        results = trailedSourceFilterTask.run(self.edgeDiaSources, self.exposure_time)
         results = trailedSourceFilterTask.run(self.edgeDiaSources, self.exposure_time)
 
         self.assertEqual(len(results.diaSources), 3)