{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# HyperLux: Simple Batch\n",
    "\n",
    "Render **1 HSSD scene × 1,024 cameras** at **512×512**. The first run downloads scenes and compiles native code.\n",
    "\n",
    "With the OptiX backend, camera inputs and rendered frames stay resident on the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%pip install -q hyperlux torch matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import hyperlux\n",
    "from hyperlux.camera import RGB\n",
    "\n",
    "BATCH_SIZE, SIZE = 1024, 512\n",
    "config = hyperlux.Config(width=SIZE, height=SIZE, batch_size=BATCH_SIZE)\n",
    "dataset = hyperlux.datasets.procthor.Train200()\n",
    "renderer = hyperlux.HyperLux(config=config)\n",
    "scene = hyperlux.load_scene(dataset[0].reference)\n",
    "renderer.init(scene)\n",
    "\n",
    "device = torch.device(\"cuda\" if renderer.backend == \"optix\" else \"cpu\")\n",
    "stream = torch.cuda.current_stream(device).cuda_stream if device.type == \"cuda\" else 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "Supply camera poses as float32 `(batch, 4, 3)` ray bases: **position, top-left ray, horizontal increment, vertical increment**, in world coordinates (+X forward, +Y left, +Z up). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "yaw = torch.arange(BATCH_SIZE, device=device) * (2 * torch.pi / BATCH_SIZE)\n",
    "forward = torch.stack((yaw.cos(), yaw.sin(), torch.zeros_like(yaw)), dim=1)\n",
    "right = torch.stack((yaw.sin(), -yaw.cos(), torch.zeros_like(yaw)), dim=1)\n",
    "up = torch.tensor([0., 0., 1.], device=device).expand_as(forward)\n",
    "positions = torch.as_tensor(hyperlux.FreeSpace(scene).sample(BATCH_SIZE, seed=0), device=device)\n",
    "camera_bases = torch.stack((positions, forward - right + up, 2 * right, -2 * up), dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "Render RGBA frames into a `(cameras, height, width, 4)` tensor. Rerun this cell after updating `camera_bases`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "frames = torch.empty((BATCH_SIZE, SIZE, SIZE, 4), dtype=torch.uint8, device=device)\n",
    "renderer.set_cameras(camera_bases, stream=stream)\n",
    "renderer.render()\n",
    "renderer.outputs[RGB()].read(out=frames, stream=stream)\n",
    "print(f\"{frames.shape}: {frames.dtype} on {frames.device}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "Preview 64 views in four rows of sixteen. Only these previews are copied to the CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib_inline.backend_inline import set_matplotlib_formats\n",
    "set_matplotlib_formats(\"jpeg\", bbox_inches=None, pil_kwargs={\"quality\": 80, \"optimize\": True})\n",
    "\n",
    "preview = frames[::BATCH_SIZE // 64, :, :, :3].cpu()\n",
    "fig, axes = plt.subplots(4, 16, figsize=(16, 4), dpi=100, squeeze=False)\n",
    "for ax, image in zip(axes.flat, preview):\n",
    "    ax.imshow(image)\n",
    "    ax.axis(\"off\")\n",
    "fig.subplots_adjust(left=0, right=1, bottom=0, top=1, wspace=0.02, hspace=0.02)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "kernelspec": {
   "display_name": ".venv (3.12.3)",
   "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": 5
}