Add wedge

src/pyransame/util.py

@@ -160,6 +160,16 @@ def _area_tetra(points: np.ndarray) -> float:
     )
 
 
+def _area_pyramid(points: np.ndarray):
+    tetra0 = [0, 1, 2, 4]
+    tetra1 = [0, 2, 3, 4]
+
+    area0 = _area_tetra(points[tetra0, :])
+    area1 = _area_tetra(points[tetra1, :])
+
+    return area0 + area1
+
+
 def _generate_points_in_pyramid(points: np.ndarray, n: int = 1) -> np.ndarray:
     tetra0 = [0, 1, 2, 4]
     tetra1 = [0, 2, 3, 4]
@@ -217,3 +227,30 @@ def _generate_points_in_pixel(
     r = pyransame.rng.random(size=(n, 2))
 
     return a + np.atleast_2d(r[:, 0]).T * v0 + np.atleast_2d(r[:, 1]).T * v1
+
+
+def _generate_points_in_wedge(points: np.ndarray, n: int = 1) -> np.ndarray:
+    tetra = [0, 2, 1, 4]
+    pyramid = [0, 2, 5, 3, 4]
+
+    area0 = _area_tetra(points[tetra, :])
+    area1 = _area_pyramid(points[pyramid, :])
+
+    areas = np.array([area0, area1])
+
+    p = areas / areas.sum()
+    out = np.empty((n, 3))
+
+    chosen_cells, unique_counts, point_indices = _random_cells(2, n, p)
+
+    for i, (chosen_cell, count) in enumerate(zip(chosen_cells, unique_counts)):
+        if chosen_cell == 0:
+            out[point_indices[i] : point_indices[i + 1], :] = _generate_points_in_tetra(
+                points[tetra, :], n=count
+            )
+        else:
+            out[
+                point_indices[i] : point_indices[i + 1], :
+            ] = _generate_points_in_pyramid(points[pyramid, :], n=count)
+
+    return out

src/pyransame/volume.py

@@ -24,6 +24,7 @@ def random_volume_points(
     - Pyramid
     - Tetrahedron
     - Voxel
+    - Wedge
 
     All cells must be convex.
 
@@ -100,6 +101,10 @@ def random_volume_points(
             points[
                 point_indices[i] : point_indices[i + 1], :
             ] = util._generate_points_in_pyramid(c.points, count)
+        elif c.type == CellType.WEDGE:
+            points[
+                point_indices[i] : point_indices[i + 1], :
+            ] = util._generate_points_in_wedge(c.points, count)
         else:
             raise NotImplementedError(
                 f"Random generation for {c.type.name} not yet supported"

tests/test_random_volume_points.py

@@ -5,6 +5,23 @@
 import pyransame
 
 
+def make_wedge():
+    points = np.array(
+        [
+            [0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.5, np.sqrt(3.0) / 2.0, 0.0],
+            [0.0, 0.0, 1.0],
+            [1.0, 0.0, 1.0],
+            [0.5, np.sqrt(3.0) / 2.0, 1.0],
+        ]
+    )
+    celltypes = [pv.CellType.WEDGE]
+    cells = [6, 0, 1, 2, 3, 4, 5]
+
+    return pv.UnstructuredGrid(cells, celltypes, points)
+
+
 def test_cell_types():
     mesh = pv.UniformGrid(dimensions=(4, 4, 4))
     assert mesh.get_cell(0).type == pv.CellType.VOXEL
@@ -18,6 +35,10 @@ def test_cell_types():
     assert mesh.get_cell(0).type == pv.CellType.PYRAMID
     pyransame.random_volume_points(mesh, 20)
 
+    mesh = make_wedge()
+    assert mesh.get_cell(0).type == pv.CellType.WEDGE
+    pyransame.random_volume_points(mesh, 20)
+
 
 def test_mixed_types():
     # as long as there are 3D cells, we should be able to sample even if there are other cell types

tests/test_util.py

@@ -17,6 +17,7 @@
     _generate_points_in_tri,
     _generate_points_in_tri_strip,
     _generate_points_in_voxel,
+    _generate_points_in_wedge,
 )
 
 
@@ -231,3 +232,19 @@ def test_uniformity_pyramid():
     points = _generate_points_in_pyramid(mesh.points, 2000000)
 
     assert np.allclose(points.mean(axis=0), center, rtol=1e-3, atol=1e-3)
+
+
+def test_uniformity_wedge():
+    a = np.array((0.0, 0.0, 0.0))
+    b = np.array((1.0, 0.0, 0.0))
+    c = np.array((0.5, np.sqrt(3.0) / 2.0, 0.0))
+    d = np.array((0.0, 0.0, 1.0))
+    e = np.array((1.0, 0.0, 1.0))
+    f = np.array((0.5, np.sqrt(3.0) / 2.0, 1.0))
+
+    center = np.array((0.5, np.sqrt(3.0) / 6.0, 0.5))
+
+    # needs a lot of points to converge
+    points = _generate_points_in_wedge(np.array([a, b, c, d, e, f]), 2000000)
+
+    assert np.allclose(points.mean(axis=0), center, rtol=1e-3, atol=1e-3)