notebooks/core_03_points.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Warp Core Tutorial: Points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install warp-lang matplotlib pyglet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "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\n",
    "\n",
    "import sys\n",
    "if \"google.colab\" in sys.modules:\n",
    "    print(\n",
    "        \"Rendering OpenGL instances with `glDrawElementsInstanced` \"\n",
    "        \"is *extremely* slow in Google Colab, so it may take a few minutes \"\n",
    "        \"to render the snippets from this notebook.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Point Queries\n",
    "\n",
    "When dealing with collections of points, a common operation is to query the neighbouring points to a given position, which can quickly become expensive as the number of points grow.\n",
    "\n",
    "To speed-up these queries, Warp exposes a hash grid spatial structure that can be initialized using `wp.HashGrid()` and then rebuilt at each step based on the current point position with a call to the `wp.HashGrid.build()` method.\n",
    "\n",
    "Instances of `wp.HashGrid` can then be passed to kernels using a `wp.uint64` integer representing their unique IDs, and the built-in `wp.hash_grid_query()` can be called there to iterate over the closest points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Compute\"\"\"\n",
    "\n",
    "@wp.func\n",
    "def contact_force(\n",
    "    n: wp.vec3,\n",
    "    v: wp.vec3,\n",
    "    c: float,\n",
    "    k_n: float,\n",
    "    k_d: float,\n",
    "    k_f: float,\n",
    "    k_mu: float,\n",
    ") -> wp.vec3:\n",
    "    vn = wp.dot(n, v)\n",
    "    jn = c * k_n\n",
    "    jd = min(vn, 0.0) * k_d\n",
    "\n",
    "    # contact force\n",
    "    fn = jn + jd\n",
    "\n",
    "    # friction force\n",
    "    vt = v - n * vn\n",
    "    vs = wp.length(vt)\n",
    "\n",
    "    if vs > 0.0:\n",
    "        vt = vt / vs\n",
    "\n",
    "    # Coulomb condition\n",
    "    ft = wp.min(vs * k_f, k_mu * wp.abs(fn))\n",
    "\n",
    "    # total force\n",
    "    return -n * fn - vt * ft\n",
    "\n",
    "\n",
    "@wp.kernel\n",
    "def update(\n",
    "    grid: wp.uint64,\n",
    "    particle_x: wp.array(dtype=wp.vec3),\n",
    "    particle_v: wp.array(dtype=wp.vec3),\n",
    "    particle_f: wp.array(dtype=wp.vec3),\n",
    "    radius: float,\n",
    "    k_contact: float,\n",
    "    k_damp: float,\n",
    "    k_friction: float,\n",
    "    k_mu: float,\n",
    "):\n",
    "    tid = wp.tid()\n",
    "\n",
    "    # order threads by cell\n",
    "    i = wp.hash_grid_point_id(grid, tid)\n",
    "\n",
    "    x = particle_x[i]\n",
    "    v = particle_v[i]\n",
    "\n",
    "    f = wp.vec3()\n",
    "\n",
    "    # ground contact\n",
    "    n = wp.vec3(0.0, 1.0, 0.0)\n",
    "    c = wp.dot(n, x)\n",
    "\n",
    "    cohesion_ground = 0.02\n",
    "    cohesion_particle = 0.0075\n",
    "\n",
    "    if c < cohesion_ground:\n",
    "        f = f + contact_force(n, v, c, k_contact, k_damp, 100.0, 0.5)\n",
    "\n",
    "    # particle contact\n",
    "    neighbors = wp.hash_grid_query(grid, x, radius * 5.0)\n",
    "\n",
    "    for index in neighbors:\n",
    "        if index != i:\n",
    "            # compute distance to point\n",
    "            n = x - particle_x[index]\n",
    "            d = wp.length(n)\n",
    "            err = d - radius * 2.0\n",
    "\n",
    "            if err <= cohesion_particle:\n",
    "                n = n / d\n",
    "                vrel = v - particle_v[index]\n",
    "\n",
    "                f = f + contact_force(n, vrel, err, k_contact, k_damp, k_friction, k_mu)\n",
    "\n",
    "    particle_f[i] = f\n",
    "\n",
    "\n",
    "@wp.kernel\n",
    "def integrate(\n",
    "    x: wp.array(dtype=wp.vec3),\n",
    "    v: wp.array(dtype=wp.vec3),\n",
    "    f: wp.array(dtype=wp.vec3),\n",
    "    gravity: wp.vec3,\n",
    "    dt: float,\n",
    "    inv_mass: float,\n",
    "):\n",
    "    tid = wp.tid()\n",
    "\n",
    "    v_new = v[tid] + f[tid] * inv_mass * dt + gravity * dt\n",
    "    x_new = x[tid] + v_new * dt\n",
    "\n",
    "    v[tid] = v_new\n",
    "    x[tid] = x_new\n",
    "\n",
    "\n",
    "\"\"\"Initialization\"\"\"\n",
    "\n",
    "def create_particle_grid(\n",
    "    dim_x: int,\n",
    "    dim_y: int,\n",
    "    dim_z: int,\n",
    "    lower: float,\n",
    "    radius: float,\n",
    "    jitter: float,\n",
    ") -> wp.array(dtype=wp.vec3):\n",
    "    points = np.meshgrid(\n",
    "        np.linspace(0, dim_x, dim_x),\n",
    "        np.linspace(0, dim_y, dim_y),\n",
    "        np.linspace(0, dim_z, dim_z),\n",
    "    )\n",
    "    points = np.array((points[0], points[1], points[2])).T * radius * 2.0 + np.array(lower)\n",
    "    points = points + np.random.rand(*points.shape) * radius * jitter\n",
    "    return wp.array(points.reshape((-1, 3)), dtype=wp.vec3)\n",
    "\n",
    "\n",
    "# Resolution of the rendered image.\n",
    "resolution = (512, 384)\n",
    "\n",
    "# Number of frames to run the sample for.\n",
    "num_frames = 200\n",
    "\n",
    "# Number of frames per second.\n",
    "fps = 60\n",
    "\n",
    "# Number of simulation step per frame.\n",
    "sim_substeps = 64\n",
    "\n",
    "# Time delta between two frames.\n",
    "frame_dt = 1.0 / fps\n",
    "\n",
    "# Time delta between two simulation steps.\n",
    "sim_dt = frame_dt / sim_substeps\n",
    "\n",
    "# Radius for each particle.\n",
    "point_radius = 0.1\n",
    "\n",
    "# Inverted mass for each particle.\n",
    "inv_mass = 64.0\n",
    "\n",
    "# Contact friction stiffness.\n",
    "k_contact = 8e3\n",
    "\n",
    "# Contact damping stiffness.\n",
    "k_damp = 2.0\n",
    "\n",
    "# Contact friction stiffness.\n",
    "k_friction = 1.0\n",
    "\n",
    "# Coefficient of friction for cohesive materials.\n",
    "k_mu = 1e5\n",
    "\n",
    "# Initialize the particle positions, velocities, and forces.\n",
    "points = create_particle_grid(8, 32, 8, (0.0, 0.5, 0.0), point_radius, 0.1)\n",
    "velocities = wp.array(((0.0, 0.0, 15.0),) * len(points), dtype=wp.vec3)\n",
    "forces = wp.empty_like(points)\n",
    "\n",
    "# Initialize a hash grid.\n",
    "grid = wp.HashGrid(128, 128, 128)\n",
    "\n",
    "# Cell size to use when building the grid.\n",
    "grid_cell_size = point_radius * 5.0\n",
    "\n",
    "# Camera settings.\n",
    "camera_pos = (-26.0, 6.0, 13.5)\n",
    "camera_front = (1.0, 0.0, 0.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",
    "    draw_grid=False,\n",
    "    draw_axis=False,\n",
    "    vsync=False,\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",
    "    # Build the hash grid from the current point positions.\n",
    "    grid.build(points, grid_cell_size)\n",
    "\n",
    "    for _ in range(sim_substeps):\n",
    "        # Update the forces of each point.\n",
    "        wp.launch(\n",
    "            kernel=update,\n",
    "            dim=points.shape,\n",
    "            inputs=(\n",
    "                grid.id,\n",
    "                points,\n",
    "                velocities,\n",
    "                forces,\n",
    "                point_radius,\n",
    "                k_contact,\n",
    "                k_damp,\n",
    "                k_friction,\n",
    "                k_mu,\n",
    "            ),\n",
    "        )\n",
    "\n",
    "        # Solve the point velocities and positions for the current substep.\n",
    "        wp.launch(\n",
    "            kernel=integrate,\n",
    "            dim=points.shape,\n",
    "            inputs=(\n",
    "                points,\n",
    "                velocities,\n",
    "                forces,\n",
    "                (0.0, -9.8, 0.0),\n",
    "                sim_dt,\n",
    "                inv_mass,\n",
    "            ),\n",
    "        )\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_points(\n",
    "        points=points.numpy(),\n",
    "        radius=point_radius,\n",
    "        name=\"points\",\n",
    "        colors=(0.8, 0.3, 0.2),\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
}