{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d1dec3f4",
   "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": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%pylab inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8470dd57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "logger = logging.getLogger(\"root\")\n",
    "\n",
    "class CheckTypesFilter(logging.Filter):\n",
    "    def filter(self, record):\n",
    "        return \"check_types\" not in record.getMessage()\n",
    "logger.addFilter(CheckTypesFilter())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7b6d2891",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "import jax_cosmo as jc\n",
    "\n",
    "from jax.experimental.ode import odeint\n",
    "from jaxpm.painting import cic_paint, cic_paint_2d\n",
    "from jaxpm.pm import linear_field, lpt, make_ode_fn\n",
    "from jaxpm.lensing import density_plane, convergence_Born"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3f5dae21",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "# Below are a few parameters that generate a low resolution version of the k-TNG simulation\n",
    "\n",
    "box_size = [200.,200.,4000.]    # Transverse comoving size of the simulation volume\n",
    "nc = [64, 64, 320]              # Number of transverse voxels in the simulation volume\n",
    "lensplane_width = 102.5         # Width of each lensplane\n",
    "field_size = 5                  # Size of the lensing field in degrees\n",
    "field_npix = 128                # Number of pixels in the lensing field\n",
    "z_source = jnp.linspace(0,2)    # Source planes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ff4b61df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining the coordinate grid for lensing map\n",
    "import astropy.units as u\n",
    "xgrid, ygrid = np.meshgrid(np.linspace(0, field_size, field_npix, endpoint=False), # range of X coordinates\n",
    "                           np.linspace(0, field_size, field_npix, endpoint=False)) # range of Y coordinates\n",
    "\n",
    "coords = np.stack([xgrid, ygrid], axis=0)*u.deg\n",
    "c = coords.reshape([2, -1]).T.to(u.rad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7682e01c",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def make_map(Omega_c=0.2589, sigma8=0.8159):\n",
    "    # Instantiates a cosmology with desired parameters\n",
    "    cosmology = jc.Planck15(Omega_c=Omega_c, sigma8=sigma8)\n",
    "\n",
    "    # Planning out the scale factor stepping to extract desired lensplanes\n",
    "    n_lens = int(box_size[-1] // lensplane_width)\n",
    "    r = jnp.linspace(0., box_size[-1], n_lens+1)\n",
    "    r_center = 0.5*(r[1:] + r[:-1])\n",
    "\n",
    "    # Retrieve the scale factor corresponding to these distances\n",
    "    a = jc.background.a_of_chi(cosmology, r)\n",
    "    a_center = jc.background.a_of_chi(cosmology, r_center)\n",
    "\n",
    "    # Then one step per lens plane\n",
    "    stages = a_center[::-1]\n",
    "\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(cosmology, 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(nc, 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 nc]),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(nc), [particles+dx, p], \n",
    "                 jnp.concatenate([jnp.atleast_1d(0.1), stages]), cosmo, rtol=1e-5, atol=1e-5)\n",
    "    \n",
    "    # Extract the lensplanes\n",
    "    lensplanes = []\n",
    "    for i in range(len(a_center)):\n",
    "        dx = box_size[0]/64\n",
    "        dz = lensplane_width\n",
    "        plane = density_plane(res[0][::-1][i],\n",
    "                              nc,\n",
    "                              (i+0.5)*lensplane_width/box_size[-1]*nc[-1],\n",
    "                              width=lensplane_width/box_size[-1]*nc[-1],\n",
    "                              plane_resolution=64\n",
    "                           )\n",
    "        lensplanes.append({'r': r_center[i], \n",
    "                           'a': stages[::-1][i], \n",
    "                           'plane': plane,\n",
    "                           'dx':dx,\n",
    "                           'dz':dz})\n",
    "    \n",
    "    # Get convergence map\n",
    "    m = convergence_Born(cosmology, \n",
    "                          lensplanes,\n",
    "                          coords=jnp.array(c).T.reshape(2,field_npix,field_npix),\n",
    "                          z_source=z_source)\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9875b889",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = make_map()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "363aa1d1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "25.6 s ± 539 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit m = make_map()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "3db7341e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1287420 characters]",
      "text/plain": [
       "<Figure size 1440x1440 with 39 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# figure(figsize=[20,20])\n",
    "# for i in range(39):\n",
    "#     subplot(7,7,i+1)\n",
    "#     imshow(lensplanes[i][-1])\n",
    "#     axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "56ad6b0a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "39"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(lensplanes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a8f7057b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d990d051",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 445676 characters]",
      "text/plain": [
       "<Figure size 720x720 with 49 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "for i in range(49):\n",
    "    subplot(7,7,i+1)\n",
    "    imshow(m[i+1])\n",
    "    axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13fd34ce",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
