notebooks/core_05_volumes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Warp Core Tutorial: Volumes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install warp-lang matplotlib pyglet usd-core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warp as wp\n",
    "\n",
    "wp.config.quiet = True\n",
    "\n",
    "# Explicitly initializing Warp is not necessary but\n",
    "# we do it here to ensure everything is good to go.\n",
    "wp.init()\n",
    "\n",
    "# Everything else is solely to visualize the results.\n",
    "import IPython\n",
    "import matplotlib\n",
    "import matplotlib.animation\n",
    "import matplotlib.pyplot\n",
    "import pyglet\n",
    "\n",
    "import warp.render\n",
    "\n",
    "matplotlib.pyplot.rc(\"animation\", html=\"jshtml\")\n",
    "pyglet.options[\"headless\"] = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Surfacing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Compute\"\"\"\n",
    "\n",
    "@wp.func\n",
    "def sdf_create_box(\n",
    "    pos: wp.vec3,\n",
    "    size: wp.vec3,\n",
    "):\n",
    "    \"\"\"Creates a SDF box primitive.\"\"\"\n",
    "    # https://iquilezles.org/articles/distfunctions\n",
    "    q = wp.vec3(\n",
    "        wp.abs(pos[0]) - size[0],\n",
    "        wp.abs(pos[1]) - size[1],\n",
    "        wp.abs(pos[2]) - size[2],\n",
    "    )\n",
    "    qp = wp.vec3(wp.max(q[0], 0.0), wp.max(q[1], 0.0), wp.max(q[2], 0.0))\n",
    "    return wp.length(qp) + wp.min(wp.max(q[0], wp.max(q[1], q[2])), 0.0)\n",
    "\n",
    "\n",
    "@wp.func\n",
    "def sdf_create_torus(\n",
    "    pos: wp.vec3,\n",
    "    major_radius: float,\n",
    "    minor_radius: float,\n",
    "):\n",
    "    \"\"\"Creates a SDF torus primitive.\"\"\"\n",
    "    # https://iquilezles.org/articles/distfunctions\n",
    "    q = wp.vec2(wp.length(wp.vec2(pos[0], pos[2])) - major_radius, pos[1])\n",
    "    return wp.length(q) - minor_radius\n",
    "\n",
    "\n",
    "@wp.func\n",
    "def sdf_translate(\n",
    "    pos: wp.vec3,\n",
    "    offset: wp.vec3,\n",
    "):\n",
    "    \"\"\"Translates a SDF position vector with an offset.\"\"\"\n",
    "    return pos - offset\n",
    "\n",
    "\n",
    "@wp.func\n",
    "def sdf_rotate(\n",
    "    pos: wp.vec3,\n",
    "    angles: wp.vec3,\n",
    "):\n",
    "    \"\"\"Rotates a SDF position vector using Euler angles.\"\"\"\n",
    "    rot = wp.quat_rpy(\n",
    "        wp.radians(angles[0]),\n",
    "        wp.radians(angles[1]),\n",
    "        wp.radians(angles[2]),\n",
    "    )\n",
    "    return wp.quat_rotate_inv(rot, pos)\n",
    "\n",
    "\n",
    "@wp.func\n",
    "def sdf_smooth_min(\n",
    "    a: float,\n",
    "    b: float,\n",
    "    radius: float,\n",
    "):\n",
    "    \"\"\"Creates a SDF torus primitive.\"\"\"\n",
    "    # https://iquilezles.org/articles/smin\n",
    "    h = wp.max(radius - wp.abs(a - b), 0.0) / radius\n",
    "    return wp.min(a, b) - h * h * h * radius * (1.0 / 6.0)\n",
    "\n",
    "\n",
    "@wp.kernel(enable_backward=False)\n",
    "def make_field(\n",
    "    torus_altitude: float,\n",
    "    torus_major_radius: float,\n",
    "    torus_minor_radius: float,\n",
    "    smooth_min_radius: float,\n",
    "    dim: int,\n",
    "    time: float,\n",
    "    out_data: wp.array3d(dtype=float),\n",
    "):\n",
    "    \"\"\"Kernel to generate a SDF volume based on primitives.\"\"\"\n",
    "    i, j, k = wp.tid()\n",
    "\n",
    "    # Retrieve the position of the current cell in a normalized [-1, 1] range\n",
    "    # for each dimension.\n",
    "    pos = wp.vec3(\n",
    "        2.0 * ((float(i) + 0.5) / float(dim)) - 1.0,\n",
    "        2.0 * ((float(j) + 0.5) / float(dim)) - 1.0,\n",
    "        2.0 * ((float(k) + 0.5) / float(dim)) - 1.0,\n",
    "    )\n",
    "\n",
    "    box = sdf_create_box(\n",
    "        sdf_translate(pos, wp.vec3(0.0, -0.7, 0.0)),\n",
    "        wp.vec3(0.9, 0.3, 0.9),\n",
    "    )\n",
    "    torus = sdf_create_torus(\n",
    "        sdf_rotate(\n",
    "            sdf_translate(pos, wp.vec3(0.0, torus_altitude, 0.0)),\n",
    "            wp.vec3(wp.sin(time) * 90.0, wp.cos(time) * 45.0, 0.0),\n",
    "        ),\n",
    "        torus_major_radius,\n",
    "        torus_minor_radius,\n",
    "    )\n",
    "    out_data[i, j, k] = sdf_smooth_min(box, torus, smooth_min_radius)\n",
    "\n",
    "\n",
    "\"\"\"Initialization\"\"\"\n",
    "\n",
    "# Resolution of the rendered image.\n",
    "resolution = (512, 384)\n",
    "\n",
    "# Number of frames to run the sample for.\n",
    "num_frames = 120\n",
    "\n",
    "# Number of frames per second.\n",
    "fps = 60\n",
    "\n",
    "dim = 64\n",
    "max_verts = int(1e6)\n",
    "max_tris = int(1e6)\n",
    "\n",
    "torus_altitude = -0.5\n",
    "torus_major_radius = 0.5\n",
    "torus_minor_radius = 0.1\n",
    "smooth_min_radius = 0.5\n",
    "\n",
    "field = wp.zeros((dim, dim, dim), dtype=float)\n",
    "mc = wp.MarchingCubes(dim, dim, dim, max_verts, max_tris)\n",
    "\n",
    "# Camera settings.\n",
    "camera_pos = (32.0, 32.0, 150.0)\n",
    "camera_front = (0.0, -0.2, -1.0)\n",
    "\n",
    "# Create a headless OpenGL renderer for our scene.\n",
    "renderer = warp.render.OpenGLRenderer(\n",
    "    fps=fps,\n",
    "    screen_width=resolution[0],\n",
    "    screen_height=resolution[1],\n",
    "    camera_pos=camera_pos,\n",
    "    camera_front=camera_front,\n",
    "    far_plane=200.0,\n",
    "    draw_grid=False,\n",
    "    draw_axis=False,\n",
    "    vsync=True,\n",
    "    headless=True,\n",
    ")\n",
    "\n",
    "# Buffer storing the pixels data to visualize the resulting 3D render.\n",
    "image = wp.empty(shape=(resolution[1], resolution[0], 3), dtype=float)\n",
    "\n",
    "\n",
    "\"\"\"Evaluation\"\"\"\n",
    "\n",
    "renders = []\n",
    "for frame in range(num_frames):\n",
    "    wp.launch(\n",
    "        make_field,\n",
    "        dim=field.shape,\n",
    "        inputs=(\n",
    "            torus_altitude,\n",
    "            torus_major_radius,\n",
    "            torus_minor_radius,\n",
    "            smooth_min_radius,\n",
    "            dim,\n",
    "            frame / fps,\n",
    "        ),\n",
    "        outputs=(field,),\n",
    "    )\n",
    "\n",
    "    mc.surface(field, 0.0)\n",
    "\n",
    "    # Use the OpenGL renderer to store an image representing the 3D scene at\n",
    "    # the current frame.\n",
    "    renderer.begin_frame(frame / num_frames)\n",
    "    renderer.render_mesh(\n",
    "        \"surface\",\n",
    "        mc.verts.numpy(),\n",
    "        mc.indices.numpy(),\n",
    "        colors=((0.35, 0.55, 0.9),) * len(mc.verts),\n",
    "        update_topology=True,\n",
    "    )\n",
    "    renderer.end_frame()\n",
    "\n",
    "    # Store the resulting render on host memory.\n",
    "    renderer.get_pixels(image, split_up_tiles=False, mode=\"rgb\")\n",
    "    renders.append(wp.clone(image, device=\"cpu\", pinned=True))\n",
    "\n",
    "# Ensure that all the kernel launches and copies to CPU have finished.\n",
    "wp.synchronize()\n",
    "\n",
    "\n",
    "\"\"\"Visualization in Matplotlib\"\"\"\n",
    "\n",
    "# Set-up Matplotlib.\n",
    "plot_fig = matplotlib.pyplot.figure(figsize=resolution, dpi=1.0)\n",
    "plot_fig.subplots_adjust(left=0, bottom=0, right=1, top=1)\n",
    "plot_img = matplotlib.pyplot.imshow(renders[0], animated=True)\n",
    "plot_img.axes.set_axis_off()\n",
    "\n",
    "# Run Matplotlib's animation.\n",
    "plot_anim = matplotlib.animation.FuncAnimation(\n",
    "    plot_fig,\n",
    "    lambda frame: plot_img.set_data(renders[frame]),\n",
    "    frames=num_frames,\n",
    "    interval=(1.0 / fps) * 1000.0,\n",
    ")\n",
    "\n",
    "# Display the result.\n",
    "IPython.display.display(plot_anim)\n",
    "matplotlib.pyplot.close()"
   ]
  }
 ],
 "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}