feat(tracking): optimize clean_swaps function

Author: ttngu207

aeon/dj_pipeline/tracking.py

@@ -1,15 +1,22 @@
 """DataJoint schema for tracking data."""
 
 import gc
+from datetime import UTC, datetime, timezone
+
 import datajoint as dj
 import matplotlib.path
 import numpy as np
 import pandas as pd
-from datetime import datetime, timezone
-
 from swc.aeon.io import api as io_api
 
-from aeon.dj_pipeline import acquisition, dict_to_uuid, fetch_stream, get_schema_name, lab, streams
+from aeon.dj_pipeline import (
+    acquisition,
+    dict_to_uuid,
+    fetch_stream,
+    get_schema_name,
+    lab,
+    streams,
+)
 from aeon.dj_pipeline.utils import tracking_utils
 
 aeon_schemas = acquisition.aeon_schemas
@@ -79,14 +86,18 @@ def insert_new_params(
     ):
         """Insert a new set of parameters for a given tracking method."""
         if tracking_paramset_id is None:
-            tracking_paramset_id = (dj.U().aggr(cls, n="max(tracking_paramset_id)").fetch1("n") or 0) + 1
+            tracking_paramset_id = (
+                dj.U().aggr(cls, n="max(tracking_paramset_id)").fetch1("n") or 0
+            ) + 1
 
         param_dict = {
             "tracking_method": tracking_method,
             "tracking_paramset_id": tracking_paramset_id,
             "paramset_description": paramset_description,
             "params": params,
-            "param_set_hash": dict_to_uuid({**params, "tracking_method": tracking_method}),
+            "param_set_hash": dict_to_uuid(
+                {**params, "tracking_method": tracking_method}
+            ),
         }
         param_query = cls & {"param_set_hash": param_dict["param_set_hash"]}
 
@@ -164,7 +175,9 @@ def key_source(self):
         return (
             acquisition.Chunk
             * (
-                streams.SpinnakerVideoSource.join(streams.SpinnakerVideoSource.RemovalTime, left=True)
+                streams.SpinnakerVideoSource.join(
+                    streams.SpinnakerVideoSource.RemovalTime, left=True
+                )
                 & "spinnaker_video_source_name='CameraTop'"
             )
             * (TrackingParamSet & "tracking_paramset_id = 1")
@@ -174,17 +187,22 @@ def key_source(self):
 
     def make(self, key):
         """Ingest SLEAP tracking data for a given chunk."""
-        chunk_start, chunk_end = (acquisition.Chunk & key).fetch1("chunk_start", "chunk_end")
+        chunk_start, chunk_end = (acquisition.Chunk & key).fetch1(
+            "chunk_start", "chunk_end"
+        )
 
         data_dirs = acquisition.Experiment.get_data_directories(key)
 
-        device_name = (streams.SpinnakerVideoSource & key).fetch1("spinnaker_video_source_name")
+        device_name = (streams.SpinnakerVideoSource & key).fetch1(
+            "spinnaker_video_source_name"
+        )
 
         devices_schema = getattr(
             aeon_schemas,
-            (acquisition.Experiment.DevicesSchema & {"experiment_name": key["experiment_name"]}).fetch1(
-                "devices_schema_name"
-            ),
+            (
+                acquisition.Experiment.DevicesSchema
+                & {"experiment_name": key["experiment_name"]}
+            ).fetch1("devices_schema_name"),
         )
 
         stream_reader = getattr(devices_schema, device_name).Pose
@@ -198,17 +216,23 @@ def make(self, key):
         )
 
         if not len(pose_data):
-            raise ValueError(f"No SLEAP data found for {key['experiment_name']} - {device_name}")
+            raise ValueError(
+                f"No SLEAP data found for {key['experiment_name']} - {device_name}"
+            )
 
         # get identity names
         class_names = np.unique(pose_data.identity)
         identity_mapping = {n: i for i, n in enumerate(class_names)}
 
         # get anchor part
         # ie the body_part with the prefix "anchor_" (there should only be one)
-        anchor_part = {part for part in pose_data.part.unique() if part.startswith("anchor_")}
+        anchor_part = {
+            part for part in pose_data.part.unique() if part.startswith("anchor_")
+        }
         if len(anchor_part) != 1:
-            raise ValueError(f"Anchor part not found or multiple anchor parts found: {anchor_part}")
+            raise ValueError(
+                f"Anchor part not found or multiple anchor parts found: {anchor_part}"
+            )
         anchor_part = anchor_part.pop()
 
         # ingest parts and classes
@@ -223,10 +247,14 @@ def make(self, key):
                 if part == anchor_part:
                     identity_likelihood = part_position.identity_likelihood.values
                     if isinstance(identity_likelihood[0], dict):
-                        identity_likelihood = np.array([v[id_name] for v in identity_likelihood])
+                        identity_likelihood = np.array(
+                            [v[id_name] for v in identity_likelihood]
+                        )
 
                     # assert no duplicate timestamps
-                    if len(part_position.index.values) != len(set(part_position.index.values)):
+                    if len(part_position.index.values) != len(
+                        set(part_position.index.values)
+                    ):
                         raise ValueError(
                             f"Duplicate timestamps found for identity {id_name} and part {part}"
                             f" - this should not happen - check for chunk-duplicate .bin files"
@@ -309,7 +337,9 @@ def key_source(self):
         ks = (
             acquisition.Chunk
             * (
-                streams.SpinnakerVideoSource.join(streams.SpinnakerVideoSource.RemovalTime, left=True)
+                streams.SpinnakerVideoSource.join(
+                    streams.SpinnakerVideoSource.RemovalTime, left=True
+                )
                 & "spinnaker_video_source_name='CameraTop'"
             )
             & "chunk_start >= spinnaker_video_source_install_time"
@@ -319,17 +349,22 @@ def key_source(self):
 
     def make(self, key):
         """Ingest blob position data for a given chunk."""
-        chunk_start, chunk_end = (acquisition.Chunk & key).fetch1("chunk_start", "chunk_end")
+        chunk_start, chunk_end = (acquisition.Chunk & key).fetch1(
+            "chunk_start", "chunk_end"
+        )
 
         data_dirs = acquisition.Experiment.get_data_directories(key)
 
-        device_name = (streams.SpinnakerVideoSource & key).fetch1("spinnaker_video_source_name")
+        device_name = (streams.SpinnakerVideoSource & key).fetch1(
+            "spinnaker_video_source_name"
+        )
 
         devices_schema = getattr(
             aeon_schemas,
-            (acquisition.Experiment.DevicesSchema & {"experiment_name": key["experiment_name"]}).fetch1(
-                "devices_schema_name"
-            ),
+            (
+                acquisition.Experiment.DevicesSchema
+                & {"experiment_name": key["experiment_name"]}
+            ).fetch1("devices_schema_name"),
         )
 
         stream_reader = devices_schema.CameraTop.Position
@@ -342,7 +377,9 @@ def make(self, key):
         )
 
         if not len(positiondata):
-            raise ValueError(f"No Blob position data found for {key['experiment_name']} - {device_name}")
+            raise ValueError(
+                f"No Blob position data found for {key['experiment_name']} - {device_name}"
+            )
 
         # replace id=NaN with -1
         positiondata.fillna({"id": -1}, inplace=True)
@@ -358,7 +395,9 @@ def make(self, key):
             & f'chunk_start <= "{chunk_start}"'
         )[:chunk_end]
         subject_visits_df = subject_visits_df[subject_visits_df.region == "Environment"]
-        subject_visits_df = subject_visits_df[~subject_visits_df.id.str.contains("Test", case=False)]
+        subject_visits_df = subject_visits_df[
+            ~subject_visits_df.id.str.contains("Test", case=False)
+        ]
         subject_names = []
         for subject_name in set(subject_visits_df.id):
             _df = subject_visits_df[subject_visits_df.id == subject_name]
@@ -416,26 +455,38 @@ class Subject(dj.Part):
         subject_name: varchar(32)
         ---
         sample_count: int      # number of data points acquired from this stream for a given chunk
+        subject_likelihood: longblob  # likelihood of the subject being identified correctly
         x:          longblob
         y:          longblob
         timestamps: longblob
-        likelihood: longblob
+        likelihood: longblob  # likelihood of the positions (x,y) being identified correctly
         """
 
-    key_source = SLEAPTracking & "experiment_name in ('social0.2-aeon3', 'social0.2-aeon4')"
+    key_source = (
+        SLEAPTracking & "experiment_name in ('social0.2-aeon3', 'social0.2-aeon4')"
+    )
 
     def make(self, key):
-        execution_time = datetime.now(timezone.utc)
+        """Processing of SLEAPTracking data to denoise and clean identity swaps."""
+        execution_time = datetime.now(UTC)
 
-        query = (SLEAPTracking.PoseIdentity.proj("identity_name")
-                 * SLEAPTracking.AnchorPart
-                 & key)
+        query = (
+            SLEAPTracking.PoseIdentity.proj("identity_name", "identity_likelihood")
+            * SLEAPTracking.AnchorPart
+            & key
+        )
         df = fetch_stream(query)
 
         subject_names = df.identity_name.unique()
 
         if len(subject_names) > 1:
-            df_clean = tracking_utils.clean_swaps(df)
+            # Get arena bounds from database
+            active_region_query = acquisition.EpochConfig.ActiveRegion & (
+                acquisition.Chunk & key
+            )
+            df_clean = tracking_utils.clean_swaps(
+                df, region_df=active_region_query.fetch(format="frame")
+            )
         else:
             df_clean = df
 
@@ -445,22 +496,27 @@ def make(self, key):
             if subj_df.empty:
                 continue
 
-            entries.append({
+            entries.append(
+                {
+                    **key,
+                    "subject_name": subj_name,
+                    "sample_count": len(subj_df.index.values),
+                    "subject_likelihood": subj_df.identity_likelihood.values,
+                    "x": subj_df.x.values,
+                    "y": subj_df.y.values,
+                    "timestamps": subj_df.index.values,
+                    "likelihood": subj_df.likelihood.values,
+                }
+            )
+
+        exec_dur = (datetime.now(UTC) - execution_time).total_seconds() / 3600
+        self.insert1(
+            {
                 **key,
-                "subject_name": subj_name,
-                "sample_count": len(subj_df.index.values),
-                "x": subj_df.x.values,
-                "y": subj_df.y.values,
-                "timestamps": subj_df.index.values,
-                "likelihood": subj_df.likelihood.values,
-            })
-
-        exec_dur = (datetime.now(timezone.utc) - execution_time).total_seconds() / 3600
-        self.insert1({
-            **key,
-            "execution_time": execution_time,
-            "execution_duration": exec_dur,
-        })
+                "execution_time": execution_time,
+                "execution_duration": exec_dur,
+            }
+        )
         self.Subject.insert(entries)
 
 
@@ -541,18 +597,24 @@ def _get_position(
     start_query = table & obj_restriction & start_restriction
     end_query = table & obj_restriction & end_restriction
     if not (start_query and end_query):
-        raise ValueError(f"No position data found for {object_name} between {start} and {end}")
+        raise ValueError(
+            f"No position data found for {object_name} between {start} and {end}"
+        )
 
     time_restriction = (
         f'{start_attr} >= "{min(start_query.fetch(start_attr))}"'
         f' AND {start_attr} < "{max(end_query.fetch(end_attr))}"'
     )
 
     # subject's position data in the time slice
-    fetched_data = (table & obj_restriction & time_restriction).fetch(*fetch_attrs, order_by=start_attr)
+    fetched_data = (table & obj_restriction & time_restriction).fetch(
+        *fetch_attrs, order_by=start_attr
+    )
 
     if not len(fetched_data[0]):
-        raise ValueError(f"No position data found for {object_name} between {start} and {end}")
+        raise ValueError(
+            f"No position data found for {object_name} between {start} and {end}"
+        )
 
     timestamp_attr = next(attr for attr in fetch_attrs if "timestamps" in attr)