notebooks/core_04_meshes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Warp Core Tutorial: Meshes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install warp-lang matplotlib"
   ]
  },
  {
   "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.pyplot\n",
    "\n",
    "matplotlib.pyplot.rc(\"animation\", html=\"jshtml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mesh Queries\n",
    "\n",
    "Triangular meshes with point positions, vertex indices, and optionally velocities, can be represented in Warp with the `wp.Mesh` structure.\n",
    "\n",
    "The main purpose of this structure is to offer facilities related to querying meshes using a fast BVH acceleration structure, which is well-suited for collision detection amongst others."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A built-in commonly used for collision detection, or to create Signed Distance Fields (SDFs), is `wp.mesh_query_point()`, which allows to gather information about the closest location on a mesh from a given point in space.\n",
    "\n",
    "The data returned contains the distance of the point to that mesh location, whether it is inside or outside the mesh, as well as the primitive index encompassing that location with its UV coordinates, which allows evaluating these into a 3D position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Triangular mesh data for a logo geometry\"\"\"\n",
    "\n",
    "GEOM_POINTS = (\n",
    "    (-3.2729452  , -0.5, -0.47547257 ), (-3.0449789  , -0.5,  0.02605331 ),\n",
    "    (-2.667206   , -0.5,  0.644819   ), (-2.2112734  , -0.5,  1.1268051  ),\n",
    "    (-1.6381009  , -0.5,  1.5110911  ), (-1.0128219  , -0.5,  1.7455708  ),\n",
    "    (-0.40708268 , -0.5,  1.7976774  ), ( 0.2963562  , -0.5,  1.7455708  ),\n",
    "    ( 0.9607153  , -0.5,  1.550171   ), ( 1.559941   , -0.5,  1.289638   ),\n",
    "    ( 2.0875204  , -0.5,  0.9900252  ), ( 2.5695062  , -0.5,  0.62527907 ),\n",
    "    ( 3.1622186  , -0.5,  0.13677979 ), ( 4.197837   , -0.5,  0.7360056  ),\n",
    "    ( 3.4618316  , -0.5,  1.2375314  ), ( 2.1135736  , -0.5,  1.810704   ),\n",
    "    ( 1.2473016  , -0.5,  2.032157   ), ( 0.32892287 , -0.5,  2.2210433  ),\n",
    "    (-0.6871557  , -0.5,  2.2666366  ), (-1.3840814  , -0.5,  2.1559103  ),\n",
    "    (-2.0875204  , -0.5,  1.8823506  ), (-2.719313   , -0.5,  1.4850378  ),\n",
    "    (-3.292485   , -0.5,  0.9379186  ), (-3.6702578  , -0.5,  0.44290602 ),\n",
    "    (-4.015464   , -0.5, -0.16934645 ), (-4.197837   , -0.5, -0.5666591  ),\n",
    "    (-4.191324   , -0.5, -0.5796858  ), (-3.7614443  , -0.5, -0.93791854 ),\n",
    "    (-3.227352   , -0.5, -1.3156914  ), (-2.5108862  , -0.5, -1.7585975  ),\n",
    "    (-1.9051471  , -0.5, -2.012617   ), (-1.2407882  , -0.5, -2.1819634  ),\n",
    "    (-0.6220224  , -0.5, -2.2666368  ), (-0.02279663 , -0.5, -2.2536101  ),\n",
    "    ( 0.6089958  , -0.5, -2.0777502  ), ( 1.1496017  , -0.5, -1.849784   ),\n",
    "    ( 1.6836942  , -0.5, -1.550171   ), ( 2.1266003  , -0.5, -1.1984516  ),\n",
    "    ( 2.5108862  , -0.5, -0.8532454  ), ( 2.764906   , -0.5, -0.58619916 ),\n",
    "    ( 1.8074474  , -0.5, -0.5406059  ), ( 1.436188   , -0.5, -0.95094526 ),\n",
    "    ( 0.934662   , -0.5, -1.3612849  ), ( 0.452676   , -0.5, -1.6022776  ),\n",
    "    (-0.042336583, -0.5, -1.7585975  ), (-0.58945584 , -0.5, -1.8237306  ),\n",
    "    (-1.1886816  , -0.5, -1.7585975  ), (-1.9181738  , -0.5, -1.5436577  ),\n",
    "    (-2.4783196  , -0.5, -1.2179915  ), (-2.9016855  , -0.5, -0.8923255  ),\n",
    "    (-2.5785387  , -0.5, -0.41622198 ), (-2.0679803  , -0.5, -0.84021866 ),\n",
    "    (-1.5664544  , -0.5, -1.0616719  ), (-1.0453885  , -0.5, -1.2310182  ),\n",
    "    (-0.53083587 , -0.5, -1.2961515  ), (-0.009769917, -0.5, -1.2114781  ),\n",
    "    ( 0.37451613 , -0.5, -1.0616719  ), ( 0.7392623  , -0.5, -0.8076521  ),\n",
    "    ( 1.0063086  , -0.5, -0.5666591  ), ( 0.23773634 , -0.5,  0.110726595),\n",
    "    (-0.10095644 , -0.5, -0.44941938 ), (-0.51780915 , -0.5, -0.7685722  ),\n",
    "    (-1.0584152  , -0.5, -0.801139   ), (-1.5208611  , -0.5, -0.65133226 ),\n",
    "    (-1.8335006  , -0.5, -0.34520614 ), (-2.1656802  , -0.5, 0.40382612  ),\n",
    "    (-1.6771808  , -0.5, 0.905352    ), (-1.1235483  , -0.5, 1.2049649   ),\n",
    "    (-0.5112959  , -0.5, 1.3222047   ), ( 0.16608977 , -0.5, 1.2505581   ),\n",
    "    ( 0.7978822  , -0.5, 1.0160785   ), ( 1.397108   , -0.5, 0.6383058   ),\n",
    "    ( 1.9312005  , -0.5, 0.27355957  ), ( 2.3415399  , -0.5, -0.11723983 ),\n",
    "    ( 1.377568   , -0.5, -0.039079905), ( 0.7718289  , -0.5,  0.4754727  ),\n",
    "    ( 0.15957642 , -0.5,  0.8011388  ), (-0.61550903 , -0.5,  0.8988387  ),\n",
    "    (-1.2277615  , -0.5,  0.6187657  ), (-1.559941   , -0.5,  0.25401962 ),\n",
    "    (-2.5785387  ,  0.5, -0.41622198 ), (-2.1656802  ,  0.5,  0.40382612 ),\n",
    "    (-1.6771808  ,  0.5,  0.905352   ), (-1.1235483  ,  0.5,  1.2049649  ),\n",
    "    (-0.5112959  ,  0.5,  1.3222047  ), ( 0.16608977 ,  0.5,  1.2505581  ),\n",
    "    ( 0.7978822  ,  0.5,  1.0160785  ), ( 1.397108   ,  0.5,  0.6383058  ),\n",
    "    ( 1.9312005  ,  0.5,  0.27355957 ), ( 2.3415399  ,  0.5, -0.11723983 ),\n",
    "    ( 2.764906   ,  0.5, -0.58619916 ), ( 2.5108862  ,  0.5, -0.8532454  ),\n",
    "    ( 2.1266003  ,  0.5, -1.1984516  ), ( 1.6836942  ,  0.5, -1.550171   ),\n",
    "    ( 1.1496017  ,  0.5, -1.849784   ), ( 0.6089958  ,  0.5, -2.0777502  ),\n",
    "    (-0.02279663 ,  0.5, -2.2536101  ), (-0.6220224  ,  0.5, -2.2666368  ),\n",
    "    (-1.2407882  ,  0.5, -2.1819634  ), (-1.9051471  ,  0.5, -2.012617   ),\n",
    "    (-2.5108862  ,  0.5, -1.7585975  ), (-3.227352   ,  0.5, -1.3156914  ),\n",
    "    (-3.7614443  ,  0.5, -0.93791854 ), (-4.191324   ,  0.5, -0.5796858  ),\n",
    "    (-4.197837   ,  0.5, -0.5666591  ), (-4.015464   ,  0.5, -0.16934645 ),\n",
    "    (-3.6702578  ,  0.5,  0.44290602 ), (-3.292485   ,  0.5,  0.9379186  ),\n",
    "    (-2.719313   ,  0.5,  1.4850378  ), (-2.0875204  ,  0.5,  1.8823506  ),\n",
    "    (-1.3840814  ,  0.5,  2.1559103  ), (-0.6871557  ,  0.5,  2.2666366  ),\n",
    "    ( 0.32892287 ,  0.5,  2.2210433  ), ( 1.2473016  ,  0.5,  2.032157   ),\n",
    "    ( 2.1135736  ,  0.5,  1.810704   ), ( 3.4618316  ,  0.5,  1.2375314  ),\n",
    "    ( 4.197837   ,  0.5,  0.7360056  ), ( 3.1622186  ,  0.5,  0.13677979 ),\n",
    "    ( 2.5695062  ,  0.5,  0.62527907 ), ( 2.0875204  ,  0.5,  0.9900252  ),\n",
    "    ( 1.559941   ,  0.5,  1.289638   ), ( 0.9607153  ,  0.5,  1.550171   ),\n",
    "    ( 0.2963562  ,  0.5,  1.7455708  ), (-0.40708268 ,  0.5,  1.7976774  ),\n",
    "    (-1.0128219  ,  0.5,  1.7455708  ), (-1.6381009  ,  0.5,  1.5110911  ),\n",
    "    (-2.2112734  ,  0.5,  1.1268051  ), (-2.667206   ,  0.5,  0.644819   ),\n",
    "    (-3.0449789  ,  0.5,  0.02605331 ), (-3.2729452  ,  0.5, -0.47547257 ),\n",
    "    (-2.9016855  ,  0.5, -0.8923255  ), (-2.4783196  ,  0.5, -1.2179915  ),\n",
    "    (-1.9181738  ,  0.5, -1.5436577  ), (-1.1886816  ,  0.5, -1.7585975  ),\n",
    "    (-0.58945584 ,  0.5, -1.8237306  ), (-0.042336583,  0.5, -1.7585975  ),\n",
    "    ( 0.452676   ,  0.5, -1.6022776  ), ( 0.934662   ,  0.5, -1.3612849  ),\n",
    "    ( 1.436188   ,  0.5, -0.95094526 ), ( 1.8074474  ,  0.5, -0.5406059  ),\n",
    "    ( 1.377568   ,  0.5, -0.039079905), ( 0.7718289  ,  0.5,  0.4754727  ),\n",
    "    ( 0.15957642 ,  0.5,  0.8011388  ), (-0.61550903 ,  0.5,  0.8988387  ),\n",
    "    (-1.2277615  ,  0.5,  0.6187657  ), (-1.559941   ,  0.5,  0.25401962 ),\n",
    "    (-1.8335006  ,  0.5, -0.34520614 ), (-1.5208611  ,  0.5, -0.65133226 ),\n",
    "    (-1.0584152  ,  0.5, -0.801139   ), (-0.51780915 ,  0.5, -0.7685722  ),\n",
    "    (-0.10095644 ,  0.5, -0.44941938 ), ( 0.23773634 ,  0.5,  0.110726595),\n",
    "    ( 1.0063086  ,  0.5, -0.5666591  ), ( 0.7392623  ,  0.5, -0.8076521  ),\n",
    "    ( 0.37451613 ,  0.5, -1.0616719  ), (-0.009769917,  0.5, -1.2114781  ),\n",
    "    (-0.53083587 ,  0.5, -1.2961515  ), (-1.0453885  ,  0.5, -1.2310182  ),\n",
    "    (-1.5664544  ,  0.5, -1.0616719  ), (-2.0679803  ,  0.5, -0.84021866 ),\n",
    ")\n",
    "\n",
    "GEOM_FACE_VERTEX_INDICES = (\n",
    "    116, 117, 115, 117, 118, 115, 115, 118, 114, 118, 119, 114,\n",
    "    119, 120, 114, 114, 120, 113, 120, 121, 113, 113, 121, 112,\n",
    "    121, 122, 112, 151, 152, 150, 152, 153, 150, 153, 154, 150,\n",
    "    154, 155, 150, 150, 155, 149, 155, 156, 149, 156, 157, 149,\n",
    "    149, 157, 148, 157, 158, 148, 148, 158, 147, 158, 159, 147,\n",
    "    147, 159, 146, 159,  80, 146,  80,  81, 146, 146,  81, 145,\n",
    "     81,  82, 145, 145,  82, 144,  82,  83, 144, 144,  83, 143,\n",
    "     83,  84, 143, 143,  84, 142,  84,  85, 142,  85,  86, 142,\n",
    "    142,  86, 141,  86,  87, 141, 141,  87, 140,  87,  88, 140,\n",
    "    140,  88, 139,  88,  89, 139, 112, 122, 111, 122, 123, 111,\n",
    "    123, 124, 111, 111, 124, 110, 124, 125, 110, 110, 125, 109,\n",
    "    125, 126, 109, 109, 126, 108, 126, 127, 108, 108, 127, 107,\n",
    "    127, 128, 107, 107, 128, 106, 106, 128, 105, 128, 129, 105,\n",
    "    105, 129, 104, 104, 129, 103, 103, 129, 102, 102, 129, 101,\n",
    "    129, 130, 101, 101, 130, 100, 130, 131, 100, 131, 132, 100,\n",
    "    100, 132,  99, 132, 133,  99,  99, 133,  98,  98, 133,  97,\n",
    "    133, 134,  97,  97, 134,  96, 134, 135,  96,  96, 135,  95,\n",
    "    135, 136,  95,  95, 136,  94, 136, 137,  94,  94, 137,  93,\n",
    "    137, 138,  93,  93, 138,  92, 138, 139,  92,  92, 139,  91,\n",
    "    139,  89,  91,  89,  90,  91,  13,  14,  12,  12,  14,  11,\n",
    "     14,  15,  11,  11,  15,  10,  10,  15,   9,  15,  16,   9,\n",
    "      9,  16,   8,  16,  17,   8,   8,  17,   7,  59,  60,  58,\n",
    "     58,  60,  57,  57,  60,  56,  56,  60,  55,  60,  61,  55,\n",
    "     55,  61,  54,  54,  61,  53,  61,  62,  53,  53,  62,  52,\n",
    "     62,  63,  52,  52,  63,  51,  63,  64,  51,  51,  64,  50,\n",
    "     50,  64,  65,  64,  79,  65,  65,  79,  66,  79,  78,  66,\n",
    "     66,  78,  67,  78,  77,  67,  67,  77,  68,  77,  76,  68,\n",
    "     68,  76,  69,  69,  76,  70,  76,  75,  70,  70,  75,  71,\n",
    "     75,  74,  71,  71,  74,  72,  74,  40,  72,  72,  40,  73,\n",
    "      7,  17,   6,  17,  18,   6,   6,  18,   5,  18,  19,   5,\n",
    "      5,  19,   4,  19,  20,   4,   4,  20,   3,  20,  21,   3,\n",
    "      3,  21,   2,  21,  22,   2,   2,  22,   1,  22,  23,   1,\n",
    "     23,  24,   1,   1,  24,   0,  24,  25,   0,  25,  26,   0,\n",
    "     26,  27,   0,  27,  28,   0,   0,  28,  49,  28,  29,  49,\n",
    "     49,  29,  48,  48,  29,  47,  29,  30,  47,  47,  30,  46,\n",
    "     30,  31,  46,  46,  31,  45,  31,  32,  45,  32,  33,  45,\n",
    "     45,  33,  44,  33,  34,  44,  44,  34,  43,  34,  35,  43,\n",
    "     43,  35,  42,  35,  36,  42,  42,  36,  41,  36,  37,  41,\n",
    "     41,  37,  40,  37,  38,  40,  40,  38,  73,  38,  39,  73,\n",
    "     50,  65,  80,  80,  65,  81,  65,  66,  81,  81,  66,  82,\n",
    "     66,  67,  82,  82,  67,  83,  67,  68,  83,  83,  68,  84,\n",
    "     68,  69,  84,  84,  69,  85,  69,  70,  85,  85,  70,  86,\n",
    "     70,  71,  86,  86,  71,  87,  71,  72,  87,  87,  72,  88,\n",
    "     72,  73,  88,  88,  73,  89,  73,  39,  89,  89,  39,  90,\n",
    "     39,  38,  90,  90,  38,  91,  38,  37,  91,  91,  37,  92,\n",
    "     37,  36,  92,  92,  36,  93,  36,  35,  93,  93,  35,  94,\n",
    "     35,  34,  94,  94,  34,  95,  34,  33,  95,  95,  33,  96,\n",
    "     33,  32,  96,  96,  32,  97,  32,  31,  97,  97,  31,  98,\n",
    "     31,  30,  98,  98,  30,  99,  30,  29,  99,  99,  29, 100,\n",
    "     29,  28, 100, 100,  28, 101,  28,  27, 101, 101,  27, 102,\n",
    "     27,  26, 102, 102,  26, 103,  26,  25, 103, 103,  25, 104,\n",
    "     25,  24, 104, 104,  24, 105,  24,  23, 105, 105,  23, 106,\n",
    "     23,  22, 106, 106,  22, 107,  22,  21, 107, 107,  21, 108,\n",
    "     21,  20, 108, 108,  20, 109,  20,  19, 109, 109,  19, 110,\n",
    "     19,  18, 110, 110,  18, 111,  18,  17, 111, 111,  17, 112,\n",
    "     17,  16, 112, 112,  16, 113,  16,  15, 113, 113,  15, 114,\n",
    "     15,  14, 114, 114,  14, 115,  14,  13, 115, 115,  13, 116,\n",
    "     13,  12, 116, 116,  12, 117,  12,  11, 117, 117,  11, 118,\n",
    "     11,  10, 118, 118,  10, 119,  10,   9, 119, 119,   9, 120,\n",
    "      9,   8, 120, 120,   8, 121,   8,   7, 121, 121,   7, 122,\n",
    "      7,   6, 122, 122,   6, 123,   6,   5, 123, 123,   5, 124,\n",
    "      5,   4, 124, 124,   4, 125,   4,   3, 125, 125,   3, 126,\n",
    "      3,   2, 126, 126,   2, 127,   2,   1, 127, 127,   1, 128,\n",
    "      1,   0, 128, 128,   0, 129,   0,  49, 129, 129,  49, 130,\n",
    "     49,  48, 130, 130,  48, 131,  48,  47, 131, 131,  47, 132,\n",
    "     47,  46, 132, 132,  46, 133,  46,  45, 133, 133,  45, 134,\n",
    "     45,  44, 134, 134,  44, 135,  44,  43, 135, 135,  43, 136,\n",
    "     43,  42, 136, 136,  42, 137,  42,  41, 137, 137,  41, 138,\n",
    "     41,  40, 138, 138,  40, 139,  40,  74, 139, 139,  74, 140,\n",
    "     74,  75, 140, 140,  75, 141,  75,  76, 141, 141,  76, 142,\n",
    "     76,  77, 142, 142,  77, 143,  77,  78, 143, 143,  78, 144,\n",
    "     78,  79, 144, 144,  79, 145,  79,  64, 145, 145,  64, 146,\n",
    "     64,  63, 146, 146,  63, 147,  63,  62, 147, 147,  62, 148,\n",
    "     62,  61, 148, 148,  61, 149,  61,  60, 149, 149,  60, 150,\n",
    "     60,  59, 150, 150,  59, 151,  59,  58, 151, 151,  58, 152,\n",
    "     58,  57, 152, 152,  57, 153,  57,  56, 153, 153,  56, 154,\n",
    "     56,  55, 154, 154,  55, 155,  55,  54, 155, 155,  54, 156,\n",
    "     54,  53, 156, 156,  53, 157,  53,  52, 157, 157,  52, 158,\n",
    "     52,  51, 158, 158,  51, 159,  51,  50, 159, 159,  50,  80,\n",
    ")\n",
    "\n",
    "\n",
    "\"\"\"Compute\"\"\"\n",
    "\n",
    "@wp.func\n",
    "def bourke_color_map(\n",
    "    low: float,\n",
    "    high: float,\n",
    "    v: float,\n",
    ") -> wp.vec3:\n",
    "    r = 1.0\n",
    "    g = 1.0\n",
    "    b = 1.0\n",
    "\n",
    "    if v < low:\n",
    "        v = low\n",
    "    if v > high:\n",
    "        v = high\n",
    "\n",
    "    dv = high - low\n",
    "\n",
    "    if v < (low + 0.25 * dv):\n",
    "        r = 0.0\n",
    "        g = 4.0 * (v - low) / dv\n",
    "    elif v < (low + 0.5 * dv):\n",
    "        r = 0.0\n",
    "        b = 1.0 + 4.0 * (low + 0.25 * dv - v) / dv\n",
    "    elif v < (low + 0.75 * dv):\n",
    "        r = 4.0 * (v - low - 0.5 * dv) / dv\n",
    "        b = 0.0\n",
    "    else:\n",
    "        g = 1.0 + 4.0 * (low + 0.75 * dv - v) / dv\n",
    "        b = 0.0\n",
    "\n",
    "    return wp.vec3(r, g, b)\n",
    "\n",
    "\n",
    "@wp.kernel\n",
    "def compute_sdf_slice(\n",
    "    pixel_size: float,\n",
    "    pixel_offset: wp.vec2,\n",
    "    geom: wp.uint64,\n",
    "    sdf_band_width: float,\n",
    "    out_texture: wp.array(dtype=wp.vec3, ndim=2),\n",
    "):\n",
    "    # Retrieve the current pixel's 2D indices.\n",
    "    j, i = wp.tid()\n",
    "\n",
    "    # Compute the position in 3D space of the pixel.\n",
    "    point = wp.vec3(\n",
    "        float(j) * pixel_size + pixel_offset[0],\n",
    "        0.0,\n",
    "        float(i) * pixel_size + pixel_offset[1],\n",
    "    )\n",
    "\n",
    "    # Query the nearest location on the mesh using BVH.\n",
    "    query = wp.mesh_query_point(geom, point, max_dist=1e6)\n",
    "    if not query.result:\n",
    "        return\n",
    "\n",
    "    # Evaluate the position of the nearest location found.\n",
    "    nearest_pos = wp.mesh_eval_position(\n",
    "        geom,\n",
    "        query.face,\n",
    "        query.u,\n",
    "        query.v,\n",
    "    )\n",
    "\n",
    "    # Compute the SDF value.\n",
    "    dist = wp.length(nearest_pos - point)\n",
    "    if query.sign >= 0:\n",
    "        sdf = dist\n",
    "    else:\n",
    "        sdf = -dist\n",
    "\n",
    "    color = bourke_color_map(\n",
    "        -pixel_size * sdf_band_width * 0.5,\n",
    "        pixel_size * sdf_band_width * 0.5,\n",
    "        sdf,\n",
    "    )\n",
    "\n",
    "    # Store the colour into our output array.\n",
    "    out_texture[i, j] = color\n",
    "\n",
    "\n",
    "\"\"\"Initialization\"\"\"\n",
    "\n",
    "# Size of the grid used to slice the mesh along the up +Y axis.\n",
    "grid_size = (10.0, 5.0)\n",
    "\n",
    "# Size for each pixel for the rendered texture.\n",
    "pixel_size = 0.025\n",
    "\n",
    "# Width of the SDF's narrow-band, in voxels.\n",
    "sdf_band_width = 0.3 / pixel_size\n",
    "\n",
    "# Resolution of the rendered texture.\n",
    "resolution = (\n",
    "    int(wp.ceil(grid_size[0] / pixel_size)),\n",
    "    int(wp.ceil(grid_size[1] / pixel_size)),\n",
    ")\n",
    "\n",
    "# Adjusted size of the grid to ensure it's a multiple of the pixel size.\n",
    "grid_size = (\n",
    "    resolution[0] * pixel_size,\n",
    "    resolution[1] * pixel_size,\n",
    ")\n",
    "\n",
    "# Amount by which to offset each pixel to account for the grid being centered\n",
    "# at the origin, and also for each query to be made from the pixel's center.\n",
    "pixel_offset = (\n",
    "    (pixel_size - grid_size[0]) * 0.5,\n",
    "    (pixel_size - grid_size[1]) * 0.5,\n",
    ")\n",
    "\n",
    "# Create a Warp mesh representing the geometry in order to perform BVH queries.\n",
    "geom_mesh = wp.Mesh(\n",
    "    points=wp.array(GEOM_POINTS, dtype=wp.vec3),\n",
    "    indices=wp.array(GEOM_FACE_VERTEX_INDICES, dtype=int),\n",
    ")\n",
    "\n",
    "# Allocate an array to store the rendered texture.\n",
    "texture = wp.empty(shape=(resolution[1], resolution[0]), dtype=wp.vec3)\n",
    "\n",
    "\n",
    "\"\"\"Evaluation\"\"\"\n",
    "\n",
    "wp.launch(\n",
    "    compute_sdf_slice,\n",
    "    dim=resolution,\n",
    "    inputs=(\n",
    "        pixel_size,\n",
    "        pixel_offset,\n",
    "        geom_mesh.id,\n",
    "        sdf_band_width,\n",
    "    ),\n",
    "    outputs=(\n",
    "        texture,\n",
    "    ),\n",
    ")\n",
    "\n",
    "\n",
    "\"\"\"Visualization in Matplotlib\"\"\"\n",
    "\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(texture.numpy())\n",
    "plot_img.axes.set_axis_off()\n",
    "\n",
    "IPython.display.display(plot_fig)\n",
    "matplotlib.pyplot.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another such built-in is `wp.mesh_query_ray()`, which allows finding intersections between a ray and a mesh geometry, as is required for raycasting-like applications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Triangular mesh data for a geometry\"\"\"\n",
    "\n",
    "GEOM_POINTS = (\n",
    "    ( 0.0     ,  0.0     , -0.5     ), ( 0.0     ,  0.28656 , -0.425325),\n",
    "    ( 0.249999,  0.081230, -0.425325), ( 0.0     ,  0.447213, -0.223606),\n",
    "    ( 0.249999,  0.344095, -0.262865), ( 0.425325,  0.138196, -0.223606),\n",
    "    ( 0.154508, -0.212662, -0.425325), ( 0.404508, -0.131432, -0.262865),\n",
    "    ( 0.262865, -0.361803, -0.223606), (-0.154508, -0.212662, -0.425325),\n",
    "    ( 0.0     , -0.425325, -0.262865), (-0.262865, -0.361803, -0.223606),\n",
    "    (-0.249999,  0.081230, -0.425325), (-0.404508, -0.131432, -0.262865),\n",
    "    (-0.425325,  0.138196, -0.223606), (-0.249999,  0.344095, -0.262865),\n",
    "    (-0.154508,  0.475528,  0.0     ), (-0.404508,  0.293892,  0.0     ),\n",
    "    (-0.262865,  0.361803,  0.223606), (-0.5     ,  0.0     ,  0.0     ),\n",
    "    (-0.404508, -0.293892,  0.0     ), (-0.425325, -0.138196,  0.223606),\n",
    "    (-0.154508, -0.475528,  0.0     ), ( 0.154508, -0.475528,  0.0     ),\n",
    "    ( 0.0     , -0.447213,  0.223606), ( 0.404508, -0.293892,  0.0     ),\n",
    "    ( 0.5     ,  0.0     ,  0.0     ), ( 0.425325, -0.138196,  0.223606),\n",
    "    ( 0.404508,  0.293892,  0.0     ), ( 0.154508,  0.475528,  0.0     ),\n",
    "    ( 0.262865,  0.361803,  0.223606), ( 0.0     ,  0.425325,  0.262865),\n",
    "    (-0.404508,  0.131432,  0.262865), (-0.249999, -0.344095,  0.262865),\n",
    "    ( 0.249999, -0.344095,  0.262865), ( 0.404508,  0.131432,  0.262865),\n",
    "    ( 0.0     ,  0.0     ,  0.5     ), ( 0.154508,  0.212662,  0.425325),\n",
    "    (-0.154508,  0.212662,  0.425325), ( 0.249999, -0.081230,  0.425325),\n",
    "    ( 0.0     , -0.262865,  0.425325), (-0.249999, -0.081230,  0.425325),\n",
    ")\n",
    "\n",
    "GEOM_FACE_VERTEX_INDICES = (\n",
    "     2,  0,  1,  4,  1,  3,  5,  2,  4,  2,  1,  4,  6,  0,  2,  7,\n",
    "     2,  5,  8,  6,  7,  6,  2,  7,  9,  0,  6, 10,  6,  8, 11,  9,\n",
    "    10,  9,  6, 10, 12,  0,  9, 13,  9, 11, 14, 12, 13, 12,  9, 13,\n",
    "     1,  0, 12, 15, 12, 14,  3,  1, 15,  1, 12, 15, 16,  3, 15, 17,\n",
    "    15, 14, 18, 16, 17, 16, 15, 17, 19, 14, 13, 20, 13, 11, 21, 19,\n",
    "    20, 19, 13, 20, 22, 11, 10, 23, 10,  8, 24, 22, 23, 22, 10, 23,\n",
    "    25,  8,  7, 26,  7,  5, 27, 25, 26, 25,  7, 26, 28,  5,  4, 29,\n",
    "     4,  3, 30, 28, 29, 28,  4, 29, 29,  3, 16, 31, 16, 18, 30, 29,\n",
    "    31, 29, 16, 31, 17, 14, 19, 32, 19, 21, 18, 17, 32, 17, 19, 32,\n",
    "    20, 11, 22, 33, 22, 24, 21, 20, 33, 20, 22, 33, 23,  8, 25, 34,\n",
    "    25, 27, 24, 23, 34, 23, 25, 34, 26,  5, 28, 35, 28, 30, 27, 26,\n",
    "    35, 26, 28, 35, 38, 36, 37, 31, 37, 30, 18, 38, 31, 38, 37, 31,\n",
    "    37, 36, 39, 35, 39, 27, 30, 37, 35, 37, 39, 35, 39, 36, 40, 34,\n",
    "    40, 24, 27, 39, 34, 39, 40, 34, 40, 36, 41, 33, 41, 21, 24, 40,\n",
    "    33, 40, 41, 33, 41, 36, 38, 32, 38, 18, 21, 41, 32, 41, 38, 32,\n",
    ")\n",
    "\n",
    "\n",
    "\"\"\"Compute\"\"\"\n",
    "\n",
    "@wp.kernel\n",
    "def draw_ray_mesh_intersection(\n",
    "    pixel_size: float,\n",
    "    pixel_offset: wp.vec2,\n",
    "    geom: wp.uint64,\n",
    "    out_texture: wp.array(dtype=wp.vec3, ndim=2),\n",
    "):\n",
    "    j, i = wp.tid()\n",
    "\n",
    "    ray_origin = wp.vec3(0.0, 0.0, 2.0)\n",
    "    ray_dir = wp.normalize(\n",
    "        wp.vec3(\n",
    "            float(j) * pixel_size + pixel_offset[0],\n",
    "            float(i) * pixel_size + pixel_offset[1],\n",
    "            -1.0,\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    color = wp.vec3(0.0, 0.0, 0.0)\n",
    "\n",
    "    query = wp.mesh_query_ray(geom, ray_origin, ray_dir, 1.0e6)\n",
    "    if query.result:\n",
    "        color = query.normal * 0.5 + wp.vec3(0.5, 0.5, 0.5)\n",
    "\n",
    "    out_texture[i, j] = color\n",
    "\n",
    "\n",
    "\"\"\"Initialization\"\"\"\n",
    "\n",
    "# Size of the grid to cast rays from.\n",
    "grid_size = (1.5, 1.5)\n",
    "\n",
    "# Size for each pixel for the rendered texture.\n",
    "pixel_size = 0.005\n",
    "\n",
    "# Resolution of the rendered texture.\n",
    "resolution = (\n",
    "    int(wp.ceil(grid_size[0] / pixel_size)),\n",
    "    int(wp.ceil(grid_size[1] / pixel_size)),\n",
    ")\n",
    "\n",
    "# Adjusted size of the grid to ensure it's a multiple of the pixel size.\n",
    "grid_size = (\n",
    "    resolution[0] * pixel_size,\n",
    "    resolution[1] * pixel_size,\n",
    ")\n",
    "\n",
    "# Amount by which to offset each pixel to account for the grid being centered\n",
    "# at the origin, and also for each query to be made from the pixel's center.\n",
    "pixel_offset = (\n",
    "    (pixel_size - grid_size[0]) * 0.5,\n",
    "    (pixel_size - grid_size[1]) * 0.5,\n",
    ")\n",
    "\n",
    "# Create a Warp mesh representing the geometry in order to perform BVH queries.\n",
    "geom_mesh = wp.Mesh(\n",
    "    points=wp.array(GEOM_POINTS, dtype=wp.vec3),\n",
    "    indices=wp.array(GEOM_FACE_VERTEX_INDICES, dtype=int),\n",
    ")\n",
    "\n",
    "# Allocate an array to store the rendered texture.\n",
    "texture = wp.empty(shape=(resolution[1], resolution[0]), dtype=wp.vec3)\n",
    "\n",
    "\n",
    "\"\"\"Evaluation\"\"\"\n",
    "\n",
    "wp.launch(\n",
    "    draw_ray_mesh_intersection,\n",
    "    dim=resolution,\n",
    "    inputs=(\n",
    "        pixel_size,\n",
    "        pixel_offset,\n",
    "        geom_mesh.id,\n",
    "    ),\n",
    "    outputs=(\n",
    "        texture,\n",
    "    ),\n",
    ")\n",
    "\n",
    "\n",
    "# \"\"\"Visualization in Matplotlib\"\"\"\n",
    "\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(texture.numpy(), origin=\"lower\")\n",
    "plot_img.axes.set_axis_off()\n",
    "\n",
    "IPython.display.display(plot_fig)\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
}