def test_ElectrodeArray_add_electrode(): earray = ElectrodeArray([]) npt.assert_equal(earray.n_electrodes, 0) with pytest.raises(TypeError): earray.add_electrode('A01', ElectrodeArray([])) # Add an electrode: key0 = 'A04' earray.add_electrode(key0, PointSource(0, 1, 2)) npt.assert_equal(earray.n_electrodes, 1) # Both numeric and string index should work: for key in [key0, 0]: npt.assert_equal(isinstance(earray[key], PointSource), True) npt.assert_almost_equal(earray[key].x, 0) npt.assert_almost_equal(earray[key].y, 1) npt.assert_almost_equal(earray[key].z, 2) with pytest.raises(ValueError): # Can't add the same electrode twice: earray.add_electrode(key0, PointSource(0, 1, 2)) # Add another electrode: key1 = 'A01' earray.add_electrode(key1, DiskElectrode(4, 5, 6, 7)) npt.assert_equal(earray.n_electrodes, 2) # Both numeric and string index should work: for key in [key1, 1]: npt.assert_equal(isinstance(earray[key], DiskElectrode), True) npt.assert_almost_equal(earray[key].x, 4) npt.assert_almost_equal(earray[key].y, 5) npt.assert_almost_equal(earray[key].z, 6) npt.assert_almost_equal(earray[key].r, 7) # We can also get a list of electrodes: for keys in [[key0, key1], [0, key1], [key0, 1], [0, 1]]: selected = earray[keys] npt.assert_equal(isinstance(selected, list), True) npt.assert_equal(isinstance(selected[0], PointSource), True) npt.assert_equal(isinstance(selected[1], DiskElectrode), True)
def test_ElectrodeArray_remove_electrode(): earray1 = ElectrodeArray([]) earray2 = ElectrodeArray([]) npt.assert_equal(earray1.n_electrodes, 0) # Can't remove electrodes from empty electrodeArray with pytest.raises(ValueError): earray1.remove_electrode(None) with pytest.raises(ValueError): earray1.remove_electrode("foo") key = [0] * 4 key[0] = 'D03' key[1] = 'A02' key[2] = 'F10' key[3] = 'E12' earray1.add_electrode(key[0], PointSource(0, 1, 2)) earray1.add_electrode(key[1], PointSource(3, 4, 5)) earray1.add_electrode(key[2], PointSource(6, 7, 8)) earray1.add_electrode(key[3], PointSource(9, 10, 11)) npt.assert_equal(earray1.n_electrodes, 4) earray2.add_electrode(key[0], PointSource(0, 1, 2)) earray2.add_electrode(key[1], PointSource(3, 4, 5)) earray2.add_electrode(key[2], PointSource(6, 7, 8)) earray2.add_electrode(key[3], PointSource(9, 10, 11)) npt.assert_equal(earray2.n_electrodes, 4) # Remove one electrode key[1] from the electrodeArray earray1.remove_electrode(key[0]) npt.assert_equal(earray1.n_electrodes, 3) # Can't remove an electrode that has been removed with pytest.raises(ValueError): earray1.remove_electrode(key[0]) # List keeps order: npt.assert_equal(earray1[0], earray1[key[1]]) npt.assert_equal(earray1[1], earray1[key[2]]) npt.assert_equal(earray1[2], earray1[key[3]]) # Other electrodes stay the same for k in [key[1], key[2], key[3]]: npt.assert_equal(earray1[k].x, earray2[k].x) npt.assert_equal(earray1[k].y, earray2[k].y) npt.assert_equal(earray1[k].z, earray2[k].z) # Remove two more electrodes from the electrodeArray # List keeps order earray1.remove_electrode(key[1]) earray1.remove_electrode(key[2]) npt.assert_equal(earray1.n_electrodes, 1) npt.assert_equal(earray1[0], earray1[key[3]]) # The last electrode stays the same for key in [key[3]]: npt.assert_equal(earray1[key].x, earray2[key].x) npt.assert_equal(earray1[key].y, earray2[key].y) npt.assert_equal(earray1[key].z, earray2[key].z)