aind-behavior-utils 0.3.1__py3-none-any.whl

aind_behavior_utils/__init__.py

@@ -0,0 +1,3 @@
+"""Init package"""
+
+__version__ = "0.3.1"

aind_behavior_utils/plotting/__init__.py

@@ -0,0 +1 @@
+"""Plotting utilities for behavior data visualization."""

aind_behavior_utils/plotting/plots.py

@@ -0,0 +1,60 @@
+"""Plotting utilities for behavior data."""
+
+from typing import Optional, Tuple, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.figure import Figure
+
+
+def plot_array(
+    data_array: np.ndarray,
+    ylim: Optional[Tuple[float, float]] = None,
+    xlim: Optional[Tuple[float, float]] = None,
+    ylabel: str = "",
+    xlabel: str = "",
+    title: str = "",
+    aspect: Optional[Union[str, float]] = None,
+) -> Figure:
+    """Plot data, return the figure.
+
+    Parameters
+    ----------
+    data_array : numpy.ndarray
+        An array of data to plot.
+    ylim : Optional[Tuple[float, float]]
+        The y-axis limits as a tuple (lower, upper).
+    xlim : Optional[Tuple[float, float]]
+        The x-axis limits as a tuple (lower, upper).
+    ylabel : str
+        The y-axis label.
+    xlabel : str
+        The x-axis label.
+    title : str
+        The plot title.
+    aspect : Optional[Union[str, float]]
+        The aspect ratio of the plot.
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+        A matplotlib figure.
+    """
+    fig, ax = plt.subplots()
+
+    if ylim:
+        ax.set_ylim(*ylim)
+    if xlim:
+        ax.set_xlim(*xlim)
+
+    ax.set_ylabel(ylabel)
+    ax.set_xlabel(xlabel)
+    ax.set_title(title)
+
+    ax.plot(data_array)
+
+    if aspect:
+        ax.set_aspect(aspect)
+
+    fig.tight_layout()
+    return fig

aind_behavior_utils/stimulus/__init__.py

@@ -0,0 +1,5 @@
+"""Stimulus pickle file parsing and analysis utilities."""
+
+from aind_behavior_utils.stimulus.camstim_dataset import CamstimDataset
+
+__all__ = ["CamstimDataset"]

aind_behavior_utils/stimulus/camstim_dataset.py

@@ -0,0 +1,267 @@
+"""Stimulus pickle file parsing utilities.
+
+Provides utilities for loading and parsing Camstim stimulus pickle files,
+extracting frame timing, wheel encoder data, and quality control metrics.
+
+The primary interface is the :class:`CamstimDataset` class, which wraps
+a loaded stimulus pickle dictionary and resolves its internal structure
+(foraging vs behavior item groups) once at construction.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import numpy as np
+import pandas as pd
+
+# Wheel encoder calibration constant (radius in cm)
+WHEEL_RADIUS = 5.5036
+
+# Default fallback frame rate (Hz) when not explicitly specified or computed
+DEFAULT_FPS = 60.0
+
+# Conversion factor from milliseconds to seconds
+MS_TO_S = 0.001
+
+
+class CamstimDataset:
+    """Wrapper around a camstim stimulus pickle file dictionary.
+
+    Resolves the internal data layout (``foraging`` vs ``behavior``
+    item groups) once at construction so that downstream accessors
+    never need to repeat the lookup.
+
+    Parameters
+    ----------
+    data : Dict[str, Any]
+        The loaded stimulus pickle file dictionary.
+
+    Examples
+    --------
+    >>> dset = CamstimDataset(pkl_data)
+    >>> print(dset.fps)
+    60.0
+    >>> print(dset.stim_frame_count)
+    120
+    >>> speed = dset.running_speed_array
+    >>> artifacts = dset.get_nb_wheel_artifacts()
+    """
+
+    def __init__(self, data: Dict[str, Any]) -> None:
+        """Initialize CamstimDataset."""
+        self.data = data
+        self._items: Optional[Dict[str, Any]] = self._resolve_items()
+
+    @classmethod
+    def from_file(cls, path: str) -> CamstimDataset:
+        """Load a pickle file and return a CamstimDataset instance.
+
+        Parameters
+        ----------
+        path : str
+            The path to the pickle file.
+
+        Returns
+        -------
+        CamstimDataset
+            A new instance wrapping the loaded data.
+        """
+        with open(path, "rb") as f:
+            data = pd.read_pickle(f)
+        return cls(data)
+
+    def _resolve_items(self) -> Optional[Dict[str, Any]]:
+        """Return the inner item-group dict (foraging or behavior).
+
+        Returns
+        -------
+        Optional[Dict[str, Any]]
+            The resolved item group dictionary, or None if not found.
+        """
+        items = self.data.get("items", {})
+        for group in ("foraging", "behavior"):
+            if group in items:
+                return items[group]
+        return None
+
+    @property
+    def fps(self) -> float:
+        """Frames per second.
+
+        Reads from the top-level ``"fps"`` key when available,
+        otherwise computes from ``intervalsms``. Falls back to
+        :data:`DEFAULT_FPS`.
+
+        Returns
+        -------
+        float
+            The frames per second.
+        """
+        try:
+            return float(self.data["fps"])
+        except KeyError:
+            if self._items is not None:
+                try:
+                    mean_interval_ms = np.mean(self._items["intervalsms"])
+                    return round(1 / (mean_interval_ms * MS_TO_S), 1)
+                except KeyError:
+                    return DEFAULT_FPS
+            return DEFAULT_FPS
+
+    @property
+    def stage(self) -> Optional[str]:
+        """Stimulus stage name.
+
+        Returns
+        -------
+        Optional[str]
+            The stage name from ``data["params"]["stage"]``, or
+            ``None`` if absent.
+        """
+        try:
+            return self.data["params"]["stage"]
+        except (KeyError, TypeError):
+            return None
+
+    @property
+    def intervals_ms(self) -> list:
+        """Inter-frame intervals in milliseconds.
+
+        Returns
+        -------
+        list
+            The inter-frame intervals in milliseconds.
+
+        Raises
+        ------
+        KeyError
+            If no ``intervalsms`` data is found.
+        """
+        if self._items is not None and "intervalsms" in self._items:
+            return self._items["intervalsms"]
+        raise KeyError("Could not find intervalsms in pickle file.")
+
+    @property
+    def stim_frame_count(self) -> int:
+        """Number of stimulus frames.
+
+        The frame count is the length of the ``intervalsms`` array plus one,
+        since each interval sits between two frames.
+
+        Returns
+        -------
+        int
+            The number of stimulus frames.
+
+        Raises
+        ------
+        KeyError
+            If no ``intervalsms`` data is found.
+        """
+        return len(self.intervals_ms) + 1
+
+    @property
+    def running_speed_array(self) -> np.ndarray:
+        """Running speed in cm/s derived from wheel encoder data.
+
+        Locates the wheel encoder ``dx`` array, scales it by FPS and
+        wheel radius to compute instantaneous speed.
+
+        Returns
+        -------
+        numpy.ndarray
+            Array of running speed values in cm/s.
+
+        Raises
+        ------
+        KeyError
+            If no encoder ``dx`` data is found in the pickle file.
+        NotImplementedError
+            If the pickle file format is unrecognised.
+        """
+        speed_dtheta = self._resolve_encoder_dx()
+        return speed_dtheta * self.fps * (2 * np.pi * WHEEL_RADIUS / 360)
+
+    def _resolve_encoder_dx(self) -> np.ndarray:
+        """Locate the encoder ``dx`` array.
+
+        Searches under the resolved item group (foraging or behavior)
+        encoders, or at the top level.
+
+        Returns
+        -------
+        numpy.ndarray
+            The encoder ``dx`` array.
+
+        Raises
+        ------
+        KeyError
+            If no encoder ``dx`` data is found.
+        NotImplementedError
+            If the pickle file format is unrecognised.
+        """
+        if self._items is not None:
+            try:
+                return np.array(self._items["encoders"][0]["dx"])
+            except (KeyError, IndexError, TypeError):
+                raise KeyError(
+                    "Could not find running speed data in pickle file."
+                )
+        if "dx" in self.data:
+            return np.array(self.data["dx"])
+        raise NotImplementedError(
+            "Encountered unknown format for stimulus pickle file."
+        )
+
+    def get_nb_wheel_artifacts(self, threshold: float = 100) -> int:
+        """Count speed values exceeding threshold.
+
+        Artifacts are defined as absolute speed values that exceed
+        the given threshold, typically indicating encoder glitches or
+        physical wheel slips.
+
+        Parameters
+        ----------
+        threshold : float, optional
+            Speed threshold in cm/s. Default is 100.
+
+        Returns
+        -------
+        int
+            The number of points exceeding the threshold.
+        """
+        return int(np.sum(np.abs(self.running_speed_array) > threshold))
+
+
+def load_pkl_file(path: str) -> Dict[str, Any]:
+    """Load a stimulus pickle file and return the raw data dictionary.
+
+    Parameters
+    ----------
+    path : str
+        The path to the pickle file.
+
+    Returns
+    -------
+    Dict[str, Any]
+        The raw data dictionary from the pickle file.
+    """
+    with open(path, "rb") as f:
+        return pd.read_pickle(f)
+
+
+def get_stim_frame_count(pkl_data: Dict[str, Any]) -> int:
+    """Get stimulus frame count from a raw pickle data dictionary.
+
+    Parameters
+    ----------
+    pkl_data : Dict[str, Any]
+        The raw stimulus pickle data dictionary.
+
+    Returns
+    -------
+    int
+        The number of stimulus frames.
+    """
+    return CamstimDataset(pkl_data).stim_frame_count

aind_behavior_utils/stimulus/wheel_utils.py

@@ -0,0 +1,131 @@
+"""Wheel QC image and metric calculation utilities.
+
+This module provides high-level functions for quality control analysis of
+wheel rotation data from stimulus pickle files. These are convenience wrappers
+around the CamstimDataset class (data parsing and metrics) and the plotting
+module (visualization).
+
+The module focuses on wheel encoder data, which is typically stored in the
+stimulus pickle file under items.foraging or items.behavior. Functions here
+combine multiple core utilities to produce either visual plots or numerical
+metrics suitable for QC assessment.
+"""
+
+from typing import Any, Dict, Union
+
+
+import aind_behavior_utils.plotting.plots as plots
+from aind_behavior_utils.stimulus.camstim_dataset import CamstimDataset
+
+
+def _resolve_pkl(
+    pkl_input: Union[str, Dict[str, Any], CamstimDataset],
+) -> CamstimDataset:
+    """Resolve pickle input to a CamstimDataset instance.
+
+    Parameters
+    ----------
+    pkl_input : Union[str, Dict[str, Any], CamstimDataset]
+        A file path, an already-loaded dictionary, or an existing
+        CamstimDataset instance.
+
+    Returns
+    -------
+    CamstimDataset
+        A CamstimDataset instance.
+    """
+    if isinstance(pkl_input, CamstimDataset):
+        return pkl_input
+    if isinstance(pkl_input, str):
+        return CamstimDataset.from_file(pkl_input)
+    return CamstimDataset(pkl_input)
+
+
+def calculate_qc_images(
+    pkl_input: Union[str, Dict[str, Any], CamstimDataset],
+) -> dict:
+    """Calculate quality control images from stimulus pickle data.
+
+    This is a high-level wrapper that combines core utilities from
+    CamstimDataset and the plotting module to generate visual QC plots
+    for wheel rotation data.
+
+    Parameters
+    ----------
+    pkl_input : Union[str, Dict[str, Any], CamstimDataset]
+        Path to a stimulus pickle file, an already-loaded pickle data
+        dictionary, or an existing CamstimDataset instance.
+
+    Returns
+    -------
+    dict
+        Dictionary containing two matplotlib Figure objects:
+        - 'wheel_speed_plot': Line plot of instantaneous wheel
+          speed in cm/s across all stimulus frames.
+        - 'wheel_traveled_distance_plot': Line plot of cumulative
+          wheel traveled distance in meters.
+
+    Notes
+    -----
+    Uses the following utilities:
+    - CamstimDataset.running_speed_array: Converts encoder data to speed (cm/s)
+    - CamstimDataset.fps: Retrieves frame rate from stimulus metadata
+    - plotting.plot_array: Creates matplotlib figures
+    """
+    dset = _resolve_pkl(pkl_input)
+    running_speed_array = dset.running_speed_array
+    wheel_speed_plot = plots.plot_array(
+        running_speed_array,
+        xlabel="Frame #",
+        ylabel="Wheel Speed (cm/s)",
+        title="wheel_speed_plot",
+    )
+
+    wheel_travel_plot = plots.plot_array(
+        running_speed_array.cumsum() / (100 * dset.fps),
+        xlabel="Frame #",
+        ylabel="Traveled Distance (m)",
+        title="wheel_traveled_distance_plot",
+    )
+    return {
+        "wheel_speed_plot": wheel_speed_plot,
+        "wheel_traveled_distance_plot": wheel_travel_plot,
+    }
+
+
+def calculate_qc_metrics(
+    pkl_input: Union[str, Dict[str, Any], CamstimDataset],
+) -> dict:
+    """Calculate quality control metrics from stimulus pickle data.
+
+    This is a high-level wrapper that combines core utilities from
+    CamstimDataset to compute numerical QC metrics for wheel rotation
+    data.
+
+    Parameters
+    ----------
+    pkl_input : Union[str, Dict[str, Any], CamstimDataset]
+        Path to a stimulus pickle file, an already-loaded pickle data
+        dictionary, or an existing CamstimDataset instance.
+
+    Returns
+    -------
+    dict
+        Dictionary containing QC metrics:
+        - 'wheel_artifacts': Count of abnormal speed values
+          (absolute value exceeding 100 cm/s), which may
+          indicate encoder glitches or physical wheel slips.
+
+    Notes
+    -----
+    Uses the following utilities:
+    - CamstimDataset.running_speed_array: Converts encoder data to speed (cm/s)
+    - CamstimDataset.get_nb_wheel_artifacts: Counts speed outliers
+
+    See Also
+    --------
+    calculate_qc_images : Generate visual QC plots
+    """
+    dset = _resolve_pkl(pkl_input)
+    metrics = {"wheel_artifacts": dset.get_nb_wheel_artifacts()}
+    return metrics

aind_behavior_utils/sync/__init__.py

@@ -0,0 +1 @@
+"""Sync HDF5 dataset parsing and analysis utilities."""

aind_behavior_utils/sync/legacy_line_labels.py

@@ -0,0 +1,97 @@
+"""Utilities for resolving legacy sync line label variants.
+
+Different experimental setups use different names for the same signal
+(e.g. ``"stim_vsync"`` vs ``"vsync_stim"``).  This module maps
+canonical signal names to their known variants so that downstream code
+can refer to signals by a single, stable name.
+"""
+
+from typing import Dict, List, Optional
+
+LINE_LABEL_VARIANTS: Dict[str, List[str]] = {
+    "behavior_monitoring": [
+        "behavior_monitoring",
+        "cam1_exposure",
+        "cam1",
+        "beh_cam_frame_readout",
+    ],
+    "eye_tracking": [
+        "eye_tracking",
+        "cam2_exposure",
+        "cam2",
+        "eye_cam_frame_readout",
+    ],
+    "face_tracking": ["face_tracking", "face_cam_frame_readout"],
+    "photodiode": ["photodiode", "stim_photodiode"],
+    "physio": ["2p_vsync", "vsync_2p"],
+    "visual_stim": ["stim_vsync", "vsync_stim"],
+}
+
+
+def build_line_label_map(
+    line_labels: List[str],
+    variants: Optional[Dict[str, List[str]]] = None,
+) -> Dict[str, str]:
+    """Build a mapping from canonical names to actual file labels.
+
+    For each canonical name, the first matching variant found in
+    ``line_labels`` is used.
+
+    Parameters
+    ----------
+    line_labels : List[str]
+        Line labels present in a sync file.
+    variants : Optional[Dict[str, List[str]]]
+        Canonical-name-to-variant-list mapping.  Defaults to
+        ``LINE_LABEL_VARIANTS``.
+
+    Returns
+    -------
+    Dict[str, str]
+        Canonical name to the actual label found in ``line_labels``.
+    """
+    if variants is None:
+        variants = LINE_LABEL_VARIANTS
+    result: Dict[str, str] = {}
+    for canonical, variant_list in variants.items():
+        for variant in variant_list:
+            if variant in line_labels:
+                result[canonical] = variant
+                break
+    return result
+
+
+def resolve_line_label(
+    line: str,
+    label_map: Dict[str, str],
+    line_labels: List[str],
+) -> str:
+    """Resolve a line name to the actual label in a sync file.
+
+    Accepts either a canonical name (looked up via ``label_map``) or a
+    direct label that exists in ``line_labels``.
+
+    Parameters
+    ----------
+    line : str
+        Canonical name or direct line label.
+    label_map : Dict[str, str]
+        Mapping from canonical names to actual labels, as returned by
+        :func:`build_line_label_map`.
+    line_labels : List[str]
+        Line labels present in the sync file.
+
+    Returns
+    -------
+    str
+        The actual line label found in ``line_labels``.
+
+    Raises
+    ------
+    ValueError
+        If ``line`` cannot be resolved to a label in ``line_labels``.
+    """
+    resolved = label_map.get(line, line)
+    if resolved in line_labels:
+        return resolved
+    raise ValueError(f"'{line}' not found in line labels or label map.")