Inno4Scale.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "45734234-7021-48e2-b3e0-3fab28ec47d7",
   "metadata": {},
   "source": [
    "# Vlasiator - A Global Hybrid-Vlasov Simulation Model \n",
    "Vlasiator [@palmroth2018] is an open-source simulation software used to model the behavior of plasma in the Earth's magnetosphere, a region of space where the solar wind interacts with the Earth’s magnetic field. Vlasiator models collisionless space plasma dynamics by solving the 6-dimensional Vlasov equation, using a hybrid-Vlasov approach. It uses a 3D Cartesian grid in real space, with each cell storing another 3D Cartesian grid in velocity space. The velocity mesh contained in each spatial cell in the simulation domain has been represented so far by a sparse grid approach, fundamentally based on an associative container such as a key-value hashtable. Storing a 3D VDF at every spatial cells increases the memory requirements exponentially both during runtime and for storing purposes. Our proposal revolves around developing innovative solutions to compressing the VDFs during runtmime."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "796e2f0c-f6d8-4600-9e57-80a57358bc7d",
   "metadata": {},
   "source": [
    "![title](images/egi.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ac9891a-860b-4411-8326-1dd98e4b5717",
   "metadata": {},
   "source": [
    "# VDF Compression"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ffaa767-6fdc-42f4-a9c1-93009df7cfeb",
   "metadata": {},
   "source": [
    "## Initialization\n",
    "Let's read in a vdf from a sample file and see what that looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aae7d622-c42c-470b-8d9d-509efb1719e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys,os\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "# sys.path.append('/home/mjalho/analysator')\n",
    "import tools as project_tools\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "mpl.rcParams.update(mpl.rcParamsDefault)\n",
    "import matplotlib.colors as colors\n",
    "# plt.rcParams['figure.figsize'] = [7, 7]\n",
    "import ctypes\n",
    "import pyzfp,zlib\n",
    "import mlp_compress\n",
    "from skimage import measure\n",
    "from mpl_toolkits.mplot3d.art3d import Poly3DCollection\n",
    "\n",
    "import pytools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be5f3337-4d10-4115-9927-2d456b737d2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "file=\"assets/bulk1.0001280.vlsv\";cid=356780649;\n",
    "#Read the VDF into a 3D uniform mesh and plot it\n",
    "vdf=project_tools.extract_vdf(file,cid,25)\n",
    "# np.save(\"sample_vdf.bin\",np.array(vdf,dtype=np.double));\n",
    "np.array(vdf,dtype=np.double).tofile(\"sample_vdf.bin\")\n",
    "nx,ny,nz=np.shape(vdf)\n",
    "print(f\"VDF shape = {np.shape(vdf)}\")\n",
    "fig = plt.figure(figsize=[7,4], dpi=300);\n",
    "ax1 = plt.subplot(221)\n",
    "ax2 = plt.subplot(222)\n",
    "ax3 = plt.subplot(223)\n",
    "ax4 = plt.subplot(224, projection='3d')\n",
    "# , ax2, ax3) = plt.subplots(1, 3)\n",
    "\n",
    "\n",
    "cmap = 'viridis'\n",
    "cmap = 'hawaii_r'\n",
    "norm = colors.LogNorm(vmin=1e-15,vmax=3e-13)\n",
    "im1 = ax1.pcolormesh(vdf[:,:,nz//2],norm=norm,rasterized=False,cmap=cmap)\n",
    "# im1 = ax1.imshow(vdf[:,:,nz//2],norm=norm,cmap=cmap, origin='lower')\n",
    "im2 = ax2.pcolormesh(vdf[nx//2,:,:].T,norm=norm,rasterized=False,cmap=cmap)\n",
    "im3 = ax3.pcolormesh(vdf[:,ny//2,:],norm=norm,rasterized=False,cmap=cmap)\n",
    "ax1.axis('equal')\n",
    "ax1.set_xlabel('vx')\n",
    "ax1.set_ylabel('vy')\n",
    "ax2.axis('equal')\n",
    "ax2.set_xlabel('vz')\n",
    "ax2.set_ylabel('vy')\n",
    "ax3.axis('equal')\n",
    "ax3.set_xlabel('vx')\n",
    "ax3.set_ylabel('vz')\n",
    "\n",
    "level = 1e-14\n",
    "for level in [3e-15,1e-14,3e-14,1e-13,2e-13]:\n",
    "    verts, faces, normals, values = measure.marching_cubes(vdf, level)\n",
    "    mesh = Poly3DCollection(verts[faces], shade=True, facecolors=mpl.colormaps[cmap](norm(level)), alpha = norm(level)**2)\n",
    "    ax4.add_collection3d(mesh)\n",
    "\n",
    "ax4.set_xlim(nx/4,3*nx/4)\n",
    "ax4.set_ylim(ny/4,3*ny/4)\n",
    "ax4.set_zlim(nz/4,3*nz/4)\n",
    "ax4.set_xlabel('vx')\n",
    "ax4.set_ylabel('vy')\n",
    "ax4.set_zlabel('vz')\n",
    "\n",
    "# cax = fig.add_axes([0.7,0.1,0.2,0.05])\n",
    "# fig.colorbar(im1, cax=cax, location='bottom')\n",
    "fig.suptitle(\"Original VDF\")\n",
    "plt.show()\n",
    "\n",
    "project_tools.plot_vdf_discrete_laplacians(vdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f75b96f-3ba2-49d0-b3bf-c30966de5abe",
   "metadata": {},
   "source": [
    "The vdf shown above is sampled on a uniform 3D velocity mesh and contains 64bit floating point numbers that represent the phase space density. We can calculate the total size of this VDF is bytes using ```sys.getsizeof(vdf)```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e456c54-ab54-4dd4-8112-ebe75da96c96",
   "metadata": {},
   "outputs": [],
   "source": [
    "vdf_mem=sys.getsizeof(vdf)\n",
    "num_stored_elements=len(vdf[vdf>1e-15])\n",
    "print(f\"VDF takes {vdf_mem} B.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74ae8266-9333-445b-9a07-2702795fa0fd",
   "metadata": {},
   "source": [
    "Now in Vlasiator we have countlesss VDFs since there is one per spatial cell. It would be great if we could compress them efficiently.\n",
    "\n",
    "## Compression algorithms\n",
    "\n",
    "### Zlib: lossless floating point compression\n",
    "We can try to do so by using zlib which is a form of lossless compression. The reconstruction is accurate, by definition, but the compression ratio is small."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed5662d9-8f88-43f7-b1f6-563d8fde0f81",
   "metadata": {},
   "outputs": [],
   "source": [
    "compressed_vdf = zlib.compress(vdf)\n",
    "compressed_vdf_mem=len(compressed_vdf)\n",
    "compression_ratio=vdf_mem/compressed_vdf_mem\n",
    "print(f\"Achieved compression ratio using zlib= {round(compression_ratio,2)}.\")\n",
    "decompressed_vdf = zlib.decompress(compressed_vdf)\n",
    "recon = np.frombuffer(decompressed_vdf, dtype=vdf.dtype).reshape(vdf.shape)\n",
    "project_tools.plot_vdfs(vdf,recon)\n",
    "project_tools.print_comparison_stats(vdf,recon)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2e4811c-6d2b-482f-bdeb-6c60e6453e1e",
   "metadata": {},
   "source": [
    "### zfp: lossy floating point compression\n",
    "We can use a lossy compression method like zfp[@zfp] to get even higher compression ratios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4c6ba0b-b56a-4ef5-8c97-7b289762e570",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Compresses a VDF  using ZFP (Zstandard Compressed FP)\n",
    "Input:VDF - numpy array \n",
    "Output: recon (Reconstructed VDF) - numpy array \n",
    "\"\"\"\n",
    "tolerance = 1e-13\n",
    "compressed_vdf = pyzfp.compress(vdf, tolerance=tolerance)\n",
    "compressed_vdf_mem=len(compressed_vdf)\n",
    "compression_ratio=vdf_mem/compressed_vdf_mem\n",
    "print(f\"Achieved compression ratio using zfp= {round(compression_ratio,2)}.\")\n",
    "recon = pyzfp.decompress(compressed_vdf,vdf.shape,vdf.dtype,tolerance)\n",
    "project_tools.plot_vdfs(vdf,recon)\n",
    "project_tools.print_comparison_stats(vdf,recon)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e59226f-a344-47f7-aa2b-d016b9e31ce3",
   "metadata": {},
   "source": [
    "### Multilevel Perceptron (MLP)\n",
    "\n",
    "This is based on [@park2019]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "169f5b61-cda9-48d0-a550-b103f6355e5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Compresses a VDF using an MLP (Multilayer Perceptron).\n",
    " Input: \"sample_vdf.bin\" - Binary file containing the VDF data\n",
    "        order - Order of the fourier features\n",
    "        epochs - Number of training epochs for the MLP model\n",
    "        n_layers - Number of layers in the MLP model\n",
    "        n_neurons - Number of neurons in each layer of the MLP model\n",
    " Output: recon (Reconstructed VDF) - NumPy array representing the reconstructed volume data\n",
    "\"\"\"\n",
    "order=0\n",
    "epochs=20\n",
    "n_layers=4\n",
    "n_neurons=25\n",
    "nx,ny,nz=np.shape(vdf)\n",
    "recon=mlp_compress.compress_mlp(\"sample_vdf.bin\",order,epochs,n_layers,n_neurons, nx)\n",
    "recon=np.array(recon,dtype=np.double)\n",
    "recon= np.reshape(recon,np.shape(vdf),order='C')\n",
    "project_tools.plot_vdfs(vdf,recon)\n",
    "project_tools.print_comparison_stats(vdf,recon)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c753cd6c-eb1c-4f88-bb5b-2781a3e4d351",
   "metadata": {},
   "source": [
    "### MLP with Fourier features \n",
    "\n",
    "We will compress the VDF using an MLP with Fourier Features, which significantly improves MLP performance [@2020fourier]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afc424ea-8b00-4dd3-bb3c-af0a20daeba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Compresses a VDF using an MLP (Multilayer Perceptron).\n",
    " Input: \"sample_vdf.bin\" - Binary file containing the VDF data\n",
    "        order - Order of the fourier features\n",
    "        epochs - Number of training epochs for the MLP model\n",
    "        n_layers - Number of layers in the MLP model\n",
    "        n_neurons - Number of neurons in each layer of the MLP model\n",
    " Output: recon (Reconstructed VDF) - NumPy array representing the reconstructed volume data\n",
    "\"\"\"\n",
    "order=16\n",
    "epochs=20\n",
    "n_layers=4\n",
    "n_neurons=25\n",
    "nx,ny,nz=np.shape(vdf)\n",
    "recon=mlp_compress.compress_mlp(\"sample_vdf.bin\",order,epochs,n_layers,n_neurons,nx)\n",
    "recon=np.array(recon,dtype=np.double)\n",
    "recon= np.reshape(recon,np.shape(vdf),order='C')\n",
    "project_tools.plot_vdfs(vdf,recon)\n",
    "project_tools.print_comparison_stats(vdf,recon)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e8837a4-50cd-422d-8c70-6df8771f514c",
   "metadata": {},
   "source": [
    "### Spherical Harmonic Decomposition\n",
    "\n",
    "Spherical harmonics have been suggested as a usable approximation for VDFs in our subject domain [@vinas_gurgiolo_2009]. Here we investigate if the method can be employed as a compression method. This prototype should be improved by using e.g., the Misner [@misner_2004] method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3815300e-e181-4705-8e21-58e7d5b533a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Compresses a VDF using a spherical harmonic decompostion\n",
    "Input: \"sample_vdf.bin\" - Binary file containing the VDF data\n",
    "       degree - Degree of the spherical harmonic decomposition (l)\n",
    "Output: recon (Reconstructed VDF) - NumPy array representing the reconstructed volume data\n",
    "\"\"\"\n",
    "\n",
    "degree=10\n",
    "nx,ny,nz=np.shape(vdf)\n",
    "recon=mlp_compress.compress_sph(\"sample_vdf.bin\",degree,nx)\n",
    "recon=np.array(recon,dtype=np.double)\n",
    "recon= np.reshape(recon,np.shape(vdf),order='C')\n",
    "project_tools.plot_vdfs(vdf,recon)\n",
    "project_tools.print_comparison_stats(vdf,recon)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "136ceeae-4283-42f0-84e5-c60017b17e80",
   "metadata": {},
   "source": [
    "### Convolutional Neural Network (CNN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad5fb16b-01ba-450b-81cc-0f9a65717323",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Function: train_and_reconstruct\n",
    "\n",
    "Description:\n",
    "This function takes an input array and trains a Convolutional Neural Network (CNN) model to reconstruct the input array.\n",
    "It uses Mean Squared Error (MSE) loss and the Adam optimizer for training.\n",
    "\n",
    "Inputs:\n",
    "- input_array (numpy array): The input array to be reconstructed.\n",
    "- num_epochs (int, optional): The number of training epochs.\n",
    "- learning_rate (float, optional): The learning rate for the Adam optimize.\n",
    "Outputs:\n",
    "    Reconstructed vdf array \n",
    "    Size of model in bytes\n",
    "\"\"\"\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CNN, self).__init__()\n",
    "        self.conv1 = nn.Conv3d(1, 16, kernel_size=3, padding=1)\n",
    "        self.conv2 = nn.Conv3d(16, 32, kernel_size=3, padding=1)\n",
    "        self.conv3 = nn.Conv3d(32, 64, kernel_size=3, padding=1)\n",
    "        self.conv4 = nn.Conv3d(64, 1, kernel_size=3, padding=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.conv1(x))\n",
    "        x = self.relu(self.conv2(x))\n",
    "        x = self.relu(self.conv3(x))\n",
    "        x = self.conv4(x)\n",
    "        return x\n",
    "\n",
    "def train_and_reconstruct(input_array, num_epochs=30, learning_rate=0.001):\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    input_tensor = torch.tensor(input_array, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)  # Add batch and channel dimensions, move to device\n",
    "    model = CNN().to(device)  # Move model to device\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    for epoch in range(num_epochs):\n",
    "        optimizer.zero_grad()\n",
    "        output_tensor = model(input_tensor)\n",
    "        loss = criterion(output_tensor, input_tensor)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if (epoch+1) % 100== 0:\n",
    "            print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        output_tensor = model(input_tensor)\n",
    "    reconstructed_array = output_tensor.squeeze(0).squeeze(0).cpu().numpy()\n",
    "    \n",
    "    param_size = 0\n",
    "    for param in model.parameters():\n",
    "        param_size += param.nelement() * param.element_size()\n",
    "    buffer_size = 0\n",
    "    for buffer in model.buffers():\n",
    "        buffer_size += buffer.nelement() * buffer.element_size()\n",
    "    size = (param_size + buffer_size)   \n",
    "    return reconstructed_array, size\n",
    "\n",
    "\n",
    "vdf_temp=vdf.copy()\n",
    "vdf_temp[vdf_temp<1e-16]=1e-16\n",
    "vdf_temp = np.log10(vdf_temp)\n",
    "input_array=vdf_temp\n",
    "recon,total_size= train_and_reconstruct(input_array,100)\n",
    "recon = 10 ** recon\n",
    "recon[recon <= 1e-16] = 0\n",
    "vdf_size=nx*ny*nz*8\n",
    "print(f\"Compresion achieved using a CNN = {round(vdf_size/total_size,2)}\")\n",
    "project_tools.plot_vdfs(vdf,recon)\n",
    "project_tools.print_comparison_stats(vdf,recon)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "805d6d35-924d-4f20-b337-6efc0af0746a",
   "metadata": {},
   "source": [
    "#### CNN with minibatches\n",
    "\n",
    "\n",
    "Here we still use a CNN but this time we use minibatch training and batch normalization layers to try and improve performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3dcc16f-ddbd-4e2d-afcf-2a4172ce37e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Function: train_and_reconstruct\n",
    "\n",
    "Description:\n",
    "This function takes an input array and trains a Convolutional Neural Network (CNN) model to reconstruct the input array.\n",
    "It uses Mean Squared Error (MSE) loss and the Adam optimizer for training.\n",
    "\n",
    "Inputs:\n",
    "- input_array (numpy array): The input array to be reconstructed.\n",
    "- num_epochs (int, optional): The number of training epochs.\n",
    "- learning_rate (float, optional): The learning rate for the Adam optimizer\n",
    "Outputs:\n",
    "    Reconstructed vdf array \n",
    "    Size of model in bytes\n",
    "\"\"\"\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(CNN, self).__init__()\n",
    "        self.conv1 = nn.Conv3d(1, 16, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm3d(16)\n",
    "        self.conv2 = nn.Conv3d(16, 32, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm3d(32)\n",
    "        self.conv3 = nn.Conv3d(32, 64, kernel_size=3, padding=1)\n",
    "        self.bn3 = nn.BatchNorm3d(64)\n",
    "        self.conv4 = nn.Conv3d(64, 1, kernel_size=3, padding=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.bn1(self.conv1(x)))\n",
    "        x = self.relu(self.bn2(self.conv2(x)))\n",
    "        x = self.relu(self.bn3(self.conv3(x)))\n",
    "        x = self.conv4(x)\n",
    "        return x\n",
    "\n",
    "def train_and_reconstruct(input_array, num_epochs=30, learning_rate=0.001, batch_size=32):\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    input_tensor = torch.tensor(input_array, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)  # Move input tensor to device\n",
    "    model = CNN().to(device) \n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    \n",
    "    for epoch in range(num_epochs):\n",
    "        for i in range(0, input_tensor.size(0), batch_size):\n",
    "            optimizer.zero_grad()\n",
    "            batch_input = input_tensor[i:i+batch_size]\n",
    "            output_tensor = model(batch_input)\n",
    "            loss = criterion(output_tensor, batch_input)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        output_tensor = model(input_tensor)\n",
    "    reconstructed_array = output_tensor.squeeze(0).squeeze(0).cpu().numpy()\n",
    "    \n",
    "    param_size = 0\n",
    "    for param in model.parameters():\n",
    "        param_size += param.nelement() * param.element_size()\n",
    "    buffer_size = 0\n",
    "    for buffer in model.buffers():\n",
    "        buffer_size += buffer.nelement() * buffer.element_size()\n",
    "    size = (param_size + buffer_size)   \n",
    "    return reconstructed_array, size\n",
    "\n",
    "vdf_temp = vdf.copy()\n",
    "vdf_temp[vdf_temp < 1e-16] = 1e-16\n",
    "vdf_temp = np.log10(vdf_temp)\n",
    "input_array = vdf_temp\n",
    "recon, total_size = train_and_reconstruct(input_array, 100)\n",
    "\n",
    "recon = 10 ** recon\n",
    "recon[recon <= 1e-16] = 0\n",
    "vdf_size = nx * ny * nz * 8\n",
    "print(f\"Compression achieved using a CNN = {round(vdf_size / total_size, 2)}\")\n",
    "\n",
    "project_tools.plot_vdfs(vdf, recon)\n",
    "project_tools.print_comparison_stats(vdf, recon)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a69fd4a-b5d2-4561-b6f7-b0475729714b",
   "metadata": {},
   "source": [
    "### Hermite Decomposition "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e681c831-3cd0-425f-8ddb-ad5a0057fc32",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Loads the original 3D VDF and fits it to a Maxwellian distribution.\n",
    "  Input: vdf - numpy array representing the original 3D VDF\n",
    "  Output: vdf_herm_3d Reconstructed VDF using Hermite Decomposition\n",
    "\"\"\"\n",
    "### load original 3d vdf and fit Maxwellian\n",
    "vdf_3d=vdf.copy()\n",
    "print('loading done')\n",
    "vdf_size=nx*ny*nz*8\n",
    "\n",
    "#### Fit Maxwellian\n",
    "v_min,v_max,n_bins=0,nx,nx ### define limits and size of velocity axes\n",
    "\n",
    "amp,ux,uy,uz,vthx,vthy,vthz=1e-14,nx,nx,nx,10,10,10 ### initial guess for scipy curve fit\n",
    "guess=amp,ux,uy,uz,vthx,vthy,vthz ### initial guess for scipy curve fit\n",
    "\n",
    "max_fit_3d,params=project_tools.max_fit(vdf_3d,v_min,v_max,n_bins,guess) ### fitting\n",
    "print('Maxwell fit done')\n",
    "\n",
    "\n",
    "#### forward transform ####\n",
    "mm=15 ### PUT THE NUMBER OF HARMONICS\n",
    "norm_amp,u,vth=params[0],params[1:4],params[4:7] ### get the maxwellin fit parameters of thermal and bulk velocity\n",
    "\n",
    "vdf_3d_norm=vdf_3d/norm_amp ### normalize data\n",
    "vdf_3d_flat= vdf_3d_norm.flatten() ### flatten data\n",
    "\n",
    "v_xyz=project_tools.get_flat_mesh(v_min,v_max,n_bins) ### flattening the mesh nodes coordinates\n",
    "herm_array=np.array(project_tools.herm_mpl_arr(m_pol=mm,v_ax=v_xyz,u=params[1:4],vth=params[4:7])) ### create array of hermite polynomials\n",
    "\n",
    "hermite_matrix=project_tools.coefficient_matrix(vdf_3d_flat,mm,herm_array,v_xyz) ### calculate the coefficients of the Hermite transform\n",
    "print('Forward transform done')\n",
    "total_size =5*8+8*np.prod(np.shape(hermite_matrix))\n",
    "\n",
    "#### inverse transform ####\n",
    "inv_herm_flat=project_tools.inv_herm_trans(hermite_matrix, herm_array, v_xyz) ### inverse Hermite transform\n",
    "vdf_herm_3d = (np.reshape(inv_herm_flat,(n_bins,n_bins,n_bins)))*norm_amp ### reshaping back to 3d array and renormalization\n",
    "print('Inverse transform done')\n",
    "print(f\"Compresion achieved using Hermite = {round(vdf_size/total_size,2)}\")\n",
    "project_tools.plot_vdfs(vdf,vdf_herm_3d)\n",
    "project_tools.print_comparison_stats(vdf,vdf_herm_3d)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66f07c18-105f-4f3e-8489-c90bf18cc962",
   "metadata": {},
   "source": [
    "### Gaussian Mixture Model (GMM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2728768-2967-4669-ba52-b201c7284027",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    " Loads the original 3D VDF and performs Gaussian Mixture Model (GMM) decomposition.\n",
    " Input: vdf - NumPy array representing the original 3D VDF\n",
    " Output: vdf_rec Reconstructed VDF using GMM\n",
    "\"\"\"\n",
    "#### load original 3d vdf\n",
    "vdf_3d=vdf.copy()\n",
    "\n",
    "### define number of populations and normalization parameter\n",
    "n_pop=15\n",
    "norm_range=300\n",
    "\n",
    "### RUN GMM\n",
    "means,weights,covs,norm_unit=project_tools.run_gmm(vdf_3d,n_pop,norm_range)\n",
    "### reconstruction resolution and limits of v_space axes\n",
    "n_bins=nx\n",
    "v_min,v_max=0,nx\n",
    "\n",
    "### reconstruction of the vdf \n",
    "vdf_rec=project_tools.reconstruct_vdf(n_pop,means,covs,weights,n_bins,v_min,v_max)\n",
    "vdf_rec=vdf_rec*norm_unit*norm_range\n",
    "total_size =5*8+8*np.prod(np.shape(np.array(covs)))+8*np.prod(np.shape(np.array(weights)))+8*np.prod(np.shape(np.array(means)))\n",
    "print(f\"Compresion achieved using GMM = {round(vdf_size/total_size,2)}\")\n",
    "\n",
    "project_tools.plot_vdfs(vdf,vdf_rec)\n",
    "project_tools.print_comparison_stats(vdf,vdf_rec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f133b29-a0d2-4d1e-9383-247690083c6d",
   "metadata": {},
   "source": [
    "### Octree based polynomial approximation\n",
    "\n",
    "- Modify `maxiter` parameter or `tol` parameter to get different levels of accuracy.\n",
    "- This is pretty slow since its just a single thread cpu implementation with very conservative refinement policy. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07c85a05-29b7-4de6-8be4-2b2bcb0b5bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from juliacall import Main as jl\n",
    "jl.Pkg.activate(\"src/jl_env\")\n",
    "jl.Pkg.instantiate()\n",
    "jl.include(\"src/octree.jl\")\n",
    "vdf_3d = vdf.copy()\n",
    "A, b, img, reco, cell, tree = jl.VDFOctreeApprox.compress(vdf_3d, maxiter=500, alpha=0.0, beta=1.0, nu=2, tol=3e-1, verbose=False)\n",
    "project_tools.plot_vdfs(vdf, reco)\n",
    "project_tools.print_comparison_stats(vdf,reco)\n",
    "vdf_size = nx * ny * nz * 8\n",
    "print(f\"Number of leaves: {len(tree)} * basis size: {len(b)} = {len(tree)*len(b)}\")\n",
    "print(f\"Compression achieved with Octree = {round((len(b)*len(tree)*3)/vdf_size,3)}. Assuming 8 octals per leaf for geometry representation.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "908dd835",
   "metadata": {},
   "source": [
    "### Discrete Cosine Transform (DCT)\n",
    "\n",
    "DCT is simple to implement and is based on linear operations. This is widely used as the JPEG compression standard for images, but it can be extended to 3D distributions as well. Blockyness is apparent, and there is little room for fine-tuning: block size and number of retained DCT components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f22279-7511-4ff9-90bd-44db73833dc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.fft import dctn, idctn\n",
    "\n",
    "vdf_3d = vdf.copy()\n",
    "orig_shape = vdf_3d.shape\n",
    "vdf_3d[np.isnan(vdf_3d)] = 0\n",
    "\n",
    "blocksize = 8\n",
    "paddings = (np.ceil(np.array(vdf_3d.shape)/8)).astype(int)*8 - vdf_3d.shape\n",
    "paddings = ((0,paddings[0]),(0,paddings[1]),(0,paddings[2]))\n",
    "vdf_3d = np.pad(vdf_3d, paddings)\n",
    "\n",
    "# dct_data = dctn(vdf_3d)\n",
    "# print(dct_data.shape)\n",
    "# vdf_rec = idctn(dct_data)\n",
    "# print(vdf_rec.shape)\n",
    "\n",
    "block_data = np.zeros_like(vdf_3d)\n",
    "for i in range(0,vdf_3d.shape[0], blocksize):\n",
    "    for j in range(0, vdf_3d.shape[1], blocksize):\n",
    "        for k in range(0, vdf_3d.shape[2], blocksize):\n",
    "            block_data[i:i+blocksize,j:j+blocksize, k:k+blocksize] = dctn(vdf_3d[i:i+blocksize,j:j+blocksize, k:k+blocksize])\n",
    "\n",
    "keep_n = 4\n",
    "zeroed = np.zeros_like(block_data)\n",
    "for i in range(keep_n):\n",
    "    for j in range(keep_n):\n",
    "        for k in range(keep_n):\n",
    "            zeroed[i::blocksize,j::blocksize,k::blocksize] = block_data[i::blocksize,j::blocksize,k::blocksize]\n",
    "\n",
    "# fig, ax = plt.subplots()\n",
    "# ax.pcolor(np.log(np.abs(zeroed[:,0,:])))\n",
    "\n",
    "\n",
    "volume_compressed = np.prod(keep_n*np.ceil(np.array(vdf_3d.shape)/8))\n",
    "volume_orig = np.prod(vdf_3d.shape)\n",
    "compression = volume_orig/volume_compressed\n",
    "\n",
    "vdf_rec = np.zeros_like(vdf_3d)\n",
    "for i in range(0,vdf_3d.shape[0], blocksize):\n",
    "    for j in range(0, vdf_3d.shape[1], blocksize):\n",
    "        for k in range(0, vdf_3d.shape[2], blocksize):\n",
    "            vdf_rec[i:i+blocksize,j:j+blocksize, k:k+blocksize] = idctn(zeroed[i:i+blocksize,j:j+blocksize, k:k+blocksize])\n",
    "\n",
    "print(\"compression:\", compression)\n",
    "# vdf_rec = idctn(dct_data)\n",
    "vdf_rec = vdf_rec[0:orig_shape[0],0:orig_shape[1],0:orig_shape[2]]\n",
    "\n",
    "project_tools.plot_vdfs(vdf,vdf_rec)\n",
    "project_tools.print_comparison_stats(vdf,vdf_rec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b527290-4290-4fdc-a2f1-4df96667b46f",
   "metadata": {},
   "source": [
    "### Principal Component Analysis (PCA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "569dd916-cc8b-485e-b1c0-d8a420bc5b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.decomposition import PCA\n",
    "vdf_temp=vdf.copy()\n",
    "vdf_temp[vdf_temp<1e-16]=1e-16\n",
    "vdf_temp = np.log10(vdf_temp)\n",
    "arr=vdf_temp.copy()\n",
    "arr = arr.reshape(arr.shape[0], -1)\n",
    "\n",
    "# perform pca\n",
    "cov_matrix = np.cov(arr, rowvar=False)\n",
    "eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)\n",
    "pca = PCA(n_components=30)\n",
    "compressed = pca.fit_transform(arr)\n",
    "#reconstruct the data\n",
    "recon = pca.inverse_transform(compressed)\n",
    "nx,ny,nz=np.shape(vdf_temp)\n",
    "recon=np.reshape(recon,(nx,ny,nz))\n",
    "recon = 10 ** recon\n",
    "recon[recon <= 1e-16] = 0\n",
    "print(f\"Compression achieved using PCA = {round(vdf_temp.size / compressed.size, 2)}\")\n",
    "project_tools.plot_vdfs(vdf, recon)\n",
    "project_tools.print_comparison_stats(vdf,recon)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b46c0119",
   "metadata": {},
   "source": [
    "### Discrete Wavelet Transform (DWT)\n",
    "\n",
    "Here we employ the PyWavelets package for prototyping compression via Discrete Wavelet Transforms [@lee_2019].\n",
    "\n",
    "First, we test quantizing the coefficients as a sanity check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5874082-9c03-420e-8708-c63215dd3a21",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pywt\n",
    "\n",
    "# print(pywt.wavelist())\n",
    "\n",
    "vdf_3d = vdf.copy()\n",
    "orig_shape = vdf_3d.shape\n",
    "vdf_3d[np.isnan(vdf_3d)] = 0\n",
    "\n",
    "comp_type = np.int8\n",
    "quant = np.iinfo(comp_type).max/3\n",
    "# quant = np.finfo(comp_type).max/3\n",
    "# quant_min = np.iinfo(comp_type).min/3\n",
    "\n",
    "# norm = np.max([np.nanmax(vdf_3d)/quant_max, np.nanmin(vdf_3d)/quant_min])\n",
    "# print(quant_max, quant_min, norm)\n",
    "\n",
    "norm = np.nanmax(vdf_3d)/quant\n",
    "\n",
    "vdf_3d /= norm\n",
    "vdf_3d[vdf_3d<0]=0\n",
    "print(np.nanmax(vdf_3d),np.nanmin(vdf_3d), norm)\n",
    "\n",
    "\n",
    "coeffs3 = pywt.dwtn(vdf_3d,'bior1.3')\n",
    "\n",
    "# print(coeffs3)\n",
    "coeffs3_compress = coeffs3.copy()\n",
    "\n",
    "for k, v in coeffs3.items():\n",
    "    print(np.min(v),np.max(v), len(v))\n",
    "    coeffs3_compress[k] = v.astype(comp_type)\n",
    "\n",
    "\n",
    "vdf_rec = pywt.idwtn(coeffs3_compress,'bior1.3')*norm\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "volume_compressed = 0\n",
    "for k,v in coeffs3_compress.items():\n",
    "    volume_compressed += sys.getsizeof(v)\n",
    "volume_orig = 0\n",
    "for k,v in coeffs3.items():\n",
    "    volume_orig += sys.getsizeof(v)\n",
    "compression = volume_orig/volume_compressed\n",
    "\n",
    "print(\"compression:\", compression)\n",
    "\n",
    "project_tools.plot_vdfs(vdf,vdf_rec)\n",
    "project_tools.print_comparison_stats(vdf,vdf_rec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fbec959",
   "metadata": {},
   "source": [
    "Next, we try a naive thresholding operation on the coefficients for reconstruction. The results are encouraging, and a more elaborate thresholding scheme could help with capturing details at the fringes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49654f5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pywt\n",
    "\n",
    "print(pywt.wavelist(kind='discrete'))\n",
    "\n",
    "vdf_3d = vdf.copy()\n",
    "orig_shape = vdf_3d.shape\n",
    "vdf_3d[np.isnan(vdf_3d)] = 0\n",
    "\n",
    "comp_type = np.float32\n",
    "# quant = np.iinfo(comp_type).max/3\n",
    "\n",
    "\n",
    "norm = 1# np.nanmax(vdf_3d)/quant\n",
    "\n",
    "vdf_3d /= norm\n",
    "vdf_3d[vdf_3d<0]=0\n",
    "print(np.nanmax(vdf_3d),np.nanmin(vdf_3d), norm)\n",
    "\n",
    "wavelet = 'db4'#'bior1.3'\n",
    "\n",
    "dwtn_mlevel = pywt.dwtn_max_level(vdf_3d.shape,wavelet)\n",
    "level_delta = 2\n",
    "print(\"Decomposing to \", dwtn_mlevel-level_delta, \"levels out of \", dwtn_mlevel)\n",
    "coeffs3 = pywt.wavedecn(vdf_3d,wavelet=wavelet, level = dwtn_mlevel-2)\n",
    "\n",
    "# print(coeffs3)\n",
    "coeffs3_comp = coeffs3.copy()\n",
    "print(type(coeffs3_comp))\n",
    "\n",
    "zeros = 0\n",
    "nonzeros = 0\n",
    "threshold = 1e-16\n",
    "for i,a in enumerate(coeffs3_comp):\n",
    "    print(type(a))\n",
    "    zero_app = False\n",
    "    # print(a.shape)\n",
    "    if(type(a) == type(np.ndarray(1))):\n",
    "        coeffs3_comp[i] = a\n",
    "        mask = np.abs(a) < threshold\n",
    "        zeros += np.sum(mask)\n",
    "        nonzeros += np.sum(~mask)\n",
    "        # nonzeros += np.prod(a.shape)\n",
    "        coeffs3_comp[i][mask] = 0\n",
    "    else:\n",
    "        for k,v in a.items():\n",
    "            mask = np.abs(v) < threshold\n",
    "            coeffs3_comp[i][k] = v\n",
    "            coeffs3_comp[i][k][mask] = 0\n",
    "            zeros += np.sum(mask)\n",
    "            nonzeros += np.sum(~mask)\n",
    "\n",
    "print(\"number of zeros:\", zeros, \"nonzeros:\", nonzeros)\n",
    "vdf_rec = pywt.waverecn(coeffs3_comp,wavelet=wavelet)*norm\n",
    "\n",
    "compression = (np.prod(vdf_3d.shape)/nonzeros)\n",
    "\n",
    "\n",
    "# volume_compressed = 0\n",
    "# for k,v in coeffs3_compress.items():\n",
    "#     volume_compressed += sys.getsizeof(v)\n",
    "# volume_orig = 0\n",
    "# for k,v in coeffs3.items():\n",
    "#     volume_orig += sys.getsizeof(v)\n",
    "# compression = volume_orig/volume_compressed\n",
    "\n",
    "print(\"compression:\", compression)\n",
    "\n",
    "project_tools.plot_vdfs(vdf,vdf_rec)\n",
    "project_tools.print_comparison_stats(vdf,vdf_rec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af891e88",
   "metadata": {},
   "source": [
    "Lastly, stationary wavelet transforms might have desirable qualities. However, computing these is very slow with the current package, and this output is disabled for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "603d5ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "if False:\n",
    "    import pywt\n",
    "\n",
    "    print(pywt.wavelist(kind='discrete'))\n",
    "\n",
    "    vdf_3d = vdf.copy()\n",
    "    orig_shape = vdf_3d.shape\n",
    "    vdf_3d[np.isnan(vdf_3d)] = 0\n",
    "\n",
    "    mlevel = 6\n",
    "    paddings = (np.ceil(np.array(vdf_3d.shape)/2**mlevel)).astype(int)*2**mlevel - vdf_3d.shape\n",
    "    paddings = ((0,paddings[0]),(0,paddings[1]),(0,paddings[2]))\n",
    "    vdf_3d = np.pad(vdf_3d, paddings)\n",
    "\n",
    "    comp_type = np.float32\n",
    "    # quant = np.iinfo(comp_type).max/3\n",
    "\n",
    "\n",
    "    norm = 1# np.nanmax(vdf_3d)/quant\n",
    "\n",
    "    vdf_3d /= norm\n",
    "    vdf_3d[vdf_3d<0]=0\n",
    "    # print(np.nanmax(vdf_3d),np.nanmin(vdf_3d), norm)\n",
    "\n",
    "    wavelet = 'db4'#'bior1.3'\n",
    "    coeffs3 = pywt.swtn(vdf_3d, wavelet, mlevel)\n",
    "\n",
    "    # print(coeffs3)\n",
    "    coeffs3_comp = coeffs3.copy()\n",
    "\n",
    "    zeros = 0\n",
    "    nonzeros = 0\n",
    "    threshold = 1e-15\n",
    "    for i,a in enumerate(coeffs3_comp):\n",
    "        # print(type(a))\n",
    "        # print(a.shape)\n",
    "        if(type(a) == type(np.ndarray(1))):\n",
    "            coeffs3_comp[i] = a\n",
    "            mask = np.abs(a) < threshold\n",
    "            zeros += np.sum(mask)\n",
    "            nonzeros += np.sum(~mask)\n",
    "            coeffs3_comp[i][mask] = 0\n",
    "        else:\n",
    "            for k,v in a.items():\n",
    "                mask = np.abs(v) < threshold\n",
    "                coeffs3_comp[i][k] = v\n",
    "                coeffs3_comp[i][k][mask] = 0\n",
    "                mask = np.abs(v) < threshold\n",
    "                zeros += np.sum(mask)\n",
    "                nonzeros += np.sum(~mask)\n",
    "\n",
    "    print(\"number of zeros:\", zeros, \"nonzeros:\", nonzeros)\n",
    "    vdf_rec = pywt.iswtn(coeffs3_comp, wavelet)*norm\n",
    "\n",
    "    compression = ((zeros+nonzeros)/nonzeros)\n",
    "\n",
    "\n",
    "    # volume_compressed = 0\n",
    "    # for k,v in coeffs3_compress.items():\n",
    "    #     volume_compressed += sys.getsizeof(v)\n",
    "    # volume_orig = 0\n",
    "    # for k,v in coeffs3.items():\n",
    "    #     volume_orig += sys.getsizeof(v)\n",
    "    # compression = volume_orig/volume_compressed\n",
    "\n",
    "    print(\"compression:\", compression)\n",
    "    vdf_rec = vdf_rec[0:orig_shape[0],0:orig_shape[1],0:orig_shape[2]]\n",
    "\n",
    "    project_tools.plot_vdfs(vdf,vdf_rec)\n",
    "    project_tools.print_comparison_stats(vdf,vdf_rec)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86bc7b80",
   "metadata": {},
   "source": [
    "### Vector Quantized-Variational AutoEncoder (VQ-VAE)\n",
    "\n",
    "This approach involves a different learning framework of first training an encoder-decoder scheme from a large set of existing data for compression and reconstruction and then packaging this scheme as a compression tool that can be called per-VDF. This is an example trained on a single VDF cell, but our use-case has already large amounts of training data available. Setting up large-scale training and assessment is to be done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67459050-d24f-44f2-90a6-a2181466529e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "import torch\n",
    "import torchvision\n",
    "from PIL import Image\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fbe7de6-8c58-40be-a45f-baa0746c6b67",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResidualStack(nn.Module):\n",
    "    def __init__(self, num_hiddens, num_residual_layers, num_residual_hiddens):\n",
    "        super().__init__()\n",
    "        # See Section 4.1 of \"Neural Discrete Representation Learning\".\n",
    "        layers = []\n",
    "        for i in range(num_residual_layers):\n",
    "            layers.append(\n",
    "                nn.Sequential(\n",
    "                    nn.ReLU(),\n",
    "                    nn.Conv3d(\n",
    "                        in_channels=num_hiddens,\n",
    "                        out_channels=num_residual_hiddens,\n",
    "                        kernel_size=3,\n",
    "                        padding=1,\n",
    "                    ),\n",
    "                    nn.ReLU(),\n",
    "                    nn.Conv3d(\n",
    "                        in_channels=num_residual_hiddens,\n",
    "                        out_channels=num_hiddens,\n",
    "                        kernel_size=1,\n",
    "                    ),\n",
    "                )\n",
    "            )\n",
    "\n",
    "        self.layers = nn.ModuleList(layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = x\n",
    "        for layer in self.layers:\n",
    "            h = h + layer(h)\n",
    "\n",
    "        # ResNet V1-style.\n",
    "        return torch.relu(h)\n",
    "\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels,\n",
    "        num_hiddens,\n",
    "        num_downsampling_layers,\n",
    "        num_residual_layers,\n",
    "        num_residual_hiddens,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        # See Section 4.1 of \"Neural Discrete Representation Learning\".\n",
    "        # The last ReLU from the Sonnet example is omitted because ResidualStack starts\n",
    "        # off with a ReLU.\n",
    "        conv = nn.Sequential()\n",
    "        for downsampling_layer in range(num_downsampling_layers):\n",
    "            if downsampling_layer == 0:\n",
    "                out_channels = num_hiddens // 2\n",
    "            elif downsampling_layer == 1:\n",
    "                (in_channels, out_channels) = (num_hiddens // 2, num_hiddens)\n",
    "\n",
    "            else:\n",
    "                (in_channels, out_channels) = (num_hiddens, num_hiddens)\n",
    "\n",
    "            conv.add_module(\n",
    "                f\"down{downsampling_layer}\",\n",
    "                nn.Conv3d(\n",
    "                    in_channels=in_channels,\n",
    "                    out_channels=out_channels,\n",
    "                    kernel_size=3,\n",
    "                    stride=1,\n",
    "                    padding=1,\n",
    "                ),\n",
    "            )\n",
    "            conv.add_module(f\"relu{downsampling_layer}\", nn.ReLU())\n",
    "\n",
    "        conv.add_module(\n",
    "            \"final_conv\",\n",
    "            nn.Conv3d(\n",
    "                in_channels=num_hiddens,\n",
    "                out_channels=num_hiddens,\n",
    "                kernel_size=5,\n",
    "                padding=1,\n",
    "            ),\n",
    "        )\n",
    "        self.conv = conv\n",
    "        self.residual_stack = ResidualStack(\n",
    "            num_hiddens, num_residual_layers, num_residual_hiddens\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.conv(x)\n",
    "        return self.residual_stack(h)\n",
    "\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        embedding_dim,\n",
    "        num_hiddens,\n",
    "        num_upsampling_layers,\n",
    "        num_residual_layers,\n",
    "        num_residual_hiddens,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        # See Section 4.1 of \"Neural Discrete Representation Learning\".\n",
    "        self.conv = nn.Conv3d(\n",
    "            in_channels=embedding_dim,\n",
    "            out_channels=num_hiddens,\n",
    "            kernel_size=1,\n",
    "        )\n",
    "\n",
    "        self.tconv = nn.ConvTranspose3d(\n",
    "            in_channels=num_hiddens,\n",
    "            out_channels=num_hiddens,\n",
    "            kernel_size=5,\n",
    "            stride=1,\n",
    "            padding=1,\n",
    "        )\n",
    "        \n",
    "        self.residual_stack = ResidualStack(\n",
    "            num_hiddens, num_residual_layers, num_residual_hiddens\n",
    "        )\n",
    "        upconv = nn.Sequential()\n",
    "        for upsampling_layer in range(num_upsampling_layers):\n",
    "            if upsampling_layer < num_upsampling_layers - 2:\n",
    "                (in_channels, out_channels) = (num_hiddens, num_hiddens)\n",
    "\n",
    "            elif upsampling_layer == num_upsampling_layers - 2:\n",
    "                (in_channels, out_channels) = (num_hiddens, num_hiddens // 2)\n",
    "\n",
    "            else:\n",
    "                (in_channels, out_channels) = (num_hiddens // 2, 1)#3)\n",
    "\n",
    "            upconv.add_module(\n",
    "                f\"up{upsampling_layer}\",\n",
    "                nn.ConvTranspose3d(\n",
    "                    in_channels=in_channels,\n",
    "                    out_channels=out_channels,\n",
    "                    kernel_size=3,\n",
    "                    stride=1,\n",
    "                    padding=1,\n",
    "                    #output_padding=(1,1,1),\n",
    "                ),\n",
    "            )\n",
    "            if upsampling_layer < num_upsampling_layers - 1:\n",
    "                upconv.add_module(f\"relu{upsampling_layer}\", nn.ReLU())\n",
    "\n",
    "        self.upconv = upconv\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.conv(x)\n",
    "        h = self.residual_stack(h)\n",
    "        h = self.tconv(h)\n",
    "        h = nn.ReLU()(h)\n",
    "        x_recon = self.upconv(h)\n",
    "        return x_recon\n",
    "\n",
    "\n",
    "class SonnetExponentialMovingAverage(nn.Module):\n",
    "    # See: https://github.com/deepmind/sonnet/blob/5cbfdc356962d9b6198d5b63f0826a80acfdf35b/sonnet/src/moving_averages.py#L25.\n",
    "    # They do *not* use the exponential moving average updates described in Appendix A.1\n",
    "    # of \"Neural Discrete Representation Learning\".\n",
    "    def __init__(self, decay, shape):\n",
    "        super().__init__()\n",
    "        self.decay = decay\n",
    "        self.counter = 0\n",
    "        self.register_buffer(\"hidden\", torch.zeros(*shape))\n",
    "        self.register_buffer(\"average\", torch.zeros(*shape))\n",
    "\n",
    "    def update(self, value):\n",
    "        self.counter += 1\n",
    "        with torch.no_grad():\n",
    "            self.hidden -= (self.hidden - value) * (1 - self.decay)\n",
    "            self.average = self.hidden / (1 - self.decay ** self.counter)\n",
    "\n",
    "    def __call__(self, value):\n",
    "        self.update(value)\n",
    "        return self.average\n",
    "\n",
    "\n",
    "class VectorQuantizer(nn.Module):\n",
    "    def __init__(self, embedding_dim, num_embeddings, use_ema, decay, epsilon):\n",
    "        super().__init__()\n",
    "        # See Section 3 of \"Neural Discrete Representation Learning\" and:\n",
    "        # https://github.com/deepmind/sonnet/blob/v2/sonnet/src/nets/vqvae.py#L142.\n",
    "\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.num_embeddings = num_embeddings\n",
    "        self.use_ema = use_ema\n",
    "        # Weight for the exponential moving average.\n",
    "        self.decay = decay\n",
    "        # Small constant to avoid numerical instability in embedding updates.\n",
    "        self.epsilon = epsilon\n",
    "\n",
    "        # Dictionary embeddings.\n",
    "        limit = 3 ** 0.5\n",
    "        e_i_ts = torch.FloatTensor(embedding_dim, num_embeddings).uniform_(\n",
    "            -limit, limit\n",
    "        )\n",
    "        if use_ema:\n",
    "            self.register_buffer(\"e_i_ts\", e_i_ts)\n",
    "        else:\n",
    "            self.register_parameter(\"e_i_ts\", nn.Parameter(e_i_ts))\n",
    "\n",
    "        # Exponential moving average of the cluster counts.\n",
    "        self.N_i_ts = SonnetExponentialMovingAverage(decay, (num_embeddings,))\n",
    "        # Exponential moving average of the embeddings.\n",
    "        self.m_i_ts = SonnetExponentialMovingAverage(decay, e_i_ts.shape)\n",
    "\n",
    "    def forward(self, x):\n",
    "        flat_x = x.permute(0, 2, 3, 4, 1).reshape(-1, self.embedding_dim)\n",
    "        distances = (\n",
    "            (flat_x ** 2).sum(1, keepdim=True)\n",
    "            - 2 * flat_x @ self.e_i_ts\n",
    "            + (self.e_i_ts ** 2).sum(0, keepdim=True)\n",
    "        )\n",
    "        encoding_indices = distances.argmin(1)\n",
    "        quantized_x = F.embedding(\n",
    "            encoding_indices.view(x.shape[0], *x.shape[2:]), self.e_i_ts.transpose(0, 1)\n",
    "        ).permute(0, 4, 1, 2, 3)\n",
    "\n",
    "        # See second term of Equation (3).\n",
    "        if not self.use_ema:\n",
    "            dictionary_loss = ((x.detach() - quantized_x) ** 2).mean()\n",
    "        else:\n",
    "            dictionary_loss = None\n",
    "\n",
    "        # See third term of Equation (3).\n",
    "        commitment_loss = ((x - quantized_x.detach()) ** 2).mean()\n",
    "        # Straight-through gradient. See Section 3.2.\n",
    "        quantized_x = x + (quantized_x - x).detach()\n",
    "\n",
    "        if self.use_ema and self.training:\n",
    "            with torch.no_grad():\n",
    "                # See Appendix A.1 of \"Neural Discrete Representation Learning\".\n",
    "\n",
    "                # Cluster counts.\n",
    "                encoding_one_hots = F.one_hot(\n",
    "                    encoding_indices, self.num_embeddings\n",
    "                ).type(flat_x.dtype)\n",
    "                n_i_ts = encoding_one_hots.sum(0)\n",
    "                # Updated exponential moving average of the cluster counts.\n",
    "                # See Equation (6).\n",
    "                self.N_i_ts(n_i_ts)\n",
    "\n",
    "                # Exponential moving average of the embeddings. See Equation (7).\n",
    "                embed_sums = flat_x.transpose(0, 1) @ encoding_one_hots\n",
    "                self.m_i_ts(embed_sums)\n",
    "\n",
    "                # This is kind of weird.\n",
    "                # Compare: https://github.com/deepmind/sonnet/blob/v2/sonnet/src/nets/vqvae.py#L270\n",
    "                # and Equation (8).\n",
    "                N_i_ts_sum = self.N_i_ts.average.sum()\n",
    "                N_i_ts_stable = (\n",
    "                    (self.N_i_ts.average + self.epsilon)\n",
    "                    / (N_i_ts_sum + self.num_embeddings * self.epsilon)\n",
    "                    * N_i_ts_sum\n",
    "                )\n",
    "                self.e_i_ts = self.m_i_ts.average / N_i_ts_stable.unsqueeze(0)\n",
    "\n",
    "        return (\n",
    "            quantized_x,\n",
    "            dictionary_loss,\n",
    "            commitment_loss,\n",
    "            encoding_indices.view(x.shape[0], -1),\n",
    "        )\n",
    "\n",
    "\n",
    "class VQVAE(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_channels,\n",
    "        num_hiddens,\n",
    "        num_downsampling_layers,\n",
    "        num_residual_layers,\n",
    "        num_residual_hiddens,\n",
    "        embedding_dim,\n",
    "        num_embeddings,\n",
    "        use_ema,\n",
    "        decay,\n",
    "        epsilon,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(\n",
    "            in_channels,\n",
    "            num_hiddens,\n",
    "            num_downsampling_layers,\n",
    "            num_residual_layers,\n",
    "            num_residual_hiddens,\n",
    "        )\n",
    "        self.pre_vq_conv = nn.Conv3d(\n",
    "            in_channels=num_hiddens, out_channels=embedding_dim, kernel_size=1\n",
    "        )\n",
    "        self.vq = VectorQuantizer(\n",
    "            embedding_dim, num_embeddings, use_ema, decay, epsilon\n",
    "        )\n",
    "        self.decoder = Decoder(\n",
    "            embedding_dim,\n",
    "            num_hiddens,\n",
    "            num_downsampling_layers,\n",
    "            num_residual_layers,\n",
    "            num_residual_hiddens,\n",
    "        )\n",
    "\n",
    "    def quantize(self, x):\n",
    "        z = self.pre_vq_conv(self.encoder(x))\n",
    "        (z_quantized, dictionary_loss, commitment_loss, encoding_indices) = self.vq(z)\n",
    "        return (z_quantized, dictionary_loss, commitment_loss, encoding_indices)\n",
    "\n",
    "    def forward(self, x):\n",
    "        (z_quantized, dictionary_loss, commitment_loss, _) = self.quantize(x)\n",
    "        x_recon = self.decoder(z_quantized)\n",
    "        return {\n",
    "            \"dictionary_loss\": dictionary_loss,\n",
    "            \"commitment_loss\": commitment_loss,\n",
    "            \"x_recon\": x_recon,\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d730c83-2ccb-465e-bbac-f43e27573744",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('Using device:', device)\n",
    "\n",
    "use_tb = True # Use Tensorboard (optional)\n",
    "\n",
    "# Initialize model\n",
    "use_ema = True # Use exponential moving average\n",
    "model_args = {\n",
    "    \"in_channels\": 1,\n",
    "    \"num_hiddens\": 128,\n",
    "    \"num_downsampling_layers\": 2,\n",
    "    \"num_residual_layers\": 2,\n",
    "    \"num_residual_hiddens\": 32,\n",
    "    \"embedding_dim\": 64,\n",
    "    \"num_embeddings\": 512,\n",
    "    \"use_ema\": use_ema,\n",
    "    \"decay\": 0.99,\n",
    "    \"epsilon\": 1e-5,\n",
    "}\n",
    "model = VQVAE(**model_args).to(device)\n",
    "\n",
    "# Initialize dataset\n",
    "batch_size = 1\n",
    "workers = 1\n",
    "\n",
    "input_norm = (input_array - input_array.min())/(input_array.max() - input_array.min()) # MinMax normalization\n",
    "input_tensor = torch.tensor(input_norm, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)  # Add batch and channel dimensions, move to device\n",
    "train_dataset = input_tensor\n",
    "train_loader = DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True,\n",
    "    num_workers=workers,\n",
    ")\n",
    "\n",
    "# Multiplier for commitment loss. See Equation (3) in \"Neural Discrete Representation Learning\"\n",
    "beta = 0.25\n",
    "\n",
    "# Initialize optimizer\n",
    "train_params = [params for params in model.parameters()]\n",
    "lr = 3e-4\n",
    "optimizer = optim.Adam(train_params, lr=lr)\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "# Train model\n",
    "epochs = 30\n",
    "eval_every = 1\n",
    "best_train_loss = float(\"inf\")\n",
    "model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b350359d-3ae0-46a9-be43-2a23a5ff720e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Util for Tensorboard plotting in color\n",
    "def prep_fig(img):\n",
    "    fig, ax = plt.subplots()\n",
    "    im = ax.imshow(img, interpolation='none')\n",
    "    #plt.colorbar(im, ax=ax)\n",
    "\n",
    "    fig.canvas.draw()\n",
    "    plt.close()\n",
    "\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab2df2e3-320a-4ec4-9bf4-c4296d03a0fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tensorboard setup\n",
    "if use_tb:\n",
    "    writer = SummaryWriter(\"runs/ASTERIX VQ-VAE_\"+str(model_args[\"num_hiddens\"])+\"h_ema\"+str(model_args[\"use_ema\"])+\"_\"+str(model_args[\"embedding_dim\"])+\"ed_res\"+str(model_args[\"num_residual_layers\"]))\n",
    "\n",
    "# Training\n",
    "for epoch in tqdm(range(epochs)):\n",
    "    total_train_loss = 0\n",
    "    total_recon_error = 0\n",
    "    n_train = 0\n",
    "    for (batch_idx, train_tensors) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        imgs = train_tensors[0].unsqueeze(0).to(device)\n",
    "        out = model(imgs)\n",
    "        recon_error = criterion(out[\"x_recon\"], imgs) #/ train_data_variance\n",
    "        total_recon_error += recon_error.item()\n",
    "        loss = recon_error + beta * out[\"commitment_loss\"]\n",
    "        if not use_ema:\n",
    "            loss += out[\"dictionary_loss\"]\n",
    "\n",
    "        total_train_loss += loss.item()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        n_train += 1\n",
    "\n",
    "        if ((batch_idx + 1) % eval_every) == 0:\n",
    "            print(f\"epoch: {epoch}\\nbatch_idx: {batch_idx + 1}\", flush=True)\n",
    "            total_train_loss /= n_train\n",
    "            if total_train_loss < best_train_loss:\n",
    "                best_train_loss = total_train_loss\n",
    "\n",
    "            print(f\"total_train_loss: {total_train_loss}\")\n",
    "            writer.add_scalar(\"tr loss: \", total_train_loss, epoch)\n",
    "\n",
    "            print(f\"best_train_loss: {best_train_loss}\")\n",
    "            print(f\"recon_error: {total_recon_error / n_train}\\n\")\n",
    "\n",
    "            total_train_loss = 0\n",
    "            total_recon_error = 0\n",
    "            n_train = 0\n",
    "    # Visualize reconstructed VDF at each epoch\n",
    "    #plt.imshow(out[\"x_recon\"].squeeze(0).squeeze(0).detach().numpy()[:,:,25])\n",
    "    #plt.show()\n",
    "\n",
    "    if use_tb:\n",
    "        # write to tensorboard\n",
    "        writer.add_figure('reconstructed VDF', prep_fig(out[\"x_recon\"].squeeze(0).squeeze(0)[:,:,25].detach().numpy()), epoch)\n",
    "\n",
    "if use_tb:\n",
    "    writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "883d53ca-de47-426e-9bab-a58591fd9ad6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save model checkpoint (optional)\n",
    "#torch.save({\n",
    "#    'epoch': epochs,\n",
    "#    'model_state_dict': model.state_dict(),\n",
    "#    'optimizer_state_dict': optimizer.state_dict()\n",
    "#    }, \"./VQ-VAE_checkpoints/ASTERIX VQ-VAE_\"+str(model_args[\"num_hiddens\"])+\"h_ema\"+str(model_args[\"use_ema\"])+\"_\"+str(model_args[\"embedding_dim\"])+\"ed_res\"+str(model_args[\"num_residual_layers\"])+\"_e\"+str(epochs)+\".pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba5853d-832d-436a-80d2-f31dcb639ae3",
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_vdf = out[\"x_recon\"].squeeze(0).squeeze(0).detach().numpy()\n",
    "rec_vdf_rescaled = (rec_vdf*(np.amax(input_array) - np.amin(input_array))) + np.amin(input_array)\n",
    "\n",
    "fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(5, 3))\n",
    "axes[0].imshow(vdf[:,:,25])\n",
    "axes[0].set_title(\"Original VDF\")\n",
    "axes[1].imshow(rec_vdf_rescaled[:,:,25])\n",
    "axes[1].set_title(\"Reconstructed VDF\")\n",
    "fig.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "## The random initialization of the cluster centroids in the embedding space, \n",
    "## can cause very different performance, e.g. when using the same input to train \n",
    "## (since in this case we train using a single VDF, we will expand during the next phase of the project)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7fd3e48-97ff-4d05-ab4d-ee07566e1e33",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}