Added Robins adjustments for FEL

This commit is contained in:
Sebastian Kloth 2024-05-27 14:27:09 +02:00
parent b7bb8cb379
commit c09549902a

View File

@ -4,7 +4,6 @@ from typing import Optional
import numpy as np import numpy as np
from numpy.typing import ArrayLike, NDArray from numpy.typing import ArrayLike, NDArray
from numpy.polynomial.polynomial import Polynomial as Poly from numpy.polynomial.polynomial import Polynomial as Poly
import math
from scipy.spatial import KDTree from scipy.spatial import KDTree
import pandas as pd import pandas as pd
import multiprocessing as mp import multiprocessing as mp
@ -49,7 +48,7 @@ def _pbc_points_reduced(
def _build_tree(points, box, r_max, pore_geometry): def _build_tree(points, box, r_max, pore_geometry):
if np.all(np.diag(np.diag(box)) == box): if np.all(np.diag(np.diag(box)) == box):
tree = KDTree(points, boxsize=box) tree = KDTree(points % box, boxsize=box)
points_pbc_index = None points_pbc_index = None
else: else:
points_pbc, points_pbc_index = _pbc_points_reduced( 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) z_bins = np.arange(0, box[2][2] + edge_length, edge_length)
bins = [x_bins, y_bins, z_bins] bins = [x_bins, y_bins, z_bins]
# Trajectory is split for parallel computing # Trajectory is split for parallel computing
size = math.ceil(len(frame_indices) / nodes) indices = np.array_split(frame_indices, nodes)
indices = [frame_indices[i : i + size] for i in range(0, len(frame_indices), size)]
pool = mp.Pool(nodes) pool = mp.Pool(nodes)
results = pool.map( results = pool.map(
partial(_calc_histogram, trajectory=trajectory, bins=bins), indices partial(_calc_histogram, trajectory=trajectory, bins=bins), indices
@ -274,7 +272,11 @@ def distance_resolved_energies(
def find_energy_maxima( 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: ) -> pd.DataFrame:
distances = [] distances = []
energies = [] energies = []
@ -283,6 +285,9 @@ def find_energy_maxima(
x = np.array(data_d["r"]) x = np.array(data_d["r"])
y = np.array(data_d["energy"]) y = np.array(data_d["energy"])
mask = (x >= r_min) * (x <= r_max) mask = (x >= r_min) * (x <= r_max)
p3 = Poly.fit(x[mask], y[mask], deg=2) p3 = Poly.fit(x[mask], y[mask], deg=degree)
energies.append(np.max(p3(np.linspace(r_min, r_max, 1000)))) 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}) return pd.DataFrame({"d": distances, "energy": energies})