{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# How to compute a functional map?"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {},
            "outputs": [],
            "source": [
                "import geomstats.backend as gs\n",
                "\n",
                "from geomfum.dataset import NotebooksDataset\n",
                "from geomfum.descriptor.pipeline import (\n",
                "    ArangeSubsampler,\n",
                "    DescriptorPipeline,\n",
                "    L2InnerNormalizer,\n",
                ")\n",
                "from geomfum.descriptor.spectral import HeatKernelSignature, WaveKernelSignature\n",
                "from geomfum.functional_map import (\n",
                "    FactorSum,\n",
                "    LBCommutativityEnforcing,\n",
                "    OperatorCommutativityEnforcing,\n",
                "    SpectralDescriptorPreservation,\n",
                ")\n",
                "from geomfum.numerics.optimization import ScipyMinimize\n",
                "from geomfum.shape import TriangleMesh"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "[Load meshes](00_load_mesh_from_file.ipynb)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "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\"))"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "[Set Laplace eigenbasis](./02_mesh_laplacian_spectrum.ipynb) for each mesh."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {},
            "outputs": [],
            "source": [
                "mesh_a.laplacian.find_spectrum(spectrum_size=10, set_as_basis=True)\n",
                "mesh_b.laplacian.find_spectrum(spectrum_size=10, set_as_basis=True)\n",
                "\n",
                "# I decide to visualize just the first 8 eigenfunctions\n",
                "\n",
                "mesh_b.basis.use_k = 8"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Set a [descriptor pipeline](./04_descriptor_pipeline.ipynb) and apply it to both shapes."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "metadata": {},
            "outputs": [],
            "source": [
                "steps = [\n",
                "    HeatKernelSignature.from_registry(n_domain=4),\n",
                "    ArangeSubsampler(subsample_step=2),\n",
                "    WaveKernelSignature.from_registry(n_domain=3),\n",
                "    L2InnerNormalizer(),\n",
                "]\n",
                "\n",
                "pipeline = DescriptorPipeline(steps)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {},
            "outputs": [],
            "source": [
                "descr_a = pipeline.apply(mesh_a)\n",
                "descr_b = pipeline.apply(mesh_b)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Create objective function: \n",
                "\n",
                "The optimization of the functional map can be performed considering different constraints, for example:\n",
                "1) SpectralDescriptorPreservation: the functional map needs to align spectral coefficients of descriptors.\n",
                "2) LBCommutativityEnforcing: the functional map needs to commute with the Laplace Beltrami operator.\n",
                "3) OperatorCommutativityEnforcing: the functional map needs to commute with a chosen operator defined on meshes.\n",
                "\n",
                "Details about these energies can be found in https://dl.acm.org/doi/10.1145/3084873.3084877."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "metadata": {},
            "outputs": [],
            "source": [
                "factors = [\n",
                "    SpectralDescriptorPreservation(\n",
                "        mesh_a.basis.project(descr_a),\n",
                "        mesh_b.basis.project(descr_b),\n",
                "        weight=1.0,\n",
                "    ),\n",
                "    LBCommutativityEnforcing.from_bases(\n",
                "        mesh_a.basis,\n",
                "        mesh_b.basis,\n",
                "        weight=1e-2,\n",
                "    ),\n",
                "    OperatorCommutativityEnforcing.from_multiplication(\n",
                "        mesh_a.basis, descr_a, mesh_b.basis, descr_b, weight=1e-1\n",
                "    ),\n",
                "    OperatorCommutativityEnforcing.from_orientation(\n",
                "        mesh_a, descr_a, mesh_b, descr_b, weight=1e-1\n",
                "    ),\n",
                "]\n",
                "\n",
                "objective = FactorSum(factors)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Instantiate an `Optimizer` and solve for the functional map matrix."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {},
            "outputs": [],
            "source": [
                "optimizer = ScipyMinimize(\n",
                "    method=\"L-BFGS-B\",\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "(8, 10)"
                        ]
                    },
                    "execution_count": 8,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "x0 = gs.zeros((mesh_b.basis.spectrum_size, mesh_a.basis.spectrum_size))\n",
                "\n",
                "res = optimizer.minimize(\n",
                "    objective,\n",
                "    x0,\n",
                "    fun_jac=objective.gradient,\n",
                ")\n",
                "\n",
                "fmap = res.x.reshape(x0.shape)\n",
                "\n",
                "fmap.shape"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Further reading\n",
                "\n",
                "* [How to compute a pointwise map from a functional map?](./10_pointwise_from_functional.ipynb)\n",
                "\n",
                "* [How to refine a functional map?](./15_refine_functional_map.ipynb)\n",
                "\n",
                "* [How to use ReMatching to compute a functional map?](./13_rematching.ipynb)"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "py12",
            "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.12.3"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}