diff --git a/lindemann/index/per_atoms.py b/lindemann/index/per_atoms.py index 8ec6436..81d3f19 100644 --- a/lindemann/index/per_atoms.py +++ b/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 diff --git a/lindemann/index/per_frames.py b/lindemann/index/per_frames.py index d54ca63..25c8175 100644 --- a/lindemann/index/per_frames.py +++ b/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