Vectorize functions.

* contrib/cone-vector.py: Import where.
(planar_angle2solid_angle_fraction,
solid_angle_fraction2planar_angle): Vectorize functions.

1 files changed, 7 insertions, 10 deletions

diff --git a/contrib/cone-vector.py b/contrib/cone-vector.py
index 85c2c49..9de71d8 100644
--- a/contrib/cone-vector.py
+++ b/contrib/cone-vector.py
@@ -16,7 +16,7 @@
 # along with this program.  If not, see
 # <https://www.gnu.org/licenses/>.
 
-from numpy import arcsin, cos, dot, empty, ones, sin, sqrt, tan, pi
+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
@@ -28,19 +28,16 @@ def random_vector_on_sphere (dim):
 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)
+    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
-    if solid_angle_fraction < 1/2:
-        return arcsin(sqrt(betaincinv(alpha, beta, 2*solid_angle_fraction)))
-    else:
-        return pi - arcsin(sqrt(betaincinv(alpha, beta, 2*(1-solid_angle_fraction))))
+    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]