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
|
"""
Demo program that solves the Navier Stokes equations in a triply
periodic domain. The solution is initialized using the Taylor-Green
vortex and evolved in time with a 4'th order Runge Kutta method.
Please note that this is not an optimized solver. For fast solvers, see
http://github.com/spectralDNS/spectralDNS
"""
from time import time
import numpy as np
from mpi4py import MPI
from mpi4py_fft import PFFT, newDistArray
# Set viscosity, end time and time step
nu = 0.000625
T = 0.1
dt = 0.01
# Set global size of the computational box
M = 6
N = [2**M, 2**M, 2**M]
L = np.array([2*np.pi, 4*np.pi, 4*np.pi], dtype=float) # Needs to be (2*int)*pi in all directions (periodic) because of initialization
# Create instance of PFFT to perform parallel FFT + an instance to do FFT with padding (3/2-rule)
FFT = PFFT(MPI.COMM_WORLD, N, collapse=False)
#FFT_pad = PFFT(MPI.COMM_WORLD, N, padding=[1.5, 1.5, 1.5])
FFT_pad = FFT
# Declare variables needed to solve Navier-Stokes
U = newDistArray(FFT, False, rank=1, view=True) # Velocity
U_hat = newDistArray(FFT, rank=1, view=True) # Velocity transformed
P = newDistArray(FFT, False, view=True) # Pressure (scalar)
P_hat = newDistArray(FFT, view=True) # Pressure transformed
U_hat0 = newDistArray(FFT, rank=1, view=True) # Runge-Kutta work array
U_hat1 = newDistArray(FFT, rank=1, view=True) # Runge-Kutta work array
a = [1./6., 1./3., 1./3., 1./6.] # Runge-Kutta parameter
b = [0.5, 0.5, 1.] # Runge-Kutta parameter
dU = newDistArray(FFT, rank=1, view=True) # Right hand side of ODEs
curl = newDistArray(FFT, False, rank=1, view=True)
U_pad = newDistArray(FFT_pad, False, rank=1, view=True)
curl_pad = newDistArray(FFT_pad, False, rank=1, view=True)
def get_local_mesh(FFT, L):
"""Returns local mesh."""
X = np.ogrid[FFT.local_slice(False)]
N = FFT.global_shape()
for i in range(len(N)):
X[i] = (X[i]*L[i]/N[i])
X = [np.broadcast_to(x, FFT.shape(False)) for x in X]
return X
def get_local_wavenumbermesh(FFT, L):
"""Returns local wavenumber mesh."""
s = FFT.local_slice()
N = FFT.global_shape()
# Set wavenumbers in grid
k = [np.fft.fftfreq(n, 1./n).astype(int) for n in N[:-1]]
k.append(np.fft.rfftfreq(N[-1], 1./N[-1]).astype(int))
K = [ki[si] for ki, si in zip(k, s)]
Ks = np.meshgrid(*K, indexing='ij', sparse=True)
Lp = 2*np.pi/L
for i in range(3):
Ks[i] = (Ks[i]*Lp[i]).astype(float)
return [np.broadcast_to(k, FFT.shape(True)) for k in Ks]
X = get_local_mesh(FFT, L)
K = get_local_wavenumbermesh(FFT, L)
K = np.array(K).astype(float)
K2 = np.sum(K*K, 0, dtype=float)
K_over_K2 = K.astype(float) / np.where(K2 == 0, 1, K2).astype(float)
def cross(x, y, z):
"""Cross product z = x \times y"""
z[0] = FFT_pad.forward(x[1]*y[2]-x[2]*y[1], z[0])
z[1] = FFT_pad.forward(x[2]*y[0]-x[0]*y[2], z[1])
z[2] = FFT_pad.forward(x[0]*y[1]-x[1]*y[0], z[2])
return z
def compute_curl(x, z):
z[2] = FFT_pad.backward(1j*(K[0]*x[1]-K[1]*x[0]), z[2])
z[1] = FFT_pad.backward(1j*(K[2]*x[0]-K[0]*x[2]), z[1])
z[0] = FFT_pad.backward(1j*(K[1]*x[2]-K[2]*x[1]), z[0])
return z
def compute_rhs(rhs):
for j in range(3):
U_pad[j] = FFT_pad.backward(U_hat[j], U_pad[j])
curl_pad[:] = compute_curl(U_hat, curl_pad)
rhs = cross(U_pad, curl_pad, rhs)
P_hat[:] = np.sum(rhs*K_over_K2, 0, out=P_hat)
rhs -= P_hat*K
rhs -= nu*K2*U_hat
return rhs
# Initialize a Taylor Green vortex
U[0] = np.sin(X[0])*np.cos(X[1])*np.cos(X[2])
U[1] = -np.cos(X[0])*np.sin(X[1])*np.cos(X[2])
U[2] = 0
for i in range(3):
U_hat[i] = FFT.forward(U[i], U_hat[i])
# Integrate using a 4th order Rung-Kutta method
t = 0.0
tstep = 0
t0 = time()
while t < T-1e-8:
t += dt
tstep += 1
U_hat1[:] = U_hat0[:] = U_hat
for rk in range(4):
dU = compute_rhs(dU)
if rk < 3:
U_hat[:] = U_hat0 + b[rk]*dt*dU
U_hat1[:] += a[rk]*dt*dU
U_hat[:] = U_hat1[:]
for i in range(3):
U[i] = FFT.backward(U_hat[i], U[i])
#k = MPI.COMM_WORLD.reduce(sum(U*U)/N[0]/N[1]/N[2]/2)
#if MPI.COMM_WORLD.Get_rank() == 0:
#print("Energy = {}".format(k))
## Transform result to real physical space
#for i in range(3):
#U[i] = FFT.backward(U_hat[i], U[i])
# Check energy
k = MPI.COMM_WORLD.reduce(np.sum(U*U)/N[0]/N[1]/N[2]/2)
if MPI.COMM_WORLD.Get_rank() == 0:
print('Time = {}'.format(time()-t0))
assert round(float(k) - 0.124953117517, 7) == 0
|
