Allow drives to target individual sections

hnn_core/network.py

@@ -487,7 +487,10 @@ def add_evoked_drive(self, name, *, mu, sigma, numspikes, location,
         numspikes : int
             Number of spikes at each target cell
         location : str
-            Target location of synapses ('distal' or 'proximal')
+            Target location of synapses. Must be an element of
+            `Cell.sect_loc` such as 'proximal' or 'distal', which defines a
+            group of sections, or an existing section such as 'soma' or
+            'apical_tuft' (defined in `Cell.sections` for all targeted cells)
         n_drive_cells : int | 'n_cells'
             The number of drive cells that each contribute an independently
             sampled synaptic spike to the network according to the Gaussian
@@ -590,7 +593,10 @@ def add_poisson_drive(self, name, *, tstart=0, tstop=None, rate_constant,
             provided as a dictionary, in which a key must be present for each
             cell type with non-zero AMPA or NMDA weights.
         location : str
-            Target location of synapses ('distal' or 'proximal')
+            Target location of synapses. Must be an element of
+            `Cell.sect_loc` such as 'proximal' or 'distal', which defines a
+            group of sections, or an existing section such as 'soma' or
+            'apical_tuft' (defined in `Cell.sections` for all targeted cells)
         n_drive_cells : int | 'n_cells'
             The number of drive cells that each contribute an independently
             sampled synaptic spike to the network according to a Poisson
@@ -692,7 +698,10 @@ def add_bursty_drive(self, name, *, tstart=0, tstart_std=0, tstop=None,
             End time of burst trains (defaults to None: tstop is set to the
             end of the simulation)
         location : str
-            Target location of synapses ('distal' or 'proximal')
+            Target location of synapses. Must be an element of
+            `Cell.sect_loc` such as 'proximal' or 'distal', which defines a
+            group of sections, or an existing section such as 'soma' or
+            'apical_tuft' (defined in `Cell.sections` for all targeted cells)
         burst_rate : float
             The mean rate at which cyclic bursts occur (unit: Hz)
         burst_std : float
@@ -799,7 +808,7 @@ def _attach_drive(self, name, drive, weights_ampa, weights_nmda, location,
             ``exp(-(x / (3 * inplane_distance)) ** 2)``, where ``x`` is the
             physical distance (in um) between the connected cells in the xy
             plane.
-        synaptic_delays : dict
+        synaptic_delays : dict or float
             Synaptic delay (in ms) at the column origin, dispersed laterally as
             a function of the space_constant
         n_drive_cells : int | 'n_cells'
@@ -827,9 +836,19 @@ def _attach_drive(self, name, drive, weights_ampa, weights_nmda, location,
         """
         if name in self.external_drives:
             raise ValueError(f"Drive {name} already defined")
-        if location not in ['distal', 'proximal']:
-            raise ValueError("Allowed drive target locations are: 'distal', "
-                             f"and 'proximal', got {location}")
+
+        # cell_sections = [set(self.cell_types[cell_type].sections.keys()) for
+        #                  cell_type in drive['target_types']]
+        # sect_locs = [set(self.cell_types[cell_type].sect_loc.keys()) for
+        #              cell_type in drive['target_types']]
+
+        # valid_cell_sections = set.intersection(*cell_sections)
+        # valid_sect_locs = set.intersection(*sect_locs)
+        # valid_loc = list(valid_cell_sections) + list(valid_sect_locs)
+
+        # _check_option('location', location, valid_loc,
+        #               extra=(f" (the location '{location}' is not defined "
+        #                      "for one of the targeted cells)"))
 
         _validate_type(
             probability, (float, dict), 'probability', 'float or dict')
@@ -844,6 +863,20 @@ def _attach_drive(self, name, drive, weights_ampa, weights_nmda, location,
             raise ValueError('Allowed drive target cell types are: ',
                              f'{self.cell_types.keys()}')
 
+        # Ensure location exists for all target cells
+        cell_sections = [set(self.cell_types[cell_type].sections.keys()) for
+                         cell_type in target_populations]
+        sect_locs = [set(self.cell_types[cell_type].sect_loc.keys()) for
+                     cell_type in target_populations]
+
+        valid_cell_sections = set.intersection(*cell_sections)
+        valid_sect_locs = set.intersection(*sect_locs)
+        valid_loc = list(valid_cell_sections) + list(valid_sect_locs)
+
+        _check_option('location', location, valid_loc,
+                      extra=(f" (the location '{location}' is not defined "
+                             "for one of the targeted cells)"))
+
         if self._legacy_mode:
             # allows tests must match HNN GUI output by preserving original
             # gid assignment convention
@@ -1176,7 +1209,9 @@ def add_connection(self, src_gids, target_gids, loc, receptor,
         valid_loc = list(
             target_sect_loc.keys()) + list(target_sections.keys())
 
-        _check_option('loc', loc, valid_loc)
+        _check_option('loc', loc, valid_loc,
+                      extra=(f" (the loc '{loc}' is not defined "
+                             f"for '{target_type}' cells)"))
         conn['loc'] = loc
 
         # `loc` specifies a group of sections, all must contain the synapse

hnn_core/network_models.py

@@ -11,7 +11,8 @@
 from .externals.mne import _validate_type
 
 
-def jones_2009_model(params=None, add_drives_from_params=False):
+def jones_2009_model(params=None, add_drives_from_params=False,
+                     legacy_mode=True):
     """Instantiate the Jones et al. 2009 model.
 
     Parameters
@@ -24,6 +25,9 @@ def jones_2009_model(params=None, add_drives_from_params=False):
         If True, add drives as defined in the params-dict. NB this is mainly
         for backward-compatibility with HNN GUI, and will be deprecated in a
         future release. Default: False
+    legacy_mode : bool
+        Set to True by default to enable matching HNN GUI output when drives
+        are added suitably. Will be deprecated in a future release.
 
     Returns
     -------
@@ -45,7 +49,8 @@ def jones_2009_model(params=None, add_drives_from_params=False):
     if isinstance(params, str):
         params = read_params(params)
 
-    net = Network(params, add_drives_from_params=add_drives_from_params)
+    net = Network(params, add_drives_from_params=add_drives_from_params,
+                  legacy_mode=legacy_mode)
 
     delay = net.delay
 
@@ -159,7 +164,8 @@ def jones_2009_model(params=None, add_drives_from_params=False):
     return net
 
 
-def law_2021_model(params=None, add_drives_from_params=False):
+def law_2021_model(params=None, add_drives_from_params=False,
+                   legacy_mode=True):
     """Instantiate the beta modulated ERP network model.
 
     Returns
@@ -237,7 +243,8 @@ def law_2021_model(params=None, add_drives_from_params=False):
 
 # Remove params argument after updating examples
 # (only relevant for Jones 2009 model)
-def calcium_model(params=None, add_drives_from_params=False):
+def calcium_model(params=None, add_drives_from_params=False,
+                  legacy_mode=True):
     """Instantiate the Jones 2009 model with improved calcium dynamics.
 
     Returns

hnn_core/tests/test_drives.py

@@ -173,6 +173,28 @@ def test_add_drives():
         assert drive_conn['nc_dict']['A_weight'] == weights_ampa[target_type]
         assert drive_conn['nc_dict']['A_delay'] == syn_delays[target_type]
 
+    # Test drive targetting specific section
+    # Section present on all cells indicated
+    location = 'apical_tuft'
+    weights_ampa_tuft = {'L2_pyramidal': 1.0, 'L5_pyramidal': 2.0}
+    syn_delays_tuft = {'L2_pyramidal': 1.0, 'L5_pyramidal': 2.0}
+    net.add_bursty_drive(
+        'bursty_tuft', location=location, burst_rate=10,
+        weights_ampa=weights_ampa_tuft, synaptic_delays=syn_delays_tuft,
+        n_drive_cells=10)
+    assert net.connectivity[-1]['loc'] == location
+
+    # Section not present on cells indicated
+    location = 'apical_tuft'
+    weights_ampa_no_tuft = {'L2_pyramidal': 1.0, 'L5_basket': 2.0}
+    syn_delays_no_tuft = {'L2_pyramidal': 1.0, 'L5_basket': 2.0}
+    match = ('Invalid value for')
+    with pytest.raises(ValueError, match=match):
+        net.add_bursty_drive(
+            'bursty_no_tuft', location=location, burst_rate=10,
+            weights_ampa=weights_ampa_no_tuft,
+            synaptic_delays=syn_delays_no_tuft, n_drive_cells=n_drive_cells)
+
     # Test probabalistic drive connections.
     # drive with cell_specific=False
     n_drive_cells = 10
@@ -234,9 +256,11 @@ def test_add_drives():
 
     # Test Network._attach_drive()
     with pytest.raises(ValueError,
-                       match=r'Allowed drive target locations are'):
+                       match='Invalid value for'):
         net.add_evoked_drive('evdist1', mu=10, sigma=1, numspikes=1,
-                             location='bogus_location')
+                             location='bogus_location',
+                             weights_ampa={'L5_basket': 1.},
+                             synaptic_delays={'L5_basket': .1})
     with pytest.raises(ValueError,
                        match='Drive evoked_dist already defined'):
         net.add_evoked_drive('evoked_dist', mu=10, sigma=1, numspikes=1,
@@ -347,7 +371,7 @@ def test_add_drives():
                               weights_ampa=weights_ampa,
                               synaptic_delays=syn_delays)
     with pytest.raises(ValueError,
-                       match='Allowed drive target locations are:'):
+                       match='Invalid value for the'):
         net.add_poisson_drive('weird_poisson', location='inbetween',
                               rate_constant=10.,
                               weights_ampa=weights_ampa,