Function to generate random ellipsoid

AlexanderFabisch · May 1, 2022 · 809fd57 · 809fd57
1 parent 03d4110
commit 809fd57
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 8 deletions.
diff --git a/distance3d/random.py b/distance3d/random.py
@@ -260,6 +260,30 @@ def rand_capsule(random_state, center_scale=1.0, radius_scale=1.0,
     return capsule2origin, radius, height
 
 
+def rand_ellipsoid(random_state, min_radius=0.0):
+    """Sample ellipsoid.
+
+    Parameters
+    ----------
+    random_state : np.random.RandomState
+        Random number generator.
+
+    min_radius : float, optional (default: 0)
+        Minimum radius of ellipsoid.
+
+    Returns
+    -------
+    ellipsoid2origin : array, shape (4, 4)
+        Pose of the ellipsoid.
+
+    radii : array, shape (3,)
+        Radii of the ellipsoid within (min_radius, min_radius + 1].
+    """
+    ellipsoid2origin = pt.random_transform(random_state)
+    radii = min_radius + 1.0 - random_state.rand(3)
+    return ellipsoid2origin, radii
+
+
 def rand_cylinder(random_state, center_scale=1.0, min_radius=0.0,
                   min_length=0.0):
     """Sample cylinder.

diff --git a/distance3d/test/test_random.py b/distance3d/test/test_random.py
@@ -3,7 +3,7 @@
 from distance3d.random import (
     randn_point, randn_direction, randn_line, randn_line_segment,
     randn_plane, randn_triangle, randn_rectangle, rand_circle, rand_box,
-    rand_capsule, rand_cylinder, rand_sphere, randn_convex)
+    rand_capsule, rand_ellipsoid, rand_cylinder, rand_sphere, randn_convex)
 from pytest import approx
 from numpy.testing import assert_array_almost_equal
 
@@ -83,6 +83,15 @@ def test_rand_capsule():
     assert 0 < height <= 1
 
 
+def test_rand_ellipsoid():
+    random_state = np.random.RandomState(1080)
+    ellipsoid2origin, radii = rand_ellipsoid(random_state, min_radius=1.0)
+    pt.assert_transform(ellipsoid2origin)
+    assert 1.0 < radii[0] <= 2.0
+    assert 1.0 < radii[1] <= 2.0
+    assert 1.0 < radii[2] <= 2.0
+
+
 def test_rand_cylinder():
     random_state = np.random.RandomState(109)
     cylinder2origin, radius, length = rand_cylinder(random_state)

diff --git a/doc/source/api.rst b/doc/source/api.rst
@@ -222,6 +222,7 @@ API Reference
    ~distance3d.random.randn_triangle
    ~distance3d.random.rand_box
    ~distance3d.random.rand_capsule
+   ~distance3d.random.rand_ellipsoid
    ~distance3d.random.rand_cylinder
    ~distance3d.random.rand_sphere
    ~distance3d.random.randn_convex

diff --git a/examples/distance/plot_plane_to_ellipsoid.py b/examples/distance/plot_plane_to_ellipsoid.py
@@ -13,10 +13,8 @@
 from distance3d import random, plotting
 
 
-random_state = np.random.RandomState(5)
-# TODO move to random
-ellipsoid2origin = pt.random_transform(random_state)
-radii = np.array([0.1, 1, 2])
+random_state = np.random.RandomState(2)
+ellipsoid2origin, radii = random.rand_ellipsoid(random_state)
 
 ax = ppu.make_3d_axis(ax_s=2)
 

diff --git a/examples/distance/plot_point_to_ellipsoid.py b/examples/distance/plot_point_to_ellipsoid.py
@@ -14,9 +14,7 @@
 
 
 random_state = np.random.RandomState(3)
-# TODO move to random
-ellipsoid2origin = pt.random_transform(random_state)
-radii = 1 + random_state.rand(3)
+ellipsoid2origin, radii = random.rand_ellipsoid(random_state, min_radius=0.5)
 
 ax = ppu.make_3d_axis(ax_s=2)