{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "40ba1e08-69c0-494d-9f6f-879849813bf8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax.lax as lax\n",
    "from jax.experimental.maps import xmap\n",
    "from jax.experimental.maps import Mesh\n",
    "from jax.experimental.pjit import PartitionSpec, pjit\n",
    "from jaxpm.pm import cic_paint, cic_read, fftk\n",
    "from functools import partial\n",
    "import jax_cosmo as jc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "72822b71-95cf-452e-8e30-e3c2e77b8156",
   "metadata": {},
   "outputs": [],
   "source": [
    "nc=512\n",
    "boxsize=1024. # Mpx/h\n",
    "halo_size=16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "69c5e0ec-e82b-4ee1-94e1-4da546d26155",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cic_paint(mesh, positions):\n",
    "  \"\"\" Paints positions onto mesh\n",
    "  mesh: [nx, ny, nz]\n",
    "  positions: [npart, 3]\n",
    "  \"\"\"\n",
    "  positions = jnp.expand_dims(positions, 1)\n",
    "  floor = jnp.floor(positions)\n",
    "  connection = jnp.array([[[0, 0, 0], [1., 0, 0], [0., 1, 0], \n",
    "                           [0., 0, 1], [1., 1, 0], [1., 0, 1], \n",
    "                           [0., 1, 1], [1., 1, 1]]])\n",
    "\n",
    "  neighboor_coords = floor + connection\n",
    "  kernel = 1. - jnp.abs(positions - neighboor_coords)\n",
    "  kernel = kernel[..., 0] * kernel[..., 1] * kernel[..., 2]  \n",
    "\n",
    "  neighboor_coords = jnp.mod(neighboor_coords.reshape([-1,8,3]).astype('int32'), mesh.shape[-1])\n",
    "\n",
    "  dnums = jax.lax.ScatterDimensionNumbers(\n",
    "    update_window_dims=(),\n",
    "    inserted_window_dims=(0, 1, 2),\n",
    "    scatter_dims_to_operand_dims=(0, 1, 2))\n",
    "  mesh = lax.scatter_add(mesh, \n",
    "                         neighboor_coords, \n",
    "                         kernel.reshape([-1,8]),\n",
    "                         dnums)\n",
    "  return mesh\n",
    "\n",
    "def cic_read(mesh, positions):\n",
    "  \"\"\" Paints positions onto mesh\n",
    "  mesh: [nx, ny, nz]\n",
    "  positions: [npart, 3]\n",
    "  \"\"\"  \n",
    "  positions = jnp.expand_dims(positions, 1)\n",
    "  floor = jnp.floor(positions)\n",
    "  connection = jnp.array([[[0, 0, 0], [1., 0, 0], [0., 1, 0], \n",
    "                           [0., 0, 1], [1., 1, 0], [1., 0, 1], \n",
    "                           [0., 1, 1], [1., 1, 1]]])\n",
    "\n",
    "  neighboor_coords = floor + connection\n",
    "  kernel = 1. - jnp.abs(positions - neighboor_coords)\n",
    "  kernel = kernel[..., 0] * kernel[..., 1] * kernel[..., 2]  \n",
    "\n",
    "  #neighboor_coords = neighboor_coords.reshape([-1,8,3]).astype('int32')\n",
    "\n",
    "  neighboor_coords = jnp.mod(neighboor_coords.astype('int32'), mesh.shape[-1])\n",
    "\n",
    "  return (mesh[neighboor_coords[...,0], \n",
    "               neighboor_coords[...,1], \n",
    "               neighboor_coords[...,3]]*kernel).sum(axis=-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "56d06cdd-b1f3-434f-85e9-01ce7bb5e2e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Defining the main operations\n",
    "@partial(xmap,\n",
    "         in_axes={0:'x', 1:'y'},\n",
    "         out_axes=['x','y',...],\n",
    "         axis_sizes={'x':nc, 'y':nc},\n",
    "         axis_resources={'x': 'nx', 'y':'ny',\n",
    "                         'key_x':'nx', 'key_y':'ny'})\n",
    "def pnormal(key):\n",
    "    return jax.random.normal(key, shape=[nc])\n",
    "\n",
    "@partial(xmap,\n",
    "         in_axes={0:'x', 1:'y'},\n",
    "         out_axes=['x','y',...],\n",
    "         axis_resources={'x': 'nx', 'y': 'ny'})\n",
    "def pfft3d(mesh):\n",
    "    # [x, y, z]\n",
    "    mesh = jnp.fft.fft(mesh)\n",
    "    mesh = lax.all_to_all(mesh, 'x', 0, 0) # [z, y, x]\n",
    "    mesh = jnp.fft.fft(mesh)\n",
    "    mesh = lax.all_to_all(mesh, 'y', 0, 0) # [z, x, y]\n",
    "    return jnp.fft.fft(mesh)\n",
    "\n",
    "@partial(xmap,\n",
    "         in_axes={0:'x', 1:'y'},\n",
    "         out_axes=['x','y',...],\n",
    "         axis_resources={'x': 'nx', 'y': 'ny'})\n",
    "def pifft3d(mesh):\n",
    "    mesh = jnp.fft.ifft(mesh) \n",
    "    mesh = lax.all_to_all(mesh, 'y', 0, 0)\n",
    "    mesh = jnp.fft.ifft(mesh)\n",
    "    mesh = lax.all_to_all(mesh, 'x', 0, 0) \n",
    "    return jnp.fft.ifft(mesh).real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "b5378206-5dac-431d-9691-02366ed3470b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@partial(xmap,\n",
    "         in_axes=(['x','y',...],\n",
    "                  ['x'],\n",
    "                  ['y'],\n",
    "                  [...],[...],[...]),\n",
    "         out_axes=['x','y',...],\n",
    "         axis_resources={'x': 'nx', 'y': 'ny'})\n",
    "def cwise_fn(kfield, kx, ky, kz, k, pk):\n",
    "    kk = jnp.sqrt((kx / boxsize * nc)**2 + (ky / boxsize * nc)**2 +\n",
    "                 (kz / boxsize * nc)**2)\n",
    "    pkmesh = jc.scipy.interpolate.interp(kk, k, pk)\n",
    "    return kfield*(pkmesh*nc**3/boxsize**3)**0.5\n",
    "\n",
    "def get_initial_cond(cosmo, seed):\n",
    "    # Get real density field\n",
    "    linear = pnormal(jax.random.split(seed, nc*nc).reshape(nc,nc,-1))\n",
    "    lineark = pfft3d(linear)\n",
    "    \n",
    "    k = jnp.logspace(-4, 2, 256)\n",
    "    pk = jc.power.linear_matter_power(cosmo, k)\n",
    "    kvec = fftk([nc,nc,nc], symmetric=False)\n",
    "    \n",
    "    lineark = cwise_fn(lineark, kvec[0].squeeze(),kvec[1].squeeze(),kvec[2].squeeze(), k, pk)\n",
    "    \n",
    "    return pifft3d(lineark)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "24b1ee98-e3fa-471c-a44d-5aa6cbbaa49f",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = jax.random.PRNGKey(42)\n",
    "# keys = jax.random.split(key, 4).reshape([2,2,2])\n",
    "\n",
    "# We reshape all our devices to the mesh shape we want\n",
    "devices = np.array(jax.local_devices()).reshape((2, 2))\n",
    "\n",
    "cosmo = jc.Planck15()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f01f4d7b-39da-4479-8eb8-1ae9ab975ec4",
   "metadata": {},
   "source": [
    "Ok, cool, now let's implement LPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "5f50c23d-b269-493e-af6b-09388419ef4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "@partial(xmap,\n",
    "         in_axes=([...]),\n",
    "         out_axes={0:'sx', 2:'sy'},\n",
    "         axis_sizes={'sx':2, 'sy':2},\n",
    "         axis_resources={'sx': 'nx', 'sy': 'ny'})\n",
    "def pmeshgrid(x, y, z):\n",
    "    return jnp.stack(jnp.meshgrid(x,y,z),axis=-1)\n",
    "\n",
    "@partial(xmap,\n",
    "         in_axes=(['x','y','z'],\n",
    "                  ['x'], ['y'], ['z']),\n",
    "         out_axes=(['x','y','z'],\n",
    "                   ['x','y','z'],\n",
    "                   ['x','y','z']),\n",
    "         axis_resources={'x': 'nx', 'y': 'ny'})\n",
    "def papply_gradient_laplace(kfield, kx, ky, kz):\n",
    "    kk = (kx**2 + ky**2 + kz**2)\n",
    "    kernel = jnp.where(kk == 0, 1., 1./kk)\n",
    "    return (kfield * kernel * 1j * 1 / 6.0 * (8 * jnp.sin(ky) - jnp.sin(2 * ky)), \n",
    "            kfield * kernel * 1j * 1 / 6.0 * (8 * jnp.sin(kz) - jnp.sin(2 * kz)), \n",
    "            kfield * kernel * 1j * 1 / 6.0 * (8 * jnp.sin(kx) - jnp.sin(2 * kx))) \n",
    "\n",
    "preshape = pjit(lambda x: x.reshape([2, nc//2, 2, nc//2]+list(x.shape[2:])), \n",
    "                in_axis_resources=PartitionSpec('nx','ny'),\n",
    "                out_axis_resources=PartitionSpec('nx', None, 'ny', None))\n",
    "\n",
    "pireshape = pjit(lambda x: x.reshape([nc, nc]+list(x.shape[4:])), \n",
    "                in_axis_resources=PartitionSpec('nx', None, 'ny', None),\n",
    "                out_axis_resources=PartitionSpec('nx','ny'))\n",
    "\n",
    "pcic_read = xmap(lambda mesh, pos: cic_read(mesh, pos.reshape(-1,3)).reshape(pos.shape[:-1]),\n",
    "         in_axes=({0:'sx',2:'sy'},\n",
    "                  {0:'sx',2:'sy'}),\n",
    "         out_axes=({0:'sx',2:'sy'}),\n",
    "         axis_resources={'sx': 'nx', 'sy': 'ny'})\n",
    "\n",
    "pcic_paint = xmap(lambda mesh, pos, halo_size=halo_size: cic_paint(mesh, pos.reshape(-1,3)+jnp.array([halo_size,halo_size,0]).reshape([-1,3])),\n",
    "         in_axes=({0:'sx',2:'sy'},\n",
    "                  {0:'sx',2:'sy'}),\n",
    "         out_axes=({0:'sx',2:'sy'}),\n",
    "         axis_resources={'sx': 'nx', 'sy': 'ny'})\n",
    "\n",
    "@partial(xmap,\n",
    "         in_axes=({0:'sx',2:'sy'},[...]),\n",
    "         out_axes={0:'sx',2:'sy'},\n",
    "         axis_resources={'sx': 'nx', 'sy': 'ny'})\n",
    "def pad_mesh(mesh, halo_size=halo_size):\n",
    "    return jnp.pad(messh,[halo_size]*3)\n",
    "\n",
    "@partial(xmap,\n",
    "         in_axes=({0:'sx',2:'sy'}),\n",
    "         out_axes={0:'sx',2:'sy'},\n",
    "         axis_resources={'sx': 'nx', 'sy': 'ny'})\n",
    "def halo_reduce(mesh, halo_size=halo_size):\n",
    "    for axis_ind, axis_name in enumerate(['sx', 'sy']):        \n",
    "        # Split the array\n",
    "        left_margin, center, right_margin = mesh.split([2*halo_size, nc//2 ], axis_ind)\n",
    "        \n",
    "        # Perform halo exchange\n",
    "        left = lax.pshuffle(right_margin, perm=[1,0], axis_name=axis_name)\n",
    "        right =lax.pshuffle(left_margin, perm=[1,0], axis_name=axis_name)\n",
    "        if axis_ind==0:\n",
    "            mesh = mesh.at[:2*halo_size].add(left)\n",
    "            mesh = mesh.at[-2*halo_size:].add(right)\n",
    "        else:\n",
    "            mesh = mesh.at[:,:2*halo_size].add(left)\n",
    "            mesh = mesh.at[:,-2*halo_size:].add(right)\n",
    "            \n",
    "    # removing leftovers\n",
    "    return mesh[halo_size:-halo_size,halo_size:-halo_size]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "65edea80-9ea1-4300-842a-63ae81a5dddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "with Mesh(devices, ('nx', 'ny')):\n",
    "    \n",
    "    initial_conditions = get_initial_cond(cosmo, key)\n",
    "    \n",
    "    # Create the particles\n",
    "    pos = pmeshgrid(jnp.arange(nc//2), jnp.arange(nc//2), jnp.arange(nc))\n",
    "    \n",
    "    # Take the FFT of the field\n",
    "    lineark = pfft3d(initial_conditions)\n",
    "    \n",
    "    # Apply the laplace kernel\n",
    "    kvec = fftk([nc,nc,nc], symmetric=False)\n",
    "    kforces = papply_gradient_laplace(lineark,\n",
    "                                      kvec[0].squeeze(), \n",
    "                                      kvec[1].squeeze(), \n",
    "                                      kvec[2].squeeze())\n",
    "    \n",
    "    # Inverse Fourier Transform\n",
    "    forces_x = pcic_read(preshape(pifft3d(kforces[0])), pos)\n",
    "    forces_y = pcic_read(preshape(pifft3d(kforces[1])), pos)\n",
    "    forces_z = pcic_read(preshape(pifft3d(kforces[2])), pos)\n",
    "    \n",
    "    # Read the forces at particle positions\n",
    "    dx = xmap(lambda a,b,c: jnp.array([a,b,c]),\n",
    "              in_axes=(['sx','tx','sy','ty','z',...],\n",
    "                       ['sx','tx','sy','ty','z',...],\n",
    "                       ['sx','tx','sy','ty','z',...]),\n",
    "              out_axes=['sx','tx','sy','ty','z',...],\n",
    "              axis_resources={'sx': 'nx', 'sy': 'ny'}\n",
    "             )(forces_x, forces_y, forces_z)\n",
    "    \n",
    "    x_final = xmap(lambda a,b:a+b,\n",
    "                    in_axes=(['sx','tx','sy','ty','z',...],\n",
    "                       ['sx','tx','sy','ty','z',...]),\n",
    "                    out_axes=['sx','tx','sy','ty','z',...], \n",
    "                   axis_resources={'sx': 'nx', 'sy': 'ny'}\n",
    "                  )(dx,pos)\n",
    "    \n",
    "    # Painting final field\n",
    "    res = pcic_paint(jnp.zeros([2,256+halo_size*2,2,256+halo_size*2,512]), x_final)\n",
    "    res = halo_reduce(res)\n",
    "    res = pireshape(res).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "id": "acc4aaa7-9d45-4cbd-b560-86014de49922",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x15343c2cceb0>"
      ]
     },
     "execution_count": 160,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 657784 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "imshow(res.sum(axis=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "id": "7fba6185-ce16-471d-8043-c0d65d4cdaf4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x1535400f9700>"
      ]
     },
     "execution_count": 141,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 836648 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "figure(figsize=[10,10])\n",
    "imshow((0+initial_conditions).real.sum(axis=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b03d0661-e0ba-48d0-a225-1ae71ad5be15",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Jax",
   "language": "python",
   "name": "jax"
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
