[MRG] BUG: Update gid_ranges after using clear_drives

doc/whats_new.rst

@@ -33,11 +33,11 @@ Changelog
 - Fixed figure annotation overlap in multiple sub-plots, 
   by `Camilo Diaz`_ in :gh:`741`
 
-- Fix bug in :func:`~hnn_core/network/pick_connection` where connections are
+- Fix bug in :func:`~hnn_core.network.pick_connection` where connections are
   returned for cases when there should be no valid matches, by `George Dang`_
   in :gh:`739`
 
-- Added check for invalid Axes object in :func:`~hnn_core/viz/plot_cells` 
+- Added check for invalid Axes object in :func:`~hnn_core.viz.plot_cells` 
   function, by `Abdul Samad Siddiqui`_ in :gh:`744`.
 
 - Added kwargs options to `plot_spikes_hist` for adjusting the histogram plots 
@@ -55,7 +55,7 @@ Changelog
 - Added feature to read/write :class:`~hnn_core.Network` configurations to
   json, by `George Dang`_ and `Rajat Partani`_ in :gh:`757`
 
-- Added :class:`~hnn_core/viz/NetworkPlotter` to visualize and animate network simulations,
+- Added :class:`~hnn_core.viz.NetworkPlotter` to visualize and animate network simulations,
   by `Nick Tolley`_ in :gh:`649`.
 
 - Added GUI feature to include Tonic input drives in simulations,
@@ -65,7 +65,7 @@ Changelog
   and :func:`~plot_spikes_raster` now plotted from 0 to tstop. Inputs tmin and tmax are deprecated,
   by `Katharina Duecker`_ in :gh:`769`
 
-- Add function :func:`~hnn_core/params/convert_to_json` to convert legacy param
+- Add function :func:`~hnn_core.params.convert_to_json` to convert legacy param
   and json files to new json format, by `George Dang`_ in :gh:`772`
 
 - Add :class:`~hnn_core.BatchSimulate` for batch simulation capability,
@@ -79,7 +79,7 @@ Bug
 - Fix inconsistent connection mapping from drive gids to cell gids, by
   `Ryan Thorpe`_ in :gh:`642`.
 
-- Objective function called by :func:`~hnn_core/optimization/optimize_evoked`
+- Objective function called by :func:`~hnn_core.optimization.optimize_evoked`
   now returns a scalar instead of tuple, by `Ryan Thorpe`_ in :gh:`670`.
 
 - Fix GUI plotting bug due to deprecation of matplotlib color cycling method,
@@ -94,6 +94,9 @@ Bug
 - Fix drive seeding so that event times are unique across multiple trials,
   by `Nick Tolley`_ in :gh:`810`.
 
+- Fix bug in :func:`~hnn_core.network.clear_drives` where network object are not
+  accurately updated, by `Nick Tolley`_ in :gh:`812`.
+
 API
 ~~~
 - :func:`~hnn_core.CellResponse.write` and :func:`~hnn_core.Cell_response.read_spikes`

hnn_core/network.py

@@ -1374,12 +1374,17 @@ def clear_connectivity(self):
 
     def clear_drives(self):
         """Remove all drives defined in Network.connectivity"""
-        connectivity = list()
-        for conn in self.connectivity:
-            if conn['src_type'] not in self.external_drives.keys():
-                connectivity.append(conn)
+        self.connectivity = [conn for conn in self.connectivity if
+                             conn['src_type'] not
+                             in self.external_drives.keys()]
+
+        for cell_name in list(self.gid_ranges.keys()):
+            if cell_name in self.external_drives:
+                self._n_gids -= len(self.gid_ranges[cell_name])
+                del self.gid_ranges[cell_name]
+                del self.pos_dict[cell_name]
+
         self.external_drives = dict()
-        self.connectivity = connectivity
 
     def add_electrode_array(self, name, electrode_pos, *, conductivity=0.3,
                             method='psa', min_distance=0.5):

hnn_core/tests/test_drives.py

@@ -98,8 +98,8 @@ def test_external_drive_times():
 def test_drive_seeds(setup_net):
     """Test that unique spike times are generated across trials"""
     net = setup_net
-    weights_ampa = {'L2_basket': 0.000003, 'L2_pyramidal': 1.438840,
-                    'L5_basket': 0.008958, 'L5_pyramidal': 0.684013}
+    weights_ampa = {'L2_basket': 0.3, 'L2_pyramidal': 0.3,
+                    'L5_basket': 0.3, 'L5_pyramidal': 0.3}
     synaptic_delays = {'L2_basket': 0.1, 'L2_pyramidal': 0.1,
                        'L5_basket': 1., 'L5_pyramidal': 1.}
     net.add_evoked_drive(
@@ -116,6 +116,53 @@ def test_drive_seeds(setup_net):
     assert ~np.any(np.allclose(trial1_spikes, trial2_spikes))
 
 
+def test_clear_drives(setup_net):
+    """Test clearing drives updates Network"""
+    net = setup_net
+    weights_ampa = {'L5_pyramidal': 0.3}
+    synaptic_delays = {'L5_pyramidal': 1.}
+
+    # Test attributes after adding 2 drives
+    n_gids = net._n_gids
+    net.add_evoked_drive(
+        'prox', mu=40, sigma=8.33, numspikes=1,
+        weights_ampa=weights_ampa, location='proximal',
+        synaptic_delays=synaptic_delays, cell_specific=True)
+
+    net.add_evoked_drive(
+        'dist', mu=40, sigma=8.33, numspikes=1,
+        weights_ampa=weights_ampa, location='distal',
+        synaptic_delays=synaptic_delays, cell_specific=True)
+
+    for drive_name in ['prox', 'dist']:
+        assert len(net.external_drives) == 2
+        assert drive_name in net.external_drives
+        assert drive_name in net.gid_ranges
+        assert drive_name in net.pos_dict
+        assert net._n_gids == n_gids + len(net.gid_ranges['L5_pyramidal']) * 2
+
+    # Test attributes after clearing drives
+    net.clear_drives()
+    for drive_name in ['prox', 'dist']:
+        assert len(net.external_drives) == 0
+        assert drive_name not in net.external_drives
+        assert drive_name not in net.gid_ranges
+        assert drive_name not in net.pos_dict
+        assert net._n_gids == n_gids
+
+    # Test attributes after adding 1 drive
+    net.add_evoked_drive(
+        'prox', mu=40, sigma=8.33, numspikes=1,
+        weights_ampa=weights_ampa, location='proximal',
+        synaptic_delays=synaptic_delays, cell_specific=True)
+
+    assert len(net.external_drives) == 1
+    assert 'prox' in net.external_drives
+    assert 'prox' in net.gid_ranges
+    assert 'prox' in net.pos_dict
+    assert net._n_gids == n_gids + len(net.gid_ranges['L5_pyramidal'])
+
+
 def test_add_drives():
     """Test methods for adding drives to a Network."""
     hnn_core_root = op.dirname(hnn_core.__file__)