feat: repetitions support for ScalaNWBReader (#194)

BlueBrain · Sep 19, 2024 · 88f089a · 88f089a
1 parent 7cd403b
commit 88f089a
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 11 deletions.
diff --git a/bluepyefe/nwbreader.py b/bluepyefe/nwbreader.py
@@ -103,6 +103,7 @@ def read(self):
 
 
 class ScalaNWBReader(NWBReader):
+
     def read(self):
         """ Read and format the content of the NWB file
 
@@ -112,6 +113,11 @@ def read(self):
 
         data = []
 
+        if self.repetition:
+            repetitions_content = self.content['general']['intracellular_ephys']['intracellular_recordings']['repetition']
+            if isinstance(self.repetition, (int, str)):
+                self.repetition = [int(self.repetition)]
+
         for sweep in list(self.content['acquisition'].keys()):
             key_current = sweep.replace('Series', 'StimulusSeries')
             try:
@@ -132,12 +138,23 @@ def read(self):
             if key_current not in self.content['stimulus']['presentation']:
                 continue
 
-            data.append(self._format_nwb_trace(
-                voltage=self.content['acquisition'][sweep]['data'],
-                current=self.content['stimulus']['presentation'][key_current]['data'],
-                start_time=self.content["acquisition"][sweep]["starting_time"],
-                trace_name=sweep,
-            ))
+            if self.repetition:
+                sweep_id = int(sweep.split("_")[-1])
+                if (int(repetitions_content[sweep_id]) in self.repetition):
+                    data.append(self._format_nwb_trace(
+                        voltage=self.content['acquisition'][sweep]['data'],
+                        current=self.content['stimulus']['presentation'][key_current]['data'],
+                        start_time=self.content['acquisition'][sweep]["starting_time"],
+                        trace_name=sweep,
+                        repetition=int(repetitions_content[sweep_id])
+                    ))
+            else:
+                data.append(self._format_nwb_trace(
+                    voltage=self.content['acquisition'][sweep]['data'],
+                    current=self.content['stimulus']['presentation'][key_current]['data'],
+                    start_time=self.content["acquisition"][sweep]["starting_time"],
+                    trace_name=sweep,
+                ))
 
         return data
 

diff --git a/bluepyefe/reader.py b/bluepyefe/reader.py
@@ -206,15 +206,15 @@ def nwb_reader(in_data):
     with h5py.File(in_data["filepath"], "r") as content:
         if "data_organization" in content:
             reader = BBPNWBReader(
-                content,
-                target_protocols,
-                in_data.get("repetition", None),
-                in_data.get("v_file", None)
+                content=content,
+                target_protocols=target_protocols,
+                v_file=in_data.get("v_file", None),
+                repetition=in_data.get("repetition", None),
             )
         elif "timeseries" in content["acquisition"].keys():
             reader = AIBSNWBReader(content, target_protocols)
         else:
-            reader = ScalaNWBReader(content, target_protocols)
+            reader = ScalaNWBReader(content, target_protocols, repetition=in_data.get("repetition", None))
 
         data = reader.read()