-
Notifications
You must be signed in to change notification settings - Fork 0
/
bispectrum.py
357 lines (237 loc) · 13.4 KB
/
bispectrum.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
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
import numpy
import h5py
import pympi
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
bias, grid_size, box_size, bin_size, bins_to_use = (1.3, 2048, 1600.0, 30, 5)
density_file, momentumx_file, momentumy_file, momentumz_file = ()
deltagk_file, deltamk_file, deltavk_file = ()
num_particle, box_size, grid_size = (4096**3, 1600.0, 2048)
mpi = pympi.snapshot(num_particle=num_particle, box_size=box_size, grid_size=grid_size, comm=comm)
bins_large = numpy.linspace(4e-3, 4e-2, 5+1)
bins_small = numpy.linspace(1.0, 2.0, 5+1)
lsize = bins_large.size - 1
ssize = bins_small.size - 1
bins_large_size = lsize
bins_small_size = ssize
kfreq = numpy.fft.fftfreq(mpi.grid_size, mpi.grid_spacing) * numpy.pi * 2
kfreq = kfreq.astype(numpy.float64)
kx = numpy.empty((mpi.grids_per_core, mpi.grid_size, mpi.grid_size), dtype=numpy.float64)
ky = numpy.empty((mpi.grids_per_core, mpi.grid_size, mpi.grid_size), dtype=numpy.float64)
kz = numpy.empty((mpi.grids_per_core, mpi.grid_size, mpi.grid_size), dtype=numpy.float64)
for i in range(mpi.grids_per_core):
for j in range(mpi.grid_size):
for k in range(mpi.grid_size):
kz[i,j,k] = kfreq[k]
ky[i,j,k] = kfreq[j]
kx[i,j,k] = kfreq[mpi.start_grid + i]
knorm = numpy.sqrt(kx**2 + ky**2 + kz**2)
if rank == 0:
knorm[0,0,0] = 0.00001
def transfer(data):
recvdata = None
senddata = numpy.concatenate( (numpy.ravel(data.real), numpy.ravel(data.imag) ))
if rank == 0:
recvdata = numpy.empty(size * len(senddata), dtype=numpy.float64)
comm.Gather(senddata, recvdata, root=0)
if rank == 0:
d = recvdata.reshape(size, len(senddata)).sum(0)
d_list = numpy.split(d, 2)
data_real, data_imag = [
numpy.reshape( x, (bins_large_size, bins_small_size)) for x in alpha_list]
result = data_real + 1j * data_imag
return result
###############################################################################################################
alpha_mnn = numpy.zeros((bins_large_size, bins_small_size), dtype=numpy.complex128)
alpha_nmn = numpy.zeros((bins_large_size, bins_small_size), dtype=numpy.complex128)
alpha_final = numpy.zeros((bins_large_size, bins_small_size), dtype=numpy.complex128)
alpha_final_inv = numpy.zeros((bins_large_size, bins_small_size), dtype=numpy.complex128)
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu)
function3 = mpi.ifft(function1 * function1)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
alpha_mnn[m,n] = numpy.sum(function3 * W_mu * knorm**2)
function2 = mpi.ifft(W_nu * knorm**2)
function3 = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
alpha_nmn[m,n] = numpy.sum(function3 * W_mu)
alpha_final = transfer(alpha_mnn)
alpha_final_inv = transfer(alpha_nmn)
gamma_x_mnn = numpy.zeros((bins_large.size-1,bins_small.size-1), dtype=numpy.complex128)
gamma_y_mnn = numpy.zeros((bins_large.size-1,bins_small.size-1), dtype=numpy.complex128)
gamma_z_mnn = numpy.zeros((bins_large.size-1,bins_small.size-1), dtype=numpy.complex128)
gamma_mnn = numpy.zeros((bins_large.size-1,bins_small.size-1), dtype=numpy.complex128)
gamma_final = numpy.zeros((bins_large.size-1,bins_small.size-1), dtype=numpy.complex128)
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu)
function2x = mpi.ifft(W_nu * kx)
function2y = mpi.ifft(W_nu * ky)
function2z = mpi.ifft(W_nu * kz)
productx = mpi.ifft(function1 * function2x)
producty = mpi.ifft(function1 * function2y)
productz = mpi.ifft(function1 * function2z)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
gamma_x_mnn[m,n] = numpy.sum(kx * W_mu * productx)
gamma_y_mnn[m,n] = numpy.sum(ky * W_mu * producty)
gamma_z_mnn[m,n] = numpy.sum(kz * W_mu * productz)
gamma_mnn = gamma_x_mnn + gamma_y_mnn + gamma_z_mnn
gamma_final = transfer(gamma_mnn)
#########################################################################################################
B_x_mnn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
B_y_mnn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
B_z_mnn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
B_mnn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
B_x_nmn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
B_y_nmn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
B_z_nmn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
B_nmn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
B_final = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
B_final_inv = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
density_k = h5py.File(deltagk_file, "r+")['d'][mpi.start_grid:mpi.end_grid,:,:]
momentum_kx = h5py.File(momentumx_file, "r+")['d'][mpi.start_grid:mpi.end_grid, :, :]
momentum_ky = h5py.File(momentumy_file, "r+")['d'][mpi.start_grid:mpi.end_grid, :, :]
momentum_kz = h5py.File(momentumz_file, "r+")['d'][mpi.start_grid:mpi.end_grid, :, :]
###############################################################################################################
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu * density_k)
function2 = mpi.ifft(W_nu * momentum_kx)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
B_x_mnn[m,n] = numpy.sum(1j * kx * W_mu * density_k * product)
function1 = mpi.ifft(1j * kx * W_nu * density_k)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
B_x_nmn[m,n] = numpy.sum(W_mu * density_k * product)
###############################################################################################################
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu * density_k)
function2 = mpi.ifft(W_nu * momentum_ky)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
B_y_mnn[m,n] = numpy.sum(1j * ky * W_mu * density_k * product)
function1 = mpi.ifft(1j * ky * W_nu * density_k)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
B_y_nmn[m,n] = numpy.sum(W_mu * density_k * product)
###############################################################################################################
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu * density_k)
function2 = mpi.ifft(W_nu * momentum_kz)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
B_z_mnn[m,n] = numpy.sum(1j * kz * W_mu * density_k * product)
function1 = mpi.ifft(1j * kz * W_nu * density_k)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
B_z_nmn[m,n] = numpy.sum(W_mu * density_k * product)
B_mnn = -(B_x_mnn + B_y_mnn + B_z_mnn)
B_nmn = -(B_x_nmn + B_y_nmn + B_z_nmn)
B_final = transfer(B_mnn)
B_final_inv = transfer(B_nmn)
#################################################################################
omega_m = 0.295037918
kmps_to_kpcgyr = 1.022712
fah_factor = (omega_m**0.55)*100*1.0*kmps_to_kpcgyr
faH = fah_factor/knorm
deltaGk = density_k
deltaMk = h5py.File(density_file, "r+")['d'][mpi.start_grid:mpi.end_grid,:,:]
deltaVk = faH * deltaMk
Bv_x_mnn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
Bv_y_mnn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
Bv_z_mnn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
Bv_mnn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
Bv_x_nmn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
Bv_y_nmn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
Bv_z_nmn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
Bv_nmn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
Bv_final = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
Bv_final_inv = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
B1_mnn = numpy.zeros((lsize, ssize), dtype=numpy.complex128)
B1_nmn = numpy.zeros((lsize,ssize), dtype=numpy.complex128)
###################################################################################################################################
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu * deltaGk * numpy.conjugate(deltaMk))
function2 = mpi.ifft(W_nu)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
B1_mnn[m,n] = numpy.sum( knorm * W_mu * deltaGk * numpy.conjugate(deltaVk) * product)
function1 = mpi.ifft(W_nu * knorm * deltaGk * numpy.conjugate(deltaVk))
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
B1_nmn[m,n] = numpy.sum( W_mu * deltaGk * numpy.conjugate(deltaMk) * product)
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu * kx * deltaGk * numpy.conjugate(deltaVk)/knorm)
function2 = mpi.ifft(W_nu)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
Bv_x_mnn[m,n] = numpy.sum(W_mu * kx * deltaGk * numpy.conjugate(deltaMk) * product)
function1 = mpi.ifft(W_nu * kx * deltaGk * numpy.conjugate(deltaMk))
product = mpi.ifft( function1 * function2 )
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
Bv_x_nmn[m,n] = numpy.sum(W_mu * kx * deltaGk * numpy.conjugate(deltaVk)/knorm * product)
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu * ky * deltaGk * numpy.conjugate(deltaVk)/knorm)
function2 = mpi.ifft(W_nu)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
Bv_y_mnn[m,n] = numpy.sum(W_mu * ky * deltaGk * numpy.conjugate(deltaMk) * product)
function1 = mpi.ifft(W_nu * ky * deltaGk * numpy.conjugate(deltaMk))
product = mpi.ifft( function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
Bv_y_nmn[m,n] = numpy.sum(W_mu * ky * deltaGk * numpy.conjugate(deltaVk)/knorm * product)
for n in xrange(ssize):
W_nu = numpy.logical_and(knorm > bins_small[n], knorm <= bins_small[n+1]).astype(numpy.int16)
function1 = mpi.ifft(W_nu * kz * deltaGk * numpy.conjugate(deltaVk)/knorm)
function2 = mpi.ifft(W_nu)
product = mpi.ifft(function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
Bv_z_mnn[m,n] = numpy.sum(W_mu * kz * deltaGk * numpy.conjugate(deltaMk) * product)
function1 = mpi.ifft(W_nu * kz * deltaGk * numpy.conjugate(deltaMk))
product = mpi.ifft( function1 * function2)
for m in xrange(lsize):
W_mu = numpy.logical_and(knorm > bins_large[m], knorm <= bins_large[m+1]).astype(numpy.int16)
Bv_z_nmn[m,n] = numpy.sum(W_mu * kz * deltaGk * numpy.conjugate(deltaVk)/knorm * product)
Bv_mnn = -(B1_mnn + Bv_x_mnn + Bv_y_mnn + Bv_z_mnn)
Bv_nmn = -(B1_nmn + Bv_x_nmn + Bv_y_nmn + Bv_z_nmn)
######################################################################################################################################
Bv_final = transfer(Bv_mnn)
Bv_final_inv = transfer(Bv_nmn)
"""
if rank == 0:
volume = boxsize**3
Bv1 = numpy.zeros(( binsize, binsize, binsize ), dtype=complex)
Bv2 = numpy.zeros(( binsize, binsize, binsize ), dtype=complex)
for m in xrange( binsize ):
for n in xrange( binsize ):
matrix = numpy.linalg.inv( numpy.array( [ [ alpha_mnn[m,n], gamma_mnn[m,n] ], [gamma_mnn[n,m], alpha_nmn[m,n]]] ) )
Bv1[m,n,n] = matrix[0,0] * B_final[m,n] + matrix[0,1] * B_final_inv[m,n]
Bv2[m,n,n] = matrix[0,0] * Bv_final[m,n] + matrix[0,1] * Bv_final_inv[m,n]
for i in xrange(lsize):
for j in xrange(ssize):
"""