use bottleneck for speed up of nanmean

N720720 · Jan 21, 2022 · 1932edd · 1932edd
1 parent da378eb
commit 1932edd
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/lindemann/index/per_trj.py b/lindemann/index/per_trj.py
@@ -1,5 +1,6 @@
 from typing import Any
 
+import bottleneck as bn
 import numba as nb
 import numpy as np
 import numpy.typing as npt
@@ -61,4 +62,4 @@ def calculate(frames: npt.NDArray[np.float64]) -> float:
     I cant implemnet this in the jit function for now.
     """
 
-    return np.mean(np.nanmean(lindemann_per_atom(frames), axis=1))  # type: ignore[no-any-return, no-untyped-call]
+    return np.mean(bn.nanmean(lindemann_per_atom(frames), axis=1))  # type: ignore[no-any-return, no-untyped-call]