aboutsummaryrefslogtreecommitdiff
path: root/contrib/cone-vector.py
blob: 9de71d873ed286f885dca6ddbbcfd99c5a5d08f3 (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
# 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, cos, dot, empty, ones, sin, sqrt, tan, pi, where, zeros
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 planar_angle2solid_angle_fraction (planar_angle, dim):
    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):
    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):
    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 = empty(dim)
    x[:-1] = random_vector_on_sphere(dim - 1) \
        * cos(maximum_planar_angle) \
        * tan(solid_angle_fraction2planar_angle(maximum_solid_angle_fraction*random(), dim))
    x[-1] = 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))