Merge pull request #475 from N720720/467-optimize-lindemann-index-cal…

…culation-using-distance-vector

467 optimize lindemann index calculation using distance vector

benchmarking/synthetic_benchmark.py

@@ -3,15 +3,16 @@
 import numpy as np
 
 from lindemann.index import per_atoms, per_frames, per_trj
+from lindemann.index.mem_use import in_gb
 
 
 def generate_test_data(num_frames, num_atoms):
     rng = np.random.default_rng(seed=42)
-    return rng.random((num_frames, num_atoms, 3)).astype(np.float32)
+    return rng.random((num_frames, num_atoms, 3), dtype=np.float32)
 
 
 def benchmark(function, positions, iterations=3):
-    index = function(positions)  # Warm-up
+    index = function(generate_test_data(10, 10))  # Warm-up
     times = []
     for _ in range(iterations):
         start_time = time.time()
@@ -24,8 +25,9 @@ def benchmark(function, positions, iterations=3):
 def main():
     num_frames = 5000
     num_atoms = 1103
+    print(in_gb(num_frames, num_atoms))
     positions = generate_test_data(num_frames, num_atoms)
-    iterations = 1
+    iterations = 3
 
     mean_time, std_time, index = benchmark(per_trj.calculate, positions, iterations=iterations)
     print(f"Mean time of {iterations} runs: {mean_time:.6f} s ± {std_time:.6f} s")

lindemann/index/mem_use.py

@@ -2,15 +2,34 @@
 import numpy.typing as npt
 
 
-def in_gb(frames: npt.NDArray[np.float64]) -> str:
-    """Shows the size of the array in memory in GB.
+def in_gb(nframes: int, natoms: int) -> str:
+    """
+    Calculates and shows the size of the memory allocations related to
+    the different flag options in gigabytes (GB).
 
     Args:
-        frames (npt.NDArray[np.float64]): numpy array of shape(frames,atoms)
+        nframes (int): The number of frames in the trajectory.
+        natoms (int): The number of atoms per frame in the trajectory.
 
     Returns:
-        str: Size of array in GB.
+        str: A formatted string containing the memory usage for different configurations:
+             - per_trj: Memory required when the `-t` flag is used.
+             - per_frames: Memory required when the `-f` flag is used.
+             - per_atoms: Memory required when the `-a` flag is used.
+
+    This function assumes memory calculations based on numpy's float32 data type.
     """
-    natoms = len(frames[0])
-    nframes = len(frames)
-    return f"This will use {np.round((np.zeros((natoms, natoms)).nbytes/1024**3),4)} GB"  # type: ignore[no-untyped-call]
+
+    num_distances = natoms * (natoms - 1) // 2
+    float_size = np.float32().nbytes
+    trj = nframes * natoms * 3 * float_size
+    atom_atom_array = 3 * natoms * natoms * float_size
+    atom_array = natoms * float_size
+    linde_index = nframes * natoms * float_size
+    sum_bytes = trj + atom_atom_array + atom_array + linde_index
+    per_trj = (
+        f"\nFlag -t (per_trj) will use {np.round((trj+num_distances*2*float_size)/1024**3,4)} GB\n"
+    )
+    per_frames = f"Flag -f (per_frames) will use {np.round((trj+(num_distances*2*float_size)+(nframes*float_size))/1024**3,4)} GB\n"
+    per_atoms = f"Flag -a (per_atoms) will use {np.round(sum_bytes/1024**3,4)} GB"
+    return f"{per_trj}{per_frames}{per_atoms}"

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

lindemann/index/per_trj.py

@@ -1,64 +1,29 @@
-from typing import Any
-
 import numba as nb
 import numpy as np
-import numpy.typing as npt
-
-# No typing for jit functions https://github.com/numba/numba/issues/7424
-
-
-@nb.njit(fastmath=True, error_model="numpy")  # type: ignore
-def lindemann_per_atom(frames: npt.NDArray[np.float32]) -> Any:
-    """Calculate the lindeman index
-    Args:
-        frames: numpy array of shape(frames,atoms)
-    Returns:
-        float32: returns the lindeman index
-    """
-
-    dt = frames.dtype
-    natoms = len(frames[0])
-    nframes = len(frames)
-    array_mean = np.zeros((natoms, natoms), dtype=dt)
-    array_var = np.zeros((natoms, natoms), dtype=dt)
-    iframe = dt.type(1)
-    for coords in frames:
-
-        # here we do someting similar to scipy's spatial.distance.pdist scipy.spatial.distance.pdist
-        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 = dt.type(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
-                array_mean[i, j] = mean + delta / iframe
-                array_var[i, j] = var + delta * (xn - array_mean[i, j])
-        iframe += 1.0  # type: ignore[assignment]
-        if iframe > nframes:
-            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]
+@nb.njit(fastmath=True)
+def calculate(positions):
+    num_frames, num_atoms, _ = positions.shape
+    num_distances = num_atoms * (num_atoms - 1) // 2
 
-    lindemann_indices = np.divide(np.sqrt(np.divide(array_var, nframes)), array_mean)
-    return lindemann_indices
+    mean_distances = np.zeros(num_distances, dtype=np.float32)
+    m2_distances = np.zeros(num_distances, 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
 
-def calculate(frames: npt.NDArray[np.float64]) -> float:
+                index += 1
 
-    return np.mean(np.nanmean(lindemann_per_atom(frames), axis=1))  # type: ignore[no-any-return, no-untyped-call]
+    return np.mean(np.sqrt(m2_distances / num_frames) / mean_distances)

lindemann/main.py

@@ -156,9 +156,8 @@ def main(
         typer.Exit()
 
     elif timeit and single_process:
-        # we use float32 here since float64 is not needed for my purposes and it enables us to use nb fastmath. Change to np.float64 if you need more precision.
         start = time.time()
-        linde_for_time = per_trj.calculate(tjr_frames.astype(np.float32))
+        linde_for_time = per_trj.calculate(tjr_frames)
         time_diff = time.time() - start
 
         console.print(
@@ -170,8 +169,8 @@ def main(
         typer.Exit()
 
     elif mem_useage and single_process:
-
-        mem_use_in_gb = mem_use.in_gb(tjr_frames)
+        nframes, natoms, _ = tjr_frames.shape
+        mem_use_in_gb = mem_use.in_gb(nframes, natoms)
 
         console.print(f"[magenta]memory use:[/] [bold blue]{mem_use_in_gb}[/]")
         typer.Exit()