#467 #473 Application of the new optimizations to the algorithms of p…

…er_atoms and per_frame

lindemann/index/per_atoms.py

@@ -15,58 +15,37 @@ def calculate(frames: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
     Returns:
         npt.NDArray[np.float32]: Returns 1D array with the progression of the lindeman index per frame of shape(frames, atoms)
     """
+    len_frames, natoms, _ = frames.shape
 
-    first = True
-    # natoms = natoms
-    dt = frames.dtype
-    natoms = len(frames[0])
-    nframes = len(frames)
-    len_frames = len(frames)
-    array_mean = np.zeros((natoms, natoms), dtype=dt)
-    array_var = np.zeros((natoms, natoms), dtype=dt)
-    iframe = dt.type(1)
-    lindex_array = np.zeros((len_frames, natoms), dtype=dt)
-    for q, coords in enumerate(frames):
-        n, p = coords.shape
-        array_distance = np.zeros((n, n), dtype=dt)
-        for i in range(n):
-            for j in range(i + 1, n):
-                d = 0.0
-                for k in range(p):
-                    d += (coords[i, k] - coords[j, k]) ** dt.type(2)
-                array_distance[i, j] = np.sqrt(d)
-        array_distance += array_distance.T
-
-        #################################################################################
-        # update mean and var arrays based on Welford algorithm suggested by Donald Knuth
-        #################################################################################
+    array_mean = np.zeros((natoms, natoms), dtype=np.float32)
+    array_var = np.zeros((natoms, natoms), dtype=np.float32)
+    lindex_array = np.zeros((len_frames, natoms), dtype=np.float32)
+    for frame, coords in enumerate(frames):
+        frame_count = frame + 1
         for i in range(natoms):
             for j in range(i + 1, natoms):
-                xn = array_distance[i, j]
+                dist = 0.0
+                for k in range(3):
+                    dist += (coords[i, k] - coords[j, k]) ** 2
+                dist = np.sqrt(dist)
                 mean = array_mean[i, j]
                 var = array_var[i, j]
-                delta = xn - mean
-                # update mean
-                array_mean[i, j] = mean + delta / iframe
-                # update variance
-                array_var[i, j] = var + delta * (xn - array_mean[i, j])
-        iframe += 1  # type: ignore[assignment]
-        if iframe > nframes + 1:
-            break
-
-        for i in range(natoms):
-            for j in range(i + 1, natoms):
-                array_mean[j, i] = array_mean[i, j]
-                array_var[j, i] = array_var[i, j]
-
-        if first:
-            lindemann_indices = np.zeros((natoms), dtype=dt)
-            first = False
-        else:
-            np.fill_diagonal(array_mean, 1)
-            lindemann_indices = np.zeros((natoms), dtype=dt)
-            lindemann_indices = np.divide(np.sqrt(np.divide(array_var, iframe - 1)), array_mean)
-            lindemann_indices = np.asarray([np.mean(lin[lin != 0]) for lin in lindemann_indices])
-
-        lindex_array[q] = lindemann_indices
+                delta = dist - mean
+                update_mean = mean + delta / frame_count
+                array_mean[i, j] = update_mean
+                array_mean[j, i] = update_mean
+                delta2 = dist - array_mean[i, j]
+                update_var = var + delta * delta2
+                array_var[i, j] = update_var
+                array_var[j, i] = update_var
+
+        np.fill_diagonal(array_mean, 1.0)
+        lindemann_indices = np.divide(
+            np.sqrt(np.divide(array_var, frame_count)), array_mean
+        ).astype(np.float32)
+        lindemann_indices = np.asarray(
+            [np.nanmean(lin[lin != 0]) for lin in lindemann_indices]
+        ).astype(np.float32)
+
+        lindex_array[frame] = lindemann_indices
     return lindex_array

lindemann/index/per_frames.py

@@ -1,72 +1,30 @@
-from typing import List
-
 import numba as nb
 import numpy as np
-import numpy.typing as npt
-
-
-@nb.njit(fastmath=True, error_model="numpy")  # type: ignore # , cache=True) #(parallel=True)
-def calculate(frames: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
-    """calculate the progression of the lindemann index over the frames.
-
-    Args:
-        frames: numpy array of shape(frames,atoms)
-    Returns:
-        npt.NDArray[np.float32]: Returns 1D array with the progression of the lindeman index per frame of shape(frames)
-    """
-
-    first = True
-    dt = frames.dtype
-    natoms = len(frames[0])
-    nframes = len(frames)
-    len_frames = len(frames)
-    array_mean = np.zeros((natoms, natoms), dtype=dt)
-    array_var = np.zeros((natoms, natoms), dtype=dt)
-    iframe = dt.type(1)
-    lindex_array = np.zeros((len_frames), dtype=dt)
-    for q, coords in enumerate(frames):
-        n, p = coords.shape
-        array_distance = np.zeros((n, n), dtype=dt)
-        for i in range(n):
-            for j in range(i + 1, n):
-                d = 0.0
-                for k in range(p):
-                    d += (coords[i, k] - coords[j, k]) ** dt.type(2)
-                array_distance[i, j] = np.sqrt(d)
-        array_distance += array_distance.T
-
-        #################################################################################
-        # update mean and var arrays based on Welford algorithm suggested by Donald Knuth
-        #################################################################################
-        for i in range(natoms):
-            for j in range(i + 1, natoms):
-                xn = array_distance[i, j]
-                mean = array_mean[i, j]
-                var = array_var[i, j]
-                delta = xn - mean
-                # update mean
-                array_mean[i, j] = mean + delta / iframe
-                # update variance
-                array_var[i, j] = var + delta * (xn - array_mean[i, j])
-        iframe += 1  # type: ignore[assignment]
-        if iframe > nframes + 1:
-            break
-
-        for i in range(natoms):
-            for j in range(i + 1, natoms):
-                array_mean[j, i] = array_mean[i, j]
-                array_var[j, i] = array_var[i, j]
 
-        if first:
-            lindemann_indices = 0
-            first = False
-        else:
-            np.fill_diagonal(array_mean, 1)
-            lindemann_indices = np.zeros((natoms), dtype=dt)  # type: ignore[assignment]
-            lindemann_indices = np.divide(np.sqrt(np.divide(array_var, iframe - 1)), array_mean)  # type: ignore[assignment]
-            lindemann_indices = np.mean(
-                np.asarray([np.mean(lin[lin != 0]) for lin in lindemann_indices])  # type: ignore[attr-defined]
-            )
 
-        lindex_array[q] = lindemann_indices
-    return lindex_array
+@nb.njit(fastmath=True, parallel=False)
+def calculate(positions):
+    num_frames, num_atoms, _ = positions.shape
+    num_distances = num_atoms * (num_atoms - 1) // 2
+
+    mean_distances = np.zeros(num_distances, dtype=np.float32)
+    m2_distances = np.zeros(num_distances, dtype=np.float32)
+    linde_per_frame = np.zeros(num_frames, dtype=np.float32)
+    for frame in range(num_frames):
+        index = 0
+        frame_count = frame + 1
+        for i in range(num_atoms):
+            for j in range(i + 1, num_atoms):
+                dist = 0.0
+                for k in range(3):
+                    dist += (positions[frame, i, k] - positions[frame, j, k]) ** 2
+                dist = np.sqrt(dist)
+                delta = dist - mean_distances[index]
+                mean_distances[index] += delta / frame_count
+                delta2 = dist - mean_distances[index]
+                m2_distances[index] += delta * delta2
+
+                index += 1
+        linde_per_frame[frame] = np.mean(np.sqrt(m2_distances / frame_count) / mean_distances)
+
+    return linde_per_frame