Neural Mass Models
==================

.. currentmodule:: brainmass

Neural mass models (NMMs) are mathematical descriptions of the average activity of large populations
of neurons. They provide a coarse-grained representation suitable for whole-brain network modeling,
capturing key dynamical features like oscillations, bistability, and excitability while remaining
computationally tractable.


Overview
--------

``brainmass`` provides two categories of neural mass models:

1. **Phenomenological Models**: Capture generic dynamical behaviors (oscillations, excitability)
   without direct physiological interpretation. Useful for studying synchronization, bifurcations,
   and other dynamical phenomena.

2. **Physiological Models**: Incorporate biophysical details like synaptic dynamics, membrane
   potentials, and ionic currents. Suitable for simulating realistic neural activity and linking
   to empirical data.


Model Comparison
----------------

The following table summarizes the key characteristics of each model:

.. list-table::
   :widths: 20 15 15 20 30
   :header-rows: 1

   * - Model
     - Category
     - Variables
     - Typical Use Cases
     - Key Features
   * - :class:`HopfStep`
     - Phenomenological
     - 2 (x, y)
     - Oscillation onset, rhythm generation
     - Supercritical Hopf bifurcation, normal form
   * - :class:`VanDerPolStep`
     - Phenomenological
     - 2 (x, y)
     - Nonlinear oscillations, relaxation
     - Self-sustained oscillations, limit cycle
   * - :class:`StuartLandauStep`
     - Phenomenological
     - 2 (x, y)
     - Oscillations with amplitude control
     - Complex amplitude, frequency control
   * - :class:`FitzHughNagumoStep`
     - Phenomenological
     - 2 (v, w)
     - Excitability, spike generation
     - Type II excitable system, fast-slow
   * - :class:`ThresholdLinearStep`
     - Phenomenological
     - 2 (rE, rI)
     - Fast dynamics, linear responses
     - Threshold activation, E-I populations
   * - :class:`KuramotoNetwork`
     - Phenomenological
     - 1 (θ)
     - Phase synchronization
     - Phase oscillators, order parameter
   * - :class:`JansenRitStep`
     - Physiological
     - 6 states
     - EEG generation, alpha rhythms
     - 3 populations (pyramidal, excitatory, inhibitory interneurons)
   * - :class:`WilsonCowanStep`
     - Physiological
     - 2 (rE, rI)
     - E-I dynamics, population rates
     - Classic two-population model, firing rates
   * - :class:`WongWangStep`
     - Physiological
     - 2 (S1, S2)
     - Decision making (perceptual choice)
     - Reduced two-population attractor (Wong & Wang 2006)
   * - :class:`MontbrioPazoRoxinStep`
     - Physiological
     - 2 (r, v)
     - Mean-field dynamics, theta neurons
     - Quadratic integrate-and-fire, exact reduction
   * - :class:`CoombesByrneStep`
     - Physiological
     - 2 (r, v)
     - Next-generation mean-field, conductance synapses
     - Ott-Antonsen reduction of QIF with synaptic conductance
   * - :class:`LarterBreakspearStep`
     - Physiological
     - 3 (V, W, Z)
     - Limit cycles, chaos, conductance-based dynamics
     - Modified Morris-Lecar with Na/K/Ca gating
   * - :class:`EpileptorStep`
     - Physiological
     - 6 states
     - Seizure onset/offset, epilepsy
     - Fast/slow subsystems + slow permittivity ``z``
   * - :class:`Generic2dOscillatorStep`
     - Phenomenological
     - 2 (V, W)
     - Flexible planar dynamics (excitable / bistable / oscillatory)
     - Configurable polynomial nullclines (TVB ``Generic2dOscillator``)
   * - :class:`WongWangExcInhStep`
     - Physiological
     - 2 (S_E, S_I)
     - Resting-state BOLD/FC, E-I balance
     - Two-population reduced Wong-Wang (Deco et al. 2014)
   * - :class:`LorenzStep`
     - Phenomenological
     - 3 (x, y, z)
     - Chaos, integration/coupling test fixture
     - Classic chaotic flow, positive Lyapunov exponent
   * - :class:`LinearStep`
     - Phenomenological
     - 1 (x)
     - Baseline node, coupling sanity checks
     - Damped linear dynamics (TVB ``Linear``)


Common API Pattern
------------------

All neural mass models follow a consistent interface:

.. code-block:: python

   import brainmass
   import jax.numpy as jnp
   import brainunit as u
   import brainstate

   # 1. Create model instance
   model = brainmass.WilsonCowanStep(
       in_size=10,  # 10 brain regions
       tau_E=10. * u.ms,
       tau_I=20. * u.ms,
       # ... other parameters
   )

   # 2. Initialize state variables
   model.init_all_states(batch_size=None)  # or batch_size=32 for batched simulations

   # 3. Run simulation
   def step(i):
       return model.update(rE_inp=0.1, rI_inp=0.05)

   outputs = brainstate.transform.for_loop(step, jnp.arange(1000))

   # 4. Access internal states
   excitatory_rate = model.rE.value  # current state
   inhibitory_rate = model.rI.value

**Key Methods:**

- ``__init__(in_size, **params)``: Initialize model with parameters
- ``init_state(batch_size=None)``: Initialize state variables before simulation
- ``update(**inputs)``: Advance dynamics by one time step, returns observable(s)
- ``.reset_state()``: Reset to initial conditions


Phenomenological Models
------------------------

Phenomenological models capture essential dynamical features without detailed biophysical mechanisms.


Oscillator Models
^^^^^^^^^^^^^^^^^

.. autosummary::
   :toctree: generated/
   :nosignatures:
   :template: classtemplate.rst

   XY_Oscillator
   HopfStep
   VanDerPolStep
   StuartLandauStep
   KuramotoNetwork

:class:`XY_Oscillator` is the shared two-variable ``(x, y)`` planar-oscillator base
that :class:`HopfStep`, :class:`StuartLandauStep` and :class:`VanDerPolStep` specialise;
instantiate the concrete subclasses for simulations.


**Example: Hopf Oscillator**

The Hopf oscillator exhibits a supercritical Hopf bifurcation, transitioning from a fixed point
to limit cycle oscillations as the bifurcation parameter crosses zero:

.. code-block:: python

   hopf = brainmass.HopfStep(
       in_size=5,
       w=0.1,  # intrinsic frequency (dimensionless)
       a=0.1,  # bifurcation parameter > 0 for oscillations
   )
   hopf.init_all_states()

   # Simulate
   ts = brainstate.transform.for_loop(
       lambda i: hopf.update(),
       jnp.arange(1000)
   )


Excitable Models
^^^^^^^^^^^^^^^^

.. autosummary::
   :toctree: generated/
   :nosignatures:
   :template: classtemplate.rst

   FitzHughNagumoStep
   ThresholdLinearStep


**Example: FitzHugh-Nagumo**

The FitzHugh-Nagumo model is a two-variable simplification of the Hodgkin-Huxley model,
exhibiting excitability and spike generation:

.. code-block:: python

   fhn = brainmass.FitzHughNagumoStep(
       in_size=1,
       tau=12.5 * u.ms,
       a=0.7,
       b=0.8,
   )
   fhn.init_all_states()

   # Apply external input to trigger spike
   output = fhn.update(inp=0.5)


Physiological Models
--------------------

Physiological models incorporate biophysical details and are suitable for linking to empirical data.


Rate-Based Models
^^^^^^^^^^^^^^^^^

.. autosummary::
   :toctree: generated/
   :nosignatures:
   :template: classtemplate.rst

   WilsonCowanStep
   JansenRitStep
   JansenRitTR
   WongWangStep

:class:`JansenRitTR` wraps :class:`JansenRitStep` with an inner dt-substep loop that
emits one output sample per acquisition repetition time (TR) -- the in-package
precedent for TR-rate observation.


**Example: Wilson-Cowan Model**

The Wilson-Cowan model describes the dynamics of excitatory and inhibitory populations:

.. code-block:: python

   wc = brainmass.WilsonCowanStep(
       in_size=(10,),  # 10 regions
       tau_E=10. * u.ms,
       tau_I=20. * u.ms,
       c_EE=12.0, c_EI=-4.0,
       c_IE=12.0, c_II=-3.0,
   )
   wc.init_all_states()

   # Add noise
   wc.noise_E = brainmass.OUProcess(in_size=(10,), sigma=0.3 * u.Hz, tau=20. * u.ms)
   wc.noise_I = brainmass.OUProcess(in_size=(10,), sigma=0.3 * u.Hz, tau=20. * u.ms)

   # Simulate
   ts = brainstate.transform.for_loop(
       lambda i: wc.update(rE_inp=0.5, rI_inp=0.2),
       jnp.arange(2000)
   )


**Example: Jansen-Rit Model**

The Jansen-Rit model simulates EEG-like signals from a cortical column with three neural populations:

.. code-block:: python

   jr = brainmass.JansenRitStep(in_size=1)
   jr.init_all_states()

   # Simulate and get EEG proxy
   eeg_signals = brainstate.transform.for_loop(
       lambda i: jr.update(p=220.),  # external input
       jnp.arange(5000)
   )

   # The output represents pyramidal membrane potential (EEG proxy)


Wilson-Cowan Variants
^^^^^^^^^^^^^^^^^^^^^^

Beyond the canonical :class:`WilsonCowanStep`, ``brainmass`` ships a family of
Wilson-Cowan variants that swap the activation nonlinearity, normalisation, or add
adaptation / explicit delays / a third population. They share the same
``(rE, rI)``-style state and ``update`` contract.

.. autosummary::
   :toctree: generated/
   :nosignatures:
   :template: classtemplate.rst

   WilsonCowanNoSaturationStep
   WilsonCowanSymmetricStep
   WilsonCowanSimplifiedStep
   WilsonCowanLinearStep
   WilsonCowanDivisiveStep
   WilsonCowanDivisiveInputStep
   WilsonCowanDelayedStep
   WilsonCowanAdaptiveStep
   WilsonCowanThreePopBase
   WilsonCowanThreePopulationStep

:class:`WilsonCowanThreePopBase` is the shared base for the three-population
variant; instantiate :class:`WilsonCowanThreePopulationStep` (or another concrete
variant) for simulations.


Spiking Network Reductions
^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. autosummary::
   :toctree: generated/
   :nosignatures:
   :template: classtemplate.rst

   MontbrioPazoRoxinStep


**Example: Montbrio-Pazo-Roxin Model (mean-field)**

The Montbrio-Pazo-Roxin model is an exact mean-field reduction of networks of quadratic integrate-and-fire neurons:

.. code-block:: python

   qif = brainmass.MontbrioPazoRoxinStep(
       in_size=5,
       tau=10. * u.ms,
       v_rest=-65. * u.mV,
       v_th=-50. * u.mV,
   )
   qif.init_all_states()

   # r = population firing rate, v = mean membrane potential
   outputs = brainstate.transform.for_loop(
       lambda i: qif.update(I_ext=2.0 * u.nA),
       jnp.arange(1000)
   )


Complex Mean-Field Models (TVB parity)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

These are literature-faithful ports of three complex mean-field models from
`The Virtual Brain <https://www.thevirtualbrain.org/>`_, validated against the TVB /
tvboptim reference equations. Full equations, default parameters and numbered
references are in each class docstring.

.. autosummary::
   :toctree: generated/
   :nosignatures:
   :template: classtemplate.rst

   CoombesByrneStep
   LarterBreakspearStep
   EpileptorStep


**Example: Coombes-Byrne (next-generation neural mass)**

A conductance-based exact mean field of QIF neurons. With the synaptic conductance
scale ``k = 0`` it reduces to :class:`MontbrioPazoRoxinStep` with ``J = 0``
(Coombes & Byrne, 2019):

.. code-block:: python

   cb = brainmass.CoombesByrneStep(in_size=1, Delta=1.0, eta=2.0, k=1.0, v_syn=-4.0)
   cb.init_all_states()
   out = brainmass.Simulator(cb, dt=0.1 * u.ms).run(40. * u.ms, monitors=["r", "v"])


**Example: Larter-Breakspear (conductance-based)**

A modified Morris-Lecar mean field whose pyramidal threshold variance ``d_V`` selects
the dynamical regime (fixed point, limit cycle, or chaos; Breakspear et al., 2003):

.. code-block:: python

   lb = brainmass.LarterBreakspearStep(in_size=1, d_V=0.57)  # limit-cycle band
   lb.init_all_states()
   out = brainmass.Simulator(lb, dt=0.1 * u.ms).run(400. * u.ms, monitors=["V"])


**Example: Epileptor (seizure dynamics)**

A six-variable model whose slow permittivity variable ``z`` autonomously ramps
seizures on and off; ``x0`` sets the epileptogenicity (Jirsa et al., 2014). The
``lfp()`` proxy (``x2 - x1``) is returned by ``update``:

.. code-block:: python

   epi = brainmass.EpileptorStep(in_size=1, x0=-1.6)  # epileptogenic node
   epi.init_all_states()
   out = brainmass.Simulator(epi, dt=0.1 * u.ms).run(
       2000. * u.ms, monitors={"lfp": lambda m: m.lfp(), "z": "z"}
   )


Canonical & Excitatory-Inhibitory Models (TVB parity)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

The remaining `The Virtual Brain <https://www.thevirtualbrain.org/>`_ node models —
the flexible planar oscillator, the two-population excitatory-inhibitory Wong-Wang
mean field, the Lorenz chaos fixture, and the linear baseline node — validated
against the TVB / tvboptim reference equations. Full equations, default parameters
and numbered references are in each class docstring.

.. autosummary::
   :toctree: generated/
   :nosignatures:
   :template: classtemplate.rst

   Generic2dOscillatorStep
   WongWangExcInhStep
   LorenzStep
   LinearStep


**Example: Generic 2D oscillator (configurable nullclines)**

A flexible planar system whose polynomial nullcline coefficients select the regime
(excitable, bistable, Morris-Lecar-like; Sanz-Leon et al., 2015):

.. code-block:: python

   # bistable-nullcline configuration
   g2d = brainmass.Generic2dOscillatorStep(in_size=1, a=1.0, b=0.0, c=-5.0, d=0.02)
   g2d.init_all_states()
   out = brainmass.Simulator(g2d, dt=0.1 * u.ms).run(200. * u.ms, monitors=["V", "W"])


**Example: Wong-Wang excitatory-inhibitory (resting-state workhorse)**

The two-population reduced Wong-Wang mean field; the feedback-inhibition weight
``J_i`` sets the local excitation-inhibition balance (Deco et al., 2014). The
population firing rates are available as ``H_e()`` / ``H_i()``:

.. code-block:: python

   rww = brainmass.WongWangExcInhStep(in_size=1, J_i=1.0, G=2.0)
   rww.init_all_states()
   out = brainmass.Simulator(rww, dt=0.1 * u.ms).run(
       3000. * u.ms, monitors={"S_e": "S_e", "rate_e": lambda m: m.H_e()}
   )


**Example: Lorenz (chaotic test fixture)**

The classic chaotic flow; use ``dt = 0.01 * u.ms`` (one natural time unit = 1 ms) and
keep trajectory comparisons short — nearby trajectories diverge exponentially
(Lorenz, 1963):

.. code-block:: python

   lz = brainmass.LorenzStep(in_size=1, sigma=10.0, rho=28.0)
   lz.init_all_states()
   out = brainmass.Simulator(lz, dt=0.01 * u.ms).run(5. * u.ms, monitors=["x", "y", "z"])


**Example: Linear node (baseline)**

A single damped linear node, ``dx/dt = gamma*x + coupling`` (distinct from the
two-population :class:`ThresholdLinearStep`); ``gamma < 0`` for stability:

.. code-block:: python

   lin = brainmass.LinearStep(in_size=1, gamma=-10.0)
   lin.init_all_states()
   out = brainmass.Simulator(lin, dt=0.1 * u.ms).run(20. * u.ms, monitors=["x"])


Adding Noise to Models
-----------------------

All models support adding stochastic noise through dedicated noise attributes:

.. code-block:: python

   # Create model
   model = brainmass.HopfStep(in_size=10, w=0.2)

   # Add Ornstein-Uhlenbeck noise
   model.noise = brainmass.OUProcess(
       in_size=10,
       sigma=0.1 * u.Hz,
       tau=50. * u.ms
   )

   model.init_all_states()

See :doc:`noise` for more details on noise processes.


Multi-Region Networks
----------------------

To create multi-region brain networks, combine models with coupling mechanisms:

.. code-block:: python

   import brainmass
   import jax.numpy as jnp
   import brainunit as u

   N = 90  # number of regions

   # Create node dynamics
   nodes = brainmass.HopfStep(in_size=N, w=0.3)

   # Create coupling
   coupling = brainmass.DiffusiveCoupling(
       conn=connectivity_matrix,  # (N, N) structural connectivity
       k=0.1,  # coupling strength
   )

   # Initialize
   nodes.init_all_states()
   coupling.init_all_states()

   # Simulation loop
   def network_step(i):
       local_activity = nodes.update()
       coupled_input = coupling(local_activity)
       nodes.x.value += coupled_input  # add coupling to node state
       return local_activity

   brainstate.transform.for_loop(network_step, jnp.arange(1000))

See :doc:`coupling` for comprehensive coupling documentation.


See Also
--------

- :doc:`noise` - Stochastic processes for adding variability
- :doc:`coupling` - Coupling mechanisms for network connectivity
- :doc:`forward` - Forward models for mapping to observed signals
- :doc:`../howto/choose_a_model` - Tutorial on selecting the right model
- :doc:`../tutorials/04_building_a_network` - Tutorial on multi-region networks