Skip to content

Instantly share code, notes, and snippets.

@houkensjtu

houkensjtu/chorin_taichi.py

Created June 9, 2023 02:36
Show Gist options
  • Save houkensjtu/dc00a6b8ae3d819b9295d68c55df0a29 to your computer and use it in GitHub Desktop.
Save houkensjtu/dc00a6b8ae3d819b9295d68c55df0a29 to your computer and use it in GitHub Desktop.
Download ZIP
A lid-driven test case using Chorin-style projection method; implemented in Taichi.
Raw
chorin_taichi.py
# Implementation based on @Ceyron's Numpy version:
# https://github.com/Ceyron/machine-learning-and-simulation/blob/main/english/simulation_scripts/lid_driven_cavity_python_simple.py
# Make sure you have Taichi installed:
# pip install -U taichi
import taichi as ti
real = ti.f32
ti.init(default_fp=real, arch=ti.cpu) # you can switch to ti.gpu if you have one
N_POINTS = 41
DOMAIN_SIZE = 1.0
TIME_STEP = 0.001
DENSITY = 1.0
NU = 0.1
TOP_VELOCITY = 1.0
MAX_ITER = 50
scalar_shape = (N_POINTS, N_POINTS)
vector_shape = (N_POINTS, N_POINTS)
u_prev = ti.field(dtype=real, shape=scalar_shape)
v_prev = ti.field(dtype=real, shape=scalar_shape)
p_prev = ti.field(dtype=real, shape=scalar_shape)
u_star = ti.field(dtype=real, shape=scalar_shape)
v_star = ti.field(dtype=real, shape=scalar_shape)
dudx = ti.field(dtype=real, shape=scalar_shape)
dudy = ti.field(dtype=real, shape=scalar_shape)
dvdx = ti.field(dtype=real, shape=scalar_shape)
dvdy = ti.field(dtype=real, shape=scalar_shape)
dpdx = ti.field(dtype=real, shape=scalar_shape)
dpdy = ti.field(dtype=real, shape=scalar_shape)
laplace_u = ti.field(dtype=real, shape=scalar_shape)
laplace_v = ti.field(dtype=real, shape=scalar_shape)
div_u = ti.field(dtype=real, shape=scalar_shape)
u_next = ti.field(dtype=real, shape=scalar_shape)
v_next = ti.field(dtype=real, shape=scalar_shape)
p_next = ti.field(dtype=real, shape=scalar_shape)
element_length = DOMAIN_SIZE / (N_POINTS - 1)
V = ti.Vector.field(2, dtype=real, shape=vector_shape)
@ti.kernel
def central_difference_x(field_in:ti.template(), field_out:ti.template()):
imax = field_in.shape[0]
jmax = field_in.shape[1]
for i, j in ti.ndrange((1, imax-1), (1, jmax-1)):
field_out[i, j] = (
(
field_in[i+1, j]
-
field_in[i-1, j]
)
/
( 2 * element_length)
)
@ti.kernel
def central_difference_y(field_in:ti.template(), field_out:ti.template()):
imax = field_in.shape[0]
jmax = field_in.shape[1]
for i, j in ti.ndrange((1, imax-1), (1, jmax-1)):
field_out[i, j] = (
(
field_in[i, j+1]
-
field_in[i, j-1]
)
/
( 2 * element_length)
)
@ti.kernel
def laplacian(field_in:ti.template(), field_out:ti.template()):
imax = field_in.shape[0]
jmax = field_in.shape[1]
for i, j in ti.ndrange((1, imax-1), (1, jmax-1)):
field_out[i, j] = - (
(
field_in[i, j] * 4
-
field_in[i-1, j]
-
field_in[i+1, j]
-
field_in[i, j-1]
-
field_in[i, j+1]
) / (element_length ** 2)
)
@ti.kernel
def advection(u_star:ti.template(),
v_star:ti.template(),
u_prev:ti.template(),
v_prev:ti.template()
):
for i, j in u_star:
u_star[i, j] = (
u_prev[i, j]
+
TIME_STEP * (
-
(
u_prev[i, j] * dudx[i, j]
+
v_prev[i, j] * dudy[i, j]
)
+
NU * laplace_u[i, j]
)
)
for i, j in v_star:
v_star[i, j] = (
v_prev[i, j]
+
TIME_STEP * (
-
(
u_prev[i, j] * dvdx[i, j]
+
v_prev[i, j] * dvdy[i, j]
)
+
NU * laplace_v[i, j]
)
)
@ti.kernel
def enforce_velocity_bc(u:ti.template(), v:ti.template()):
imax = u.shape[0]
jmax = u.shape[1]
for j in range(jmax):
u[0, j] = 0.0
u[imax-1, j] = 0.0
v[0, j] = 0.0
v[imax-1, j] = 0.0
for i in range(imax):
u[i, 0] = 0.0
v[i, 0] = 0.0
u[i, jmax-1] = TOP_VELOCITY
v[i, jmax-1] = 0.0
@ti.kernel
def divergence(u:ti.template(), v:ti.template(), div:ti.template()):
imax = u.shape[0]
jmax = u.shape[1]
for i, j in ti.ndrange((1, imax-1), (1, jmax-1)):
dudx = (
(
u[i+1, j]
-
u[i-1, j]
) /
( 2 * element_length)
)
dvdy = (
(
v[i, j+1]
-
v[i, j-1]
) /
( 2 * element_length)
)
div[i, j] = (
dudx
+
dvdy
)
@ti.kernel
def pressure_iter(p:ti.template(), p_next:ti.template(), div:ti.template()):
imax = p.shape[0]
jmax = p.shape[1]
for i, j in ti.ndrange((1, imax-1), (1, jmax-1)):
p_next[i, j] = 0.25 * (
p[i-1, j]
+
p[i+1, j]
+
p[i, j-1]
+
p[i, j+1]
-
element_length**2
*
div[i, j]
*
DENSITY
/
TIME_STEP
)
for i in range(imax):
p_next[i, 0] = p_next[i, 1]
p_next[i, jmax-1] = 0.0
for j in range(jmax):
p_next[0, j] = p_next[1, j]
p_next[imax-1, j] = p_next[imax-2, j]
@ti.kernel
def update_velocity(vs:ti.template(),
vn:ti.template(),
dpds:ti.template()):
for i, j in vs:
vn[i, j] = (
vs[i, j]
-
TIME_STEP / DENSITY
*
dpds[i, j]
)
@ti.kernel
def assemble(u:ti.template(), v:ti.template(), vf:ti.template()):
for i, j in vf:
vf[i, j] = ti.Vector([u[i, j], v[i, j]])
gui = ti.GUI('Chorin Projection', (400, 400))
step = 0
while gui.running:
# Advection
central_difference_x(u_prev, dudx)
central_difference_y(u_prev, dudy)
central_difference_x(v_prev, dvdx)
central_difference_y(v_prev, dvdy)
laplacian(u_prev, laplace_u)
laplacian(v_prev, laplace_v)
advection(u_star, v_star, u_prev, v_prev)
enforce_velocity_bc(u_star, v_star)
# Projection
divergence(u_star, v_star, div_u)
for _ in range(MAX_ITER):
pressure_iter(p_prev, p_next, div_u)
p_prev, p_next = p_next, p_prev
central_difference_x(p_next, dpdx)
central_difference_y(p_next, dpdy)
update_velocity(u_star, u_next, dpdx)
update_velocity(v_star, v_next, dpdy)
enforce_velocity_bc(u_next, v_next)
# Advance in time
u_prev, u_next = u_next, u_prev
v_prev, v_next = v_next, v_prev
p_prev, p_next = p_next, p_prev
assemble(u_next, v_next, V)
gui.contour(p_next, normalize=True)
gui.vector_field(V, arrow_spacing=2, color=0xFFFFFF)
gui.show()
print(f'>>> Current time:{step * TIME_STEP:5.3e} sec.')
step += 1
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment