Added AdExMF model

[GiM6114]

Added the biophysically-grounded AdEx neural mass model (1st order and 2nd order) adapted for JAX.

[mapi1]

Hey and thanks a lot for your contribution, this model is quite interesting and more involved than the typical model.

A few things still remain before merging it:

• Exporting them from dynamics/tvb/__init__.py
• Add it to the docs docs/_quarto.yml and docs/network_dynamics/dynamics.qmd
• A note to CHANGELOG.md
• See commets

Optionally:

• There are unit tests that test against TVB tests/test_network_dynamics/test_tvb_comparison.py, as this model is the same (except for some parameters) as zerlaut.py it could be added there.

[mapi1]

Overwrite if init need to pass kwargs along: def __init__(self, **kwargs): super().__init__(**kwargs)

[mapi1]

Nice use of the JAX transforms here! One latent issue: in_axes=(0, 0, 0, None) maps
fe, fi, and W by node but shares params across all nodes. If a parameter
used by TF (E_L_e, Q_e, g_L_e, etc.) has one value per node, TF_e
returns a vector and grad fails because its output must be scalar:

TypeError: Gradient only defined for scalar-output functions. Output had shape: (68,).

Mapping params by node is tricky because arrays such as the 10-element
P_e/P_i coefficients are indistinguishable from node-wise parameters by
shape alone. A simpler option might be to differentiate the sum across nodes. Since
nodes do not interact inside TF, this gives the same derivative for each node:

def _elementwise_derivatives(TF):
    """Derivatives of a node-wise transfer function.

    Each output node must depend only on the corresponding entries of fe, fi,
    and W. Nonlocal transfer functions are not supported.
    """
    def total(fe, fi, W, params):
        return TF(fe, fi, W, params)[0].sum()

    d_fe = jax.grad(total, argnums=0)
    d_fi = jax.grad(total, argnums=1)

    def total_d_fe(fe, fi, W, params):
        return d_fe(fe, fi, W, params).sum()

    def total_d_fi(fe, fi, W, params):
        return d_fi(fe, fi, W, params).sum()

    d2_ee = jax.grad(total_d_fe, argnums=0)
    d2_ei = jax.grad(total_d_fe, argnums=1)
    d2_ii = jax.grad(total_d_fi, argnums=1)

    def call(fe, fi, W, params):
        args = (fe, fi, W, params)
        return d_fe(*args), d_fi(*args), d2_ee(*args), d2_ei(*args), d2_ii(*args)

    return call

Then the call site could be sth. like:

self.deriv_e = _elementwise_derivatives(self.TF_e)

d_fe, d_fi, d2_ee, d2_ei, d2_ii = self.deriv_e(fe, fi, W, params)

The mixed second derivatives are equal, so the dynamics can use
2 * C_ei * d2_ei instead of computing both.

This relies on the current TF being node-wise. If it later includes something
nonlocal, such as normalization or electrical coupling across nodes, the sum
would mix cross-node derivatives and we would need the full Jacobian/Hessian.

I checked this on CPU and GPU with 1–10,000 nodes: results match, and node-wise
parameters work. One remaining limitation is that node-wise polynomial
coefficients (P_e shaped [10, n_nodes]) are still unsupported because
threshold_func ends with features @ P.

[mapi1]

Docstring could be completed, to contain basic equations etc.

[mapi1]

Type annotation is incorrect, a simple array is returned

[mapi1]

Comments are swapped here