{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyMZiQw19pYNdaeM1UKzfEd7",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ColmTalbot/wcosmo/blob/timing-notebook/examples/wcosmo_timing.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Compare `wcosmo` and `astropy` timing.\n",
        "\n",
        "The two primary pieces of functionality we use are converting from luminosity distance to redshift, calculating the distance to redshift jacobian, and calculating the differential comoving volume.\n",
        "\n",
        "Timing the `wcosmo` implementation is non trivial as we rely on JIT compilation and also need to make sure we wait until the evaluation is complete.\n",
        "The steps are:\n",
        "\n",
        "- jit compile a wrapper function to call.\n",
        "- burn an evaluation for the compilation.\n",
        "- run the function and use `block_until_ready` to ensure we get the full time.\n",
        "\n",
        "We also time `wcosmo` with the `numpy` and `cupy` backends.\n",
        "Note that `cupy` also requires burning a call to compile the underlying `CUDA` code.\n",
        "\n",
        "We manually switch backends, although this can be done automatically using `GWPopulation`."
      ],
      "metadata": {
        "id": "nHKQkswiTuBQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "kcb71ObqYwC8"
      },
      "outputs": [],
      "source": [
        "!pip install wcosmo --quiet --progress-bar off"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import wcosmo\n",
        "\n",
        "\n",
        "def set_backend(backend):\n",
        "    from importlib import import_module\n",
        "    np_modules = dict(\n",
        "        numpy=\"numpy\",\n",
        "        jax=\"jax.numpy\",\n",
        "        cupy=\"cupy\",\n",
        "    )\n",
        "    linalg_modules = dict(\n",
        "        numpy=\"scipy.linalg\",\n",
        "        jax=\"jax.scipy.linalg\",\n",
        "        cupy=\"cupyx.scipy.linalg\",\n",
        "    )\n",
        "    setattr(wcosmo.wcosmo, \"xp\", import_module(np_modules[backend]))\n",
        "    setattr(wcosmo.utils, \"xp\", import_module(np_modules[backend]))\n",
        "    toeplitz = getattr(import_module(linalg_modules[backend]), \"toeplitz\")\n",
        "    setattr(wcosmo.utils, \"toeplitz\", toeplitz)\n",
        "\n",
        "\n",
        "ndata = np.random.uniform(1, 10, 1000000)"
      ],
      "metadata": {
        "id": "z-o1RrCPdXRz"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### wcosmo + jax + GPU"
      ],
      "metadata": {
        "id": "dseFs4ZVXXNu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "from jax import jit\n",
        "\n",
        "\n",
        "set_backend(\"jax\")\n",
        "\n",
        "jdata = jnp.array(ndata)\n",
        "\n",
        "\n",
        "@jit\n",
        "def time_jax_redshift(jdata):\n",
        "    return wcosmo.z_at_value(wcosmo.FlatwCDM(67, 0.3, -1).luminosity_distance, jdata)\n",
        "\n",
        "\n",
        "@jit\n",
        "def time_jax_dvcdz(jdata):\n",
        "    return wcosmo.FlatwCDM(67, 0.3, -1).differential_comoving_volume(jdata)\n",
        "\n",
        "\n",
        "burn_vals = time_jax_redshift(jdata)\n",
        "burn_vals = time_jax_dvcdz(jdata)"
      ],
      "metadata": {
        "id": "IY_u2ZI3PRua"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "\n",
        "_ = time_jax_redshift(jdata).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xjppj7LCPU_B",
        "outputId": "fb0c6ff9-96e6-4771-cb6b-67f1d5e40b85"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 1.56 ms, sys: 0 ns, total: 1.56 ms\n",
            "Wall time: 4.83 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "\n",
        "_ = time_jax_dvcdz(jdata).block_until_ready()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JgS3_9RVU1nH",
        "outputId": "ce2f0c85-0161-4f45-9cd7-3494077e7a83"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 632 µs, sys: 173 µs, total: 805 µs\n",
            "Wall time: 3.69 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### astropy + cpu\n",
        "\n",
        "Note that this is very slow in this case so we only use one percent of the full data.\n",
        "Since this is `numpy`-based the time scales linearly with the amount of data.\n",
        "\n",
        "In practice, most people when using `astropy` use intepolation to evaluate `z_at_value` at many points as is done in `wcosmo`."
      ],
      "metadata": {
        "id": "ozTDHo9iXgWi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from astropy import cosmology, units"
      ],
      "metadata": {
        "id": "_kOnSJMbXn9Q"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "\n",
        "_ = cosmology.z_at_value(\n",
        "    cosmology.FlatwCDM(67, 0.3, -1).luminosity_distance,\n",
        "    ndata[:10000] * units.Mpc,\n",
        ").value"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V-mMVWAhPgfp",
        "outputId": "2272fbab-0bdb-4cad-be8d-65bb6ecc249e"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 40.9 s, sys: 418 ms, total: 41.3 s\n",
            "Wall time: 52 s\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "\n",
        "_ = cosmology.FlatwCDM(67, 0.3, -1).differential_comoving_volume(\n",
        "    ndata[:10000],\n",
        ").value"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vF6Ab5jMVBdX",
        "outputId": "3f176351-bd05-4a09-f665-183cecb2e105"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 106 ms, sys: 27 µs, total: 106 ms\n",
            "Wall time: 107 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### wcosmo + numpy + cpu"
      ],
      "metadata": {
        "id": "3KTFoNDkXzTp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "set_backend(\"numpy\")"
      ],
      "metadata": {
        "id": "qGl50ks_PpGQ"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "\n",
        "_ = wcosmo.z_at_value(\n",
        "    wcosmo.FlatwCDM(67, 0.3, -1).luminosity_distance, ndata\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bQl0yn4uVSXp",
        "outputId": "6901b6ba-de14-43b0-dca7-6909127dfed9"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 89.5 ms, sys: 75.2 ms, total: 165 ms\n",
            "Wall time: 86.7 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "\n",
        "_ = wcosmo.FlatwCDM(67, 0.3, -1).differential_comoving_volume(ndata)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XjsGLM2sVXC-",
        "outputId": "1fc8817c-a0a1-44fa-dcf0-05b8f724b76f"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 65.8 ms, sys: 89.1 ms, total: 155 ms\n",
            "Wall time: 79.4 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### wcosmo + cupy + gpu\n",
        "\n",
        "The final test is using the `cupy` backend on the GPU.\n",
        "Typically this is much faster than `numpy` but slower than the `JAX` GPU code.\n",
        "Especially, not tested here is transfer between CPU/GPU which can be quite slow for `cupy`."
      ],
      "metadata": {
        "id": "Vn27QdGSYC3Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cupy\n",
        "\n",
        "set_backend(\"cupy\")\n",
        "\n",
        "cdata = cupy.asarray(ndata)\n",
        "\n",
        "_ = wcosmo.z_at_value(\n",
        "    wcosmo.FlatwCDM(67, 0.3, -1).luminosity_distance, cdata\n",
        ")\n",
        "_ = wcosmo.FlatwCDM(67, 0.3, -1).differential_comoving_volume(cdata)"
      ],
      "metadata": {
        "id": "wU3XF5PpVbwn"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "\n",
        "_ = wcosmo.z_at_value(\n",
        "    wcosmo.FlatwCDM(67, 0.3, -1).luminosity_distance, cdata\n",
        ")\n",
        "cupy.cuda.stream.get_current_stream().synchronize()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q4LUkMDuWDsk",
        "outputId": "a874883f-54e6-41e0-c8fc-50ce7a05039e"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 5.42 ms, sys: 17 µs, total: 5.44 ms\n",
            "Wall time: 5.45 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "\n",
        "_ = wcosmo.FlatwCDM(67, 0.3, -1).differential_comoving_volume(cdata)\n",
        "cupy.cuda.stream.get_current_stream().synchronize()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BCr3UfbaWEgJ",
        "outputId": "fefccd5c-5683-4e66-ecb0-d6635d1e5f80"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "CPU times: user 73.6 ms, sys: 0 ns, total: 73.6 ms\n",
            "Wall time: 83.6 ms\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "TjD_9bTFboru"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}