Added derive_surface and get_positions with tests. (#40)

* Refactored code with tests.

membrane_curvature/core.py

@@ -5,7 +5,7 @@
 Handles the primary functions
 """
 
-from .lib.mods import core_fast_leaflet, gaussian_curvature, mean_curvature
+from .lib.mods import derive_surface, gaussian_curvature, mean_curvature
 import time
 import MDAnalysis as mda
 import math
@@ -32,7 +32,7 @@ def main():
 
     # 4. Save pickles zpo4
     z_Ref = np.zeros([n_cells, n_cells])
-    z_coords = core_fast_leaflet(u, z_Ref, n_cells, selection, max_width)
+    z_coords = derive_surface(n_cells, selection, max_width)
 
     # 5. Calculate curvature
     K = gaussian_curvature(z_coords)

membrane_curvature/lib/mods.py

@@ -1,89 +1,110 @@
-import itertools as it
 import numpy as np
 
 
-def grid_map(coords, factor):
-    """ Maps (x,y) coordinates to unit cell in grid.
+def derive_surface(n_cells_x, n_cells_y, selection, max_width_x, max_width_y):
+    """
+    Derive surface from AtomGroup positions.
 
     Parameters
     ----------
-    (x, y):  tuple
-        Value of (x, y) coordinates
-    factor:  float
-        Mappign factor to assign grid.
+    n_cells_x : int.
+        number of cells in the grid of size `max_width_x`, `x` axis.
+    n_cells_y : int.
+        number of cells in the grid of size `max_width_y`, `y` axis.
+    selection : AtomGroup.
+        AtomGroup of reference selection to define the surface
+        of the membrane.
+    max_width_x: float.
+        Maximum width of simulation box in x axis. (Determined by simulation box dimensions)
+    max_width_y: float.
+        Maximum width of simulation box in y axis. (Determined by simulation box dimensions)
 
-        Returns
-        -------
-        Returns l, m  with l,m as int.
+    Returns
+    -------
+    z_coordinates: numpy.ndarray
+        Average z-coordinate values. Return Numpy array of floats of
+        shape `(n_cells_x, n_cells_y)`.
 
     """
+    coordinates = selection.positions
+    return get_z_surface(coordinates, n_x_bins=n_cells_x, n_y_bins=n_cells_y,
+                         x_range=(0, max_width_x), y_range=(0, max_width_y))
 
-    index_grid_l = int(abs(coords[0]) * factor)
-    index_grid_m = int(abs(coords[1]) * factor)
-
-    return index_grid_l, index_grid_m
 
-
-def core_fast_leaflet(universe, z_Ref, n_cells, selection, max_width):
+def get_z_surface(coordinates, n_x_bins=10, n_y_bins=10, x_range=(0, 100), y_range=(0, 100)):
     """
-    Runs core_fast_leaflet for selected surface
+    Derive surface from distribution of z coordinates in grid.
 
     Parameters
     ----------
-    universe: MDA Universe
-        Provides coordinates and trajectory of the coordinates.
-    z_Ref : dict.
-        Dictionary to store values in every iteration over frames.
-    n_cells : int.
-        number of cells in the grid of size `max_width`.
-    selection : AtomGroup
-        AtomGroup of reference selection to define the surface
-        of the membrane.
-    max_width : int.
-        Maximum width of simulation box.
+    coordinates : numpy.ndarray 
+        Coordinates of AtomGroup. Numpy array of shape=(n_atoms, 3).
+    n_x_bins : int.
+        Number of bins in grid in the `x` dimension. 
+    n_y_bins : int.
+        Number of bins in grid in the `y` dimension. 
+    x_range : tuple of (float, float)
+        Range of coordinates (min, max) in the `x` dimension with shape=(2,).
+    y_range : tuple of (float, float)
+        Range of coordinates (min, max) in the `y` dimension with shape=(2,). 
 
     Returns
     -------
-
-    Returns a dictionary with tuples of [i,j] as keys, where 0<i<`max_width`,
-    and 0<j<`max_width`. The resulting values for each key are a list of
-    `len=(n_frames)` containing the extracted z_coordinate a given [i,j] cell
-    in each frame.
+    z_surface: numpy.ndarray 
+        Surface derived from set of coordinates in grid of `x_range, y_range` dimensions.
+        Returns Numpy array of floats of shape (`n_x_bins`, `n_y_bins`)
 
     """
 
-    grid_count_frames = np.zeros([n_cells, n_cells])
+    grid_z_coordinates = np.zeros((n_x_bins, n_y_bins))
+    grid_norm_unit = np.zeros((n_x_bins, n_y_bins))
+
+    x_factor = n_x_bins / (x_range[1] - x_range[0])
+    y_factor = n_y_bins / (y_range[1] - y_range[0])
 
-    for ts in universe.trajectory:
+    x_coords, y_coords, z_coords = coordinates.T
 
-        grid_z_coordinates = np.zeros([n_cells, n_cells])
-        grid_norm_unit = np.zeros([n_cells, n_cells])
+    cell_x_floor = np.floor(x_coords * x_factor).astype(int)
+    cell_y_floor = np.floor(y_coords * y_factor).astype(int)
 
-        factor = np.float32(n_cells / max_width)
+    for l, m, z in zip(cell_x_floor, cell_y_floor, z_coords):
 
-        for bead in selection:
-            x, y, z = bead.position/10
+        try:
+            grid_z_coordinates[l, m] += z
+            grid_norm_unit[l, m] += 1
 
-            l, m = grid_map((x, y), factor)
+        except IndexError:
+            print("Atom outside grid boundaries. Skipping atom.")
 
-            try:
-                grid_z_coordinates[l, m] += z
-                grid_norm_unit[l, m] += 1
+    z_surface = normalized_grid(grid_z_coordinates, grid_norm_unit)
 
-            except ValueError:
-                print("Atom outside grid boundaries. Skipping atom {}".format(bead))
+    return z_surface
 
-        for i, j in it.product(range(n_cells), range(n_cells)):
-            if grid_norm_unit[i, j] > 0:
-                z_Ref[i, j] += grid_z_coordinates[i, j] / grid_norm_unit[i, j]
-                grid_count_frames[i, j] += 1
 
-    for i, j in it.product(range(n_cells), range(n_cells)):
-        if grid_count_frames[i, j] > 0:
-            z_Ref[i, j] /= grid_count_frames[i, j]
+def normalized_grid(grid_z_coordinates, grid_norm_unit):
+    """
+    Calculates average z coordinate in unit cell
+
+    Parameters
+    ----------
+
+    z_ref: np.array
+        Empty array of (l,m) 
+    grid_z_coordinates: np.array
+        Array of size (l,m) with z coordinates stored in unit cell.
+    grid_norm_unit: np.array
+        Array of size (l,m) with number of atoms in unit cell.
+
+    Returns
+    -------
+    Returns nornalized set of z coordinates in grid.
+
+    """
+
+    grid_norm_unit = np.where(grid_norm_unit > 0, grid_norm_unit, np.nan)
+    z_normalized = grid_z_coordinates / grid_norm_unit
 
-        else:
-            z_Ref[i, j] = np.nan
+    return z_normalized
 
 
 def gaussian_curvature(Z):

membrane_curvature/tests/test_mdakit_membcurv.py

@@ -2,9 +2,8 @@
 Unit and regression test for the membrane_curvature package.
 """
 
-# Import package, test suite, and other packages as needed
 import pytest
-from ..lib.mods import core_fast_leaflet, mean_curvature, gaussian_curvature, grid_map
+from ..lib.mods import mean_curvature, gaussian_curvature, normalized_grid, derive_surface, get_z_surface
 import numpy as np
 from numpy.testing import assert_almost_equal
 import MDAnalysis as mda
@@ -120,6 +119,13 @@
 }
 
 
+@pytest.fixture
+def small_grofile():
+    u = mda.Universe(GRO_PO4_SMALL)
+    sel = u.select_atoms('name PO4')
+    return sel
+
+
 def test_gaussian_curvature_small():
     K_test = gaussian_curvature(MEMBRANE_CURVATURE_DATA['z_ref']['small'])
     for k, k_test in zip(MEMBRANE_CURVATURE_DATA['gaussian_curvature']['small'], K_test):
@@ -144,44 +150,57 @@ def test_mean_curvature_all():
         assert_almost_equal(h, h_test)
 
 
-def test_core_fast_leaflets():
-    n_cells, max_width = 3, 3
-    z_calc = np.zeros([n_cells, n_cells])
-    u = mda.Universe(GRO_PO4_SMALL, XTC_PO4_SMALL)
-    selection = u.select_atoms('index 0:3')
-    core_fast_leaflet(u, z_calc, n_cells, selection, max_width)
-    for z, z_test in zip(MEMBRANE_CURVATURE_DATA['grid']['small']['upper'], z_calc):
-        print(z, z_test)
-        assert_almost_equal(z, z_test)
-
-
-@pytest.mark.parametrize('factor', [1, 2])
-@pytest.mark.parametrize('dummy_coord', [
-    # dummy coordinates (x,y) in grid of 3x3
-    (0, 0), (1, 0), (2, 0),
-    (0, 1), (1, 1), (2, 1),
-    (0, 2), (1, 2), (2, 2),
-    # dummy coordinates (x,y) in grid of 5x5
-    (0, 0), (1, 0), (2, 0), (3, 0), (4, 0),
-    (0, 1), (1, 1), (2, 1), (3, 1), (4, 1),
-    (0, 2), (1, 2), (2, 2), (3, 2), (4, 2),
-    (0, 3), (1, 3), (2, 3), (3, 3), (4, 3),
-    (0, 4), (1, 4), (2, 4), (3, 4), (4, 4)
+@pytest.mark.parametrize('n_cells, grid_z_coords', [
+    (3, np.full((3, 3), 10.)),
+    (3, np.array([[10., 20., 30.], [10., 20., 30.], [10., 20., 30.]], dtype=float))
+])
+def test_normalized_grid_identity_other_values(n_cells, grid_z_coords):
+    unit = np.ones([n_cells, n_cells])
+    z_avg = normalized_grid(grid_z_coords, unit)
+    assert_almost_equal(z_avg, grid_z_coords)
+
+
+def test_normalized_grid_more_beads():
+    # sum of z coordinate in grid,
+    grid_z_coords = np.full((3, 3), 10.)
+    # grid number of beads per unit
+    norm_grid = np.array([[2., 1., 1.], [1., 2., 1.], [1., 1., 2.]])
+    # avg z coordinate in grid
+    expected_normalized_surface = np.array([[5., 10., 10.], [10., 5., 10.], [10., 10., 5.]])
+    average_surface = normalized_grid(grid_z_coords, norm_grid)
+    assert_almost_equal(average_surface, expected_normalized_surface)
+
+
+def test_derive_surface(small_grofile):
+    n_cells, max_width = 3, 30
+    expected_surface = np.array(([150., 150., 120.], [150., 120., 120.], [150., 120., 120.]))
+    max_width_x = max_width_y = max_width
+    surface = derive_surface(n_cells, n_cells, small_grofile, max_width_x, max_width_y)
+    assert_almost_equal(surface, expected_surface)
+
+
+def test_derive_surface_from_numpy():
+    dummy_array = np.array([[0., 0., 150.], [100., 0., 150.], [200., 0., 150.],
+                            [0., 100., 150.], [100., 100., 120.], [200., 100., 120.],
+                            [0., 200., 120.], [100., 200., 120.], [200., 200., 120.]])
+    x_bin = y_bin = 3
+    x_range = y_range = (0, 300)
+    expected_surface = np.array(([150., 150., 120.], [150., 120., 120.], [150., 120., 120.]))
+    surface = get_z_surface(dummy_array, x_bin, y_bin, x_range, y_range)
+    assert_almost_equal(surface, expected_surface)
+
+
+@pytest.mark.parametrize('x_bin, y_bin, x_range, y_range, expected_surface', [
+    (3, 3, (0, 300), (0, 300), np.array(([150., np.nan, 150.],
+                                         [np.nan, 150., 150.],
+                                         [150., 150., 150.]))),
+    (3, 4, (0, 300), (0, 400), np.array([[150.,  np.nan, 150.,  np.nan],
+                                         [np.nan, 150., 150.,  np.nan],
+                                         [150., 150., 150.,  np.nan]]))
 ])
-def test_grid_map_grids(dummy_coord, factor):
-    cell = (dummy_coord[0] * factor, dummy_coord[1] * factor)
-    assert grid_map(dummy_coord, factor) == cell
-
-
-@pytest.mark.parametrize('dummy_coords, expected', [
-    # dummy coordinates (x,y)
-    (0, 0), (-1, 0), (2, 0),
-    (0, 1), (-1, 1), (2, 1),
-    (0, 2), (-1, 2), (2, 2),
-    # should map to
-    (0, 0), (2, 0), (2, 0),
-    (0, 1), (2, 1), (2, 1),
-    (0, 2), (2, 2), (2, 2)])
-@pytest.mark.xfail(reason='Incorrect mapping of negative coordinates')
-def test_grid_map_9grid_negative(dummy_coords, expected):
-    assert grid_map(dummy_coords, 1) == expected
+def test_get_z_surface(x_bin, y_bin, x_range, y_range, expected_surface):
+    dummy_array = np.array([[0., 0., 150.], [0., 0., 150.], [200., 0., 150.],
+                            [0., 0., 150.], [100., 100., 150.], [200., 100., 150.],
+                            [0., 200., 150.], [100., 200., 150.], [200., 200., 150.]])
+    surface = get_z_surface(dummy_array, x_bin, y_bin, x_range, y_range)
+    assert_almost_equal(surface, expected_surface)