From f70138b190e4fefcaec5f6499657f9d530d3cc41 Mon Sep 17 00:00:00 2001
From: bjhardcastle <ben.hardcastle@alleninstitute.org>
Date: Tue, 22 Oct 2024 21:34:13 -0700
Subject: [PATCH] Use solenoid time for `trials.reward_times` if available -
 fixes #147

---
 src/npc_sessions/sessions.py                  |  3 ++-
 .../trials/TaskControl/DynamicRouting1.py     | 24 ++++++++++++++-----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/npc_sessions/sessions.py b/src/npc_sessions/sessions.py
index 53b4114..5f7143f 100644
--- a/src/npc_sessions/sessions.py
+++ b/src/npc_sessions/sessions.py
@@ -1140,7 +1140,8 @@ def set_lazy_eval(
                 kwargs |= {"sync": self.sync_data}
             if self.is_ephys and self.is_sync:
                 kwargs |= {"ephys_recording_dirs": self.ephys_recording_dirs}
-
+            if (reward_times := getattr(self, "_reward_times_with_duration", None)) is not None:
+                kwargs |= {"reward_times_with_duration": reward_times.timestamps}
             # set items in LazyDict for postponed evaluation
             if "RFMapping" in stim_filename:
                 # create two separate trials tables
diff --git a/src/npc_sessions/trials/TaskControl/DynamicRouting1.py b/src/npc_sessions/trials/TaskControl/DynamicRouting1.py
index 9a0edbd..171e2fb 100644
--- a/src/npc_sessions/trials/TaskControl/DynamicRouting1.py
+++ b/src/npc_sessions/trials/TaskControl/DynamicRouting1.py
@@ -540,9 +540,20 @@ def response_time(self) -> npt.NDArray[np.float64]:
     @npc_io.cached_property
     def reward_time(self) -> npt.NDArray[np.floating]:
         """delivery time of water reward, for contingent and non-contingent rewards"""
-        all_reward_times = npc_stim.safe_index(self._flip_times, self._sam.rewardFrames)
-        all_reward_times = all_reward_times[all_reward_times <= self.stop_time[-1]]
-        all_reward_trials = (
+        if (
+            (solenoid_times := getattr(self, "_reward_times_with_duration", None)) is not None
+            and len(solenoid_times) >= len(np.where(self.is_rewarded)[0])
+        ):
+            logger.info(f'Using solenoid opening time on sync for `reward_time`')
+            all_reward_times = solenoid_times
+        else:
+            logger.info(f'Using flip time of each TaskControl frame for `reward_time`')
+            all_reward_times = npc_stim.safe_index(self._flip_times, self._sam.rewardFrames)
+        all_reward_times = all_reward_times[
+            (all_reward_times >= self.start_time[0]) &
+            (all_reward_times <= self.stop_time[-1])
+        ]
+        trial_idx_from_rewards = (
             np.searchsorted(
                 self.start_time,
                 all_reward_times,
@@ -550,10 +561,11 @@ def reward_time(self) -> npt.NDArray[np.floating]:
             )
             - 1
         )
+        assert len(is_rewarded := np.where(self.is_rewarded)[0]) <= len(trial_idx_from_rewards)
         reward_time = np.full(self._len, np.nan)
-        if np.all(np.where(self.is_rewarded)[0] == all_reward_trials):
-            # expected single reward per trial
-            reward_time[all_reward_trials] = all_reward_times
+        if np.all(is_rewarded == trial_idx_from_rewards):
+            # expected case: single reward per trial
+            reward_time[trial_idx_from_rewards] = all_reward_times
         else:
             # mismatch between reward times and trials that are marked as having rewards
             for trial_idx in np.where(self.is_rewarded)[0]: