aboutsummaryrefslogtreecommitdiff
path: root/python/nsmc.py
blob: 27186e03a9194523d5d5b3a1f979420b995c15aa (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
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