From c09549902a06f3241bafa2d6c7a61e2e9cf7b15f Mon Sep 17 00:00:00 2001 From: Sebastian Kloth Date: Mon, 27 May 2024 14:27:09 +0200 Subject: [PATCH] Added Robins adjustments for FEL --- src/mdevaluate/extra/free_energy_landscape.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/src/mdevaluate/extra/free_energy_landscape.py b/src/mdevaluate/extra/free_energy_landscape.py index ffb01f0..9c3dc44 100644 --- a/src/mdevaluate/extra/free_energy_landscape.py +++ b/src/mdevaluate/extra/free_energy_landscape.py @@ -4,7 +4,6 @@ from typing import Optional import numpy as np from numpy.typing import ArrayLike, NDArray from numpy.polynomial.polynomial import Polynomial as Poly -import math from scipy.spatial import KDTree import pandas as pd import multiprocessing as mp @@ -49,7 +48,7 @@ def _pbc_points_reduced( def _build_tree(points, box, r_max, pore_geometry): if np.all(np.diag(np.diag(box)) == box): - tree = KDTree(points, boxsize=box) + tree = KDTree(points % box, boxsize=box) points_pbc_index = None else: points_pbc, points_pbc_index = _pbc_points_reduced( @@ -79,8 +78,7 @@ def occupation_matrix( z_bins = np.arange(0, box[2][2] + edge_length, edge_length) bins = [x_bins, y_bins, z_bins] # Trajectory is split for parallel computing - size = math.ceil(len(frame_indices) / nodes) - indices = [frame_indices[i : i + size] for i in range(0, len(frame_indices), size)] + indices = np.array_split(frame_indices, nodes) pool = mp.Pool(nodes) results = pool.map( partial(_calc_histogram, trajectory=trajectory, bins=bins), indices @@ -274,7 +272,11 @@ def distance_resolved_energies( def find_energy_maxima( - energy_df: pd.DataFrame, r_min: float, r_max: float + energy_df: pd.DataFrame, + r_min: float, + r_max: float, + r_eval: float = None, + degree: int = 2, ) -> pd.DataFrame: distances = [] energies = [] @@ -283,6 +285,9 @@ def find_energy_maxima( x = np.array(data_d["r"]) y = np.array(data_d["energy"]) mask = (x >= r_min) * (x <= r_max) - p3 = Poly.fit(x[mask], y[mask], deg=2) - energies.append(np.max(p3(np.linspace(r_min, r_max, 1000)))) + p3 = Poly.fit(x[mask], y[mask], deg=degree) + if r_eval is None: + energies.append(np.max(p3(np.linspace(r_min, r_max, 1000)))) + else: + energies.append(p3(r_eval)) return pd.DataFrame({"d": distances, "energy": energies})