Coverage for /Users/estefania/python_playground/membrane-curvature/membrane_curvature/surface.py : 100%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1r"""
2.. role:: raw-math(raw) :format: latex html
4--------------------
5Surface
6--------------------
8Calculation of curvature requires a surface of reference. In MembraneCurvature,
9the surface of reference is defined by the `z` position of the `atoms` in `AtomGroup`.
12Functions
13---------
16"""
18import numpy as np
19import warnings
20import MDAnalysis
21import logging
23MDAnalysis.start_logging()
24logger = logging.getLogger("MDAnalysis.MDAKit.membrane_curvature")
27def derive_surface(atoms, n_cells_x, n_cells_y, max_width_x, max_width_y):
28 """
29 Derive surface from `atom` positions in `AtomGroup`.
31 Parameters
32 ----------
33 atoms: AtomGroup.
34 AtomGroup of reference selection to define the surface
35 of the membrane.
36 n_cells_x: int.
37 number of cells in the grid of size `max_width_x`, `x` axis.
38 n_cells_y: int.
39 number of cells in the grid of size `max_width_y`, `y` axis.
40 max_width_x: float.
41 Maximum width of simulation box in x axis. (Determined by simulation box dimensions)
42 max_width_y: float.
43 Maximum width of simulation box in y axis. (Determined by simulation box dimensions)
45 Returns
46 -------
47 z_coordinates: numpy.ndarray
48 Average z-coordinate values. Return Numpy array of floats of
49 shape `(n_cells_x, n_cells_y)`.
51 """
52 coordinates = atoms.positions
53 return get_z_surface(coordinates, n_x_bins=n_cells_x, n_y_bins=n_cells_y,
54 x_range=(0, max_width_x), y_range=(0, max_width_y))
57def get_z_surface(coordinates, n_x_bins=10, n_y_bins=10, x_range=(0, 100), y_range=(0, 100)):
58 """
59 Derive surface from distribution of z coordinates in grid.
61 Parameters
62 ----------
63 coordinates : numpy.ndarray
64 Coordinates of AtomGroup. Numpy array of shape=(n_atoms, 3).
65 n_x_bins : int.
66 Number of bins in grid in the `x` dimension.
67 n_y_bins : int.
68 Number of bins in grid in the `y` dimension.
69 x_range : tuple of (float, float)
70 Range of coordinates (min, max) in the `x` dimension with shape=(2,).
71 y_range : tuple of (float, float)
72 Range of coordinates (min, max) in the `y` dimension with shape=(2,).
74 Returns
75 -------
76 z_surface: np.ndarray
77 Surface derived from set of coordinates in grid of `x_range, y_range` dimensions.
78 Returns Numpy array of floats of shape (`n_x_bins`, `n_y_bins`)
80 """
82 grid_z_coordinates = np.zeros((n_x_bins, n_y_bins))
83 grid_norm_unit = np.zeros((n_x_bins, n_y_bins))
85 x_factor = n_x_bins / (x_range[1] - x_range[0])
86 y_factor = n_y_bins / (y_range[1] - y_range[0])
88 x_coords, y_coords, z_coords = coordinates.T
90 cell_x_floor = np.floor(x_coords * x_factor).astype(int)
91 cell_y_floor = np.floor(y_coords * y_factor).astype(int)
93 for l, m, z in zip(cell_x_floor, cell_y_floor, z_coords):
95 try:
96 # negative coordinates
97 if l < 0 or m < 0:
98 msg = ("Atom with negative coordinates falls "
99 "outside grid boundaries. Element "
100 "({},{}) in grid can't be assigned."
101 " Skipping atom.").format(l, m)
102 warnings.warn(msg)
103 logger.warning(msg)
104 continue
106 grid_z_coordinates[l, m] += z
107 grid_norm_unit[l, m] += 1
109 # too large positive coordinates
110 except IndexError:
111 msg = ("Atom coordinates exceed size of grid "
112 "and element ({},{}) can't be assigned. "
113 "Maximum (x,y) coordinates must be < ({}, {}). "
114 "Skipping atom.").format(l, m, x_range[1], y_range[1])
115 warnings.warn(msg)
116 logger.warning(msg)
118 z_surface = normalized_grid(grid_z_coordinates, grid_norm_unit)
120 return z_surface
123def normalized_grid(grid_z_coordinates, grid_norm_unit):
124 """
125 Calculates average `z` coordinates in unit cell.
127 Parameters
128 ----------
130 z_ref: np.array
131 Empty array of `(l,m)`
132 grid_z_coordinates: np.array
133 Array of size `(l,m)` with `z` coordinates stored in unit cell.
134 grid_norm_unit: np.array
135 Array of size `(l,m)` with number of atoms in unit cell.
137 Returns
138 -------
139 z_surface: np.ndarray
140 Normalized `z` coordinates in grid.
141 Returns Numpy array of floats of shape (`n_x_bins`, `n_y_bins`)
143 """
145 grid_norm_unit = np.where(grid_norm_unit > 0, grid_norm_unit, np.nan)
146 z_normalized = grid_z_coordinates / grid_norm_unit
148 return z_normalized