{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nested schemas: point-cloud statistics\n",
    "\n",
    "<div class=\"admonition note alert alert-warning\">\n",
    "<p class=\"admonition-title\">Note</p>\n",
    "\n",
    "All examples are expected to run from the `examples/<example_name>` directory of the [Tesseract-Torch repository](https://github.com/pasteurlabs/tesseract-torch).\n",
    "</div>\n",
    "\n",
    "Real Tesseracts rarely take a flat list of arrays. Inputs and outputs are usually grouped into sub-models, some fields are differentiable arrays and others are plain settings such as strings and numbers. Tesseract-Torch's `apply_tesseract()` mirrors that structure: you pass nested dictionaries in, you get nested dictionaries out, and autograd flows through whichever leaves are differentiable.\n",
    "\n",
    "In this example, you will learn how to:\n",
    "1. Build a Tesseract whose inputs and outputs are nested Pydantic models.\n",
    "1. Call it through `apply_tesseract()` with nested dictionaries that mix tensors and plain Python values.\n",
    "1. Take reverse-mode gradients through a nested output field with `.backward()`.\n",
    "1. Do the same in forward mode with `torch.autograd.forward_ad`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Build + serve example Tesseract\n",
    "\n",
    "The example Tesseract summarizes a weighted point cloud. Its input schema has two sub-models: `cloud` holds the differentiable arrays (`positions` and `weights`), and `settings` holds plain options (`label` and `scale`). Its output schema is nested as well: `statistics` carries the differentiable results, `summary` echoes back non-differentiable bookkeeping. Here is the computation at the heart of it (see `pointstats_torch/tesseract_api.py`):\n",
    "\n",
    "```python\n",
    "def evaluate(inputs: dict) -> dict:\n",
    "    positions = inputs[\"settings\"][\"scale\"] * inputs[\"cloud\"][\"positions\"]\n",
    "    weights = inputs[\"cloud\"][\"weights\"]\n",
    "    total_weight = weights.sum()\n",
    "\n",
    "    barycenter = (weights[:, None] * positions).sum(dim=0) / total_weight\n",
    "    offsets = positions - barycenter\n",
    "    radius_of_gyration = torch.sqrt(\n",
    "        (weights * (offsets * offsets).sum(dim=1)).sum() / total_weight\n",
    "    )\n",
    "\n",
    "    return {\n",
    "        \"statistics\": {\n",
    "            \"barycenter\": barycenter,\n",
    "            \"radius_of_gyration\": radius_of_gyration,\n",
    "        },\n",
    "        \"summary\": {\n",
    "            \"label\": inputs[\"settings\"][\"label\"],\n",
    "            \"n_points\": positions.shape[0],\n",
    "        },\n",
    "    }\n",
    "```\n",
    "\n",
    "You may build the example Tesseract either via the command line, or running the cell below (you can skip running this if already built).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Build pointstats_torch Tesseract so we can use it below\n",
    "tesseract build pointstats_torch/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To interact with the Tesseract, we use the Python SDK from `tesseract_core` to load the built image and start a server container."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tesseract_core import Tesseract\n",
    "\n",
    "pointstats = Tesseract.from_image(\"pointstats_torch\")\n",
    "pointstats.serve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Invoke the Tesseract via Tesseract-Torch\n",
    "\n",
    "The input dictionary follows the schema one level at a time: `cloud` and `settings` are dictionaries themselves, and only the leaves are tensors or plain values. Let's summarize the four corners of a unit square, all with the same weight:\n",
    "\n",
    "$$\\mathbf{p}_1 = (0, 0, 0), \\quad \\mathbf{p}_2 = (1, 0, 0), \\quad \\mathbf{p}_3 = (1, 1, 0), \\quad \\mathbf{p}_4 = (0, 1, 0)$$\n",
    "\n",
    "`settings` takes a string and a float, which `apply_tesseract()` passes through to the Tesseract untouched."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "import torch\n",
    "\n",
    "from tesseract_torch import apply_tesseract\n",
    "\n",
    "cloud = {\n",
    "    \"positions\": torch.tensor(\n",
    "        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]]\n",
    "    ),\n",
    "    \"weights\": torch.tensor([1.0, 1.0, 1.0, 1.0]),\n",
    "}\n",
    "settings = {\"label\": \"unit square\", \"scale\": 1.0}\n",
    "\n",
    "outputs = apply_tesseract(pointstats, inputs={\"cloud\": cloud, \"settings\": settings})\n",
    "pprint(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output is nested like the schema: `outputs[\"statistics\"]` holds tensors, `outputs[\"summary\"]` holds the label and point count as plain Python values. As expected, the barycenter is $(0.5, 0.5, 0)$ and every corner sits at distance $\\sqrt{0.5} \\approx 0.707$ from it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Gradients through a nested output\n",
    "\n",
    "Because the positions are a differentiable leaf, we can take gradients of any function of the nested outputs with respect to them. Autodifferentiation is dispatched to the underlying Tesseract's `vector_jacobian_product` endpoint, which only sees the paths that actually need gradients (`cloud.positions` here, not `settings.scale`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reverse-mode AD via .backward()\n",
    "cloud_grad = {\n",
    "    \"positions\": torch.tensor(\n",
    "        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]],\n",
    "        requires_grad=True,\n",
    "    ),\n",
    "    \"weights\": torch.tensor([1.0, 1.0, 1.0, 1.0]),\n",
    "}\n",
    "\n",
    "result = apply_tesseract(pointstats, inputs={\"cloud\": cloud_grad, \"settings\": settings})\n",
    "result[\"statistics\"][\"barycenter\"].sum().backward()\n",
    "print(\"Reverse-mode gradient of cloud['positions']:\")\n",
    "print(cloud_grad[\"positions\"].grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each point contributes its weight divided by the total weight to every coordinate of the barycenter, so with four equal weights the gradient is $0.25$ everywhere.\n",
    "\n",
    "Weights are differentiable too. Let's make them uneven and ask how the radius of gyration responds to each weight, this time via `torch.autograd.grad`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cloud_w = {\n",
    "    \"positions\": torch.tensor(\n",
    "        [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]]\n",
    "    ),\n",
    "    \"weights\": torch.tensor([4.0, 1.0, 1.0, 1.0], requires_grad=True),\n",
    "}\n",
    "\n",
    "result = apply_tesseract(pointstats, inputs={\"cloud\": cloud_w, \"settings\": settings})\n",
    "(grad_w,) = torch.autograd.grad(\n",
    "    result[\"statistics\"][\"radius_of_gyration\"], cloud_w[\"weights\"]\n",
    ")\n",
    "print(\"d(radius_of_gyration) / d(weights):\")\n",
    "print(grad_w)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Forward-mode AD works the same way through `torch.autograd.forward_ad`, dispatching to the Tesseract's `jacobian_vector_product` endpoint. Here we push a tangent of ones through the positions and read the tangent off the nested `barycenter` output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.autograd.forward_ad as fwAD\n",
    "\n",
    "tangent = torch.ones_like(cloud[\"positions\"])\n",
    "\n",
    "with fwAD.dual_level():\n",
    "    cloud_dual = {\n",
    "        \"positions\": fwAD.make_dual(cloud[\"positions\"], tangent),\n",
    "        \"weights\": cloud[\"weights\"],\n",
    "    }\n",
    "    result = apply_tesseract(\n",
    "        pointstats, inputs={\"cloud\": cloud_dual, \"settings\": settings}\n",
    "    )\n",
    "    _, jvp_result = fwAD.unpack_dual(result[\"statistics\"][\"barycenter\"])\n",
    "\n",
    "print(\"Forward-mode JVP of the barycenter:\")\n",
    "print(jvp_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Shifting every point by $(1, 1, 1)$ shifts the barycenter by $(1, 1, 1)$, which is exactly what the JVP reports."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step N+1: Clean-up and conclusions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we kept the Tesseract alive using `.serve()`, we need to manually stop it using `.teardown()` to avoid leaking resources.\n",
    "\n",
    "This is not necessary when using `Tesseract` in a `with` statement, as it will automatically clean up when the context is exited."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pointstats.teardown()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's it!\n",
    "Nested schemas need no special handling on the PyTorch side: mirror the schema with nested dictionaries, and autograd follows the differentiable leaves wherever they sit in the tree."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "science",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}