use plot utils in examples

jeertmans · Jan 4, 2024 · acf9134 · acf9134
1 parent 1e3fc4b
commit acf9134
Showing 1 changed file with 17 additions and 20 deletions.
diff --git a/docs/source/notebooks/quickstart.ipynb b/docs/source/notebooks/quickstart.ipynb
@@ -17,6 +17,7 @@
    "source": [
     "from pathlib import Path\n",
     "\n",
+    "import numpy as np\n",
     "import jax.numpy as jnp\n",
     "import plotly.graph_objects as go\n",
     "\n",
@@ -29,9 +30,17 @@
     "    consecutive_vertices_are_on_same_side_of_mirrors,\n",
     "    image_method,\n",
     ")\n",
-    "from differt.rt.utils import generate_all_path_candidates, rays_intersect_triangles\n",
-    "\n",
-    "dplt.use(\"plotly\")"
+    "from differt.rt.utils import generate_all_path_candidates, rays_intersect_triangles"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c7764529-0550-40b5-ac3c-ca9cc0cb31bb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dplt.use(\"plotly\")  # Let's use the Plotly backend"
    ]
   },
   {
@@ -66,20 +75,7 @@
     "tx = jnp.array([0.0, 4.9352, 22.0])\n",
     "rx = jnp.array([0.0, 10.034, 1.50])\n",
     "\n",
-    "fig.add_traces(\n",
-    "    go.Scatter3d(\n",
-    "        x=[tx[0], rx[0]],\n",
-    "        y=[tx[1], rx[1]],\n",
-    "        z=[tx[2], rx[2]],\n",
-    "        marker=dict(\n",
-    "            size=7,\n",
-    "            color=\"red\",\n",
-    "        ),\n",
-    "        mode=\"markers+text\",\n",
-    "        text=[\"tx\", \"rx\"],\n",
-    "        name=\"nodes\",\n",
-    "    )\n",
-    ")"
+    "dplt.draw_markers(np.array([tx, rx]), labels=[\"tx\", \"rx\"], figure=fig, name=\"nodes\")"
    ]
   },
   {
@@ -91,14 +87,15 @@
    "source": [
     "select = [8, 9, 22, 23]\n",
     "\n",
-    "x, y, z = mesh.vertices.T\n",
-    "i, j, k = triangles = mesh.triangles[select, :].T\n",
+    "# differt.plotting (dplt) works with NumPy arrays, not JAX arrays\n",
+    "vertices = np.asarray(mesh.vertices)\n",
+    "triangles = np.asarray(mesh.triangles[select, :])\n",
     "\n",
     "select = jnp.array(\n",
     "    select[::2], dtype=int\n",
     ")  # We actually only need one triangle per plane, so [8, 22]\n",
     "\n",
-    "fig.add_traces(go.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, color=\"red\"))"
+    "dplt.draw_mesh(vertices, triangles, figure=fig, color=\"red\")"
    ]
   },
   {