{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph Construction\n",
    "- We shall be creating weighted KNN graphs on pointclouds with Faiss library.\n",
    "- By default the graphs shall be directed but they can be symmetrized.\n",
    "- Our graph representation shall be same as that of torch_geometric (*i.e.* edge_index and edge_attr).\n",
    "- edge_index is a (2, E) size tensor, which contains all the edges. \n",
    "- edge_attr is a (E,1) size tensor, which contains the scalar weights associated with the edges.\n",
    "- The graphs constructed in these notebooks *shall be used for pointcloud shape and color processing in the Primal_Dual_TV notebook and Tikhonov_reg notebook, respectively*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the required libs\n",
    "import numpy as np\n",
    "import faiss\n",
    "import torch\n",
    "from utilities import *\n",
    "from torch_geometric.data import Data\n",
    "from torch_geometric.utils import is_undirected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;33m[Open3D WARNING] geometry::TriangleMesh appears to be a geometry::PointCloud (only contains vertices, but no triangles).\u001b[0;m\n"
     ]
    }
   ],
   "source": [
    "#load the pcd/mesh \n",
    "(pos, tex, fac) = getPositionTexture(\"data/noisy_girl_skate.ply\") # Color Processing.\n",
    "#(pos, tex, fac) = getPositionTexture(\"data/noisy_3d_signal.ply\")# Shape Processing.\n",
    "pos = (pos - np.min(pos))/(np.max(pos)-np.min(pos)) #a good practice.\n",
    "displaySur(**dict(position=pos, texture=tex)) # Open3d visualization works only locally!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "- Create a knn graph using the position coords of the pointcloud.\n",
    "- Assign scalar weigths to the edges.\n",
    "\"\"\"\n",
    "\n",
    "#Faiss graph construction\n",
    "res = faiss.StandardGpuResources()\n",
    "index = faiss.IndexFlatL2(pos.shape[1])\n",
    "gpu_index_flat = faiss.index_cpu_to_gpu(res,0,index)\n",
    "gpu_index_flat.add(pos.astype(np.float32))\n",
    "k = 8\n",
    "D, I = gpu_index_flat.search(pos.astype(np.float32),k+1)\n",
    "#Convert to torch_geometric Data class \n",
    "edge_index = np.vstack((I[:,1:].flatten(), np.repeat(I[:,0].flatten(),k)))\n",
    "\n",
    "### Shape Processing ###\n",
    "#edge_attr = np.ones(edge_index.shape[1]) # Lets keep the weights equal to 1 !\n",
    "### Shape Processing ###\n",
    "\n",
    "### Color Processing ###\n",
    "# RBF kernel\n",
    "edge_attr = np.exp(-np.sum(((tex[I]-tex[I][:,0,None])**2), axis=2)/(0.2)**2)[:,1:].flatten() \n",
    "### Color Porcessing ###\n",
    "\n",
    "edge_index = torch.from_numpy(edge_index).type(torch.long) # it is important to convert to torch.long\n",
    "edge_attr = torch.from_numpy(edge_attr).type(torch.float32)\n",
    "edge_attr = edge_attr.view(-1,1)\n",
    "#getWgStats(edge_attr.numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph symmetrization\n",
    "- Let $A$ be the adjacency representation of the graph. We shall be symmerizing the graph using the following transformation:\n",
    "$$ A = \\frac{ A + A^{T}}{2}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "#Check if the graph is symmetric and create a temporary graph.\n",
    "print(is_undirected(edge_index))\n",
    "tmp_graph = Data(edge_index = edge_index, edge_attr=edge_attr, num_nodes=len(pos))\n",
    "\n",
    "# graph symmetrization by converting to a sparse tensor.\n",
    "tst = ToSparseTensor()\n",
    "nG = tst(tmp_graph).adj_t.to_symmetric().to_torch_sparse_coo_tensor()\n",
    "new_edge_index = torch.stack((nG.coalesce().indices()[1], nG.coalesce().indices()[0]))\n",
    "new_edge_attr = nG.coalesce().values()\n",
    "\n",
    "# Create a new graph\n",
    "graph = Data(edge_index = new_edge_index.type(torch.long), edge_attr=new_edge_attr/2, \n",
    "             x=torch.from_numpy(pos).type(torch.float32), tex=torch.from_numpy(tex).type(torch.float32))\n",
    "print(is_undirected(graph.edge_index))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save it\n",
    "torch.save(graph, \"./data/girl_skate.pt\") # Color Processing\n",
    "#torch.save(graph, \"./data/3d_signal.pt\") # Shape Processing"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}