From e16ba2445007a5e860097f132d91fdfc00c5a1e6 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Thu, 17 Oct 2024 22:15:39 -0400
Subject: [PATCH 1/5] fixed race conditioning in indices_boundary_masker due to
 duplicate bc indices at corners and edges.

---
 examples/cfd/flow_past_sphere_3d.py           | 24 ++++++++--
 examples/cfd/lid_driven_cavity_2d.py          |  2 +-
 examples/cfd/windtunnel_3d.py                 |  4 +-
 .../bc_grads_approximation.py                 |  1 -
 .../boundary_condition/boundary_condition.py  |  1 +
 .../indices_boundary_masker.py                | 46 ++++++++-----------
 xlb/utils/utils.py                            |  8 ++--
 7 files changed, 49 insertions(+), 37 deletions(-)

diff --git a/examples/cfd/flow_past_sphere_3d.py b/examples/cfd/flow_past_sphere_3d.py
index 2e0df95a..994d18b0 100644
--- a/examples/cfd/flow_past_sphere_3d.py
+++ b/examples/cfd/flow_past_sphere_3d.py
@@ -80,10 +80,9 @@ def setup_boundary_conditions(self):
         bc_outlet = ExtrapolationOutflowBC(indices=outlet)
         bc_sphere = HalfwayBounceBackBC(indices=sphere)
 
-        self.boundary_conditions = [bc_left, bc_outlet, bc_sphere, bc_walls]
-        # Note: it is important to add bc_walls to be after bc_outlet/bc_inlet because
+        self.boundary_conditions = [bc_walls, bc_left, bc_outlet, bc_sphere]
+        # Note: it is important to add bc_walls before bc_outlet/bc_inlet because
         # of the corner nodes. This way the corners are treated as wall and not inlet/outlet.
-        # TODO: how to ensure about this behind in the src code?
 
     def setup_boundary_masker(self):
         indices_boundary_masker = IndicesBoundaryMasker(
@@ -105,6 +104,8 @@ def run(self, num_steps, post_process_interval=100):
             self.f_0, self.f_1 = self.stepper(self.f_0, self.f_1, self.bc_mask, self.missing_mask, i)
             self.f_0, self.f_1 = self.f_1, self.f_0
 
+            if i == 0:
+                self.check_boundary_mask()
             if i % post_process_interval == 0 or i == num_steps - 1:
                 self.post_process(i)
                 end_time = time.time()
@@ -134,6 +135,23 @@ def post_process(self, i):
 
         # save_fields_vtk(fields, timestep=i)
         save_image(fields["u_magnitude"][:, self.grid_shape[1] // 2, :], timestep=i)
+        return
+
+    def check_boundary_mask(self):
+        # Write the results. We'll use JAX backend for the post-processing
+        if not isinstance(self.f_0, jnp.ndarray):
+            bmask = wp.to_jax(self.bc_mask)[0]
+        else:
+            bmask = self.bc_mask[0]
+
+        # save_fields_vtk(fields, timestep=i)
+        save_image(bmask[0, :, :], prefix="00_left")
+        save_image(bmask[self.grid_shape[0] - 1, :, :], prefix="00_right")
+        save_image(bmask[:, :, self.grid_shape[2] - 1], prefix="00_top")
+        save_image(bmask[:, :, 0], prefix="00_bottom")
+        save_image(bmask[:, 0, :], prefix="00_front")
+        save_image(bmask[:, self.grid_shape[1] - 1, :], prefix="00_back")
+        save_image(bmask[:, self.grid_shape[1] // 2, :], prefix="00_middle")
 
 
 if __name__ == "__main__":
diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py
index 20f3b7c4..300614da 100644
--- a/examples/cfd/lid_driven_cavity_2d.py
+++ b/examples/cfd/lid_driven_cavity_2d.py
@@ -50,7 +50,7 @@ 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_walls = HalfwayBounceBackBC(indices=walls)
-        self.boundary_conditions = [bc_top, bc_walls]
+        self.boundary_conditions = [bc_walls, bc_top]
 
     def setup_boundary_masker(self):
         indices_boundary_masker = IndicesBoundaryMasker(
diff --git a/examples/cfd/windtunnel_3d.py b/examples/cfd/windtunnel_3d.py
index 140e756e..b0ee5b9d 100644
--- a/examples/cfd/windtunnel_3d.py
+++ b/examples/cfd/windtunnel_3d.py
@@ -104,7 +104,9 @@ def setup_boundary_conditions(self):
         # bc_car = HalfwayBounceBackBC(mesh_vertices=car)
         bc_car = GradsApproximationBC(mesh_vertices=car)
         # bc_car = FullwayBounceBackBC(mesh_vertices=car)
-        self.boundary_conditions = [bc_left, bc_do_nothing, bc_walls, bc_car]
+        self.boundary_conditions = [bc_walls, bc_left, bc_do_nothing, bc_car]
+        # Note: it is important to add bc_walls before bc_outlet/bc_inlet because
+        # of the corner nodes. This way the corners are treated as wall and not inlet/outlet.
 
     def setup_boundary_masker(self):
         indices_boundary_masker = IndicesBoundaryMasker(
diff --git a/xlb/operator/boundary_condition/bc_grads_approximation.py b/xlb/operator/boundary_condition/bc_grads_approximation.py
index 3d608798..bc298515 100644
--- a/xlb/operator/boundary_condition/bc_grads_approximation.py
+++ b/xlb/operator/boundary_condition/bc_grads_approximation.py
@@ -47,7 +47,6 @@ def __init__(
         indices=None,
         mesh_vertices=None,
     ):
-
         # TODO: the input velocity must be suitably stored elesewhere when mesh is moving.
         self.u = (0, 0, 0)
 
diff --git a/xlb/operator/boundary_condition/boundary_condition.py b/xlb/operator/boundary_condition/boundary_condition.py
index 6d72fc03..008cc9aa 100644
--- a/xlb/operator/boundary_condition/boundary_condition.py
+++ b/xlb/operator/boundary_condition/boundary_condition.py
@@ -15,6 +15,7 @@
 from xlb import DefaultConfig
 from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry
 
+
 # Enum for implementation step
 class ImplementationStep(Enum):
     COLLISION = auto()
diff --git a/xlb/operator/boundary_masker/indices_boundary_masker.py b/xlb/operator/boundary_masker/indices_boundary_masker.py
index fdc43310..10d72c43 100644
--- a/xlb/operator/boundary_masker/indices_boundary_masker.py
+++ b/xlb/operator/boundary_masker/indices_boundary_masker.py
@@ -66,7 +66,7 @@ def jax_implementation(self, bclist, bc_mask, missing_mask, start_index=None):
             start_index = (0,) * dim
 
         domain_shape = bc_mask[0].shape
-        for bc in bclist:
+        for bc in reversed(bclist):
             assert bc.indices is not None, f"Please specify indices associated with the {bc.__class__.__name__} BC!"
             assert bc.mesh_vertices is None, f"Please use MeshBoundaryMasker operator if {bc.__class__.__name__} is imposed on a mesh (e.g. STL)!"
             id_number = bc.id
@@ -103,6 +103,11 @@ def _construct_warp(self):
         _c = self.velocity_set.c
         _q = wp.constant(self.velocity_set.q)
 
+        @wp.func
+        def check_index_bounds(index: wp.vec3i, shape: wp.vec3i):
+            is_in_bounds = index[0] >= 0 and index[0] < shape[0] and index[1] >= 0 and index[1] < shape[1] and index[2] >= 0 and index[2] < shape[2]
+            return is_in_bounds
+
         # Construct the warp 2D kernel
         @wp.kernel
         def kernel2d(
@@ -173,14 +178,8 @@ def kernel3d(
             index[2] = indices[2, ii] - start_index[2]
 
             # Check if index is in bounds
-            if (
-                index[0] >= 0
-                and index[0] < missing_mask.shape[1]
-                and index[1] >= 0
-                and index[1] < missing_mask.shape[2]
-                and index[2] >= 0
-                and index[2] < missing_mask.shape[3]
-            ):
+            shape = wp.vec3i(missing_mask.shape[1], missing_mask.shape[2], missing_mask.shape[3])
+            if check_index_bounds(index, shape):
                 # Stream indices
                 for l in range(_q):
                     # Get the index of the streaming direction
@@ -195,27 +194,12 @@ def kernel3d(
 
                     # check if pull index is out of bound
                     # These directions will have missing information after streaming
-                    if (
-                        pull_index[0] < 0
-                        or pull_index[0] >= missing_mask.shape[1]
-                        or pull_index[1] < 0
-                        or pull_index[1] >= missing_mask.shape[2]
-                        or pull_index[2] < 0
-                        or pull_index[2] >= missing_mask.shape[3]
-                    ):
+                    if not check_index_bounds(pull_index, shape):
                         # Set the missing mask
                         missing_mask[l, index[0], index[1], index[2]] = True
 
                     # handling geometries in the interior of the computational domain
-                    elif (
-                        is_interior[ii]
-                        and push_index[0] >= 0
-                        and push_index[0] < missing_mask.shape[1]
-                        and push_index[1] >= 0
-                        and push_index[1] < missing_mask.shape[2]
-                        and push_index[2] >= 0
-                        and push_index[2] < missing_mask.shape[3]
-                    ):
+                    elif check_index_bounds(pull_index, shape) and is_interior[ii]:
                         # Set the missing mask
                         missing_mask[l, push_index[0], push_index[1], push_index[2]] = True
                         bc_mask[0, push_index[0], push_index[1], push_index[2]] = id_number[ii]
@@ -241,8 +225,14 @@ def warp_implementation(self, bclist, bc_mask, missing_mask, start_index=None):
             # We are done with bc.indices. Remove them from BC objects
             bc.__dict__.pop("indices", None)
 
-        indices = wp.array2d(index_list, dtype=wp.int32)
-        id_number = wp.array1d(id_list, dtype=wp.uint8)
+        # Remove duplicates indices to avoid race conditioning
+        index_arr, unique_loc = np.unique(index_list, axis=-1, return_index=True)
+        id_arr = np.array(id_list)[unique_loc]
+        is_interior = np.array(is_interior)[unique_loc]
+
+        # convert to warp arrays
+        indices = wp.array2d(index_arr, dtype=wp.int32)
+        id_number = wp.array1d(id_arr, dtype=wp.uint8)
         is_interior = wp.array1d(is_interior, dtype=wp.bool)
 
         if start_index is None:
diff --git a/xlb/utils/utils.py b/xlb/utils/utils.py
index 074177e0..0a9858a5 100644
--- a/xlb/utils/utils.py
+++ b/xlb/utils/utils.py
@@ -44,7 +44,7 @@ def downsample_field(field, factor, method="bicubic"):
         return jnp.stack(downsampled_components, axis=-1)
 
 
-def save_image(fld, timestep, prefix=None):
+def save_image(fld, timestep=None, prefix=None, **kwargs):
     """
     Save an image of a field at a given timestep.
 
@@ -74,7 +74,8 @@ def save_image(fld, timestep, prefix=None):
     else:
         fname = prefix
 
-    fname = fname + "_" + str(timestep).zfill(4)
+    if timestep is not None:
+        fname = fname + "_" + str(timestep).zfill(4)
 
     if len(fld.shape) > 3:
         raise ValueError("The input field should be 2D!")
@@ -82,7 +83,8 @@ def save_image(fld, timestep, prefix=None):
         fld = np.sqrt(fld[0, ...] ** 2 + fld[0, ...] ** 2)
 
     plt.clf()
-    plt.imsave(fname + ".png", fld.T, cmap=cm.nipy_spectral, origin="lower")
+    kwargs.pop("cmap", None)
+    plt.imsave(fname + ".png", fld.T, cmap=cm.nipy_spectral, origin="lower", **kwargs)
 
 
 def save_fields_vtk(fields, timestep, output_dir=".", prefix="fields"):

From bd62072fb03cc1d7ad8edd8aa525aefa07ffd4ea Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Fri, 18 Oct 2024 09:36:18 -0400
Subject: [PATCH 2/5] fixed a bug in 2d kbc warp

---
 xlb/operator/collision/kbc.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/xlb/operator/collision/kbc.py b/xlb/operator/collision/kbc.py
index cc2fb04a..bbc5d158 100644
--- a/xlb/operator/collision/kbc.py
+++ b/xlb/operator/collision/kbc.py
@@ -178,7 +178,7 @@ def decompose_shear_d2q9_jax(self, fneq):
 
     def _construct_warp(self):
         # Raise error if velocity set is not supported
-        if not isinstance(self.velocity_set, D3Q27):
+        if not (isinstance(self.velocity_set, D3Q27) or isinstance(self.velocity_set, D2Q9)):
             raise NotImplementedError("Velocity set not supported for warp backend: {}".format(type(self.velocity_set)))
 
         # Set local constants TODO: This is a hack and should be fixed with warp update
@@ -192,7 +192,7 @@ def _construct_warp(self):
         def decompose_shear_d2q9(fneq: Any):
             pi = self.momentum_flux.warp_functional(fneq)
             N = pi[0] - pi[1]
-            s = wp.vec9(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
+            s = _f_vec()
             s[3] = N
             s[6] = N
             s[2] = -N

From ec68bfc095091802f8774f27fe244363f85479b7 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Fri, 18 Oct 2024 16:27:39 -0400
Subject: [PATCH 3/5] moved np.unique to examples and a helper function

---
 examples/cfd/flow_past_sphere_3d.py           | 23 ++++++++-----------
 examples/cfd/lid_driven_cavity_2d.py          | 17 ++++++++------
 examples/cfd/turbulent_channel_3d.py          |  4 ++--
 examples/cfd/windtunnel_3d.py                 | 22 ++++++++----------
 examples/performance/mlups_3d.py              | 13 ++++-------
 xlb/grid/grid.py                              |  1 -
 xlb/helper/__init__.py                        |  1 +
 xlb/helper/check_boundary_overlaps.py         | 22 ++++++++++++++++++
 .../indices_boundary_masker.py                |  9 ++------
 9 files changed, 61 insertions(+), 51 deletions(-)
 create mode 100644 xlb/helper/check_boundary_overlaps.py

diff --git a/examples/cfd/flow_past_sphere_3d.py b/examples/cfd/flow_past_sphere_3d.py
index 994d18b0..8b958fc8 100644
--- a/examples/cfd/flow_past_sphere_3d.py
+++ b/examples/cfd/flow_past_sphere_3d.py
@@ -1,7 +1,7 @@
 import xlb
 from xlb.compute_backend import ComputeBackend
 from xlb.precision_policy import PrecisionPolicy
-from xlb.helper import create_nse_fields, initialize_eq
+from xlb.helper import create_nse_fields, initialize_eq, check_bc_overlaps
 from xlb.operator.stepper import IncompressibleNavierStokesStepper
 from xlb.operator.boundary_condition import (
     FullwayBounceBackBC,
@@ -48,15 +48,12 @@ def _setup(self, omega):
         self.setup_stepper(omega)
 
     def define_boundary_indices(self):
-        inlet = self.grid.boundingBoxIndices["left"]
-        outlet = self.grid.boundingBoxIndices["right"]
-        walls = [
-            self.grid.boundingBoxIndices["bottom"][i]
-            + self.grid.boundingBoxIndices["top"][i]
-            + self.grid.boundingBoxIndices["front"][i]
-            + self.grid.boundingBoxIndices["back"][i]
-            for i in range(self.velocity_set.d)
-        ]
+        box = self.grid.bounding_box_indices()
+        box_noedge = self.grid.bounding_box_indices(remove_edges=True)
+        inlet = box_noedge["left"]
+        outlet = box_noedge["right"]
+        walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(self.velocity_set.d)]
+        walls = np.unique(np.array(walls), axis=-1).tolist()
 
         sphere_radius = self.grid_shape[1] // 12
         x = np.arange(self.grid_shape[0])
@@ -79,12 +76,12 @@ def setup_boundary_conditions(self):
         # bc_outlet = DoNothingBC(indices=outlet)
         bc_outlet = ExtrapolationOutflowBC(indices=outlet)
         bc_sphere = HalfwayBounceBackBC(indices=sphere)
-
         self.boundary_conditions = [bc_walls, bc_left, bc_outlet, bc_sphere]
-        # Note: it is important to add bc_walls before bc_outlet/bc_inlet because
-        # of the corner nodes. This way the corners are treated as wall and not inlet/outlet.
 
     def setup_boundary_masker(self):
+        # check boundary condition list for duplicate indices before creating bc mask
+        check_bc_overlaps(self.boundary_conditions, self.velocity_set.d, self.backend)
+
         indices_boundary_masker = IndicesBoundaryMasker(
             velocity_set=self.velocity_set,
             precision_policy=self.precision_policy,
diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py
index 300614da..a77cc626 100644
--- a/examples/cfd/lid_driven_cavity_2d.py
+++ b/examples/cfd/lid_driven_cavity_2d.py
@@ -1,15 +1,16 @@
 import xlb
 from xlb.compute_backend import ComputeBackend
 from xlb.precision_policy import PrecisionPolicy
-from xlb.helper import create_nse_fields, initialize_eq
+from xlb.helper import create_nse_fields, initialize_eq, check_bc_overlaps
 from xlb.operator.boundary_masker import IndicesBoundaryMasker
 from xlb.operator.stepper import IncompressibleNavierStokesStepper
 from xlb.operator.boundary_condition import HalfwayBounceBackBC, EquilibriumBC
 from xlb.operator.macroscopic import Macroscopic
 from xlb.utils import save_fields_vtk, save_image
+import xlb.velocity_set
 import warp as wp
 import jax.numpy as jnp
-import xlb.velocity_set
+import numpy as np
 
 
 class LidDrivenCavity2D:
@@ -39,11 +40,11 @@ def _setup(self, omega):
         self.setup_stepper(omega)
 
     def define_boundary_indices(self):
-        lid = self.grid.boundingBoxIndices["top"]
-        walls = [
-            self.grid.boundingBoxIndices["bottom"][i] + self.grid.boundingBoxIndices["left"][i] + self.grid.boundingBoxIndices["right"][i]
-            for i in range(self.velocity_set.d)
-        ]
+        box = self.grid.bounding_box_indices()
+        box_noedge = self.grid.bounding_box_indices(remove_edges=True)
+        lid = box_noedge["top"]
+        walls = [box["bottom"][i] + box["left"][i] + box["right"][i] for i in range(self.velocity_set.d)]
+        walls = np.unique(np.array(walls), axis=-1).tolist()
         return lid, walls
 
     def setup_boundary_conditions(self):
@@ -53,6 +54,8 @@ def setup_boundary_conditions(self):
         self.boundary_conditions = [bc_walls, bc_top]
 
     def setup_boundary_masker(self):
+        # check boundary condition list for duplicate indices before creating bc mask
+        check_bc_overlaps(self.boundary_conditions, self.velocity_set.d, self.backend)
         indices_boundary_masker = IndicesBoundaryMasker(
             velocity_set=self.velocity_set,
             precision_policy=self.precision_policy,
diff --git a/examples/cfd/turbulent_channel_3d.py b/examples/cfd/turbulent_channel_3d.py
index 2ec55604..eb73fdc6 100644
--- a/examples/cfd/turbulent_channel_3d.py
+++ b/examples/cfd/turbulent_channel_3d.py
@@ -77,8 +77,8 @@ def _setup(self):
 
     def define_boundary_indices(self):
         # top and bottom sides of the channel are no-slip and the other directions are periodic
-        boundingBoxIndices = self.grid.bounding_box_indices(remove_edges=True)
-        walls = [boundingBoxIndices["bottom"][i] + boundingBoxIndices["top"][i] for i in range(self.velocity_set.d)]
+        box = self.grid.bounding_box_indices(remove_edges=True)
+        walls = [box["bottom"][i] + box["top"][i] for i in range(self.velocity_set.d)]
         return walls
 
     def setup_boundary_conditions(self):
diff --git a/examples/cfd/windtunnel_3d.py b/examples/cfd/windtunnel_3d.py
index b0ee5b9d..a7567d06 100644
--- a/examples/cfd/windtunnel_3d.py
+++ b/examples/cfd/windtunnel_3d.py
@@ -3,7 +3,7 @@
 import time
 from xlb.compute_backend import ComputeBackend
 from xlb.precision_policy import PrecisionPolicy
-from xlb.helper import create_nse_fields, initialize_eq
+from xlb.helper import create_nse_fields, initialize_eq, check_bc_overlaps
 from xlb.operator.stepper import IncompressibleNavierStokesStepper
 from xlb.operator.boundary_condition import (
     FullwayBounceBackBC,
@@ -67,15 +67,12 @@ def voxelize_stl(self, stl_filename, length_lbm_unit):
         return mesh_matrix, pitch
 
     def define_boundary_indices(self):
-        inlet = self.grid.boundingBoxIndices["left"]
-        outlet = self.grid.boundingBoxIndices["right"]
-        walls = [
-            self.grid.boundingBoxIndices["bottom"][i]
-            + self.grid.boundingBoxIndices["top"][i]
-            + self.grid.boundingBoxIndices["front"][i]
-            + self.grid.boundingBoxIndices["back"][i]
-            for i in range(self.velocity_set.d)
-        ]
+        box = self.grid.bounding_box_indices()
+        box_noedge = self.grid.bounding_box_indices(remove_edges=True)
+        inlet = box_noedge["left"]
+        outlet = box_noedge["right"]
+        walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(self.velocity_set.d)]
+        walls = np.unique(np.array(walls), axis=-1).tolist()
 
         # Load the mesh
         stl_filename = "examples/cfd/stl-files/DrivAer-Notchback.stl"
@@ -105,10 +102,11 @@ def setup_boundary_conditions(self):
         bc_car = GradsApproximationBC(mesh_vertices=car)
         # bc_car = FullwayBounceBackBC(mesh_vertices=car)
         self.boundary_conditions = [bc_walls, bc_left, bc_do_nothing, bc_car]
-        # Note: it is important to add bc_walls before bc_outlet/bc_inlet because
-        # of the corner nodes. This way the corners are treated as wall and not inlet/outlet.
 
     def setup_boundary_masker(self):
+        # check boundary condition list for duplicate indices before creating bc mask
+        check_bc_overlaps(self.boundary_conditions, self.velocity_set.d, self.backend)
+
         indices_boundary_masker = IndicesBoundaryMasker(
             velocity_set=self.velocity_set,
             precision_policy=self.precision_policy,
diff --git a/examples/performance/mlups_3d.py b/examples/performance/mlups_3d.py
index 1812d957..4c337e85 100644
--- a/examples/performance/mlups_3d.py
+++ b/examples/performance/mlups_3d.py
@@ -48,15 +48,10 @@ def create_grid_and_fields(cube_edge):
 
 
 def define_boundary_indices(grid):
-    lid = grid.boundingBoxIndices["top"]
-    walls = [
-        grid.boundingBoxIndices["bottom"][i]
-        + grid.boundingBoxIndices["left"][i]
-        + grid.boundingBoxIndices["right"][i]
-        + grid.boundingBoxIndices["front"][i]
-        + grid.boundingBoxIndices["back"][i]
-        for i in range(len(grid.shape))
-    ]
+    box = grid.bounding_box_indices()
+    box_noedge = grid.bounding_box_indices(remove_edges=True)
+    lid = box_noedge["top"]
+    walls = [box["bottom"][i] + box["left"][i] + box["right"][i] + box["front"][i] + box["back"][i] for i in range(len(grid.shape))]
     return lid, walls
 
 
diff --git a/xlb/grid/grid.py b/xlb/grid/grid.py
index 7494c3ee..53139fc1 100644
--- a/xlb/grid/grid.py
+++ b/xlb/grid/grid.py
@@ -25,7 +25,6 @@ def __init__(self, shape: Tuple[int, ...], compute_backend: ComputeBackend):
         self.shape = shape
         self.dim = len(shape)
         self.compute_backend = compute_backend
-        self.boundingBoxIndices = self.bounding_box_indices()
         self._initialize_backend()
 
     @abstractmethod
diff --git a/xlb/helper/__init__.py b/xlb/helper/__init__.py
index 92d3583c..4c63aa6b 100644
--- a/xlb/helper/__init__.py
+++ b/xlb/helper/__init__.py
@@ -1,2 +1,3 @@
 from xlb.helper.nse_solver import create_nse_fields as create_nse_fields
 from xlb.helper.initializers import initialize_eq as initialize_eq
+from xlb.helper.check_boundary_overlaps import check_bc_overlaps as check_bc_overlaps
diff --git a/xlb/helper/check_boundary_overlaps.py b/xlb/helper/check_boundary_overlaps.py
new file mode 100644
index 00000000..18e17a16
--- /dev/null
+++ b/xlb/helper/check_boundary_overlaps.py
@@ -0,0 +1,22 @@
+import numpy as np
+from xlb.compute_backend import ComputeBackend
+
+
+def check_bc_overlaps(bclist, dim, backend):
+    index_list = [[] for _ in range(dim)]
+    for bc in bclist:
+        if bc.indices is None:
+            continue
+        # Detect duplicates within bc.indices
+        index_arr = np.unique(bc.indices, axis=-1)
+        if index_arr.shape[-1] != len(bc.indices[0]):
+            if backend == ComputeBackend.WARP:
+                raise ValueError(f"Boundary condition {bc.__class__.__name__} has duplicate indices!")
+        for d in range(dim):
+            index_list[d] += bc.indices[d]
+
+    # Detect duplicates within bclist
+    index_arr = np.unique(index_list, axis=-1)
+    if index_arr.shape[-1] != len(index_list[0]):
+        if backend == ComputeBackend.WARP:
+            raise ValueError("Boundary condition list containes duplicate indices!")
diff --git a/xlb/operator/boundary_masker/indices_boundary_masker.py b/xlb/operator/boundary_masker/indices_boundary_masker.py
index 10d72c43..848d74e5 100644
--- a/xlb/operator/boundary_masker/indices_boundary_masker.py
+++ b/xlb/operator/boundary_masker/indices_boundary_masker.py
@@ -225,14 +225,9 @@ def warp_implementation(self, bclist, bc_mask, missing_mask, start_index=None):
             # We are done with bc.indices. Remove them from BC objects
             bc.__dict__.pop("indices", None)
 
-        # Remove duplicates indices to avoid race conditioning
-        index_arr, unique_loc = np.unique(index_list, axis=-1, return_index=True)
-        id_arr = np.array(id_list)[unique_loc]
-        is_interior = np.array(is_interior)[unique_loc]
-
         # convert to warp arrays
-        indices = wp.array2d(index_arr, dtype=wp.int32)
-        id_number = wp.array1d(id_arr, dtype=wp.uint8)
+        indices = wp.array2d(index_list, dtype=wp.int32)
+        id_number = wp.array1d(id_list, dtype=wp.uint8)
         is_interior = wp.array1d(is_interior, dtype=wp.bool)
 
         if start_index is None:

From 8f869d23d0d7d0a19c5e995eb673395a401451bd Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Fri, 18 Oct 2024 17:02:15 -0400
Subject: [PATCH 4/5] minor

---
 xlb/operator/boundary_masker/indices_boundary_masker.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/xlb/operator/boundary_masker/indices_boundary_masker.py b/xlb/operator/boundary_masker/indices_boundary_masker.py
index 848d74e5..e7e50d74 100644
--- a/xlb/operator/boundary_masker/indices_boundary_masker.py
+++ b/xlb/operator/boundary_masker/indices_boundary_masker.py
@@ -66,7 +66,7 @@ def jax_implementation(self, bclist, bc_mask, missing_mask, start_index=None):
             start_index = (0,) * dim
 
         domain_shape = bc_mask[0].shape
-        for bc in reversed(bclist):
+        for bc in bclist:
             assert bc.indices is not None, f"Please specify indices associated with the {bc.__class__.__name__} BC!"
             assert bc.mesh_vertices is None, f"Please use MeshBoundaryMasker operator if {bc.__class__.__name__} is imposed on a mesh (e.g. STL)!"
             id_number = bc.id

From cee77b9691be618bec7203831545be31740d36e4 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Fri, 18 Oct 2024 17:31:38 -0400
Subject: [PATCH 5/5] addressing PR reviews

---
 examples/cfd/flow_past_sphere_3d.py         | 25 +++------------------
 examples/cfd/lid_driven_cavity_2d.py        |  4 ++--
 examples/cfd/windtunnel_3d.py               |  6 ++---
 examples/performance/mlups_3d.py            |  4 ++--
 xlb/distribute/__init__.py                  |  2 +-
 xlb/experimental/ooc/__init__.py            |  4 ++--
 xlb/helper/__init__.py                      |  6 ++---
 xlb/helper/check_boundary_overlaps.py       |  2 ++
 xlb/operator/__init__.py                    |  4 ++--
 xlb/operator/boundary_condition/__init__.py | 22 +++++++++---------
 xlb/operator/boundary_masker/__init__.py    | 12 +++-------
 xlb/operator/collision/__init__.py          |  8 +++----
 xlb/operator/equilibrium/__init__.py        |  5 +----
 xlb/operator/force/__init__.py              |  4 ++--
 xlb/operator/precision_caster/__init__.py   |  2 +-
 xlb/operator/stepper/__init__.py            |  4 ++--
 xlb/operator/stream/__init__.py             |  2 +-
 xlb/utils/__init__.py                       | 14 ++++++------
 xlb/velocity_set/__init__.py                |  8 +++----
 19 files changed, 55 insertions(+), 83 deletions(-)

diff --git a/examples/cfd/flow_past_sphere_3d.py b/examples/cfd/flow_past_sphere_3d.py
index 8b958fc8..1b5905e3 100644
--- a/examples/cfd/flow_past_sphere_3d.py
+++ b/examples/cfd/flow_past_sphere_3d.py
@@ -49,9 +49,9 @@ def _setup(self, omega):
 
     def define_boundary_indices(self):
         box = self.grid.bounding_box_indices()
-        box_noedge = self.grid.bounding_box_indices(remove_edges=True)
-        inlet = box_noedge["left"]
-        outlet = box_noedge["right"]
+        box_no_edge = self.grid.bounding_box_indices(remove_edges=True)
+        inlet = box_no_edge["left"]
+        outlet = box_no_edge["right"]
         walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(self.velocity_set.d)]
         walls = np.unique(np.array(walls), axis=-1).tolist()
 
@@ -101,8 +101,6 @@ def run(self, num_steps, post_process_interval=100):
             self.f_0, self.f_1 = self.stepper(self.f_0, self.f_1, self.bc_mask, self.missing_mask, i)
             self.f_0, self.f_1 = self.f_1, self.f_0
 
-            if i == 0:
-                self.check_boundary_mask()
             if i % post_process_interval == 0 or i == num_steps - 1:
                 self.post_process(i)
                 end_time = time.time()
@@ -132,23 +130,6 @@ def post_process(self, i):
 
         # save_fields_vtk(fields, timestep=i)
         save_image(fields["u_magnitude"][:, self.grid_shape[1] // 2, :], timestep=i)
-        return
-
-    def check_boundary_mask(self):
-        # Write the results. We'll use JAX backend for the post-processing
-        if not isinstance(self.f_0, jnp.ndarray):
-            bmask = wp.to_jax(self.bc_mask)[0]
-        else:
-            bmask = self.bc_mask[0]
-
-        # save_fields_vtk(fields, timestep=i)
-        save_image(bmask[0, :, :], prefix="00_left")
-        save_image(bmask[self.grid_shape[0] - 1, :, :], prefix="00_right")
-        save_image(bmask[:, :, self.grid_shape[2] - 1], prefix="00_top")
-        save_image(bmask[:, :, 0], prefix="00_bottom")
-        save_image(bmask[:, 0, :], prefix="00_front")
-        save_image(bmask[:, self.grid_shape[1] - 1, :], prefix="00_back")
-        save_image(bmask[:, self.grid_shape[1] // 2, :], prefix="00_middle")
 
 
 if __name__ == "__main__":
diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py
index a77cc626..dfb10928 100644
--- a/examples/cfd/lid_driven_cavity_2d.py
+++ b/examples/cfd/lid_driven_cavity_2d.py
@@ -41,8 +41,8 @@ def _setup(self, omega):
 
     def define_boundary_indices(self):
         box = self.grid.bounding_box_indices()
-        box_noedge = self.grid.bounding_box_indices(remove_edges=True)
-        lid = box_noedge["top"]
+        box_no_edge = self.grid.bounding_box_indices(remove_edges=True)
+        lid = box_no_edge["top"]
         walls = [box["bottom"][i] + box["left"][i] + box["right"][i] for i in range(self.velocity_set.d)]
         walls = np.unique(np.array(walls), axis=-1).tolist()
         return lid, walls
diff --git a/examples/cfd/windtunnel_3d.py b/examples/cfd/windtunnel_3d.py
index a7567d06..c83e2c9e 100644
--- a/examples/cfd/windtunnel_3d.py
+++ b/examples/cfd/windtunnel_3d.py
@@ -68,9 +68,9 @@ def voxelize_stl(self, stl_filename, length_lbm_unit):
 
     def define_boundary_indices(self):
         box = self.grid.bounding_box_indices()
-        box_noedge = self.grid.bounding_box_indices(remove_edges=True)
-        inlet = box_noedge["left"]
-        outlet = box_noedge["right"]
+        box_no_edge = self.grid.bounding_box_indices(remove_edges=True)
+        inlet = box_no_edge["left"]
+        outlet = box_no_edge["right"]
         walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(self.velocity_set.d)]
         walls = np.unique(np.array(walls), axis=-1).tolist()
 
diff --git a/examples/performance/mlups_3d.py b/examples/performance/mlups_3d.py
index 4c337e85..2001fb2f 100644
--- a/examples/performance/mlups_3d.py
+++ b/examples/performance/mlups_3d.py
@@ -49,8 +49,8 @@ def create_grid_and_fields(cube_edge):
 
 def define_boundary_indices(grid):
     box = grid.bounding_box_indices()
-    box_noedge = grid.bounding_box_indices(remove_edges=True)
-    lid = box_noedge["top"]
+    box_no_edge = grid.bounding_box_indices(remove_edges=True)
+    lid = box_no_edge["top"]
     walls = [box["bottom"][i] + box["left"][i] + box["right"][i] + box["front"][i] + box["back"][i] for i in range(len(grid.shape))]
     return lid, walls
 
diff --git a/xlb/distribute/__init__.py b/xlb/distribute/__init__.py
index 25fa0af9..dd9f33dd 100644
--- a/xlb/distribute/__init__.py
+++ b/xlb/distribute/__init__.py
@@ -1 +1 @@
-from .distribute import distribute as distribute
+from .distribute import distribute
diff --git a/xlb/experimental/ooc/__init__.py b/xlb/experimental/ooc/__init__.py
index 801683d1..5206cc18 100644
--- a/xlb/experimental/ooc/__init__.py
+++ b/xlb/experimental/ooc/__init__.py
@@ -1,2 +1,2 @@
-from xlb.experimental.ooc.out_of_core import OOCmap as OOCmap
-from xlb.experimental.ooc.ooc_array import OOCArray as OOCArray
+from xlb.experimental.ooc.out_of_core import OOCmap
+from xlb.experimental.ooc.ooc_array import OOCArray
diff --git a/xlb/helper/__init__.py b/xlb/helper/__init__.py
index 4c63aa6b..d52f2063 100644
--- a/xlb/helper/__init__.py
+++ b/xlb/helper/__init__.py
@@ -1,3 +1,3 @@
-from xlb.helper.nse_solver import create_nse_fields as create_nse_fields
-from xlb.helper.initializers import initialize_eq as initialize_eq
-from xlb.helper.check_boundary_overlaps import check_bc_overlaps as check_bc_overlaps
+from xlb.helper.nse_solver import create_nse_fields
+from xlb.helper.initializers import initialize_eq
+from xlb.helper.check_boundary_overlaps import check_bc_overlaps
diff --git a/xlb/helper/check_boundary_overlaps.py b/xlb/helper/check_boundary_overlaps.py
index 18e17a16..1adceb2b 100644
--- a/xlb/helper/check_boundary_overlaps.py
+++ b/xlb/helper/check_boundary_overlaps.py
@@ -12,6 +12,7 @@ def check_bc_overlaps(bclist, dim, backend):
         if index_arr.shape[-1] != len(bc.indices[0]):
             if backend == ComputeBackend.WARP:
                 raise ValueError(f"Boundary condition {bc.__class__.__name__} has duplicate indices!")
+            print(f"WARNING: there are duplicate indices in {bc.__class__.__name__} and hence the order in bc list matters!")
         for d in range(dim):
             index_list[d] += bc.indices[d]
 
@@ -20,3 +21,4 @@ def check_bc_overlaps(bclist, dim, backend):
     if index_arr.shape[-1] != len(index_list[0]):
         if backend == ComputeBackend.WARP:
             raise ValueError("Boundary condition list containes duplicate indices!")
+        print("WARNING: there are duplicate indices in the boundary condition list and hence the order in this list matters!")
diff --git a/xlb/operator/__init__.py b/xlb/operator/__init__.py
index c88ef83a..02b8a590 100644
--- a/xlb/operator/__init__.py
+++ b/xlb/operator/__init__.py
@@ -1,2 +1,2 @@
-from xlb.operator.operator import Operator as Operator
-from xlb.operator.parallel_operator import ParallelOperator as ParallelOperator
+from xlb.operator.operator import Operator
+from xlb.operator.parallel_operator import ParallelOperator
diff --git a/xlb/operator/boundary_condition/__init__.py b/xlb/operator/boundary_condition/__init__.py
index 925dfdc6..4782ea0e 100644
--- a/xlb/operator/boundary_condition/__init__.py
+++ b/xlb/operator/boundary_condition/__init__.py
@@ -1,12 +1,10 @@
-from xlb.operator.boundary_condition.boundary_condition import BoundaryCondition as BoundaryCondition
-from xlb.operator.boundary_condition.boundary_condition_registry import (
-    BoundaryConditionRegistry as BoundaryConditionRegistry,
-)
-from xlb.operator.boundary_condition.bc_equilibrium import EquilibriumBC as EquilibriumBC
-from xlb.operator.boundary_condition.bc_do_nothing import DoNothingBC as DoNothingBC
-from xlb.operator.boundary_condition.bc_halfway_bounce_back import HalfwayBounceBackBC as HalfwayBounceBackBC
-from xlb.operator.boundary_condition.bc_fullway_bounce_back import FullwayBounceBackBC as FullwayBounceBackBC
-from xlb.operator.boundary_condition.bc_zouhe import ZouHeBC as ZouHeBC
-from xlb.operator.boundary_condition.bc_regularized import RegularizedBC as RegularizedBC
-from xlb.operator.boundary_condition.bc_extrapolation_outflow import ExtrapolationOutflowBC as ExtrapolationOutflowBC
-from xlb.operator.boundary_condition.bc_grads_approximation import GradsApproximationBC as GradsApproximationBC
+from xlb.operator.boundary_condition.boundary_condition import BoundaryCondition
+from xlb.operator.boundary_condition.boundary_condition_registry import BoundaryConditionRegistry
+from xlb.operator.boundary_condition.bc_equilibrium import EquilibriumBC
+from xlb.operator.boundary_condition.bc_do_nothing import DoNothingBC
+from xlb.operator.boundary_condition.bc_halfway_bounce_back import HalfwayBounceBackBC
+from xlb.operator.boundary_condition.bc_fullway_bounce_back import FullwayBounceBackBC
+from xlb.operator.boundary_condition.bc_zouhe import ZouHeBC
+from xlb.operator.boundary_condition.bc_regularized import RegularizedBC
+from xlb.operator.boundary_condition.bc_extrapolation_outflow import ExtrapolationOutflowBC
+from xlb.operator.boundary_condition.bc_grads_approximation import GradsApproximationBC
diff --git a/xlb/operator/boundary_masker/__init__.py b/xlb/operator/boundary_masker/__init__.py
index fbe851dc..d03636a8 100644
--- a/xlb/operator/boundary_masker/__init__.py
+++ b/xlb/operator/boundary_masker/__init__.py
@@ -1,9 +1,3 @@
-from xlb.operator.boundary_masker.indices_boundary_masker import (
-    IndicesBoundaryMasker as IndicesBoundaryMasker,
-)
-from xlb.operator.boundary_masker.mesh_boundary_masker import (
-    MeshBoundaryMasker as MeshBoundaryMasker,
-)
-from xlb.operator.boundary_masker.mesh_distance_boundary_masker import (
-    MeshDistanceBoundaryMasker as MeshDistanceBoundaryMasker,
-)
+from xlb.operator.boundary_masker.indices_boundary_masker import IndicesBoundaryMasker
+from xlb.operator.boundary_masker.mesh_boundary_masker import MeshBoundaryMasker
+from xlb.operator.boundary_masker.mesh_distance_boundary_masker import MeshDistanceBoundaryMasker
diff --git a/xlb/operator/collision/__init__.py b/xlb/operator/collision/__init__.py
index 0526c8af..2f92bb4d 100644
--- a/xlb/operator/collision/__init__.py
+++ b/xlb/operator/collision/__init__.py
@@ -1,4 +1,4 @@
-from xlb.operator.collision.collision import Collision as Collision
-from xlb.operator.collision.bgk import BGK as BGK
-from xlb.operator.collision.kbc import KBC as KBC
-from xlb.operator.collision.forced_collision import ForcedCollision as ForcedCollision
+from xlb.operator.collision.collision import Collision
+from xlb.operator.collision.bgk import BGK
+from xlb.operator.collision.kbc import KBC
+from xlb.operator.collision.forced_collision import ForcedCollision
diff --git a/xlb/operator/equilibrium/__init__.py b/xlb/operator/equilibrium/__init__.py
index b9f9f088..987aa74a 100644
--- a/xlb/operator/equilibrium/__init__.py
+++ b/xlb/operator/equilibrium/__init__.py
@@ -1,4 +1 @@
-from xlb.operator.equilibrium.quadratic_equilibrium import (
-    Equilibrium as Equilibrium,
-    QuadraticEquilibrium as QuadraticEquilibrium,
-)
+from xlb.operator.equilibrium.quadratic_equilibrium import Equilibrium, QuadraticEquilibrium
diff --git a/xlb/operator/force/__init__.py b/xlb/operator/force/__init__.py
index 2f3e3da9..ba8a13c3 100644
--- a/xlb/operator/force/__init__.py
+++ b/xlb/operator/force/__init__.py
@@ -1,2 +1,2 @@
-from xlb.operator.force.momentum_transfer import MomentumTransfer as MomentumTransfer
-from xlb.operator.force.exact_difference_force import ExactDifference as ExactDifference
+from xlb.operator.force.momentum_transfer import MomentumTransfer
+from xlb.operator.force.exact_difference_force import ExactDifference
diff --git a/xlb/operator/precision_caster/__init__.py b/xlb/operator/precision_caster/__init__.py
index c333ab7a..a027c524 100644
--- a/xlb/operator/precision_caster/__init__.py
+++ b/xlb/operator/precision_caster/__init__.py
@@ -1 +1 @@
-from xlb.operator.precision_caster.precision_caster import PrecisionCaster as PrecisionCaster
+from xlb.operator.precision_caster.precision_caster import PrecisionCaster
diff --git a/xlb/operator/stepper/__init__.py b/xlb/operator/stepper/__init__.py
index 528375d0..e5d159c6 100644
--- a/xlb/operator/stepper/__init__.py
+++ b/xlb/operator/stepper/__init__.py
@@ -1,2 +1,2 @@
-from xlb.operator.stepper.stepper import Stepper as Stepper
-from xlb.operator.stepper.nse_stepper import IncompressibleNavierStokesStepper as IncompressibleNavierStokesStepper
+from xlb.operator.stepper.stepper import Stepper
+from xlb.operator.stepper.nse_stepper import IncompressibleNavierStokesStepper
diff --git a/xlb/operator/stream/__init__.py b/xlb/operator/stream/__init__.py
index 2f5b2f32..9093da71 100644
--- a/xlb/operator/stream/__init__.py
+++ b/xlb/operator/stream/__init__.py
@@ -1 +1 @@
-from xlb.operator.stream.stream import Stream as Stream
+from xlb.operator.stream.stream import Stream
diff --git a/xlb/utils/__init__.py b/xlb/utils/__init__.py
index 6f1f61ab..3c8032e2 100644
--- a/xlb/utils/__init__.py
+++ b/xlb/utils/__init__.py
@@ -1,9 +1,9 @@
 from .utils import (
-    downsample_field as downsample_field,
-    save_image as save_image,
-    save_fields_vtk as save_fields_vtk,
-    save_BCs_vtk as save_BCs_vtk,
-    rotate_geometry as rotate_geometry,
-    voxelize_stl as voxelize_stl,
-    axangle2mat as axangle2mat,
+    downsample_field,
+    save_image,
+    save_fields_vtk,
+    save_BCs_vtk,
+    rotate_geometry,
+    voxelize_stl,
+    axangle2mat,
 )
diff --git a/xlb/velocity_set/__init__.py b/xlb/velocity_set/__init__.py
index c1338dbf..5b7b737f 100644
--- a/xlb/velocity_set/__init__.py
+++ b/xlb/velocity_set/__init__.py
@@ -1,4 +1,4 @@
-from xlb.velocity_set.velocity_set import VelocitySet as VelocitySet
-from xlb.velocity_set.d2q9 import D2Q9 as D2Q9
-from xlb.velocity_set.d3q19 import D3Q19 as D3Q19
-from xlb.velocity_set.d3q27 import D3Q27 as D3Q27
+from xlb.velocity_set.velocity_set import VelocitySet
+from xlb.velocity_set.d2q9 import D2Q9
+from xlb.velocity_set.d3q19 import D3Q19
+from xlb.velocity_set.d3q27 import D3Q27