membrane-curvature/benchmarks/membranecurvature.py at f197846faf5d9bf8e97fca1bee58329ddf6d31f4 · MDAnalysis/membrane-curvature · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
import MDAnalysis as mda
import numpy as np

from membrane_curvature.base import MembraneCurvature
from membrane_curvature.tests.datafiles import GRO_PO4, GRO_PO4_SMALL, XTC_PO4

PO4_SELECT = 'name PO4'


def po4_universe(gro=GRO_PO4):
    """Universe with first frame loaded"""
    universe = mda.Universe(gro)
    universe.trajectory[0]
    return universe


def po4_frame0_context(gro=GRO_PO4):
    """Positions and box ranges for one PO4 frame"""
    universe = po4_universe(gro)
    atomgroup = universe.select_atoms(PO4_SELECT)
    return {
        'positions': atomgroup.positions.copy(),
        'x_range': (0.0, float(universe.dimensions[0])),
        'y_range': (0.0, float(universe.dimensions[1])),
    }


def dummy_height_grid(n_bins=100, seed=0):
    """Smooth height field for gradient / Monge curvature micro-benchmarks."""
    rng = np.random.default_rng(seed)
    x = np.linspace(0, 2 * np.pi, n_bins)
    y = np.linspace(0, 2 * np.pi, n_bins)
    xx, yy = np.meshgrid(x, y)
    z = np.sin(xx) * np.cos(yy) + 0.1 * rng.standard_normal((n_bins, n_bins))
    return z, 0.1, 0.1


class MembraneCurvatureSmallBenchmark:
    def setup(self):
        self.universe = po4_universe(GRO_PO4_SMALL)

    def time_run_fourier_low_modes(self):
        # Nine PO4 atoms support at most fourier_m = fourier_n = 1.
        MembraneCurvature(self.universe, select=PO4_SELECT, fourier_m=1, fourier_n=1).run()

    def time_run_binning_default_wrap(self):
        MembraneCurvature(
            self.universe,
            select=PO4_SELECT,
            surface_method='binning',
            n_x_bins=3,
            n_y_bins=3,
        ).run()

    def time_run_binning_no_wrap(self):
        MembraneCurvature(
            self.universe,
            select=PO4_SELECT,
            surface_method='binning',
            n_x_bins=3,
            n_y_bins=3,
            wrap=False,
        ).run()


class MembraneCurvatureBenchmark:
    """Benchmark for PO4-only membrane (~900 atoms), regular grid sizes."""

    params = ([25, 50, 100],)
    param_names = ['n_bins']

    def setup(self, n_bins):
        self.universe = po4_universe(GRO_PO4)

    def time_run_fourier(self, n_bins):
        MembraneCurvature(
            self.universe,
            select=PO4_SELECT,
            n_x_bins=n_bins,
            n_y_bins=n_bins,
        ).run()

    def time_run_binning_wrap(self, n_bins):
        MembraneCurvature(
            self.universe,
            select=PO4_SELECT,
            surface_method='binning',
            n_x_bins=n_bins,
            n_y_bins=n_bins,
        ).run()

    def time_run_binning_no_wrap(self, n_bins):
        MembraneCurvature(
            self.universe,
            select=PO4_SELECT,
            surface_method='binning',
            n_x_bins=n_bins,
            n_y_bins=n_bins,
            wrap=False,
        ).run()


class MembraneCurvatureFourierModesBenchmark:
    """Benchmarks for Fourier modes with PO4 selection."""

    params = ([(2, 2), (3, 3), (5, 5)],)
    param_names = ['modes']

    def setup(self, modes):
        self.universe = po4_universe(GRO_PO4)
        self.fourier_m, self.fourier_n = modes

    def time_run_fourier(self, modes):
        MembraneCurvature(
            self.universe,
            select=PO4_SELECT,
            fourier_m=self.fourier_m,
            fourier_n=self.fourier_n,
        ).run()


class MembraneCurvatureTrajectoryBenchmark:
    """Multi-frame trajectory"""

    def setup(self):
        self.universe = mda.Universe(GRO_PO4, XTC_PO4)

    def time_run_fourier_trajectory(self):
        MembraneCurvature(self.universe, select=PO4_SELECT).run()

    def time_run_binning_trajectory(self):
        MembraneCurvature(
            self.universe,
            select=PO4_SELECT,
            surface_method='binning',
            n_x_bins=3,
            n_y_bins=3,
        ).run()

    def peakmem_run_fourier_trajectory(self):
        MembraneCurvature(self.universe, select=PO4_SELECT).run()

    def peakmem_run_binning_trajectory(self):
        MembraneCurvature(
            self.universe,
            select=PO4_SELECT,
            surface_method='binning',
            n_x_bins=3,
            n_y_bins=3,
        ).run()