fixes bug in unitVector

hexrd/core/matrixutil.py

@@ -74,13 +74,25 @@ def rowNorm(a):
 
 
 def unitVector(a):
+    """ normalize array of column vectors (hstacked, axis = 0)
+
+    Parameters
+    ----------
+    a: numpy array
+       array to normalize; shape should be (m, n) or (m,)
+
+    Returns
+    -------
+    anrm: numpy array
+       array with shape (m, n)  with each column normalized to magnitude 1.0
     """
-    normalize array of column vectors (hstacked, axis = 0)
-    """
-    assert a.ndim in [
-        1,
-        2,
-    ], "incorrect arg shape; must be 1-d or 2-d, yours is %d-d" % (a.ndim)
+    if a.ndim not in [1, 2]:
+        raise ValueError(
+            "incorrect arg shape; must be 1-d or 2-d, yours is %d-d" % (a.ndim)
+        )
+
+    if a.ndim == 1:
+        a = a.reshape(len(a), 1)
 
     ztol = constants.ten_epsf
 

tests/core/test_matrixutil.py

@@ -58,6 +58,15 @@ def test_unitVector(random_vectors):
     res_z = matrixutil.unitVector(zeros)
     np.testing.assert_allclose(res_z, 0.0)
 
+    # Test on a single 1D vector.
+    x = np.array([1.0, 2.0, 3.0])
+    u = matrixutil.unitVector(x)
+    assert u.shape == (3, 1)
+
+    # Test dimensions of input.
+    x = np.zeros((2, 3, 4))
+    with pytest.raises(ValueError):
+        u = matrixutil.unitVector(x)
 
 # --- Null Space ---