.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/06_custom_coupling.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_06_custom_coupling.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_06_custom_coupling.py:


Custom coupling functions: no registry needed
=================================================

lyapax coupling is a plain callable
``coupling_fn(cvar_state, weights, params) -> coupling`` -- not a fixed set
of named coupling "kinds" dispatched through hardcoded if/elif branches
internally. That means adding a new coupling rule never requires touching
the library's source: any function with the right signature is a
first-class coupling, exactly like the built-in ``linear_coupling``/
``sigmoidal_coupling``/``kuramoto_coupling`` in ``lyapax.coupling``. This
demo writes one from scratch -- no import from ``lyapax.coupling`` at all
-- and plugs it directly into ``network_problem``, reproducing the
exact linear-network result from
:ref:`04_linear_network.py <sphx_glr_auto_examples_04_linear_network.py>`.

**The system.** Same idea as
:ref:`04_linear_network.py <sphx_glr_auto_examples_04_linear_network.py>` --
scalar
linear node dynamics ``x_i' = gamma * x_i + c_i`` coupled linearly, so
the whole network reduces to ``x' = (gamma*I + G*W) x`` -- but scaled
down to 2 nodes joined by a single edge (``weights = [[0,1],[1,0]]``).
That adjacency matrix has eigenvalues ``{1, -1}``, so ``A``'s eigenvalues
are exactly ``gamma +/- G = {-0.8, -2.2}``. The only thing this example
changes versus
:ref:`04_linear_network.py <sphx_glr_auto_examples_04_linear_network.py>` is
the coupling term itself:
``my_linear_coupling`` below reimplements ``lyapax.coupling.linear_coupling``
by hand (same ``G * W @ x`` formula) purely to demonstrate that
``network_problem`` accepts *any* callable of the right signature --
the result should (and does) match
:ref:`04_linear_network.py <sphx_glr_auto_examples_04_linear_network.py>`'s
exact-eigenvalue check
bit-for-bit in spirit, just for this smaller network.

.. GENERATED FROM PYTHON SOURCE LINES 37-52

.. code-block:: Python

    import os

    os.environ["JAX_PLATFORMS"] = "cpu"

    import jax
    import jax.numpy as jnp
    import numpy as np

    jax.config.update("jax_enable_x64", True)

    from lyapax.core import lyapunov_spectrum
    from lyapax.network import Network, network_problem
    from lyapax.simulator import ModelSpec, Parameter, StateVar, build_jax_dfun









.. GENERATED FROM PYTHON SOURCE LINES 53-61

.. code-block:: Python

    def my_linear_coupling(cvar_state, weights, params):
        """A user-written coupling -- deliberately not imported from
        lyapax.coupling, to prove the extension point needs no library
        changes and no registration step."""
        G = params["G"]
        return G * jnp.einsum("ts,cs->ct", weights, cvar_state)









.. GENERATED FROM PYTHON SOURCE LINES 62-91

.. code-block:: Python

    weights = np.array([[0., 1.], [1., 0.]])
    gamma, G = -1.5, 0.7
    A = gamma * np.eye(2) + G * weights
    expected = np.sort(np.linalg.eigvalsh(A))[::-1]

    model = ModelSpec(
        name="linear_node",
        state_variables=(StateVar("x", default_init=0.0),),
        parameters=(Parameter("gamma", gamma),),
        cvar=("x",),
        dfun_str={"x": "gamma * x + c"},
    )
    dfun = build_jax_dfun(model)
    params = {"gamma": gamma, "G": G}
    dt = 1e-3
    network = Network(weights=jnp.array(weights), cvar_indices=model.cvar_indices)
    problem = network_problem(
        dfun, network, my_linear_coupling,  # <- user function, not a lyapax builder
        params=params, state0=jnp.array([0.2, -0.3]), dt=dt,
    )

    result = lyapunov_spectrum(
        problem, n_steps=20_000, renorm_every=10, t_transient=5.0,
    )

    estimate = np.array(result.exponents)
    print(f"exact eigenvalues:  {expected}")
    print(f"lyapax estimate:    {estimate}")
    print(f"max abs error:      {np.max(np.abs(estimate - expected)):.2e}")




.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    exact eigenvalues:  [-0.8 -2.2]
    lyapax estimate:    [-0.79999995 -2.19999819]
    max abs error:      1.81e-06





.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 0.665 seconds)


.. _sphx_glr_download_auto_examples_06_custom_coupling.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: 06_custom_coupling.ipynb <06_custom_coupling.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: 06_custom_coupling.py <06_custom_coupling.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: 06_custom_coupling.zip <06_custom_coupling.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_