IPKM/mdevaluate
6
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge branch 'mdeval_dev'

Browse Source
This commit is contained in:
Sebastian Kloth 2024-03-05 13:58:15 +01:00
commit 7b9f8b6773
1 changed files with 23 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

49
src/mdevaluate/extra/free_energy_landscape.py
View File

@ -47,6 +47,21 @@ def _pbc_points_reduced(
return coordinates_pbc, indices return coordinates_pbc, indices
def _build_tree(points, box, r_max, pore_geometry):
if np.all(np.diag(np.diag(box)) == box):
tree = KDTree(points, boxsize=box)
points_pbc_index = None
else:
points_pbc, points_pbc_index = _pbc_points_reduced(
points,
pore_geometry,
box,
thickness=r_max + 0.01,
)
tree = KDTree(points_pbc)
return tree, points_pbc_index
def occupation_matrix( def occupation_matrix(
trajectory: Coordinates, trajectory: Coordinates,
edge_length: float = 0.05, edge_length: float = 0.05,
@ -113,23 +128,14 @@ def find_maxima(
maxima_df = occupation_df.copy() maxima_df = occupation_df.copy()
maxima_df["maxima"] = None maxima_df["maxima"] = None
points = np.array(maxima_df[["x", "y", "z"]]) points = np.array(maxima_df[["x", "y", "z"]])
if np.all(np.diag(np.diag(box)) == box): tree, points_pbc_index = _build_tree(points, box, radius, pore_geometry)
tree = KDTree(points, boxsize=box)
all_neighbors = tree.query_ball_point(points, radius)
else:
points_pbc, points_pbc_index = _pbc_points_reduced(
points,
pore_geometry,
box,
thickness=radius + 0.01,
)
tree = KDTree(points_pbc)
all_neighbors = tree.query_ball_point(points, radius)
all_neighbors = points_pbc_index[all_neighbors]
for i in range(len(maxima_df)): for i in range(len(maxima_df)):
if maxima_df.loc[i, "maxima"] is not None: if maxima_df.loc[i, "maxima"] is not None:
continue continue
neighbors = np.array(all_neighbors[i]) maxima_pos = maxima_df.loc[i, ["x", "y", "z"]]
neighbors = np.array(tree.query_ball_point(maxima_pos, radius))
if points_pbc_index is not None:
neighbors = points_pbc_index[neighbors]
neighbors = neighbors[neighbors != i] neighbors = neighbors[neighbors != i]
if len(neighbors) == 0: if len(neighbors) == 0:
maxima_df.loc[i, "maxima"] = True maxima_df.loc[i, "maxima"] = True
@ -154,16 +160,7 @@ def _calc_energies(
nodes: int = 8, nodes: int = 8,
) -> NDArray: ) -> NDArray:
points = np.array(maxima_df[["x", "y", "z"]]) points = np.array(maxima_df[["x", "y", "z"]])
if np.all(np.diag(np.diag(box)) == box): tree, points_pbc_index = _build_tree(points, box, bins[-1], pore_geometry)
tree = KDTree(points, boxsize=box)
else:
points_pbc, points_pbc_index = _pbc_points_reduced(
points,
pore_geometry,
box,
thickness=bins[-1] + 0.01,
)
tree = KDTree(points_pbc)
maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]] maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]]
maxima_occupations = np.array(maxima_df.loc[maxima_indices, "occupation"]) maxima_occupations = np.array(maxima_df.loc[maxima_indices, "occupation"])
num_of_neighbors = np.max( num_of_neighbors = np.max(
@ -187,7 +184,7 @@ def _calc_energies(
all_occupied_bins_hist = [] all_occupied_bins_hist = []
if distances.ndim == 1: if distances.ndim == 1:
current_distances = distances[1:][distances[1:] <= bins[-1]] current_distances = distances[1:][distances[1:] <= bins[-1]]
if np.all(np.diag(np.diag(box)) == box): if points_pbc_index is None:
current_indices = indices[1:][distances[1:] <= bins[-1]] current_indices = indices[1:][distances[1:] <= bins[-1]]
else: else:
current_indices = points_pbc_index[indices[1:][distances[1:] <= bins[-1]]] current_indices = points_pbc_index[indices[1:][distances[1:] <= bins[-1]]]
@ -201,7 +198,7 @@ def _calc_energies(
return result return result
for i, maxima_occupation in enumerate(maxima_occupations): for i, maxima_occupation in enumerate(maxima_occupations):
current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]] current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]]
if np.all(np.diag(np.diag(box)) == box): if points_pbc_index is None:
current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]] current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
else: else:
current_indices = points_pbc_index[ current_indices = points_pbc_index[