Simplify trials_id

allensdk/brain_observatory/behavior/stimulus_processing.py

@@ -687,7 +687,8 @@ def add_active_flag(
         return stim_pres_table
 
 def compute_trials_id_for_stimulus(
-    stim_pres_table: pd.DataFrame, trials_table: pd.DataFrame
+        stim_pres_table: pd.DataFrame,
+        trials_table: pd.DataFrame
 ) -> pd.Series:
     """Add an id to allow for merging of the stimulus presentations
     table with the trials table.
@@ -714,26 +715,24 @@ def compute_trials_id_for_stimulus(
     passive stimulus/replay blocks that contain the same image ordering and
     length.
     """
-    stim_pres_sorted = stim_pres_table.sort_values("start_time")
-    trials_sorted = trials_table.sort_values("start_time")
     # Create a placeholder for the trials_id.
     trials_ids = pd.Series(
-        data=np.full(len(stim_pres_sorted), INT_NULL, dtype=int),
-        index=stim_pres_sorted.index,
+        data=np.full(len(stim_pres_table), INT_NULL, dtype=int),
+        index=stim_pres_table.index,
         name="trials_id",
     ).astype('int')
     # Return input frame if the stimulus_block or active is not available.
     if "stimulus_block" not in stim_pres_table.columns \
             or "active" not in stim_pres_table.columns:
         return trials_ids
-    active_sorted = stim_pres_sorted.active
+    active_sorted = stim_pres_table.active
 
     # Find stimulus blocks that start within a trial. Copy the trial_id
     # into our new trials_ids series.
-    for idx, trial in trials_sorted.iterrows():
-        stim_mask = (stim_pres_sorted.start_time > trial.start_time) & (
-            stim_pres_sorted.start_time < trial.stop_time
-        ) & (~stim_pres_sorted.image_name.isna())
+    for idx, trial in trials_table.iterrows():
+        stim_mask = (stim_pres_table.start_time > trial.start_time) & (
+            stim_pres_table.start_time < trial.stop_time
+        ) & (~stim_pres_table.image_name.isna())
         trials_ids[stim_mask] = idx
 
     # The code below finds all stimulus blocks that contain images/trials
@@ -742,8 +741,8 @@ def compute_trials_id_for_stimulus(
     # copying the active stimulus block data into the passive stimulus block.
 
     # Get the block ids for the behavior trial presentations
-    stim_blocks = stim_pres_sorted.stimulus_block
-    stim_image_names = stim_pres_sorted.image_name
+    stim_blocks = stim_pres_table.stimulus_block
+    stim_image_names = stim_pres_table.image_name
     active_stim_blocks = stim_blocks[active_sorted].unique()
     # Find passive blocks that show images for potential copying of the active
     # into a passive stimulus block.