From e04f6ebc68e139a34975f384791fbc3368e44b38 Mon Sep 17 00:00:00 2001 From: Hesam Salehipour Date: Fri, 25 Oct 2024 13:18:26 -0400 Subject: [PATCH] Fixed 2d bugs in Warp backeneds in pytest, KBC and bc_regularized --- examples/cfd/lid_driven_cavity_2d.py | 8 ++++---- .../mask/test_bc_indices_masker_warp.py | 10 ++++++---- .../boundary_condition/bc_regularized.py | 11 +++++++--- xlb/operator/collision/kbc.py | 20 +++++++++++-------- 4 files changed, 30 insertions(+), 19 deletions(-) diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py index d790c2c..bcc60b8 100644 --- a/examples/cfd/lid_driven_cavity_2d.py +++ b/examples/cfd/lid_driven_cavity_2d.py @@ -67,7 +67,7 @@ def initialize_fields(self): self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.precision_policy, self.backend) def setup_stepper(self, omega): - self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions) + self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions, collision_type='KBC') def run(self, num_steps, post_process_interval=100): for i in range(num_steps): @@ -100,7 +100,7 @@ def post_process(self, i): fields = {"rho": rho[0], "u_x": u[0], "u_y": u[1], "u_magnitude": u_magnitude} - save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity") + # save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity") save_image(fields["u_magnitude"], timestep=i, prefix="lid_driven_cavity") @@ -112,7 +112,7 @@ def post_process(self, i): precision_policy = PrecisionPolicy.FP32FP32 velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend) - omega = 1.6 + omega = 1.99 simulation = LidDrivenCavity2D(omega, grid_shape, velocity_set, backend, precision_policy) - simulation.run(num_steps=5000, post_process_interval=1000) + simulation.run(num_steps=50000, post_process_interval=1000) diff --git a/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py b/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py index f56c9ce..cebc23f 100644 --- a/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py +++ b/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py @@ -60,7 +60,6 @@ def test_indices_masker_warp(dim, velocity_set, grid_shape): [test_bc], bc_mask, missing_mask, - start_index=(0, 0, 0) if dim == 3 else (0, 0), ) assert missing_mask.dtype == xlb.Precision.BOOL.wp_dtype @@ -69,9 +68,12 @@ def test_indices_masker_warp(dim, velocity_set, grid_shape): bc_mask = bc_mask.numpy() missing_mask = missing_mask.numpy() - assert bc_mask.shape == (1,) + grid_shape - - assert missing_mask.shape == (velocity_set.q,) + grid_shape + if len(grid_shape) == 2: + assert bc_mask.shape == (1,) + grid_shape + (1,), "bc_mask shape is incorrect got {}".format(bc_mask.shape) + assert missing_mask.shape == (velocity_set.q,) + grid_shape + (1,), "missing_mask shape is incorrect got {}".format(missing_mask.shape) + else: + assert bc_mask.shape == (1,) + grid_shape, "bc_mask shape is incorrect got {}".format(bc_mask.shape) + assert missing_mask.shape == (velocity_set.q,) + grid_shape, "missing_mask shape is incorrect got {}".format(missing_mask.shape) if dim == 2: assert np.all(bc_mask[0, indices[0], indices[1]] == test_bc.id) diff --git a/xlb/operator/boundary_condition/bc_regularized.py b/xlb/operator/boundary_condition/bc_regularized.py index af4c783..5fd71f6 100644 --- a/xlb/operator/boundary_condition/bc_regularized.py +++ b/xlb/operator/boundary_condition/bc_regularized.py @@ -162,9 +162,14 @@ def _get_fsum( def get_normal_vectors( missing_mask: Any, ): - for l in range(_q): - if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1: - return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l]) + if wp.static(_d == 3): + for l in range(_q): + if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1: + return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l]) + else: + for l in range(_q): + if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) == 1: + return -_u_vec(_c_float[0, l], _c_float[1, l]) @wp.func def bounceback_nonequilibrium( diff --git a/xlb/operator/collision/kbc.py b/xlb/operator/collision/kbc.py index 5c5c29e..b841baa 100644 --- a/xlb/operator/collision/kbc.py +++ b/xlb/operator/collision/kbc.py @@ -67,7 +67,7 @@ def jax_implementation( fneq = f - feq if isinstance(self.velocity_set, D2Q9): shear = self.decompose_shear_d2q9_jax(fneq) - delta_s = shear * rho / 4.0 # TODO: Check this + delta_s = shear * rho / 4.0 elif isinstance(self.velocity_set, D3Q27): shear = self.decompose_shear_d3q27_jax(fneq) delta_s = shear * rho @@ -191,16 +191,16 @@ def _construct_warp(self): @wp.func def decompose_shear_d2q9(fneq: Any): pi = self.momentum_flux.warp_functional(fneq) - N = pi[0] - pi[1] + N = pi[0] - pi[2] s = _f_vec() s[3] = N s[6] = N s[2] = -N s[1] = -N - s[8] = pi[2] - s[4] = -pi[2] - s[5] = -pi[2] - s[7] = pi[2] + s[8] = pi[1] + s[4] = -pi[1] + s[5] = -pi[1] + s[7] = pi[1] return s # Construct functional for decomposing shear @@ -271,8 +271,12 @@ def functional( ): # Compute shear and delta_s fneq = f - feq - shear = decompose_shear_d3q27(fneq) - delta_s = shear * rho # TODO: Check this + if wp.static(self.velocity_set.d == 3): + shear = decompose_shear_d3q27(fneq) + delta_s = shear * rho + else: + shear = decompose_shear_d2q9(fneq) + delta_s = shear * rho / self.compute_dtype(4.0) # Perform collision delta_h = fneq - delta_s