{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working with Grids\n",
    "\n",
    "Grids are the foundation of iso-surface extraction. They define **where** in 3D space we sample scalar values, and **store** those values for the extraction algorithms.\n",
    "\n",
    "`isoext` provides two grid types:\n",
    "- **UniformGrid** — Dense regular lattice, simple and fast\n",
    "- **SparseGrid** — Only stores active cells, memory-efficient for large sparse domains\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-05T04:10:48.502029Z",
     "iopub.status.busy": "2026-07-05T04:10:48.501944Z",
     "iopub.status.idle": "2026-07-05T04:10:49.131346Z",
     "shell.execute_reply": "2026-07-05T04:10:49.130975Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import isoext\n",
    "from isoext import viewer\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## UniformGrid\n",
    "\n",
    "A `UniformGrid` divides a 3D bounding box into a regular lattice of cells.\n",
    "\n",
    "### Creating a Grid\n",
    "\n",
    "```python\n",
    "UniformGrid(shape, aabb_min=[-1,-1,-1], aabb_max=[1,1,1])\n",
    "```\n",
    "\n",
    "- `shape`: Number of cells in each dimension `[nx, ny, nz]`\n",
    "- `aabb_min` / `aabb_max`: Axis-aligned bounding box corners\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-05T04:10:49.133071Z",
     "iopub.status.busy": "2026-07-05T04:10:49.132935Z",
     "iopub.status.idle": "2026-07-05T04:10:49.263221Z",
     "shell.execute_reply": "2026-07-05T04:10:49.262606Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cells: 250,047\n",
      "Points: 262,144\n"
     ]
    }
   ],
   "source": [
    "# Create a 64³ grid spanning [-1, 1]³ (default bounds)\n",
    "grid = isoext.UniformGrid([64, 64, 64])\n",
    "\n",
    "print(f\"Cells: {grid.get_num_cells():,}\")\n",
    "print(f\"Points: {grid.get_num_points():,}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting Point Positions\n",
    "\n",
    "`get_points()` returns the 3D coordinates of all grid points as a PyTorch tensor. This is essential for computing your scalar field values.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-05T04:10:49.277971Z",
     "iopub.status.busy": "2026-07-05T04:10:49.277838Z",
     "iopub.status.idle": "2026-07-05T04:10:49.388797Z",
     "shell.execute_reply": "2026-07-05T04:10:49.387794Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape: torch.Size([64, 64, 64, 3])\n",
      "Device: cuda:0\n",
      "Dtype: torch.float32\n",
      "\n",
      "Min corner: tensor([-1., -1., -1.], device='cuda:0')\n",
      "Max corner: tensor([1., 1., 1.], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "points = grid.get_points()\n",
    "\n",
    "print(f\"Shape: {points.shape}\")  # (nx, ny, nz, 3)\n",
    "print(f\"Device: {points.device}\")\n",
    "print(f\"Dtype: {points.dtype}\")\n",
    "\n",
    "# Points range from aabb_min to aabb_max\n",
    "print(f\"\\nMin corner: {points[0, 0, 0]}\")\n",
    "print(f\"Max corner: {points[-1, -1, -1]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting Values\n",
    "\n",
    "`set_values()` stores your scalar field on the grid. The values can come from **anywhere** — you're not limited to the built-in SDF utilities.\n",
    "\n",
    "The iso-surface is extracted where values equal zero:\n",
    "- **Negative** values → inside the surface\n",
    "- **Positive** values → outside the surface\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-05T04:10:49.390414Z",
     "iopub.status.busy": "2026-07-05T04:10:49.390325Z",
     "iopub.status.idle": "2026-07-05T04:10:49.499456Z",
     "shell.execute_reply": "2026-07-05T04:10:49.499222Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sphere: 9,168 vertices\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"420\"\n",
       "            src=\"_static/viser/index.html?playbackPath=../scenes/a2ce8dd973f21c15.viser\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            style=\"border: 1px solid #8888; border-radius: 4px;\"\n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x766a2ed96630>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example 1: Simple sphere using raw PyTorch\n",
    "# No SDF classes needed — just compute distance from origin minus radius\n",
    "points = grid.get_points()\n",
    "radius = 0.7\n",
    "values = points.norm(dim=-1) - radius  # Signed distance to sphere\n",
    "\n",
    "grid.set_values(values)\n",
    "v, f = isoext.marching_cubes(grid)\n",
    "print(f\"Sphere: {v.shape[0]:,} vertices\")\n",
    "viewer.embed(v, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-05T04:10:49.500463Z",
     "iopub.status.busy": "2026-07-05T04:10:49.500342Z",
     "iopub.status.idle": "2026-07-05T04:10:49.571880Z",
     "shell.execute_reply": "2026-07-05T04:10:49.571593Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gyroid: 38,760 vertices\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"420\"\n",
       "            src=\"_static/viser/index.html?playbackPath=../scenes/fa78756a8cc02f13.viser\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            style=\"border: 1px solid #8888; border-radius: 4px;\"\n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x766a2c0e0b30>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example 2: Gyroid — a triply periodic minimal surface\n",
    "points = grid.get_points()\n",
    "x, y, z = points[..., 0], points[..., 1], points[..., 2]\n",
    "scale = 6.0\n",
    "\n",
    "gyroid = (\n",
    "    torch.sin(scale * x) * torch.cos(scale * y) +\n",
    "    torch.sin(scale * y) * torch.cos(scale * z) +\n",
    "    torch.sin(scale * z) * torch.cos(scale * x)\n",
    ")\n",
    "\n",
    "grid.set_values(gyroid)\n",
    "v, f = isoext.marching_cubes(grid)\n",
    "print(f\"Gyroid: {v.shape[0]:,} vertices\")\n",
    "viewer.embed(v, f, color=\"gold\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Neural Networks\n",
    "\n",
    "Since `get_points()` returns a standard PyTorch tensor on CUDA, you can feed it directly to a neural network:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-05T04:10:49.573266Z",
     "iopub.status.busy": "2026-07-05T04:10:49.573130Z",
     "iopub.status.idle": "2026-07-05T04:10:49.574663Z",
     "shell.execute_reply": "2026-07-05T04:10:49.574415Z"
    }
   },
   "outputs": [],
   "source": [
    "# Pseudocode for neural SDF extraction:\n",
    "#\n",
    "# grid = isoext.UniformGrid([128, 128, 128])\n",
    "# points = grid.get_points()  # (128, 128, 128, 3)\n",
    "#\n",
    "# # Reshape for batch processing\n",
    "# points_flat = points.reshape(-1, 3)  # (N, 3)\n",
    "#\n",
    "# # Query your neural network\n",
    "# with torch.no_grad():\n",
    "#     values_flat = model(points_flat)  # (N, 1) or (N,)\n",
    "#\n",
    "# # Reshape back and set on grid\n",
    "# values = values_flat.reshape(points.shape[:-1])\n",
    "# grid.set_values(values)\n",
    "#\n",
    "# vertices, faces = isoext.marching_cubes(grid)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SparseGrid\n",
    "\n",
    "For large domains where the surface occupies only a small region, `SparseGrid` saves memory by only storing active cells.\n",
    "\n",
    "This is useful for:\n",
    "- High-resolution extraction in large scenes\n",
    "- Adaptive refinement workflows\n",
    "- Memory-constrained applications\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-05T04:10:49.575284Z",
     "iopub.status.busy": "2026-07-05T04:10:49.575220Z",
     "iopub.status.idle": "2026-07-05T04:10:49.578239Z",
     "shell.execute_reply": "2026-07-05T04:10:49.577895Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initially: 0 active cells\n",
      "After adding: 1000 active cells\n"
     ]
    }
   ],
   "source": [
    "# Create a sparse grid with the same logical resolution\n",
    "sparse_grid = isoext.SparseGrid([64, 64, 64])\n",
    "\n",
    "print(f\"Initially: {sparse_grid.get_num_cells()} active cells\")\n",
    "\n",
    "# Add cells near the surface we want to extract\n",
    "# Cell indices are linearized: idx = x + y*nx + z*nx*ny\n",
    "cell_indices = torch.arange(0, 1000, device=\"cuda\", dtype=torch.int32)\n",
    "sparse_grid.add_cells(cell_indices)\n",
    "\n",
    "print(f\"After adding: {sparse_grid.get_num_cells()} active cells\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SparseGrid Workflow\n",
    "\n",
    "With `SparseGrid`, you work with cells explicitly:\n",
    "\n",
    "1. **Add cells** where you expect the surface\n",
    "2. **Get points** for those cells\n",
    "3. **Compute and set values**\n",
    "4. **Optionally filter** to cells that actually cross the surface\n",
    "5. **Extract** the mesh\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-05T04:10:49.579261Z",
     "iopub.status.busy": "2026-07-05T04:10:49.579174Z",
     "iopub.status.idle": "2026-07-05T04:10:54.220111Z",
     "shell.execute_reply": "2026-07-05T04:10:54.219479Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing 10738 chunks\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Active cells: 2416778\n",
      "Extracted: 2,416,776 vertices\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"420\"\n",
       "            src=\"_static/viser/index.html?playbackPath=../scenes/db1dbb3f795d929f.viser\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            style=\"border: 1px solid #8888; border-radius: 4px;\"\n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x766a2ed96930>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Complete SparseGrid example\n",
    "sparse_grid = isoext.SparseGrid([1024, 1024, 1024])\n",
    "\n",
    "# For this demo, we'll use get_potential_cell_indices to get all possible cells in chunks\n",
    "# In practice, you'd use spatial hashing or other methods to find relevant cells\n",
    "# If you know the active indices, just set them using sparse_grid.add_cells(active)\n",
    "chunks = sparse_grid.get_potential_cell_indices(100000)\n",
    "print(f\"Processing {len(chunks)} chunks\")\n",
    "\n",
    "for chunk in chunks:\n",
    "    # Get corner positions for these cells\n",
    "    points = sparse_grid.get_points_by_cell_indices(chunk)  # (N, 8, 3)\n",
    "    \n",
    "    # Compute values at corners (sphere SDF)\n",
    "    values = points.norm(dim=-1) - 0.7  # (N, 8)\n",
    "    \n",
    "    # Filter to cells that cross the surface (have both + and - values)\n",
    "    active = sparse_grid.filter_cell_indices(chunk, values, level=0.0)\n",
    "    \n",
    "    if active is not None and active.numel() > 0:\n",
    "        sparse_grid.add_cells(active)\n",
    "\n",
    "print(f\"Active cells: {sparse_grid.get_num_cells()}\")\n",
    "\n",
    "# Recompute values for the cells we're keeping\n",
    "points = sparse_grid.get_points()\n",
    "values = points.norm(dim=-1) - 0.7\n",
    "# Note: set_values expects (num_cells, 8) tensor\n",
    "# We set values for all currently active cells\n",
    "sparse_grid.set_values(values)\n",
    "\n",
    "v, f = isoext.marching_cubes(sparse_grid)\n",
    "print(f\"Extracted: {v.shape[0]:,} vertices\")\n",
    "\n",
    "# Embed a coarser preview: the full-resolution mesh would weigh ~30 MB as a\n",
    "# static scene file, which is too heavy for a documentation page.\n",
    "preview = isoext.UniformGrid([128, 128, 128])\n",
    "preview.set_values(preview.get_points().norm(dim=-1) - 0.7)\n",
    "v_lo, f_lo = isoext.marching_cubes(preview)\n",
    "viewer.embed(v_lo, f_lo, color=\"steelblue\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "**Key points:**\n",
    "- Use `get_points()` to get coordinates for computing your scalar field\n",
    "- Values can come from anywhere: math formulas, SDFs, neural networks, simulations\n",
    "- Negative values = inside, positive = outside, zero = surface\n",
    "- Use `UniformGrid` for most cases, `SparseGrid` for large sparse domains\n",
    "\n",
    "**Next:** [Marching Cubes](marching_cubes.ipynb) or [Dual Contouring](dual_contouring.ipynb)\n"
   ]
  }
 ],
 "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
