brainstate 0.1.9__py2.py3-none-any.whl → 0.2.0__py2.py3-none-any.whl

brainstate/nn/_synaptic_projection.py

@@ -1,423 +0,0 @@
-# Copyright 2025 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-# -*- coding: utf-8 -*-
-
-
-from typing import Callable, Union, Tuple
-
-import brainunit as u
-
-from brainstate import init
-from brainstate._state import ParamState
-from brainstate.typing import ArrayLike
-from ._dynamics import Dynamics, Projection
-
-__all__ = [
-    'SymmetryGapJunction',
-    'AsymmetryGapJunction',
-]
-
-
-class align_pre_ltp(Projection):
-    pass
-
-
-class align_post_ltp(Projection):
-    pass
-
-
-def get_gap_junction_post_key(i: int):
-    return f'gap_junction_post_{i}'
-
-
-def get_gap_junction_pre_key(i: int):
-    return f'gap_junction_pre_{i}'
-
-
-class SymmetryGapJunction(Projection):
-    """
-    Implements a symmetric electrical coupling (gap junction) between neuron populations.
-
-    This class represents electrical synapses where the conductance is identical in both
-    directions. Gap junctions allow bidirectional flow of electrical current directly between
-    neurons, with the current magnitude proportional to the voltage difference between
-    connected neurons.
-
-    Parameters
-    ----------
-    couples : Union[Tuple[Dynamics, Dynamics], Dynamics]
-        Either a single Dynamics object (when pre and post populations are the same)
-        or a tuple of two Dynamics objects (pre, post) representing the coupled neuron populations.
-    states : Union[str, Tuple[str, str]]
-        Either a single string (when pre and post states are the same)
-        or a tuple of two strings (pre_state, post_state) representing the state variables
-        to use for calculating voltage differences (typically membrane potentials).
-    conn : Callable
-        Connection function that returns pre_ids and post_ids arrays defining connections.
-    weight : Union[Callable, ArrayLike]
-        Conductance weights for the gap junctions. The same weight applies in both directions
-        of the connection.
-    param_type : type, optional
-        The parameter state type to use for weights, defaults to ParamState.
-
-    Notes
-    -----
-    The symmetric gap junction applies identical conductance in both directions between
-    connected neurons, ensuring balanced electrical coupling in the network.
-
-    See Also
-    --------
-    AsymmetryGapJunction : For gap junctions with different conductances in each direction.
-    """
-
-    def __init__(
-        self,
-        couples: Union[Tuple[Dynamics, Dynamics], Dynamics],
-        states: Union[str, Tuple[str, str]],
-        conn: Callable,
-        weight: Union[Callable, ArrayLike],
-        param_type: type = ParamState
-    ):
-        super().__init__()
-
-        if isinstance(states, str):
-            pre_state = states
-            post_state = states
-        else:
-            pre_state, post_state = states
-        if isinstance(couples, Dynamics):
-            pre = couples
-            post = couples
-        else:
-            pre, post = couples
-        assert isinstance(pre_state, str), "pre_state must be a string representing the pre-synaptic state"
-        assert isinstance(post_state, str), "post_state must be a string representing the post-synaptic state"
-        assert isinstance(pre, Dynamics), "pre must be a Dynamics object"
-        assert isinstance(post, Dynamics), "post must be a Dynamics object"
-        self.pre_state = pre_state
-        self.post_state = post_state
-        self.pre = pre
-        self.post = post
-        self.pre_ids, self.post_ids = conn(pre.out_size, post.out_size)
-        self.weight = param_type(init.param(weight, (len(self.pre_ids),)))
-
-    def update(self, *args, **kwargs):
-        if not hasattr(self.pre, self.pre_state):
-            raise ValueError(f"pre_state {self.pre_state} not found in pre-synaptic neuron group")
-        if not hasattr(self.post, self.post_state):
-            raise ValueError(f"post_state {self.post_state} not found in post-synaptic neuron group")
-        pre = getattr(self.pre, self.pre_state).value
-        post = getattr(self.post, self.post_state).value
-
-        return symmetry_gap_junction_projection(
-            pre=self.pre,
-            pre_value=pre,
-            post=self.post,
-            post_value=post,
-            pre_ids=self.pre_ids,
-            post_ids=self.post_ids,
-            weight=self.weight.value,
-        )
-
-
-def symmetry_gap_junction_projection(
-    pre: Dynamics,
-    pre_value: ArrayLike,
-    post: Dynamics,
-    post_value: ArrayLike,
-    pre_ids: ArrayLike,
-    post_ids: ArrayLike,
-    weight: ArrayLike,
-):
-    """
-    Calculate symmetrical electrical coupling through gap junctions between neurons.
-
-    This function implements bidirectional gap junction coupling where the same weight is
-    applied in both directions. It computes the electrical current flowing between
-    connected neurons based on their potential differences and updates both pre-synaptic
-    and post-synaptic neuron groups with appropriate input currents.
-
-    Parameters
-    ----------
-    pre : Dynamics
-        The pre-synaptic neuron group dynamics object.
-    pre_value : ArrayLike
-        State values (typically membrane potentials) of the pre-synaptic neurons.
-    post : Dynamics
-        The post-synaptic neuron group dynamics object.
-    post_value : ArrayLike
-        State values (typically membrane potentials) of the post-synaptic neurons.
-    pre_ids : ArrayLike
-        Indices of pre-synaptic neurons that form gap junctions.
-    post_ids : ArrayLike
-        Indices of post-synaptic neurons that form gap junctions,
-        where each pre_ids[i] is connected to post_ids[i].
-    weight : ArrayLike
-        Conductance weights for the gap junctions. Can be a scalar (same weight for all connections)
-        or an array with length matching pre_ids.
-
-    Returns
-    -------
-    ArrayLike
-        The input currents that were added to the pre-synaptic neuron group.
-
-    Notes
-    -----
-    The electrical coupling is implemented as I = g(V_pre - V_post), where:
-    - I is the current flowing from pre to post neuron
-    - g is the gap junction conductance (weight)
-    - V_pre and V_post are the membrane potentials of connected neurons
-
-    Equal but opposite currents are applied to both connected neurons, ensuring
-    conservation of current in the network.
-
-    Raises
-    ------
-    AssertionError
-        If weight dimensionality is incorrect or pre_ids and post_ids have different lengths.
-    """
-    assert u.math.ndim(weight) == 0 or weight.shape[0] == len(pre_ids), \
-        "weight must be a scalar or have the same length as pre_ids"
-    assert len(pre_ids) == len(post_ids), "pre_ids and post_ids must have the same length"
-    # Calculate the voltage difference between connected pre-synaptic and post-synaptic neurons
-    # and multiply by the connection weights
-    diff = (pre_value[pre_ids] - post_value[post_ids]) * weight
-
-    # add to post-synaptic neuron group
-    # Initialize the input currents for the post-synaptic neuron group
-    inputs = u.math.zeros(post.out_size, unit=u.get_unit(diff))
-    # Add the calculated current to the corresponding post-synaptic neurons
-    inputs = inputs.at[post_ids].add(diff)
-    # Generate a unique key for the post-synaptic input currents
-    key = get_gap_junction_post_key(0 if post.current_inputs is None else len(post.current_inputs))
-    # Add the input currents to the post-synaptic neuron group
-    post.add_current_input(key, inputs)
-
-    # add to pre-synaptic neuron group
-    # Initialize the input currents for the pre-synaptic neuron group
-    inputs = u.math.zeros(pre.out_size, unit=u.get_unit(diff))
-    # Add the calculated current to the corresponding pre-synaptic neurons
-    inputs = inputs.at[pre_ids].add(diff)
-    # Generate a unique key for the pre-synaptic input currents
-    key = get_gap_junction_pre_key(0 if pre.current_inputs is None else len(pre.current_inputs))
-    # Add the input currents to the pre-synaptic neuron group with opposite polarity
-    pre.add_current_input(key, -inputs)
-    return inputs
-
-
-class AsymmetryGapJunction(Projection):
-    """
-    Implements an asymmetric electrical coupling (gap junction) between neuron populations.
-
-    This class represents electrical synapses where the conductance in one direction can differ
-    from the conductance in the opposite direction. Unlike chemical synapses, gap junctions
-    allow bidirectional flow of electrical current directly between neurons, with the current
-    magnitude proportional to the voltage difference between connected neurons.
-
-    Parameters
-    ----------
-    pre : Dynamics
-        The pre-synaptic neuron group dynamics object.
-    pre_state : str
-        Name of the state variable in pre-synaptic neurons (typically membrane potential).
-    post : Dynamics
-        The post-synaptic neuron group dynamics object.
-    post_state : str
-        Name of the state variable in post-synaptic neurons (typically membrane potential).
-    conn : Callable
-        Connection function that returns pre_ids and post_ids arrays defining connections.
-    weight : Union[Callable, ArrayLike]
-        Conductance weights for the gap junctions. Must have shape [..., 2] where the last
-        dimension contains [pre_weight, post_weight] for each connection, defining
-        potentially different conductances in each direction.
-    param_type : type, optional
-        The parameter state type to use for weights, defaults to ParamState.
-
-    Examples
-    --------
-    >>> import brainstate
-    >>> import brainunit as u
-    >>> import numpy as np
-    >>>
-    >>> # Create two neuron populations
-    >>> n_neurons = 100
-    >>> pre_pop = brainstate.nn.LIF(n_neurons, V_rest=-70*u.mV, V_threshold=-50*u.mV)
-    >>> post_pop = brainstate.nn.LIF(n_neurons, V_rest=-70*u.mV, V_threshold=-50*u.mV)
-    >>> pre_pop.init_state()
-    >>> post_pop.init_state()
-    >>>
-    >>> # Create asymmetric gap junction with different weights in each direction
-    >>> weights = np.ones((n_neurons, 2)) * u.nS
-    >>> weights[:, 0] *= 2.0  # Double weight in pre->post direction
-    >>>
-    >>> gap_junction = brainstate.nn.AsymmetryGapJunction(
-    ...     pre=pre_pop,
-    ...     pre_state='V',
-    ...     post=post_pop,
-    ...     post_state='V',
-    ...     conn=one_to_one,
-    ...     weight=weights
-    ... )
-
-    Notes
-    -----
-    The asymmetric gap junction allows for different conductances in each direction between
-    the same pair of neurons. This can model rectifying electrical synapses that preferentially
-    allow current to flow in one direction.
-
-    See Also
-    --------
-    SymmetryGapJunction : For gap junctions with identical conductance in both directions.
-    """
-
-    def __init__(
-        self,
-        pre: Dynamics,
-        pre_state: str,
-        post: Dynamics,
-        post_state: str,
-        conn: Callable,
-        weight: Union[Callable, ArrayLike],
-        param_type: type = ParamState
-    ):
-        super().__init__()
-
-        assert isinstance(pre_state, str), "pre_state must be a string representing the pre-synaptic state"
-        assert isinstance(post_state, str), "post_state must be a string representing the post-synaptic state"
-        self.pre_state = pre_state
-        self.post_state = post_state
-        self.pre = pre
-        self.post = post
-        self.pre_ids, self.post_ids = conn(pre.out_size, post.out_size)
-        self.weight = param_type(init.param(weight, (len(self.pre_ids), 2)))
-
-    def update(self, *args, **kwargs):
-        if not hasattr(self.pre, self.pre_state):
-            raise ValueError(f"pre_state {self.pre_state} not found in pre-synaptic neuron group")
-        if not hasattr(self.post, self.post_state):
-            raise ValueError(f"post_state {self.post_state} not found in post-synaptic neuron group")
-        pre = getattr(self.pre, self.pre_state).value
-        post = getattr(self.post, self.post_state).value
-
-        return asymmetry_gap_junction_projection(
-            pre=self.pre,
-            pre_value=pre,
-            post=self.post,
-            post_value=post,
-            pre_ids=self.pre_ids,
-            post_ids=self.post_ids,
-            weight=self.weight.value,
-        )
-
-
-def asymmetry_gap_junction_projection(
-    pre: Dynamics,
-    pre_value: ArrayLike,
-    post: Dynamics,
-    post_value: ArrayLike,
-    pre_ids: ArrayLike,
-    post_ids: ArrayLike,
-    weight: ArrayLike,
-):
-    """
-    Calculate asymmetrical electrical coupling through gap junctions between neurons.
-
-    This function implements bidirectional gap junction coupling where different weights
-    can be applied in each direction. It computes the electrical current flowing between
-    connected neurons based on their potential differences and updates both pre-synaptic
-    and post-synaptic neuron groups with appropriate input currents.
-
-    Parameters
-    ----------
-    pre : Dynamics
-        The pre-synaptic neuron group dynamics object.
-    pre_value : ArrayLike
-        State values (typically membrane potentials) of the pre-synaptic neurons.
-    post : Dynamics
-        The post-synaptic neuron group dynamics object.
-    post_value : ArrayLike
-        State values (typically membrane potentials) of the post-synaptic neurons.
-    pre_ids : ArrayLike
-        Indices of pre-synaptic neurons that form gap junctions.
-    post_ids : ArrayLike
-        Indices of post-synaptic neurons that form gap junctions,
-        where each pre_ids[i] is connected to post_ids[i].
-    weight : ArrayLike
-        Conductance weights for the gap junctions. Must have shape [..., 2], where
-        the last dimension contains [pre_weight, post_weight] for each connection.
-        Can be a 1D array [pre_weight, post_weight] (same weights for all connections)
-        or a 2D array with shape [len(pre_ids), 2] for connection-specific weights.
-
-    Returns
-    -------
-    ArrayLike
-        The input currents that were added to the pre-synaptic neuron group.
-
-    Notes
-    -----
-    The electrical coupling is implemented with direction-specific conductances:
-    - I_pre2post = g_pre * (V_pre - V_post) flowing from pre to post neuron
-    - I_post2pre = g_post * (V_pre - V_post) flowing from post to pre neuron
-    where g_pre and g_post can be different, allowing for asymmetrical coupling.
-
-    Raises
-    ------
-    AssertionError
-        If weight dimensionality is incorrect or pre_ids and post_ids have different lengths.
-    ValueError
-        If weight shape is incompatible with asymmetrical gap junction requirements.
-    """
-    assert weight.shape[-1] == 2, 'weight must be a 2-element array for asymmetry gap junctions'
-    assert len(pre_ids) == len(post_ids), "pre_ids and post_ids must have the same length"
-    if u.math.ndim(weight) == 1:
-        # If weight is a 1D array, it should have two elements for pre and post weights
-        assert weight.shape[0] == 2, "weight must be a 2-element array for asymmetry gap junctions"
-        pre_weight = weight[0]
-        post_weight = weight[1]
-    elif u.math.ndim(weight) == 2:
-        # If weight is a 2D array, it should have two rows for pre and post weights
-        pre_weight = weight[:, 0]
-        post_weight = weight[:, 1]
-        assert pre_weight.shape[0] == len(pre_ids), "pre_weight must have the same length as pre_ids"
-        assert post_weight.shape[0] == len(post_ids), "post_weight must have the same length as post_ids"
-    else:
-        raise ValueError("weight must be a 1D or 2D array for asymmetry gap junctions")
-
-    # Calculate the voltage difference between connected pre-synaptic and post-synaptic neurons
-    # and multiply by the connection weights
-    diff = pre_value[pre_ids] - post_value[post_ids]
-    pre2post_current = diff * pre_weight
-    post2pre_current = diff * post_weight
-
-    # add to post-synaptic neuron group
-    # Initialize the input currents for the post-synaptic neuron group
-    inputs = u.math.zeros(post.out_size, unit=u.get_unit(pre2post_current))
-    # Add the calculated current to the corresponding post-synaptic neurons
-    inputs = inputs.at[post_ids].add(pre2post_current)
-    # Generate a unique key for the post-synaptic input currents
-    key = get_gap_junction_post_key(0 if post.current_inputs is None else len(post.current_inputs))
-    # Add the input currents to the post-synaptic neuron group
-    post.add_current_input(key, inputs)
-
-    # add to pre-synaptic neuron group
-    # Initialize the input currents for the pre-synaptic neuron group
-    inputs = u.math.zeros(pre.out_size, unit=u.get_unit(post2pre_current))
-    # Add the calculated current to the corresponding pre-synaptic neurons
-    inputs = inputs.at[pre_ids].add(post2pre_current)
-    # Generate a unique key for the pre-synaptic input currents
-    key = get_gap_junction_pre_key(0 if pre.current_inputs is None else len(pre.current_inputs))
-    # Add the input currents to the pre-synaptic neuron group with opposite polarity
-    pre.add_current_input(key, -inputs)
-    return inputs

brainstate/nn/_synouts.py

@@ -1,162 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-# -*- coding: utf-8 -*-
-
-import brainunit as u
-import jax.numpy as jnp
-
-from brainstate.mixin import BindCondData
-from brainstate.typing import ArrayLike
-from ._module import Module
-
-__all__ = [
-    'SynOut', 'COBA', 'CUBA', 'MgBlock',
-]
-
-
-class SynOut(Module, BindCondData):
-    """
-    Base class for synaptic outputs.
-
-    :py:class:`~.SynOut` is also subclass of :py:class:`~.ParamDesc` and :py:class:`~.BindCondData`.
-    """
-
-    __module__ = 'brainstate.nn'
-
-    def __init__(self, ):
-        super().__init__()
-        self._conductance = None
-
-    def __call__(self, *args, **kwargs):
-        if self._conductance is None:
-            raise ValueError(f'Please first pack conductance data at the current step using '
-                             f'".{BindCondData.bind_cond.__name__}(data)". {self}')
-        ret = self.update(self._conductance, *args, **kwargs)
-        return ret
-
-    def update(self, conductance, potential):
-        raise NotImplementedError
-
-
-class COBA(SynOut):
-    r"""
-    Conductance-based synaptic output.
-
-    Given the synaptic conductance, the model output the post-synaptic current with
-
-    .. math::
-
-       I_{syn}(t) = g_{\mathrm{syn}}(t) (E - V(t))
-
-    Parameters
-    ----------
-    E: ArrayLike
-      The reversal potential.
-
-    See Also
-    --------
-    CUBA
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(self, E: ArrayLike):
-        super().__init__()
-
-        self.E = E
-
-    def update(self, conductance, potential):
-        return conductance * (self.E - potential)
-
-
-class CUBA(SynOut):
-    r"""Current-based synaptic output.
-
-    Given the conductance, this model outputs the post-synaptic current with a identity function:
-
-    .. math::
-
-       I_{\mathrm{syn}}(t) = g_{\mathrm{syn}}(t)
-
-    Parameters
-    ----------
-    scale: ArrayLike
-      The scaling factor for the conductance. Default 1. [mV]
-
-    See Also
-    --------
-    COBA
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(self, scale: ArrayLike = u.volt):
-        super().__init__()
-        self.scale = scale
-
-    def update(self, conductance, potential=None):
-        return conductance * self.scale
-
-
-class MgBlock(SynOut):
-    r"""Synaptic output based on Magnesium blocking.
-
-    Given the synaptic conductance, the model output the post-synaptic current with
-
-    .. math::
-
-       I_{syn}(t) = g_{\mathrm{syn}}(t) (E - V(t)) g_{\infty}(V,[{Mg}^{2+}]_{o})
-
-    where The fraction of channels :math:`g_{\infty}` that are not blocked by magnesium can be fitted to
-
-    .. math::
-
-       g_{\infty}(V,[{Mg}^{2+}]_{o}) = (1+{e}^{-\alpha V} \frac{[{Mg}^{2+}]_{o}} {\beta})^{-1}
-
-    Here :math:`[{Mg}^{2+}]_{o}` is the extracellular magnesium concentration.
-
-    Parameters
-    ----------
-    E: ArrayLike
-      The reversal potential for the synaptic current. [mV]
-    alpha: ArrayLike
-      Binding constant. Default 0.062
-    beta: ArrayLike
-      Unbinding constant. Default 3.57
-    cc_Mg: ArrayLike
-      Concentration of Magnesium ion. Default 1.2 [mM].
-    V_offset: ArrayLike
-      The offset potential. Default 0. [mV]
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        E: ArrayLike = 0.,
-        cc_Mg: ArrayLike = 1.2,
-        alpha: ArrayLike = 0.062,
-        beta: ArrayLike = 3.57,
-        V_offset: ArrayLike = 0.,
-    ):
-        super().__init__()
-
-        self.E = E
-        self.V_offset = V_offset
-        self.cc_Mg = cc_Mg
-        self.alpha = alpha
-        self.beta = beta
-
-    def update(self, conductance, potential):
-        norm = (1 + self.cc_Mg / self.beta * jnp.exp(self.alpha * (self.V_offset - potential)))
-        return conductance * (self.E - potential) / norm

brainstate/nn/_synouts_test.py

@@ -1,57 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-
-import unittest
-
-import brainunit as u
-import jax.numpy as jnp
-import numpy as np
-
-import brainstate
-
-
-class TestSynOutModels(unittest.TestCase):
-    def setUp(self):
-        self.conductance = jnp.array([0.5, 1.0, 1.5])
-        self.potential = jnp.array([-70.0, -65.0, -60.0])
-        self.E = jnp.array([-70.0])
-        self.alpha = jnp.array([0.062])
-        self.beta = jnp.array([3.57])
-        self.cc_Mg = jnp.array([1.2])
-        self.V_offset = jnp.array([0.0])
-
-    def test_COBA(self):
-        model = brainstate.nn.COBA(E=self.E)
-        output = model.update(self.conductance, self.potential)
-        expected_output = self.conductance * (self.E - self.potential)
-        np.testing.assert_array_almost_equal(output, expected_output)
-
-    def test_CUBA(self):
-        model = brainstate.nn.CUBA()
-        output = model.update(self.conductance)
-        expected_output = self.conductance * model.scale
-        self.assertTrue(u.math.allclose(output, expected_output))
-
-    def test_MgBlock(self):
-        model = brainstate.nn.MgBlock(E=self.E, cc_Mg=self.cc_Mg, alpha=self.alpha, beta=self.beta, V_offset=self.V_offset)
-        output = model.update(self.conductance, self.potential)
-        norm = (1 + self.cc_Mg / self.beta * jnp.exp(self.alpha * (self.V_offset - self.potential)))
-        expected_output = self.conductance * (self.E - self.potential) / norm
-        np.testing.assert_array_almost_equal(output, expected_output)
-
-
-if __name__ == '__main__':
-    unittest.main()