From e04f6ebc68e139a34975f384791fbc3368e44b38 Mon Sep 17 00:00:00 2001 From: Hesam Salehipour Date: Fri, 25 Oct 2024 13:18:26 -0400 Subject: [PATCH 1/2] 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 From 39d62a794acadc055e3cc89d80b9a0872292f393 Mon Sep 17 00:00:00 2001 From: Hesam Salehipour Date: Tue, 29 Oct 2024 10:11:02 -0400 Subject: [PATCH 2/2] Addressing PR reveiew comments --- examples/cfd/lid_driven_cavity_2d.py | 19 +++++++++++++------ .../cfd/lid_driven_cavity_2d_distributed.py | 16 +++++++++++----- .../bc_extrapolation_outflow.py | 12 ++++++++---- .../boundary_condition/bc_regularized.py | 1 - xlb/operator/boundary_condition/bc_zouhe.py | 11 ++++++++--- xlb/velocity_set/velocity_set.py | 1 - 6 files changed, 40 insertions(+), 20 deletions(-) diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py index bcc60b8..f8bd65a 100644 --- a/examples/cfd/lid_driven_cavity_2d.py +++ b/examples/cfd/lid_driven_cavity_2d.py @@ -14,7 +14,7 @@ class LidDrivenCavity2D: - def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy): + def __init__(self, omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy): # initialize backend xlb.init( velocity_set=velocity_set, @@ -29,6 +29,7 @@ def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy): self.grid, self.f_0, self.f_1, self.missing_mask, self.bc_mask = create_nse_fields(grid_shape) self.stepper = None self.boundary_conditions = [] + self.prescribed_vel = prescribed_vel # Setup the simulation BC, its initial conditions, and the stepper self._setup(omega) @@ -49,7 +50,7 @@ def define_boundary_indices(self): def setup_boundary_conditions(self): lid, walls = self.define_boundary_indices() - bc_top = EquilibriumBC(rho=1.0, u=(0.02, 0.0), indices=lid) + bc_top = EquilibriumBC(rho=1.0, u=(self.prescribed_vel, 0.0), indices=lid) bc_walls = HalfwayBounceBackBC(indices=walls) self.boundary_conditions = [bc_walls, bc_top] @@ -67,7 +68,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, collision_type='KBC') + self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions) def run(self, num_steps, post_process_interval=100): for i in range(num_steps): @@ -100,7 +101,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 +113,13 @@ def post_process(self, i): precision_policy = PrecisionPolicy.FP32FP32 velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend) - omega = 1.99 - simulation = LidDrivenCavity2D(omega, grid_shape, velocity_set, backend, precision_policy) + # Setting fluid viscosity and relaxation parameter. + Re = 200.0 + prescribed_vel = 0.05 + clength = grid_shape[0] - 1 + visc = prescribed_vel * clength / Re + omega = 1.0 / (3.0 * visc + 0.5) + + simulation = LidDrivenCavity2D(omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy) simulation.run(num_steps=50000, post_process_interval=1000) diff --git a/examples/cfd/lid_driven_cavity_2d_distributed.py b/examples/cfd/lid_driven_cavity_2d_distributed.py index a7d43b4..cdc9027 100644 --- a/examples/cfd/lid_driven_cavity_2d_distributed.py +++ b/examples/cfd/lid_driven_cavity_2d_distributed.py @@ -7,8 +7,8 @@ class LidDrivenCavity2D_distributed(LidDrivenCavity2D): - def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy): - super().__init__(omega, grid_shape, velocity_set, backend, precision_policy) + def __init__(self, omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy): + super().__init__(omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy) def setup_stepper(self, omega): stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions) @@ -29,7 +29,13 @@ def setup_stepper(self, omega): precision_policy = PrecisionPolicy.FP32FP32 velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend) - omega = 1.6 - simulation = LidDrivenCavity2D_distributed(omega, grid_shape, velocity_set, backend, precision_policy) - simulation.run(num_steps=5000, post_process_interval=1000) + # Setting fluid viscosity and relaxation parameter. + Re = 200.0 + prescribed_vel = 0.05 + clength = grid_shape[0] - 1 + visc = prescribed_vel * clength / Re + omega = 1.0 / (3.0 * visc + 0.5) + + simulation = LidDrivenCavity2D_distributed(omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy) + simulation.run(num_steps=50000, post_process_interval=1000) diff --git a/xlb/operator/boundary_condition/bc_extrapolation_outflow.py b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py index f47fe2f..8a2a482 100644 --- a/xlb/operator/boundary_condition/bc_extrapolation_outflow.py +++ b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py @@ -134,7 +134,6 @@ def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask): def _construct_warp(self): # Set local constants sound_speed = self.compute_dtype(1.0 / wp.sqrt(3.0)) - _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) _c = self.velocity_set.c _q = self.velocity_set.q _opp_indices = self.velocity_set.opp_indices @@ -143,9 +142,14 @@ def _construct_warp(self): 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 -wp.vec3i(_c[0, l], _c[1, l], _c[2, l]) + if wp.static(self.velocity_set.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 -wp.vec3i(_c[0, l], _c[1, l], _c[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 -wp.vec2i(_c[0, l], _c[1, l]) # Construct the functionals for this BC @wp.func diff --git a/xlb/operator/boundary_condition/bc_regularized.py b/xlb/operator/boundary_condition/bc_regularized.py index 5fd71f6..d1abd4e 100644 --- a/xlb/operator/boundary_condition/bc_regularized.py +++ b/xlb/operator/boundary_condition/bc_regularized.py @@ -133,7 +133,6 @@ def _construct_warp(self): # Set local constants TODO: This is a hack and should be fixed with warp update # _u_vec = wp.vec(_d, dtype=self.compute_dtype) # compute Qi tensor and store it in self - _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype) _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype) _rho = self.compute_dtype(rho) _u = _u_vec(u[0], u[1], u[2]) if _d == 3 else _u_vec(u[0], u[1]) diff --git a/xlb/operator/boundary_condition/bc_zouhe.py b/xlb/operator/boundary_condition/bc_zouhe.py index a92d909..3e89992 100644 --- a/xlb/operator/boundary_condition/bc_zouhe.py +++ b/xlb/operator/boundary_condition/bc_zouhe.py @@ -207,9 +207,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/velocity_set/velocity_set.py b/xlb/velocity_set/velocity_set.py index 41db991..33b2331 100644 --- a/xlb/velocity_set/velocity_set.py +++ b/xlb/velocity_set/velocity_set.py @@ -94,7 +94,6 @@ def _init_jax_properties(self): self.w = jnp.array(self._w, dtype=dtype) self.opp_indices = jnp.array(self._opp_indices, dtype=jnp.int32) self.cc = jnp.array(self._cc, dtype=dtype) - self.c_float = jnp.array(self._c_float, dtype=dtype) self.qi = jnp.array(self._qi, dtype=dtype) def _init_backend_constants(self):