tables_nerf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8b553f8a",
   "metadata": {},
   "source": [
    "# TILTED paper table generation\n",
    "\n",
    "Let's start with imports and types..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "975935d0-e1d6-4352-9d12-167d9469a406",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from typing import Any, Callable, Dict, Generic, Literal, Tuple, TypeVar\n",
    "from pathlib import Path\n",
    "from tqdm.auto import tqdm\n",
    "import dataclasses\n",
    "import yaml\n",
    "import numpy as onp\n",
    "import itertools\n",
    "\n",
    "EvalMetricName = Literal[\"lpips\", \"mse\", \"psnr\", \"ssim\"]\n",
    "TrainMetricName = Literal[\n",
    "    \"iterations_per_second\",\n",
    "    \"train/distortion_loss\",\n",
    "    \"train/grad_norm\",\n",
    "    \"train/interlevel_loss\",\n",
    "    \"train/iterations_per_sec\",  # Duplicate, oops!\n",
    "    \"train/l12_reg\",\n",
    "    \"train/loss\",\n",
    "    \"train/mse\",\n",
    "    \"train/psnr\",\n",
    "    \"train/tv_reg_l1\",\n",
    "    \"train/tv_reg_l2\",\n",
    "]\n",
    "\n",
    "EvalMetrics = Dict[EvalMetricName, float]\n",
    "TrainMetrics = Dict[TrainMetricName, float]\n",
    "ConfigDict = Dict[str, Any]\n",
    "\n",
    "T = TypeVar(\"T\")\n",
    "\n",
    "@dataclasses.dataclass\n",
    "class Result(Generic[T]):\n",
    "    config: ConfigDict\n",
    "    metrics: T\n",
    "    path: Path\n",
    "\n",
    "class Results(Tuple[Result[T], ...], Generic[T]):\n",
    "    def filter(self, condition: Callable[[Result], bool]) -> Results[T]:\n",
    "        return Results(filter(condition, self))\n",
    "\n",
    "    def filter_config_eq(self, config: Dict[str, Any]) -> Results[T]:\n",
    "        return Results(\n",
    "            x for x in self\n",
    "            if all(map(lambda k: x.config[k] == config[k], config.keys()))\n",
    "        )\n",
    "\n",
    "    def filter_config_neq(self, config: Dict[str, Any]) -> Results[T]:\n",
    "        return Results(\n",
    "            x for x in self\n",
    "            if all(map(lambda k: x.config[k] != config[k], config.keys()))\n",
    "        )\n",
    "    \n",
    "    def assert_length(self, length: int) -> Results[T]:\n",
    "        assert len(self) == length, f\"Expected length {length}, but got {len(self)}\"\n",
    "        return self\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c015e49",
   "metadata": {},
   "source": [
    "Load our NeRF results..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "10150726",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5b1a3b9648341f7916eb89f7916c757",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1458 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3a79377d51a54ff9aaa3f54d66b86460",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1458 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_output_dirs = [\n",
    "    Path(\"./paper_results/\"),\n",
    "]\n",
    "\n",
    "eval_results = Results[EvalMetrics](\n",
    "    Result(yaml.load((p / \"config_dict.yaml\").read_text(), yaml.Loader), yaml.load((p / \"eval_metrics_fixed.yaml\").read_text(), yaml.Loader), p)\n",
    "    for p in tqdm(list(itertools.chain(*[d.iterdir() for d in _output_dirs])))\n",
    "    if (p / \"eval_metrics_fixed.yaml\").exists()\n",
    ")\n",
    "train_results = Results[TrainMetrics](\n",
    "    Result(yaml.load((p / \"config_dict.yaml\").read_text(), yaml.Loader), yaml.load((p / \"train_metrics.yaml\").read_text(), yaml.Loader), p)\n",
    "    for p in tqdm(list(itertools.chain(*[d.iterdir() for d in _output_dirs])))\n",
    "    if (p / \"train_metrics.yaml\").exists()\n",
    ")\n",
    "\n",
    "del _output_dirs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c09fa9cd",
   "metadata": {},
   "source": [
    "Make some helper functions..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0d955bf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bold(text: str) -> str:\n",
    "    return r\"\\textbf{\" + text + \"}\"\n",
    "\n",
    "def get_metric_moments(results: Results, metric: str) -> str:\n",
    "    metric_from_dataset = {}\n",
    "    for result in results:\n",
    "        dataset = result.config[\"dataset_path\"]\n",
    "        if dataset not in metric_from_dataset:\n",
    "            metric_from_dataset[dataset] = []\n",
    "        metric_from_dataset[dataset].append(result.metrics[metric])\n",
    "\n",
    "    means = []\n",
    "    stderrs = []\n",
    "    for x in metric_from_dataset.values():\n",
    "        means.append(onp.mean(x))\n",
    "        stderrs.append(onp.std(x) / onp.sqrt(len(x) - 1))\n",
    "        \n",
    "    mean = onp.mean(means)\n",
    "    stderr = onp.sqrt(onp.sum(onp.array(stderrs) ** 2)) / len(stderrs)\n",
    "    return f\"{mean:.2f}\" + r\"{\\scriptsize$\\pm$\" + f\"{stderr:.2f}\" + r\"}\"\n",
    "\n",
    "def get_psnr(results: Results[EvalMetrics]) -> str:\n",
    "    return(get_metric_moments(results, \"psnr\"))\n",
    "\n",
    "def get_train_psnr(results: Results[TrainMetrics]) -> str:\n",
    "    return(get_metric_moments(results, \"train/psnr\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdc165f6",
   "metadata": {},
   "source": [
    "### How do naive K-Planes / VM compare on axis-aligned vs randomly rotated synthetic data?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2cffaa78",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{tabular}{lllll}\n",
      "\\toprule\n",
      "& K-Planes & VM\\\\\n",
      "\\cmidrule(lr){2-2} \\cmidrule(l){3-3}\n",
      "Lego & 35.31{\\scriptsize$\\pm$0.02} $\\to$ 33.29{\\scriptsize$\\pm$0.11} & 34.24{\\scriptsize$\\pm$0.04} $\\to$ 32.63{\\scriptsize$\\pm$0.01}\\\\\n",
      "Average & 32.12{\\scriptsize$\\pm$0.02} $\\to$ 31.62{\\scriptsize$\\pm$0.04} & 31.30{\\scriptsize$\\pm$0.03} $\\to$ 30.76{\\scriptsize$\\pm$0.03}\\\\\n",
      "\\bottomrule\n",
      "\\end{tabular}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def _cell(lego: bool, kplanes: bool, axis_aligned: bool) -> str:\n",
    "    return get_psnr(\n",
    "        eval_results.filter_config_eq({\n",
    "            \"dataset_type\": \"blender\",\n",
    "            \"field.primary_channels\": 64 if kplanes else 192,\n",
    "            \"primary_transform_count\": None,\n",
    "        }).filter_config_eq({\n",
    "            \"dataset_path\": Path(\"data/nerf_synthetic/lego\"),\n",
    "        } if lego else {}).filter_config_eq({ \n",
    "            # Only include axis-aligned scenes.\n",
    "            \"render_config.global_rotate_seed\": None,\n",
    "        } if axis_aligned else {}).filter_config_neq({\n",
    "            # Randomly rotated scenes only.\n",
    "            \"render_config.global_rotate_seed\": None,\n",
    "        } if not axis_aligned else {}).assert_length(3 if lego else 24)\n",
    "    )\n",
    "\n",
    "\n",
    "print(\n",
    "    r\"\\begin{tabular}{lllll}\",\n",
    "    r\"\\toprule\",\n",
    "    r\"& K-Planes & VM\\\\\",\n",
    "    r\"\\cmidrule(lr){2-2} \\cmidrule(l){3-3}\",\n",
    "    f\"Lego & {_cell(lego=True, kplanes=True, axis_aligned=True)} $\\\\to$ {_cell(lego=True, kplanes=True, axis_aligned=False)} & {_cell(lego=True, kplanes=False, axis_aligned=True)} $\\\\to$ {_cell(lego=True, kplanes=False, axis_aligned=False)}\" + r\"\\\\\",\n",
    "    f\"Average & {_cell(lego=False, kplanes=True, axis_aligned=True)} $\\\\to$ {_cell(lego=False, kplanes=True, axis_aligned=False)} & {_cell(lego=False, kplanes=False, axis_aligned=True)} $\\\\to$ {_cell(lego=False, kplanes=False, axis_aligned=False)}\" + r\"\\\\\",\n",
    "    r\"\\bottomrule\",\n",
    "    r\"\\end{tabular}\",\n",
    "    sep=\"\\n\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c6209ca",
   "metadata": {},
   "source": [
    "### What impact does incorporating TILTED have on randomly rotated synthetic data?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4913261d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{tabular}{lllll}\n",
      "\\toprule\n",
      "& K-Planes & w/ TILTED & VM & w/ TILTED \\\\\n",
      "\\cmidrule(lr){2-3} \\cmidrule(l){4-5}\n",
      "Lego & 33.29{\\scriptsize$\\pm$0.11} & \\textbf{34.35{\\scriptsize$\\pm$0.07}} & 32.63{\\scriptsize$\\pm$0.01} & \\textbf{33.90{\\scriptsize$\\pm$0.06}}\\\\\n",
      "Average & 31.62{\\scriptsize$\\pm$0.04} & \\textbf{31.91{\\scriptsize$\\pm$0.04}} & 30.76{\\scriptsize$\\pm$0.03} & \\textbf{31.08{\\scriptsize$\\pm$0.02}}]\\\\\n",
      "\\bottomrule\n",
      "\\end{tabular}\n"
     ]
    }
   ],
   "source": [
    "def _cell(lego: bool, kplanes: bool, tilted: bool) -> str:\n",
    "    return get_psnr(\n",
    "        eval_results.filter_config_eq({\n",
    "            \"dataset_type\": \"blender\",\n",
    "            \"field.primary_channels\": 64 if kplanes else 192,\n",
    "            \"primary_transform_count\": 8 if tilted else None,\n",
    "            \"field.grid_type\": \"kplane\" if kplanes else \"vm\"\n",
    "        }).filter_config_eq({\n",
    "            \"dataset_path\": Path(\"data/nerf_synthetic/lego\"),\n",
    "        } if lego else {}).filter_config_neq({\n",
    "            # For axis-aligned approaches, only include randomly rotated scenes.\n",
    "            # For TILTED, we include both the axis-aligned and the randomly rotated options.\n",
    "            \"render_config.global_rotate_seed\": None,\n",
    "        } if not tilted else {}).filter(\n",
    "            lambda result: result.path.name.endswith(\"-bottleneck\") if tilted else True\n",
    "        ).assert_length(\n",
    "            (3 if lego else 24) * (2 if tilted else 1)\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "print(\n",
    "    r\"\\begin{tabular}{lllll}\",\n",
    "    r\"\\toprule\",\n",
    "    r\"& K-Planes & w/ TILTED & VM & w/ TILTED \\\\\",\n",
    "    r\"\\cmidrule(lr){2-3} \\cmidrule(l){4-5}\",\n",
    "    f\"Lego & {_cell(True, kplanes=True, tilted=False)} & {bold(_cell(True, kplanes=True, tilted=True))} & {_cell(True, kplanes=False, tilted=False)} & {bold(_cell(True, kplanes=False, tilted=True))}\" r\"\\\\\",\n",
    "    f\"Average & {_cell(False, kplanes=True, tilted=False)} & {bold(_cell(False, kplanes=True, tilted=True))} & {_cell(False, kplanes=False, tilted=False)} & {bold(_cell(False, kplanes=False, tilted=True))}\" r\"]\\\\\",\n",
    "    r\"\\bottomrule\",\n",
    "    r\"\\end{tabular}\",\n",
    "    sep=\"\\n\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31ac2528",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e7017a17",
   "metadata": {},
   "source": [
    "### Real-world experiments: how does TILTED compare?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bf234363",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "How often does TILTED improve performance?\n",
      "100.0 %\n",
      "\n",
      "How often does adding TILTED outperform _doubling the channel count_?\n",
      "63.888888888888886 %\n",
      "\n",
      "How many experiments have we run?\n",
      "72\n",
      "\n",
      "What are the runtimes of each method?\n",
      "32c 9:16\n",
      "32c TILTED 11:4\n",
      "64c 14:46\n",
      "64c TILTED 17:36\n"
     ]
    }
   ],
   "source": [
    "dataset_names = set(\n",
    "    [\n",
    "        p.config[\"dataset_path\"].name for p in\n",
    "        train_results.filter_config_eq({\n",
    "            \"dataset_type\": \"nerfstudio\",\n",
    "        })\n",
    "    ]\n",
    ")\n",
    "\n",
    "def get_train_psnr(dname: str, channels: int, grid_type: Literal[\"kplane\", \"vm\"], tilted: bool) -> float:\n",
    "    return train_results.filter(\n",
    "        lambda p: p.config[\"dataset_path\"].name == dname\n",
    "    ).filter_config_eq(\n",
    "        {\n",
    "            \"primary_transform_count\": 8 if tilted else None,\n",
    "            \"field.primary_channels\": channels,\n",
    "            \"field.grid_type\": grid_type,\n",
    "        }\n",
    "    ).assert_length(1)[0].metrics[\"train/psnr\"]\n",
    "\n",
    "count = 0\n",
    "tilted_better_count = 0\n",
    "tilted_better_than_naive_double_count = 0\n",
    "tolerance = 0.0\n",
    "\n",
    "for dname in dataset_names:\n",
    "    naive = get_train_psnr(dname, 32, 'kplane', False)\n",
    "    tilted = get_train_psnr(dname, 32, 'kplane', True)\n",
    "    naive_double = get_train_psnr(dname, 64, 'kplane', False)\n",
    "    tilted_double = get_train_psnr(dname, 64, 'kplane', True)\n",
    "\n",
    "    if tilted > naive:\n",
    "        tilted_better_count += 1\n",
    "    if tilted_double > naive_double:\n",
    "        tilted_better_count += 1\n",
    "    if tilted > naive_double - tolerance:\n",
    "        tilted_better_than_naive_double_count += 1\n",
    "    count += 2\n",
    "\n",
    "    naive = get_train_psnr(dname, 96, 'vm', False)\n",
    "    tilted = get_train_psnr(dname, 96, 'vm', True)\n",
    "    naive_double = get_train_psnr(dname, 192, 'vm', False)\n",
    "    tilted_double = get_train_psnr(dname, 192, 'vm', True)\n",
    "\n",
    "    if tilted > naive:\n",
    "        tilted_better_count += 1\n",
    "    if tilted_double > naive_double:\n",
    "        tilted_better_count += 1\n",
    "    if tilted > naive_double - tolerance:\n",
    "        tilted_better_than_naive_double_count += 1\n",
    "    count += 2\n",
    "\n",
    "print(\"How often does TILTED improve performance?\")\n",
    "print(tilted_better_count / count* 100, \"%\")\n",
    "print()\n",
    "\n",
    "print(\"How often does adding TILTED outperform _doubling the channel count_?\")\n",
    "print(tilted_better_than_naive_double_count / (count / 2) * 100, \"%\")\n",
    "print()\n",
    "\n",
    "print(\"How many experiments have we run?\")\n",
    "print(count)\n",
    "print()\n",
    "\n",
    "print(\"What are the runtimes of each method?\")\n",
    "\n",
    "def minutes(it_per_sec: float) -> str:\n",
    "    sec = (30_000 / it_per_sec)\n",
    "    return f\"{int(sec // 60)}:{int(sec % 60)}\"\n",
    "\n",
    "print(\"32c\", minutes(train_results.filter_config_eq({\n",
    "    \"dataset_type\": \"nerfstudio\",\n",
    "    \"primary_transform_count\": None,\n",
    "    \"field.grid_type\": \"kplane\",\n",
    "    \"field.primary_channels\": 32\n",
    "})[0].metrics[\"train/iterations_per_sec\"]))\n",
    "\n",
    "print(\"32c TILTED\", minutes(train_results.filter_config_eq({\n",
    "    \"dataset_type\": \"nerfstudio\",\n",
    "    \"primary_transform_count\": 8,\n",
    "    \"field.grid_type\": \"kplane\",\n",
    "    \"field.primary_channels\": 32\n",
    "})[0].metrics[\"train/iterations_per_sec\"]))\n",
    "\n",
    "print(\"64c\", minutes(train_results.filter_config_eq({\n",
    "    \"dataset_type\": \"nerfstudio\",\n",
    "    \"primary_transform_count\": None,\n",
    "    \"field.grid_type\": \"kplane\",\n",
    "    \"field.primary_channels\": 64\n",
    "})[0].metrics[\"train/iterations_per_sec\"]))\n",
    "\n",
    "print(\"64c TILTED\", minutes(train_results.filter_config_eq({\n",
    "    \"dataset_type\": \"nerfstudio\",\n",
    "    \"primary_transform_count\": 8,\n",
    "    \"field.grid_type\": \"kplane\",\n",
    "    \"field.primary_channels\": 64\n",
    "})[0].metrics[\"train/iterations_per_sec\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1315da33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{tabular}{lcc}\n",
      "\\toprule\n",
      "Dataset & K-Plane / 2x / TILTED & VM / 2x / TILTED\\\\\n",
      "\\cmidrule(r){1-1} \\cmidrule(lr){2-2} \\cmidrule(l){3-3}\n",
      "Kitchen         & 25.95 / 26.91 / \\textbf{27.12} & 25.63 / 26.54 / \\textbf{26.90} \\\\\n",
      "Floating Tree   & 24.58 / \\textbf{25.17} / 25.06 & 24.03 / 24.70 / \\textbf{25.04} \\\\\n",
      "Poster          & 33.14 / 33.71 / \\textbf{33.79} & 32.84 / 33.49 / \\textbf{33.61} \\\\\n",
      "Redwoods2       & 23.55 / 24.08 / \\textbf{24.12} & 23.22 / 23.81 / \\textbf{23.85} \\\\\n",
      "Stump           & 26.82 / \\textbf{27.29} / 27.28 & 26.33 / 26.83 / \\textbf{26.97} \\\\\n",
      "Vegetation      & 21.62 / 22.10 / \\textbf{22.10} & 21.11 / 21.55 / \\textbf{21.73} \\\\\n",
      "Bww_Entrance    & 24.64 / \\textbf{25.06} / 24.95 & 24.22 / 24.75 / \\textbf{24.80} \\\\\n",
      "Library         & 25.24 / 25.68 / \\textbf{25.78} & 25.50 / 25.78 / \\textbf{25.84} \\\\\n",
      "Storefront      & 29.71 / \\textbf{30.12} / 29.87 & 29.15 / 29.77 / \\textbf{29.87} \\\\\n",
      "Dozer           & 22.37 / \\textbf{22.88} / 22.69 & 21.91 / \\textbf{22.46} / 22.40 \\\\\n",
      "Egypt           & 20.69 / 21.10 / \\textbf{21.12} & 20.84 / \\textbf{21.17} / 21.09 \\\\\n",
      "Person          & 24.83 / 24.93 / \\textbf{25.36} & 25.28 / 25.38 / \\textbf{25.39} \\\\\n",
      "Giannini Hall   & 20.51 / \\textbf{20.90} / 20.82 & 20.27 / \\textbf{20.64} / 20.60 \\\\\n",
      "Sculpture       & 23.20 / 23.40 / \\textbf{23.40} & 22.86 / 23.07 / \\textbf{23.28} \\\\\n",
      "Plane           & 22.75 / 23.00 / \\textbf{23.01} & 22.53 / \\textbf{22.84} / 22.74 \\\\\n",
      "Aspen           & 15.99 / 16.20 / \\textbf{16.21} & 15.98 / 16.15 / \\textbf{16.20} \\\\\n",
      "Desolation      & 22.14 / \\textbf{22.40} / 22.25 & 21.88 / 22.11 / \\textbf{22.12} \\\\\n",
      "Campanile       & 24.27 / \\textbf{24.64} / 24.37 & 23.97 / \\textbf{24.35} / 24.19 \\\\\n",
      "\\cmidrule(r){1-1} \\cmidrule(lr){2-2} \\cmidrule(l){3-3}\n",
      "Best # & 0 / 8 / 10 & 0 / 5 / 13\\\\\n",
      "\\bottomrule\n",
      "\\end{tabular}\n"
     ]
    }
   ],
   "source": [
    "lines = []\n",
    "deltas = []\n",
    "\n",
    "kplane_best_count = 0\n",
    "kplane_tilted_best_count = 0\n",
    "kplane_double_best_count = 0\n",
    "\n",
    "vm_best_count = 0\n",
    "vm_tilted_best_count = 0\n",
    "vm_double_best_count = 0\n",
    "\n",
    "for dname in dataset_names:\n",
    "    kplane = get_train_psnr(dname, 32, 'kplane', False)\n",
    "    kplane_tilted = get_train_psnr(dname, 32, 'kplane', True)\n",
    "    kplane_double = get_train_psnr(dname, 64, 'kplane', False)\n",
    "\n",
    "    vm = get_train_psnr(dname, 96, 'vm', False)\n",
    "    vm_tilted = get_train_psnr(dname, 96, 'vm', True)\n",
    "    vm_double = get_train_psnr(dname, 192, 'vm', False)\n",
    "\n",
    "    if max(kplane, kplane_tilted, kplane_double) == kplane:\n",
    "        kplane_best_count += 1\n",
    "    if max(kplane, kplane_tilted, kplane_double) == kplane_tilted:\n",
    "        kplane_tilted_best_count += 1\n",
    "    if max(kplane, kplane_tilted, kplane_double) == kplane_double:\n",
    "        kplane_double_best_count += 1\n",
    "\n",
    "    if max(vm, vm_tilted, vm_double) == vm:\n",
    "        vm_best_count += 1\n",
    "    if max(vm, vm_tilted, vm_double) == vm_tilted:\n",
    "        vm_tilted_best_count += 1\n",
    "    if max(vm, vm_tilted, vm_double) == vm_double:\n",
    "        vm_double_best_count += 1\n",
    "\n",
    "\n",
    "    lines.append(\n",
    "        f\"{dname.replace('-', ' ').title().ljust(15)}\"\n",
    "        f\" & {kplane:.2f}\"\n",
    "        + \" / \" + (bold if kplane_double > kplane_tilted else lambda x: x)(f\"{kplane_double:.2f}\")\n",
    "        + \" / \" + (bold if kplane_double <= kplane_tilted else lambda x: x)(f\"{kplane_tilted:.2f}\")\n",
    "        + f\" & {vm:.2f}\"\n",
    "        + \" / \" + (bold if vm_double > vm_tilted else lambda x: x)(f\"{vm_double:.2f}\")\n",
    "        + \" / \" + (bold if vm_double <= vm_tilted else lambda x: x)(f\"{vm_tilted:.2f}\")\n",
    "        + r\" \\\\\"\n",
    "    )\n",
    "    # lines.append(f\"{dname.replace('-', ' ').title().ljust(15)} & {kplane:.2f} & \\\\textbf{'{'}{kplane_tilted:.2f} (+{kplane_tilted - kplane:.2f}){'}'} & {vm:.2f} & \\\\textbf{'{'}{vm_tilted:.2f} (+{vm_tilted - vm:.2f}){'}'}\")\n",
    "    deltas.append(kplane_tilted - kplane + vm_tilted - vm)\n",
    "\n",
    "print(r\"\\begin{tabular}{lcc}\")\n",
    "print(r\"\\toprule\")\n",
    "print(r\"Dataset & K-Plane / 2x / TILTED & VM / 2x / TILTED\\\\\")\n",
    "print(r\"\\cmidrule(r){1-1} \\cmidrule(lr){2-2} \\cmidrule(l){3-3}\")\n",
    "for i in sorted(range(len(lines)), key=lambda i: -deltas[i]):\n",
    "    print(lines[i])\n",
    "print(r\"\\cmidrule(r){1-1} \\cmidrule(lr){2-2} \\cmidrule(l){3-3}\")\n",
    "print(f\"Best # & {kplane_best_count} / {kplane_double_best_count} / {kplane_tilted_best_count} & {vm_best_count} / {vm_double_best_count} / {vm_tilted_best_count}\" + r\"\\\\\")\n",
    "print(r\"\\bottomrule\")\n",
    "print(r\"\\end{tabular}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36b746b9",
   "metadata": {},
   "source": [
    "### Ablations\n",
    "\n",
    "Let's compare:\n",
    "- 4 transform, 8 transform TILTED\n",
    "- 2-phase registration vs direct optimization\n",
    "\n",
    "We'll use Lego, since that's where the ``correct'' alignment might be the most obvious."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "12589578",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{tabular}{lll}\n",
      "\\toprule\n",
      "& 8 transforms & 4 transforms\\\\\n",
      "\\midrule\\\\\n",
      "Two phase & 34.35{\\scriptsize$\\pm$0.07} & 34.19{\\scriptsize$\\pm$0.22}  \\\\\n",
      "w/o & 33.95{\\scriptsize$\\pm$0.15} & 33.83{\\scriptsize$\\pm$0.08}  \\\\\n",
      "\\bottomrule\n",
      "\\end{tabular}\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    r\"\\begin{tabular}{lll}\",\n",
    "    r\"\\toprule\",\n",
    "    sep=\"\\n\"\n",
    ")\n",
    "print(r\"& 8 transforms & 4 transforms\\\\\")\n",
    "print(r\"\\cmidrule(lr){2-2} \\cmidrule(l){3-3}\")\n",
    "for two_phase_enabled in (True, False):\n",
    "    print(\"Two phase\" if two_phase_enabled else \"w/o\", end=\" \")\n",
    "    for primary_transform_count in (8, 4):\n",
    "        print(\n",
    "            \"&\",\n",
    "            get_psnr(\n",
    "                eval_results.filter_config_eq(\n",
    "                    {\n",
    "                        \"primary_transform_count\": primary_transform_count,\n",
    "                        \"field.grid_type\": \"kplane\",\n",
    "                        \"field.primary_channels\": 64,\n",
    "                    }\n",
    "                ).filter_config_neq({\n",
    "                    # Since TILTED is transform invariant, we can include both rotated and axis-aligned scenes.\n",
    "                    # \"render_config.global_rotate_seed\": None,\n",
    "                }).filter(\n",
    "                    lambda p: p.config[\"dataset_path\"].name == \"lego\"\n",
    "                ).filter(\n",
    "                    lambda p: p.path.name.endswith(\"-bottleneck\") if two_phase_enabled else not p.path.name.endswith(\"-bottleneck\")\n",
    "                ).assert_length(6)\n",
    "            ),\n",
    "            end=\" \"\n",
    "        )\n",
    "    print(r\" \\\\\")\n",
    "\n",
    "print(\n",
    "    r\"\\bottomrule\",\n",
    "    r\"\\end{tabular}\",\n",
    "    sep=\"\\n\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aa100f0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe28a81b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36bf582e",
   "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}