Reworked while loops to for loops in free_energy_landscape.py
This commit is contained in:
parent
fa2e793354
commit
4efb221a80
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user