load probe lfp metadata without needing to load lfp

allensdk/brain_observatory/behavior/behavior_project_cache/project_apis/data_io/behavior_neuropixels_project_cloud_api.py

@@ -9,6 +9,7 @@
 from allensdk.brain_observatory.ecephys.behavior_ecephys_session import (
     BehaviorEcephysSession,
 )
+from allensdk.brain_observatory.ecephys._probe import ProbeWithLFPMeta
 
 
 class VisualBehaviorNeuropixelsProjectCloudApi(ProjectCloudApiBase):
@@ -125,18 +126,24 @@ def f():
             return f
 
         # Backwards compatibility check for VBN data that doesn't contain
-        # the LFP, probes dataset.
-        if not probes_meta.empty and "file_id" in probes_meta.columns:
-            probe_data_path_map = {
-                p.name: make_lazy_load_filepath_function(
-                    file_id=str(int(getattr(p, self.cache.file_id_column)))
+        # the LFP dataset.
+        has_probe_file = self.cache.file_id_column in probes_meta.columns
+
+        if not probes_meta.empty and has_probe_file:
+            probe_meta = {
+                p.name: ProbeWithLFPMeta(
+                    lfp_csd_filepath=make_lazy_load_filepath_function(
+                        file_id=str(int(getattr(
+                            p, self.cache.file_id_column)))
+                        ),
+                    lfp_sampling_rate=p.lfp_sampling_rate
                 )
                 for p in probes_meta.itertuples(index=False)
             }
         else:
-            probe_data_path_map = None
+            probe_meta = None
         return BehaviorEcephysSession.from_nwb_path(
-            str(session_data_path), probe_data_path_map=probe_data_path_map
+            str(session_data_path), probe_meta=probe_meta
         )
 
     def _get_ecephys_session_table(self):

allensdk/brain_observatory/ecephys/_probe.py

@@ -1,5 +1,7 @@
+import dataclasses
 import logging
 from datetime import datetime
+from pathlib import Path
 from typing import Optional, Union, Callable, Tuple
 
 import numpy as np
@@ -22,6 +24,26 @@
     NwbWritableInterface, NwbReadableInterface
 
 
+@dataclasses.dataclass
+class ProbeWithLFPMeta:
+    """
+    Metadata for a single probe which has LFP data associated with it
+
+    Attributes:
+
+    - lfp_csd_filepath --> Either a path to the NWB file containing the LFP 
+    and CSD data or a callable which returns it. The nwb file is loaded 
+    separately from the main session nwb file in order to load the LFP data 
+    on the fly rather than with the main session NWB file. This is to speed 
+    up download of the NWB for users who don't wish to load the LFP data (it 
+    is large).
+    - lfp_sampling_rate --> LFP sampling rate
+    """  # noqa E402
+
+    lfp_csd_filepath: Union[Path, Callable[[], Path]]
+    lfp_sampling_rate: float
+
+
 class Probe(DataObject, JsonReadableInterface, NwbWritableInterface,
             NwbReadableInterface):
     """A single probe"""
@@ -33,7 +55,7 @@ def __init__(
             units: Units,
             sampling_rate: float = 30000.0,
             lfp: Optional[LFP] = None,
-            probe_nwb_path: Optional[Union[str, Callable[[], str]]] = None,
+            lfp_meta: Optional[ProbeWithLFPMeta] = None,
             current_source_density: Optional[CurrentSourceDensity] = None,
             location: str = 'See electrode locations',
             temporal_subsampling_factor: Optional[float] = 2.0
@@ -54,11 +76,8 @@ def __init__(
             probe sampling rate
         lfp:
             probe LFP
-        probe_nwb_path
-            The file at this path should contain LFP, CSD data for this probe
-            Can be one of the following:
-                Path to the probe NWB file
-                Callable that returns path to the probe NWB file
+        lfp_meta
+            `ProbeWithLFPMeta`
         current_source_density
             probe current source density
         location:
@@ -72,7 +91,7 @@ def __init__(
         self._units = units
         self._sampling_rate = sampling_rate
         self._lfp = lfp
-        self._probe_nwb_path = probe_nwb_path
+        self._lfp_meta = lfp_meta
         self._current_source_density = current_source_density
         self._location = location
         self._temporal_subsampling_factor = temporal_subsampling_factor
@@ -103,7 +122,7 @@ def sampling_rate(self) -> float:
     @property
     def lfp(self) -> Optional[DataArray]:
         if self._lfp is None:
-            if self._probe_nwb_path is None:
+            if self._lfp_meta is None:
                 return None
             lfp = self._read_lfp_from_nwb()
             self._lfp = lfp
@@ -114,7 +133,7 @@ def lfp(self) -> Optional[DataArray]:
     @property
     def current_source_density(self) -> Optional[DataArray]:
         if self._current_source_density is None:
-            if self._probe_nwb_path is None:
+            if self._lfp_meta is None:
                 return None
             csd = self._read_csd_data_from_nwb()
             self._current_source_density = csd
@@ -169,7 +188,7 @@ def from_nwb(
             cls,
             nwbfile: NWBFile,
             probe_name: str,
-            probe_nwb_path: Optional[str] = None,
+            lfp_meta: Optional[ProbeWithLFPMeta] = None
     ) -> "Probe":
         """
 
@@ -178,9 +197,8 @@ def from_nwb(
         nwbfile
         probe_name
             Probe name
-        probe_nwb_path
-            Path to load probe NWB file, which should contain LFP and CSD data,
-            if LFP data exists
+        lfp_meta
+            `ProbeWithLFPMeta`
 
         Returns
         -------
@@ -196,7 +214,7 @@ def from_nwb(
             sampling_rate=probe.device.sampling_rate,
             channels=channels,
             units=units,
-            probe_nwb_path=probe_nwb_path
+            lfp_meta=lfp_meta
         )
 
     def to_nwb(
@@ -354,11 +372,11 @@ def _add_csd_to_nwb(
         return nwbfile
 
     def _read_lfp_from_nwb(self) -> LFP:
-        if isinstance(self._probe_nwb_path, Callable):
+        if isinstance(self._lfp_meta.lfp_csd_filepath, Callable):
             logging.info('Fetching LFP NWB file')
-            path = self._probe_nwb_path()
+            path = self._lfp_meta.lfp_csd_filepath()
         else:
-            path = self._probe_nwb_path
+            path = self._lfp_meta.lfp_csd_filepath
         with pynwb.NWBHDF5IO(path, 'r', load_namespaces=True) as f:
             nwbfile = f.read()
             probe = nwbfile.electrode_groups[self._name]
@@ -378,11 +396,11 @@ def _read_lfp_from_nwb(self) -> LFP:
         )
 
     def _read_csd_data_from_nwb(self) -> CurrentSourceDensity:
-        if isinstance(self._probe_nwb_path, Callable):
+        if isinstance(self._lfp_meta.lfp_csd_filepath, Callable):
             logging.info('Fetching LFP NWB file')
-            path = self._probe_nwb_path()
+            path = self._lfp_meta.lfp_csd_filepath()
         else:
-            path = self._probe_nwb_path
+            path = self._lfp_meta.lfp_csd_filepath
         with pynwb.NWBHDF5IO(path, 'r', load_namespaces=True) as f:
             nwbfile = f.read()
             csd_mod = nwbfile.get_processing_module(
@@ -404,12 +422,21 @@ def _read_csd_data_from_nwb(self) -> CurrentSourceDensity:
             )
 
     def to_dict(self) -> dict:
+        has_lfp_data = False
+        lfp_sampling_rate = None
+
+        if self._lfp is not None:
+            lfp_sampling_rate = self._lfp.sampling_rate
+            has_lfp_data = True
+        elif self._lfp_meta is not None:
+            lfp_sampling_rate = self._lfp_meta.lfp_sampling_rate
+            has_lfp_data = True
+
         return {
             'id': self._id,
             'name': self._name,
             'location': self._location,
             'sampling_rate': self._sampling_rate,
-            'lfp_sampling_rate':
-                self._lfp.sampling_rate if self._lfp is not None else None,
-            'has_lfp_data': self._lfp is not None
+            'lfp_sampling_rate': lfp_sampling_rate,
+            'has_lfp_data': has_lfp_data
         }

allensdk/brain_observatory/ecephys/behavior_ecephys_session.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import pandas as pd
@@ -33,6 +33,7 @@
 from allensdk.brain_observatory.ecephys.data_objects.trials import VBNTrials
 from allensdk.brain_observatory.ecephys.optotagging import OptotaggingTable
 from allensdk.brain_observatory.ecephys.probes import Probes
+from allensdk.brain_observatory.ecephys._probe import ProbeWithLFPMeta
 from pynwb import NWBFile
 from xarray import DataArray
 
@@ -498,7 +499,8 @@ def from_json(
             ),
         )
         probes = Probes.from_json(
-            probes=session_data["probes"], skip_probes=skip_probes
+            probes=session_data["probes"],
+            skip_probes=skip_probes
         )
         optotagging_table = OptotaggingTable.from_json(dict_repr=session_data)
 
@@ -535,8 +537,8 @@ def to_nwb(self) -> Tuple[NWBFile, Dict[str, Optional[NWBFile]]]:
     def from_nwb(
         cls,
         nwbfile: NWBFile,
-        probe_data_path_map: Optional[
-            Dict[str, Union[str, Callable[[], str]]]
+        probe_meta: Optional[
+            Dict[str, ProbeWithLFPMeta]
         ] = None,
         **kwargs,
     ) -> "BehaviorEcephysSession":
@@ -545,14 +547,8 @@ def from_nwb(
         Parameters
         ----------
         nwbfile
-        probe_data_path_map
-            Maps the probe name to the path to the probe nwb file, or a
-            callable that returns the nwb path. This file should contain
-            LFP and CSD data. The nwb file is loaded
-            separately from the main session nwb file in order to load the LFP
-            data on the fly rather than with the main
-            session NWB file. This is to speed up download of the NWB
-            for users who don't wish to load the LFP data (it is large).
+        probe_meta
+            Maps the probe name to the `ProbeWithLFPMeta`
         kwargs: kwargs sent to `BehaviorSession.from_nwb`
 
         Returns
@@ -566,7 +562,7 @@ def from_nwb(
         return BehaviorEcephysSession(
             behavior_session=behavior_session,
             probes=Probes.from_nwb(
-                nwbfile=nwbfile, probe_data_path_map=probe_data_path_map
+                nwbfile=nwbfile, probe_lfp_meta_map=probe_meta
             ),
             optotagging_table=OptotaggingTable.from_nwb(nwbfile=nwbfile),
             metadata=BehaviorEcephysMetadata.from_nwb(nwbfile=nwbfile),

allensdk/brain_observatory/ecephys/probes.py

@@ -1,12 +1,12 @@
 import logging
-from typing import List, Dict, Any, Optional, Union, Callable, Tuple
+from typing import List, Dict, Any, Optional, Tuple
 
 import numpy as np
 import pandas as pd
 import pynwb
 from pynwb import NWBFile
 
-from allensdk.brain_observatory.ecephys._probe import Probe
+from allensdk.brain_observatory.ecephys._probe import Probe, ProbeWithLFPMeta
 from allensdk.brain_observatory.ecephys.nwb_util import \
     add_ragged_data_to_dynamic_table
 from allensdk.core import DataObject, JsonReadableInterface, \
@@ -182,28 +182,28 @@ def to_dataframe(self):
     def from_nwb(
             cls,
             nwbfile: NWBFile,
-            probe_data_path_map: Optional[
-                Dict[str, Union[str, Callable[[], str]]]] = None
+            probe_lfp_meta_map: Optional[
+                Dict[str, ProbeWithLFPMeta]] = None
     ) -> "Probes":
         """
 
         Parameters
         ----------
         nwbfile
-        probe_data_path_map
+        probe_lfp_meta_map
             See description in `BehaviorEcephysSession.from_nwb`
 
         Returns
         -------
         `NWBFile` with probes added
         """
-        if probe_data_path_map is None:
-            probe_data_path_map = dict()
+        if probe_lfp_meta_map is None:
+            probe_lfp_meta_map = dict()
         probes = [
             Probe.from_nwb(
                 nwbfile=nwbfile,
                 probe_name=probe_name,
-                probe_nwb_path=probe_data_path_map.get(probe_name)
+                lfp_meta=probe_lfp_meta_map.get(probe_name)
             )
             for probe_name in nwbfile.electrode_groups]
         return Probes(probes=probes)

allensdk/test/brain_observatory/behavior/behavior_project_cache/conftest.py

@@ -281,13 +281,15 @@ def vbn_s3_cloud_cache_data():
             "ecephys_probe_id": 5111,
             "ecephys_session_id": 5111,
             "has_lfp_data": True,
+            "lfp_sampling_rate": 1.0,
             "name": "probeA",
             "file_id": 1024123123,
         },
         {
             "ecephys_probe_id": 5222,
             "ecephys_session_id": 5112,
             "has_lfp_data": True,
+            "lfp_sampling_rate": 2.0,
             "name": "probeA",
             "file_id": 1024123124,
         },
@@ -398,13 +400,15 @@ def vbn_s3_cloud_cache_data():
             "ecephys_probe_id": 5411,
             "ecephys_session_id": 222,
             "has_lfp_data": True,
+            "lfp_sampling_rate": 1.0,
             "name": "probeA",
             "file_id": 1024123125,
         },
         {
             "ecephys_probe_id": 5422,
             "ecephys_session_id": 222,
             "has_lfp_data": True,
+            "lfp_sampling_rate": 2.0,
             "name": "probeB",
             "file_id": 1024123126,
         },