-
Notifications
You must be signed in to change notification settings - Fork 1
/
mpi_run_analyses.py
183 lines (132 loc) · 6.72 KB
/
mpi_run_analyses.py
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
# Script to run Eyring Model analyses
import numpy as np
import matplotlib.pyplot as plt
from mpi4py import MPI
from time import perf_counter as timer
from eyring_model import EyringModel, Path
# Define global constants
global kB
kB = 1.380649 * 10**-23 # Boltzmann (m^2 kg / s^2 K)
global h
h = 6.62607 * 10**-34 # Planck (m^2 kg / s)
global R
R = 1.9858775 * 10**-3 # universal gas (kcal / mol K)
def path_convergence(realizations, dH_barrier, dS_barrier, dH_sigma, dS_sigma, dist, T=300, multi=True, large_n_paths=5000, method='equal_sums'):
if dist in ['normal', 'N', 'norm']:
params = {'mu' : np.array([dH_barrier, dS_barrier]),
'cov' : np.array([[dH_sigma**2,0],
[0,dS_sigma**2]])}
elif dist in ['exponential', 'exp']:
params = {'beta' : np.array([dH_barrier, dS_barrier])}
elif dist in ['equal', 'single', 'none', None]:
params = {'mu' : np.array([dH_barrier, dS_barrier])}
# run this on all available processors
def calculate_barriers(N, dist, dist_params, T=T, multi=multi):
model = EyringModel(T=T)
for n in range(N):
model.add_Path(n_jumps=200, lam=10, area=model.area/N) # 2 Angstrom jump lengths over 200 jumps, no entropic penalty
model.paths[n].generate_membrane_barriers(dist=dist, dist_params=dist_params, multi=multi)
return model.calculate_effective_barrier()
########### DISTRIBUTE THE WORK ACROSS PROCESSORS ##########
comm = MPI.COMM_WORLD
nprocs = comm.Get_size()
rank = comm.Get_rank()
my_start = timer()
if rank == 0:
print(f'\n{nprocs} processors calculating the path convergence of {dist} distributions of barriers')
fig1, ax1 = plt.subplots(1,1, figsize=(3.55,3.55))
# Increasing number of paths to see how effective barrier changes (when it stabilizes)
n_paths = np.array([50,100,200,300,400,500,600,700,800,900,1000,1100,1200,1300,1400,1500,
1600,1700,1800,1900,2000,2100,2200,2300,2400,2500,2600,2700,2800,2900,
3000,3100,3200,3300,3400,3500,3600,3700,3800,3900,4000,4100,4200,4300,
4400,4500,4600,4700,4800,4900,5000,5500,6000,6500,7000,
7500,8000,8500,9000,9500,10_000], dtype='i')
# split up n_paths for each proc
if method == 'naive_split':
ave, res = divmod(n_paths.size, nprocs)
count = np.array([ave + 1 if p < res else ave for p in range(nprocs)], dtype='i')
elif method == 'large_single_proc':
large_array = np.where(n_paths >= large_n_paths)[0] # if number of paths is large, assign to its own processor
n_large = len(large_array)
rest = np.where(n_paths < large_n_paths)[0]
ave, res = divmod(len(rest), nprocs-n_large) # split the rest among the remaining processors
if n_large > 15:
raise ValueError(f'{n_large} paths are considered large. Please choose a larger cutoff value or a different method.')
count = np.zeros(nprocs, dtype='i')
count[:nprocs-n_large] = np.array([ave+1 if p < res else ave for p in range(nprocs-n_large)], dtype='i')
count[nprocs-n_large:] = np.ones(n_large)
elif method == 'equal_sums':
per_proc = n_paths.sum() / nprocs
np.random.shuffle(n_paths)
smaller_arrays = []
start = 0
end = 0
for i in range(nprocs):
end += np.searchsorted(np.cumsum(n_paths[start:], dtype=float), per_proc)
smaller_arrays.append(n_paths[start:end].tolist())
start = end
if start != len(n_paths):
for i, n in enumerate(n_paths[start:]):
smaller_arrays[i].append(n)
count = np.array([len(small) for small in smaller_arrays], dtype='i')
# save the indices where each processor starts
displ = np.array([sum(count[:p]) for p in range(nprocs)])
# save to calculate average across realizations
avg_barrier = np.zeros(len(n_paths))
avg_permeability = np.zeros(len(n_paths))
else:
# initialize n_paths, count, displ on workers
n_paths = None
count = np.zeros(nprocs, dtype='i')
displ = None
# send count to all processors
comm.Bcast(count, root=0)
# initialize my_paths on all processors
my_paths = np.zeros(count[rank], dtype='i')
# break the array up and send to processors
comm.Scatterv([n_paths, count, displ, MPI.INT], my_paths, root=0)
print(f'Rank {rank} is tasked with n_paths: {my_paths}, total: {my_paths.sum()}')
# calculate effective barriers and permeabilities on each processor
effective_barriers = np.zeros(len(my_paths))
permeabilities = np.zeros(len(my_paths))
########## RUN A BUNCH OF REALIZATIONS ##########
for r in range(realizations):
if rank == 0:
print(f'Realization {r+1}...')
for i,N in enumerate(my_paths):
effective_barriers[i] = calculate_barriers(N, dist, params)
# create big arrays for Gather
all_barriers = np.zeros(sum(count))
# gather all the barriers and permeabilities back to master for plotting
comm.Gatherv(effective_barriers, [all_barriers, count, displ, MPI.DOUBLE], root=0)
if rank == 0:
# save for average across realizations
avg_barrier += all_barriers
ax1.scatter(n_paths, all_barriers, alpha=0.1, c='k')
if rank == 0:
data = np.column_stack((n_paths, avg_barrier/realizations))
np.savetxt(f'avg_convergence_{dist}_{realizations}iter.csv', data, delimiter=',')
ax1.scatter(n_paths, avg_barrier / realizations, c='r', label='average over paths')
ax1.set_xlabel('Number of Paths')
ax1.set_ylabel('$\Delta G_{eff}^{\ddag}$ (kcal/mol)')
fig1.savefig(f'figs/effective_barrier_convergence_{dist}.pdf')
plt.show()
my_end = timer()
print(f'Total time for rank {rank}: {my_end - my_start:.2f} s')
if __name__ == "__main__":
plt.rcParams['text.usetex'] = True
plt.rcParams['font.family'] = 'Helvetica'
plt.rcParams['font.size'] = 8
plt.rc('text.latex', preamble=r'\usepackage[cm]{sfmath}')
# Inputs
iters = 100
T = 300
multi = True
dH_barrier = 3.5
dS_barrier = -9/T
dH_sigma = dH_barrier/3
dS_sigma = -dS_barrier/3
# normal distributions
path_convergence(iters, dH_barrier, dS_barrier, dH_sigma, dS_sigma, dist='norm', T=T, multi=multi)
# exponential distributions
path_convergence(iters, dH_barrier, dS_barrier, dH_sigma, dS_sigma, dist='exp', T=T, multi=multi)