Added new function for finding maxima of occupation

This commit is contained in:
Sebastian Kloth 2023-12-29 16:59:56 +01:00
parent f48dff2ede
commit 62d05423c5

View File

@ -44,7 +44,7 @@ def occupation_matrix(trajectory, edge_length=0.05, segments=1000, skip=0.1, nod
occupation_df = pd.DataFrame(
{"x": coords[0], "y": coords[1], "z": coords[2], "occupation": matbin_new}
)
occupation_df = occupation_df.query("occupation != 0")
occupation_df = occupation_df.query("occupation != 0").reset_index(drop=True)
return occupation_df
@ -64,6 +64,32 @@ def _calc_histogram(numberlist, trajectory, bins):
return matbin
def find_maxima(occupation_df, box, edge_length=0.05):
maxima_df = occupation_df.copy()
maxima_df["maxima"] = None
points = np.array(maxima_df[["x", "y", "z"]])
tree = KDTree(points, boxsize=box)
all_neighbors = tree.query_ball_point(
points, edge_length * 3 ** (1 / 2) + edge_length / 100
)
for i in range(len(maxima_df)):
if maxima_df.loc[i, "maxima"] is not None:
continue
neighbors = np.array(all_neighbors[i])
neighbors = neighbors[neighbors != i]
if len(neighbors) == 0:
maxima_df.loc[i, "maxima"] = True
elif (
maxima_df.loc[neighbors, "occupation"].max()
< maxima_df.loc[i, "occupation"]
):
maxima_df.loc[neighbors, "maxima"] = False
maxima_df.loc[i, "maxima"] = True
else:
maxima_df.loc[i, "maxima"] = False
return maxima_df
def get_fel(
traj,
path,