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

.. only:: html

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

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

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

.. _sphx_glr_basic_usage_example_eigenvalues.py:

Eigenvalues
===============

This example demonstrates how to compute a subset of eigenvalues of a linear
operator, using :func:`scipy.sparse.linalg.eigsh`. Concretely, we will compute
leading eigenvalues of the Hessian.

As always, imports go first.

.. GENERATED FROM PYTHON SOURCE LINES 10-24

.. code-block:: Python


    import numpy
    import scipy
    import torch
    from torch import nn

    from curvlinops import HessianLinearOperator
    from curvlinops.examples.functorch import functorch_hessian
    from curvlinops.examples.utils import report_nonclose

    # make deterministic
    torch.manual_seed(0)
    numpy.random.seed(0)








.. GENERATED FROM PYTHON SOURCE LINES 25-30

Setup
-----

We will use synthetic data, consisting of two mini-batches, a small MLP, and
mean-squared error as loss function.

.. GENERATED FROM PYTHON SOURCE LINES 31-54

.. code-block:: Python


    N = 20
    D_in = 7
    D_hidden = 5
    D_out = 3

    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    X1, y1 = torch.rand(N, D_in).to(DEVICE), torch.rand(N, D_out).to(DEVICE)
    X2, y2 = torch.rand(N, D_in).to(DEVICE), torch.rand(N, D_out).to(DEVICE)

    model = nn.Sequential(
        nn.Linear(D_in, D_hidden),
        nn.ReLU(),
        nn.Linear(D_hidden, D_hidden),
        nn.Sigmoid(),
        nn.Linear(D_hidden, D_out),
    ).to(DEVICE)
    params = [p for p in model.parameters() if p.requires_grad]

    loss_function = nn.MSELoss(reduction="mean").to(DEVICE)









.. GENERATED FROM PYTHON SOURCE LINES 55-59

Linear operator
------------------

We are ready to setup the linear operator. In this example, we will use the Hessian.

.. GENERATED FROM PYTHON SOURCE LINES 60-64

.. code-block:: Python


    data = [(X1, y1), (X2, y2)]
    H = HessianLinearOperator(model, loss_function, params, data)








.. GENERATED FROM PYTHON SOURCE LINES 65-70

Leading eigenvalues
-------------------

Through :func:`scipy.sparse.linalg.eigsh`, we can obtain the leading
:math:`k=3` eigenvalues.

.. GENERATED FROM PYTHON SOURCE LINES 71-78

.. code-block:: Python


    k = 3
    which = "LA"  # largest algebraic
    top_k_evals, _ = scipy.sparse.linalg.eigsh(H, k=k, which=which)

    print(f"Leading {k} Hessian eigenvalues: {top_k_evals}")





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

 .. code-block:: none

    Leading 3 Hessian eigenvalues: [1.4109129 1.4281708 1.4943049]




.. GENERATED FROM PYTHON SOURCE LINES 79-85

Verifying results
-----------------

To double-check this result, let's compute the Hessian with
:code:`functorch`, compute all its eigenvalues with
:func:`scipy.linalg.eigh`, then extract the top :math:`k`.

.. GENERATED FROM PYTHON SOURCE LINES 86-95

.. code-block:: Python


    H_functorch = (
        functorch_hessian(model, loss_function, params, data).detach().cpu().numpy()
    )
    evals_functorch, _ = scipy.linalg.eigh(H_functorch)
    top_k_evals_functorch = evals_functorch[-k:]

    print(f"Leading {k} Hessian eigenvalues (functorch): {top_k_evals_functorch}")





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

 .. code-block:: none

    Leading 3 Hessian eigenvalues (functorch): [1.4109125 1.4281707 1.4943047]




.. GENERATED FROM PYTHON SOURCE LINES 96-97

Both results should match.

.. GENERATED FROM PYTHON SOURCE LINES 98-102

.. code-block:: Python


    print(f"Comparing leading {k} Hessian eigenvalues (linear operator vs. functorch).")
    report_nonclose(top_k_evals, top_k_evals_functorch)





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

 .. code-block:: none

    Comparing leading 3 Hessian eigenvalues (linear operator vs. functorch).
    Compared arrays match.




.. GENERATED FROM PYTHON SOURCE LINES 103-106

:func:`scipy.sparse.linalg.eigsh` can also compute other subsets of
eigenvalues, and also their associated eigenvectors. Check out its
documentation for more!


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

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


.. _sphx_glr_download_basic_usage_example_eigenvalues.py:

.. only:: html

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

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

      :download:`Download Jupyter notebook: example_eigenvalues.ipynb <example_eigenvalues.ipynb>`

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

      :download:`Download Python source code: example_eigenvalues.py <example_eigenvalues.py>`

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

      :download:`Download zipped: example_eigenvalues.zip <example_eigenvalues.zip>`


.. only:: html

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

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