Skip to content

Commit

Permalink
Merge pull request #16 from hsalehipour/main
Browse files Browse the repository at this point in the history 
modified two benchmark examples and fixed a bug in moving bounce-back
  • Loading branch information
@mehdiataei
mehdiataei authored Sep 20, 2023
2 parents 4ba5dc3 + 21268fb commit 3cd7fd0
Show file tree
Hide file tree
Showing 4 changed files with 91 additions and 62 deletions.
1 change: 1 addition & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
*zraw
*mhd
*png
*pdf
# Exclusions
!jax-lbm.png
!airfoil.png
Expand Down
23 changes: 14 additions & 9 deletions examples/CFD/channel3d.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ def __init__(self, **kwargs):
def set_boundary_conditions(self):
# top and bottom sides of the channel are no-slip and the other directions are periodic
wall = np.concatenate((self.boundingBoxIndices['bottom'], self.boundingBoxIndices['top']))
self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
self.BCs.append(Regularized(tuple(wall.T), self.gridInfo, self.precisionPolicy, 'velocity', np.zeros((wall.shape[0], 3))))
return

def initialize_macroscopic_fields(self):
Expand Down Expand Up @@ -107,8 +107,11 @@ def output_data(self, **kwargs):
dns_dic = get_dns_data()
plt.clf()
plt.semilogx(yplus, uplus,'r.', yplus, uplus_loglaw, 'k:', dns_dic['y+'], dns_dic['Umean'], 'b-')
fname = "uplus_" + str(timestep).zfill(4) + '.png'
plt.savefig(fname)
ax = plt.gca()
ax.set_xlim([0.1, 300])
ax.set_ylim([0, 20])
fname = "uplus_" + str(timestep//10000).zfill(5) + '.pdf'
plt.savefig(fname, format='pdf')
fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]}
save_fields_vtk(timestep, fields)

Expand All @@ -117,17 +120,19 @@ def output_data(self, **kwargs):
if __name__ == "__main__":
precision = "f64/f64"
lattice = LatticeD3Q27(precision)

# h: channel half-width
h = 10
h = 50

# Define channel geometry based on h
nx = 6*h
ny = 3*h
nz = 2*h

# Define flow regime
Re_tau = 180
u_tau = 0.004
# DeltaPlus = Re_tau/h # DeltaPlus = u_tau / nu * Delta where u_tau / nu = Re_tau/h
u_tau = 0.001
DeltaPlus = Re_tau/h # DeltaPlus = u_tau / nu * Delta where u_tau / nu = Re_tau/h
visc = u_tau * h / Re_tau
omega = 1.0 / (3.0 * visc + 0.5)

Expand All @@ -146,8 +151,8 @@ def output_data(self, **kwargs):
'ny': ny,
'nz': nz,
'precision': precision,
'io_rate': 20000,
'print_info_rate': 20000
'io_rate': 500000,
'print_info_rate': 100000
}
sim = turbulentChannel(**kwargs)
sim.run(4000000)
sim.run(10000000)
109 changes: 62 additions & 47 deletions examples/CFD/taylor_green_vortex.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,16 +6,13 @@


from src.boundary_conditions import *
from jax.config import config
from src.utils import *
import numpy as np
from src.lattice import LatticeD2Q9
from src.models import BGKSim, KBCSim, AdvectionDiffusionBGK
import jax.numpy as jnp
from jax.experimental import mesh_utils
from jax.sharding import PositionalSharding
import os
import matplotlib.pyplot as plt
import json

# Use 8 CPU devices
# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
Expand All @@ -25,8 +22,8 @@
jax.config.update('jax_enable_x64', True)

def taylor_green_initial_fields(xx, yy, u0, rho0, nu, time):
ux = -u0 * np.cos(xx) * np.sin(yy) * np.exp(-2 * nu * time)
uy = u0 * np.sin(xx) * np.cos(yy) * np.exp(-2 * nu * time)
ux = u0 * np.sin(xx) * np.cos(yy) * np.exp(-2 * nu * time)
uy = -u0 * np.cos(xx) * np.sin(yy) * np.exp(-2 * nu * time)
rho = 1.0 - rho0 * u0 ** 2 / 12. * (np.cos(2. * xx) + np.cos(2. * yy)) * np.exp(-4 * nu * time)
return ux, uy, np.expand_dims(rho, axis=-1)

Expand All @@ -51,7 +48,7 @@ def initialize_populations(self, rho, u):
ADE = AdvectionDiffusionBGK(**kwargs)
ADE.initialize_macroscopic_fields = self.initialize_macroscopic_fields
print("Initializing the distribution functions using the specified macroscopic fields....")
f = ADE.run(20000)
f = ADE.run(int(20000*32/nx))
return f

def output_data(self, **kwargs):
Expand All @@ -64,49 +61,67 @@ def output_data(self, **kwargs):
time = timestep * (kx**2 + ky**2)/2.
ux_th, uy_th, rho_th = taylor_green_initial_fields(xx, yy, vel_ref, 1, visc, time)
vel_err_L2 = np.sqrt(np.sum((u[..., 0]-ux_th)**2 + (u[..., 1]-uy_th)**2) / np.sum(ux_th**2 + uy_th**2))
print("error= {:07.6f}".format(vel_err_L2))
rho_err_L2 = np.sqrt(np.sum((rho - rho_th)**2) / np.sum(rho_th**2))
print("Vel error= {:07.6f}, Pressure error= {:07.6f}".format(vel_err_L2, rho_err_L2))
if timestep == endTime:
ErrL2ResList.append(vel_err_L2)
ErrL2ResListRho.append(rho_err_L2)
# save_image(timestep, u)


if __name__ == "__main__":
precision = "f64/f64"
lattice = LatticeD2Q9(precision)

resList = [32, 64, 128, 256, 512]
ErrL2ResList = []

for nx in resList:
print("Running at nx = ny = {:07.6f}".format(nx))
ny = nx
twopi = 2.0 * np.pi
coord = np.array([(i, j) for i in range(nx) for j in range(ny)])
xx, yy = coord[:, 0], coord[:, 1]
kx, ky = twopi / nx, twopi / ny
xx = xx.reshape((nx, ny)) * kx
yy = yy.reshape((nx, ny)) * ky

Re = 1000.0
vel_ref = 0.04*32/nx

visc = vel_ref * nx / Re
omega = 1.0 / (3.0 * visc + 0.5)
print("omega = ", omega)
os.system("rm -rf ./*.vtk && rm -rf ./*.png")
kwargs = {
'lattice': lattice,
'omega': omega,
'nx': nx,
'ny': ny,
'nz': 0,
'precision': precision,
'io_rate': 500,
'print_info_rate': 500
}
sim = TaylorGreenVortex(**kwargs)
endTime = int(20000*nx/32.0)
sim.run(endTime)
plt.loglog(resList, ErrL2ResList, '-o')
plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--')
plt.savefig('Error.png'); plt.savefig('Error.pdf', format='pdf')
precision_list = ["f32/f32", "f64/f32", "f64/f64"]
resList = [32, 64, 128, 256, 512, 1024]
result_dict = dict.fromkeys(precision_list)
result_dict['resolution_list'] = resList

for precision in precision_list:
lattice = LatticeD2Q9(precision)
ErrL2ResList = []
ErrL2ResListRho = []
result_dict[precision] = dict.fromkeys(['vel_error', 'rho_error'])
for nx in resList:
print("Running at nx = ny = {:07.6f}".format(nx))
ny = nx
twopi = 2.0 * np.pi
coord = np.array([(i, j) for i in range(nx) for j in range(ny)])
xx, yy = coord[:, 0], coord[:, 1]
kx, ky = twopi / nx, twopi / ny
xx = xx.reshape((nx, ny)) * kx
yy = yy.reshape((nx, ny)) * ky

Re = 1600.0
vel_ref = 0.04*32/nx

visc = vel_ref * nx / Re
omega = 1.0 / (3.0 * visc + 0.5)
print("omega = ", omega)
os.system("rm -rf ./*.vtk && rm -rf ./*.png")
kwargs = {
'lattice': lattice,
'omega': omega,
'nx': nx,
'ny': ny,
'nz': 0,
'precision': precision,
'io_rate': 5000,
'print_info_rate': 1000
}
sim = TaylorGreenVortex(**kwargs)
tc = 2.0/(2. * visc * (kx**2 + ky**2))
endTime = int(0.05*tc)
sim.run(endTime)
result_dict[precision]['vel_error'] = ErrL2ResList
result_dict[precision]['rho_error'] = ErrL2ResListRho

with open('data.json', 'w') as fp:
json.dump(result_dict, fp)

# plt.loglog(resList, ErrL2ResList, '-o')
# plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--')
# plt.savefig('ErrorVel.png'); plt.savefig('ErrorVel.pdf', format='pdf')

# plt.figure()
# plt.loglog(resList, ErrL2ResListRho, '-o')
# plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--')
# plt.savefig('ErrorRho.png'); plt.savefig('ErrorRho.pdf', format='pdf')
20 changes: 14 additions & 6 deletions src/boundary_conditions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -445,8 +445,8 @@ def apply(self, fout, fin, time):
"""
indices, vel = self.update_function(time)
c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
cu = 3.0 * jnp.dot(vel, c)
return fout.at[indices].set(fin[indices][..., self.lattice.opp_indices] - 6.0 * self.lattice.w * cu)
cu = 6.0 * self.lattice.w * jnp.dot(vel, c)
return fout.at[indices].set(fin[indices][..., self.lattice.opp_indices] - cu)


class BounceBackHalfway(BoundaryCondition):
Expand All @@ -467,13 +467,16 @@ class BounceBackHalfway(BoundaryCondition):
Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True.
isSolid : bool
Whether the boundary condition represents a solid boundary. For this class, it is True.
vel : array-like
The prescribed value of velocity vector for the boundary condition. No-slip BC is assumed if vel=None (default).
"""
def __init__(self, indices, gridInfo, precision_policy):
def __init__(self, indices, gridInfo, precision_policy, vel=None):
super().__init__(indices, gridInfo, precision_policy)
self.name = "BounceBackHalfway"
self.implementationStep = "PostStreaming"
self.needsExtraConfiguration = True
self.isSolid = True
self.vel = vel

def configure(self, boundaryBitmask):
"""
Expand Down Expand Up @@ -523,7 +526,12 @@ def apply(self, fout, fin):
nbd = len(self.indices[0])
bindex = np.arange(nbd)[:, None]
fbd = fout[self.indices]
fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
if self.vel is not None:
c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
cu = 6.0 * self.lattice.w * jnp.dot(self.vel, c)
fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown] - cu[bindex, self.iknown])
else:
fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])

return fbd

Expand Down Expand Up @@ -926,8 +934,8 @@ def configure(self, boundaryBitmask):
Parameters
----------
connectivity_bitmask : np.ndarray
The connectivity bitmask for the lattice.
boundaryBitmask : np.ndarray
The connectivity bitmask for the boundary voxels.
"""
shiftDir = ~boundaryBitmask[:, self.lattice.opp_indices]
idx = np.array(self.indices).T
Expand Down

0 comments on commit 3cd7fd0

Please sign in to comment.