[ENH] Fixed bug in clear_connectivity and improved tests for deleting drives

examples/howto/plot_connectivity.py

@@ -71,7 +71,7 @@
 # directly.
 def get_network(probability=1.0):
     net = jones_2009_model(add_drives_from_params=True)
-    net.clear_connectivity()
+    net.connectivity = list()
 
     # Pyramidal cell connections
     location, receptor = 'distal', 'ampa'

hnn_core/network.py

@@ -1253,23 +1253,49 @@ def add_connection(self, src_gids, target_gids, loc, receptor,
 
         self.connectivity.append(deepcopy(conn))
 
-    def clear_connectivity(self):
-        """Remove all connections defined in Network.connectivity
+    def clear_connectivity(self, src_types=None):
+        """Remove connections with src_type in Network.connectivity.
+
+        Parameters
+        ----------
+        src_types : list | None
+            Source types of connections to be removed
         """
-        connectivity = list()
+        if src_types is None:
+            src_types = list()
+            # Storing all external drives
+            src_types_external_drives = self.external_drives.keys()
+            for conn in self.connectivity:
+                src_type = conn['src_type']
+                if src_type not in src_types_external_drives:
+                    src_types.append(src_type)  # Storing drives to be deleted
+            src_types = list(set(src_types))  # Removing duplicate entries
+        connectivity = list()  # Initialize empty list
         for conn in self.connectivity:
-            if conn['src_type'] in self.external_drives.keys():
+            # Removes connections in src_types
+            if conn['src_type'] not in src_types:
                 connectivity.append(conn)
         self.connectivity = connectivity
 
-    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.external_drives = dict()
-        self.connectivity = connectivity
+    def clear_drives(self, drive_names='all'):
+        """Remove all drives defined in Network.connectivity.
+
+        Parameters
+        ----------
+        drive_names : list | 'all'
+            The drive_names to remove
+        """
+        if drive_names == 'all':
+            drive_names = list(self.external_drives.keys())
+        _validate_type(drive_names, (list,))
+        for drive_name in drive_names:
+            del self.external_drives[drive_name]
+        self.clear_connectivity(src_types=drive_names)
+
+    @property
+    def drive_names(self):
+        """Returns a list containing names of all external drives."""
+        return list(self.external_drives.keys())
 
     def add_electrode_array(self, name, electrode_pos, *, conductivity=0.3,
                             method='psa', min_distance=0.5):

hnn_core/tests/test_network.py

@@ -784,8 +784,26 @@ def test_network_connectivity():
     # Needs to be updated if number of drives change in preceeding tests
     net.clear_connectivity()
     assert len(net.connectivity) == 4  # 2 drives x 2 target cell types
+
+    # Only external drive connections are left
+    # Testing deletion of a custom number of drives
+
+    # Deleting one external drive
+    all_drives = net.drive_names
+    drives_to_be_deleted = all_drives[0:1]
+    connectivities_to_be_deleted = 0
+    for drive in drives_to_be_deleted:
+        drive_connections = len(pick_connection(net, src_gids=drive))
+        connectivities_to_be_deleted = (connectivities_to_be_deleted +
+                                        drive_connections)
+    expected_connectivity_length = (len(net.connectivity) -
+                                    connectivities_to_be_deleted)
+    net.clear_drives(drive_names=drives_to_be_deleted)
+    assert len(net.connectivity) == expected_connectivity_length
+
+    # Deleting all remaining external drives
     net.clear_drives()
-    assert len(net.connectivity) == 0
+    assert len(net.connectivity) == 0  # All connections are deleted
 
     with pytest.warns(UserWarning, match='No connections'):
         simulate_dipole(net, tstop=10)