diff --git a/amep/cluster.py b/amep/cluster.py index e4dced0..2052b10 100644 --- a/amep/cluster.py +++ b/amep/cluster.py @@ -940,7 +940,8 @@ def identify( sorted_clusters : list List of lists, where each list contains the indices of the particles that belong to the same cluster. The list is sorted by the number of - particles in each cluster. + particles in each cluster. Single particles are also included as + clusters in this list. idx : numpy.ndarray Array of shape (N,) containing the cluster ID for each particle. N is the total number of particles. @@ -1036,18 +1037,18 @@ def identify( particle_idx = set(np.arange(len(coords))) clustered_idx = set() - n = 0 + cluster_id = 0 for cl in sorted_clusters: - for i in cl: - idx[i] = n + idx[i] = cluster_id clustered_idx.add(i) - - n += 1 + cluster_id += 1 # add leftover single particles for k in particle_idx - clustered_idx: - idx[k] = n - n += 1 + idx[k] = cluster_id + sorted_clusters.append([k]) + cluster_id += 1 + idx=np.array(idx, dtype=int) return sorted_clusters, idx