jaxleyverse / jaxley

Differentiable neuron simulations with biophysical detail on CPU, GPU, or TPU.
https://jaxley.readthedocs.io
Apache License 2.0
64 stars 11 forks source link

API for recording extracellular fields #515

Open ntolley opened 1 week ago

ntolley commented 1 week ago

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 commented 1 week ago

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 commented 1 week ago

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 #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 commented 1 week ago

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 https://github.com/jaxleyverse/jaxley/blob/32068e9e7d463f8170e1e86f02c46c18cbaaefeb/jaxley/modules/network.py#L563. (see https://github.com/jaxleyverse/jaxley/blob/32068e9e7d463f8170e1e86f02c46c18cbaaefeb/jaxley/modules/base.py#L699 for how it is done for channels). Then one could amend https://github.com/jaxleyverse/jaxley/blob/32068e9e7d463f8170e1e86f02c46c18cbaaefeb/jaxley/modules/base.py#L1198 with the synapse currents that we can pull from synapse_current_names and make sure https://github.com/jaxleyverse/jaxley/blob/32068e9e7d463f8170e1e86f02c46c18cbaaefeb/jaxley/modules/network.py#L254 / 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 commented 1 week ago

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 commented 1 week ago

Awesome! :)