API for recording extracellular fields

[ntolley]

As mentioned in #363, I'd be interested in adding the ability to record local field potentials (and possibly EEG/current dipole). Something important to decide on is what the API would look like.

At the most minimal level there wouldn't even be an API, but instead just a tutorial notebook showing how it can be calculated (perhaps with helper functions built intoutils)

For an example of a real API, we use the following for hnn-core:

depths = list(range(-325, 2150, 100))
electrode_pos = [(135, 135, dep) for dep in depths]
net.add_electrode_array('shank1', electrode_pos)

For this repo I think net.record_lfp() might be a better fit considering the existing function names

Alternatively this could be merged with the net.record() functionality that currently exists. Something like

net.record(state='lfp', electrode_pos=electrode_pos)

This is a little confusing since the LFP isn't a state, and it'd add an obligatory argument to this method when state='lfp'.

In any case would love to hear your thoughts on this!

[jnsbck]

Cool that you are thinking about a record_lfp implementation! I think the ability to obtain lfps will be really useful. However, rather than adding this into the jaxley modules directly, I would be for making this functionality more of a util, i.e. utils/lfp.py, since I'd argue that lfps are not a core part of jaxley. In the long-run this could even become sth like jaxley-lfp if it ends up being super big. Also since lfps are essentially just a function of voltage (correct me if I am wrong), I think this should be easily separable from the rest of jaxley, i.e. via sth. like a lfp_of_voltage(voltage, rec_pos, electrode_pos). Lemme know if you disagree @michaeldeistler @ntolley or have other ideas.

[ntolley]

I definitely agree that a utility function seems right for this feature, this is essentially what happened with NEURON and LFPy.

The only information that needs to be tracked is membrane currents (so ion channels and synapses). At the moment I'm enabling it with a code snippet like:

channel_currents = [name for name in net._get_state_names()[0] if name.startswith('i_')]
synapse_currents = [name for name in net._get_state_names()[1]]

for current_name in channel_currents:
    net.record(state=current_name, verbose=False)

for synapse_name in synapse_currents:
    net.record(state=synapse_name, verbose=False)

However, I'm not totally certain about a few things:

1. One of the items in net._get_state_names() is simply 'i'. You can add it through net.record(state='i') but an error is thrown when you run the simulation (possibly a bug?). Not sure what this corresponds to but I'm assuming it must be one of membrane, axial, or stimulation currents?

2. I'm not actually sure how to record synapse currents, has this changed since the discussion in record synaptic currents #363? The code above just pulls out a variable called IonotropicSynapse_s (using default network built in the tutorials), which looks more like the synapse conductance/alpha function. Is there a way to enable recording something like IonotropicSynapse_current'?

Thanks a bunch for the feedback!!

[jnsbck]

Cool, the membrane currents we can just record, for the synapses is currently not implemented, which why #363 is still open.

Regarding your Qs:

1. Good catch! "i" is just the external current and is part of _get_state_names to be able reuse it in stimulate and clamp. In record it throws an error. We should either raise a more specific error in record or rm it from _get_state_names and append it manually in clamp / stimulate.
2. This has not changed since the discussion. To do this we would probably need to add a self.synapse_current_names to Module, similar to membrane_current_names for channels, which we would need to populate in

   jaxley/jaxley/modules/network.py

   Line 563 in 32068e9

   def _append_multiple_synapses(self, pre_nodes, post_nodes, synapse_type):

   . (see

   jaxley/jaxley/modules/base.py

   Line 699 in 32068e9

   def _gather_channels_from_constituents(self, constituents: List):

   for how it is done for channels). Then one could amend

   jaxley/jaxley/modules/base.py

   Line 1198 in 32068e9

   def _get_state_names(self) -> Tuple[List, List]:

   with the synapse currents that we can pull from synapse_current_names and make sure

   jaxley/jaxley/modules/network.py

   Line 254 in 32068e9

   def _step_synapse_state(

   / recording works correctly. :) In general you can implement it much like channel currents are handled since synapses and channels are handled pretty similarly.

Lemme know in case you need any further explanations or have questions along the way :)

[ntolley]

Awesome! This is super helpful, in that case I think the first PR should be restricted to just recording synapse currents as you outlined above. I'll give it a shot and tag you as a reviewer once it's ready

[jnsbck]

Awesome! :)

[ntolley]

@michaeldeistler would a generic callback function be possible to implement? i.e. a function that gets called at every time step during integrate. I'm recalling that this is how it's done in HNN and LFPy, rather than storing all the currents and calculating the results post-hoc.

It may remove the need for fixing #561 for efficiently calculating LFPs

[michaeldeistler]

Yes, I think that should be possible, and might be a nice solution to solve this.

Another option might be to support this only via build_init_and_step_fn. This is demonstrated here. The user can completely modify recordings here:

    recs = jnp.asarray(
        [
            states[rec_state][rec_ind]
            for rec_state, rec_ind in zip(rec_states, rec_inds)
        ]
    )

Michael