brainmass 0.0.1__py2.py3-none-any.whl

brainmass/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+# ==============================================================================
+
+__version__ = "0.0.1"
+
+__all__ = [
+    'DiffusiveCoupling',
+    'AdditiveCoupling',
+    'WilsonCowanModel',
+    'OUProcess',
+    'BOLDSignal',
+]
+
+from .bold import *
+from .coupling import *
+from .noise import *
+from .wilson_cowan import *

brainmass/bold.py

@@ -0,0 +1,143 @@
+# 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.
+# ==============================================================================
+
+from typing import Union, Callable
+
+import brainstate
+import jax.numpy as jnp
+
+from .integration import ode_rk2_step
+
+__all__ = [
+    'BOLDSignal',
+]
+
+
+class BOLDSignal(brainstate.nn.Dynamics):
+    r"""
+    Balloon-Windkessel hemodynamic model of Friston et al. (2003) [1]_.
+
+    The Balloon-Windkessel model describes the coupling of perfusion to BOLD signal, with a
+    dynamical model of the transduction of neuronal activity into perfusion changes. The
+    model assumes that the BOLD signal is a static nonlinear function of the normalized total
+    deoxyhemoglobin voxel content, normalized venous volume, resting net oxygen extraction
+    fraction by the capillary bed, and resting blood volume fraction. The BOLD-signal estimation
+    for each brain area is computed by the level of synaptic activity in that particular cortical
+    area, noted $z_i$ for a given cortical are $i$.
+
+    For the i-th region, synaptic activity $z_i$ causes an increase in a vasodilatory signal $x_i$
+    that is subject to autoregulatory feedback. Inflow $f_i$ responds in proportion to this signal
+    with concomitant changes in blood volume $v_i$ and deoxyhemoglobin content $q_i$. The equations
+    relating these biophysical processes are as follows:
+
+    $$
+    \begin{gathered}
+    \dot{x}_i=z_i-k_i x_i-\gamma_i\left(f_i-1\right) \\
+    \dot{f}_i=x_i \\
+    \tau_i \dot{v}_i=f_i-v_i^{1 / \alpha} \\
+    \tau_i \dot{q}_i=\frac{f_i}{\rho}\left[1-(1-\rho)^{1 / f_i}\right]-q_i v_i^{1 / \alpha-1},
+    \end{gathered}
+    $$
+
+    where $\rho$ is the resting oxygen extraction fraction. The BOLD signal is given by the following:
+
+    $$
+    \mathrm{BOLD}_i=V_0\left[k_1\left(1-q_i\right)+k_2\left(1-q_i / v_i\right)+k_3\left(1-v_i\right)\right],
+    $$
+
+    where $V_0 = 0.02, k1 = 7\rho, k2 = 2$, and $k3 = 2\rho − 0.2$. All biophysical parameters were taken
+    as in Friston et al. (2003) [1]_. The BOLD model converts the local synaptic activity of a given cortical
+    area into an observable BOLD signal and does not actively couple the signals from other cortical areas.
+
+    Parameters
+    ----------
+    in_size : int
+        Size of the input vector (number of brain regions).
+    gamma : float or callable, optional
+        Rate of signal decay (default is 0.41).
+    k : float or callable, optional
+        Rate of flow-dependent elimination (default is 0.65).
+    alpha : float or callable, optional
+        Grubb's exponent (default is 0.32).
+    tau : float or callable, optional
+        Hemodynamic transit time (default is 0.98).
+    rho : float or callable, optional
+        Resting oxygen extraction fraction (default is 0.34).
+    V0 : float, optional
+        Resting blood volume fraction (default is 0.02).
+
+    References
+    ----------
+    .. [1] Friston KJ, Harrison L, Penny W (2003) Dynamic causal modelling. Neuroimage 19:1273–1302,
+           doi:10.1016/S1053-8119(03)00202-7
+    """
+
+    def __init__(
+        self,
+        in_size,
+        gamma: Union[brainstate.typing.ArrayLike, Callable] = 0.41,
+        k: Union[brainstate.typing.ArrayLike, Callable] = 0.65,
+        alpha: Union[brainstate.typing.ArrayLike, Callable] = 0.32,
+        tau: Union[brainstate.typing.ArrayLike, Callable] = 0.98,
+        rho: Union[brainstate.typing.ArrayLike, Callable] = 0.34,
+        V0: float = 0.02,
+    ):
+        super().__init__(in_size)
+
+        self.gamma = brainstate.init.param(gamma, self.varshape)
+        self.k = brainstate.init.param(k, self.varshape)
+        self.alpha = brainstate.init.param(alpha, self.varshape)
+        self.tau = brainstate.init.param(tau, self.varshape)
+        self.rho = brainstate.init.param(rho, self.varshape)
+
+        self.V0 = V0
+        self.k1 = 7 * self.rho
+        self.k2 = 2.
+        self.k3 = 2 * self.rho - 0.2
+
+        self.init = brainstate.init.Constant(1.)
+
+    def init_state(self, batch_size=None, **kwargs):
+        self.x = brainstate.HiddenState(brainstate.init.param(self.init, self.varshape, batch_size))
+        self.f = brainstate.HiddenState(brainstate.init.param(self.init, self.varshape, batch_size))
+        self.v = brainstate.HiddenState(brainstate.init.param(self.init, self.varshape, batch_size))
+        self.q = brainstate.HiddenState(brainstate.init.param(self.init, self.varshape, batch_size))
+
+    def reset_state(self, batch_size=None, **kwargs):
+        self.x.value = brainstate.init.param(self.init, self.varshape, batch_size)
+        self.f.value = brainstate.init.param(self.init, self.varshape, batch_size)
+        self.v.value = brainstate.init.param(self.init, self.varshape, batch_size)
+        self.q.value = brainstate.init.param(self.init, self.varshape, batch_size)
+
+    def derivative(self, y, t, z):
+        x, f, v, q = y
+        dx = z - self.k * x - self.gamma * (f - 1)
+        df = x
+        dv = (f - jnp.power(v, 1 / self.alpha)) / self.tau
+        E = 1 - jnp.power(1 - self.rho, 1 / f)
+        dq = (f * E / self.rho - jnp.power(v, 1 / self.alpha) * q / v) / self.tau
+        return dx, df, dv, dq
+
+    def update(self, z):
+        x, f, v, q = ode_rk2_step(self.derivative, (self.x.value, self.f.value, self.v.value, self.q.value), 0., z)
+        self.x.value = x
+        self.f.value = f
+        self.v.value = v
+        self.q.value = q
+
+    def bold(self):
+        return self.V0 * (self.k1 * (1 - self.q.value) +
+                          self.k2 * (1 - self.q.value / self.rho) +
+                          self.k3 * (1 - self.v.value))

brainmass/coupling.py

@@ -0,0 +1,171 @@
+# 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.
+# ==============================================================================
+
+
+from typing import Union
+
+import brainstate
+import brainunit as u
+from brainstate.nn._dynamics import maybe_init_prefetch
+
+Prefetch = Union[brainstate.nn.PrefetchDelayAt, brainstate.nn.PrefetchDelay, brainstate.nn.Prefetch]
+
+__all__ = [
+    'DiffusiveCoupling',
+    'AdditiveCoupling',
+]
+
+
+class DiffusiveCoupling(brainstate.nn.Module):
+    r"""
+    Diffusive coupling.
+
+    This class implements a diffusive coupling mechanism for neural network modules.
+    It simulates the following model:
+
+    $$
+    \mathrm{current}_i = k * \sum_j (g_{ij} * x_{D_{ij}} - y_i)
+    $$
+
+    where:
+        - $\mathrm{current}_i$: the output current for neuron $i$
+        - $g_{ij}$: the connection strength between neuron $i$ and neuron $j$
+        - $x_{D_{ij}}$: the delayed state variable for neuron $j$, as seen by neuron $i$
+        - $y_i$: the state variable for neuron i
+
+    Parameters
+    ----------
+    x : Prefetch
+        The delayed state variable for the source units.
+    y : Prefetch
+        The delayed state variable for the target units.
+    conn : brainstate.typing.Array
+        The connection matrix (1D or 2D array) specifying the coupling strengths between units.
+    k: float
+        The global coupling strength. Default is 1.0.
+
+    Attributes
+    ----------
+    x : Prefetch
+        The delayed state variable for the source units.
+    y : Prefetch
+        The delayed state variable for the target units.
+    conn : Array
+        The connection matrix.
+    """
+
+    def __init__(
+        self,
+        x: Prefetch,
+        y: Prefetch,
+        conn: brainstate.typing.Array,
+        k: float = 1.0
+    ):
+        super().__init__()
+        assert isinstance(x, Prefetch), f'The first element must be a Prefetch. But got {type(x)}.'
+        assert isinstance(y, Prefetch), f'The second element must be a Prefetch. But got {type(y)}.'
+        self.x = x
+        self.y = y
+        self.k = k
+
+        # Connection matrix
+        self.conn = u.math.asarray(conn)
+        assert self.conn.ndim in (1, 2), f'Only support 1d, 2d connection matrix. But we got {self.conn.ndim}d.'
+
+    @brainstate.nn.call_order(2)
+    def init_state(self, *args, **kwargs):
+        maybe_init_prefetch(self.x)
+        maybe_init_prefetch(self.y)
+
+    def update(self):
+        delayed_x = self.x()
+        y = u.math.expand_dims(self.y(), axis=1)  # (..., 1)
+        if self.conn.ndim == 1:
+            assert self.conn.size == delayed_x.shape[-1], (
+                f'Connection matrix size {self.conn.size} does not '
+                f'match the variable size {delayed_x.shape[-1]}.'
+            )
+            diffusive = (self.conn * delayed_x).reshape(y.shape[0], -1) - y
+        elif self.conn.ndim == 2:
+            delayed_x = delayed_x.reshape(y.shape[0], -1)
+            assert self.conn.shape == delayed_x.shape, (f'Connection matrix shape {self.conn.shape} does not '
+                                                        f'match the variable shape {delayed_x.shape}.')
+            diffusive = (self.conn * delayed_x) - y
+        else:
+            raise NotImplementedError(f'Only support 1d, 2d connection matrix. But we got {self.conn.ndim}d.')
+        return self.k * diffusive.sum(axis=1)
+
+
+class AdditiveCoupling(brainstate.nn.Module):
+    r"""
+    Additive coupling.
+
+    This class implements an additive coupling mechanism for neural network modules.
+    It simulates the following model:
+
+    $$
+    \mathrm{current}_i = k * \sum_j g_{ij} * x_{D_{ij}}
+    $$
+
+    where:
+        - $\mathrm{current}_i$: the output current for neuron $i$
+        - $g_{ij}$: the connection strength between neuron $i$ and neuron $j$
+        - $x_{D_{ij}}$: the delayed state variable for neuron $j$, as seen by neuron $i$
+
+    Parameters
+    ----------
+    x : Prefetch
+        The delayed state variable for the source units.
+    conn : brainstate.typing.Array
+        The connection matrix (1D or 2D array) specifying the coupling strengths between units.
+    k: float
+        The global coupling strength. Default is 1.0.
+
+    Attributes
+    ----------
+    x : Prefetch
+        The delayed state variable for the source units.
+    conn : Array
+        The connection matrix.
+    """
+
+    def __init__(
+        self,
+        x: Prefetch,
+        conn: brainstate.typing.Array,
+        k: float = 1.0
+    ):
+        super().__init__()
+        assert isinstance(x, Prefetch), f'The first element must be a Prefetch. But got {type(x)}.'
+        self.x = x
+        self.k = k
+
+        # Connection matrix
+        self.conn = u.math.asarray(conn)
+        assert self.conn.ndim == 2, f'Only support 2d connection matrix. But we got {self.conn.ndim}d.'
+
+    @brainstate.nn.call_order(2)
+    def init_state(self, *args, **kwargs):
+        maybe_init_prefetch(self.x)
+
+    def update(self):
+        delayed_x = self.x()
+        assert self.conn.size == delayed_x.size, (
+            f'Connection matrix size {self.conn.size} does not '
+            f'match the variable size {delayed_x.size}.'
+        )
+        delayed_x = delayed_x.reshape(self.conn.shape)
+        diffusive = self.conn * delayed_x
+        return self.k * diffusive.sum(axis=1)

brainmass/integration.py

@@ -0,0 +1,253 @@
+# 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.
+# ==============================================================================
+
+from typing import Callable
+
+import brainstate
+import jax
+import jax.numpy as jnp
+from brainstate.typing import PyTree
+
+__all__ = [
+    'ode_euler_step',
+    'ode_rk2_step',
+    'ode_rk3_step',
+    'ode_rk4_step',
+    'sde_euler_step',
+    'sde_milstein_step',
+]
+
+ODE = Callable[[PyTree, float, ...], PyTree]
+
+
+def ode_euler_step(f: ODE, y: PyTree, t, *args):
+    """
+    Euler method for solving ordinary differential equations.
+    
+    The Euler method is the simplest numerical method for solving ODEs of the form:
+    dy/dt = f(y, t)
+    
+    The method approximates the solution using:
+    y_{n+1} = y_n + dt * f(y_n, t_n)
+    
+    Args:
+        f (ODE): The differential equation function dy/dt = f(y, t, *args)
+        y (PyTree): Current state vector
+        t (float): Current time
+        *args: Additional arguments passed to function f
+        
+    Returns:
+        PyTree: Updated state vector y_{n+1}
+        
+    Note:
+        This is a first-order method with O(dt) local truncation error.
+        It's the least accurate but most computationally efficient method.
+    """
+    dt = brainstate.environ.get_dt()
+    k1 = f(y, t, *args)
+    return jax.tree.map(lambda x, _k1: x + dt * _k1, y, k1)
+
+
+def ode_rk2_step(f: ODE, y: PyTree, t, *args):
+    """
+    Second-order Runge-Kutta method (RK2) for solving ODEs.
+    
+    Also known as the midpoint method or Heun's method, this method provides
+    better accuracy than Euler by using two function evaluations:
+    
+    k1 = f(y_n, t_n)
+    k2 = f(y_n + dt*k1, t_n + dt)
+    y_{n+1} = y_n + dt/2 * (k1 + k2)
+    
+    Args:
+        f (ODE): The differential equation function dy/dt = f(y, t, *args)
+        y (PyTree): Current state vector
+        t (float): Current time
+        *args: Additional arguments passed to function f
+        
+    Returns:
+        PyTree: Updated state vector y_{n+1}
+        
+    Note:
+        This is a second-order method with O(dt²) local truncation error.
+        More accurate than Euler with only one additional function evaluation.
+    """
+    dt = brainstate.environ.get_dt()
+    k1 = f(y, t, *args)
+    k2 = f(jax.tree.map(lambda x, k: x + dt * k, y, k1), t + dt, *args)
+    return jax.tree.map(lambda x, _k1, _k2: x + dt / 2 * (_k1 + _k2), y, k1, k2)
+
+
+def ode_rk3_step(f: ODE, y: PyTree, t, *args):
+    """
+    Third-order Runge-Kutta method (RK3) for solving ODEs.
+    
+    This method uses three function evaluations to achieve third-order accuracy:
+    
+    k1 = f(y_n, t_n)
+    k2 = f(y_n + dt/2*k1, t_n + dt/2)
+    k3 = f(y_n - dt*k1 + 2*dt*k2, t_n + dt)
+    y_{n+1} = y_n + dt/6 * (k1 + 4*k2 + k3)
+    
+    Args:
+        f (ODE): The differential equation function dy/dt = f(y, t, *args)
+        y (PyTree): Current state vector
+        t (float): Current time
+        *args: Additional arguments passed to function f
+        
+    Returns:
+        PyTree: Updated state vector y_{n+1}
+        
+    Note:
+        This is a third-order method with O(dt³) local truncation error.
+        More accurate than RK2 but requires one additional function evaluation.
+    """
+    dt = brainstate.environ.get_dt()
+    k1 = f(y, t, *args)
+    k2 = f(jax.tree.map(lambda x, k: x + dt / 2 * k, y, k1), t + dt / 2, *args)
+    k3 = f(jax.tree.map(lambda x, k1_val, k2_val: x - dt * k1_val + 2 * dt * k2_val, y, k1, k2), t + dt, *args)
+    return jax.tree.map(lambda x, _k1, _k2, _k3: x + dt / 6 * (_k1 + 4 * _k2 + _k3), y, k1, k2, k3)
+
+
+def ode_rk4_step(f: ODE, y: PyTree, t, *args):
+    """
+    Fourth-order Runge-Kutta method (RK4) for solving ODEs.
+    
+    The classic RK4 method uses four function evaluations to achieve fourth-order accuracy:
+    
+    k1 = f(y_n, t_n)
+    k2 = f(y_n + dt/2*k1, t_n + dt/2)
+    k3 = f(y_n + dt/2*k2, t_n + dt/2)
+    k4 = f(y_n + dt*k3, t_n + dt)
+    y_{n+1} = y_n + dt/6 * (k1 + 2*k2 + 2*k3 + k4)
+    
+    Args:
+        f (ODE): The differential equation function dy/dt = f(y, t, *args)
+        y (PyTree): Current state vector
+        t (float): Current time
+        *args: Additional arguments passed to function f
+        
+    Returns:
+        PyTree: Updated state vector y_{n+1}
+        
+    Note:
+        This is a fourth-order method with O(dt⁴) local truncation error.
+        The most commonly used method due to excellent accuracy/cost trade-off.
+    """
+    dt = brainstate.environ.get_dt()
+    k1 = f(y, t, *args)
+    k2 = f(jax.tree.map(lambda x, k: x + dt / 2 * k, y, k1), t + dt / 2, *args)
+    k3 = f(jax.tree.map(lambda x, k: x + dt / 2 * k, y, k2), t + dt / 2, *args)
+    k4 = f(jax.tree.map(lambda x, k: x + dt * k, y, k3), t + dt, *args)
+    return jax.tree.map(
+        lambda x, _k1, _k2, _k3, _k4: x + dt / 6 * (_k1 + 2 * _k2 + 2 * _k3 + _k4),
+        y, k1, k2, k3, k4
+    )
+
+
+def sde_euler_step(df, dg, y, t, sde_type='ito', **kwargs):
+    """
+    Euler-Maruyama method for solving stochastic differential equations (SDEs).
+    
+    Solves SDEs of the form:
+    dy = f(y, t)dt + g(y, t)dW
+    
+    where f is the drift term, g is the diffusion term, and dW is Wiener noise.
+    
+    The Euler-Maruyama scheme approximates:
+    y_{n+1} = y_n + f(y_n, t_n)*dt + g(y_n, t_n)*ΔW_n
+    
+    where ΔW_n ~ N(0, dt) is the Wiener increment.
+    
+    Args:
+        df (Callable): Drift function f(y, t, **kwargs) -> PyTree
+        dg (Callable): Diffusion function g(y, t, **kwargs) -> PyTree  
+        y (PyTree): Current state vector
+        t (float): Current time
+        sde_type (str): Type of SDE interpretation ('ito' only supported)
+        **kwargs: Additional arguments passed to df and dg
+        
+    Returns:
+        PyTree: Updated state vector y_{n+1}
+        
+    Note:
+        This method has strong convergence order 0.5 and weak convergence order 1.0.
+        Only Itô interpretation is currently supported.
+    """
+    assert sde_type in ['ito', ]
+
+    dt = brainstate.environ.get_dt()
+    dt_sqrt = jnp.sqrt(dt)
+    y_bars = jax.tree.map(
+        lambda y0, drift, diffusion: y0 + drift * dt + diffusion * brainstate.random.randn_like(y0) * dt_sqrt,
+        y, df(y, t, **kwargs), dg(y, t, **kwargs)
+    )
+    return y_bars
+
+
+def sde_milstein_step(df, dg, y, t, sde_type='ito', **kwargs):
+    """
+    Milstein method for solving stochastic differential equations (SDEs).
+    
+    Solves SDEs of the form:
+    dy = f(y, t)dt + g(y, t)dW
+    
+    The Milstein scheme includes an additional correction term for higher accuracy:
+    y_{n+1} = y_n + f(y_n, t_n)*dt + g(y_n, t_n)*ΔW_n + 
+              (1/2)*g(y_n, t_n)*∂g/∂y(y_n, t_n)*((ΔW_n)² - dt)
+    
+    This method approximates the derivative ∂g/∂y using finite differences:
+    ∂g/∂y ≈ (g(y + g*√dt) - g(y)) / √dt
+    
+    Args:
+        df (Callable): Drift function f(y, t, **kwargs) -> PyTree
+        dg (Callable): Diffusion function g(y, t, **kwargs) -> PyTree
+        y (PyTree): Current state vector
+        t (float): Current time
+        sde_type (str): SDE interpretation ('ito' or 'stra' for Stratonovich)
+        **kwargs: Additional arguments passed to df and dg
+        
+    Returns:
+        PyTree: Updated state vector y_{n+1}
+        
+    Note:
+        This method has strong convergence order 1.0, better than Euler-Maruyama.
+        Supports both Itô and Stratonovich interpretations.
+        The finite difference approximation is used for the derivative term.
+    """
+    assert sde_type in ['ito', 'stra']
+
+    dt = brainstate.environ.get_dt()
+    dt_sqrt = jnp.sqrt(dt)
+
+    # drift values
+    drifts = df(y, t, **kwargs)
+
+    # diffusion values
+    diffusions = dg(y, t, **kwargs)
+
+    # intermediate results
+    y_bars = jax.tree.map(lambda y0, drift, diffusion: y0 + drift * dt + diffusion * dt_sqrt, y, drifts, diffusions)
+    diffusion_bars = dg(y_bars, t, **kwargs)
+
+    # integral results
+    def f_integral(y0, drift, diffusion, diffusion_bar):
+        noise = brainstate.random.randn_like(y0) * dt_sqrt
+        noise_p2 = (noise ** 2 - dt) if sde_type == 'ito' else noise ** 2
+        minus = (diffusion_bar - diffusion) / 2 / dt_sqrt
+        return y0 + drift * dt + diffusion * noise + minus * noise_p2
+
+    integrals = jax.tree.map(f_integral, y, drifts, diffusions, diffusion_bars)
+    return integrals

brainmass/noise.py

@@ -0,0 +1,77 @@
+# 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.
+# ==============================================================================
+
+
+import brainstate
+import brainunit as u
+
+__all__ = [
+    'OUProcess',
+]
+
+
+class OUProcess(brainstate.nn.Dynamics):
+    r"""
+    The Ornstein–Uhlenbeck process.
+
+    The Ornstein–Uhlenbeck process :math:`x_{t}` is defined by the following
+    stochastic differential equation:
+
+    .. math::
+
+       \tau dx_{t}=-\theta \,x_{t}\,dt+\sigma \,dW_{t}
+
+    where :math:`\theta >0` and :math:`\sigma >0` are parameters and :math:`W_{t}`
+    denotes the Wiener process.
+
+    Parameters
+    ==========
+    in_size: int, sequence of int
+      The model size.
+    mean: ArrayLike
+      The noise mean value.
+    sigma: ArrayLike
+      The noise amplitude.
+    tau: ArrayLike
+      The decay time constant.
+    """
+
+    def __init__(
+        self,
+        in_size: brainstate.typing.Size,
+        mean: brainstate.typing.ArrayLike = 0.,  # noise mean value
+        sigma: brainstate.typing.ArrayLike = 1.,  # noise amplitude
+        tau: brainstate.typing.ArrayLike = 10.,  # time constant
+    ):
+        super().__init__(in_size=in_size)
+
+        # parameters
+        self.mean = mean
+        self.sigma = sigma
+        self.tau = tau
+
+    def init_state(self, batch_size=None, **kwargs):
+        size = self.in_size if batch_size is None else (batch_size, *self.in_size)
+        self.x = brainstate.HiddenState(u.math.zeros(size, unit=u.get_unit(self.mean)))
+
+    def reset_state(self, batch_size=None, **kwargs):
+        size = self.in_size if batch_size is None else (batch_size, *self.in_size)
+        self.x.value = u.math.zeros(size, unit=u.get_unit(self.mean))
+
+    def update(self):
+        df = lambda x: (self.mean - x) / self.tau
+        dg = lambda x: self.sigma / u.math.sqrt(self.tau)
+        self.x.value = brainstate.nn.exp_euler_step(df, dg, self.x.value)
+        return self.x.value