JaxPM_highres/notebooks/01-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/01-Introduction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Installing JaxPM\n",
        "!pip install --quiet --upgrade jax\n",
        "!pip install --quiet jaxpm"
      ],
      "metadata": {
        "id": "r46BAlndVOeI"
      },
      "id": "r46BAlndVOeI",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "c5f42bbe",
      "metadata": {
        "id": "c5f42bbe"
      },
      "outputs": [],
      "source": [
        "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": 6,
      "id": "281b4d3b",
      "metadata": {
        "id": "281b4d3b"
      },
      "outputs": [],
      "source": [
        "mesh_shape = [128, 128, 128]\n",
        "box_size = [128., 128., 128.]\n",
        "snapshots = jnp.array([0.1, 0.5, 1.0])\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: jnp.interp(x.reshape([-1]), k, pk).reshape(x.shape)\n",
        "\n",
        "    # Create initial conditions and particles\n",
        "    initial_conditions = linear_field(mesh_shape, box_size, pk_fn, seed=jax.random.PRNGKey(0))\n",
        "    particles = jnp.stack(jnp.meshgrid(*[jnp.arange(s) for s in mesh_shape]),axis=-1).reshape([-1,3])\n",
        "\n",
        "    # Initial displacement\n",
        "    cosmo = jc.Planck15(Omega_c=omega_c, sigma8=sigma8)\n",
        "    dx, p, f = lpt(cosmo, initial_conditions, particles, a=0.1)\n",
        "\n",
        "    # Evolve the simulation forward\n",
        "    res = odeint(make_ode_fn(mesh_shape), [particles + dx, p], snapshots, cosmo, rtol=1e-8, atol=1e-8)\n",
        "\n",
        "    # Return the simulation volume at requested timesteps\n",
        "    return initial_conditions ,  particles + dx , res[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "826be667",
      "metadata": {
        "id": "826be667"
      },
      "outputs": [],
      "source": [
        "initial_conditions , lpt_particles , ode_particles = run_simulation(0.25, 0.8)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "id": "f9c05d34",
      "metadata": {
        "id": "f9c05d34",
        "outputId": "ed862b4d-4712-4b59-8582-e438246b73b7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 478
        }
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 2000x500 with 4 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAB8UAAAH/CAYAAADOlQwMAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOydeZhcZZX/z721L11dvW9JurMnhBAwEAQDYZNNRUTg58qiiCIM7uM2KijoqLiMOiLquIzbjICCzriBgAbZIewQsnQ6S3d6r15qr7rv74/Ydb7n7a4QMJBlzud58uR21Vv3vvddznnvvVXfr2OMMaQoiqIoiqIoiqIoiqIoiqIoiqIoiqIoByHuvq6AoiiKoiiKoiiKoiiKoiiKoiiKoiiKorxU6ENxRVEURVEURVEURVEURVEURVEURVEU5aBFH4oriqIoiqIoiqIoiqIoiqIoiqIoiqIoBy36UFxRFEVRFEVRFEVRFEVRFEVRFEVRFEU5aNGH4oqiKIqiKIqiKIqiKIqiKIqiKIqiKMpBiz4UVxRFURRFURRFURRFURRFURRFURRFUQ5a9KG4oiiKoiiKoiiKoiiKoiiKoiiKoiiKctCiD8UVRVEURVEURVEURVEURVEURVEURVGUgxZ9KK4oiqIoiqIoiqIoiqIoiqIoiqIoiqIctOhDceWAxHEcuuqqq/aobFdXF1100UUv+Bhbtmwhx3HoRz/60Qv+7P7AXXfdRY7j0F133VV57aKLLqKurq49+vxVV11FjuO8NJU7gNiX4+CF9JeiKIqy/zFTLt5TfvSjH5HjOLRly5bnLfti1zovJfuqTi+k3RRFUZT9kwcffJCOPfZYisVi5DgOPfroo//Q9emeXlcd6PcAno9/ZF3yj3LCCSfQCSec8LIfV1EURXn50Pz90qD3phVl76IPxZV9wtQNy4ceemiv7O+ee+6hq666ilKp1F7Z34uhv7+fPvzhD9OSJUsoGo1SLBajlStX0jXXXLNP67U7MpkMXXXVVfvkoljZRW9vL1111VX06KOP7uuqTOM//uM/aOnSpRQOh2nhwoX0zW9+c19XSVEUZY/WEFMXjVP/fD4fzZkzh97whjdU4u1FF10kylT7t7sHu1MX+DP9+853vrOXz1xBPv/5z9Mtt9yyr6sxjXvuuYdWr15N0WiUWltb6corr6TJycl9XS1FUZQDhmKxSOeddx6NjIzQ1772NfrJT35CnZ2d+7paz4teOxE9/fTTdNVVV+13X0zzPI++9KUv0dy5cykcDtNhhx1Gv/jFL/Z1tRRFUQ4qDsT8ff3119N5551Hc+bMed5r/4MZvTet/F/Dv68roCgvhmw2S34/D9977rmHrr76arrooosomUyKsuvXryfXfWm///Hggw/SmWeeSZOTk/S2t72NVq5cSUREDz30EP3rv/4r/fWvf6U//elPL2kd9oTvfe975Hle5e9MJkNXX301EdG0b23/y7/8C33sYx97Oau3X9LZ2UnZbJYCgcBLsv/e3l66+uqrqauriw4//HDxnt1fLyc33HADvec976E3vvGN9MEPfpDWrl1LV155JWUyGfroRz+6T+qkKIryQnnzm99MZ555JpXLZXrmmWfo+uuvp9///vd033330bvf/W465ZRTKmW7u7vp05/+NF166aV03HHHVV6fP3/+8x7n+uuvp3g8Ll47+uijaf78+ZTNZikYDO69kzpAeKnXX5///Ofp3HPPpbPPPlu8/va3v53e9KY3USgUesmOXY1HH32UTj75ZFq6dCl99atfpe3bt9N1111HGzZsoN///vcve30URVEORDZt2kQ9PT30ve99jy655JLK6/vz9emBcu10/PHHv6TrkqeffpquvvpqOuGEE6b9qmxf3g/55Cc/Sf/6r/9K73rXu+ioo46iW2+9ld7ylreQ4zj0pje9aZ/VS1EU5WDiQMzfX/ziF2liYoJWrVpFfX19+7o6VdF70/v3+ko58NCH4soBSTgc3uOyL/VN0VQqRW94wxvI5/PRunXraMmSJeL9a6+9lr73ve+9pHXYU15I8vT7/eKLBwcLxhjK5XIUiUT2qLzjOC9ovO1NXqrFzvORzWbpk5/8JL3mNa+hm266iYiI3vWud5HnefS5z32OLr30Uqqrq9sndVMURXkhvOIVr6C3ve1tlb9f9apX0VlnnUXXX3893XDDDXTMMcdU3nvooYfo05/+NB1zzDHiM3vCueeeS42NjTO+t69yyN4mnU5TLBbb4/L74qE0EZHP5yOfz7dPjv2JT3yC6urq6K677qJEIkFEu2Tk3/Wud9Gf/vQnOvXUU/dJvRRFUQ4kBgYGiIimfdl9f70+3ZfXTp7nUaFQ2OO1huu6+2xdsq++ILhjxw76yle+Qpdffjl961vfIiKiSy65hNasWUMf+chH6Lzzzttn6wZFUZSDiQMtfxMR/eUvf6n8Stz+kvtLid6bfn703rTyUqLy6cp+w0UXXUTxeJx27NhBZ599NsXjcWpqaqIPf/jDVC6XRVn0FL/qqqvoIx/5CBERzZ07tyJbOiXZZXtajoyM0Ic//GFavnw5xeNxSiQSdMYZZ9Bjjz32oup9ww030I4dO+irX/3qtAfiREQtLS30L//yL+K1b3/727Rs2TIKhULU3t5Ol19++TSJ9RNOOIEOPfRQevrpp+nEE0+kaDRKHR0d9KUvfWnaMbZv305nn302xWIxam5upg984AOUz+enlUMfkC1btlBTUxMREV199dWVdsN2tT1fSqUSfe5zn6P58+dTKBSirq4u+sQnPjHtWF1dXfTa176W7r77blq1ahWFw2GaN28e/ed//qcoVywW6eqrr6aFCxdSOBymhoYGWr16Nd12223TGxqYks7961//Su9+97upoaGBEokEXXDBBTQ6OjpjXf74xz/SkUceSZFIhG644QYiItq8eTOdd955VF9fT9FolF75ylfS//7v/4rPV/NtefbZZ+ncc8+l+vp6CofDdOSRR9JvfvObaXVNpVL0gQ98gLq6uigUCtGsWbPoggsuoKGhIbrrrrvoqKOOIiKiiy++uNIHU8eaybclnU7Thz70IZo9ezaFQiFavHgxXXfddWSMEeUcx6ErrriCbrnlFjr00EMpFArRsmXL6A9/+MNu25aI6M4776Th4WF673vfK16//PLLKZ1OT2sjRVGUA4WTTjqJiHb9KvzloJp35/3330+nn3461dbWUjQapTVr1tDf/va3592fMYauueYamjVrFkWjUTrxxBPpqaee2qO6TOWz6667jr72ta9RZ2cnRSIRWrNmDT355JOi7NSabNOmTXTmmWdSTU0NvfWtbyWiPc9DM3mKp1Ipev/731/57IIFC+iLX/zitG+ee55H//Zv/0bLly+ncDhMTU1NdPrpp1ck8x3HoXQ6TT/+8Y+nydxX8xTf22svm/HxcbrtttvobW97W+WBOBHRBRdcQPF4nH75y18+7z4URVH+r3PRRRfRmjVriIjovPPOI8dxKopm1TxJf/rTn9LKlSspEolQfX09velNb6Jt27Y977FSqRRddNFFVFtbS8lkki688MIXZXv2j147TZ3Xs88+S+effz4lEglqaGig973vfZTL5UTZqWu8n/3sZ5WcNnV9t27dOjrjjDMokUhQPB6nk08+me677z7x+X90XbJjxw565zvfSe3t7RQKhWju3Ll02WWXUaFQoB/96Ed03nnnERHRiSeeWMnPU8eayVN8YGCA3vnOd1JLSwuFw2FasWIF/fjHPxZlcP3y3e9+t3Iv4qijjqIHH3xwt21LRHTrrbdSsVgU/eM4Dl122WW0fft2uvfee593H4qiKMruORDzN9GuX2C/WL9zvTet96aVA5f982s6yv9ZyuUynXbaaXT00UfTddddR7fffjt95Stfofnz59Nll10242fOOecceu655+gXv/gFfe1rX6v8Umvqga/N5s2b6ZZbbqHzzjuP5s6dS/39/XTDDTfQmjVr6Omnn6b29vYXVOff/OY3FIlE6Nxzz92j8ldddRVdffXVdMopp9Bll11G69evp+uvv54efPBB+
          },
          "metadata": {}
        }
      ],
      "source": [
        "from jaxpm.plotting import plot_fields_single_projection\n",
        "\n",
        "fields = {\"Initial Conditions\" : initial_conditions , \"LPT Field\" : jnp.log10(cic_paint(jnp.zeros(mesh_shape) ,lpt_particles) + 1)}\n",
        "for i , field in enumerate(ode_particles[1:]):\n",
        "    fields[f\"field_{i}\"] = jnp.log10(cic_paint(jnp.zeros(mesh_shape) , field)+1)\n",
        "plot_fields_single_projection(fields)"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "c2cvuO2VZBHU"
      },
      "id": "c2cvuO2VZBHU",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Introduction.ipynb",
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}