Add reward and opt out times as separate columns

src/constantinople_lab_to_nwb/schierek_embargo_2024/interfaces/schierek_embargo_2024_processedbehaviorinterface.py

@@ -161,6 +161,96 @@ def add_to_nwbfile(
             data=center_port_offset_times + time_shift,
         )
 
+        side_port_columns = ["Cled", "Lled", "Rled", "l_opt", "r_opt"]
+        missing_columns = [col for col in side_port_columns if col not in data]
+        if missing_columns:
+            raise ValueError(f"Missing required columns in data: {', '.join(missing_columns)}")
+
+        # During the delay between the center light turning off and the reward arriving, the side light turns on.
+        # The side light turns off when the reward is available, then stays off until the animal collects the reward.
+        # When the animal nose pokes to collect the reward, the light flashes on/off.
+        reward_side_light_onset_times = []
+        reward_side_light_offset_times = []
+        reward_side_light_flash_onset_times = []
+        reward_side_light_flash_offset_times = []
+
+        opt_out_side_light_onset_times = []
+        opt_out_side_light_offset_times = []
+        opt_out_reward_port_turns_off = []
+        opt_out_reward_port_light_turns_off = []
+
+        for i in range(num_trials):
+            rewarded_side = data["RewardedSide"][i]
+            if rewarded_side == "Left":
+                side_port_column_name = "Lled"
+                # the opt-out port is the opposite of the rewarded side
+                opt_out_port_column_name = "r_opt"
+            elif rewarded_side == "Right":
+                side_port_column_name = "Rled"
+                opt_out_port_column_name = "l_opt"
+            else:
+                raise ValueError(f"Invalid rewarded side '{rewarded_side}'.")
+
+            reward_side_light_onset_times.append(data[side_port_column_name][i][0])
+            reward_side_light_offset_times.append(data[side_port_column_name][i][1])
+            reward_side_light_flash_onset_times.append(data[side_port_column_name][i][2])
+            reward_side_light_flash_offset_times.append(data[side_port_column_name][i][3])
+
+            opt_out_side_light_onset_times.append(data[opt_out_port_column_name][i][0])
+            opt_out_side_light_offset_times.append(data[opt_out_port_column_name][i][1])
+            opt_out_reward_port_turns_off.append(data[side_port_column_name][i][3])
+            opt_out_reward_port_light_turns_off.append(data[opt_out_port_column_name][i][3])
+
+        trials_table.add_column(
+            name="rewarded_port_onset_time",
+            description="The time of reward port light on for each trial. During the delay between the center light turning off and the reward arriving, the side light turns on.",
+            data=reward_side_light_onset_times + time_shift,
+        )
+
+        trials_table.add_column(
+            name="rewarded_port_offset_time",
+            description="The time of reward port light off for each trial. The side light turns off when the reward is available, then stays off until the animal collects the reward.",
+            data=reward_side_light_offset_times + time_shift,
+        )
+
+        trials_table.add_column(
+            name="rewarded_port_flash_onset_time",
+            description="The time of reward port light flash on for each trial. When the animal nose pokes to collect the reward, the light flashes on/off.",
+            data=reward_side_light_flash_onset_times + time_shift,
+        )
+
+        trials_table.add_column(
+            name="rewarded_port_flash_offset_time",
+            description="The time of reward port light flash off for each trial. When the animal nose pokes to collect the reward, the light flashes on/off.",
+            data=reward_side_light_flash_offset_times + time_shift,
+        )
+
+        trials_table.add_column(
+            name="opt_out_port_onset_time",
+            description=f"The time of side light turns on when the animal opts out by poking into the port opposite to the rewarded side.",
+            data=opt_out_side_light_onset_times + time_shift,
+        )
+
+        trials_table.add_column(
+            name="opt_out_port_offset_time",
+            description=f"The time of side light turns off when the animal opts out by poking into the port opposite to the rewarded side.",
+            data=opt_out_side_light_offset_times + time_shift,
+        )
+
+        trials_table.add_column(
+            name=f"opt_out_reward_port_offset_time",
+            description="The time of rewarded port turns off when the animal opts out by poking into the port opposite to the rewarded side.",
+            data=opt_out_reward_port_turns_off + time_shift,
+        )
+
+        trials_table.add_column(
+            name=f"opt_out_reward_port_light_offset_time",
+            description="The time of rewarded port light turns off when the animal opts out by poking into the port opposite to the rewarded side.",
+            data=opt_out_reward_port_light_turns_off + time_shift,
+        )
+
+        # filter columns to add, these columns were added separately
+        columns_to_add = [column for column in columns_to_add if column not in side_port_columns]
         for column_name in columns_to_add:
             name = column_name_mapping.get(column_name, column_name) if column_name_mapping is not None else column_name
             description = (