{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to compute a Neural Adjoint Map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A neural adjoint map can be seen a functional map plus a non linear module. Given a correspondence, we can compute a neural adjoint map as we do for functional maps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import geomstats.backend as gs\n",
    "\n",
    "from geomfum.dataset import NotebooksDataset\n",
    "from geomfum.refine import NeuralZoomOut, ZoomOut\n",
    "from geomfum.shape import TriangleMesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = NotebooksDataset()\n",
    "\n",
    "mesh_a = TriangleMesh.from_file(dataset.get_filename(\"cat-00\"))\n",
    "mesh_b = TriangleMesh.from_file(dataset.get_filename(\"lion-00\"))\n",
    "\n",
    "mesh_a.n_vertices, mesh_b.n_vertices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh_a.laplacian.find_spectrum(spectrum_size=50, set_as_basis=True)\n",
    "mesh_b.laplacian.find_spectrum(spectrum_size=50, set_as_basis=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by estimating a correspondence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from geomfum.convert import NearestNeighbors\n",
    "\n",
    "finder = NearestNeighbors(n_neighbors=1)\n",
    "finder.fit(mesh_b.vertices)\n",
    "p2p = finder.kneighbors(mesh_a.vertices)[1].flatten()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we convert it into a NAM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from geomfum.convert import NamFromP2pConverter\n",
    "\n",
    "mesh_a.basis.use_k = 10\n",
    "mesh_b.basis.use_k = 10\n",
    "nam_converter = NamFromP2pConverter(device=\"cpu\")\n",
    "\n",
    "nam = nam_converter(p2p, mesh_a.basis, mesh_b.basis)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(nam)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the linear part of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fmap = nam.linear_module.weight.detach().cpu().numpy()\n",
    "plt.imshow(fmap)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given a NAM, we can obtain a correspondence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from geomfum.convert import P2pFromNamConverter\n",
    "\n",
    "p2p_from_nam = P2pFromNamConverter()\n",
    "\n",
    "p2p_from_nam(nam, mesh_a.basis, mesh_b.basis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As in ZoomOut, we can perform spectral upsampling on NAMS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nzo = NeuralZoomOut(nit=2, step=2, device=\"cpu\")\n",
    "\n",
    "nam_ref = nzo(nam, mesh_a.basis, mesh_b.basis)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "VENV",
   "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}