diff --git a/devtools/conda-envs/test_env.yaml b/devtools/conda-envs/test_env.yaml index 1ccbb07..cf74e87 100644 --- a/devtools/conda-envs/test_env.yaml +++ b/devtools/conda-envs/test_env.yaml @@ -6,16 +6,13 @@ dependencies: - python - pip - numpy - - MDAnalysis - - mdtraj # temporary; remove after refactoring - sphinx - - # Testing + # Testing - pytest - pytest-cov - codecov # Pip-only installs - #- pip: - # - codecov + - pip: + - MDAnalysis==2.0.0b0 diff --git a/membrane_curvature/base.py b/membrane_curvature/base.py new file mode 100644 index 0000000..e8e1ef9 --- /dev/null +++ b/membrane_curvature/base.py @@ -0,0 +1,174 @@ +""" +MembraneCurvature +======================================= + +:Author: Estefania Barreto-Ojeda +:Year: 2021 +:Copyright: GNU Public License v3 + +MembraneCurvature calculates the mean and Gaussian curvature of +surfaces derived from a selection of reference. +""" + +import numpy as np +import warnings +from .surface import get_z_surface +from .curvature import mean_curvature, gaussian_curvature + +import MDAnalysis +from MDAnalysis.analysis.base import AnalysisBase + +import logging +MDAnalysis.start_logging() + +logger = logging.getLogger("MDAnalysis.MDAKit.membrane_curvature") + + +class MembraneCurvature(AnalysisBase): + r""" + MembraneCurvature is a tool to calculate membrane curvature. + + Parameters + ---------- + universe : Universe or AtomGroup + An MDAnalysis Universe object. + 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. + n_x_bins : int, optional, default: '100' + Number of bins in grid in the x dimension. + n_y_bins : int, optional, default: '100' + Number of bins in grid in the y dimension. + x_range : tuple of (float, float), optional, default: (0, `universe.dimensions[0]`) + Range of coordinates (min, max) in the x dimension. + y_range : tuple of (float, float), optional, default: (0, `universe.dimensions[1]`) + Range of coordinates (min, max) in the y dimension. + + Attributes + ---------- + results.z_surface : ndarray + Surface derived from atom selection in every frame. + Array of shape (`n_frames`, `n_x_bins`, `n_y_bins`) + results.mean_curvature : ndarray + Mean curvature associated to the surface. + Array of shape (`n_frames`, `n_x_bins`, `n_y_bins`) + 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`. + Each array has shape (`n_x_bins`, `n_y_bins`) + 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 + Average of the array elements in `gaussian_curvature`. + Each array has shape (`n_x_bins`, `n_y_bins`) + + + Notes + ----- + The derived surface and calculated curvatures are available in the + :attr:`results` attributes. + + The attribute :attr:`~MembraneCurvature.results.z_surface` contains the + derived surface averaged over the `n_frames` of the trajectory. + + The attributes :attr:`~MembraneCurvature.results.mean_curvature` and + :attr:`~MembraneCurvature.results.gaussian_curvature` contain the computed + values of mean and Gaussian curvature averaged over the `n_frames` of the + trajectory. + + Example + ----------- + You can pass a universe containing your selection of reference:: + + import MDAnalysis as mda + from MDAnalysis.analysis import MembraneCurvature + + u = mda.Universe(coordinates, trajectory) + mc = MembraneCurvature(u).run() + + surface = mc.results.average_z_surface + 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. + + A simple plot using `imshow` can be obtained by + + import matplotlib.pyplot as plt + fig, ax = plt.subplots() + ax.imshow(mean_curvature, cmap='bwr', interpolation='gaussian', origin='lower') + ax.set_title('Mean Curvature') + plt.show() + + As an alternative, you can use contour plots using `contourf`: + + import matplotlib.pyplot as plt + fig, ax = plt.subplots() + ax.contourf(mean_curvature, cmap='bwr, origin='lower') + ax.set_title('Mean Curvature') + plt.show() + + + + """ + + def __init__(self, universe, select='all', + n_x_bins=100, n_y_bins=100, + x_range=None, + y_range=None, + pbc=True, **kwargs): + + super().__init__(universe.universe.trajectory, **kwargs) + self.ag = universe.select_atoms(select) + self.pbc = pbc + 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]) + self.y_range = y_range if y_range else (0, universe.dimensions[1]) + + # Raise if selection doesn't exist + if len(self.ag) == 0: + raise ValueError("Invalid selection. Empty AtomGroup.") + + # Only checks the first frame. NPT simulations not properly covered here. + # Warning message if range doesn't cover entire dimensions of simulation box + for dim_string, dim_range, num in [('x', self.x_range, 0), ('y', self.y_range, 1)]: + if (dim_range[1] < universe.dimensions[num]): + msg = (f"Grid range in {dim_string} does not cover entire " + "dimensions of simulation box.\n Minimum dimensions " + "must be equal to simulation box.") + warnings.warn(msg) + 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) + self.results.mean = np.full((self.n_frames, + self.n_x_bins, + self.n_y_bins), np.nan) + self.results.gaussian = np.full((self.n_frames, + self.n_x_bins, + self.n_y_bins), np.nan) + + def _single_frame(self): + # 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, + n_y_bins=self.n_y_bins, + x_range=self.x_range, + 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]) + + 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) diff --git a/membrane_curvature/core.py b/membrane_curvature/core.py deleted file mode 100644 index 8883003..0000000 --- a/membrane_curvature/core.py +++ /dev/null @@ -1,46 +0,0 @@ -""" -core.py -MDAkit for Membrane Curvature - -Handles the primary functions -""" - -from membrane_curvature.surface import normalized_grid, derive_surface, get_z_surface -from membrane_curvature.curvature import mean_curvature, gaussian_curvature -import time -import MDAnalysis as mda -import math -from pathlib import Path -__author__ = "Estefania Barreto-Ojeda" -version = 0.1 - - -def main(): - - - - # 1. Populate universe with coordinates and trajectory - u = mda.Universe(topology, trajectory) - - # 2 Set grid: Extract box dimension from MD sim, - # set grid max width, set number of unit cells - box_size = u.dimensions[0] - max_width = box_size * 0.1 - n_cells = math.ceil(max_width / unit_width * 10) - - # 3. Assign lipids in upper and lower leaflet - selection = u.select_atoms(atoms_upper, atoms_lower) - - # 4. Save pickles zpo4 - z_Ref = np.zeros([n_cells, n_cells]) - z_coords = derive_surface(n_cells, selection, max_width) - - # 5. Calculate curvature - K = gaussian_curvature(z_coords) - H = mean_curvature(z_coords) - - return - - -if __name__ == '__main__': - main() diff --git a/membrane_curvature/surface.py b/membrane_curvature/surface.py index 17c3b76..7ec3251 100644 --- a/membrane_curvature/surface.py +++ b/membrane_curvature/surface.py @@ -1,19 +1,19 @@ import numpy as np -def derive_surface(n_cells_x, n_cells_y, selection, max_width_x, max_width_y): +def derive_surface(atoms, n_cells_x, n_cells_y, max_width_x, max_width_y): """ Derive surface from AtomGroup positions. Parameters ---------- + atoms : AtomGroup. + AtomGroup of reference selection to define the surface + of the membrane. 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. @@ -26,7 +26,7 @@ def derive_surface(n_cells_x, n_cells_y, selection, max_width_x, max_width_y): shape `(n_cells_x, n_cells_y)`. """ - coordinates = selection.positions + coordinates = atoms.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)) @@ -105,6 +105,3 @@ def normalized_grid(grid_z_coordinates, grid_norm_unit): z_normalized = grid_z_coordinates / grid_norm_unit return z_normalized - - - diff --git a/membrane_curvature/tests/test_membrane_curvature.py b/membrane_curvature/tests/test_membrane_curvature.py index 0a5efb8..acb1d80 100644 --- a/membrane_curvature/tests/test_membrane_curvature.py +++ b/membrane_curvature/tests/test_membrane_curvature.py @@ -9,6 +9,7 @@ from numpy.testing import assert_almost_equal import MDAnalysis as mda from membrane_curvature.tests.datafiles import (GRO_PO4_SMALL, XTC_PO4_SMALL) +from membrane_curvature.base import MembraneCurvature # Reference data from datafile MEMBRANE_CURVATURE_DATA = { @@ -176,7 +177,7 @@ 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) + surface = derive_surface(small_grofile, n_cells, n_cells, max_width_x, max_width_y) assert_almost_equal(surface, expected_surface) @@ -205,3 +206,104 @@ def test_get_z_surface(x_bin, y_bin, x_range, y_range, expected_surface): [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) + + +class TestMembraneCurvature(object): + @pytest.fixture() + def universe(self): + return mda.Universe(GRO_PO4_SMALL, XTC_PO4_SMALL) + + @pytest.fixture() + def universe_dummy(self): + a = 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.]]) + + u = mda.Universe(a, n_atoms=9) + u.dimensions = [3, 3, 3, 90., 90., 90.] + + return u + + def test_invalid_selection(self, universe): + with pytest.raises(ValueError, match=r'Invalid selection'): + MembraneCurvature(universe, select='name P') + + def test_grid_bigger_than_simulation_box_x_dim(self, universe): + regex = (r"Grid range in x does not cover entire " + r"dimensions of simulation box.\n Minimum dimensions " + r"must be equal to simulation box.") + with pytest.warns(UserWarning, match=regex): + MembraneCurvature(universe, select='name PO4', x_range=(0, 10)) + + def test_grid_bigger_than_simulation_box_y_dim(self, universe): + regex = (r"Grid range in y does not cover entire " + r"dimensions of simulation box.\n Minimum dimensions " + r"must be equal to simulation box.") + with pytest.warns(UserWarning, match=regex): + MembraneCurvature(universe, select='name PO4', y_range=(0, 10)) + + @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_analysis_get_z_surface_dummy(self, universe_dummy, x_bin, y_bin, x_range, y_range, expected_surface): + u = universe_dummy + 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_almost_equal(avg_surface, expected_surface) + + @pytest.mark.parametrize('x_bin, y_bin, expected_surface', [ + (3, 3, + np.array([[150., 150., 120.], + [150., 120., 120.], + [150., 120., 120.]])), + (4, 4, + np.array([[150., 150., 135., 120.], + [150., 120., 120., np.nan], + [150., 120., 120., 120.], + [150., np.nan, 120., 120.]])), + (5, 5, + np.array([[150., 150., 150., 120., 120.], + [150., 120., np.nan, 120., np.nan], + [150., np.nan, 120., np.nan, 120.], + [150., 120., np.nan, 120., np.nan], + [150., np.nan, 120., np.nan, 120.]])) + ]) + def test_analysis_get_z_surface(self, universe, x_bin, y_bin, expected_surface): + mc = MembraneCurvature(universe, + select='name PO4', + 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): + 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).run() + avg_mean = mc.results.average_mean + assert_almost_equal(avg_mean, expected_mean) + + + def test_analysis_mean_default_selection(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, + 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