Add ibl annotation property

Author: bjhardcastle

src/aind_session/extensions/ecephys.py

@@ -2,11 +2,12 @@
 
 import concurrent.futures
 import contextlib
+import datetime
 import functools
 import itertools
 import json
 import logging
-from collections.abc import Iterator
+from collections.abc import Iterator, Mapping
 from typing import Any, ClassVar, Literal
 
 import codeocean.computation
@@ -16,6 +17,7 @@
 
 import aind_session.extension
 import aind_session.utils.codeocean_utils
+import aind_session.utils.docdb_utils
 
 logger = logging.getLogger(__name__)
 
@@ -100,6 +102,232 @@ class EcephysExtension(aind_session.extension.ExtensionBaseClass):
 
     DEFAULT_SORTING_PIPELINE_ID: ClassVar[str] = "1f8f159a-7670-47a9-baf1-078905fc9c2e"
     DEFAULT_TRIGGER_CAPSULE_ID: ClassVar[str] = "eb5a26e4-a391-4d79-9da5-1ab65b71253f"
+    IBL_ALIGNMENT_EVALUATION_PREFIX: ClassVar[str] = "Probe Alignment for"
+
+    @property
+    def latest_ibl_annotations(self) -> dict[str, dict[str, Any]]:
+        """Latest IBL probe-alignment annotations in DocDB, keyed by probe name.
+
+        The annotations are stored as QC evaluations on the latest derived ecephys
+        DocDB asset, not as separate DocDB records per probe.
+
+        For a given channel record with xyz coordinates, convert to ccf with:
+            ccf_ap: y * 1000
+            ccf_ml: x * -1000
+            ccf_dv: z * -1000
+
+        Examples
+        --------
+        >>> session = aind_session.Session('ecephys_795555_2025-08-26_11-29-20')
+        >>> annotations = session.ecephys.latest_ibl_annotations
+        >>> sorted(annotations)
+        ['ProbeA_0', 'ProbeB_0', 'ProbeC_0', 'ProbeD_0', 'ProbeE_0']
+        >>> annotations['ProbeA_0']['ccf_channel_results']['channel_0']
+        {'x': -4.172946453094482, 'y': 3.727851629257202, 'z': -5.629019260406494, 'axial': 0.0, 'lateral': 16.0, 'brain_region_id': 698, 'brain_region': 'OLF', 'channel_number': 0, 'ccf_ap': 3727.851629257202, 'ccf_ml': 4172.946453094482, 'ccf_dv': 5629.019260406494}
+        """
+        return EcephysExtension.get_latest_ibl_annotations(self._base.id)
+
+    @staticmethod
+    def get_latest_ibl_annotations(
+        session_id: str,
+    ) -> dict[str, dict[str, Any]]:
+        """Return the latest IBL probe-alignment annotations for a session.
+
+        Results are keyed by probe name, such as ``ProbeA_0``. Each value includes
+        asset metadata, evaluation metadata, and the parsed curation payload written
+        by the IBL ephys alignment GUI.
+        """
+        docdb_api_client = aind_session.utils.docdb_utils.get_docdb_api_client()
+
+        records = docdb_api_client.retrieve_docdb_records(
+            filter_query= {
+                "data_description.data_level": "derived",
+                "data_description.name": {"$regex": session_id},
+                "data_description.modality.abbreviation": "ecephys",
+                "quality_control.evaluations.name": {
+                    "$regex": f"{EcephysExtension.IBL_ALIGNMENT_EVALUATION_PREFIX} .*{session_id}"
+                },
+            },
+            projection={
+                "_id": 1,
+                "name": 1,
+                "created": 1,
+                "last_modified": 1,
+                "data_description.name": 1,
+                "quality_control.evaluations.name": 1,
+                "quality_control.evaluations.created": 1,
+                "quality_control.evaluations.latest_status": 1,
+                "quality_control.evaluations.metrics.name": 1,
+                "quality_control.evaluations.metrics.status_history": 1,
+                "quality_control.evaluations.metrics.value.curation_history": 1,
+                "quality_control.evaluations.metrics.value.curations": 1,
+            },
+            sort={"created": 1},
+        )
+        if not records:
+            raise KeyError(
+                f"No IBL probe-alignment evaluations found in DocDB for {session_id!r}"
+            )
+
+        latest_record = records[-1]
+        latest_by_probe: dict[str, dict[str, Any]] = {}
+        for evaluation in EcephysExtension._iter_ibl_alignment_evaluations(
+            latest_record, session_id
+        ):
+            probe = EcephysExtension._get_probe_name_from_alignment_evaluation_name(
+                evaluation.get("name", ""), session_id
+            )
+            if probe is None:
+                continue
+            previous = latest_by_probe.get(probe)
+            if previous is not None and EcephysExtension._parse_docdb_timestamp(
+                previous["created"]
+            ) >= EcephysExtension._parse_docdb_timestamp(evaluation.get("created")):
+                continue
+            latest_by_probe[probe] = (
+                EcephysExtension._format_ibl_alignment_annotation(
+                    record=latest_record,
+                    evaluation=evaluation,
+                )
+            )
+
+        if not latest_by_probe:
+            raise KeyError(
+                f"No IBL probe-alignment evaluations found in latest DocDB asset for {session_id!r}"
+            )
+
+        return dict(sorted(latest_by_probe.items()))
+
+    @staticmethod
+    def _iter_ibl_alignment_evaluations(
+        record: dict[str, Any], session_id: str
+    ) -> Iterator[dict[str, Any]]:
+        evaluations = (record.get("quality_control") or {}).get("evaluations") or []
+        for evaluation in evaluations:
+            name = evaluation.get("name", "")
+            if (
+                isinstance(name, str)
+                and name.startswith(EcephysExtension.IBL_ALIGNMENT_EVALUATION_PREFIX)
+                and session_id in name
+            ):
+                yield evaluation
+
+    @staticmethod
+    def _get_probe_name_from_alignment_evaluation_name(
+        evaluation_name: str, session_id: str
+    ) -> str | None:
+        marker = f"{session_id}_"
+        if marker not in evaluation_name:
+            return None
+        return evaluation_name.rsplit(marker, maxsplit=1)[-1]
+
+    @staticmethod
+    def _format_ibl_alignment_annotation(
+        record: dict[str, Any], evaluation: dict[str, Any]
+    ) -> dict[str, Any]:
+        metric = next(iter(evaluation.get("metrics") or []), {})
+        value = metric.get("value") or {}
+        curation_history = value.get("curation_history") or []
+        latest_curation = curation_history[-1] if curation_history else {}
+        status_history = metric.get("status_history") or []
+        latest_status = status_history[-1] if status_history else {}
+        curation = EcephysExtension._parse_latest_curation(value)
+        channel_results = curation.get("channel_results")
+
+        return {
+            "asset_id": record.get("_id"),
+            "asset_name": record.get("name"),
+            "asset_data_description_name": (record.get("data_description") or {}).get(
+                "name"
+            ),
+            "asset_created": record.get("created"),
+            "asset_last_modified": record.get("last_modified"),
+            "name": evaluation.get("name"),
+            "created": evaluation.get("created"),
+            "latest_status": evaluation.get("latest_status"),
+            "curator": latest_curation.get("curator") or latest_status.get("evaluator"),
+            "curation_timestamp": latest_curation.get("timestamp"),
+            "channel_results": channel_results,
+            "previous_alignments": curation.get("previous_alignments"),
+            "ccf_channel_results": (
+                EcephysExtension._parse_ibl_annotation_channel_records(channel_results)
+                if isinstance(channel_results, Mapping)
+                else curation.get("ccf_channel_results")
+            ),
+        }
+
+    @staticmethod
+    def _parse_ibl_annotation_channel_records(
+        channel_results: Mapping[str, Any],
+    ) -> dict[str, dict[str, Any]]:
+        return {
+            str(channel_name): EcephysExtension.parse_ibl_annotation_channel_record(
+                str(channel_name), channel_record
+            )
+            for channel_name, channel_record in channel_results.items()
+            if isinstance(channel_record, Mapping)
+        }
+
+    @staticmethod
+    def parse_ibl_annotation_channel_record(
+        channel_name: str, channel_record: Mapping[str, Any]
+    ) -> dict[str, Any]:
+        """Return an IBL annotation channel record with channel number and CCF coords.
+
+        ``x``, ``y``, and ``z`` values are stored in mm by the IBL ephys alignment
+        GUI. The returned CCF coordinates are in microns:
+        ``ccf_ap = y * 1000``, ``ccf_ml = x * -1000``, and
+        ``ccf_dv = z * -1000``.
+        """
+        parsed = dict(channel_record)
+        with contextlib.suppress(ValueError):
+            parsed["channel_number"] = int(channel_name.rsplit("_", maxsplit=1)[-1])
+
+        coordinates = {
+            key: value
+            for key, value in (
+                ("x", parsed.get("x")),
+                ("y", parsed.get("y")),
+                ("z", parsed.get("z")),
+            )
+            if isinstance(value, (int, float))
+        }
+        # sometimes values are stored as mm, sometimes microns - need to detect:
+        scale = 1000.0 if all(abs(v) < 10 for v in coordinates.values()) else 1.0
+        if len(coordinates) == 3:
+            parsed["ccf_ap"] = coordinates["y"] * scale
+            parsed["ccf_ml"] = coordinates["x"] * -scale
+            parsed["ccf_dv"] = coordinates["z"] * -scale
+        return parsed
+
+    @staticmethod
+    def _parse_latest_curation(value: dict[str, Any]) -> dict[str, Any]:
+        curations = value.get("curations") or []
+        if not curations:
+            return {}
+        try:
+            curation = json.loads(curations[-1])
+        except json.JSONDecodeError:
+            return {}
+        if not isinstance(curation, dict):
+            return {}
+        return curation
+
+    @staticmethod
+    def _parse_docdb_timestamp(value: object) -> datetime.datetime:
+        if value is None:
+            return datetime.datetime.min.replace(tzinfo=datetime.timezone.utc)
+        if isinstance(value, datetime.datetime):
+            return value if value.tzinfo else value.replace(tzinfo=datetime.timezone.utc)
+
+        text = str(value).strip().replace(" ", "T")
+        if text.endswith("Z"):
+            text = f"{text[:-1]}+00:00"
+        try:
+            timestamp = datetime.datetime.fromisoformat(text)
+        except ValueError:
+            return datetime.datetime.min.replace(tzinfo=datetime.timezone.utc)
+        return timestamp if timestamp.tzinfo else timestamp.replace(tzinfo=datetime.timezone.utc)
 
     @property
     def clipped_dir(self) -> upath.UPath: