diff --git a/src/mdevaluate/distribution.py b/src/mdevaluate/distribution.py index f2d1982..d3c113e 100644 --- a/src/mdevaluate/distribution.py +++ b/src/mdevaluate/distribution.py @@ -2,6 +2,7 @@ import numpy as np from numpy.typing import ArrayLike from scipy import spatial from scipy.spatial import KDTree +from scipy.sparse.csgraph import connected_components from .coordinates import rotate_axis, polar_coordinates, Coordinates, CoordinateFrame from .atoms import next_neighbors @@ -432,3 +433,19 @@ def hbonds( ) else: return pairs[is_bond] + +def calc_cluster_sizes(frame, r_max=0.35): + frame_PBC, indices_PBC = pbc_points( + frame, frame.box, thickness=r_max + 0.1, index=True + ) + tree = KDTree(frame_PBC) + matrix = tree.sparse_distance_matrix(tree, r_max, output_type="ndarray") + new_matrix = np.zeros((len(frame), len(frame))) + for entry in matrix: + if entry[2] > 0: + new_matrix[indices_PBC[entry[0]], indices_PBC[entry[1]]] = 1 + n_components, labels = connected_components(new_matrix, directed=False) + cluster_sizes = [] + for i in range(0, np.max(labels) + 1): + cluster_sizes.append(np.sum(labels == i)) + return np.array(cluster_sizes).flatten() diff --git a/src/mdevaluate/utils.py b/src/mdevaluate/utils.py index 42ad521..e2d3d29 100644 --- a/src/mdevaluate/utils.py +++ b/src/mdevaluate/utils.py @@ -498,5 +498,5 @@ def timing(function): time_needed = end_time - start_time print(f"Finished in {int(time_needed // 60)} min " f"{int(time_needed % 60)} s") return result + return wrap - return wrap \ No newline at end of file