diff --git a/membrane_curvature/base.py b/membrane_curvature/base.py index e90740e..1f3d671 100644 --- a/membrane_curvature/base.py +++ b/membrane_curvature/base.py @@ -12,7 +12,7 @@ import numpy as np import warnings -from .surface import get_z_surface +from .surface import get_z_surface, surface_interpolation from .curvature import mean_curvature, gaussian_curvature import MDAnalysis @@ -32,7 +32,7 @@ class MembraneCurvature(AnalysisBase): ---------- universe : Universe or AtomGroup An MDAnalysis Universe object. - select : str or iterable of str, optional. + select : str or iterable of str, optional. The selection string of an atom selection to use as a reference to derive a surface. wrap : bool, optional @@ -57,13 +57,13 @@ class MembraneCurvature(AnalysisBase): results.gaussian_curvature : ndarray Gaussian curvature associated to the surface. Arrays of shape (`n_frames`, `n_x_bins`, `n_y_bins`) - results.average_z_surface : ndarray - Average of the array elements in `z_surface`. + results.average_z_surface : ndarray + Average of the array elements in `z_surface`. Each array has shape (`n_x_bins`, `n_y_bins`) - results.average_mean_curvature : ndarray + results.average_mean_curvature : ndarray Average of the array elements in `mean_curvature`. Each array has shape (`n_x_bins`, `n_y_bins`) - results.average_gaussian_curvature: ndarray + results.average_gaussian_curvature: ndarray Average of the array elements in `gaussian_curvature`. Each array has shape (`n_x_bins`, `n_y_bins`) @@ -120,7 +120,8 @@ def __init__(self, universe, select='all', n_x_bins=100, n_y_bins=100, x_range=None, y_range=None, - wrap=True, **kwargs): + wrap=True, + interpolation=False, **kwargs): super().__init__(universe.universe.trajectory, **kwargs) self.ag = universe.select_atoms(select) @@ -129,6 +130,7 @@ def __init__(self, universe, select='all', self.n_y_bins = n_y_bins self.x_range = x_range if x_range else (0, universe.dimensions[0]) self.y_range = y_range if y_range else (0, universe.dimensions[1]) + self.interpolation = interpolation # Raise if selection doesn't exist if len(self.ag) == 0: @@ -156,7 +158,6 @@ def __init__(self, universe, select='all', logger.warn(msg) def _prepare(self): - # Initialize empty np.array with results self.results.z_surface = np.full((self.n_frames, self.n_x_bins, self.n_y_bins), np.nan) @@ -166,6 +167,9 @@ def _prepare(self): self.results.gaussian = np.full((self.n_frames, self.n_x_bins, self.n_y_bins), np.nan) + self.results.interpolated_z_surface = np.full((self.n_frames, + self.n_x_bins, + self.n_y_bins), np.nan) def _single_frame(self): # Apply wrapping coordinates @@ -179,8 +183,20 @@ def _single_frame(self): y_range=self.y_range) self.results.mean[self._frame_index] = mean_curvature(self.results.z_surface[self._frame_index]) self.results.gaussian[self._frame_index] = gaussian_curvature(self.results.z_surface[self._frame_index]) + # Perform interpolation fo surface + if self.interpolation: + # Populate a slice with np.arrays of interpolated surface + self.results.interpolated_z_surface[self._frame_index] = surface_interpolation( + self.results.z_surface[self._frame_index]) def _conclude(self): self.results.average_z_surface = np.nanmean(self.results.z_surface, axis=0) self.results.average_mean = np.nanmean(self.results.mean, axis=0) self.results.average_gaussian = np.nanmean(self.results.gaussian, axis=0) + if self.interpolation: + self.results.average_interpolated_z_surface = np.nanmean( + self.results.interpolated_z_surface[self._frame_index], axis=0) + self.results.average_interpolated_mean = np.nanmean(mean_curvature( + self.results.interpolated_z_surface[self._frame_index]), axis=0) + self.results.average_interpolated_gaussian = np.nanmean(gaussian_curvature( + self.results.interpolated_z_surface[self._frame_index]), axis=0) diff --git a/membrane_curvature/surface.py b/membrane_curvature/surface.py index a5afc44..e9729ef 100644 --- a/membrane_curvature/surface.py +++ b/membrane_curvature/surface.py @@ -39,7 +39,7 @@ def derive_surface(atoms, n_cells_x, n_cells_y, max_width_x, max_width_y): Returns ------- - z_coordinates: numpy.ndarray + z_coordinates: np.ndarray Average z-coordinate values. Return Numpy array of floats of shape `(n_cells_x, n_cells_y)`. @@ -55,8 +55,8 @@ def get_z_surface(coordinates, n_x_bins=10, n_y_bins=10, x_range=(0, 100), y_ran Parameters ---------- - coordinates : numpy.ndarray - Coordinates of AtomGroup. Numpy array of shape=(n_atoms, 3). + coordinates : np.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. @@ -70,7 +70,7 @@ def get_z_surface(coordinates, n_x_bins=10, n_y_bins=10, x_range=(0, 100), y_ran ------- z_surface: np.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`) + Returns NumPy array of floats of shape (`n_x_bins`, `n_y_bins`) """ @@ -131,3 +131,60 @@ def normalized_grid(grid_z_coordinates, grid_norm_unit): z_normalized = grid_z_coordinates / grid_norm_unit return z_normalized + + +def interpolation_by_array(array_surface): + """ + Interpolates values contained in `array_surface` over axis + + Parameters + ---------- + + array_surface: np.ndarray + Array of floats of shape (`n_x_bins`, `n_y_bins`) + + Returns + ------- + interpolated_array: np.ndarray + Returns interpolated array. + Array of shape (`n_x_bins` x `n_y_bins`,) + + """ + + # create mask for nans + mask_nans = ~np.isnan(array_surface) + + # index of array_surface + index_array = np.arange(array_surface.shape[0]) + + # interpolate values in array + interpolated_array = np.interp(index_array, + np.flatnonzero(mask_nans), + array_surface[mask_nans]) + + return interpolated_array + + +def surface_interpolation(array_surface): + """ + Calculates interpolation of arrays. + + Parameters + ---------- + + array_surface: np.ndarray + Array of floats of shape (`n_x_bins`, `n_y_bins`) + + Returns + ------- + + interpolated_surface: np.ndarray + Interpolated surface derived from set of coordinates + in grid of `x_range, y_range` dimensions. + Array of floats of shape (`n_x_bins`, `n_y_bins`) + + """ + + interpolated_surface = np.apply_along_axis(interpolation_by_array, axis=0, arr=array_surface) + + return interpolated_surface diff --git a/membrane_curvature/tests/test_membrane_curvature.py b/membrane_curvature/tests/test_membrane_curvature.py index a9bca29..971e18d 100644 --- a/membrane_curvature/tests/test_membrane_curvature.py +++ b/membrane_curvature/tests/test_membrane_curvature.py @@ -4,12 +4,13 @@ import pytest -from membrane_curvature.surface import normalized_grid, derive_surface, get_z_surface +from membrane_curvature.surface import (normalized_grid, derive_surface, get_z_surface, + surface_interpolation) from membrane_curvature.curvature import mean_curvature, gaussian_curvature import numpy as np -from numpy.testing import assert_almost_equal +from numpy.testing import assert_almost_equal, assert_allclose import MDAnalysis as mda -from membrane_curvature.tests.datafiles import (GRO_PO4_SMALL, XTC_PO4_SMALL) +from membrane_curvature.tests.datafiles import (GRO_PO4_SMALL) from membrane_curvature.base import MembraneCurvature # Reference data from datafile @@ -132,19 +133,19 @@ def small_grofile(): 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): - assert_almost_equal(k, k_test) + assert_allclose(k, k_test) def test_mean_curvature_small(): H_test = mean_curvature(MEMBRANE_CURVATURE_DATA['z_ref']['small']) for h, h_test in zip(MEMBRANE_CURVATURE_DATA['mean_curvature']['small'], H_test): - assert_almost_equal(h, h_test) + assert_allclose(h, h_test) def test_gaussian_curvature_all(): K_test = gaussian_curvature(MEMBRANE_CURVATURE_DATA['z_ref']['all']) for k, k_test in zip(MEMBRANE_CURVATURE_DATA['gaussian_curvature']['all'], K_test): - assert_almost_equal(k, k_test) + assert_allclose(k, k_test, rtol=1e-4) def test_mean_curvature_all(): @@ -160,7 +161,7 @@ def test_mean_curvature_all(): 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) + assert_allclose(z_avg, grid_z_coords) def test_normalized_grid_more_beads(): @@ -171,7 +172,7 @@ def test_normalized_grid_more_beads(): # 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) + assert_allclose(average_surface, expected_normalized_surface) def test_derive_surface(small_grofile): @@ -179,7 +180,7 @@ def test_derive_surface(small_grofile): expected_surface = np.array(([150., 150., 120.], [150., 120., 120.], [150., 120., 120.])) max_width_x = max_width_y = max_width surface = derive_surface(small_grofile, n_cells, n_cells, max_width_x, max_width_y) - assert_almost_equal(surface, expected_surface) + assert_allclose(surface, expected_surface) def test_derive_surface_from_numpy(): @@ -190,7 +191,7 @@ def test_derive_surface_from_numpy(): 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) + assert_allclose(surface, expected_surface) @pytest.mark.parametrize('x_bin, y_bin, x_range, y_range, expected_surface', [ @@ -206,7 +207,38 @@ def test_get_z_surface(x_bin, y_bin, x_range, y_range, expected_surface): [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) + assert_allclose(surface, expected_surface) + + +@pytest.mark.parametrize('dummy_surface, expected_interpolated_surface', [ + # array 3x3 with all 150 and one nan + (np.array(([150., 150., 150.], + [150., np.nan, 150.], + [150., 150., 150.])), + np.full((3, 3), 150.)), + # array 3x4 with all 150 and two nans + (np.array([[150., 150, 150., 150.], + [150., np.nan, np.nan, 150.], + [150., 150., 150., 150.]]), + np.full((3, 4), 150.)), + # array 4x4 with all 150 and two nans + (np.array([[150., 150, 150., 150.], + [150., np.nan, np.nan, 150.], + [150., 130., 140., 150.], + [150., 150., 150., 150.]]), + np.array([[150., 150, 150., 150.], + [150., 140., 145., 150.], + [150., 130., 140., 150.], + [150., 150., 150., 150.]])), + # array 3x3 with lots of nans + (np.array([[np.nan, np.nan, 150.], + [150, np.nan, 150.], + [np.nan, 150., np.nan]]), + np.full((3, 3), 150.)) +]) +def test_surface_interpolation(dummy_surface, expected_interpolated_surface): + surface = surface_interpolation(dummy_surface) + assert_allclose(surface, expected_interpolated_surface) class TestMembraneCurvature(object): @@ -321,7 +353,7 @@ def test_analysis_get_z_surface_dummy(self, universe_dummy, x_bin, y_bin, x_rang y_range=y_range).run() avg_surface = mc.results.average_z_surface - assert_almost_equal(avg_surface, expected_surface) + assert_allclose(avg_surface, expected_surface) @pytest.mark.xfail(reason="Wrapping coordinates not applied.") @pytest.mark.parametrize('x_bin, y_bin, expected_surface', [ @@ -347,7 +379,7 @@ def test_analysis_get_z_surface(self, universe, x_bin, y_bin, expected_surface): n_x_bins=x_bin, n_y_bins=y_bin).run() avg_surface = mc.results.average_z_surface - assert_almost_equal(avg_surface, expected_surface) + assert_allclose(avg_surface, expected_surface) # test using wrap=True with test grofile def test_analysis_mean_wrap(self, universe): @@ -359,7 +391,7 @@ def test_analysis_mean_wrap(self, universe): n_x_bins=3, n_y_bins=3).run() avg_mean = mc.results.average_mean - assert_almost_equal(avg_mean, expected_mean) + assert_allclose(avg_mean, expected_mean) # test using wrap=False with test grofile def test_analysis_mean_no_wrap(self, universe): @@ -372,7 +404,7 @@ def test_analysis_mean_no_wrap(self, universe): n_y_bins=3, wrap=False).run() avg_mean = mc.results.average_mean - assert_almost_equal(avg_mean, expected_mean) + assert_allclose(avg_mean, expected_mean) # test using dummy Universe with atoms out of boounds # with wrap=True (default) @@ -387,7 +419,7 @@ def test_analysis_mean_no_wrap(self, universe): # test surface in universe with atoms out of bounds in x def test_analysis_get_z_surface_wrap(self, curvature_unwrapped_universe, dummy_surface): avg_surface = curvature_unwrapped_universe.results.average_z_surface - assert_almost_equal(avg_surface, dummy_surface) + assert_allclose(avg_surface, dummy_surface) # test surface in universe with atoms out of bounds in x and y def test_analysis_get_z_surface_wrap_xy(self, universe_dummy_wrap_xy, dummy_surface): @@ -396,29 +428,29 @@ def test_analysis_get_z_surface_wrap_xy(self, universe_dummy_wrap_xy, dummy_surf n_x_bins=x_bin, n_y_bins=y_bin).run() avg_surface = mc.results.average_z_surface - assert_almost_equal(avg_surface, dummy_surface) + assert_allclose(avg_surface, dummy_surface) # test mean curvature def test_analysis_mean_wrap(self, curvature_unwrapped_universe, dummy_surface): avg_mean = curvature_unwrapped_universe.results.average_mean expected_mean = mean_curvature(dummy_surface) - assert_almost_equal(avg_mean, expected_mean) + assert_allclose(avg_mean, expected_mean) def test_analysis_mean_wrap_xy(self, curvature_unwrapped_universe, dummy_surface): avg_mean = curvature_unwrapped_universe.results.average_mean expected_mean = mean_curvature(dummy_surface) - assert_almost_equal(avg_mean, expected_mean) + assert_allclose(avg_mean, expected_mean) # test gaussian def test_analysis_gaussian_wrap(self, curvature_unwrapped_universe, dummy_surface): avg_gaussian = curvature_unwrapped_universe.results.average_gaussian expected_gaussian = gaussian_curvature(dummy_surface) - assert_almost_equal(avg_gaussian, expected_gaussian) + assert_allclose(avg_gaussian, expected_gaussian) def test_analysis_mean_gaussian_wrap_xy(self, curvature_unwrapped_universe, dummy_surface): avg_gaussian = curvature_unwrapped_universe.results.average_gaussian expected_gaussian = gaussian_curvature(dummy_surface) - assert_almost_equal(avg_gaussian, expected_gaussian) + assert_allclose(avg_gaussian, expected_gaussian) # test using dummy Universe with atoms out of boounds # with wrap=False @@ -441,7 +473,7 @@ def test_analysis_get_z_surface_no_wrap(self, universe_dummy_wrap): n_y_bins=y_bin, wrap=False).run() avg_surface = mc.results.average_z_surface - assert_almost_equal(avg_surface, expected_surface) + assert_allclose(avg_surface, expected_surface) # test surface in universe with atoms out of bounds in x and y def test_analysis_get_z_surface_no_wrap_xy(self, universe_dummy_wrap_xy): @@ -454,7 +486,7 @@ def test_analysis_get_z_surface_no_wrap_xy(self, universe_dummy_wrap_xy): n_y_bins=y_bin, wrap=False).run() avg_surface = mc.results.average_z_surface - assert_almost_equal(avg_surface, expected_surface) + assert_allclose(avg_surface, expected_surface) # test mean def test_analysis_mean_no_wrap(self, universe_dummy_wrap): @@ -465,7 +497,7 @@ def test_analysis_mean_no_wrap(self, universe_dummy_wrap): n_y_bins=y_bin, wrap=False).run() avg_mean = mc.results.average_mean - assert_almost_equal(avg_mean, expected_mean) + assert_allclose(avg_mean, expected_mean) def test_analysis_mean_no_wrap(self, universe_dummy_wrap_xy): expected_mean = np.array(np.full((3, 3), np.nan)) @@ -475,7 +507,7 @@ def test_analysis_mean_no_wrap(self, universe_dummy_wrap_xy): n_y_bins=y_bin, wrap=False).run() avg_mean = mc.results.average_mean - assert_almost_equal(avg_mean, expected_mean) + assert_allclose(avg_mean, expected_mean) # test gaussian def test_analysis_gaussian_no_wrap(self, universe_dummy_wrap): @@ -486,7 +518,7 @@ def test_analysis_gaussian_no_wrap(self, universe_dummy_wrap): n_y_bins=y_bin, wrap=False).run() avg_gaussian = mc.results.average_gaussian - assert_almost_equal(avg_gaussian, expected_gaussian) + assert_allclose(avg_gaussian, expected_gaussian) def test_analysis_gaussian_no_wrap(self, universe_dummy_wrap_xy): expected_gaussian = np.array(np.full((3, 3), np.nan)) @@ -496,7 +528,7 @@ def test_analysis_gaussian_no_wrap(self, universe_dummy_wrap_xy): n_y_bins=y_bin, wrap=False).run() avg_gaussian = mc.results.average_gaussian - assert_almost_equal(avg_gaussian, expected_gaussian) + assert_allclose(avg_gaussian, expected_gaussian) @pytest.mark.parametrize('x_bin, y_bin, expected_surface', [ (3, 3, @@ -521,7 +553,7 @@ def test_analysis_get_z_surface(self, universe_dummy_full, x_bin, y_bin, expecte n_x_bins=x_bin, n_y_bins=y_bin).run() avg_surface = mc.results.average_z_surface - assert_almost_equal(avg_surface, expected_surface) + assert_allclose(avg_surface, expected_surface) def test_analysis_mean(self, universe_dummy_full): expected_mean = np.array([[-5.54630914e-04, - 1.50000000e+01, 8.80203593e-02], @@ -531,7 +563,7 @@ def test_analysis_mean(self, universe_dummy_full): n_x_bins=3, n_y_bins=3).run() avg_mean = mc.results.average_mean - assert_almost_equal(avg_mean, expected_mean) + assert_allclose(avg_mean, expected_mean) @pytest.mark.parametrize('x_bin, y_bin, box_dim, dummy_array, expected_surface', [ # test with negative z coordinates with 3 bins @@ -563,9 +595,45 @@ def test_analysis_wrapping_coordinates(self, x_bin, y_bin, box_dim, dummy_array, y_range=y_range).run() avg_surface = mc.results.average_z_surface # assert if default values of wrapped coords in z_surface returns correctly - assert_almost_equal(avg_surface, expected_surface) + assert_allclose(avg_surface, expected_surface) def test_test_analysis_no_wrapping(self, universe): regex = (r"`wrap == False` may result in inaccurate calculation") with pytest.warns(UserWarning, match=regex): MembraneCurvature(universe, wrap=False) + + # test curvature with interpolated surface + @pytest.mark.parametrize('dim_x, dim_y, x_bins, y_bins, dummy_array, expected_interp_surf', [ + # array 3x3 with all 150 and one nan + (300, 300, 3, 3, np.array([[0., 0., 150.], [100., 0., 150.], [200., 0., 150.], + [0., 100., 150.], [100., 100., np.nan], [200., 100., 150.], + [0., 200., 150.], [100., 200., 150.], [200., 200., 150.]]), + np.full((1, 3, 3), 150.)), + # array 3x3 with all 150 and one nan + (300, 300, 3, 3, np.array([[0., 0., 150.], [100., 0., 150.], [200., 0., 150.], + [0., 100., 150.], [100., 100., np.nan], [200., 100., 150.], + [0., 200., np.nan], [100., 200., 150.], [200., 200., 150.]]), + np.full((1, 3, 3), 150.)), + # array 3x4 with all 150 and three nans + (300, 400, 3, 4, np.array([[0., 0., 150.], [100., 0., 150.], [200., 0., 150.], + [0., 100., 150.], [100., 100., np.nan], [200., 100., 150.], + [0., 200., 150], [100., 200., np.nan], [200., 200., np.nan], + [0., 300., 150.], [100., 300., 150.], [200., 300., 150.]]), + np.full((1, 3, 4), 150.)), + # array 4x4 with all 120 and many nans + (400, 400, 4, 4, np.array([[0., 0., np.nan], [100., 0., 120.], [200., 0., 120.], [300., 0., np.nan], + [0., 100., 120.], [100., 100., np.nan], [200., 100., 120.], [300., 100., 120.], + [0., 200., 120], [100., 200., np.nan], [200., 200., np.nan], [300., 200., 120.], + [0., 300., np.nan], [100., 300., 120.], [200., 300., 120.], [300., 300., np.nan]]), + np.full((1, 4, 4), 120.)) + ]) + def test_analysis_interpolates_surface(self, dim_x, dim_y, x_bins, y_bins, dummy_array, expected_interp_surf): + u = mda.Universe(dummy_array, n_atoms=len(dummy_array)) + u.dimensions = [dim_x, dim_y, 300, 90., 90., 90.] + mc = MembraneCurvature(u, + n_x_bins=x_bins, + n_y_bins=y_bins, + wrap=True, + interpolation=True).run() + surface = mc.results.interpolated_z_surface + assert_allclose(surface, expected_interp_surf)