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
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)
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})