add hexagonal prism (#58)

src/pyransame/util.py

@@ -322,3 +322,32 @@ def _generate_points_in_pentagonal_prism(points: np.ndarray, n: int = 1) -> np.n
             ] = _generate_points_in_hexahedron(points[hexahedron, :], n=count)
 
     return out
+
+
+def _generate_points_in_hexagonal_prism(points: np.ndarray, n: int = 1) -> np.ndarray:
+    hexahedron0 = [0, 1, 2, 5, 6, 7, 8, 11]
+    hexahedron1 = [5, 2, 3, 4, 11, 8, 9, 10]
+
+    areas = np.array(
+        [
+            _area_hexahedron(points[hexahedron0, :]),
+            _area_hexahedron(points[hexahedron1, :]),
+        ]
+    )
+
+    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_hexahedron(points[hexahedron0, :], n=count)
+        else:
+            out[
+                point_indices[i] : point_indices[i + 1], :
+            ] = _generate_points_in_hexahedron(points[hexahedron1, :], n=count)
+
+    return out

src/pyransame/volume.py

@@ -21,6 +21,7 @@ def random_volume_points(
 
     Supported cell types:
 
+    - Hexagonal Prism
     - Hexahedron
     - Pentagonal Prism
     - Pyramid
@@ -115,6 +116,10 @@ def random_volume_points(
             points[
                 point_indices[i] : point_indices[i + 1], :
             ] = util._generate_points_in_pentagonal_prism(c.points, count)
+        elif c.type == CellType.HEXAGONAL_PRISM:
+            points[
+                point_indices[i] : point_indices[i + 1], :
+            ] = util._generate_points_in_hexagonal_prism(c.points, count)
         else:
             raise NotImplementedError(
                 f"Random generation for {c.type.name} not yet supported"

tests/test_random_volume_points.py

@@ -54,6 +54,20 @@ def make_pentagonal_prism():
     return pv.UnstructuredGrid(cells, celltypes, p)
 
 
+def make_hexagonal_prism():
+    angles = np.arange(6) * 2 * np.pi / 6  # angles of regular pentagon in radians
+    p = np.zeros(shape=(12, 3))
+    np.sin(angles, out=p[0:6, 0])
+    np.cos(angles, out=p[0:6, 1])
+
+    np.sin(angles, out=p[6:, 0])
+    np.cos(angles, out=p[6:, 1])
+    p[6:, 2] = 1.0
+    celltypes = [pv.CellType.HEXAGONAL_PRISM]
+    cells = [12, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
+    return pv.UnstructuredGrid(cells, celltypes, p)
+
+
 def test_cell_types():
     mesh = pv.UniformGrid(dimensions=(4, 4, 4))
     assert mesh.get_cell(0).type == pv.CellType.VOXEL
@@ -79,6 +93,10 @@ def test_cell_types():
     assert mesh.get_cell(0).type == pv.CellType.PENTAGONAL_PRISM
     pyransame.random_volume_points(mesh, 20)
 
+    mesh = make_hexagonal_prism()
+    assert mesh.get_cell(0).type == pv.CellType.HEXAGONAL_PRISM
+    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

@@ -9,6 +9,7 @@
 from hypothesis.extra.numpy import arrays
 
 from pyransame.util import (
+    _generate_points_in_hexagonal_prism,
     _generate_points_in_hexahedron,
     _generate_points_in_pentagonal_prism,
     _generate_points_in_pixel,
@@ -285,3 +286,18 @@ def test_uniformity_pentagonal_prism():
     center = np.array([0.0, 0.0, 0.5])
     points = _generate_points_in_pentagonal_prism(p, 2000000)
     assert np.allclose(points.mean(axis=0), center, rtol=1e-3, atol=1e-3)
+
+
+def test_uniformity_hexagonal_prism():
+    angles = np.arange(6) * 2 * np.pi / 6  # angles of regular hexagon in radians
+    p = np.zeros(shape=(12, 3))
+    np.sin(angles, out=p[0:6, 0])
+    np.cos(angles, out=p[0:6, 1])
+
+    np.sin(angles, out=p[6:, 0])
+    np.cos(angles, out=p[6:, 1])
+    p[6:, 2] = 1.0
+
+    center = np.array([0.0, 0.0, 0.5])
+    points = _generate_points_in_hexagonal_prism(p, 2000000)
+    assert np.allclose(points.mean(axis=0), center, rtol=1e-3, atol=1e-3)