Consistent parameters of NWBIO.write_block (#1166)

* Consistent write_block

* try to fix failing test

* fix another source of test failure

* Move building of global annotations to a method, since it needs to be performed when `write_all_blocks()` is called, not on class initialization.

Co-authored-by: Andrew Davison <[email protected]>

Author: legouee

neo/io/nwbio.py

@@ -24,7 +24,6 @@
 
 import numpy as np
 import quantities as pq
-
 from neo.core import (Segment, SpikeTrain, Epoch, Event, AnalogSignal,
                       IrregularlySampledSignal, Block, ImageSequence)
 from neo.io.baseio import BaseIO
@@ -43,8 +42,7 @@
     "experiment_description", "session_id", "institution", "keywords", "notes",
     "pharmacology", "protocol", "related_publications", "slices", "source_script",
     "source_script_file_name", "data_collection", "surgery", "virus", "stimulus_notes",
-    "lab", "session_description",
-    "rec_datetime",
+    "lab", "session_description", "rec_datetime",
 )
 
 POSSIBLE_JSON_FIELDS = (
@@ -207,7 +205,7 @@ class NWBIO(BaseIO):
     is_writable = True
     is_streameable = False
 
-    def __init__(self, filename, mode='r'):
+    def __init__(self, filename, mode='r', **annotations):
         """
         Arguments:
             filename : the filename
@@ -218,6 +216,9 @@ def __init__(self, filename, mode='r'):
         self.filename = filename
         self.blocks_written = 0
         self.nwb_file_mode = mode
+        self._blocks = {}
+        self.annotations = annotations
+        self._io_nwb = None
 
     def read_all_blocks(self, lazy=False, **kwargs):
         """
@@ -383,17 +384,11 @@ def _read_acquisition_group(self, lazy):
     def _read_stimulus_group(self, lazy):
         self._read_timeseries_group("stimulus", lazy)
 
-    def write_all_blocks(self, blocks, **kwargs):
-        """
-        Write list of blocks to the file
-        """
-        import pynwb
-
-        # todo: allow metadata in NWBFile constructor to be taken from kwargs
+    def _build_global_annotations(self, blocks):
         annotations = defaultdict(set)
         for annotation_name in GLOBAL_ANNOTATIONS:
-            if annotation_name in kwargs:
-                annotations[annotation_name] = kwargs[annotation_name]
+            if annotation_name in self.annotations:
+                annotations[annotation_name] = self.annotations[annotation_name]
             else:
                 for block in blocks:
                     if annotation_name in block.annotations:
@@ -411,65 +406,86 @@ def write_all_blocks(self, blocks, **kwargs):
                             "We don't yet support multiple values for {}".format(annotation_name))
                     # take single value from set
                     annotations[annotation_name], = annotations[annotation_name]
+
         if "identifier" not in annotations:
-            annotations["identifier"] = self.filename
+            annotations["identifier"] = str(self.filename)
         if "session_description" not in annotations:
-            annotations["session_description"] = blocks[0].description or self.filename
+            annotations["session_description"] = blocks[0].description or str(self.filename)
             # todo: concatenate descriptions of multiple blocks if different
-        if "session_start_time" not in annotations:
-            annotations["session_start_time"] = blocks[0].rec_datetime
-            if annotations["session_start_time"] is None:
+        if annotations.get("session_start_time", None) is None:
+            if "rec_datetime" in annotations:
+                annotations["session_start_time"] = annotations["rec_datetime"]
+            else:
                 raise Exception("Writing to NWB requires an annotation 'session_start_time'")
-        self.annotations = {"rec_datetime": "rec_datetime"}
-        self.annotations["rec_datetime"] = blocks[0].rec_datetime
-        # todo: handle subject
-        nwbfile = pynwb.NWBFile(**annotations)
-        assert self.nwb_file_mode in ('w',)  # possibly expand to 'a'ppend later
-        if self.nwb_file_mode == "w" and os.path.exists(self.filename):
-            os.remove(self.filename)
-        io_nwb = pynwb.NWBHDF5IO(self.filename, mode=self.nwb_file_mode)
+        return annotations
+
+    def write_all_blocks(self, blocks, validate=True, **kwargs):
+        """
+        Write list of blocks to the file
+        """
+        import pynwb
+
+        global_annotations = self._build_global_annotations(blocks)
+        self._nwbfile = pynwb.NWBFile(**global_annotations)
 
         if sum(statistics(block)["SpikeTrain"]["count"] for block in blocks) > 0:
-            nwbfile.add_unit_column('_name', 'the name attribute of the SpikeTrain')
+            self._nwbfile.add_unit_column('_name', 'the name attribute of the SpikeTrain')
             # nwbfile.add_unit_column('_description',
             # 'the description attribute of the SpikeTrain')
-            nwbfile.add_unit_column(
+            self._nwbfile.add_unit_column(
                 'segment', 'the name of the Neo Segment to which the SpikeTrain belongs')
-            nwbfile.add_unit_column(
+            self._nwbfile.add_unit_column(
                 'block', 'the name of the Neo Block to which the SpikeTrain belongs')
 
         if sum(statistics(block)["Epoch"]["count"] for block in blocks) > 0:
-            nwbfile.add_epoch_column('_name', 'the name attribute of the Epoch')
+            self._nwbfile.add_epoch_column('_name', 'the name attribute of the Epoch')
             # nwbfile.add_epoch_column('_description', 'the description attribute of the Epoch')
-            nwbfile.add_epoch_column(
+            self._nwbfile.add_epoch_column(
                 'segment', 'the name of the Neo Segment to which the Epoch belongs')
-            nwbfile.add_epoch_column('block',
+            self._nwbfile.add_epoch_column('block',
                                      'the name of the Neo Block to which the Epoch belongs')
 
         for i, block in enumerate(blocks):
-            self.write_block(nwbfile, block)
-        io_nwb.write(nwbfile)
+            self._write_block(block)
+
+        assert self.nwb_file_mode in ('w',)  # possibly expand to 'a'ppend later
+        if self.nwb_file_mode == "w" and os.path.exists(self.filename):
+            os.remove(self.filename)
+        io_nwb = pynwb.NWBHDF5IO(self.filename, mode=self.nwb_file_mode)
+        io_nwb.write(self._nwbfile)
         io_nwb.close()
 
+        if validate:
+            self.validate_file()
+
+    def validate_file(self):
+        import pynwb
+
         with pynwb.NWBHDF5IO(self.filename, "r") as io_validate:
             errors = pynwb.validate(io_validate, namespace="core")
             if errors:
                 raise Exception(f"Errors found when validating {self.filename}")
 
-    def write_block(self, nwbfile, block, **kwargs):
+    def write_block(self, block, **kwargs):
+        """
+        Write a single Block to the file
+            :param block: Block to be written
+        """
+        return self.write_all_blocks([block], **kwargs)
+
+    def _write_block(self, block):
         """
         Write a Block to the file
             :param block: Block to be written
-            :param nwbfile: Representation of an NWB file
         """
-        electrodes = self._write_electrodes(nwbfile, block)
+        electrodes = self._write_electrodes(self._nwbfile, block)
         if not block.name:
             block.name = "block%d" % self.blocks_written
         for i, segment in enumerate(block.segments):
             assert segment.block is block
             if not segment.name:
                 segment.name = "%s : segment%d" % (block.name, i)
-            self._write_segment(nwbfile, segment, electrodes)
+            self._write_segment(self._nwbfile, segment, electrodes)
         self.blocks_written += 1
 
     def _write_electrodes(self, nwbfile, block):
@@ -485,10 +501,10 @@ def _write_electrodes(self, nwbfile, block):
                         if elec_meta["device"]["name"] in devices:
                             device = devices[elec_meta["device"]["name"]]
                         else:
-                            device = nwbfile.create_device(**elec_meta["device"])
+                            device = self._nwbfile.create_device(**elec_meta["device"])
                             devices[elec_meta["device"]["name"]] = device
                         elec_meta.pop("device")
-                        electrodes[elec_meta["name"]] = nwbfile.create_icephys_electrode(
+                        electrodes[elec_meta["name"]] = self._nwbfile.create_icephys_electrode(
                             device=device, **elec_meta
                         )
         return electrodes
@@ -503,13 +519,13 @@ def _write_segment(self, nwbfile, segment, electrodes):
                 logging.warning("Warning signal name exists. New name: %s" % (signal.name))
             else:
                 signal.name = "%s : analogsignal%s %i" % (segment.name, signal.name, i)
-            self._write_signal(nwbfile, signal, electrodes)
+            self._write_signal(self._nwbfile, signal, electrodes)
 
         for i, train in enumerate(segment.spiketrains):
             assert train.segment is segment
             if not train.name:
                 train.name = "%s : spiketrain%d" % (segment.name, i)
-            self._write_spiketrain(nwbfile, train)
+            self._write_spiketrain(self._nwbfile, train)
 
         for i, event in enumerate(segment.events):
             assert event.segment is segment
@@ -518,12 +534,12 @@ def _write_segment(self, nwbfile, segment, electrodes):
                 logging.warning("Warning event name exists. New name: %s" % (event.name))
             else:
                 event.name = "%s : event%s %d" % (segment.name, event.name, i)
-            self._write_event(nwbfile, event)
+            self._write_event(self._nwbfile, event)
 
         for i, epoch in enumerate(segment.epochs):
             if not epoch.name:
                 epoch.name = "%s : epoch%d" % (segment.name, i)
-            self._write_epoch(nwbfile, epoch)
+            self._write_epoch(self._nwbfile, epoch)
 
     def _write_signal(self, nwbfile, signal, electrodes):
         import pynwb
@@ -572,8 +588,8 @@ def _write_signal(self, nwbfile, signal, electrodes):
                     signal.__class__.__name__))
         nwb_group = signal.annotations.get("nwb_group", "acquisition")
         add_method_map = {
-            "acquisition": nwbfile.add_acquisition,
-            "stimulus": nwbfile.add_stimulus
+            "acquisition": self._nwbfile.add_acquisition,
+            "stimulus": self._nwbfile.add_stimulus
         }
         if nwb_group in add_method_map:
             add_time_series = add_method_map[nwb_group]
@@ -586,7 +602,8 @@ def _write_spiketrain(self, nwbfile, spiketrain):
         segment = spiketrain.segment
         if hasattr(spiketrain, 'proxy_for') and spiketrain.proxy_for is SpikeTrain:
             spiketrain = spiketrain.load()
-        nwbfile.add_unit(spike_times=spiketrain.rescale('s').magnitude,
+        self._nwbfile.add_unit(
+                         spike_times=spiketrain.rescale('s').magnitude,
                          obs_intervals=[[float(spiketrain.t_start.rescale('s')),
                                          float(spiketrain.t_stop.rescale('s'))]],
                          _name=spiketrain.name,
@@ -596,7 +613,7 @@ def _write_spiketrain(self, nwbfile, spiketrain):
         # todo: handle annotations (using add_unit_column()?)
         # todo: handle Neo Units
         # todo: handle spike waveforms, if any (see SpikeEventSeries)
-        return nwbfile.units
+        return self._nwbfile.units
 
     def _write_event(self, nwbfile, event):
         import pynwb
@@ -611,7 +628,7 @@ def _write_event(self, nwbfile, event):
             timestamps=event.times.rescale('second').magnitude,
             description=event.description or "",
             comments=json.dumps(hierarchy))
-        nwbfile.add_acquisition(tS_evt)
+        self._nwbfile.add_acquisition(tS_evt)
         return tS_evt
 
     def _write_epoch(self, nwbfile, epoch):
@@ -621,11 +638,11 @@ def _write_epoch(self, nwbfile, epoch):
         for t_start, duration, label in zip(epoch.rescale('s').magnitude,
                                             epoch.durations.rescale('s').magnitude,
                                             epoch.labels):
-            nwbfile.add_epoch(t_start, t_start + duration, [label], [],
+            self._nwbfile.add_epoch(t_start, t_start + duration, [label], [],
                               _name=epoch.name,
                               segment=segment.name,
                               block=segment.block.name)
-        return nwbfile.epochs
+        return self._nwbfile.epochs
 
 
 class AnalogSignalProxy(BaseAnalogSignalProxy):