Initial commit

1 files changed, 100 insertions, 0 deletions

diff --git a/src/samball/samball.py b/src/samball/samball.py
new file mode 100644
index 0000000..e5d84ad
--- /dev/null
+++ b/src/samball/samball.py
@@ -0,0 +1,100 @@
+# samball --- Sample n-dimensional balls
+# 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, cos, dot, empty, log, ones, sin, sqrt, pi, where
+from numpy.random import randn, random
+from numpy.linalg import norm
+from scipy.special import betainc, betaincinv
+
+def random_vector_on_sphere(dim):
+    """Return a random vector uniformly distributed on the unit sphere."""
+    x = randn(dim)
+    return x / norm(x)
+
+def planar_angle2solid_angle_fraction(planar_angle, dim):
+    """Return the solid angle fraction for a given planar angle."""
+    alpha = (dim - 1) / 2
+    beta = 1/2
+    return where(planar_angle < pi/2,
+                 0.5*betainc(alpha, beta, sin(planar_angle)**2),
+                 1 - 0.5*betainc(alpha, beta, sin(planar_angle)**2))
+
+def solid_angle_fraction2planar_angle(solid_angle_fraction, dim):
+    """Return the planar angle for a given solid angle fraction."""
+    alpha = (dim - 1) / 2
+    beta = 1/2
+    return where(solid_angle_fraction < 1/2,
+                 arcsin(sqrt(betaincinv(alpha, beta, 2*solid_angle_fraction))),
+                 pi - arcsin(sqrt(betaincinv(alpha, beta, 2*(1-solid_angle_fraction)))))
+
+def rotate_from_nth_canonical(x, axis):
+    """Rotate vector from around the nth canonical axis to the given 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_planar_angle_cdf(maximum_planar_angle, dim):
+    """Return a random planar angle using inverse transform sampling."""
+    return solid_angle_fraction2planar_angle(
+        planar_angle2solid_angle_fraction(maximum_planar_angle, dim)*random(),
+        dim)
+
+def random_planar_angle_pdf(maximum_planar_angle, dim):
+    """Return a random planar angle using rejection sampling."""
+    # We apply the log function just to prevent the floats from
+    # underflowing.
+    box_height = (dim-2)*log(sin(min(maximum_planar_angle, pi/2)))
+    while True:
+        theta = maximum_planar_angle*random()
+        f = box_height + log(random())
+        if f < (dim-2)*log(sin(theta)):
+            return theta
+
+def random_vector_on_disk(axis, planar_angle):
+    """Return a random vector uniformly distributed on the periphery of a
+disk."""
+    dim = axis.size
+    x = empty(dim)
+    x[:-1] = sin(planar_angle) * random_vector_on_sphere(dim - 1)
+    x[-1] = cos(planar_angle)
+    return rotate_from_nth_canonical(x, axis)
+
+def random_vector_on_spherical_cap_cdf(axis, maximum_planar_angle):
+    """Return a random vector uniformly distributed on a spherical
+cap. The random planar angle is generated using inverse transform
+sampling."""
+    return random_vector_on_disk(axis, random_planar_angle_cdf(maximum_planar_angle, axis.size))
+
+def random_vector_on_spherical_cap_pdf(axis, maximum_planar_angle):
+    """Return a random vector uniformly distributed on a spherical
+cap. The random planar angle is generated using rejection sampling.
+
+This function is more numerically stable than
+random_vector_on_spherical_cap_cdf for large dimensions and small
+angles.
+
+    """
+    return random_vector_on_disk(axis, random_planar_angle_pdf(maximum_planar_angle, axis.size))