diff --git a/src/mdevaluate/coordinates.py b/src/mdevaluate/coordinates.py index c8888a0..bd16541 100755 --- a/src/mdevaluate/coordinates.py +++ b/src/mdevaluate/coordinates.py @@ -360,6 +360,7 @@ class CoordinatesMap: self._description = self.function.func.__name__ else: self._description = self.function.__name__ + self._slice = slice(None) def __iter__(self): for frame in self.coordinates: diff --git a/src/mdevaluate/free_energy_landscape.py b/src/mdevaluate/free_energy_landscape.py index 726d48e..fc1753a 100644 --- a/src/mdevaluate/free_energy_landscape.py +++ b/src/mdevaluate/free_energy_landscape.py @@ -4,7 +4,7 @@ import os.path import numpy as np import math import scipy -from scipy import spatial as sp +from scipy.spatial import KDTree import cmath import pandas as pd import multiprocessing as mp @@ -52,7 +52,6 @@ def get_fel( """ if geometry not in VALID_GEOMETRY: raise ValueError("results: status must be one of %r." % VALID_GEOMETRY) - EnergyDifference = [] if (os.path.exists(f"{path}/radiiData.npy")) & (not overwrite): data = np.load(f"{path}/radiiData.npy") @@ -60,17 +59,19 @@ def get_fel( # Here the different standard geometries are inserted else: if geometry == "cylindrical": - bins, data = shortAllRadii( + bins, data = short_all_radii( traj, path, edge=edge, radiusmin=radiusmin, radiusmax=radiusmax, z=z ) - if geometry == "slab": - bins, data = shortAllRadiiSlab( + elif geometry == "slab": + bins, data = short_all_radii_slab( traj, path, edge=edge, radiusmin=radiusmin, radiusmax=radiusmax, z=z ) np.save(f"{path}/radiiData", data) np.save(f"{path}/radiiBins", bins) - return CalculateEnergyDifference(data, bins, temperature) + energy_differences = np.array(calculate_energy_difference(data, bins, temperature)) + r = bins[1:] - (bins[1] - bins[0]) / 2 + return r, energy_differences def fill_bins(traj, path, edge=0.05): @@ -83,7 +84,7 @@ def fill_bins(traj, path, edge=0.05): size = math.ceil(len(traj) / 8) # Trajectory is split for parallel computing - indices = list(Chunksplit(np.arange(0, len(traj), 1), size)) + indices = list(chunksplit(np.arange(0, len(traj), 1), size)) # indices = list(Chunksplit(np.arange(len(traj)-80, len(traj), 1), size)) fill = partial(help_fill, traj=traj) results = pool.map(fill, indices) @@ -153,9 +154,6 @@ def help_fill(numberlist, traj, edge=0.05): def calculate_maxima(matbin): - i = 0 - j = 0 - k = 0 maxima = [] z = [-1, 0, 1] @@ -164,13 +162,13 @@ def calculate_maxima(matbin): max_j = matbin.shape[1] max_k = matbin.shape[2] - # For each element in the matrix all sourrounding 26 elements are compared to + # For each element in the matrix all surrounding 26 elements are compared to # determine the minimum. - # Algorithm can definitely improved but is so fast that its not necessary. + # Algorithm can definitely be improved but is so fast that it's not necessary. - while i < max_i: - while j < max_j: - while k < max_k: + for i in range(max_i): + for j in range(max_j): + for k in range(max_k): a = matbin[i][j][k] b = True for l in z: @@ -186,16 +184,10 @@ def calculate_maxima(matbin): if b: maxima.append([i, j, k]) - - k += 1 - k = 0 - j += 1 - j = 0 - i += 1 return maxima -def ListOfCoordinates(matbin): +def list_of_coordinates(matbin): # Matrix elements are translated back into postion vectors max_i = matbin.shape[0] max_j = matbin.shape[1] @@ -203,28 +195,15 @@ def ListOfCoordinates(matbin): coord = np.zeros((max_i, max_j, max_k, 3)) - i = 0 - j = 0 - k = 0 - while i < max_i: - while j < max_j: - while k < max_k: + for i in range(max_i): + for j in range(max_j): + for k in range(max_k): coord[i][j][k] = [i, j, k] - k += 1 - k = 0 - j += 1 - j = 0 - i += 1 + return coord -def calculateTree(coord, box): - coordlist = np.reshape(coord, (-1, 3)) - tree = sp.cKDTree(coordlist, boxsize=box) - return tree - - -def SphereQuotient( +def sphere_quotient( matbin, maxima, coordlist, @@ -293,7 +272,7 @@ def SphereQuotient( return b * edge, quotient -def SphereQuotientSlab( +def sphere_quotient_slab( matbin, maxima, coordlist, tree, radius, box, edge, unitdist, z=[-np.inf, np.inf] ): # Same as SphereQuotient bur for Slabs @@ -349,18 +328,18 @@ def SphereQuotientSlab( return z * edge, quotient -def shortAllRadii( +def short_all_radii( traj, path, edge=0.05, radiusmin=0.05, radiusmax=2.05, z=[-np.inf, np.inf] ): # Shorthand function cylindrical systems matbin = fill_bins(traj, path, edge) maxima = calculate_maxima(matbin) - coordinates = ListOfCoordinates(matbin) - tree = calculateTree(coordinates, box=matbin.shape) + coordinates = list_of_coordinates(matbin) + tree = KDTree(np.reshape(coordinates, (-1, 3)), boxsize=matbin.shape) bins = [] data = [] for radius in np.arange(radiusmin, radiusmax, 0.1): - b, d = SphereQuotient( + b, d = sphere_quotient( matbin, maxima, coordinates, @@ -376,18 +355,18 @@ def shortAllRadii( return bins, data -def shortAllRadiiSlab( +def short_all_radii_slab( traj, path, edge=0.05, radiusmin=0.05, radiusmax=2.05, z=[-np.inf, np.inf] ): # Shorthand function for Slab systems matbin = fill_bins(traj, path, edge) maxima = calculate_maxima(matbin) - coordinates = ListOfCoordinates(matbin) - tree = calculateTree(coordinates, box=matbin.shape) + coordinates = list_of_coordinates(matbin) + tree = KDTree(np.reshape(coordinates, (-1, 3)), boxsize=matbin.shape) bins = [] data = [] for radius in np.arange(radiusmin, radiusmax, 0.1): - c, d = SphereQuotientSlab( + c, d = sphere_quotient_slab( matbin, maxima, coordinates, @@ -403,7 +382,7 @@ def shortAllRadiiSlab( return bins, data -def CalculateEnergyDifference(data, bins, temperature): +def calculate_energy_difference(data, bins, temperature): """ Calculates the energy difference between local energy minimum and the minimum in the center @@ -413,39 +392,39 @@ def CalculateEnergyDifference(data, bins, temperature): while i < len(data): # q = ( - GetMinimum(data[0], bins[0])[1] - GetMinimum(data[i], bins[0])[1] + get_minimum(data[0], bins[0])[1] - get_minimum(data[i], bins[0])[1] ) * temperature difference.append(q) i += 1 return difference -def Chunksplit(list_a, chunk_size): +def chunksplit(list_a, chunk_size): for i in range(0, len(list_a), chunk_size): yield list_a[i : i + chunk_size] -def GetMinimum(data, bins): +def get_minimum(data, bins): # Fits a polynom of order 3 to determine the energy minimum and analytically # calculates the minimum location y = data[:10] x = bins[1:11] - popt, pcov = scipy.optimize.curve_fit(Pol3, x, y, maxfev=80000) + popt, _ = scipy.optimize.curve_fit(pol3, x, y, maxfev=80000) - a, b = SolveQuadratic(3 * popt[0], 2 * popt[1], popt[2]) + a, b = solve_quadratic(3 * popt[0], 2 * popt[1], popt[2]) if 6 * popt[0] * a + 2 * popt[1] > 0: if (np.real(a) < 0) or (np.real(a) > 0.3): return np.real(a), np.average(y) - return np.real(a), np.real(Pol3(a, *popt)) + return np.real(a), np.real(pol3(a, *popt)) else: if (np.real(b) < 0) or (np.real(b) > 0.3): return np.real(b), np.average(y) - return np.real(b), np.real(Pol3(b, *popt)) + return np.real(b), np.real(pol3(b, *popt)) -def SolveQuadratic(a, b, c): +def solve_quadratic(a, b, c): d = (b**2) - (4 * a * c) sol1 = (-b - cmath.sqrt(d)) / (2 * a) @@ -453,5 +432,5 @@ def SolveQuadratic(a, b, c): return sol1, sol2 -def Pol3(x, a, b, c, d): +def pol3(x, a, b, c, d): return a * x**3 + b * x**2 + c * x + d diff --git a/test/test_free_energy_landscape.py b/test/test_free_energy_landscape.py index dc51631..843bcb5 100644 --- a/test/test_free_energy_landscape.py +++ b/test/test_free_energy_landscape.py @@ -39,7 +39,7 @@ def test_get_fel(trajectory): ) oxygens_water = trajectory.subset(atom_name="OW", residue_name="SOL") - data = fel.get_fel( + r, energy_differences = fel.get_fel( oxygens_water, os.path.join(os.path.dirname(__file__), 'data/pore'), "cylindrical", @@ -51,4 +51,4 @@ def test_get_fel(trajectory): overwrite=True, ) - assert (np.array(data) == test_array).all() + assert (energy_differences == test_array).all()