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
22def derive_surface(atoms, n_cells_x, n_cells_y, max_width_x, max_width_y):
23 """
24 Derive surface from `atom` positions in `AtomGroup`.
26 Parameters
27 ----------
28 atoms: AtomGroup.
29 AtomGroup of reference selection to define the surface
30 of the membrane.
31 n_cells_x: int.
32 number of cells in the grid of size `max_width_x`, `x` axis.
33 n_cells_y: int.
34 number of cells in the grid of size `max_width_y`, `y` axis.
35 max_width_x: float.
36 Maximum width of simulation box in x axis. (Determined by simulation box dimensions)
37 max_width_y: float.
38 Maximum width of simulation box in y axis. (Determined by simulation box dimensions)
40 Returns
41 -------
42 z_coordinates: np.ndarray
43 Average z-coordinate values. Return Numpy array of floats of
44 shape `(n_cells_x, n_cells_y)`.
46 """
47 coordinates = atoms.positions
48 return get_z_surface(coordinates, n_x_bins=n_cells_x, n_y_bins=n_cells_y,
49 x_range=(0, max_width_x), y_range=(0, max_width_y))
52def get_z_surface(coordinates, n_x_bins=10, n_y_bins=10, x_range=(0, 100), y_range=(0, 100)):
53 """
54 Derive surface from distribution of z coordinates in grid.
56 Parameters
57 ----------
58 coordinates : np.ndarray
59 Coordinates of AtomGroup. NumPy array of shape=(n_atoms, 3).
60 n_x_bins : int.
61 Number of bins in grid in the `x` dimension.
62 n_y_bins : int.
63 Number of bins in grid in the `y` dimension.
64 x_range : tuple of (float, float)
65 Range of coordinates (min, max) in the `x` dimension with shape=(2,).
66 y_range : tuple of (float, float)
67 Range of coordinates (min, max) in the `y` dimension with shape=(2,).
69 Returns
70 -------
71 z_surface: np.ndarray
72 Surface derived from set of coordinates in grid of `x_range, y_range` dimensions.
73 Returns NumPy array of floats of shape (`n_x_bins`, `n_y_bins`)
75 """
77 grid_z_coordinates = np.zeros((n_x_bins, n_y_bins))
78 grid_norm_unit = np.zeros((n_x_bins, n_y_bins))
80 x_factor = n_x_bins / (x_range[1] - x_range[0])
81 y_factor = n_y_bins / (y_range[1] - y_range[0])
83 x_coords, y_coords, z_coords = coordinates.T
85 cell_x_floor = np.floor(x_coords * x_factor).astype(int)
86 cell_y_floor = np.floor(y_coords * y_factor).astype(int)
88 for l, m, z in zip(cell_x_floor, cell_y_floor, z_coords):
90 try:
91 if l < 0 or m < 0:
92 msg = ("Atom outside grid boundaries. Skipping atom.")
93 warnings.warn(msg)
94 continue
96 grid_z_coordinates[l, m] += z
97 grid_norm_unit[l, m] += 1
99 except IndexError:
100 msg = ("Atom outside grid boundaries. Skipping atom.")
101 warnings.warn(msg)
103 z_surface = normalized_grid(grid_z_coordinates, grid_norm_unit)
105 return z_surface
108def normalized_grid(grid_z_coordinates, grid_norm_unit):
109 """
110 Calculates average `z` coordinates in unit cell.
112 Parameters
113 ----------
115 z_ref: np.array
116 Empty array of `(l,m)`
117 grid_z_coordinates: np.array
118 Array of size `(l,m)` with `z` coordinates stored in unit cell.
119 grid_norm_unit: np.array
120 Array of size `(l,m)` with number of atoms in unit cell.
122 Returns
123 -------
124 z_surface: np.ndarray
125 Normalized `z` coordinates in grid.
126 Returns Numpy array of floats of shape (`n_x_bins`, `n_y_bins`)
128 """
130 grid_norm_unit = np.where(grid_norm_unit > 0, grid_norm_unit, np.nan)
131 z_normalized = grid_z_coordinates / grid_norm_unit
133 return z_normalized
136def interpolation_by_array(array_surface):
137 """
138 Interpolates values contained in `array_surface` over axis
140 Parameters
141 ----------
143 array_surface: np.ndarray
144 Array of floats of shape (`n_x_bins`, `n_y_bins`)
146 Returns
147 -------
148 interpolated_array: np.ndarray
149 Returns interpolated array.
150 Array of shape (`n_x_bins` x `n_y_bins`,)
152 """
154 # create mask for nans
155 mask_nans = ~np.isnan(array_surface)
157 # index of array_surface
158 index_array = np.arange(array_surface.shape[0])
160 # interpolate values in array
161 interpolated_array = np.interp(index_array,
162 np.flatnonzero(mask_nans),
163 array_surface[mask_nans])
165 return interpolated_array
168def surface_interpolation(array_surface):
169 """
170 Calculates interpolation of arrays.
172 Parameters
173 ----------
175 array_surface: np.ndarray
176 Array of floats of shape (`n_x_bins`, `n_y_bins`)
178 Returns
179 -------
181 interpolated_surface: np.ndarray
182 Interpolated surface derived from set of coordinates
183 in grid of `x_range, y_range` dimensions.
184 Array of floats of shape (`n_x_bins`, `n_y_bins`)
186 """
188 interpolated_surface = np.apply_along_axis(interpolation_by_array, axis=0, arr=array_surface)
190 return interpolated_surface