Add Python implementation of cone sampling.

* contrib/cone-vector.py: New file.

1 files changed, 76 insertions, 0 deletions

diff --git a/contrib/cone-vector.py b/contrib/cone-vector.py
new file mode 100644
index 0000000..6effe53
--- /dev/null
+++ b/contrib/cone-vector.py
@@ -0,0 +1,76 @@
+# cone-vector.py --- A Python implementation of cone sampling
+# Copyright © 2021 Arun I <arunisaac@systemreboot.net>
+# Copyright © 2021 Murugesan Venkatapathi <murugesh@iisc.ac.in>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see
+# <https://www.gnu.org/licenses/>.
+
+from numpy import arcsin, concatenate, cos, dot, ones, sin, sqrt, tan, pi
+from numpy.random import randn, random
+from numpy.linalg import norm
+from scipy.special import betainc, betaincinv, gamma
+
+def random_vector_on_sphere (dim):
+    x = randn(dim)
+    return x / norm(x)
+
+def surface_area_of_ball (dim):
+    return 2 * pi**(dim/2) / gamma(dim/2)
+
+def planar_angle2solid_angle_fraction (planar_angle, dim):
+    alpha = (dim - 1) / 2
+    beta = 1/2
+    x = sin(planar_angle)**2
+    if planar_angle < pi/2:
+        return 0.5*betainc(alpha, beta, x)
+    else:
+        return 1 - 0.5*betainc(alpha, beta, x)
+
+def solid_angle_fraction2planar_angle (solid_angle_fraction, dim):
+    alpha = (dim - 1) / 2
+    beta = 1/2
+    sn = surface_area_of_ball(dim)
+    if solid_angle_fraction < 1/2:
+        planar_angle = betaincinv(alpha, beta, 2*solid_angle_fraction)
+    else:
+        planar_angle = betaincinv(alpha, beta, 2*(1-solid_angle_fraction))
+    return arcsin(sqrt(planar_angle))
+
+def rotate_from_nth_canonical (x, axis):
+    xn = x[-1]
+    axisn = axis[-1]
+    if axisn != 1:
+        b = norm(axis[:-1])
+        a = (dot(x, axis) - xn*axisn) / b
+        s = sqrt(1 - axisn**2)
+        x = x + (xn*s + a*(axisn - 1))/b * axis
+        x[-1] = x[-1] + xn*(axisn - 1) - a*s \
+            - axisn*(xn*s + a*(axisn - 1))/b
+    return x
+
+def random_vector_on_spherical_cap (axis, maximum_planar_angle):
+    dim = axis.size
+    maximum_solid_angle_fraction = planar_angle2solid_angle_fraction(maximum_planar_angle, dim)
+    x = random_vector_on_sphere(dim - 1) \
+        * cos(maximum_planar_angle) \
+        * tan(solid_angle_fraction2planar_angle(maximum_solid_angle_fraction*random(), dim))
+    x = concatenate([x, [cos(maximum_planar_angle)]])
+    return rotate_from_nth_canonical(x / norm(x), axis)
+
+# Sample code exercising the above functions
+dim = 100
+maximum_planar_angle = pi/3
+axis = ones(dim)
+axis = axis/norm(axis)
+print(random_vector_on_spherical_cap(axis, maximum_planar_angle))