diff --git a/examples/interfaces/ldc.py b/examples/interfaces/ldc.py index ced31fe..d75ce8a 100644 --- a/examples/interfaces/ldc.py +++ b/examples/interfaces/ldc.py @@ -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) @@ -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, @@ -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) @@ -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) @@ -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) @@ -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) @@ -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": @@ -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)): @@ -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() @@ -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) diff --git a/examples/interfaces/taylor_green.py b/examples/interfaces/taylor_green.py index 4d97476..f842107 100644 --- a/examples/interfaces/taylor_green.py +++ b/examples/interfaces/taylor_green.py @@ -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) diff --git a/xlb/operator/boundary_masker/planar_boundary_masker.py b/xlb/operator/boundary_masker/planar_boundary_masker.py index 4cd405c..832a8b2 100644 --- a/xlb/operator/boundary_masker/planar_boundary_masker.py +++ b/xlb/operator/boundary_masker/planar_boundary_masker.py @@ -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 diff --git a/xlb/operator/stepper/nse.py b/xlb/operator/stepper/nse.py index 63c32ee..9c6d56e 100644 --- a/xlb/operator/stepper/nse.py +++ b/xlb/operator/stepper/nse.py @@ -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 @@ -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)