#476 Memory optimization using OVITO pipeline for lindemann index per frame

Author: N720720

lindemann/index/mem_use.py

@@ -31,8 +31,9 @@ def in_gb(nframes: int, natoms: int) -> str:
         f"\nFlag -t (per_trj) will use {np.round((trj+num_distances*2*float_size)/1024**3,4)} GB\n"
     )
     online_per_trj = (
-        f"\nFlag -ot (per_trj) will use {np.round((num_distances*2*float_size)/1024**3,4)} GB\n"
+        f"Flag -ot (per_trj) will use {np.round((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"
+    online_per_frames = f"Flag -of (per_frames) will use {np.round(((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}{online_per_trj}{per_frames}{per_atoms}"
+    return f"{per_trj}{online_per_trj}{per_frames}{online_per_frames}{per_atoms}"

lindemann/index/online_frames.py

@@ -0,0 +1,50 @@
+from typing import Optional
+
+import numba as nb
+import numpy as np
+from ovito.data import DataCollection
+from ovito.pipeline import Pipeline
+
+import lindemann
+from lindemann.trajectory import read
+
+
+@nb.njit(fastmath=True, parallel=False)
+def calculate_frame(positions, mean_distances, m2_distances, frame: int, num_atoms: int):
+
+    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[i, k] - positions[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
+    return np.mean(np.sqrt(m2_distances / frame_count) / mean_distances)
+
+
+def calculate(pipeline: Pipeline, data: DataCollection, nframes: Optional[int] = None):
+    num_particle = data.particles.count
+    num_frame = pipeline.source.num_frames
+    if nframes is None:
+        nframes = num_frame
+    elif nframes > num_frame:
+        raise ValueError(f"Requested {nframes} frames, but only {num_frame} frames are available.")
+
+    num_distances = num_particle * (num_particle - 1) // 2
+    mean_distances = np.zeros(num_distances, dtype=np.float32)
+    m2_distances = np.zeros(num_distances, dtype=np.float32)
+    lindemann_index_array = np.zeros(nframes, dtype=np.float32)
+    for frame in range(nframes):
+        data = pipeline.compute(frame)
+        lindemann_index_array[frame] = calculate_frame(
+            data.particles["Position"].array, mean_distances, m2_distances, frame, num_particle
+        )
+    return lindemann_index_array