AnalysisBase (#43)

* Added AnalysisBase with docstrings and tests.

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
 

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)

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