Merge branch 'ticket/PSB-155/dev' into rc/2.16.0

allensdk/brain_observatory/behavior/behavior_session.py

@@ -285,6 +285,7 @@ def from_json(
                     path=session_data["stim_table_file"],
                     behavior_session_id=session_data["behavior_session_id"],
                     exclude_columns=stimulus_presentation_exclude_columns,
+                    trials=trials,
                 ),
                 templates=Templates.from_stimulus_file(
                     stimulus_file=stimulus_file_lookup.behavior_stimulus_file
@@ -357,7 +358,7 @@ def from_lims(
         date_of_acquisition: Optional[DateOfAcquisition] = None,
         eye_tracking_z_threshold: float = 3.0,
         eye_tracking_dilation_frames: int = 2,
-        load_stimulus_movie: bool = True
+        load_stimulus_movie: bool = True,
     ) -> "BehaviorSession":
         """
 
@@ -453,7 +454,7 @@ def from_lims(
             project_code=ProjectCode.from_lims(
                 behavior_session_id=behavior_session_id.value, lims_db=lims_db
             ),
-            load_stimulus_movie=load_stimulus_movie
+            load_stimulus_movie=load_stimulus_movie,
         )
 
         if date_of_acquisition is None:
@@ -1114,8 +1115,10 @@ def stimulus_presentations(self) -> pd.DataFrame:
         table = table.drop(columns=["image_set", "index"], errors="ignore")
         table = table.rename(columns={"stop_time": "end_time"})
 
-        if "trials_id" not in table.columns \
-                and 'stimulus_block' in table.columns:
+        if (
+            "trials_id" not in table.columns
+            and "stimulus_block" in table.columns
+        ):
             table["trials_id"] = compute_trials_id_for_stimulus(
                 table, self.trials
             )
@@ -1172,9 +1175,7 @@ def stimulus_natural_movie_template(self) -> Optional[pd.DataFrame]:
         if self._stimuli.templates.fingerprint_movie_template_key is not None:
             return self._stimuli.templates.value[
                 self._stimuli.templates.fingerprint_movie_template_key
-            ].to_dataframe(
-                index_name='frame_number',
-                index_type='int')
+            ].to_dataframe(index_name="frame_number", index_type="int")
         else:
             return None
 
@@ -1412,7 +1413,7 @@ def _read_stimuli(
         trials: Trials,
         stimulus_presentation_columns: Optional[List[str]] = None,
         project_code: Optional[ProjectCode] = None,
-        load_stimulus_movie: bool = False
+        load_stimulus_movie: bool = False,
     ) -> Stimuli:
         """
         Construct the Stimuli data object for this session
@@ -1431,7 +1432,7 @@ def _read_stimuli(
             presentation_columns=stimulus_presentation_columns,
             project_code=project_code,
             trials=trials,
-            load_stimulus_movie=load_stimulus_movie
+            load_stimulus_movie=load_stimulus_movie,
         )
 
     @classmethod
@@ -1511,7 +1512,7 @@ def _read_data_from_stimulus_file(
         include_stimuli: bool = True,
         stimulus_presentation_columns: Optional[List[str]] = None,
         project_code: Optional[ProjectCode] = None,
-        load_stimulus_movie: bool = False
+        load_stimulus_movie: bool = False,
     ):
         """Helper method to read data from stimulus file"""
 
@@ -1548,7 +1549,7 @@ def _read_data_from_stimulus_file(
                 trials=trials,
                 stimulus_presentation_columns=stimulus_presentation_columns,
                 project_code=project_code,
-                load_stimulus_movie=load_stimulus_movie
+                load_stimulus_movie=load_stimulus_movie,
             )
         else:
             stimuli = None

allensdk/brain_observatory/behavior/data_objects/stimuli/presentations.py

@@ -20,6 +20,8 @@
     Trials,
 )
 from allensdk.brain_observatory.behavior.stimulus_processing import (
+    add_active_flag,
+    compute_is_sham_change,
     compute_trials_id_for_stimulus,
     fix_omitted_end_frame,
     get_flashes_since_change,
@@ -58,22 +60,61 @@ def __init__(
         columns_to_rename: Optional[Dict[str, str]] = None,
         column_list: Optional[List[str]] = None,
         sort_columns: bool = True,
+        trials: Optional[Trials] = None,
     ):
         """
 
         Parameters
         ----------
-        presentations: The stimulus presentations table
-        columns_to_rename: Optional dict mapping
-            old column name -> new column name
+        presentations : pandas.DataFrame
+            The stimulus presentations table
+        columns_to_rename : Optional dict mapping
+            Mapping to rename columns. old column name -> new column name
         column_list: Optional list of columns to include.
             This will reorder the columns.
-        sort_columns: Whether to sort the columns by name
+        sort_columns: bool
+            Whether to sort the columns by name
+        trials : Optional Trials object.
+            allensdk Trials object for the same session as the presentations
+            table.
         """
         if columns_to_rename is not None:
             presentations = presentations.rename(columns=columns_to_rename)
         if column_list is not None:
             presentations = presentations[column_list]
+        presentations = enforce_df_int_typing(
+            presentations,
+            [
+                "flashes_since_change",
+                "image_index",
+                "movie_frame_index",
+                "repeat",
+                "stimulus_index",
+            ],
+        )
+        presentations = presentations.reset_index(drop=True)
+        presentations.index = pd.Index(
+            range(presentations.shape[0]),
+            name="stimulus_presentations_id",
+            dtype="int",
+        )
+        if trials is not None:
+            if "active" not in presentations.columns:
+                # Add column marking where the mouse is engaged in active,
+                # trained behavior.
+                presentations = add_active_flag(presentations, trials.data)
+            if "trials_id" not in presentations.columns:
+                # Add trials_id to presentations df to allow for joining of the
+                # two tables.
+                presentations["trials_id"] = compute_trials_id_for_stimulus(
+                    presentations, trials.data
+                )
+            if "is_sham_change" not in presentations.columns:
+                # Mark changes in active and replay stimulus that are
+                # #sham-changes
+                presentations = compute_is_sham_change(
+                    presentations, trials.data
+                )
         if sort_columns:
             presentations = enforce_df_column_order(
                 presentations,
@@ -97,23 +138,6 @@ def __init__(
                     "trials_id",
                 ],
             )
-        presentations = presentations.reset_index(drop=True)
-        presentations = enforce_df_int_typing(
-            presentations,
-            [
-                "flashes_since_change",
-                "image_index",
-                "movie_frame_index",
-                "movie_repeat",
-                "stimulus_index",
-            ],
-        )
-        presentations.index = pd.Index(
-            range(presentations.shape[0]),
-            name="stimulus_presentations_id",
-            dtype="int",
-        )
-
         super().__init__(name="presentations", value=presentations)
 
     def to_nwb(
@@ -181,6 +205,7 @@ def from_nwb(
         cls,
         nwbfile: NWBFile,
         add_is_change: bool = True,
+        add_trials_dependent_values: bool = True,
         column_list: Optional[List[str]] = None,
     ) -> "Presentations":
         """
@@ -240,7 +265,13 @@ def from_nwb(
             table["flashes_since_change"] = get_flashes_since_change(
                 stimulus_presentations=table
             )
-        return Presentations(presentations=table, column_list=column_list)
+        trials = None
+        if add_trials_dependent_values and nwbfile.trials is not None:
+            trials = Trials.from_nwb(nwbfile)
+
+        return Presentations(
+            presentations=table, column_list=column_list, trials=trials
+        )
 
     @classmethod
     def from_stimulus_file(
@@ -291,7 +322,8 @@ def from_stimulus_file(
         )
         raw_stim_pres_df = raw_stim_pres_df.drop(columns=["index"])
         raw_stim_pres_df = cls._check_for_errant_omitted_stimulus(
-            input_df=raw_stim_pres_df)
+            input_df=raw_stim_pres_df
+        )
 
         # Fill in nulls for image_name
         # This makes two assumptions:
@@ -394,14 +426,8 @@ def from_stimulus_file(
                 stim_pres_df, stimulus_file.session_type, project_code.value
             )
 
-        # Add trials_id to presentations df to allow for joining of the two
-        # tables.
-        stim_pres_df["trials_id"] = compute_trials_id_for_stimulus(
-            stim_pres_df, trials.data
-        )
-
         return Presentations(
-            presentations=stim_pres_df, column_list=column_list
+            presentations=stim_pres_df, column_list=column_list, trials=trials
         )
 
     @classmethod
@@ -412,6 +438,7 @@ def from_path(
         exclude_columns: Optional[List[str]] = None,
         columns_to_rename: Optional[Dict[str, str]] = None,
         sort_columns: bool = True,
+        trials: Optional[Trials] = None,
     ) -> "Presentations":
         """
         Reads the table directly from a precomputed csv
@@ -446,6 +473,7 @@ def from_path(
             presentations=df,
             columns_to_rename=columns_to_rename,
             sort_columns=sort_columns,
+            trials=trials,
         )
 
     @classmethod
@@ -547,7 +575,7 @@ def _postprocess(
 
     @staticmethod
     def _check_for_errant_omitted_stimulus(
-        input_df: pd.DataFrame
+        input_df: pd.DataFrame,
     ) -> pd.DataFrame:
         """Check if the first entry in the DataFrame is an omitted stimulus.