1 files changed, 6 insertions, 6 deletions
diff --git a/photon.py b/photon.py
index d1574ea..f8ece6a 100644
--- a/photon.py
+++ b/photon.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-def uniform_sphere(size=None):
+def uniform_sphere(size=None, dtype=np.double):
     """
     Generate random points isotropically distributed across the unit sphere.
 
@@ -17,13 +17,13 @@ def uniform_sphere(size=None):
 
     c = np.sqrt(1-u**2)
 
-    points = np.empty((x.size, 3))
+    if size is None:
+        return np.array([c*np.cos(theta), c*np.sin(theta), u])
+
+    points = np.empty((size, 3), dtype)
 
     points[:,0] = c*np.cos(theta)
     points[:,1] = c*np.sin(theta)
     points[:,2] = u
 
-    if size is None:
-        return points[0]
-    else:
-        return points
+    return points