Repository: glotzerlab/rowan

Issue:
to_axis_angle not working for multiple quaternions?
I have several quaternions and want to get axes/angles. This is a traceback of what happens. I'm using rowan v1.2.0.
```
>>> rowan.to_axis_angle(rowan.normalize([[0.707, 0, 0, 0.707], [0, 0, 0, 1]]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/bdice/anaconda3/envs/dice/lib/python3.6/site-packages/rowan/functions.py", line 1035, in to_axis_angle
    q[..., 1:]/sines,
ValueError: operands could not be broadcast together with shapes (2,3) (2,)
```

Discussion:

Patch:
diff --git a/rowan/functions.py b/rowan/functions.py
index 2f6a0c4..8a960dc 100644
--- a/rowan/functions.py
+++ b/rowan/functions.py
@@ -1031,8 +1031,8 @@ def to_axis_angle(q):
     sines = np.sin(angles/2)
     # Avoid divide by zero issues; these values will not be used
     sines[sines == 0] = 1
-    axes = np.where(angles != 0,
-                    q[..., 1:]/sines,
+    axes = np.where(angles[..., np.newaxis] != 0,
+                    q[..., 1:]/sines[..., np.newaxis],
                     0)
 
     return axes, angles


Test patch:
diff --git a/tests/test_axis_angle.py b/tests/test_axis_angle.py
index 3387061..b81cfa0 100644
--- a/tests/test_axis_angle.py
+++ b/tests/test_axis_angle.py
@@ -85,8 +85,19 @@ class TestFromAxisAngle(unittest.TestCase):
 class TestToAxisAngle(unittest.TestCase):
     """Test converting to axis angle representation"""
 
+    def test_div_zero(self):
+        q = (1, 0, 0, 0)
+        axes, angles = rowan.to_axis_angle(q)
+        self.assertTrue(np.allclose(axes, [0, 0, 0]))
+        self.assertEqual(angles, 0)
+
     def test_to_axis_angle(self):
-        axes, angles = rowan.to_axis_angle(
-            np.array((np.sqrt(2)/2, np.sqrt(2)/2, 0, 0)))
+        q = (np.sqrt(2)/2, np.sqrt(2)/2, 0, 0)
+        axes, angles = rowan.to_axis_angle(q)
         self.assertTrue(np.allclose(axes, np.array([1, 0, 0])))
         self.assertTrue(np.allclose(angles, np.pi/2))
+
+        q2 = np.stack((q, q), axis=0)
+        axes, angles = rowan.to_axis_angle(q2)
+        self.assertTrue(np.allclose(axes, np.array([[1, 0, 0], [1, 0, 0]])))
+        self.assertTrue(np.allclose(angles, [np.pi/2, np.pi/2]))