scripts/_mkl/notebooks/05b - Gaussian Approximation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "# Trimesh to Gaussians\n",
       "> Pretty much self-explanatory\n",
       "\n",
       "**Example:**\n",
       "```python\n",
       "from bayes3d._mkl.trimesh_to_gaussians import (\n",
       "    patch_trimesh, \n",
       "    uniformly_sample_from_mesh, \n",
       "    ellipsoid_embedding, \n",
       "    get_mean_colors, \n",
       "    pack_transform\n",
       ")\n",
       "import trimesh\n",
       "import numpy as np\n",
       "import jax.numpy as jnp\n",
       "import jax\n",
       "from jax import jit, vmap\n",
       "from sklearn.mixture import GaussianMixture\n",
       "from bayes3d._mkl.utils import keysplit\n",
       "\n",
       "# SEED\n",
       "key = jax.random.PRNGKey(0)\n",
       "\n",
       "# LOAD MESH\n",
       "# -------------------\n",
       "mesh = load_mesh(...)\n",
       "mesh = patch_trimesh(mesh)\n",
       "\n",
       "# SAMPLE FROM MESH\n",
       "# ----------------\n",
       "key = keysplit(key)\n",
       "n = 20_000\n",
       "xs, cs = uniformly_sample_from_mesh(key, n, mesh, with_color=True)\n",
       "\n",
       "# GMM CONFIG\n",
       "# ----------\n",
       "key = keysplit(key)\n",
       "n_components = 150\n",
       "noise        = 0.0; \n",
       "X            = xs + np.random.randn(*xs.shape)*noise\n",
       "means_init   = np.array(uniformly_sample_from_mesh(key, n_components, mesh, with_color=False)[0]);\n",
       "\n",
       "# FIT THE GMM\n",
       "# -----------\n",
       "gm = GaussianMixture(n_components=n_components, \n",
       "                     tol=1e-3, max_iter=100, \n",
       "                     covariance_type=\"full\", \n",
       "                     means_init=means_init).fit(X)\n",
       "\n",
       "mus        = gm.means_\n",
       "covs       = gm.covariances_\n",
       "labels     = gm.predict(X)\n",
       "choleskys  = vmap(ellipsoid_embedding)(covs)\n",
       "transforms = vmap(pack_transform, (0,0,None))(mus, choleskys, 2.0)\n",
       "mean_colors, nums = get_mean_colors(cs, gm.n_components, labels)\n",
       "```\n"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Markdown as md\n",
    "import bayes3d._mkl.trimesh_to_gaussians\n",
    "\n",
    "md(bayes3d._mkl.trimesh_to_gaussians._doc_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bayes3d._mkl.trimesh_to_gaussians import (\n",
    "    patch_trimesh, \n",
    "    sample_from_mesh,\n",
    "    barycentric_to_mesh as _barycentric_to_mesh,\n",
    "    uniformly_sample_from_mesh, \n",
    "    ellipsoid_embedding, \n",
    "    get_mean_colors, \n",
    "    pack_transform,\n",
    "    transform_from_gaussian\n",
    ")\n",
    "import trimesh\n",
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "import jax\n",
    "from jax import jit, vmap\n",
    "from sklearn.mixture import GaussianMixture\n",
    "from bayes3d._mkl.utils import keysplit\n",
    "\n",
    "\n",
    "# SEED\n",
    "key = jax.random.PRNGKey(0)\n",
    "\n",
    "barycentric_to_mesh = vmap(_barycentric_to_mesh, (0,0,None))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import traceviz.client\n",
    "import numpy as np\n",
    "from traceviz.proto import  viz_pb2\n",
    "import json\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_mesh(t):\n",
    "    mesh = trimesh.load(f\"data/flag_objs/flag_t_{t}.obj\")\n",
    "    mesh = patch_trimesh(mesh)\n",
    "    return mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LOAD MESH\n",
    "# -------------------\n",
    "mesh = load_mesh(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2.797654 2.322004 0.816412]\n"
     ]
    }
   ],
   "source": [
    "print(mesh.vertices.max(axis=0) - mesh.vertices.min(axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SAMPLE FROM MESH\n",
    "# ----------------\n",
    "key = keysplit(key)\n",
    "n = 100\n",
    "xs, ps, fs = sample_from_mesh(key, n, mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit(key, mesh, means_init, precisions_init, covariance_type=\"full\", iter=20, noise=0.0):\n",
    "    \n",
    "    # SAMPLE FROM MESH\n",
    "    # ----------------\n",
    "    _, key = keysplit(key, 1, 1)\n",
    "    n = 20_000\n",
    "    xs, cs = uniformly_sample_from_mesh(key, n, mesh, with_color=True)\n",
    "\n",
    "    # GMM CONFIG\n",
    "    # ----------\n",
    "    key = keysplit(key)\n",
    "    n_components = means_init.shape[0]\n",
    "    X            = xs + np.random.randn(*xs.shape)*noise\n",
    "\n",
    "    # FIT THE GMM\n",
    "    # -----------\n",
    "    gm = GaussianMixture(n_components=n_components, \n",
    "                        tol=1e-3, max_iter=iter, \n",
    "                        covariance_type=covariance_type, \n",
    "                        means_init=means_init,\n",
    "                        precisions_init=precisions_init).fit(X)\n",
    "\n",
    "    mus    = gm.means_\n",
    "    if gm.covariance_type == \"spherical\":\n",
    "        covs = gm.covariances_[:,None,None]*jnp.eye(3)[None,:,:]\n",
    "    else:\n",
    "        covs       = gm.covariances_\n",
    "    labels = gm.predict(X)\n",
    "    mean_colors, nums = get_mean_colors(cs, gm.n_components, labels)\n",
    "\n",
    "    return mus, covs, mean_colors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "n: 10000\n",
      "n_components: 100\n",
      "noise: 0.01\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# SAMPLE FROM MESH\n",
    "# ----------------\n",
    "key = keysplit(key)\n",
    "n = 10_000\n",
    "xs, cs = uniformly_sample_from_mesh(key, n, mesh, with_color=True)\n",
    "\n",
    "# GMM CONFIG\n",
    "# ----------\n",
    "key = keysplit(key)\n",
    "n_components = 100\n",
    "noise        = 0.01; \n",
    "X            = xs + np.random.randn(*xs.shape)*noise\n",
    "means_init   = np.array(uniformly_sample_from_mesh(key, n_components, mesh, with_color=False)[0]);\n",
    "\n",
    "# FIT THE GMM\n",
    "# -----------\n",
    "gm = GaussianMixture(n_components=n_components, \n",
    "                     tol=1e-3, max_iter=100, \n",
    "                     covariance_type=\"spherical\", \n",
    "                     means_init=means_init).fit(X)\n",
    "\n",
    "\n",
    "mus        = gm.means_\n",
    "if gm.covariance_type == \"spherical\":\n",
    "    covs = gm.covariances_[:,None,None]*jnp.eye(3)[None,:,:]\n",
    "else:\n",
    "    covs       = gm.covariances_\n",
    "labels     = gm.predict(X)\n",
    "transforms = vmap(transform_from_gaussian, (0,0,None))(mus, covs, 2.0)\n",
    "mean_colors, nums = get_mean_colors(cs, gm.n_components, labels)\n",
    "\n",
    "print(f\"\"\"\n",
    "n: {n}\n",
    "n_components: {n_components}\n",
    "noise: {noise}\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "response:  listener_identifiers: \"2023-11-22T15:51:30.513873 :: ipv4:127.0.0.1:53162\"\n",
      "\n",
      "response:  listener_identifiers: \"2023-11-22T15:51:30.513873 :: ipv4:127.0.0.1:53162\"\n",
      "\n"
     ]
    }
   ],
   "source": [
    "msg = viz_pb2.Message()\n",
    "msg.payload.json = json.dumps({\"type\": \"setup\"})\n",
    "msg.payload.data.MergeFrom(traceviz.client.to_pytree_msg([]))\n",
    "stub = traceviz.client.connect()\n",
    "print('response: ', stub.Broadcast(msg))\n",
    "\n",
    "msg = viz_pb2.Message()\n",
    "msg.payload.json = json.dumps({\"type\": \"Gaussians2\"})\n",
    "msg.payload.data.MergeFrom(traceviz.client.to_pytree_msg({\n",
    "        'transforms': np.array(transforms )[nums>0], \n",
    "        'colors':     np.array(mean_colors)[nums>0]  \n",
    "}))\n",
    "stub = traceviz.client.connect()\n",
    "print('response: ', stub.Broadcast(msg))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit(key, mesh, means_init, precisions_init, covariance_type=\"full\", iter=20, noise=0.0):\n",
    "    \n",
    "    # SAMPLE FROM MESH\n",
    "    # ----------------\n",
    "    _, key = keysplit(key, 1, 1)\n",
    "    n = 20_000\n",
    "    xs, cs = uniformly_sample_from_mesh(key, n, mesh, with_color=True)\n",
    "\n",
    "    # GMM CONFIG\n",
    "    # ----------\n",
    "    key = keysplit(key)\n",
    "    n_components = means_init.shape[0]\n",
    "    X            = xs + np.random.randn(*xs.shape)*noise\n",
    "\n",
    "    # FIT THE GMM\n",
    "    # -----------\n",
    "    gm = GaussianMixture(n_components=n_components, \n",
    "                        tol=1e-3, max_iter=iter, \n",
    "                        covariance_type=covariance_type, \n",
    "                        means_init=means_init,\n",
    "                        precisions_init=precisions_init).fit(X)\n",
    "\n",
    "    mus    = gm.means_\n",
    "    if gm.covariance_type == \"spherical\":\n",
    "        covs = gm.covariances_[:,None,None]*jnp.eye(3)[None,:,:]\n",
    "    else:\n",
    "        covs       = gm.covariances_\n",
    "    labels = gm.predict(X)\n",
    "    mean_colors, nums = get_mean_colors(cs, gm.n_components, labels)\n",
    "\n",
    "    return mus, covs, mean_colors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 3, 3)"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "covs.shape\n",
    "prec = vmap(jnp.linalg.inv)(CVs[-1])\n",
    "prec.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/dcolmap/lib/python3.11/site-packages/sklearn/mixture/_base.py:268: ConvergenceWarning: Initialization 1 did not converge. Try different init parameters, or increase max_iter, tol or check for degenerate data.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/dcolmap/lib/python3.11/site-packages/sklearn/mixture/_base.py:268: ConvergenceWarning: Initialization 1 did not converge. Try different init parameters, or increase max_iter, tol or check for degenerate data.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/dcolmap/lib/python3.11/site-packages/sklearn/mixture/_base.py:268: ConvergenceWarning: Initialization 1 did not converge. Try different init parameters, or increase max_iter, tol or check for degenerate data.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\r"
     ]
    }
   ],
   "source": [
    "MUs = [mus]\n",
    "CVs = [covs]\n",
    "CLs = [mean_colors]\n",
    "\n",
    "for t in range(1,5):\n",
    "    print(t, end=\"\\r\")\n",
    "    mesh = load_mesh(t)\n",
    "    key = keysplit(key)\n",
    "\n",
    "    mus, covs, mean_colors = fit(key, mesh, MUs[-1], vmap(jnp.linalg.inv)(CVs[-1])[:,0,0], covariance_type=\"spherical\", iter=10, noise=0.0)\n",
    "    MUs.append(mus)\n",
    "    CVs.append(covs)\n",
    "    CLs.append(mean_colors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "response:  listener_identifiers: \"2023-11-22T16:16:10.810408 :: ipv4:127.0.0.1:38474\"\n",
      "\n",
      "response:  listener_identifiers: \"2023-11-22T16:16:10.810408 :: ipv4:127.0.0.1:38474\"\n",
      "\n"
     ]
    }
   ],
   "source": [
    "t = 4\n",
    "mus = MUs[t]\n",
    "covs = CVs[t]\n",
    "mean_colors = CLs[t]\n",
    "transforms = vmap(transform_from_gaussian, (0,0,None))(mus, covs, 3.0)\n",
    "\n",
    "colors = 0.4*jnp.ones_like(mean_colors)\n",
    "colors = colors.at[:,3].set(1.)\n",
    "\n",
    "\n",
    "msg = viz_pb2.Message()\n",
    "msg.payload.json = json.dumps({\"type\": \"setup\"})\n",
    "msg.payload.data.MergeFrom(traceviz.client.to_pytree_msg([]))\n",
    "stub = traceviz.client.connect()\n",
    "print('response: ', stub.Broadcast(msg))\n",
    "\n",
    "msg = viz_pb2.Message()\n",
    "msg.payload.json = json.dumps({\"type\": \"Gaussians2\"})\n",
    "msg.payload.data.MergeFrom(traceviz.client.to_pytree_msg({\n",
    "        'transforms': np.array(transforms ), \n",
    "        'colors':     np.array(colors)\n",
    "}))\n",
    "stub = traceviz.client.connect()\n",
    "print('response: ', stub.Broadcast(msg))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dcolmap",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}