generated from taichiCourse01/taichi_course_homework
-
Notifications
You must be signed in to change notification settings - Fork 0
/
fft_fluid_solver.py
110 lines (96 loc) · 3.37 KB
/
fft_fluid_solver.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
from numpy.core.fromnumeric import shape
import taichi as ti
import numpy as np
lin_iters = 20
N = 64
dt = 0.1
diff = 0.0
visc = 0.0
force = 5e5
source = 100.0
dvel = False
v = ti.Vector.field(2, float, shape=(N + 2, N + 2), offset = (-1, -1))
v_prev = ti.Vector.field(2, float, shape=(N + 2, N + 2), offset = (-1, -1))
dens = ti.field(float, shape=(N + 2, N + 2), offset = (-1, -1))
dens_prev = ti.field(float, shape=(N + 2, N + 2), offset = (-1, -1))
div = ti.field(float, shape=(N + 2, N + 2), offset = (-1, -1))
p = ti.field(float, shape=(N + 2, N + 2), offset = (-1, -1))
pixels = ti.field(float, shape=(N, N))
@ti.kernel
def add_source(a : ti.template(), b : ti.template()):
for i, j in a:
a[i, j] += dt * b[i, j]
@ti.kernel
def swap(a : ti.template(), b : ti.template()):
for i, j in a:
a[i, j], b[i, j] = b[i, j], a[i, j]
@ti.func
def set_bnd(x : ti.template()):
for i in range(N):
x[-1, i] = x[0, i]
x[N, i] = x[N - 1, i]
x[i, -1] = x[i, 0]
x[i, N] = x[i, N - 1]
x[-1, -1] = (x[0, -1] + x[-1, 0]) / 2.0
x[-1, N] = (x[0, N] + x[-1, N - 1]) / 2.0
x[N, -1] = (x[N - 1, -1] + x[N, 0]) / 2.0
x[N, N] = (x[N - 1, N] + x[N, N - 1]) / 2.0
@ti.kernel
def lin_solve(x : ti.template(), x0 : ti.template(), a : float, c : float):
for i, j in ti.ndrange(N, N):
x[i, j] = (x0[i, j] + a * (x[i - 1, j] + x[i + 1, j] + x[i, j - 1] + x[i, j + 1])) / c
set_bnd(x)
def diffuse(a, a_prev, diff):
k = dt * diff * N * N
for t in range(lin_iters):
lin_solve(a, a_prev, k, 1.0 + 4.0 * k)
@ti.kernel
def advect(d : ti.template(), d0 : ti.template(), v : ti.template() ):
dt0 = dt * N
for i, j in ti.ndrange(N, N):
x, y = i - dt0 * v[i, j][0], j - dt0 * v[i, j][1]
if (x < 0.5): x = 0.5
if (x > N + 0.5): x = N + 0.5
i0, i1 = int(x), int(x) + 1
if (y < 0.5): y = 0.5
if (y > N + 0.5): y = N + 0.5
j0, j1 = int(y), int(y) + 1
s1, s0, t1, t0 = x - i0, i1 - x, y - j0, j1 - y
d[i, j] = s0 * (t0 * d0[i0, j0] + t1 * d0[i0, j1]) + s1 * (t0 * d0[i1, j0] + t1 * d0[i1, j1])
set_bnd(d)
def fft_project(v):
u0 = np.zeros(shape = (N + 2, N))
v0 = np.zeros(shape = (N + 2, N))
for i, j in ti.ndrange(N, N):
u0[i, j], v0[i, j] = v[i, j][0], v[i, j][1]
u0 = np.fft.fft2(u0)
v0 = np.fft.fft2(v0)
for i, j in ti.ndrange(N + 2, N):
x, y = i, j
if j > N // 2 : j = j - N
r = x * x + y * y
if (r == 0.0): continue
f = ti.exp(-r*dt*visc)
U, V = u0[i,j], v0[i,j]
u0[i, j] = f * np.complex((1-x*x/r)*U.real+(-x*y/r)*V.real, (1-x*x/r)*U.imag+(-x*y/r)*V.imag)
v0[i, j] = f * np.complex((-y*x/r)*U.real+(1-y*y/r)*V.real,(-y*x/r)*U.imag+(1-y*y/r)*V.imag)
u0 = np.fft.ifft2(u0)
v0 = np.fft.ifft2(v0)
f = 1.0/(N*N)
for i, j in ti.ndrange(N, N):
v[i, j][0], v[i, j][1] = f * u0[i, j], f * v0[i, j]
print("Okay")
def dens_step():
add_source(dens, dens_prev)
swap(dens, dens_prev)
diffuse(dens, dens_prev, diff)
swap(dens, dens_prev)
advect(dens, dens_prev, v)
def vel_step():
add_source(v, v_prev)
swap(v, v_prev)
diffuse(v, v_prev, visc)
fft_project(v)
swap(v, v_prev)
advect(v, v_prev, v_prev)
fft_project(v)