Enforce int typing for VBN columns.

AllenInstitute · Aug 11, 2023 · 1f76626 · 1f76626
1 parent da33e16
commit 1f76626
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/...ior/behavior_project_cache/project_apis/data_io/behavior_neuropixels_project_cloud_api.py b/...ior/behavior_project_cache/project_apis/data_io/behavior_neuropixels_project_cloud_api.py
@@ -10,6 +10,14 @@
     BehaviorEcephysSession,
 )
 from allensdk.brain_observatory.ecephys._probe import ProbeWithLFPMeta
+from allensdk.core.dataframe_utils import (
+    enforce_df_int_typing
+)
+
+INTEGER_COLUMNS = [
+    "prior_exposures_to_image_set", "ecephys_session_id", "unit_count",
+    "probe_count", "channel_count"
+]
 
 
 class VisualBehaviorNeuropixelsProjectCloudApi(ProjectCloudApiBase):
@@ -149,6 +157,7 @@ def f():
     def _get_ecephys_session_table(self):
         session_table_path = self._get_metadata_path(fname="ecephys_sessions")
         df = pd.read_csv(session_table_path)
+        df = enforce_df_int_typing(df, INTEGER_COLUMNS)
         self._ecephys_session_table = df.set_index("ecephys_session_id")
 
     def get_ecephys_session_table(self) -> pd.DataFrame:
@@ -161,6 +170,7 @@ def get_ecephys_session_table(self) -> pd.DataFrame:
     def _get_behavior_session_table(self):
         session_table_path = self._get_metadata_path(fname="behavior_sessions")
         df = pd.read_csv(session_table_path)
+        df = enforce_df_int_typing(df, INTEGER_COLUMNS)
         self._behavior_session_table = df.set_index("behavior_session_id")
 
     def get_behavior_session_table(self) -> pd.DataFrame: