Port to python.

* python/nsmc.py: New file.

1 files changed, 193 insertions, 0 deletions

diff --git a/python/nsmc.py b/python/nsmc.py
new file mode 100644
index 0000000..27186e0
--- /dev/null
+++ b/python/nsmc.py
@@ -0,0 +1,193 @@
+# nsmc --- n-sphere Monte Carlo method
+# Copyright © 2022 Arun I <arunisaac@systemreboot.net>
+# Copyright © 2022 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 math import exp, fabs, inf, isclose, log, pi, sqrt
+from numpy import arange, arcsin, polyval, prod, random, sin, sum, where
+from numpy.linalg import norm
+
+from scipy.integrate import quad
+from scipy.special import betaincinv, gamma, gammainc, gammaln
+from sambal import random_on_disk, random_on_sphere
+
+def ln_ball_volume(n):
+    return n/2*log(pi) - gammaln(1 + n/2)
+
+def ball_volume(n):
+    return exp(ln_ball_volume(n))
+
+def ln_ball_surface_area(n):
+    return log(n) + ln_ball_volume(n)
+
+def ball_surface_area(n):
+    return n * ball_volume(n)
+
+def lower_incomplete_gamma(a, x):
+    return gamma(a) * gammainc(a, x)
+
+# Extent oracles and associated true volumes
+
+def make_bernoulli_extent_oracle(p, r0, r1):
+    return lambda x: r1 if binomial(1, p) else r0
+
+def bernoulli_true_volume(p, r0, r1, n):
+    return ball_volume(n) * (p*r1**n + (1-p)*r0**n)
+
+def make_uniform_extent_oracle(a, b):
+    return lambda x: random.uniform(a, b)
+
+# TODO: Verify the accuracy of this function for non-trivial a, b.
+def uniform_true_volume(a, b, n):
+    return ball_volume(n) * (b**(n+1)/(n+1)) - (a**(n+1)/(n+1))
+
+def make_beta_extent_oracle(alpha, beta):
+    return lambda x: random.beta(alpha, beta)
+
+def beta_true_volume(alpha, beta, n):
+    r = arange(n)
+    return ball_volume(n) * prod((alpha + r) / (alpha + beta + r))
+
+def make_cube_extent_oracle(n, edge):
+    return lambda x: edge / 2 / norm(x, inf)
+
+def cube_true_volume(n, edge):
+    return edge**n
+
+def make_spheroid_extent_oracle(n, eccentricity):
+    return lambda x: 1 / sqrt((x[0]/eccentricity)**2 + norm(x[1:])**2)
+
+def spheroid_true_volume(n, eccentricity):
+    return ball_volume(n) * eccentricity
+
+def make_ellipsoid_extent_oracle(n, axes):
+    return lambda x: 1 / norm(x / axes)
+
+def ellipsoid_true_volume(n, axes):
+    return ball_volume(n) * prod(axes)
+
+# Integrands and associated true integrals
+
+def gaussian_integrand():
+    return lambda r, x: exp(-r*r/2)
+
+def gaussian_true_integral(n, max_extent):
+    return (2**(n/2 - 1) * ball_surface_area(n) * \
+        (max_extent*lower_incomplete_gamma(n/2, max_extent**2/2) \
+         - sqrt(2)*lower_incomplete_gamma((n+1)/2, max_extent**2/2))) \
+         / max_extent
+
+def polynomial_integrand(coefficients):
+    return lambda r, x: polyval(coefficients, r)
+
+def polynomial_true_integral(n, max_extent, coefficients):
+    m = len(coefficients)
+    return ball_surface_area(n) * max_extent**n * \
+        polyval([ak/(n+m-1-index)/(n+m-index) for index, ak in enumerate(coefficients)],
+                max_extent)
+
+def x_coordinate_integrand():
+    return lambda r, x: fabs(r*x[0])
+
+def x_coordinate_true_integral(n, max_extent):
+    return max_extent**(n+1) * ball_surface_area(n+3) / 2 / pi**2 / (n+2)
+
+# Volume estimation
+
+def accumulate_stats(x, n, mean):
+    return n + 1, (mean*n + x)/(n + 1)
+
+def volume(extent_oracle, true_volume, n, rtol, window_length=1000):
+    vn = ln_ball_volume(n)
+    vol = 0
+    samples = 0
+    accurate_estimates = 0
+    while accurate_estimates < window_length:
+        x = random_on_sphere(n)
+        samples, vol = accumulate_stats(exp(vn + n*log(extent_oracle(x))), samples, vol)
+        accurate_estimates = accurate_estimates + 1 if isclose(vol, true_volume, rel_tol=rtol) else 0
+    return vol, samples
+
+# Integral estimation
+
+def integral_per_direction(integrand, direction, n, radius):
+    return ball_surface_area(n) * quad(lambda r: r**(n-1) * integrand(r, direction),
+                                       0, radius)[0]
+
+def integral(integrand, extent_oracle, true_integral, n, rtol, window_length=1000):
+    accurate_estimates = 0
+    integ = 0
+    samples = 0
+    while accurate_estimates < window_length:
+        x = random_on_sphere(n)
+        samples, integ = accumulate_stats(integral_per_direction(integrand, x, n, extent_oracle(x)),
+                                          samples, integ)
+        accurate_estimates = accurate_estimates + 1 if isclose(integ, true_integral, rel_tol=rtol) else 0
+    return integ, samples
+
+# Volume estimation with importance sampling
+
+def solid_angle_fraction_to_planar_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 hollow_cone_random_vector(axis, minimum_planar_angle, maximum_planar_angle):
+    """Return a random vector uniformly distributed on a spherical
+cap. The random planar angle is generated using rejection sampling.
+
+    """
+    # We apply the log function just to prevent the floats from
+    # underflowing.
+    dim = axis.size
+    box_height = (dim-2)*log(sin(min(maximum_planar_angle, pi/2)))
+    while True:
+        theta = random.uniform(minimum_planar_angle, maximum_planar_angle)
+        f = box_height + log(random.random())
+        if f < (dim-2)*log(sin(theta)):
+            break
+    return random_on_disk(axis, theta)
+
+def volume_cone(extent_oracle, mean, omega_min, omega_max, samples):
+    n = mean.size
+    vn = ln_ball_volume(n)
+    theta_min = solid_angle_fraction_to_planar_angle(omega_min, n)
+    theta_max = solid_angle_fraction_to_planar_angle(omega_max, n)
+    samples_done = 0
+    vol = 0
+    for i in range(samples):
+        x = hollow_cone_random_vector(mean, theta_min, theta_max)
+        samples_done, vol = accumulate_stats(exp(vn + n*log(extent_oracle(x))),
+                                             samples_done, vol)
+    return vol * (omega_max - omega_min)
+
+def volume_importance(extent_oracle, mean, samples_per_cone, solid_angle_factor,
+                      solid_angle_threshold_exponent_factor):
+    vol = 0
+    omega_max = 0.5
+    samples = 0
+    n = mean.size
+    while omega_max > 2**(-n*solid_angle_threshold_exponent_factor):
+        omega_min = omega_max / solid_angle_factor
+        vol = vol + volume_cone(extent_oracle, mean, omega_min, omega_max, samples_per_cone)
+        samples = samples + samples_per_cone
+        omega_max = omega_min
+    vol = vol + volume_cone(extent_oracle, mean, 0, omega_max, samples_per_cone)
+    samples = samples + samples_per_cone
+    return vol, samples