# This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
# licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
from typing import Callable, Dict, Optional, Tuple
import jax.numpy as jnp
from jax import Array
from jax.nn import sigmoid
from jaxley.solver_gate import exponential_euler, save_exp
from jaxley.synapses.synapse import Synapse
[docs]
class DynamicSynapse(Synapse):
r"""A state-based synapse with fixed time constant.
Unlike the ``ConductanceSynapse``, this synapse contains a synaptic state. However,
unlike in the ``IonotropicSynapse``, the synaptic state approaches its steady-state
with a constant (i.e., not voltage dependent) time constant.
This synapse implements the following equations:
.. math::
I = \overline{g}\, \cdot s\, \cdot (E - V_{\text{post}})
.. math::
\tau \frac{\text{d}s}{\text{d}t} = \sigma\!\left(\frac{V_{\text{pre}} - V_{\text{thr}}}{\Delta}\right) - s,
where :math:`\mathrm{\sigma}(\cdot)` is a nonlinearity such as a ReLU, Sigmoid,
or TanH. By default, it is a sigmoid, but it can be modified by the user.
More informally:
This synapse has a state which defines its conductance. The state approaches
a nonlinear map of the voltage with a time constant which is not voltage-dependent.
The current is conductance-based, i.e., it depends on a reversal potential.
The synaptic parameters are:
- ``gS``: the maximal conductance :math:`\overline{g}` (uS).
- ``tau``: the time constant :math:`\tau` (:math:`ms`).
- ``v_th``: the threshold at which the synapse becomes active
:math:`V_{\text{thr}}` (mV).
- ``delta``: The inverse of the slope of the activation :math:`\Delta` (mV).
The inserted cellular parameters are:
- ``e_syn``: The synaptic reversal potential :math:`E` (mV). This synapse uses
the pre-synaptic reveral potential to compute the current, thereby directly
enforcing Dale's law.
The synaptic state is:
- ``s``: the activity level of the synapse.
.. rubric:: Example usage
Insert a synapse with a sigmoid nonlinearity (the default) and change parameters
and initial state.
::
import jaxley as jx
from jaxley.connect import connect
from jaxley.synapses import DynamicSynapse
cell = jx.Cell()
net = jx.Network([cell for _ in range(2)])
# Connect neurons with the `DynamicSynapse`.
connect(net.cell(0), net.cell(1), DynamicSynapse())
# Set parameters.
net.set("DynamicSynapse_gS", 0.0001) # Maximal conductance.
net.set("DynamicSynapse_e_syn", 10.0) # Reversal potential.
net.set("DynamicSynapse_tau", 4.0) # Time constant.
net.set("DynamicSynapse_v_th", -40.0) # Threshold.
net.set("DynamicSynapse_delta", 10.0) # 1 / slope of activation.
# Set the initial state.
net.set("DynamicSynapse_s", 0.1)
Insert a synapse with a ReLU nonlinearity.
::
import jaxley as jx
from jaxley.connect import connect
from jaxley.synapses import DynamicSynapse
from jax.nn import relu
cell = jx.Cell()
net = jx.Network([cell for _ in range(2)])
# Connect neurons with the `DynamicSynapse`.
connect(net.cell(0), net.cell(1), DynamicSynapse(relu))
Insert a synapse with a custom nonlinearity.
::
import jaxley as jx
from jaxley.connect import connect
from jaxley.synapses import DynamicSynapse
cell = jx.Cell()
net = jx.Network([cell for _ in range(2)])
def nonlinearity(x):
return x ** 2
# Connect neurons with the `DynamicSynapse`.
connect(net.cell(0), net.cell(1), DynamicSynapse(nonlinearity))
"""
def __init__(self, nonlinearity: Callable = sigmoid, name: Optional[str] = None):
super().__init__(name)
prefix = self._name
self.synapse_params = {
f"{prefix}_gS": 1e-4, # uS
f"{prefix}_tau": 5.0, # ms
f"{prefix}_v_th": -35.0, # mV
f"{prefix}_delta": 10.0, # mV
}
self.synapse_states = {f"{prefix}_s": 0.0}
self.node_params = {f"{prefix}_e_syn": 0.0}
self.node_states = {}
self.nonlinearity = nonlinearity
[docs]
def update_states(
self,
synapse_states: dict[str, Array],
synapse_params: dict[str, Array],
pre_voltage: Array,
post_voltage: Array,
pre_states: dict[str, Array],
post_states: dict[str, Array],
pre_params: dict[str, Array],
post_params: dict[str, Array],
delta_t: float,
) -> Dict:
"""Return updated synapse state and current."""
prefix = self._name
v_th = synapse_params[f"{prefix}_v_th"]
delta = synapse_params[f"{prefix}_delta"]
s = synapse_states[f"{prefix}_s"]
s_inf = self.nonlinearity((pre_voltage - v_th) / delta)
s_tau = synapse_params[f"{prefix}_tau"]
new_s = exponential_euler(s, delta_t, s_inf, s_tau)
return {f"{prefix}_s": new_s}
[docs]
def compute_current(
self,
synapse_states: dict[str, Array],
synapse_params: dict[str, Array],
pre_voltage: Array,
post_voltage: Array,
pre_states: dict[str, Array],
post_states: dict[str, Array],
pre_params: dict[str, Array],
post_params: dict[str, Array],
delta_t: float,
) -> float:
prefix = self._name
g_syn = synapse_params[f"{prefix}_gS"] * synapse_states[f"{prefix}_s"]
return g_syn * (post_voltage - pre_params[f"{prefix}_e_syn"])
