From 8a5229505f91155cc3d4fb6f4c063233dc6c97b8 Mon Sep 17 00:00:00 2001 From: Estefania Barreto-Ojeda Date: Mon, 2 Aug 2021 15:32:05 -0600 Subject: [PATCH] Added wrap transformations (#48) * Added coordinate wrapping with tests. --- membrane_curvature/__init__.py | 5 +- membrane_curvature/base.py | 42 ++- membrane_curvature/data/test_po4_small.gro | 2 +- membrane_curvature/surface.py | 9 +- .../tests/test_membrane_curvature.py | 274 +++++++++++++++++- 5 files changed, 316 insertions(+), 16 deletions(-) diff --git a/membrane_curvature/__init__.py b/membrane_curvature/__init__.py index 106aa2c..1f3b51b 100644 --- a/membrane_curvature/__init__.py +++ b/membrane_curvature/__init__.py @@ -11,8 +11,9 @@ # # Add imports here -from membrane_curvature.surface import normalized_grid, derive_surface, get_z_surface -from membrane_curvature.curvature import mean_curvature, gaussian_curvature +from .surface import normalized_grid, derive_surface, get_z_surface +from .curvature import mean_curvature, gaussian_curvature +from .base import MembraneCurvature # Handle versioneer from ._version import get_versions diff --git a/membrane_curvature/base.py b/membrane_curvature/base.py index a90cf26..e90740e 100644 --- a/membrane_curvature/base.py +++ b/membrane_curvature/base.py @@ -35,8 +35,8 @@ class MembraneCurvature(AnalysisBase): select : str or iterable of str, optional. The selection string of an atom selection to use as a reference to derive a surface. - pbc : bool, optional - Apply periodic boundary conditions. + wrap : bool, optional + Apply coordinate wrapping to pack atoms into the primary unit cell. n_x_bins : int, optional, default: '100' Number of bins in grid in the x dimension. n_y_bins : int, optional, default: '100' @@ -68,8 +68,22 @@ class MembraneCurvature(AnalysisBase): Each array has shape (`n_x_bins`, `n_y_bins`) + See also + -------- + `MDAnalysis.transformations.wrap + `_ + Notes ----- + Use `wrap=True` to translates the atoms of your `mda.Universe` back + in the unit cell. Use `wrap=False` for processed trajectories where + rotational/translational fit is performed. + + For more details on when to use `wrap=True`, check the `Usage + `_ + page. + + The derived surface and calculated curvatures are available in the :attr:`results` attributes. @@ -95,8 +109,10 @@ class MembraneCurvature(AnalysisBase): mean_curvature = mc.results.average_mean_curvature gaussian_curvature = mc.results.average_gaussian_curvature - The respective 2D curvature plots can be obtained using the `matplotlib` package for - data visualization via `imshow`. We recommend using the `gaussian` interpolation. + The respective 2D curvature plots can be obtained using the `matplotlib` + package for data visualization via `imshow`. We recommend using the + `gaussian` interpolation. + """ @@ -104,11 +120,11 @@ def __init__(self, universe, select='all', n_x_bins=100, n_y_bins=100, x_range=None, y_range=None, - pbc=True, **kwargs): + wrap=True, **kwargs): super().__init__(universe.universe.trajectory, **kwargs) self.ag = universe.select_atoms(select) - self.pbc = pbc + self.wrap = wrap self.n_x_bins = n_x_bins self.n_y_bins = n_y_bins self.x_range = x_range if x_range else (0, universe.dimensions[0]) @@ -128,6 +144,17 @@ def __init__(self, universe, select='all', warnings.warn(msg) logger.warn(msg) + # Apply wrapping coordinates + if not self.wrap: + # Warning + msg = (" `wrap == False` may result in inaccurate calculation " + "of membrane curvature. Surfaces will be derived from " + "a reduced number of atoms. \n " + " Ignore this warning if your trajectory has " + " rotational/translational fit rotations! ") + warnings.warn(msg) + logger.warn(msg) + def _prepare(self): # Initialize empty np.array with results self.results.z_surface = np.full((self.n_frames, @@ -141,6 +168,9 @@ def _prepare(self): self.n_y_bins), np.nan) def _single_frame(self): + # Apply wrapping coordinates + if self.wrap: + self.ag.wrap() # Populate a slice with np.arrays of surface, mean, and gaussian per frame self.results.z_surface[self._frame_index] = get_z_surface(self.ag.positions, n_x_bins=self.n_x_bins, diff --git a/membrane_curvature/data/test_po4_small.gro b/membrane_curvature/data/test_po4_small.gro index bff4d61..e5e60d3 100644 --- a/membrane_curvature/data/test_po4_small.gro +++ b/membrane_curvature/data/test_po4_small.gro @@ -9,4 +9,4 @@ Test file 9 lipids in grid 2708POPC PO4 7 0.000 2.000 12.000 2713POPC PO4 8 1.000 2.000 12.000 2999POPC PO4 9 2.000 2.000 12.000 - 3.00000 3.00000 3.00000 + 3.00000 3.00000 30.00000 diff --git a/membrane_curvature/surface.py b/membrane_curvature/surface.py index 6fbd6fa..a5afc44 100644 --- a/membrane_curvature/surface.py +++ b/membrane_curvature/surface.py @@ -16,6 +16,7 @@ """ import numpy as np +import warnings def derive_surface(atoms, n_cells_x, n_cells_y, max_width_x, max_width_y): @@ -87,11 +88,17 @@ def get_z_surface(coordinates, n_x_bins=10, n_y_bins=10, x_range=(0, 100), y_ran for l, m, z in zip(cell_x_floor, cell_y_floor, z_coords): try: + if l < 0 or m < 0: + msg = ("Atom outside grid boundaries. Skipping atom.") + warnings.warn(msg) + continue + grid_z_coordinates[l, m] += z grid_norm_unit[l, m] += 1 except IndexError: - print("Atom outside grid boundaries. Skipping atom.") + msg = ("Atom outside grid boundaries. Skipping atom.") + warnings.warn(msg) z_surface = normalized_grid(grid_z_coordinates, grid_norm_unit) diff --git a/membrane_curvature/tests/test_membrane_curvature.py b/membrane_curvature/tests/test_membrane_curvature.py index acb1d80..a9bca29 100644 --- a/membrane_curvature/tests/test_membrane_curvature.py +++ b/membrane_curvature/tests/test_membrane_curvature.py @@ -2,6 +2,7 @@ Unit and regression test for the membrane_curvature package. """ + import pytest from membrane_curvature.surface import normalized_grid, derive_surface, get_z_surface from membrane_curvature.curvature import mean_curvature, gaussian_curvature @@ -211,7 +212,7 @@ def test_get_z_surface(x_bin, y_bin, x_range, y_range, expected_surface): class TestMembraneCurvature(object): @pytest.fixture() def universe(self): - return mda.Universe(GRO_PO4_SMALL, XTC_PO4_SMALL) + return mda.Universe(GRO_PO4_SMALL) @pytest.fixture() def universe_dummy(self): @@ -220,10 +221,71 @@ def universe_dummy(self): [0., 200., 150.], [100., 200., 150.], [200., 200., 150.]]) u = mda.Universe(a, n_atoms=9) - u.dimensions = [3, 3, 3, 90., 90., 90.] + u.dimensions = [300, 300, 300, 90., 90., 90.] + + return u + + @pytest.fixture() + def universe_dummy_full(self): + a = np.array([[0., 0., 150.], [100., 0., 150.], [200., 0., 150.], + [0., 200., 150.], [100., 200., 120.], [200., 200., 120.], + [0., 100., 120.], [100., 100., 120.], [200., 100., 120.]]) + + u = mda.Universe(a, n_atoms=9) + u.dimensions = [300, 300, 300, 90., 90., 90.] + + return u + @pytest.fixture() + def universe_dummy_wrap(self): + # Atoms out of bounds in x + # +-----------+ + # | | 8 | 9 | 7 + # +-----------+ + # 6 | 4 | 5 | | + # +-----------+ + # | | 2 | 3 | 1 + # +-----------+ + a = np.array([[300., 0., 110.], [100., 0., 150.], [200., 0., 150.], + [0., 100., 150.], [100., 100., 150.], [-100., 100., 150.], + [300., 200., 110.], [100., 200., 150.], [200., 200., 150.]]) + + u = mda.Universe(a, n_atoms=9) + u.dimensions = [300, 300, 300, 90., 90., 90.] return u + # Equivalent to universe_dummy_wrap when wrapping is applied. + # Atoms out of bounds in x and y + @pytest.fixture() + def universe_dummy_wrap_xy(self): + a = np.array([[300., 0., 110.], [100., 0., 150.], [200., 0., 150.], + [0., -200., 150.], [100., -200., 150.], [-100., 100., 150.], + [300., 200., 110.], [100., 200., 150.], [200., 200., 150.]]) + + u = mda.Universe(a, n_atoms=9) + u.dimensions = [300, 300, 300, 90., 90., 90.] + + return u + + @pytest.fixture() + def dummy_surface(self): + surface = np.array([[110., 150., 110.], + [150., 150., 150.], + [150., 150., 150.]]) + return surface + + @pytest.fixture() + def curvature_unwrapped_universe(self, universe_dummy_wrap): + return MembraneCurvature(universe_dummy_wrap, + n_x_bins=3, + n_y_bins=3).run() + + @pytest.fixture() + def curvature_unwrapped_universe_xy(self, universe_dummy_wrap_xy): + return MembraneCurvature(universe_dummy_wrap_xy, + n_x_bins=3, + n_y_bins=3).run() + def test_invalid_selection(self, universe): with pytest.raises(ValueError, match=r'Invalid selection'): MembraneCurvature(universe, select='name P') @@ -261,6 +323,7 @@ def test_analysis_get_z_surface_dummy(self, universe_dummy, x_bin, y_bin, x_rang avg_surface = mc.results.average_z_surface assert_almost_equal(avg_surface, expected_surface) + @pytest.mark.xfail(reason="Wrapping coordinates not applied.") @pytest.mark.parametrize('x_bin, y_bin, expected_surface', [ (3, 3, np.array([[150., 150., 120.], @@ -286,7 +349,8 @@ def test_analysis_get_z_surface(self, universe, x_bin, y_bin, expected_surface): avg_surface = mc.results.average_z_surface assert_almost_equal(avg_surface, expected_surface) - def test_analysis_mean(self, universe): + # test using wrap=True with test grofile + def test_analysis_mean_wrap(self, universe): expected_mean = np.array([[7.50000000e+00, 1.33985392e-01, 2.77315457e-04], [-2.77315457e-04, -3.53944270e-01, -7.50000000e+00], [-2.77315457e-04, -5.01100068e-01, -7.50000000e+00]]) @@ -297,13 +361,211 @@ def test_analysis_mean(self, universe): avg_mean = mc.results.average_mean assert_almost_equal(avg_mean, expected_mean) - - def test_analysis_mean_default_selection(self, universe): + # test using wrap=False with test grofile + def test_analysis_mean_no_wrap(self, universe): expected_mean = np.array([[7.50000000e+00, 1.33985392e-01, 2.77315457e-04], [-2.77315457e-04, -3.53944270e-01, -7.50000000e+00], [-2.77315457e-04, -5.01100068e-01, -7.50000000e+00]]) mc = MembraneCurvature(universe, + select='name PO4', + n_x_bins=3, + n_y_bins=3, + wrap=False).run() + avg_mean = mc.results.average_mean + assert_almost_equal(avg_mean, expected_mean) + + # test using dummy Universe with atoms out of boounds + # with wrap=True (default) + # +-----------+ +-----------+ + # | | 8 | 9 | 7 | 7 | 8 | 9 | + # +-----------+ +-----------+ + # 6 | 4 | 5 | | ---> | 4 | 5 | 6 | + # +-----------+ +-----------+ + # | | 2 | 3 | 1 | 1 | 2 | 3 | + # +-----------+ +-----------+ + # + # 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) + + # 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): + x_bin = y_bin = 3 + mc = MembraneCurvature(universe_dummy_wrap_xy, + 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) + + # 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) + + 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) + + # 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) + + 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) + + # test using dummy Universe with atoms out of boounds + # with wrap=False + # +-----------+ + # | | 8 | 9 | 7 + # +-----------+ + # 6 | 4 | 5 | | + # +-----------+ + # | | 2 | 3 | 1 + # +-----------+ + # test surface + # with wrap=False in universe with atoms out of bounds in x + def test_analysis_get_z_surface_no_wrap(self, universe_dummy_wrap): + expected_surface = [[np.nan, 150., np.nan], + [150., 150., 150.], + [150., np.nan, 150.]] + x_bin = y_bin = 3 + mc = MembraneCurvature(universe_dummy_wrap, + n_x_bins=x_bin, + n_y_bins=y_bin, + wrap=False).run() + avg_surface = mc.results.average_z_surface + assert_almost_equal(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): + expected_surface = [[np.nan, np.nan, np.nan], + [150., np.nan, 150.], + [150., np.nan, 150.]] + x_bin = y_bin = 3 + mc = MembraneCurvature(universe_dummy_wrap_xy, + n_x_bins=x_bin, + n_y_bins=y_bin, + wrap=False).run() + avg_surface = mc.results.average_z_surface + assert_almost_equal(avg_surface, expected_surface) + + # test mean + def test_analysis_mean_no_wrap(self, universe_dummy_wrap): + expected_mean = np.array(np.full((3, 3), np.nan)) + x_bin = y_bin = 3 + mc = MembraneCurvature(universe_dummy_wrap, + n_x_bins=x_bin, + n_y_bins=y_bin, + wrap=False).run() + avg_mean = mc.results.average_mean + assert_almost_equal(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)) + x_bin = y_bin = 3 + mc = MembraneCurvature(universe_dummy_wrap_xy, + n_x_bins=x_bin, + n_y_bins=y_bin, + wrap=False).run() + avg_mean = mc.results.average_mean + assert_almost_equal(avg_mean, expected_mean) + + # test gaussian + def test_analysis_gaussian_no_wrap(self, universe_dummy_wrap): + expected_gaussian = np.array(np.full((3, 3), np.nan)) + x_bin = y_bin = 3 + mc = MembraneCurvature(universe_dummy_wrap, + n_x_bins=x_bin, + n_y_bins=y_bin, + wrap=False).run() + avg_gaussian = mc.results.average_gaussian + assert_almost_equal(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)) + x_bin = y_bin = 3 + mc = MembraneCurvature(universe_dummy_wrap_xy, + n_x_bins=x_bin, + n_y_bins=y_bin, + wrap=False).run() + avg_gaussian = mc.results.average_gaussian + assert_almost_equal(avg_gaussian, expected_gaussian) + + @pytest.mark.parametrize('x_bin, y_bin, expected_surface', [ + (3, 3, + np.array([[150., 120., 150.], + [150., 120., 120.], + [150., 120., 120.]])), + (4, 4, + np.array([[150., 120., 150., np.nan], + [150., 120., 120., np.nan], + [150., 120., 120., np.nan], + [np.nan, np.nan, np.nan, np.nan]])), + (5, 5, + np.array([[150., 120., np.nan, 150., np.nan], + [150., 120., np.nan, 120., np.nan], + [np.nan, np.nan, np.nan, np.nan, np.nan], + [150., 120., np.nan, 120., np.nan], + [np.nan, np.nan, np.nan, np.nan, np.nan]])) + + ]) + def test_analysis_get_z_surface(self, universe_dummy_full, x_bin, y_bin, expected_surface): + mc = MembraneCurvature(universe_dummy_full, + 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) + + def test_analysis_mean(self, universe_dummy_full): + expected_mean = np.array([[-5.54630914e-04, - 1.50000000e+01, 8.80203593e-02], + [-2.77315457e-04, - 2.20748929e-03, - 5.01100068e-01], + [-2.77315457e-04, - 2.20748929e-03, - 1.50000000e+01]]) + mc = MembraneCurvature(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) \ No newline at end of file + assert_almost_equal(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 + (3, 3, 300, 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.]]), + np.array([[150., 150., 120.], + [150., 120., 120.], + [150., 120., 120.]])), + # test with negative z coordinates with 4 bins + (4, 4, 400, np.array([[0., 0., -150.], [100., 0., -150.], [200., 0., -150.], [300., 0., -150.], + [0., 100., -150.], [100., 100., 120.], [200., 100., 120.], [300., 100., -150.], + [0., 200., 120.], [100., 200., 120.], [200., 200., 120.], [300., 200., -150.], + [0., 300., -150.], [100., 300., -150.], [200., 300., -150.], [300., 300., -150.]]), + np.array([[150., 150., 120., 150.], + [150., 120., 120., 150.], + [150., 120., 120., 150.], + [150., 150., 150., 150.]])) + ]) + def test_analysis_wrapping_coordinates(self, x_bin, y_bin, box_dim, dummy_array, expected_surface): + x_range, y_range = (0, box_dim), (0, box_dim) + u = mda.Universe(dummy_array, n_atoms=len(dummy_array)) + u.dimensions = [box_dim, box_dim, 300, 90., 90., 90.] + # Check with wrapped coords in base + mc = MembraneCurvature(u, select='all', + n_x_bins=x_bin, + n_y_bins=y_bin, + x_range=x_range, + 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) + + 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)