diff --git a/src/mdevaluate/extra/free_energy_landscape.py b/src/mdevaluate/extra/free_energy_landscape.py index 17f1b9d..189a47b 100644 --- a/src/mdevaluate/extra/free_energy_landscape.py +++ b/src/mdevaluate/extra/free_energy_landscape.py @@ -4,15 +4,24 @@ import os.path import numpy as np import math import scipy +from numpy.typing import ArrayLike, NDArray from scipy.spatial import KDTree import cmath import pandas as pd import multiprocessing as mp +from ..coordinates import Coordinates + VALID_GEOMETRY = {"cylindrical", "slab"} -def occupation_matrix(trajectory, edge_length=0.05, segments=1000, skip=0.1, nodes=8): +def occupation_matrix( + trajectory: Coordinates, + edge_length: float = 0.05, + segments: int = 1000, + skip: float = 0.1, + nodes: int = 8, +) -> pd.DataFrame: frame_indices = np.unique( np.int_(np.linspace(len(trajectory) * skip, len(trajectory) - 1, num=segments)) ) @@ -48,7 +57,9 @@ def occupation_matrix(trajectory, edge_length=0.05, segments=1000, skip=0.1, nod return occupation_df -def _calc_histogram(numberlist, trajectory, bins): +def _calc_histogram( + numberlist: ArrayLike, trajectory: Coordinates, bins: ArrayLike +) -> NDArray: matbin = None for index in range(0, len(numberlist), 1000): try: @@ -64,7 +75,9 @@ def _calc_histogram(numberlist, trajectory, bins): return matbin -def find_maxima(occupation_df, box, edge_length=0.05): +def find_maxima( + occupation_df: pd.DataFrame, box: ArrayLike, edge_length: float = 0.05 +) -> pd.DataFrame: maxima_df = occupation_df.copy() maxima_df["maxima"] = None points = np.array(maxima_df[["x", "y", "z"]]) @@ -80,8 +93,8 @@ def find_maxima(occupation_df, box, edge_length=0.05): 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, "occupation"].max() + < maxima_df.loc[i, "occupation"] ): maxima_df.loc[neighbors, "maxima"] = False maxima_df.loc[i, "maxima"] = True @@ -90,6 +103,44 @@ def find_maxima(occupation_df, box, edge_length=0.05): return maxima_df +def calc_energies(maxima_indices, maxima_df, bins): + points = np.array(maxima_df[["x", "y", "z"]]) + tree = KDTree(points, boxsize=box) + maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]] + maxima_occupations = np.array(maxima_df.loc[maxima_indices, "occupation"]) + num_of_neighbors = np.max( + tree.query_ball_point(maxima, bins[-1], return_length=True) + ) + distances, indices = tree.query( + maxima, k=num_of_neighbors, distance_upper_bound=bins[-1] + ) + all_energy_hist = [] + all_occupied_bins_hist = [] + if distances.ndim == 1: + current_distances = distances[1:][distances[1:] <= bins[-1]] + current_indices = indices[1:][distances[1:] <= bins[-1]] + energy = -np.log( + maxima_df.loc[current_indices, "occupation"] / maxima_occupations + ) + 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 + 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]] + + energy = -np.log( + maxima_df.loc[current_indices, "occupation"] / maxima_occupation + ) + energy_hist = np.histogram(current_distances, bins=bins, weights=energy)[0] + occupied_bins_hist = np.histogram(current_distances, bins=bins)[0] + all_energy_hist.append(energy_hist) + all_occupied_bins_hist.append(occupied_bins_hist) + result = np.sum(all_energy_hist, axis=0) / np.sum(all_occupied_bins_hist, axis=0) + return result + + def get_fel( traj, path, @@ -309,7 +360,7 @@ def sphere_quotient( # Distances between maxima and other cubes in the system coordlist = coordlist.reshape(-1, 3) - numOfNeigbour = tree.query_ball_point(maxima[0], unitdist, return_length=True) + numOfNeigbour = np.max(tree.query_ball_point(maxima, unitdist, return_length=True)) d, neighbourlist = tree.query(maxima_masked, k=numOfNeigbour, workers=-1) i = 0 @@ -365,7 +416,7 @@ def sphere_quotient_slab( maxima_masked = np.array(maxima)[mask] coordlist = coordlist.reshape(-1, 3) - numOfNeigbour = tree.query_ball_point(maxima[0], unitdist, return_length=True) + numOfNeigbour = np.max(tree.query_ball_point(maxima, unitdist, return_length=True)) d, neighbourlist = tree.query(maxima_masked, k=numOfNeigbour, workers=-1) i = 0