update recording filtering to utilize sample frame id (#513)

Author: BryonLewis

bats_ai/core/utils/grts_utils.py

@@ -1,5 +1,17 @@
 from __future__ import annotations
 
+from django.db.models import Case, F, IntegerField, Value, When
+
+
+def recording_effective_sample_frame_id_case() -> Case:
+    """Return a ``Case`` matching ``normalize_sample_frame_id`` for ``Recording`` rows."""
+    return Case(
+        When(sample_frame_id__isnull=True, then=Value(14)),
+        When(sample_frame_id=19, then=Value(20)),
+        default=F("sample_frame_id"),
+        output_field=IntegerField(),
+    )
+
 
 def normalize_sample_frame_id(sample_frame_id: int | None) -> int | None:
     """Normalize sample frame IDs for AKCAN compatibility.
@@ -12,4 +24,6 @@ def normalize_sample_frame_id(sample_frame_id: int | None) -> int | None:
     """
     if sample_frame_id == 19:
         return 20
+    if sample_frame_id is None:
+        return 14
     return sample_frame_id

bats_ai/core/views/recording.py

@@ -19,7 +19,6 @@
 from bats_ai.core.models import (
     Annotations,
     CompressedSpectrogram,
-    GRTSCells,
     PulseMetadata,
     Recording,
     RecordingAnnotation,
@@ -29,7 +28,7 @@
     Spectrogram,
 )
 from bats_ai.core.tasks.tasks import recording_compute_spectrogram
-from bats_ai.core.views.recording_location import _parse_bbox
+from bats_ai.core.views.recording_location import _parse_bbox, filter_recordings_by_map_bbox
 from bats_ai.core.views.species import SpeciesSchema
 
 if TYPE_CHECKING:
@@ -525,13 +524,7 @@ def get_recordings(  # noqa: C901
     if q.bbox and q.bbox.strip():
         min_lon, min_lat, max_lon, max_lat = _parse_bbox(q.bbox)
         bbox_poly = Polygon.from_bbox((min_lon, min_lat, max_lon, max_lat))
-        # Need to check the GRTSCells centroids as well as the recording_location
-        grts_cell_ids = GRTSCells.objects.filter(centroid_4326__intersects=bbox_poly).values_list(
-            "grts_cell_id", flat=True
-        )
-        queryset = queryset.filter(
-            Q(recording_location__intersects=bbox_poly) | Q(grts_cell_id__in=grts_cell_ids)
-        )
+        queryset = filter_recordings_by_map_bbox(queryset, bbox_poly)
 
     sort_field = q.sort_by or "created"
     order_prefix = "" if q.sort_direction == "asc" else "-"
@@ -610,13 +603,7 @@ def apply_filters_and_sort(qs: QuerySet[Recording]) -> QuerySet[Recording]:
         if bbox and bbox.strip():
             min_lon, min_lat, max_lon, max_lat = _parse_bbox(bbox)
             bbox_poly = Polygon.from_bbox((min_lon, min_lat, max_lon, max_lat))
-            # Need to check the GRTSCells centroids as well as the recording_location
-            grts_cell_ids = GRTSCells.objects.filter(
-                centroid_4326__intersects=bbox_poly
-            ).values_list("grts_cell_id", flat=True)
-            qs = qs.filter(
-                Q(recording_location__intersects=bbox_poly) | Q(grts_cell_id__in=grts_cell_ids)
-            )
+            qs = filter_recordings_by_map_bbox(qs, bbox_poly)
         order_prefix = "" if sort_direction == "asc" else "-"
         if sort_by == "owner_username":
             qs = qs.order_by(f"{order_prefix}owner__username")

bats_ai/core/views/recording_location.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Any, Literal
 
 from django.contrib.gis.geos import Polygon
-from django.db.models import Q, QuerySet
+from django.db.models import Exists, OuterRef, Q, QuerySet
 from ninja import Query, Router, Schema
 from ninja.errors import HttpError
 
@@ -14,7 +14,10 @@
     Recording,
     RecordingAnnotation,
 )
-from bats_ai.core.utils.grts_utils import normalize_sample_frame_id
+from bats_ai.core.utils.grts_utils import (
+    normalize_sample_frame_id,
+    recording_effective_sample_frame_id_case,
+)
 
 if TYPE_CHECKING:
     from django.http import HttpRequest
@@ -60,14 +63,33 @@ def _split_tags(tags: str | None) -> list[str]:
     return [t.strip() for t in tags.split(",") if t.strip()]
 
 
-def _apply_recording_filters_and_sort(  # noqa: PLR0913
+def filter_recordings_by_map_bbox(
+    qs: QuerySet[Recording],
+    bbox_poly: Polygon,
+) -> QuerySet[Recording]:
+    """Keep recordings whose point lies in the bbox or whose GRTS cell centroid matches.
+
+    Cell matching uses ``(grts_cell_id, sample_frame_id)`` on ``GRTSCells``, with the
+    same effective sample frame rules as ``normalize_sample_frame_id``.
+    """
+    cell_centroid_in_bbox = GRTSCells.objects.filter(
+        centroid_4326__intersects=bbox_poly,
+        grts_cell_id=OuterRef("grts_cell_id"),
+        sample_frame_id=OuterRef("_map_bbox_effective_sf"),
+    )
+    return qs.annotate(_map_bbox_effective_sf=recording_effective_sample_frame_id_case()).filter(
+        Q(recording_location__intersects=bbox_poly)
+        | (Q(grts_cell_id__isnull=False) & Exists(cell_centroid_in_bbox))
+    )
+
+
+def _apply_recording_filters_and_sort(
     *,
     qs: QuerySet[Recording],
     exclude_submitted: bool,
     submitted_by_user: QuerySet[int] | None,
     tags: str | None,
     bbox_poly: Polygon | None,
-    grts_cell_ids: QuerySet[int] | None,
 ) -> QuerySet[Recording]:
     if exclude_submitted and submitted_by_user is not None:
         qs = qs.exclude(pk__in=submitted_by_user)
@@ -78,10 +100,8 @@ def _apply_recording_filters_and_sort(  # noqa: PLR0913
             qs = qs.filter(tags__text=tag)
         qs = qs.distinct()
 
-    if bbox_poly is not None and grts_cell_ids is not None:
-        qs = qs.filter(
-            Q(recording_location__intersects=bbox_poly) | Q(grts_cell_id__in=grts_cell_ids)
-        )
+    if bbox_poly is not None:
+        qs = filter_recordings_by_map_bbox(qs, bbox_poly)
 
     # Keep deterministic ordering even though we don't expose sorting params.
     return qs.order_by("-created")
@@ -177,28 +197,22 @@ def get_recording_locations(
 
     bbox = _parse_bbox(q.bbox)
     bbox_poly: Polygon | None = None
-    grts_cell_ids: QuerySet[int] | None = None
     if bbox is not None:
         bbox_poly = Polygon.from_bbox((bbox[0], bbox[1], bbox[2], bbox[3]))
-        grts_cell_ids = GRTSCells.objects.filter(centroid_4326__intersects=bbox_poly).values_list(
-            "grts_cell_id", flat=True
-        )
 
     my_qs = _apply_recording_filters_and_sort(
         qs=my_qs,
         exclude_submitted=exclude_submitted,
         submitted_by_user=submitted_by_user,
         tags=q.tags,
         bbox_poly=bbox_poly,
-        grts_cell_ids=grts_cell_ids,
     )
     shared_qs = _apply_recording_filters_and_sort(
         qs=shared_qs,
         exclude_submitted=exclude_submitted,
         submitted_by_user=submitted_by_user,
         tags=q.tags,
         bbox_poly=bbox_poly,
-        grts_cell_ids=grts_cell_ids,
     )
 
     my_list = list(
@@ -226,7 +240,7 @@ def get_recording_locations(
     sample_frame_cell_id_pairs = {
         (normalize_sample_frame_id(r.sample_frame_id), r.grts_cell_id)
         for r in recordings
-        if r.sample_frame_id is not None and r.grts_cell_id is not None
+        if r.grts_cell_id is not None
     }
     cell_centroids_by_sample_frame_id = _precompute_grts_cell_centroids(sample_frame_cell_id_pairs)