Reworked while loops to for loops in free_energy_landscape.py

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:

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":
+        elif geometry == "slab":
-            bins, data = shortAllRadiiSlab(
+            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:
+    for i in range(max_i):
-        while j < max_j:
+        for j in range(max_j):
-            while k < max_k:
+            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
+    for i in range(max_i):
-    j = 0
+        for j in range(max_j):
-    k = 0
+            for k in range(max_k):
-    while i < max_i:
-        while j < max_j:
-            while k < 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):
+def sphere_quotient(
-    coordlist = np.reshape(coord, (-1, 3))
-    tree = sp.cKDTree(coordlist, boxsize=box)
-    return tree
-
-
-def SphereQuotient(
     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)
+    coordinates = list_of_coordinates(matbin)
-    tree = calculateTree(coordinates, box=matbin.shape)
+    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)
+    coordinates = list_of_coordinates(matbin)
-    tree = calculateTree(coordinates, box=matbin.shape)
+    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

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