feat(lib): add random sampling and ML features

differt-core/src/rt/graph.rs

@@ -117,6 +117,9 @@ pub mod complete {
     /// A complete graph, i.e.,
     /// a simple undirected graph in which every pair of
     /// distinc nodes is connected by a unique edge.
+    ///
+    /// Args:
+    ///     num_nodes (int): The number of nodes.
     #[pyclass]
     #[derive(Clone, Debug)]
     pub struct CompleteGraph {
@@ -547,6 +550,12 @@ pub mod directed {
     }
 
     impl DiGraph {
+        #[inline]
+        pub fn empty(num_nodes: usize) -> Self {
+            Self {
+                edges_list: vec![vec![]; num_nodes],
+            }
+        }
         #[inline]
         pub fn get_adjacent_nodes(&self, node: NodeId) -> &[NodeId] {
             self.edges_list[node].as_ref()
@@ -574,6 +583,23 @@ pub mod directed {
 
     #[pymethods]
     impl DiGraph {
+        /// Create an edgeless directed graph with ``num_nodes`` nodes.
+        ///
+        /// This is equivalent to creating a directed graph from
+        /// an adjacency matrix will all entries equal to :py:data:`False`.
+        ///
+        /// Args:
+        ///     num\_nodes (int): The number of nodes.
+        ///
+        /// Return:
+        ///     DiGraph: A directed graph.
+        #[classmethod]
+        #[pyo3(name = "empty")]
+        #[pyo3(text_signature = "(cls, num_nodes)")]
+        fn py_empty(_: Bound<'_, PyType>, num_nodes: usize) -> Self {
+            Self::empty(num_nodes)
+        }
+
         /// Create a directed graph from an adjacency matrix.
         ///
         /// Each row of the adjacency matrix ``M`` contains boolean
@@ -1111,6 +1137,16 @@ mod tests {
         assert_eq!(got, expected);
     }
 
+    #[rstest]
+    #[case(9, 2)]
+    #[case(3, 3)]
+    fn test_empty_di_graph_returns_all_paths(#[case] num_nodes: usize, #[case] depth: usize) {
+        let mut graph = DiGraph::empty(num_nodes);
+        let (from, to) = graph.insert_from_and_to_nodes(false);
+
+        assert_eq!(graph.all_paths(from, to, depth + 2, true).count(), 0);
+    }
+
     #[rstest]
     #[case(9, 2)]
     #[case(3, 3)]

differt/src/differt/geometry/triangle_mesh.py

@@ -4,10 +4,11 @@
 from typing import Any
 
 import equinox as eqx
+import jax
 import jax.numpy as jnp
 import numpy as np
 from beartype import beartype as typechecker
-from jaxtyping import Array, Bool, Float, UInt, jaxtyped
+from jaxtyping import Array, Bool, Float, PRNGKeyArray, UInt, jaxtyped
 
 import differt_core.geometry.triangle_mesh
 
@@ -134,6 +135,13 @@ def diffraction_edges(self) -> UInt[Array, "num_edges 3"]:
         """The diffraction edges."""
         raise NotImplementedError
 
+    @cached_property
+    def bounding_box(self) -> Float[Array, "2 3"]:
+        """The bounding box (min. and max. coordinates)."""
+        return jnp.vstack(
+            (jnp.min(self.vertices, axis=0), jnp.max(self.vertices, axis=0))
+        )
+
     @classmethod
     def empty(cls) -> "TriangleMesh":
         """
@@ -207,3 +215,26 @@ def plot(self, **kwargs: Any) -> Any:
             triangles=np.asarray(self.triangles),
             **kwargs,
         )
+
+    @eqx.filter_jit
+    def sample(
+        self, size: int, replace: bool = False, *, key: PRNGKeyArray
+    ) -> "TriangleMesh":
+        """
+        Generate a new mesh by randomly sampling triangles from this geometry.
+
+        Args:
+            size: The size of the sample, i.e., the number of triangles.
+            replace: Whether to sample with or without replacement.
+            key: The :func:`jax.random.PRNGKey` to be used.
+
+        Return:
+            A new random mesh.
+        """
+        triangles = self.triangles[
+            jax.random.choice(
+                key, self.triangles.shape[0], shape=(size,), replace=replace
+            ),
+            :,
+        ]
+        return TriangleMesh(vertices=self.vertices, triangles=triangles)

differt/src/differt/plotting/_core.py

@@ -1,9 +1,7 @@
 """Core plotting implementations."""
 
-from __future__ import annotations
-
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any
+from typing import Any, Optional, Union
 
 import numpy as np
 from jaxtyping import Float, Num, UInt
@@ -15,18 +13,34 @@
     process_vispy_kwargs,
 )
 
-if TYPE_CHECKING:
+# We cannot use from __future__ import annotations because
+#   otherwise array annotations do not render correctly.
+# We cannot rely on TYPE_CHECKING-guarded annotation
+#   because Sphinx will fail to import this NumPy or Jax typing
+# Hence, we prefer to silence pyright instead.
+
+try:
     from matplotlib.figure import Figure as MplFigure
+except ImportError:
+    MplFigure = Any
+
+try:
     from plotly.graph_objects import Figure
+except ImportError:
+    Figure = Any
+
+try:
     from vispy.scene.canvas import SceneCanvas as Canvas
+except ImportError:
+    Canvas = Any
 
 
-@dispatch  # type: ignore
+@dispatch
 def draw_mesh(
     vertices: Float[np.ndarray, "num_vertices 3"],
     triangles: UInt[np.ndarray, "num_triangles 3"],
     **kwargs: Any,
-) -> Canvas | MplFigure | Figure:  # type: ignore
+) -> Union[Canvas, MplFigure, Figure]:  # type: ignore[reportInvalidTypeForm]
     """
     Plot a 3D mesh made of triangles.
 
@@ -72,7 +86,7 @@ def _(
     vertices: Float[np.ndarray, "num_vertices 3"],
     triangles: UInt[np.ndarray, "num_triangles 3"],
     **kwargs: Any,
-) -> Canvas:
+) -> Canvas:  # type: ignore[reportInvalidTypeForm]
     from vispy.scene.visuals import Mesh
 
     canvas, view = process_vispy_kwargs(kwargs)
@@ -88,7 +102,7 @@ def _(
     vertices: Float[np.ndarray, "num_vertices 3"],
     triangles: UInt[np.ndarray, "num_triangles 3"],
     **kwargs: Any,
-) -> MplFigure:
+) -> MplFigure:  # type: ignore[reportInvalidTypeForm]
     fig, ax = process_matplotlib_kwargs(kwargs)
 
     x, y, z = vertices.T
@@ -102,7 +116,7 @@ def _(
     vertices: Float[np.ndarray, "num_vertices 3"],
     triangles: UInt[np.ndarray, "num_triangles 3"],
     **kwargs: Any,
-) -> Figure:
+) -> Figure:  # type: ignore[reportInvalidTypeForm]
     fig = process_plotly_kwargs(kwargs)
 
     x, y, z = vertices.T
@@ -111,10 +125,10 @@ def _(
     return fig.add_mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, **kwargs)
 
 
-@dispatch  # type: ignore
+@dispatch
 def draw_paths(
-    paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any
-) -> Canvas | MplFigure | Figure:  # type: ignore
+    paths: Float[np.ndarray, r"\*batch path_length 3"], **kwargs: Any
+) -> Union[Canvas, MplFigure, Figure]:  # type: ignore[reportInvalidTypeForm]
     """
     Plot a batch of paths of the same length.
 
@@ -137,12 +151,12 @@ def draw_paths(
             >>> from differt.plotting import draw_paths
             >>>
             >>> def rotation(angle: float) -> np.ndarray:
-            ...     c = np.cos(angle)
-            ...     s = np.sin(angle)
+            ...     co = np.cos(angle)
+            ...     si = np.sin(angle)
             ...     return np.array(
             ...         [
-            ...             [+c, -s, 0.0],
-            ...             [+s, +c, 0.0],
+            ...             [+co, -si, 0.0],
+            ...             [+si, +co, 0.0],
             ...             [0.0, 0.0, 1.0],
             ...         ]
             ...     )
@@ -172,7 +186,7 @@ def draw_paths(
 
 
 @draw_paths.register("vispy")
-def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> Canvas:
+def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> Canvas:  # type: ignore[reportInvalidTypeForm]
     from vispy.scene.visuals import LinePlot
 
     canvas, view = process_vispy_kwargs(kwargs)
@@ -186,7 +200,7 @@ def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> Canvas
 
 
 @draw_paths.register("matplotlib")
-def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> MplFigure:
+def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> MplFigure:  # type: ignore[reportInvalidTypeForm]
     fig, ax = process_matplotlib_kwargs(kwargs)
 
     for i in np.ndindex(paths.shape[:-2]):
@@ -196,7 +210,7 @@ def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> MplFig
 
 
 @draw_paths.register("plotly")
-def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> Figure:
+def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> Figure:  # type: ignore[reportInvalidTypeForm]
     fig = process_plotly_kwargs(kwargs)
 
     for i in np.ndindex(paths.shape[:-2]):
@@ -206,13 +220,13 @@ def _(paths: Float[np.ndarray, "*batch path_length 3"], **kwargs: Any) -> Figure
     return fig
 
 
-@dispatch  # type: ignore
+@dispatch
 def draw_markers(
     markers: Float[np.ndarray, "num_markers 3"],
-    labels: Sequence[str] | None = None,
-    text_kwargs: Mapping[str, Any] | None = None,
+    labels: Optional[Sequence[str]] = None,
+    text_kwargs: Optional[Mapping[str, Any]] = None,
     **kwargs: Any,
-) -> Canvas | MplFigure | Figure:  # type: ignore
+) -> Union[Canvas, MplFigure, Figure]:  # type: ignore[reportInvalidTypeForm]
     """
     Plot markers and, optionally, their label.
 
@@ -259,10 +273,10 @@ def draw_markers(
 @draw_markers.register("vispy")
 def _(
     markers: Float[np.ndarray, "num_markers 3"],
-    labels: Sequence[str] | None = None,
-    text_kwargs: Mapping[str, Any] | None = None,
+    labels: Optional[Sequence[str]] = None,
+    text_kwargs: Optional[Mapping[str, Any]] = None,
     **kwargs: Any,
-) -> Canvas:
+) -> Canvas:  # type: ignore[reportInvalidTypeForm]
     from vispy.scene.visuals import Markers, Text
 
     canvas, view = process_vispy_kwargs(kwargs)
@@ -280,20 +294,20 @@ def _(
 @draw_markers.register("matplotlib")
 def _(
     markers: Float[np.ndarray, "num_markers 3"],
-    labels: Sequence[str] | None = None,
-    text_kwargs: Mapping[str, Any] | None = None,
+    labels: Optional[Sequence[str]] = None,
+    text_kwargs: Optional[Mapping[str, Any]] = None,
     **kwargs: Any,
-) -> MplFigure:
+) -> MplFigure:  # type: ignore[reportInvalidTypeForm]
     raise NotImplementedError  # TODO
 
 
 @draw_markers.register("plotly")
 def _(
     markers: Float[np.ndarray, "num_markers 3"],
-    labels: Sequence[str] | None = None,
-    text_kwargs: Mapping[str, Any] | None = None,
+    labels: Optional[Sequence[str]] = None,
+    text_kwargs: Optional[Mapping[str, Any]] = None,
     **kwargs: Any,
-) -> Figure:
+) -> Figure:  # type: ignore[reportInvalidTypeForm]
     fig = process_plotly_kwargs(kwargs)
 
     if labels:
@@ -309,14 +323,18 @@ def _(
     )
 
 
-@dispatch  # type: ignore
+@dispatch
 def draw_image(
-    data: Num[np.ndarray, "m n"] | Num[np.ndarray, "m n 3"] | Num[np.ndarray, "m n 4"],
-    x: Float[np.ndarray, " *m"] | None = None,
-    y: Float[np.ndarray, " *n"] | None = None,
+    data: Union[
+        Num[np.ndarray, "rows cols"],
+        Num[np.ndarray, "rows cols 3"],
+        Num[np.ndarray, "rows cols 4"],
+    ],
+    x: Optional[Float[np.ndarray, " rows"]] = None,
+    y: Optional[Float[np.ndarray, " cols"]] = None,
     z0: float = 0.0,
     **kwargs: Any,
-) -> Canvas | MplFigure | Figure:  # type: ignore
+) -> Union[Canvas, MplFigure, Figure]:  # type: ignore[reportInvalidTypeForm]
     """
     Plot a 2D image on a 3D canvas, at using a fixed z-coordinate.
 
@@ -368,12 +386,16 @@ def draw_image(
 
 @draw_image.register("vispy")
 def _(
-    data: Num[np.ndarray, "m n"] | Num[np.ndarray, "m n 3"] | Num[np.ndarray, "m n 4"],
-    x: Float[np.ndarray, " ..."] | None = None,
-    y: Float[np.ndarray, " ..."] | None = None,
+    data: Union[
+        Num[np.ndarray, "rows cols"],
+        Num[np.ndarray, "rows cols 3"],
+        Num[np.ndarray, "rows cols 4"],
+    ],
+    x: Optional[Float[np.ndarray, " rows"]] = None,
+    y: Optional[Float[np.ndarray, " cols"]] = None,
     z0: float = 0.0,
     **kwargs: Any,
-) -> Canvas:
+) -> Canvas:  # type: ignore[reportInvalidTypeForm]
     from vispy.scene.visuals import Image
     from vispy.visuals.transforms import STTransform
 
@@ -416,12 +438,16 @@ def _(
 
 @draw_image.register("matplotlib")
 def _(
-    data: Num[np.ndarray, "m n"] | Num[np.ndarray, "m n 3"] | Num[np.ndarray, "m n 4"],
-    x: Float[np.ndarray, " ..."] | None = None,
-    y: Float[np.ndarray, " ..."] | None = None,
+    data: Union[
+        Num[np.ndarray, "rows cols"],
+        Num[np.ndarray, "rows cols 3"],
+        Num[np.ndarray, "rows cols 4"],
+    ],
+    x: Optional[Float[np.ndarray, " rows"]] = None,
+    y: Optional[Float[np.ndarray, " cols"]] = None,
     z0: float = 0.0,
     **kwargs: Any,
-) -> MplFigure:
+) -> MplFigure:  # type: ignore[reportInvalidTypeForm]
     fig, ax = process_matplotlib_kwargs(kwargs)
 
     ax.plot_surface(X=x, Y=y, Z=np.full_like(data, z0), color=data, **kwargs)
@@ -431,12 +457,16 @@ def _(
 
 @draw_image.register("plotly")
 def _(
-    data: Num[np.ndarray, "m n"] | Num[np.ndarray, "m n 3"] | Num[np.ndarray, "m n 4"],
-    x: Float[np.ndarray, " ..."] | None = None,
-    y: Float[np.ndarray, " ..."] | None = None,
+    data: Union[
+        Num[np.ndarray, "rows cols"],
+        Num[np.ndarray, "rows cols 3"],
+        Num[np.ndarray, "rows cols 4"],
+    ],
+    x: Optional[Float[np.ndarray, " rows"]] = None,
+    y: Optional[Float[np.ndarray, " cols"]] = None,
     z0: float = 0.0,
     **kwargs: Any,
-) -> Figure:
+) -> Figure:  # type: ignore[reportInvalidTypeForm]
     fig = process_plotly_kwargs(kwargs)
 
     return fig.add_surface(