Merge pull request #85 from hsalehipour/main

fixed a few 2D bugs

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]
 
@@ -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.6
 
-    simulation = LidDrivenCavity2D(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(omega, prescribed_vel, grid_shape, velocity_set, backend, precision_policy)
+    simulation.run(num_steps=50000, post_process_interval=1000)

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)

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)

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

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])
@@ -162,9 +161,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(

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(

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

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):