Be more conistent about typing.

AllenInstitute · Oct 28, 2023 · eabc4ee · eabc4ee
1 parent 0eb740c
commit eabc4ee
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 11 deletions.
diff --git a/...rvatory/behavior/behavior_project_cache/project_apis/data_io/behavior_project_lims_api.py b/...rvatory/behavior/behavior_project_cache/project_apis/data_io/behavior_project_lims_api.py
@@ -437,9 +437,7 @@ def _get_ophys_cells_table(self):
         df = self.lims_engine.select(query)
 
         # NaN's for invalid cells force this to float, push to int
-        df["cell_specimen_id"] = pd.array(
-            df["cell_specimen_id"], dtype="Int64"
-        )
+        df = enforce_df_int_typing(df, VBO_INTEGER_COLUMNS)
         return df
 
     def get_ophys_cells_table(self):
@@ -509,10 +507,7 @@ def get_ophys_session_table(self) -> pd.DataFrame:
 
         # Fill NaN values of imaging_plane_group_count with zero to match
         # the behavior of the BehaviorOphysExperiment object.
-        im_plane_count = (
-            table["imaging_plane_group_count"].astype("Int64")
-        )
-        table["imaging_plane_group_count"] = im_plane_count
+        table = enforce_df_int_typing(table, VBO_INTEGER_COLUMNS)
         return table
 
     def get_behavior_session(
@@ -544,7 +539,7 @@ def get_ophys_experiment_table(self) -> pd.DataFrame:
         df = self._get_ophys_experiment_table()
         # Set type to pandas.Int64 to enforce integer typing and not revert to
         # float.
-        df["imaging_plane_group"] = df["imaging_plane_group"].astype("Int64")
+        df = enforce_df_int_typing(df, VBO_INTEGER_COLUMNS)
         return df.set_index("ophys_experiment_id")
 
     def get_behavior_session_table(self) -> pd.DataFrame:
@@ -559,13 +554,13 @@ def get_behavior_session_table(self) -> pd.DataFrame:
         acquisition date for behavior sessions (only in the stimulus pkl file)
         """
         summary_tbl = self._get_behavior_summary_table()
-        # Query returns float typing of age_in_days. Convert to int to match
-        # typing of the Age data_object.
-        summary_tbl["age_in_days"] = summary_tbl["age_in_days"].astype("Int64")
         # Add UTC time zone to match timezone from DateOfAcquisition object.
         summary_tbl["date_of_acquisition"] = pd.to_datetime(
             summary_tbl["date_of_acquisition"], utc=True
         )
+        # Query returns float typing of age_in_days. Convert to int to match
+        # typing of the Age data_object.
+        summary_tbl = enforce_df_int_typing(summary_tbl, VBO_INTEGER_COLUMNS)
 
         return summary_tbl.set_index("behavior_session_id")
 

diff --git a/allensdk/brain_observatory/ophys/project_constants.py b/allensdk/brain_observatory/ophys/project_constants.py
@@ -60,4 +60,5 @@
     "targeted_areas",
     "num_depths_per_area",
     "num_targeted_structures",
+    "cell_specimen_id"
 ]