Merge pull request #1246 from NeurodataWithoutBorders/fix/elec_table

Author: rly

src/pynwb/file.py

@@ -525,9 +525,12 @@ def add_electrode(self, **kwargs):
                     ('rel_y', 'the y coordinate within the electrode group'),
                     ('rel_z', 'the z coordinate within the electrode group'),
                     ('reference', 'Description of the reference used for this electrode.')]
+        # add column if the arg is supplied and column does not yet exist
+        # do not pass arg to add_row if arg is not supplied
         for col_name, col_doc in new_cols:
-            if kwargs[col_name] is not None and col_name not in self.electrodes:
-                self.electrodes.add_column(col_name, col_doc)
+            if kwargs[col_name] is not None:
+                if col_name not in self.electrodes:
+                    self.electrodes.add_column(col_name, col_doc)
             else:
                 d.pop(col_name)  # remove args from d if not set
 

tests/integration/hdf5/test_nwbfile.py

@@ -444,17 +444,86 @@ def setUpContainer(self):
     def addContainer(self, nwbfile):
         """ Add electrodes and related objects to the given NWBFile """
         self.dev1 = nwbfile.create_device(name='dev1')
-        self.group = nwbfile.create_electrode_group(name='tetrode1', description='tetrode description',
-                                                    location='tetrode location', device=self.dev1)
+        self.group = nwbfile.create_electrode_group(
+            name='tetrode1',
+            description='tetrode description',
+            location='tetrode location',
+            device=self.dev1
+        )
 
-        nwbfile.add_electrode(id=1, x=1.0, y=2.0, z=3.0, imp=-1.0, location='CA1', filtering='none', group=self.group,
-                              group_name='tetrode1')
-        nwbfile.add_electrode(id=2, x=1.0, y=2.0, z=3.0, imp=-2.0, location='CA1', filtering='none', group=self.group,
-                              group_name='tetrode1')
-        nwbfile.add_electrode(id=3, x=1.0, y=2.0, z=3.0, imp=-3.0, location='CA1', filtering='none', group=self.group,
-                              group_name='tetrode1')
-        nwbfile.add_electrode(id=4, x=1.0, y=2.0, z=3.0, imp=-4.0, location='CA1', filtering='none', group=self.group,
-                              group_name='tetrode1')
+        nwbfile.add_electrode(
+            id=1,
+            x=1.0, y=2.0, z=3.0,
+            imp=-1.0,
+            location='CA1',
+            filtering='none',
+            group=self.group,
+            group_name='tetrode1'
+        )
+        nwbfile.add_electrode(
+            id=2,
+            x=1.0, y=2.0, z=3.0,
+            imp=-2.0,
+            location='CA1',
+            filtering='none',
+            group=self.group,
+            group_name='tetrode1'
+        )
+
+        self.container = nwbfile.electrodes  # override self.container which has the placeholder
+
+    def getContainer(self, nwbfile):
+        """ Return the test electrodes table from the given NWBFile """
+        return nwbfile.electrodes
+
+    def test_roundtrip(self):
+        super().test_roundtrip()
+        # When comparing the pandas dataframes for the row we drop the 'group' column since the
+        # ElectrodeGroup object after reading will naturally have a different address
+        pd.testing.assert_frame_equal(self.read_container[0].drop('group', axis=1),
+                                      self.container[0].drop('group', axis=1))
+
+
+class TestElectrodesOptColumns(NWBH5IOMixin, TestCase):
+
+    def setUpContainer(self):
+        """
+        Return placeholder table for electrodes. Tested electrodes are added directly to the NWBFile in addContainer
+        """
+        return DynamicTable('electrodes', 'a placeholder table')
+
+    def addContainer(self, nwbfile):
+        """ Add electrodes and related objects to the given NWBFile """
+        self.dev1 = nwbfile.create_device(name='dev1')
+        self.group = nwbfile.create_electrode_group(
+            name='tetrode1',
+            description='tetrode description',
+            location='tetrode location',
+            device=self.dev1
+        )
+
+        nwbfile.add_electrode(
+            id=1,
+            x=1.0, y=2.0, z=3.0,
+            imp=-1.0,
+            location='CA1',
+            filtering='none',
+            group=self.group,
+            group_name='tetrode1',
+            rel_x=4.0, rel_y=5.0, rel_z=6.0,
+            reference='ref1'
+        )
+        nwbfile.add_electrode(
+            id=2,
+            x=1.0, y=2.0, z=3.0,
+            imp=-2.0,
+            location='CA1',
+            filtering='none',
+            group=self.group,
+            group_name='tetrode1',
+            rel_x=4.0, rel_y=5.0, rel_z=6.0,
+            reference='ref1'
+        )
 
         self.container = nwbfile.electrodes  # override self.container which has the placeholder
 

tests/unit/test_file.py

@@ -242,9 +242,21 @@ def test_add_invalid_time_w_ts(self):
                                                timeseries=ts, tags=('hi', 'there'))
 
     def test_add_electrode(self):
-        dev1 = self.nwbfile.create_device('dev1')
-        group = self.nwbfile.create_electrode_group('tetrode1', 'tetrode description', 'tetrode location', dev1)
-        self.nwbfile.add_electrode(1.0, 2.0, 3.0, -1.0, 'CA1', 'none', group=group, id=1)
+        dev1 = self.nwbfile.create_device(name='dev1')
+        group = self.nwbfile.create_electrode_group(
+            name='tetrode1',
+            description='tetrode description',
+            location='tetrode location',
+            device=dev1
+        )
+        self.nwbfile.add_electrode(
+            x=1.0, y=2.0, z=3.0,
+            imp=-1.0,
+            location='CA1',
+            filtering='none',
+            group=group,
+            id=1
+        )
         elec = self.nwbfile.electrodes[0]
         self.assertEqual(elec.index[0], 1)
         self.assertEqual(elec.iloc[0]['x'], 1.0)
@@ -254,6 +266,45 @@ def test_add_electrode(self):
         self.assertEqual(elec.iloc[0]['filtering'], 'none')
         self.assertEqual(elec.iloc[0]['group'], group)
 
+    def test_add_electrode_some_opt(self):
+        dev1 = self.nwbfile.create_device(name='dev1')
+        group = self.nwbfile.create_electrode_group(
+            name='tetrode1',
+            description='tetrode description',
+            location='tetrode location',
+            device=dev1
+        )
+        self.nwbfile.add_electrode(
+            x=1.0, y=2.0, z=3.0,
+            imp=-1.0,
+            location='CA1',
+            filtering='none',
+            group=group,
+            id=1,
+            rel_x=4.0, rel_y=5.0, rel_z=6.0,
+            reference='ref1'
+        )
+        self.nwbfile.add_electrode(
+            x=1.0, y=2.0, z=3.0,
+            imp=-1.0,
+            location='CA1',
+            filtering='none',
+            group=group,
+            id=2,
+            rel_x=7.0, rel_y=8.0, rel_z=9.0,
+            reference='ref2'
+        )
+        elec = self.nwbfile.electrodes[0]
+        self.assertEqual(elec.iloc[0]['rel_x'], 4.0)
+        self.assertEqual(elec.iloc[0]['rel_y'], 5.0)
+        self.assertEqual(elec.iloc[0]['rel_z'], 6.0)
+        self.assertEqual(elec.iloc[0]['reference'], 'ref1')
+        elec = self.nwbfile.electrodes[1]
+        self.assertEqual(elec.iloc[0]['rel_x'], 7.0)
+        self.assertEqual(elec.iloc[0]['rel_y'], 8.0)
+        self.assertEqual(elec.iloc[0]['rel_z'], 9.0)
+        self.assertEqual(elec.iloc[0]['reference'], 'ref2')
+
     def test_all_children(self):
         ts1 = TimeSeries('test_ts1', [0, 1, 2, 3, 4, 5], 'grams', timestamps=[0.0, 0.1, 0.2, 0.3, 0.4, 0.5])
         ts2 = TimeSeries('test_ts2', [0, 1, 2, 3, 4, 5], 'grams', timestamps=[0.0, 0.1, 0.2, 0.3, 0.4, 0.5])