JaxPM/notebooks/Introduction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c5f42bbe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
      "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": 2,
   "id": "281b4d3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_shape= [128, 128, 128]\n",
    "box_size  = [128.,128.,128.]\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": 3,
   "id": "826be667",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/francois/.local/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py:6655: UserWarning: Explicitly requested dtype <class 'jax._src.numpy.lax_numpy.int64'> requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.\n",
      "  lax._check_user_dtype_supported(dtype, \"astype\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "958 ms ± 15.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "res = run_simulation(0.25, 0.8)\n",
    "%timeit res = run_simulation(0.25, 0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4e012ce8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f1594106680>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXcAAAC6CAYAAABVwQ0gAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAEAAElEQVR4nOz9e5Bl93Xfh3728+yzz6v79GOmB4MWOC+MAZIgARIjihQJhZQLQ8kWcxPaEUXZiVxXvklkKpXceyXnZbpurktVTlKxoiTXqlgu2aR0bTmWZF8RsEzGlERaHlB8AgRBYGYINgbTM/08z3322c/7x/6tX/96zxlQMh+AXfhVTU336f387X3WWr/v+q7vssqy5PXx+nh9vD5eH/92DfvVvoDXx+vj9fH6eH1858frxv318fp4fbw+/i0crxv318fr4/Xx+vi3cLxu3F8fr4/Xx+vj38LxunF/fbw+Xh+vj38Lx+vG/fXx+nh9vD7+LRzfNeNuWdbjlmV93bKsq5Zl/fx36zyvj9fH93K8/l6/Pv5NGdZ3g+duWZYDPA/8MHAD+Bzw42VZPvsdP9nr4/XxPRqvv9evj3+Txncrcn8UuFqW5fWyLBPg/wv82HfpXK+P18f3arz+Xr8+/o0Z7nfpuPcALxm/3wAumRtYlvXTwE8DOE7jkTBcxbJsbNvFsmygxLJsHMfX+5RlQZ6nlGWBbTvYtodtO1iWjWU5FEWGZVlYlk1ZFhRFTlGkpGlEmie4tovntfC80DiHQ54nlGUh10WazqrJcRtYlqO2BSjVNtW5yjKnKDKKIgNQP+fEZY5TFmSWTdOyyYsMx3aNe7PUdTtASVkWlGWJZVnGPTm4bo7jDLHtlDzvkGUdfa2y4pJ5qOagAHIsqyTPXebzIXmeYtsOjtPAdQNcN8d1BzjOSM2pT56H5HmPovDvuK/quLa6LhcoyfOEJJkC4HlNHKeh7i2jKGz9HOQaq5+rZ1cU6bH7r55X9T9Y6rllajvU884BS12Ho+cxy+bk+ZykyMnUM/LLnGaRA5BZFjPbJXQaDIff3CvLcu1f52U2xrd8r9X96nfbhkc8IDP+7gINrOpny8IFHGAOZGVJYFm0sHCsaptRWTApS9qWhQMEWDTUNpn6+35ZkKljZ0Bb7ZupY7qWRQM4aVdf+VlZMC1LMkp93pZlEWDRsiwcLKaUjNT3Ii5LMiDnaKXvYBGo64/LkinHUQDX2E5Gz7LI1XXFZamPJ9u4ak7kml2gb9k0sPS9HZQFGSWucVwXaFgWDcshsWxKy8IrC2ZFxrws6an5mKrrjNW7Oa9ds1xHTomDRc+y8Pw2h36HaVnwpmiPnTLX95CpeUHNeUPNn9z7flkw5/hoAD3LxjX+ntWux5wzuZ5q32rOHcDFYrvMScvy+Ma1+f9Oj0UnOzaLZVn+MvDLAMvLZ8r3vOe/odns02h0AYjjAZ4X0mz2sW2XsixI0xlRtMtsdoDj+LTbJ2m3N+h0NgiCJeJ4QFFkeF5IUWTMZgcMBi+ys/MM8/kI32+ztvYAvd4mSTLBcXxarXWiaF8bk7IsiKI9bNul19skDFexbY88n5PnibrGMeOxxWh0gyjaI4r2SJIJ8/mILIv5wyymlyekWKS2w3tHL7PV6PBcsMS7m33a7ZOE4SplWTCfj0jTCIBWax3XDXAcn273NMvLfVZWfpOVld8kSU6xs/OT7O6eZzB4kSjaw7JslpfP0G6fxLZddawZvt8CYGvrMwwGLwIW3e49rK5eZHn5DOvrT7Ky8psEwXXyvMNg8D52dn+c6aRNHB8Sx0Nms32iaA/PC/H9Nt3uaXy/TVkWTKe7HB5ew/NClpbuo9VaJwxdXPeQ2ewEs9kBtu2SphGO4+N51fVMJtuMxzeZz8c4jqcdt++3AbSxTtOILItx3SaTyTa/P9witWze3eyztHQfvd4mS0v3cfv2V7h160t8eXKLyHZ52PG5UmQ8NN1hNYtJsXguWOINy2/giSc+8s3vxXtdf7cDyypbyvH0raOFct+yWbFsztmu/vxqkXFQFpyzXc7bLj/hhdhuwF4a8bE04rztHtv/vV5IZLv8Tjzg42nEvjLEACuWTd+yOVCf9dW53u8GALygzrVfFjyVJwA86visWDaX1TZPZDEvqMDloCz09uYxAX1dckzz77LNimWzXxZ8xG9z3na5kid8LI30trL9eTUfB2Wh9+1bNpccn8fdgNSy+Y1kov9unuuyG5B7IdteiFcWbKQRvxMPjl3jlTw5dk9X8kQfy3w+cs2X3YD39Db5eysXGNoef2X3q5ydj/it+Ygnsliff78s9JxfcnwuOT77ZcFVda6+ZfNCkR37/bzt6jmT63nB+Ls55Bovu4Gey6tFxj9StmPR+G4Z9xvAvcbvp4Gbd9vYsmzCcFVH1FVkXkWKEhWbURtURjjLYpJkTBw3cZwq6nQcn2bTpixDAMbjbRqNLs1mH8fxsW2XOB4wHt+k0eiS5wlpGpGrF9yybFz1clfGJtGRYp4nylgd4vub+H6bNI2wLJs0jfiDeIBXFlxyfBwvJM9TvpDP8Si5GA8I85S57dJun9THS9OIOB4CEIarANrAJ4nLaPQuwvCrdDpXOHHiV8jzn2U+7xOrl9bzQnWtMVG0x2x2gOsGhOEqvt8mCJaRVVCWxcxmB8xmF4iiByiKgCxbZjJ5hDQ5gecVZNmMsjzQcyKONUkmJMkU23Zw3QbLy2fIshiAPE8oijaOM8b3OySJr56hi+sG+H4L1x3hOJu4bsB0ukMcD5jPR+R5oo+TppFaORRYlo3vt2i3T/L22QFlWeB5IZ4X4roBtu3h+208ZeBS9dze4wbMLJsvJBMAtr2QN3wbL3Jt/IneaxnnbFd/Ic1hGr6+Zesvvoxly2bbdtl0fM7nyTFjBuCVBXtuwCXH50qesK/eYeDYuUzD9Qk11yvqdzHYK+rYK9adSG3f2FZ+N434C0Wm96tvWx/nbZcHvZC+ZWvDKte66Nz1Y6XKeJr7y3mXLZs0T0gtG68s6BnzcUX9XDem5lgx5umS4/OXG11Sy2a7yPih8U2+EK7yd1Yv8mODF7msnIQYXfPZ7SunsV8W2qmIM3/KuN/6+c1nZjo38/pkO5n/eoRvju+Wcf8ccN6yrDcALwP/AfChu21sWTZBsEyaTkmSCZbl0Gh01JfY1V921w3wvFDBAA5lWRDHA/I80YZoefkMYfh5LCvDth9kOl0jTac6Uq5WAJE6T7W0F6gBqiiy2ewzn4+Iol3KMqfZ7OO6gXY0ZenRaIzpdE6R5wmj0Q2yLObh6S4ATwdLPOyF2LaLl0akWKzmCZvAlfmIXhZTloVyIhZ5PqcoMrIsVlHwGp7XIkkmxHET1/1zdDpXWFv7ddJ0nfn85zRsEgRL5HlCHA+YTneZzfbxvBDb9mg0uvR6prP0yPOE2ewC4/H3M50+xHy+yXh8hiQ5YGVlh6I4w2RySx+/mnuHLIvJsjmNRocwXKXd7nJ4eEsb56IAxxnTaGwxnW5SlgWO4yvoqRqNRoxlndSrsKKoDDyA4/j6Gdi2i+P4+H6bMFzVz9f32zQaXSzLZjLZ1nP4sNMASr2954VcSGdAyQUsXLfxqrzXUC2vxRhJpLho9GvG8ZwyCr08YaLeu6tFph0FwJNZzH3BEptuwPks5mqRaUO9yLivqOjxoCxAHX+RQRdjbUbmLxTZMWNjGveDsuCgFvnWDZNcz35ZMHR8NlUQ9EKR8WjNOJrHNa89Vcdc9UL23IDzswO97Xnb5ZqsfNSxPmesTuqOc9GKQVZMAD/Z7PPb3dOERYarrjmyHH5r6T7eNDvgTV7I+9V8n7ddnshi/Yz3DQMu8ylzaK585PNztqvn0zTqi4z/VeM59C2bwPoeG/eyLDPLsn4G+GdUcOKvlGX51bttXxQZjUaHokjJshjbduh2T2vDHkV7FEWG77dptU4wn4+YzQ70/pZlM53u4LoBq6ubOM4E257RaGyxvHwR1w1oNvvKWA6ORZxBsKSvwbZdwnCNTmeD/f3ndYQ/n49pt08QBEuE4XPKcczI82UAZrMDrk93eFcyxityNrIZW8m
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "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": {},
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}