{
  "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": "[encoded data, 350028 characters]\n"
          },
          "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
}