From 9d8965ace99e8d745976cef816036ca7f7fb2c26 Mon Sep 17 00:00:00 2001 From: Sebastian Kloth Date: Thu, 4 Jan 2024 16:25:23 +0100 Subject: [PATCH] Added new function for adding distances from the pore center --- src/mdevaluate/extra/free_energy_landscape.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/src/mdevaluate/extra/free_energy_landscape.py b/src/mdevaluate/extra/free_energy_landscape.py index 189a47b..673b917 100644 --- a/src/mdevaluate/extra/free_energy_landscape.py +++ b/src/mdevaluate/extra/free_energy_landscape.py @@ -103,7 +103,7 @@ def find_maxima( return maxima_df -def calc_energies(maxima_indices, maxima_df, bins): +def calc_energies(maxima_indices, maxima_df, bins, box): points = np.array(maxima_df[["x", "y", "z"]]) tree = KDTree(points, boxsize=box) maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]] @@ -125,7 +125,7 @@ def calc_energies(maxima_indices, maxima_df, bins): energy_hist = np.histogram(current_distances, bins=bins, weights=energy)[0] occupied_bins_hist = np.histogram(current_distances, bins=bins)[0] result = energy_hist / occupied_bins_hist - return r, result + return result for i, maxima_occupation in enumerate(maxima_occupations): current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]] current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]] @@ -141,6 +141,22 @@ def calc_energies(maxima_indices, maxima_df, bins): return result +def add_distances(occupation_df, pore_geometry, origin): + new_df = occupation_df.copy() + if pore_geometry is "cylindrical": + new_df["distance"] = ( + (new_df["x"] - origin[0]) ** 2 + (new_df["y"] - origin[1]) ** 2 + ) ** (1 / 2) + elif pore_geometry is "slit": + new_df["distance"] = np.abs(new_df["z"] - origin[2]) + else: + raise ValueError( + f"pore_geometry is {pore_geometry}, should either be " + f"'cylindrical' or 'slit'" + ) + return new_df + + def get_fel( traj, path,