DOCS: math

brunobeltran · Jun 3, 2019 · 4fb8776 · 4fb8776
1 parent 8a11519
commit 4fb8776
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nuc_chain/math.py b/nuc_chain/math.py
@@ -4,7 +4,7 @@
 def r2(r, axis=None):
     """Calculate the instantaneous r^2 along a given axis.
 
-    r.shape == (N,3) ==> r2(r, axis=0) would give the r^2"""
+    r.shape == (N,3) ==> r2(r, axis=0) would give the r^2 of shape (N,)"""
     zero = [slice(None)]*len(r.shape) # [:,:,...,:]
     zero[axis] = slice(0, 1) # [:,:,..,0,..,:]
     return np.sum(np.power(r - r[zero], 2), axis=1-axis)