Skip to content

Commit

Permalink
Working!
Browse files Browse the repository at this point in the history
  • Loading branch information
@loliverhennigh
loliverhennigh committed Mar 29, 2024
1 parent 5c4732f commit 0387247
Show file tree
Hide file tree
Showing 4 changed files with 67 additions and 93 deletions.
72 changes: 16 additions & 56 deletions examples/interfaces/ldc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def run_ldc(backend, compute_mlup=True):
velocity_set = xlb.velocity_set.D3Q19()

# Make grid
nr = 256
nr = 128
shape = (nr, nr, nr)
if backend == "jax":
grid = xlb.grid.JaxGrid(shape=shape)
Expand Down Expand Up @@ -115,32 +115,17 @@ def run_ldc(backend, compute_mlup=True):
precision_policy=precision_policy,
compute_backend=compute_backend,
)
do_nothing_bc = xlb.operator.boundary_condition.DoNothingBC(
velocity_set=velocity_set,
precision_policy=precision_policy,
compute_backend=compute_backend,
)
half_way_bc = xlb.operator.boundary_condition.HalfwayBounceBackBC(
velocity_set=velocity_set,
precision_policy=precision_policy,
compute_backend=compute_backend,
)
full_way_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(
velocity_set=velocity_set,
precision_policy=precision_policy,
compute_backend=compute_backend,
)
stepper = xlb.operator.stepper.IncompressibleNavierStokesStepper(
collision=collision,
equilibrium=equilibrium,
macroscopic=macroscopic,
stream=stream,
boundary_conditions=[equilibrium_bc, do_nothing_bc, half_way_bc, full_way_bc],
)
indices_boundary_masker = xlb.operator.boundary_masker.IndicesBoundaryMasker(
velocity_set=velocity_set,
precision_policy=precision_policy,
compute_backend=compute_backend,
boundary_conditions=[equilibrium_bc, half_way_bc],
)
planar_boundary_masker = xlb.operator.boundary_masker.PlanarBoundaryMasker(
velocity_set=velocity_set,
Expand Down Expand Up @@ -170,8 +155,7 @@ def run_ldc(backend, compute_mlup=True):
lower_bound,
upper_bound,
direction,
#do_nothing_bc.id,
full_way_bc.id,
half_way_bc.id,
boundary_id,
missing_mask,
(0, 0, 0)
Expand All @@ -185,8 +169,7 @@ def run_ldc(backend, compute_mlup=True):
lower_bound,
upper_bound,
direction,
#half_way_bc.id,
full_way_bc.id,
half_way_bc.id,
boundary_id,
missing_mask,
(0, 0, 0)
Expand All @@ -200,8 +183,7 @@ def run_ldc(backend, compute_mlup=True):
lower_bound,
upper_bound,
direction,
#half_way_bc.id,
full_way_bc.id,
half_way_bc.id,
boundary_id,
missing_mask,
(0, 0, 0)
Expand All @@ -215,8 +197,7 @@ def run_ldc(backend, compute_mlup=True):
lower_bound,
upper_bound,
direction,
#half_way_bc.id,
full_way_bc.id,
half_way_bc.id,
boundary_id,
missing_mask,
(0, 0, 0)
Expand All @@ -230,34 +211,11 @@ def run_ldc(backend, compute_mlup=True):
lower_bound,
upper_bound,
direction,
#half_way_bc.id,
full_way_bc.id,
boundary_id,
missing_mask,
(0, 0, 0)
)

# Set full way bc (sphere)
"""
sphere_radius = nr // 8
x = np.arange(nr)
y = np.arange(nr)
z = np.arange(nr)
X, Y, Z = np.meshgrid(x, y, z)
indices = np.where(
(X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2
< sphere_radius**2
)
indices = np.array(indices).T
indices = wp.from_numpy(indices, dtype=wp.int32)
boundary_id, missing_mask = indices_boundary_masker(
indices,
full_way_bc.id,
half_way_bc.id,
boundary_id,
missing_mask,
(0, 0, 0)
)
"""

# Set initial conditions
if backend == "warp":
Expand All @@ -271,7 +229,7 @@ def run_ldc(backend, compute_mlup=True):
plot_freq = 512
save_dir = "ldc"
os.makedirs(save_dir, exist_ok=True)
num_steps = nr * 32
num_steps = nr * 512
start = time.time()

for _ in tqdm(range(num_steps)):
Expand All @@ -288,20 +246,22 @@ def run_ldc(backend, compute_mlup=True):
rho, u = macroscopic(f0, rho, u)
local_rho = rho.numpy()
local_u = u.numpy()
local_boundary_id = boundary_id.numpy()
elif backend == "jax":
local_rho, local_u = macroscopic(f0)
local_boundary_id = boundary_id

# Plot the velocity field, rho and boundary id side by side
plt.subplot(1, 3, 1)
plt.imshow(np.linalg.norm(u[:, :, nr // 2, :], axis=0))
plt.imshow(np.linalg.norm(local_u[:, :, nr // 2, :], axis=0))
plt.colorbar()
plt.subplot(1, 3, 2)
plt.imshow(rho[0, :, nr // 2, :])
plt.imshow(local_rho[0, :, nr // 2, :])
plt.colorbar()
plt.subplot(1, 3, 3)
plt.imshow(boundary_id[0, :, nr // 2, :])
plt.imshow(local_boundary_id[0, :, nr // 2, :])
plt.colorbar()
plt.savefig(f"{save_dir}/{str(_).zfill(6)}.png")
plt.savefig(f"{save_dir}/{backend}_{str(_).zfill(6)}.png")
plt.close()

wp.synchronize()
Expand All @@ -314,6 +274,6 @@ def run_ldc(backend, compute_mlup=True):

# Run the LDC example
backends = ["warp", "jax"]
#backends = ["jax"]
compute_mlup = False
for backend in backends:
run_ldc(backend, compute_mlup=True)
run_ldc(backend, compute_mlup=compute_mlup)
6 changes: 3 additions & 3 deletions examples/interfaces/taylor_green.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,7 @@ def run_taylor_green(backend, compute_mlup=True):
if __name__ == "__main__":

# Run Taylor-Green vortex on different backends
#backends = ["warp", "jax"]
backends = ["jax"]
backends = ["warp", "jax"]
#backends = ["jax"]
for backend in backends:
run_taylor_green(backend, compute_mlup=False)
run_taylor_green(backend, compute_mlup=True)
74 changes: 44 additions & 30 deletions xlb/operator/boundary_masker/planar_boundary_masker.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,41 +41,55 @@ def jax_implementation(
mask,
start_index=(0, 0, 0),
):
# Get plane dimensions
# TODO: Optimize this

# x plane
if direction[0] != 0:
dim = (
upper_bound[1] - lower_bound[1] + 1,
upper_bound[2] - lower_bound[2] + 1,
)

# Set boundary id
boundary_id = boundary_id.at[0, lower_bound[0], lower_bound[1] : upper_bound[1] + 1, lower_bound[2] : upper_bound[2] + 1].set(id_number)

# Set mask
for l in range(self.velocity_set.q):
d_dot_c = (
direction[0] * self.velocity_set.c[0, l]
+ direction[1] * self.velocity_set.c[1, l]
+ direction[2] * self.velocity_set.c[2, l]
)
if d_dot_c >= 0:
mask = mask.at[l, lower_bound[0], lower_bound[1] : upper_bound[1] + 1, lower_bound[2] : upper_bound[2] + 1].set(True)

# y plane
elif direction[1] != 0:
dim = (
upper_bound[0] - lower_bound[0] + 1,
upper_bound[2] - lower_bound[2] + 1,
)
elif direction[2] != 0:
dim = (
upper_bound[0] - lower_bound[0] + 1,
upper_bound[1] - lower_bound[1] + 1,
)

# Get the constants
_c = self.velocity_set.wp_c
_q = self.velocity_set.q
# Set boundary id
boundary_id = boundary_id.at[0, lower_bound[0] : upper_bound[0] + 1, lower_bound[1], lower_bound[2] : upper_bound[2] + 1].set(id_number)

# Get the mask
for i in range(dim[0]):
for j in range(dim[1]):
for k in range(_q):
d_dot_c = (
direction[0] * _c[0, k]
+ direction[1] * _c[1, k]
+ direction[2] * _c[2, k]
)
if d_dot_c >= 0:
mask[k, i, j] = True
# Set mask
for l in range(self.velocity_set.q):
d_dot_c = (
direction[0] * self.velocity_set.c[0, l]
+ direction[1] * self.velocity_set.c[1, l]
+ direction[2] * self.velocity_set.c[2, l]
)
if d_dot_c >= 0:
mask = mask.at[l, lower_bound[0] : upper_bound[0] + 1, lower_bound[1], lower_bound[2] : upper_bound[2] + 1].set(True)

# z plane
elif direction[2] != 0:

# Get the boundary id
boundary_id[:, :, :] = id_number
# Set boundary id
boundary_id = boundary_id.at[0, lower_bound[0] : upper_bound[0] + 1, lower_bound[1] : upper_bound[1] + 1, lower_bound[2]].set(id_number)

# Set mask
for l in range(self.velocity_set.q):
d_dot_c = (
direction[0] * self.velocity_set.c[0, l]
+ direction[1] * self.velocity_set.c[1, l]
+ direction[2] * self.velocity_set.c[2, l]
)
if d_dot_c >= 0:
mask = mask.at[l, lower_bound[0] : upper_bound[0] + 1, lower_bound[1] : upper_bound[1] + 1, lower_bound[2]].set(True)

return boundary_id, mask

Expand Down
8 changes: 4 additions & 4 deletions xlb/operator/stepper/nse.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,10 @@ def kernel(
_missing_mask[l] = wp.uint8(0)

# Apply streaming boundary conditions
if _boundary_id == _equilibrium_bc:
if (_boundary_id == wp.uint8(0)) or _boundary_id == _fullway_bounce_back_bc:
# Regular streaming
f_post_stream = self.stream.warp_functional(f_0, index)
elif _boundary_id == _equilibrium_bc:
# Equilibrium boundary condition
f_post_stream = self.equilibrium_bc.warp_functional(
f_0, _missing_mask, index
Expand All @@ -181,9 +184,6 @@ def kernel(
f_post_stream = self.halfway_bounce_back_bc.warp_functional(
f_0, _missing_mask, index
)
else:
# Regular streaming
f_post_stream = self.stream.warp_functional(f_0, index)

# Compute rho and u
rho, u = self.macroscopic.warp_functional(f_post_stream)
Expand Down

0 comments on commit 0387247

Please sign in to comment.