chore(lib): fmt and tests

jeertmans · Jul 25, 2024 · 0d4decd · 0d4decd
1 parent ab54ad7
commit 0d4decd
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 4 deletions.
diff --git a/differt/src/differt/scene/triangle_scene.py b/differt/src/differt/scene/triangle_scene.py
@@ -101,9 +101,7 @@ def load_xml(cls, file: str) -> "TriangleScene":
         scene = differt_core.scene.triangle_scene.TriangleScene.load_xml(file)
 
         meshes = map(
-            lambda mesh: TriangleMesh(
-                vertices=mesh.vertices, triangles=mesh.triangles
-            ),
+            lambda mesh: TriangleMesh(vertices=mesh.vertices, triangles=mesh.triangles),
             scene.meshes,
         )
 

diff --git a/differt/tests/scene/test_triangle_scene.py b/differt/tests/scene/test_triangle_scene.py
@@ -1,3 +1,5 @@
+import equinox as eqx
+import jax.numpy as jnp
 import pytest
 
 from differt.scene.sionna import (
@@ -20,3 +22,15 @@ def test_load_xml(self, scene_name: str) -> None:
         sionna_scene = SionnaScene.load_xml(file)
 
         assert len(scene.meshes) == len(sionna_scene.shapes)
+
+    def test_plot(self) -> None:
+        file = get_sionna_scene("simple_street_canyon")
+        scene = TriangleScene.load_xml(file)
+
+        tx = jnp.array([[0.0, 0.0, 0.0]])
+        rx = jnp.array([[1.0, 1.0, 1.0]])
+
+        scene = eqx.tree_at(lambda s: s.transmitters, scene, tx)
+        scene = eqx.tree_at(lambda s: s.receivers, scene, rx)
+
+        scene.plot()
diff --git a/docs/source/notebooks/quickstart.ipynb b/docs/source/notebooks/quickstart.ipynb
@@ -40,8 +40,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import jax.numpy as jnp\n",
     "import equinox as eqx\n",
+    "import jax.numpy as jnp\n",
     "\n",
     "from differt.scene.sionna import download_sionna_scenes, get_sionna_scene\n",
     "from differt.scene.triangle_scene import TriangleScene\n",