JaxPM/notebooks/Introduction.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/DifferentiableUniverseInitiative/JaxPM/blob/main/notebooks/Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install --quiet git+https://github.com/DifferentiableUniverseInitiative/JaxPM.git"
      ],
      "metadata": {
        "id": "9Jy5BL1XiK1s"
      },
      "id": "9Jy5BL1XiK1s",
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "c5f42bbe",
      "metadata": {
        "id": "c5f42bbe",
        "outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Populating the interactive namespace from numpy and matplotlib\n"
          ]
        }
      ],
      "source": [
        "%pylab inline\n",
        "\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jax_cosmo as jc\n",
        "\n",
        "from jax.experimental.ode import odeint\n",
        "\n",
        "from jaxpm.painting import cic_paint\n",
        "from jaxpm.pm import linear_field, lpt, make_ode_fn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "281b4d3b",
      "metadata": {
        "id": "281b4d3b"
      },
      "outputs": [],
      "source": [
        "mesh_shape= [256, 256, 256]\n",
        "box_size  = [256.,256.,256.]\n",
        "snapshots = jnp.linspace(0.1,1.,2)\n",
        "\n",
        "@jax.jit\n",
        "def run_simulation(omega_c, sigma8):\n",
        "    # Create a small function to generate the matter power spectrum\n",
        "    k = jnp.logspace(-4, 1, 128)\n",
        "    pk = jc.power.linear_matter_power(jc.Planck15(Omega_c=omega_c, sigma8=sigma8), k)\n",
        "    pk_fn = lambda x: jc.scipy.interpolate.interp(x.reshape([-1]), k, pk).reshape(x.shape)\n",
        "\n",
        "    # Create initial conditions\n",
        "    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=jax.random.PRNGKey(0))\n",
        "\n",
        "    # Create particles\n",
        "    particles = jnp.stack(jnp.meshgrid(*[jnp.arange(s) for s in mesh_shape]),axis=-1).reshape([-1,3])\n",
        "\n",
        "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
        "    \n",
        "    # Initial displacement\n",
        "    dx, p, f = lpt(cosmo, initial_conditions, particles, 0.1)\n",
        "    \n",
        "    # Evolve the simulation forward\n",
        "    res = odeint(make_ode_fn(mesh_shape), [particles+dx, p], snapshots, cosmo, rtol=1e-5, atol=1e-5)\n",
        "    \n",
        "    # Return the simulation volume at requested \n",
        "    return res[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "826be667",
      "metadata": {
        "id": "826be667",
        "outputId": "dc43b5c4-a004-41bf-f2c8-128c17cf4de1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1 loop, best of 5: 29.2 s per loop\n"
          ]
        }
      ],
      "source": [
        "res = run_simulation(0.25, 0.8)\n",
        "%timeit res = run_simulation(0.25, 0.8)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "id": "4e012ce8",
      "metadata": {
        "id": "4e012ce8",
        "outputId": "75390318-8072-481f-ffb9-ec09cd71cb1d",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 323
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<matplotlib.image.AxesImage at 0x7f3502382c90>"
            ]
          },
          "metadata": {},
          "execution_count": 11
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAloAAAEgCAYAAABsCt3QAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nOy9eZAc133n+cmzsrKy7upGdwMoQd0NECIIniJbMjkyaYoKgj5k71i7lixbEfIh74St2Q3HLD0T3pUcoXDsxqwjbNobjnHYGluhwzP0ruWxRHJ0rGgNZQmkRIGkSIG4BDYafdZ95H3sH/kyWYBsjxRD2l5H/yIQQDeqMl++9/K9b36/398vpSRJ2I/92I/92I/92I/92I/XPuR/6Absx37sx37sx37sx378U419oLUf+7Ef+7Ef+7Ef+/E6xT7Q2o/92I/92I/92I/9eJ1iH2jtx37sx37sx37sx368TrEPtPZjP/ZjP/ZjP/ZjP16n2Ada+7Ef+7Ef+7Ef+7Efr1O8bkBLkqQHJUl6WZKkC5Ik/drrdZ792I/92I/XOvbXr/3Yj/14rUJ6PepoSZKkAOeAB4AN4Bng3UmSvPSan2w/9mM/9uM1jP31az/2Yz9ey3i9GK27gAtJklxKksQH/hR45+t0rv3Yj/3Yj9cy9tev/diP/XjNQn2djnsQuDLz8waw9rd9uFIpJPX6AaKoDCjEsZ42Th2g61tAQhyXiGONJCkQhnXCMESSJBRFAmSiKCJJYgDiOEKWVRRFR9e3kKSIJFGJohJRVCaOA6LIR5ZVZFkjSWJkWSGKAsLQRZIkZFlFkhRkWUGSFBRliix7RFEJ348JQw9JkpAkOf9enETocUQiSQRSimEjJEzdQtNMACRJRlFkkkQijqP8GHEcEkU+cZxeR9YmTTNR1RiQSBKZOJaI4xCQUBQdSBlJRXGQpIAkUQnDIlHkkyQRIKHrBTRtlzBs4LoBcRySJDGKoqEoBZIkQpJkZFlDlhNUdUAQNIjjEFlWieOIMHTz48mygixrKIqOokyI4wJhmORtSZKYOA4JQ0+MgSXGMCSKAmQ5nXayrIjPRhSLHooyJo6LxLFJHBeRpABIiCJN9JUsxiXIx9T3faLIR1UNFEUTbdYwjA1UdUAUWUhSiCx7JImMLLuAQhRZyLKDLAd43gK+vyT6XCKOE5IkJoq8vK/TuZWIPlJIkkS0J/13HKf9mvWZqhoABIFNFPliDmnis6HoQ5WUUU6PWyio6PoGsuwQxyZB0CJJCkiShySFxHGRKErHPDtnkiSoagFV9dL5FplAgiRFM3NjgqKMUZQpmhajKBZJEhGGfcKwBsQkiSr63RDnUcV9EpAk2ZxMSJJX76103kZ5e6LIxzBqaJpLFJXEPePjuhZJEqNpOpLkI0kRsmyzvr7eSZJk7vtcW/4+4vtavwCakpyUJJmWrOBJCjIJl/UyehKz5A3ZTWJCoIREXcx/KYmJSJgmCYokUZIUzhlVdLXIwekuTuQDUJAkVCCRVRIkjCSiH0eUJAlkFS2JiZDwZAUjjvCSiBAwgGGSUJUkpknCkORvXPALSISiHaokMS/JlJAIJZmprCKTUIwCRklMXVaYyGr6hB75TJMET8yzkHRDUQAVKT9XQZIoIBHJSr4uluKQrTgEoCrJqIAsKYwVjRiJZuRxOQ4JkyQ/7pKsoIrzZOeLZhSZgiQBME0SQhIioCLJmJICQJxEDMXnsz5FjEkkSezEIcMkwRDHQZy3hJS3oSBJ7CYxrvhcFYmhuP4S6fe8mZ/DmeNUZBVfkkniABWQxNjtibE0Rd/044iCJGFIMpEk40syHdVIxz4OaUQehThiomhokU9X7HseEIoxzNqS9VdBkqhIMqH4vULCOImZiu+qSHgkFJDyvvGShCnXKl7Z2GaR9fPs2AMYSCiiH7MxqskKV7USjqzSCF0CSUYloRiHqEmMFIdpW5FAkokliWimvQkwljVUEqwoIAESSUJKEgJJRktiFBLi/F4I6cYRIQlqdl2KTjEO6cRRfr3zkpyPYdZeU9HwJQUl8hgmMR5Q4NrwgU4cSXwP8XoBrf9qSJL0i8AvAszNafz6r//P9Ho/Qq32BXx/iW73x6jVvsSBAx/Ftm/EcW4AQJYdHOcYtn0Cx0lYWvoMzeZnsO0b2dh4mOFwk0KhQqv1MseOvY9KxcG2dabTG+n3T7Gz8zOMRgnj8SaGUcMwakiSwnh8ldFog729b6PrJer1Zer1FVx3gK6XqFQOUSpdxvPabG2dZTi8giwr6HqZIJiys/M8Uu8CZhQQyApmHGLGIc+aLVYX72B+/gS+P6VUmqPVWkWSQiaTEZPJNkkS43ljJpOtdFDUIoZRxTDqmGaTcnmeJJHx/Slh6OB5I9GHcg7gDKOGrodiQzbQtD0kKcR134Cm71AyX6LX+2d0Oi+jaUUkSaZQqKBpJsPhOuXyQVZWHkHXt0gSlVde+QidzjYAhUJFAMoUCGZAtVw+SKFQpNu9xGSyTQbCosjHdQdIkky12kbXLYLAJkligsDGshYoFCqYZg1F6RMEdZrNJ2g0/hLHuYF+/0F8f4kgsNC0PoXCFnFs4PtLGMYloqiMpu2SJBqdzhEGg8uoqoGmmaJvJNpti2bzzymVnsM0X6Rcfpbd3XfT7f4IhcIWxeLLyLKLab7IcHgvL7/8IWx7jySJ8f0JrjtgMtlmPN4kDF0m402GssZq6waq1TbT6a6YGxal0jy+P2E63cXzRui6RaOxSqk0Txi6TKe7BIGdA7HpdA9dL9FqHRcgFHTdYn7+JpaW/phi8Rye12YyuYMoKlMuf41i8RzD4b1sb/+gAP8yg8FlPG9Eo7FKrTakWHwZgPH4DlR1kt8vxeI5VLVPufwUR458kVLpGNvbz7Gz82v4/pIAmSqGcYk4NphOb6HTOUm/f0nMRyMfP9cd4HkjosjHNFv4/gTb7gBwpnuOtxy8i5tvuRsATe1Trz/Bzs77uXr1MgcO3Mzq6i8QRXWCYI5f+IWHX3k915jXO2bXsLak8FL1MADrusUBpwdApJl8R1Z4LHS5EIesKTrvN2q8GNicj0POxyG9JGZN0Wk1j/HX73+KSuUb3PTb/5JO9xz3lOYBWAxsTpfmqUY+7mSbhiTT0kwuFiq0QpeLhQrfKjZ4YLRB25/wUXfA6cjnlGpwn9nC8Cd82BtxVIC805FPT2yyAKvi9w1J5qis8uOFCp+vHALggdEGn/ZGeTuP6RZaEvNvnR69JM7/ZMfJzgHwy7rFlmZixiGPVdv81oGbuX26x29uPsP/Nd0F4IO6xYU4pCuOfzryOS1J1xwXYE3RWZVV1hSdgqLzZX9CL4lpSHL+/V4S5+c/H4ecUg3uF2vkmdDldOTTkGTeo5kEksx65LOoGlwsVFjxRrxrkq55dynpw352zefjkKYkcz4O83M1JZmG+HP0uv7L2pKdrynJ+Xe6icYHdYsv15dphi6F4ToAdyo6/STmQhxeM75PWousF8ocdwfcZndYm+7ye/4kPYf0qij1eOiClLbhlGrwkHjY+3hg80HdQjNqLAY2ABMx7wDaik4gybwY2DQkmapuMfQnPOJPOC3AflOSuUvR8z7IrgmgK9qcRTYma6IPs35ZU3S2Kof4s/oy90y2ucnpUY18FgObvmhLXZJ5JvI5oZlsaSat0KUa+diyylDRMeOQQJIJJBlbAFUtiXms2kaLIx4aXaEVunRUAy2J6YhrsmSVM6Gbj8GqrPKIP2FN0bm8eDvL3oj7xpv5Z+LmMX7rwM386s7z/Oed52mIa50d8//FHfzNC8PfEK8X0LoKHJ75+ZD4XR5JkvwB8AcAy8uNRFX7lErPE8dFut0fY3Pz64ThW6lWnySKynjeYYbD25BllUrlq1QqTyFJb8G2b0HT+mjaHrq+RaVyiPn5T3D06M/xhjf8j+x1PovjLOK6y8SxAWgUixaqaoiNcUypNEcQ2AwGl3HdQQ5EbHuP0egqhlHDshZQ1T62vYznjYEEVTXw/THT6S5PeyPuiwJsRUNLYqphyjBsqUVWAdvu5KyQL
            "text/plain": [
              "<Figure size 720x360 with 2 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          }
        }
      ],
      "source": [
        "figure(figsize=[10,5])\n",
        "subplot(121)\n",
        "imshow(cic_paint(jnp.zeros(mesh_shape), res[0]).sum(axis=0),cmap='gist_stern')\n",
        "subplot(122)\n",
        "imshow(cic_paint(jnp.zeros(mesh_shape), res[1]).sum(axis=0),cmap='gist_stern')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4e050871",
      "metadata": {
        "id": "4e050871"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.2"
    },
    "colab": {
      "name": "Introduction.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}