diff --git a/src/mdevaluate/extra/free_energy_landscape.py b/src/mdevaluate/extra/free_energy_landscape.py index 429f951..17f1b9d 100644 --- a/src/mdevaluate/extra/free_energy_landscape.py +++ b/src/mdevaluate/extra/free_energy_landscape.py @@ -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,