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

.. only:: html

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

        :ref:`Go to the end <sphx_glr_download_examples_benchmarks_plot_kaggle.py>`
        to download the full example code. or to run this example in your browser via Binder

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

.. _sphx_glr_examples_benchmarks_plot_kaggle.py:


Benchmark
=========

Numerical backends performace comparison, with GPU acceleration, as run on Kaggle (https://www.kaggle.com/code/benjaminvial/nannos-cpu)

.. GENERATED FROM PYTHON SOURCE LINES 15-22

.. code-block:: Python



    import os

    import matplotlib.pyplot as plt
    import numpy as np








.. GENERATED FROM PYTHON SOURCE LINES 23-24

Load data

.. GENERATED FROM PYTHON SOURCE LINES 24-43

.. code-block:: Python

    absolute_path = os.path.dirname("__file__")
    full_path = os.path.join(absolute_path, "kaggle/nannos_cpu/results_kaggle_cpu.npz")
    arch = np.load(full_path, allow_pickle=True)
    num_harmo = arch["num_harmo_real"]
    timings = arch["timedict"].item()
    # cases = arch["cases"].tolist()


    full_path = os.path.join(absolute_path, "kaggle/nannos_gpu/results_kaggle_gpu.npz")
    arch = np.load(full_path, allow_pickle=True)
    timings.update(arch["timedict"].item())
    # cases += arch["cases"].tolist()

    skip_first = True

    ifirst = 1 if skip_first else 0
    cases = ["numpy", "scipy", "autograd", "torch cpu", "torch gpu", "jax cpu", "jax gpu"]









.. GENERATED FROM PYTHON SOURCE LINES 44-46

Time vs. number of harmonics
------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 46-80

.. code-block:: Python

    tav = {}
    tstd = {}
    for case in cases:
        t = np.array(timings[case])[:, ifirst:]
        tav[case] = np.mean(t, axis=1)
        tstd[case] = np.std(t, axis=1)

    # colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
    colors = ["#467ccc", "#8746cc", "#cc4646", "#e4731d", "#46cc96", "#525252"]

    colors = colors[:3] + [colors[4], colors[4]] + [colors[5], colors[5]]

    plt.figure(figsize=(2, 2))
    for i, case in enumerate(cases):
        ls = "--" if case.split(" ")[-1] == "gpu" else "-"
        plt.errorbar(
            num_harmo,
            tav[case],
            yerr=tstd[case],
            label=case,
            ls=ls,
            color=colors[i],
            fmt="s",
            capsize=1,
            ms=1,
        )

    plt.xlabel("Number of harmonics")
    plt.ylabel("CPU time (s)")
    plt.xscale("log")
    plt.yscale("log")
    plt.legend()
    plt.tight_layout()




.. image-sg:: /examples/benchmarks/images/sphx_glr_plot_kaggle_001.png
   :alt: plot kaggle
   :srcset: /examples/benchmarks/images/sphx_glr_plot_kaggle_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 81-83

Speedup vs. number of harmonics
-----------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 83-111

.. code-block:: Python


    speedup_av = {}
    speedup_std = {}

    for case in cases[1:]:
        s = np.array(timings["numpy"])[:, ifirst:] / np.array(timings[case])[:, ifirst:]
        speedup_av[case] = np.mean(s, axis=1)
        speedup_std[case] = np.std(s, axis=1)


    plt.figure(figsize=(2, 2))
    for i, case in enumerate(cases[1:]):
        ls = "--" if case.split(" ")[-1] == "gpu" else "-"
        plt.errorbar(
            num_harmo,
            speedup_av[case],
            yerr=speedup_std[case],
            label=case,
            ls=ls,
            color=colors[i + 1],
            fmt="s",
            capsize=1,
            ms=1,
        )
    plt.xlabel("number of harmonics")
    plt.ylabel("Speedup vs. numpy")
    plt.legend()
    plt.tight_layout()



.. image-sg:: /examples/benchmarks/images/sphx_glr_plot_kaggle_002.png
   :alt: plot kaggle
   :srcset: /examples/benchmarks/images/sphx_glr_plot_kaggle_002.png
   :class: sphx-glr-single-img






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

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

**Estimated memory usage:**  519 MB


.. _sphx_glr_download_examples_benchmarks_plot_kaggle.py:

.. only:: html

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

    .. container:: binder-badge

      .. image:: images/binder_badge_logo.svg
        :target: https://mybinder.org/v2/gh/nannos/nannos.gitlab.io/doc?filepath=notebooks/examples/benchmarks/plot_kaggle.ipynb
        :alt: Launch binder
        :width: 150 px

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

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

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

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

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

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


.. only:: html

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

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