Added new function for adding distances from the pore center

This commit is contained in:
Sebastian Kloth 2024-01-04 16:25:23 +01:00
parent ec3f6b1c67
commit 9d8965ace9

View File

@ -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,