{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2206b58f",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "9f910706",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "pdims=(1 , 1)\n",
    "#for single gpu\n",
    "# pdims = (1 , 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "8bc804ab",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(128, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "folder = f'out/final_field/pmwd/1/128_128/{pdims[0]}x{pdims[1]}/lfm/halo_0'\n",
    "folder = f'out/final_field/jaxpm/1/128_128/{pdims[0]}x{pdims[1]}/LeapfrogMidpoint/halo_32'\n",
    "folder = f'out/final_field/jaxpm/1/128_128/{pdims[0]}x{pdims[1]}/lpt/halo_32'\n",
    "\n",
    "only_final_fields = True\n",
    "\n",
    "init_field_slices = []\n",
    "field_slices = []\n",
    "nb_solutions = 1\n",
    "nb_to_plot = 1\n",
    "final_slices = []\n",
    "\n",
    "for _ in range(nb_solutions):\n",
    "    final_slices.append([])\n",
    "\n",
    "for i in range(pdims[0]):\n",
    "    row_init_field = []\n",
    "    row_field = []\n",
    "    row_final_field = []\n",
    "    for _ in range(nb_solutions):\n",
    "        row_final_field.append([])\n",
    "    \n",
    "    for j in range(pdims[1]):\n",
    "        slice_index = i * pdims[1]  + j \n",
    "        if not only_final_fields:\n",
    "            row_field.append(np.load(f'{folder}/field_{slice_index}.npy'))\n",
    "            row_init_field.append(np.load(f'{folder}/initial_conditions_{slice_index}.npy'))\n",
    "\n",
    "        for sol_indx in range((nb_solutions - nb_to_plot) , nb_solutions):\n",
    "            row_final_field[sol_indx].append(np.load(f'{folder}/final_field_{sol_indx}_{slice_index}.npy'))\n",
    "    \n",
    "    if not only_final_fields:\n",
    "        field_slices.append(np.vstack(row_field))\n",
    "        init_field_slices.append(np.vstack(row_init_field))\n",
    "\n",
    "    for sol_indx in range((nb_solutions - nb_to_plot) , nb_solutions):\n",
    "        final_slices[sol_indx].append(np.vstack(row_final_field[sol_indx]))\n",
    "\n",
    "if not only_final_fields:\n",
    "    field = np.hstack(field_slices)\n",
    "    initial_conditions = np.hstack(init_field_slices)\n",
    "final_fields = []\n",
    "\n",
    "for sol_indx in range(nb_solutions - nb_to_plot , nb_solutions):\n",
    "    final_fields.append(np.hstack(final_slices[sol_indx]))\n",
    "\n",
    "if not only_final_fields:\n",
    "    print(field.shape)\n",
    "    box_size = field.shape\n",
    "else:\n",
    "    print(final_fields[-1].shape)\n",
    "    box_size = final_fields[-1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "bfa0c523",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 253564 characters]",
      "text/plain": [
       "<Figure size 1500x500 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "sum_over = box_size[0] // 8\n",
    "\n",
    "# Function to create subplots\n",
    "def plot_subplots(proj_axis, input , row, axes, title):\n",
    "    slicing = [slice(None)] * input.ndim\n",
    "    slicing[proj_axis] = slice(None, sum_over)\n",
    "    slicing = tuple(slicing)\n",
    "\n",
    "    # Plot initial conditions\n",
    "    axes[row, proj_axis].imshow(input[slicing].sum(axis=proj_axis), cmap='magma', extent=[0, box_size + 5, 0, box_size + 5])\n",
    "    axes[row, proj_axis].set_xlabel('Mpc/h')\n",
    "    axes[row, proj_axis].set_ylabel('Mpc/h')\n",
    "    axes[row, proj_axis].set_title(title)\n",
    "\n",
    "# Initialize figure and axes\n",
    "if only_final_fields:\n",
    "    nb_rows = len(final_fields)\n",
    "    field_start = 0\n",
    "else:\n",
    "    nb_rows = 2 + len(final_fields)\n",
    "    field_start = 2\n",
    "    \n",
    "nb_cols = 3\n",
    "fig, axes = plt.subplots(nb_rows, nb_cols, figsize=(15, 5 * nb_rows))\n",
    "\n",
    "# Plot initial conditions and LPT field for each projection\n",
    "if not only_final_fields:\n",
    "    for proj_axis in range(3):\n",
    "        plot_subplots(proj_axis,initial_conditions ,  0, axes, f'Initial conditions projection {proj_axis}')\n",
    "        plot_subplots(proj_axis, field ,  1, axes, f'LPT density field projection {proj_axis}')\n",
    "\n",
    "if len(final_fields) == 1:  # Check if axes is 1-dimensional\n",
    "    axes = np.expand_dims(axes,axis=0)\n",
    "# Plot final fields for each projection\n",
    "for indx, final_field in enumerate(final_fields):\n",
    "    for proj_axis in range(3):\n",
    "        slicing = [slice(None)] * final_field.ndim\n",
    "        slicing[proj_axis] = slice(None, sum_over)\n",
    "        slicing = tuple(slicing)\n",
    "        axes[indx + field_start, proj_axis].imshow(final_fields[indx][slicing].sum(axis=proj_axis) + 1, cmap='magma', extent=[0, box_size[0] + 5, 0, box_size[0] + 5])\n",
    "        axes[indx + field_start, proj_axis].set_xlabel('Mpc/h')\n",
    "        axes[indx + field_start, proj_axis].set_ylabel('Mpc/h')\n",
    "        axes[indx + field_start, proj_axis].set_title(f'ODE Step {indx} projection {proj_axis}')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ea64b6d",
   "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"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
