Stop shifting the second TDT block's reward TTLs twice

src/lerner_lab_to_nwb/seiler_2024/seiler_2024nwbconverter.py

@@ -155,9 +155,9 @@ def temporally_align_data_interfaces(self, metadata: dict, conversion_options: d
             else:
                 ttl_timestamps = self.get_ttl_timestamps(ttl_name, tdt_photometry, len_behavior)
             if second_folder_path is not None:
-                ttl_timestamps2 = self.get_ttl_timestamps(ttl_name, tdt_photometry2, len_behavior)
-                t_end = 1 / tdt_photometry.streams["Dv1A"].fs * (len(tdt_photometry.streams["Dv1A"].data) - 1)
-                ttl_timestamps2 += t_end
+                ttl_timestamps2 = self.get_second_block_ttl_timestamps(
+                    ttl_name, tdt_photometry, tdt_photometry2, len_behavior
+                )
                 ttl_timestamps = np.concatenate((ttl_timestamps, ttl_timestamps2))
             session_dict[behavior_name] = ttl_timestamps
         if "MedPC" in self.data_interface_objects.keys():
@@ -169,6 +169,12 @@ def temporally_align_data_interfaces(self, metadata: dict, conversion_options: d
             session_dict.keys()
         )
 
+    def get_second_block_ttl_timestamps(self, ttl_name, tdt_photometry, tdt_photometry2, len_behavior):
+        """Get the TTL onsets of the second TDT block of a session, placed after the end of the first block."""
+        ttl_timestamps2 = self.get_ttl_timestamps(ttl_name, tdt_photometry2, len_behavior)
+        t_end = 1 / tdt_photometry.streams["Dv1A"].fs * (len(tdt_photometry.streams["Dv1A"].data) - 1)
+        return ttl_timestamps2 + t_end
+
     def get_ttl_timestamps(self, ttl_name, tdt_photometry, len_behavior):
         if ttl_name == "PrtN":
             ttl_timestamps = []

tests/test_split_recordings.py

@@ -0,0 +1,47 @@
+import numpy as np
+import pytest
+from tdt import StructType
+
+from lerner_lab_to_nwb.seiler_2024 import Seiler2024NWBConverter
+
+
+def tdt_block(*, ttl_name_to_onsets: dict[str, list[float]], number_of_samples: int = 101) -> StructType:
+    block = StructType()
+    block.streams = StructType()
+    block.streams["Dv1A"] = StructType()
+    block.streams["Dv1A"].fs = 1.0
+    block.streams["Dv1A"].data = np.zeros(number_of_samples)
+    block.epocs = StructType()
+    for ttl_name, onsets in ttl_name_to_onsets.items():
+        block.epocs[ttl_name] = StructType()
+        block.epocs[ttl_name].onset = np.array(onsets)
+    return block
+
+
+@pytest.fixture
+def converter() -> Seiler2024NWBConverter:
+    return Seiler2024NWBConverter(source_data={})
+
+
+@pytest.fixture
+def first_block() -> StructType:
+    # 101 samples at 1 Hz, so the second block is placed 100 s after the first
+    return tdt_block(ttl_name_to_onsets={})
+
+
+class TestSeiler2024NWBConverter:
+    def test_rewarded_nose_pokes_already_in_nose_pokes_are_not_added_again(self, converter, first_block):
+        # As in 139.298: RNnR already contains the rewarded nose poke
+        second_block = tdt_block(ttl_name_to_onsets={"RNnR": [5.0, 7.0], "RNRW": [5.0]})
+        right_reward_times = converter.get_second_block_ttl_timestamps("RNRW", first_block, second_block, 10)
+        right_nose_poke_times = converter.get_second_block_ttl_timestamps("RNnR", first_block, second_block, 10)
+        np.testing.assert_array_equal(right_reward_times, np.array([105.0]))
+        np.testing.assert_array_equal(right_nose_poke_times, np.array([105.0, 107.0]))
+
+    def test_rewarded_nose_pokes_merged_into_nose_pokes_are_shifted_once(self, converter, first_block):
+        # As in 92.246: LNnR has only the unrewarded nose pokes
+        second_block = tdt_block(ttl_name_to_onsets={"LNnR": [7.0], "LNRW": [5.0]})
+        left_reward_times = converter.get_second_block_ttl_timestamps("LNRW", first_block, second_block, 10)
+        left_nose_poke_times = converter.get_second_block_ttl_timestamps("LNnR", first_block, second_block, 10)
+        np.testing.assert_array_equal(left_reward_times, np.array([105.0]))
+        np.testing.assert_array_equal(left_nose_poke_times, np.array([105.0, 107.0]))