add get_perp() to return an arbitrary perpendicular vector and a new rotate() function which is faster and can handle different angles for each vector.

1 files changed, 21 insertions, 7 deletions

diff --git a/chroma/transform.py b/chroma/transform.py
index 76b1939..580dbf4 100644
--- a/chroma/transform.py
+++ b/chroma/transform.py
@@ -1,5 +1,11 @@
 import numpy as np
 
+def get_perp(x):
+    """Returns an arbitrary vector perpendicular to `x`."""
+    a = np.zeros(3)
+    a[np.argmin(abs(x))] = 1
+    return np.cross(a,x)
+
 def make_rotation_matrix(phi, n):
     """
     Make the rotation matrix to rotate points through an angle `phi`
@@ -9,11 +15,8 @@ def make_rotation_matrix(phi, n):
 
     Raises ValueError if n has zero magnitude
     """
-    norm = np.linalg.norm(n)
-    if norm == 0.0:
-        raise ValueError('rotation axis has zero magnitude')
-    n = np.asarray(n)/norm
-    
+    n = normalize(n)
+
     return np.cos(phi)*np.identity(3) + (1-np.cos(phi))*np.outer(n,n) + \
         np.sin(phi)*np.array([[0,n[2],-n[1]],[-n[2],0,n[0]],[n[1],-n[0],0]])
 
@@ -22,11 +25,22 @@ def rotate(x, phi, n):
     Rotate an array of points `x` through an angle phi counter-clockwise
     around the axis `n` (when looking towards +infinity).
     """
-    return np.inner(np.asarray(x),make_rotation_matrix(phi, n))
+    n = normalize(n)
+    x = np.atleast_2d(x)
+    phi = np.atleast_1d(phi)
+
+    return (x*np.cos(phi)[:,np.newaxis] + n*np.dot(x,n)[:,np.newaxis]*(1-np.cos(phi)[:,np.newaxis]) + np.cross(x,n)*np.sin(phi)[:,np.newaxis]).squeeze()
+
+def rotate_matrix(x, phi, n):
+    """Same as rotate() except uses a rotation matrix.
+
+    Can only handle a single rotation angle `phi`."""
+    rotation_matrix = make_rotation_matrix(phi, n)
+    return np.inner(np.asarray(x),rotation_matrix)
 
 def normalize(x):
     "Returns unit vectors in the direction of `x`."
-    x = np.asarray(x)
+    x = np.asarray(x, dtype=float)
 
     if x.shape[-1] != 3:
         raise ValueError('dimension of last axis must be 3.')