add normalize() function to transform module. np.apply_along_axis(np.linalg.norm,...) is really slow!

1 files changed, 18 insertions, 0 deletions

diff --git a/transform.py b/transform.py
index 4793c0f..299da46 100644
--- a/transform.py
+++ b/transform.py
@@ -18,3 +18,21 @@ def rotate(x, phi, n):
     around the axis `n` (when looking towards +infinity).
     """
     return np.inner(np.asarray(x),make_rotation_matrix(phi, n))
+
+def normalize(x):
+    "Returns unit vectors in the direction of `x`."
+    x = np.asarray(x)
+
+    if x.shape[-1] != 3:
+        raise ValueError('dimension of last axis must be 3.')
+
+    d = len(x.shape)
+
+    if d == 1:
+        norm = np.sqrt(x.dot(x))
+    elif d == 2:
+        norm = np.sqrt(np.sum(x*x, axis=1))[:,np.newaxis]
+    else:
+        raise ValueError('len(`x`.shape) must be zero or one.')
+
+    return x/norm