Fixed zouhe and regularized bcs but there are limitations in applying…

… the same bc type multiple times

xlb/operator/boundary_condition/bc_do_nothing.py

@@ -27,8 +27,6 @@ class DoNothingBC(BoundaryCondition):
     boundary nodes.
     """
 
-    id = boundary_condition_registry.register_boundary_condition(__qualname__)
-
     def __init__(
         self,
         velocity_set: VelocitySet = None,

xlb/operator/boundary_condition/bc_equilibrium.py

@@ -29,8 +29,6 @@ class EquilibriumBC(BoundaryCondition):
     Full Bounce-back boundary condition for a lattice Boltzmann method simulation.
     """
 
-    id = boundary_condition_registry.register_boundary_condition(__qualname__)
-
     def __init__(
         self,
         rho: float,

xlb/operator/boundary_condition/bc_extrapolation_outflow.py

@@ -38,8 +38,6 @@ class ExtrapolationOutflowBC(BoundaryCondition):
     doi:10.1016/j.camwa.2015.05.001.
     """
 
-    id = boundary_condition_registry.register_boundary_condition(__qualname__)
-
     def __init__(
         self,
         velocity_set: VelocitySet = None,

xlb/operator/boundary_condition/bc_fullway_bounce_back.py

@@ -27,8 +27,6 @@ class FullwayBounceBackBC(BoundaryCondition):
     Full Bounce-back boundary condition for a lattice Boltzmann method simulation.
     """
 
-    id = boundary_condition_registry.register_boundary_condition(__qualname__)
-
     def __init__(
         self,
         velocity_set: VelocitySet = None,

xlb/operator/boundary_condition/bc_grads_approximation.py

@@ -39,8 +39,6 @@ class GradsApproximationBC(BoundaryCondition):
 
     """
 
-    id = boundary_condition_registry.register_boundary_condition(__qualname__)
-
     def __init__(
         self,
         velocity_set: VelocitySet = None,
@@ -49,6 +47,7 @@ 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)
 

xlb/operator/boundary_condition/bc_halfway_bounce_back.py

@@ -29,8 +29,6 @@ class HalfwayBounceBackBC(BoundaryCondition):
     TODO: Implement moving boundary conditions for this
     """
 
-    id = boundary_condition_registry.register_boundary_condition(__qualname__)
-
     def __init__(
         self,
         velocity_set: VelocitySet = None,

xlb/operator/boundary_condition/bc_regularized.py

@@ -61,9 +61,8 @@ def __init__(
             indices,
             mesh_vertices,
         )
-
-        self.id = boundary_condition_registry.register_boundary_condition(__class__.__name__)
-
+        # Overwrite the boundary condition registry id with the bc_type in the name
+        self.id = boundary_condition_registry.register_boundary_condition(self.__class__.__name__ + "_" + bc_type)
         # The operator to compute the momentum flux
         self.momentum_flux = MomentumFlux()
 

xlb/operator/boundary_condition/bc_zouhe.py

@@ -35,6 +35,8 @@ class ZouHeBC(BoundaryCondition):
     Reynolds numbers. One needs to use "Regularized" BC at higher Reynolds.
     """
 
+
+
     def __init__(
         self,
         bc_type,
@@ -48,7 +50,6 @@ def __init__(
         # Important Note: it is critical to add id inside __init__ for this BC because different instantiations of this BC
         # may have different types (velocity or pressure).
         assert bc_type in ["velocity", "pressure"], f"type = {bc_type} not supported! Use 'pressure' or 'velocity'."
-        self.id = boundary_condition_registry.register_boundary_condition(__class__.__name__ + "_" + bc_type)
         self.bc_type = bc_type
         self.equilibrium_operator = QuadraticEquilibrium()
         self.prescribed_value = prescribed_value
@@ -63,6 +64,9 @@ def __init__(
             mesh_vertices,
         )
 
+        # Overwrite the boundary condition registry id with the bc_type in the name
+        self.id = boundary_condition_registry.register_boundary_condition(self.__class__.__name__ + "_" + bc_type)
+
         # Set the prescribed value for pressure or velocity
         dim = self.velocity_set.d
         if self.compute_backend == ComputeBackend.JAX:

xlb/operator/boundary_condition/boundary_condition.py

@@ -13,7 +13,7 @@
 from xlb.compute_backend import ComputeBackend
 from xlb.operator.operator import Operator
 from xlb import DefaultConfig
-
+from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry
 
 # Enum for implementation step
 class ImplementationStep(Enum):
@@ -35,6 +35,7 @@ def __init__(
         indices=None,
         mesh_vertices=None,
     ):
+        self.id = boundary_condition_registry.register_boundary_condition(self.__class__.__name__)
         velocity_set = velocity_set or DefaultConfig.velocity_set
         precision_policy = precision_policy or DefaultConfig.default_precision_policy
         compute_backend = compute_backend or DefaultConfig.default_backend

xlb/operator/stepper/nse_stepper.py

@@ -40,9 +40,6 @@ def __init__(self, omega, boundary_conditions=[], collision_type="BGK", forcing_
 
         operators = [self.macroscopic, self.equilibrium, self.collision, self.stream]
 
-        self.boundary_conditions = boundary_conditions
-        self.active_bcs = set(type(bc).__name__ for bc in boundary_conditions)
-
         super().__init__(operators, boundary_conditions)
 
     @Operator.register_backend(ComputeBackend.JAX)
@@ -103,6 +100,8 @@ def _construct_warp(self):
         bc_to_id = boundary_condition_registry.bc_to_id
         id_to_bc = boundary_condition_registry.id_to_bc
 
+        active_bcs = set(boundary_condition_registry.id_to_bc[bc.id] for bc in self.boundary_conditions)
+
         for bc in self.boundary_conditions:
             bc_name = id_to_bc[bc.id]
             setattr(self, bc_name, bc)
@@ -120,32 +119,40 @@ def apply_post_streaming_bc(
         ):
             f_result = f_post
 
-            if wp.static("EquilibriumBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["EquilibriumBC"]):
+            if wp.static("EquilibriumBC" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["EquilibriumBC"]):
                     f_result = self.EquilibriumBC.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
-            if wp.static("DoNothingBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["DoNothingBC"]):
+            if wp.static("DoNothingBC" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["DoNothingBC"]):
                     f_result = self.DoNothingBC.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
-            if wp.static("HalfwayBounceBackBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["HalfwayBounceBackBC"]):
+            if wp.static("HalfwayBounceBackBC" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["HalfwayBounceBackBC"]):
                     f_result = self.HalfwayBounceBackBC.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
-            if wp.static("ZouHeBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["ZouHeBC"]):
+            if wp.static("ZouHeBC_pressure" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["ZouHeBC_pressure"]):
+                    f_result = self.ZouHeBC_pressure.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
+
+            if wp.static("ZouHeBC_velocity" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["ZouHeBC_velocity"]):
                     f_result = self.ZouHeBC_velocity.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
-            if wp.static("RegularizedBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["RegularizedBC"]):
-                    f_result = self.RegularizedBC.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
+            if wp.static("RegularizedBC_pressure" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["RegularizedBC_pressure"]):
+                    f_result = self.RegularizedBC_pressure.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
+
+            if wp.static("RegularizedBC_velocity" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["RegularizedBC_velocity"]):
+                    f_result = self.RegularizedBC_velocity.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
-            if wp.static("ExtrapolationOutflowBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["ExtrapolationOutflowBC"]):
+            if wp.static("ExtrapolationOutflowBC" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["ExtrapolationOutflowBC"]):
                     f_result = self.ExtrapolationOutflowBC.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
-            if wp.static("GradsApproximationBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["GradsApproximationBC"]):
+            if wp.static("GradsApproximationBC" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["GradsApproximationBC"]):
                     f_result = self.GradsApproximationBC.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
             return f_result
@@ -163,12 +170,12 @@ def apply_post_collision_bc(
         ):
             f_result = f_post
 
-            if wp.static("FullwayBounceBackBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["FullwayBounceBackBC"]):
+            if wp.static("FullwayBounceBackBC" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["FullwayBounceBackBC"]):
                     f_result = self.FullwayBounceBackBC.warp_functional(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
-            if wp.static("ExtrapolationOutflowBC" in self.active_bcs):
-                if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["ExtrapolationOutflowBC"]):
+            if wp.static("ExtrapolationOutflowBC" in active_bcs):
+                if _boundary_id == wp.static(bc_to_id["ExtrapolationOutflowBC"]):
                     f_result = self.ExtrapolationOutflowBC.prepare_bc_auxilary_data(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
 
             return f_result
@@ -288,7 +295,7 @@ def kernel3d(
 
             # Store the result in f_1
             for l in range(self.velocity_set.q):
-                if wp.static("GradsApproximationBC" in self.active_bcs):
+                if wp.static("GradsApproximationBC" in active_bcs):
                     if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["GradsApproximationBC"]):
                         if _missing_mask[l] == wp.uint8(1):
                             f_0[_opp_indices[l], index[0], index[1], index[2]] = self.store_dtype(_f1_thread[_opp_indices[l]])