WIP: Add dipole recording for individual cells

hnn_core/cell_response.py

@@ -56,6 +56,9 @@ class CellResponse(object):
     isec : list (n_trials,) of dict, shape
         Each element of the outer list is a trial.
         Dictionary indexed by gids containing currents for cell sections.
+    dcell : list (n_trials,) of dict, shape
+        Each element of the outer list is a trial.
+        Dictionary indexed by gids containing dipoles for individual cells.
     times : array-like, shape (n_times,)
         Array of time points for samples in continuous data.
         This includes vsoma and isoma.
@@ -115,6 +118,7 @@ def __init__(self, spike_times=None, spike_gids=None, spike_types=None,
         self._spike_types = spike_types
         self._vsec = list()
         self._isec = list()
+        self._dcell = list()
         if times is not None:
             if not isinstance(times, (list, np.ndarray)):
                 raise TypeError("'times' is an np.ndarray of simulation times")
@@ -226,6 +230,10 @@ def vsec(self):
     def isec(self):
         return self._isec
 
+    @property
+    def dcell(self):
+        return self._dcell
+
     @property
     def times(self):
         return self._times
@@ -423,6 +431,10 @@ def to_dict(self):
             # Turn `int` gid keys into string values for hdf5 format
             trial = dict((str(key), val) for key, val in trial.items())
             cell_response_data['isec'].append(trial)
+        for trial in self.dcell:
+            # Turn `int` gid keys into string values for hdf5 format
+            trial = dict((str(key), val) for key, val in trial.items())
+            cell_response_data['dcell'].append(trial)
         cell_response_data['times'] = self.times
         return cell_response_data
 

hnn_core/dipole.py

@@ -16,7 +16,7 @@
 
 
 def simulate_dipole(net, tstop, dt=0.025, n_trials=None, record_vsec=False,
-                    record_isec=False, postproc=False):
+                    record_isec=False, record_dcell=False, postproc=False):
     """Simulate a dipole given the experiment parameters.
 
     Parameters
@@ -37,6 +37,8 @@ def simulate_dipole(net, tstop, dt=0.025, n_trials=None, record_vsec=False,
     record_isec : 'all' | 'soma' | False
         Option to record voltages from all sections ('all'), or just
         the soma ('soma'). Default: False.
+    record_dcell : bool
+        Option to record dipole from individual cells. Default: False.
     postproc : bool
         If True, smoothing (``dipole_smooth_win``) and scaling
         (``dipole_scalefctr``) values are read from the parameter file, and
@@ -96,6 +98,10 @@ def simulate_dipole(net, tstop, dt=0.025, n_trials=None, record_vsec=False,
 
     net._params['record_isec'] = record_isec
 
+    _check_option('record_dcell', record_dcell, [True, False])
+
+    net._params['record_dcell'] = record_dcell
+
     if postproc:
         warnings.warn('The postproc-argument is deprecated and will be removed'
                       ' in a future release of hnn-core. Please define '

hnn_core/drives.py

@@ -10,7 +10,8 @@
                      _extract_drive_specs_from_hnn_params)
 
 
-def _get_target_properties(weights_ampa, weights_nmda, synaptic_delays,
+def _get_target_properties(weights_ampa, weights_nmda, weights_gabaa,
+                           weights_gababb, synaptic_delays,
                            location, probability=1.0):
     """Retrieve drive properties associated with each target cell type
 
@@ -23,13 +24,23 @@ def _get_target_properties(weights_ampa, weights_nmda, synaptic_delays,
         weights_ampa = dict()
     if weights_nmda is None:
         weights_nmda = dict()
+    if weights_gabaa is None:
+        weights_gabaa = dict()
+    if weights_gababb is None:
+        weights_gababb = dict()
 
     weights_by_type = {cell_type: dict() for cell_type in
-                       (set(weights_ampa.keys()) | set(weights_nmda.keys()))}
+                       (set(weights_ampa.keys()) | set(weights_nmda.keys()) |
+                        set(weights_gabaa.keys()) |
+                        set(weights_gababb.keys()))}
     for cell_type in weights_ampa:
         weights_by_type[cell_type].update({'ampa': weights_ampa[cell_type]})
     for cell_type in weights_nmda:
         weights_by_type[cell_type].update({'nmda': weights_nmda[cell_type]})
+    for cell_type in weights_gabaa:
+        weights_by_type[cell_type].update({'gabaa': weights_gabaa[cell_type]})
+    for cell_type in weights_gababb:
+        weights_by_type[cell_type].update({'gabab': weights_gababb[cell_type]})
 
     target_populations = set(weights_by_type)
     if not target_populations:

hnn_core/network.py

@@ -485,6 +485,7 @@ def copy(self):
     def add_evoked_drive(self, name, *, mu, sigma, numspikes, location,
                          n_drive_cells='n_cells', cell_specific=True,
                          weights_ampa=None, weights_nmda=None,
+                         weights_gabaa=None, weights_gabab=None,
                          space_constant=3., synaptic_delays=0.1,
                          probability=1.0, event_seed=2, conn_seed=3):
         """Add an 'evoked' external drive to the network
@@ -528,6 +529,12 @@ def add_evoked_drive(self, name, *, mu, sigma, numspikes, location,
         weights_nmda : dict or None
             Synaptic weights (in uS) of NMDA receptors on each targeted cell
             type (dict keys). Cell types omitted from the dict are set to zero.
+        weights_gabaa : dict or None
+            Synaptic weights (in uS) of GABAa receptors on each targeted cell
+            type (dict keys). Cell types omitted from the dict are set to zero.
+        weights_gabab : dict or None
+            Synaptic weights (in uS) of GABAb receptors on each targeted cell
+            type (dict keys). Cell types omitted from the dict are set to zero.
         synaptic_delays : dict or float
             Synaptic delay (in ms) at the column origin, dispersed laterally as
             a function of the space_constant. If float, applies to all target
@@ -580,14 +587,16 @@ def add_evoked_drive(self, name, *, mu, sigma, numspikes, location,
         drive['dynamics'] = dict(mu=mu, sigma=sigma, numspikes=numspikes)
         drive['events'] = list()
 
-        self._attach_drive(name, drive, weights_ampa, weights_nmda, location,
+        self._attach_drive(name, drive, weights_ampa, weights_nmda,
+                           weights_gabaa, weights_gabab, location,
                            space_constant, synaptic_delays,
                            n_drive_cells, cell_specific, probability)
 
     def add_poisson_drive(self, name, *, tstart=0, tstop=None, rate_constant,
                           location, n_drive_cells='n_cells',
                           cell_specific=True, weights_ampa=None,
-                          weights_nmda=None, space_constant=100.,
+                          weights_nmda=None, weights_gabaa=None,
+                          weights_gabab=None, space_constant=100.,
                           synaptic_delays=0.1, probability=1.0, event_seed=2,
                           conn_seed=3):
         """Add a Poisson-distributed external drive to the network
@@ -635,6 +644,12 @@ def add_poisson_drive(self, name, *, tstart=0, tstop=None, rate_constant,
         weights_nmda : dict or None
             Synaptic weights (in uS) of NMDA receptors on each targeted cell
             type (dict keys). Cell types omitted from the dict are set to zero.
+        weights_gabaa : dict or None
+            Synaptic weights (in uS) of GABAa receptors on each targeted cell
+            type (dict keys). Cell types omitted from the dict are set to zero.
+        weights_gabab : dict or None
+            Synaptic weights (in uS) of GABAb receptors on each targeted cell
+            type (dict keys). Cell types omitted from the dict are set to zero.
         synaptic_delays : dict or float
             Synaptic delay (in ms) at the column origin, dispersed laterally as
             a function of the space_constant. If float, applies to all target
@@ -664,6 +679,8 @@ def add_poisson_drive(self, name, *, tstart=0, tstop=None, rate_constant,
                                       tstop=tstop)
         target_populations = _get_target_properties(weights_ampa,
                                                     weights_nmda,
+                                                    weights_gabaa,
+                                                    weights_gabab,
                                                     synaptic_delays,
                                                     location)[0]
         _check_poisson_rates(rate_constant, target_populations,
@@ -689,14 +706,16 @@ def add_poisson_drive(self, name, *, tstart=0, tstop=None, rate_constant,
                                  rate_constant=rate_constant)
         drive['events'] = list()
 
-        self._attach_drive(name, drive, weights_ampa, weights_nmda, location,
+        self._attach_drive(name, drive, weights_ampa, weights_nmda,
+                           weights_gabaa, weights_gabab, location,
                            space_constant, synaptic_delays,
                            n_drive_cells, cell_specific, probability)
 
     def add_bursty_drive(self, name, *, tstart=0, tstart_std=0, tstop=None,
                          location, burst_rate, burst_std=0, numspikes=2,
                          spike_isi=10, n_drive_cells=1, cell_specific=False,
                          weights_ampa=None, weights_nmda=None,
+                         weights_gabaa=None, weights_gabab=None,
                          synaptic_delays=0.1, space_constant=100.,
                          probability=1.0, event_seed=2, conn_seed=3):
         """Add a bursty (rhythmic) external drive to all cells of the network
@@ -751,6 +770,12 @@ def add_bursty_drive(self, name, *, tstart=0, tstart_std=0, tstop=None,
         weights_nmda : dict or None
             Synaptic weights (in uS) of NMDA receptors on each targeted cell
             type (dict keys). Cell types omitted from the dict are set to zero.
+        weights_gabaa : dict or None
+            Synaptic weights (in uS) of GABAa receptors on each targeted cell
+            type (dict keys). Cell types omitted from the dict are set to zero.
+        weights_gabab : dict or None
+            Synaptic weights (in uS) of GABAb receptors on each targeted cell
+            type (dict keys). Cell types omitted from the dict are set to zero.
         synaptic_delays : dict or float
             Synaptic delay (in ms) at the column origin, dispersed laterally as
             a function of the space_constant. If float, applies to all target
@@ -795,11 +820,13 @@ def add_bursty_drive(self, name, *, tstart=0, tstart_std=0, tstop=None,
                                  numspikes=numspikes, spike_isi=spike_isi)
         drive['events'] = list()
 
-        self._attach_drive(name, drive, weights_ampa, weights_nmda, location,
+        self._attach_drive(name, drive, weights_ampa, weights_nmda,
+                           weights_gabaa, weights_gabab, location,
                            space_constant, synaptic_delays,
                            n_drive_cells, cell_specific, probability)
 
-    def _attach_drive(self, name, drive, weights_ampa, weights_nmda, location,
+    def _attach_drive(self, name, drive, weights_ampa, weights_nmda,
+                      weights_gabaa, weights_gabab, location,
                       space_constant, synaptic_delays, n_drive_cells,
                       cell_specific, probability):
         """Attach a drive to network based on connectivity information
@@ -816,6 +843,12 @@ def _attach_drive(self, name, drive, weights_ampa, weights_nmda, location,
         weights_nmda : dict or None
             Synaptic weights (in uS) of NMDA receptors on each targeted cell
             type (dict keys). Cell types omitted from the dict are set to zero.
+        weights_gabaa : dict or None
+            Synaptic weights (in uS) of GABAa receptors on each targeted cell
+            type (dict keys). Cell types omitted from the dict are set to zero.
+        weights_gabab : dict or None
+            Synaptic weights (in uS) of GABAb receptors on each targeted cell
+            type (dict keys). Cell types omitted from the dict are set to zero.
         location : str
             Target location of synapses. Must be an element of
             `Cell.sect_loc` such as 'proximal' or 'distal', which defines a
@@ -866,7 +899,8 @@ def _attach_drive(self, name, drive, weights_ampa, weights_nmda, location,
         # allow passing weights as None, convert to dict here
         (target_populations, weights_by_type, delays_by_type,
          probability_by_type) = \
-            _get_target_properties(weights_ampa, weights_nmda, synaptic_delays,
+            _get_target_properties(weights_ampa, weights_nmda, weights_gabaa,
+                                   weights_gabab, synaptic_delays,
                                    location, probability)
 
         # weights passed must correspond to cells in the network
@@ -977,8 +1011,8 @@ def _attach_drive(self, name, drive, weights_ampa, weights_nmda, location,
                         receptor=receptor, weight=weights, delay=delays,
                         lamtha=space_constant, probability=probability,
                         conn_seed=drive['conn_seed'] + seed_increment)
-                    # Ensure that AMPA/NMDA connections target the same gids
-                    # when probability < 1
+                    # Ensure that multireceptor connections target the
+                    # same gids when probability < 1
                     if receptor_idx > 0:
                         self.connectivity[-1]['src_gids'] = \
                             self.connectivity[-2]['src_gids']

hnn_core/network_builder.py

@@ -100,6 +100,10 @@ def simulation_time():
             isec_py[gid][sec_name] = {
                 key: isec.to_python() for key, isec in isec.items()}
 
+    dcell_py = dict()
+    for gid, dcell in neuron_net._dcell.items():
+        dcell_py[gid] = dcell.to_python()
+
     dpl_data = np.c_[
         neuron_net._nrn_dipoles['L2_pyramidal'].as_numpy() +
         neuron_net._nrn_dipoles['L5_pyramidal'].as_numpy(),
@@ -119,6 +123,7 @@ def simulation_time():
             'gid_ranges': net.gid_ranges,
             'vsec': vsec_py,
             'isec': isec_py,
+            'dcell': dcell_py,
             'rec_data': rec_arr_py,
             'rec_times': rec_times_py,
             'times': times.to_python()}
@@ -291,6 +296,7 @@ def __init__(self, net, trial_idx=0):
 
         self._vsec = dict()
         self._isec = dict()
+        self._dcell = dict()
         self._nrn_rec_arrays = dict()
         self._nrn_rec_callbacks = list()
 
@@ -562,6 +568,9 @@ def aggregate_data(self, n_samples):
                 nrn_dpl = self._nrn_dipoles[_long_name(cell.name)]
                 nrn_dpl.add(cell.dipole)
 
+                if self.net._params['record_dcell']:
+                    self._dcell[cell.gid] = cell.dipole
+
             self._vsec[cell.gid] = cell.vsec
             self._isec[cell.gid] = cell.isec
 
@@ -574,6 +583,7 @@ def aggregate_data(self, n_samples):
         # aggregate the currents and voltages independently on each proc
         vsec_list = _PC.py_gather(self._vsec, 0)
         isec_list = _PC.py_gather(self._isec, 0)
+        dcell_list = _PC.py_gather(self._dcell, 0)
 
         # combine spiking data from each proc
         spike_times_list = _PC.py_gather(self._spike_times, 0)
@@ -589,6 +599,8 @@ def aggregate_data(self, n_samples):
                 self._vsec.update(vsec)
             for isec in isec_list:
                 self._isec.update(isec)
+            for dcell in dcell_list:
+                self._dcell.update(dcell)
 
         _PC.barrier()  # get all nodes to this place before continuing
 

hnn_core/parallel_backends.py

@@ -53,6 +53,7 @@ def _gather_trial_data(sim_data, net, n_trials, postproc):
         net.cell_response.update_types(net.gid_ranges)
         net.cell_response._vsec.append(sim_data[idx]['vsec'])
         net.cell_response._isec.append(sim_data[idx]['isec'])
+        net.cell_response._dcell.append(sim_data[idx]['dcell'])
 
         # extracellular array
         for arr_name, arr in net.rec_arrays.items():

hnn_core/params_default.py

@@ -149,6 +149,7 @@ def get_params_default(nprox=2, ndist=1):
         'record_isoma': 0,  # whether to record somatic currents
         'record_vsec': 0,  # whether to record voltages
         'record_isec': 0,  # whether to record currents
+        'record_dcell': 0,  # whether to record cell dipoles
 
         # numerics
         # N_trials of 1 means that seed is set by rank

hnn_core/tests/conftest.py

@@ -75,8 +75,8 @@ def pytest_runtest_setup(item):
 def run_hnn_core_fixture():
     def _run_hnn_core_fixture(backend=None, n_procs=None, n_jobs=1,
                               reduced=False, record_vsec=False,
-                              record_isec=False, postproc=False,
-                              electrode_array=None):
+                              record_isec=False, record_dcell=False,
+                              postproc=False, electrode_array=None):
         hnn_core_root = op.dirname(hnn_core.__file__)
 
         # default params
@@ -105,15 +105,18 @@ def _run_hnn_core_fixture(backend=None, n_procs=None, n_jobs=1,
             with MPIBackend(n_procs=n_procs, mpi_cmd='mpiexec'):
                 dpls = simulate_dipole(net, record_vsec=record_vsec,
                                        record_isec=record_isec,
+                                       record_dcell=record_dcell,
                                        postproc=postproc, tstop=tstop)
         elif backend == 'joblib':
             with JoblibBackend(n_jobs=n_jobs):
                 dpls = simulate_dipole(net, record_vsec=record_vsec,
                                        record_isec=record_isec,
+                                       record_dcell=record_dcell,
                                        postproc=postproc, tstop=tstop)
         else:
             dpls = simulate_dipole(net, record_vsec=record_vsec,
                                    record_isec=record_isec,
+                                   record_dcell=record_dcell,
                                    postproc=postproc, tstop=tstop)
 
         # check that the network object is picklable after the simulation

hnn_core/tests/test_cell_response.py

@@ -43,7 +43,7 @@ def test_cell_response(tmp_path):
     # reset clears all recorded variables, but leaves simulation time intact
     assert len(cell_response.times) == len(sim_times)
     sim_attributes = ['_spike_times', '_spike_gids', '_spike_types',
-                      '_vsec', '_isec']
+                      '_vsec', '_isec', '_dcell']
     net_attributes = ['_times', '_cell_type_names']  # `Network.__init__`
     # creates these check that we always know which response attributes are
     # simulated see #291 for discussion; objective is to keep cell_response