Coupled Jansen-Rit 3-column network

def _stable_expit(x):
    """Numerically stable sigmoid 1/(1+exp(-x)), scalar or array."""
    x = np.asarray(x, dtype=float)
    return np.where(
        x >= 0,
        1.0 / (1.0 + np.exp(-x)),
        np.exp(x) / (1.0 + np.exp(x)),
    )

def _jr_sigmoid(v, v_max=5.0, r=0.56, v0=6.0):
    """Jansen-Rit population sigmoid: S(v) = v_max / (1 + exp(r*(v0 - v)))."""
    return v_max * _stable_expit(r * (v - v0))

def rhs(t: float, y: np.ndarray) -> np.ndarray:
    dy = np.empty(18)

    col_outputs = np.empty(n_cols)
    for j in range(n_cols):
        off = j * 6
        col_outputs[j] = _jr_sigmoid(y[off + 1] - y[off + 2])

    total_output = col_outputs.sum()

    for j in range(n_cols):
        off = j * 6
        x0, x1, x2, x3, x4, x5 = y[off:off + 6]

        p_j = p_ext + w * (total_output - col_outputs[j])

        S_pyr = _jr_sigmoid(x1 - x2)
        S_exc = _jr_sigmoid(C1 * x0)
        S_inh = _jr_sigmoid(C3 * x0)

        dy[off + 0] = x3
        dy[off + 1] = x4
        dy[off + 2] = x5
        dy[off + 3] = A * a * S_pyr - 2.0 * a * x3 - a2 * x0
        dy[off + 4] = A * a * (p_j + C2 * S_exc) - 2.0 * a * x4 - a2 * x1
        dy[off + 5] = B * b * C4 * S_inh - 2.0 * b * x5 - b2 * x2

    return dy