Repository: Autodesk/XLB (language: python)

Issue:
Added boundary force computation and more
This PR includes major improvements for:

1. boundary force computations using momentum transfer method 
2. handling of stl (or mesh) geometries for assigning BCs
3. handling of non-mesh boundary conditions in the interior of the domain 
4. fixed syntax error in Warp for scenarios where bc list is empty (mysterious perf issue remains when bc list is empty)


Discussion / hints:

Fix patch:
diff --git a/examples/cfd/data/turbulent_channel_dns_data.json b/examples/cfd/data/turbulent_channel_dns_data.json
new file mode 100644
index 0000000..bdbee75
--- /dev/null
+++ b/examples/cfd/data/turbulent_channel_dns_data.json
@@ -0,0 +1,1 @@
+{"y": [0, 0.000301, 0.0012, 0.00271, 0.00482, 0.00752, 0.0108, 0.0147, 0.0192, 0.0243, 0.03, 0.0362, 0.0431, 0.0505, 0.0585, 0.067, 0.0761, 0.0858, 0.096, 0.107, 0.118, 0.13, 0.142, 0.155, 0.169, 0.182, 0.197, 0.212, 0.227, 0.243, 0.259, 0.276, 0.293, 0.31, 0.328, 0.347, 0.366, 0.385, 0.404, 0.424, 0.444, 0.465, 0.486, 0.507, 0.529, 0.55, 0.572, 0.595, 0.617, 0.64, 0.663, 0.686, 0.71, 0.733, 0.757, 0.781, 0.805, 0.829, 0.853, 0.878, 0.902, 0.926, 0.951, 0.975, 1], "y+": [0, 0.053648, 0.21456, 0.48263, 0.85771, 1.3396, 1.9279, 2.6224, 3.4226, 4.328, 5.3381, 6.4523, 7.67, 8.9902, 10.412, 11.936, 13.559, 15.281, 17.102, 19.019, 21.033, 23.141, 25.342, 27.635, 30.019, 32.492, 35.053, 37.701, 40.432, 43.247, 46.143, 49.118, 52.171, 55.3, 58.503, 61.778, 65.123, 68.536, 72.016, 75.559, 79.164, 82.828, 86.55, 90.327, 94.157, 98.037, 101.97, 105.94, 109.96, 114.02, 118.12, 122.25, 126.42, 130.62, 134.84, 139.1, 143.37, 147.67, 151.99, 156.32, 160.66, 165.02, 169.38, 173.75, 178.12], "Umean": [0, 0.053639, 0.21443, 0.48197, 0.85555, 1.3339, 1.9148, 2.5939, 3.3632, 4.2095, 5.1133, 6.0493, 6.9892, 7.9052, 8.7741, 9.579, 10.311, 10.967, 11.55, 12.066, 12.52, 12.921, 13.276, 13.59, 13.87, 14.121, 14.349, 14.557, 14.75, 14.931, 15.101, 15.264, 15.419, 15.569, 15.714, 15.855, 15.993, 16.128, 16.26, 16.389, 16.515, 16.637, 16.756, 16.872, 16.985, 17.094, 17.2, 17.302, 17.4, 17.494, 17.585, 17.672, 17.756, 17.835, 17.911, 17.981, 18.045, 18.103, 18.154, 18.198, 18.235, 18.264, 18.285, 18.297, 18.301], "dUmean/dy": [178, 178, 178, 178, 177, 176, 175, 173, 169, 163, 155, 144, 131, 116, 101, 87.1, 73.9, 62.2, 52.2, 43.8, 36.9, 31.1, 26.4, 22.6, 19.4, 16.9, 14.9, 13.3, 12, 10.9, 10.1, 9.38, 8.79, 8.29, 7.86, 7.49, 7.19, 6.91, 6.63, 6.35, 6.07, 5.81, 5.58, 5.36, 5.14, 4.92, 4.68, 4.45, 4.23, 4.04, 3.85, 3.66, 3.48, 3.28, 3.06, 2.81, 2.54, 2.25, 1.96, 1.67, 1.35, 1.02, 0.673, 0.33, 0], "Wmean": [0, 7.07e-05, 0.000283, 0.000636, 0.00113, 0.00176, 0.00252, 0.00339, 0.00435, 0.00538, 0.00643, 0.00751, 0.00864, 0.00986, 0.0112, 0.0126, 0.0141, 0.0156, 0.017, 0.0181, 0.0186, 0.0184, 0.0176, 0.0163, 0.0149, 0.0135, 0.0124, 0.0116, 0.0107, 0.00966, 0.00843, 0.00695, 0.00519, 0.00329, 0.00145, -0.000284, -0.00177, -0.00292, -0.00377, -0.00445, -0.00497, -0.0054, -0.00594, -0.00681, -0.0082, -0.00996, -0.0119, -0.0139, -0.0163, -0.0191, -0.0225, -0.0263, -0.0306, -0.0354, -0.0405, -0.0455, -0.05, -0.0539, -0.0577, -0.0615, -0.0653, -0.0685, -0.071, -0.0724, -0.0729], "dWmean/dy": [0.235, 0.235, 0.235, 0.234, 0.234, 0.232, 0.228, 0.22, 0.208, 0.194, 0.179, 0.168, 0.164, 0.164, 0.166, 0.167, 0.162, 0.148, 0.121, 0.076, 0.0159, -0.0439, -0.087, -0.107, -0.106, -0.0871, -0.0643, -0.0546, -0.061, -0.0707, -0.0818, -0.0958, -0.108, -0.106, -0.0989, -0.0881, -0.0697, -0.0506, -0.0379, -0.0303, -0.0221, -0.0216, -0.0314, -0.0522, -0.0756, -0.0841, -0.0884, -0.0974, -0.114, -0.136, -0.154, -0.172, -0.196, -0.214, -0.215, -0.199, -0.174, -0.156, -0.155, -0.159, -0.147, -0.118, -0.0788, -0.0387, 0], "Pmean": [6.217e-13, -7.3193e-10, -1.5832e-07, -3.7598e-06, -3.3837e-05, -0.00017683, -0.00065008, -0.001865, -0.0044488, -0.0092047, -0.017023, -0.028777, -0.045228, -0.066952, -0.094281, -0.12724, -0.16551, -0.20842, -0.25498, -0.30396, -0.35398, -0.40362, -0.45163, -0.49698, -0.5388, -0.57639, -0.60919, -0.63686, -0.6593, -0.67652, -0.68867, -0.69613, -0.69928, -0.69854, -0.69444, -0.68744, -0.67802, -0.66675, -0.65429, -0.64131, -0.62817, -0.61487, -0.60122, -0.58703, -0.57221, -0.55678, -0.5409, -0.52493, -0.50917, -0.49371, -0.47867, -0.46421, -0.4505, -0.43759, -0.4255, -0.41436, -0.40444, -0.39595, -0.389, -0.3836, -0.37966, -0.37702, -0.37542, -0.3746, -0.37436]}
\ No newline at end of file
diff --git a/examples/cfd/flow_past_sphere_3d.py b/examples/cfd/flow_past_sphere_3d.py
index 5a70d20..8c99dff 100644
--- a/examples/cfd/flow_past_sphere_3d.py
+++ b/examples/cfd/flow_past_sphere_3d.py
@@ -19,7 +19,6 @@
 import numpy as np
 import jax.numpy as jnp
 import time
-import jax
 
 
 class FlowOverSphere:
@@ -35,7 +34,7 @@ def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
         self.velocity_set = velocity_set
         self.backend = backend
         self.precision_policy = precision_policy
-        self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_map = create_nse_fields(grid_shape)
+        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 = []
 
@@ -92,7 +91,7 @@ def setup_boundary_masker(self):
             precision_policy=self.precision_policy,
             compute_backend=self.backend,
         )
-        self.boundary_map, self.missing_mask = indices_boundary_masker(self.boundary_conditions, self.boundary_map, self.missing_mask, (0, 0, 0))
+        self.bc_mask, self.missing_mask = indices_boundary_masker(self.boundary_conditions, self.bc_mask, self.missing_mask, (0, 0, 0))
 
     def initialize_fields(self):
         self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend)
@@ -103,7 +102,7 @@ def setup_stepper(self, omega):
     def run(self, num_steps, post_process_interval=100):
         start_time = time.time()
         for i in range(num_steps):
-            self.f_1 = self.stepper(self.f_0, self.f_1, self.boundary_map, self.missing_mask, i)
+            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 % post_process_interval == 0 or i == num_steps - 1:
@@ -143,9 +142,6 @@ def post_process(self, i):
     backend = ComputeBackend.WARP
     precision_policy = PrecisionPolicy.FP32FP32
 
-    if precision_policy == PrecisionPolicy.FP64FP64 or precision_policy == PrecisionPolicy.FP64FP32:
-        jax.config.update("jax_enable_x64", True)
-
     velocity_set = xlb.velocity_set.D3Q19(precision_policy=precision_policy, backend=backend)
     omega = 1.6
 
diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py
index f73fb0e..c67ce8e 100644
--- a/examples/cfd/lid_driven_cavity_2d.py
+++ b/examples/cfd/lid_driven_cavity_2d.py
@@ -8,7 +8,6 @@
 from xlb.operator.macroscopic import Macroscopic
 from xlb.utils import save_fields_vtk, save_image
 import warp as wp
-import jax
 import jax.numpy as jnp
 import xlb.velocity_set
 
@@ -26,7 +25,7 @@ def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
         self.velocity_set = velocity_set
         self.backend = backend
         self.precision_policy = precision_policy
-        self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_map = create_nse_fields(grid_shape)
+        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 = []
 
@@ -59,7 +58,7 @@ def setup_boundary_masker(self):
             precision_policy=self.precision_policy,
             compute_backend=self.backend,
         )
-        self.boundary_map, self.missing_mask = indices_boundary_masker(self.boundary_conditions, self.boundary_map, self.missing_mask)
+        self.bc_mask, self.missing_mask = indices_boundary_masker(self.boundary_conditions, self.bc_mask, self.missing_mask)
 
     def initialize_fields(self):
         self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend)
@@ -69,7 +68,7 @@ def setup_stepper(self, omega):
 
     def run(self, num_steps, post_process_interval=100):
         for i in range(num_steps):
-            self.f_1 = self.stepper(self.f_0, self.f_1, self.boundary_map, self.missing_mask, i)
+            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 % post_process_interval == 0 or i == num_steps - 1:
@@ -108,9 +107,6 @@ def post_process(self, i):
     backend = ComputeBackend.WARP
     precision_policy = PrecisionPolicy.FP32FP32
 
-    if precision_policy == PrecisionPolicy.FP64FP64 or precision_policy == PrecisionPolicy.FP64FP32:
-        jax.config.update("jax_enable_x64", True)
-
     velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend)
     omega = 1.6
 
diff --git a/examples/cfd/lid_driven_cavity_2d_distributed.py b/examples/cfd/lid_driven_cavity_2d_distributed.py
index 7efe907..a7d43b4 100644
--- a/examples/cfd/lid_driven_cavity_2d_distributed.py
+++ b/examples/cfd/lid_driven_cavity_2d_distributed.py
@@ -1,5 +1,4 @@
 import xlb
-import jax
 from xlb.compute_backend import ComputeBackend
 from xlb.precision_policy import PrecisionPolicy
 from xlb.operator.stepper import IncompressibleNavierStokesStepper
@@ -29,9 +28,6 @@ def setup_stepper(self, omega):
     backend = ComputeBackend.JAX  # Must be JAX for distributed multi-GPU computations. Distributed computations on WARP are not supported yet!
     precision_policy = PrecisionPolicy.FP32FP32
 
-    if precision_policy == PrecisionPolicy.FP64FP64 or precision_policy == PrecisionPolicy.FP64FP32:
-        jax.config.update("jax_enable_x64", True)
-
     velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend)
     omega = 1.6
 
diff --git a/examples/cfd/turbulent_channel_3d.py b/examples/cfd/turbulent_channel_3d.py
new file mode 100644
index 0000000..ea8cad7
--- /dev/null
+++ b/examples/cfd/turbulent_channel_3d.py
@@ -0,0 +1,207 @@
+import xlb
+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.operator.stepper import IncompressibleNavierStokesStepper
+from xlb.operator.boundary_condition import RegularizedBC
+from xlb.operator.macroscopic import Macroscopic
+from xlb.operator.boundary_masker import IndicesBoundaryMasker
+from xlb.utils import save_fields_vtk, save_image
+import warp as wp
+import numpy as np
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import json
+
+
+# helper functions for this benchmark example
+def vonKarman_loglaw_wall(yplus):
+    vonKarmanConst = 0.41
+    cplus = 5.5
+    uplus = np.log(yplus) / vonKarmanConst + cplus
+    return uplus
+
+
+def get_dns_data():
+    """
+    Reference: DNS of Turbulent Channel Flow up to Re_tau=590, 1999,
+    Physics of Fluids, vol 11, 943-945.
+    https://turbulence.oden.utexas.edu/data/MKM/chan180/profiles/chan180.means
+    """
+    file_name = "examples/cfd/data/turbulent_channel_dns_data.json"
+    with open(file_name, "r") as file:
+        return json.load(file)
+
+
+class TurbulentChannel3D:
+    def __init__(self, channel_half_width, Re_tau, u_tau, grid_shape, velocity_set, backend, precision_policy):
+        # initialize backend
+        xlb.init(
+            velocity_set=velocity_set,
+            default_backend=backend,
+            default_precision_policy=precision_policy,
+        )
+
+        self.channel_half_width = channel_half_width
+        self.Re_tau = Re_tau
+        self.u_tau = u_tau
+        self.visc = u_tau * channel_half_width / Re_tau
+        self.omega = 1.0 / (3.0 * self.visc + 0.5)
+        # DeltaPlus = Re_tau / channel_half_width
+        # DeltaPlus = u_tau / nu * Delta where u_tau / nu = Re_tau / channel_half_width
+
+        self.grid_shape = grid_shape
+        self.velocity_set = velocity_set
+        self.backend = backend
+        self.precision_policy = 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 = []
+
+        # Setup the simulation BC, its initial conditions, and the stepper
+        self._setup()
+
+    def get_force(self):
+        # define the external force
+        shape = (self.velocity_set.d,)
+        force = np.zeros(shape)
+        force[0] = self.Re_tau**2 * self.visc**2 / self.channel_half_width**3
+        return force
+
+    def _setup(self):
+        self.setup_boundary_conditions()
+        self.setup_boundary_masker()
+        self.initialize_fields()
+        self.setup_stepper()
+
+    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)]
+        return walls
+
+    def setup_boundary_conditions(self):
+        walls = self.define_boundary_indices()
+        bc_walls = RegularizedBC("velocity", (0.0, 0.0, 0.0), indices=walls)
+        self.boundary_conditions = [bc_walls]
+
+    def setup_boundary_masker(self):
+        indices_boundary_masker = IndicesBoundaryMasker(
+            velocity_set=self.velocity_set,
+            precision_policy=self.precision_policy,
+            compute_backend=self.backend,
+        )
+        self.bc_mask, self.missing_mask = indices_boundary_masker(self.boundary_conditions, self.bc_mask, self.missing_mask)
+
+    def initialize_fields(self):
+        shape = (self.velocity_set.d,) + (self.grid_shape)
+        np.random.seed(0)
+        u_init = np.random.random(shape)
+        if self.backend == ComputeBackend.JAX:
+            u_init = jnp.full(shape=shape, fill_value=1e-2 * u_init)
+        else:
+            u_init = wp.array(1e-2 * u_init, dtype=self.precision_policy.compute_precision.wp_dtype)
+        self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend, u=u_init)
+
+    def setup_stepper(self):
+        force = self.get_force()
+        self.stepper = IncompressibleNavierStokesStepper(
+            self.omega, boundary_conditions=self.boundary_conditions, collision_type="KBC", forcing_scheme="exact_difference", force_vector=force
+        )
+
+    def run(self, num_steps, print_interval, post_process_interval=100):
+        start_time = time.time()
+        for i in range(num_steps):
+            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 + 1) % print_interval == 0:
+                elapsed_time = time.time() - start_time
+                print(f"Iteration: {i + 1}/{num_steps} | Time elapsed: {elapsed_time:.2f}s")
+
+            if i % post_process_interval == 0 or i == num_steps - 1:
+                self.post_process(i)
+
+    def post_process(self, i):
+        # Write the results. We'll use JAX backend for the post-processing
+        if not isinstance(self.f_0, jnp.ndarray):
+            f_0 = wp.to_jax(self.f_0)
+        else:
+            f_0 = self.f_0
+
+        macro = Macroscopic(
+            compute_backend=ComputeBackend.JAX,
+            precision_policy=self.precision_policy,
+            velocity_set=xlb.velocity_set.D3Q27(precision_policy=self.precision_policy, backend=ComputeBackend.JAX),
+        )
+
+        rho, u = macro(f_0)
+
+        # compute velocity magnitude
+        u_magnitude = (u[0] ** 2 + u[1] ** 2 + u[2] ** 2) ** 0.5
+        fields = {"rho": rho[0], "u_x": u[0], "u_y": u[1], "u_z": u[2], "u_magnitude": u_magnitude}
+        save_fields_vtk(fields, timestep=i)
+        save_image(fields["u_magnitude"][:, grid_size_y // 2, :], timestep=i)
+
+        # Save monitor plot
+        self.plot_uplus(u, i)
+        return
+
+    def plot_uplus(self, u, timestep):
+        # Compute moving average of drag coefficient, 100, 1000, 10000
+        # mean streamwise velocity in wall units u^+(z)
+        # Wall distance in wall units to be used inside output_data
+        zz = np.arange(self.grid_shape[-1])
+        zz = np.minimum(zz, zz.max() - zz)
+        yplus = zz * self.u_tau / self.visc
+        uplus = np.mean(u[0], axis=(0, 1)) / self.u_tau
+        uplus_loglaw = vonKarman_loglaw_wall(yplus)
+        dns_dic = get_dns_data()
+        plt.clf()
+        plt.semilogx(yplus, uplus, "r.", yplus, uplus_loglaw, "k:", dns_dic["y+"], dns_dic["Umean"], "b-")
+        ax = plt.gca()
+        ax.set_xlim([0.1, 300])
+        ax.set_ylim([0, 20])
+        fname = "uplus_" + str(timestep // 10000).zfill(5) + ".png"
+        plt.savefig(fname, format="png")
+
+
+if __name__ == "__main__":
+    # Problem Configuration
+    # h: channel half-width
+    channel_half_width = 50
+
+    # Define channel geometry based on h
+    grid_size_x = 6 * channel_half_width
+    grid_size_y = 3 * channel_half_width
+    grid_size_z = 2 * channel_half_width
+
+    # Grid parameters
+    grid_shape = (grid_size_x, grid_size_y, grid_size_z)
+
+    # Define flow regime
+    # Set up Reynolds number and deduce relaxation time (omega)
+    Re_tau = 180
+    u_tau = 0.001
+
+    # Runtime & backend configurations
+    backend = ComputeBackend.WARP
+    precision_policy = PrecisionPolicy.FP64FP64
+    velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, backend=backend)
+    num_steps = 10000000
+    print_interval = 100000
+
+    # Print simulation info
+    print("\n" + "=" * 50 + "\n")
+    print("Simulation Configuration:")
+    print(f"Grid size: {grid_size_x} x {grid_size_y} x {grid_size_z}")
+    print(f"Backend: {backend}")
+    print(f"Velocity set: {velocity_set}")
+    print(f"Precision policy: {precision_policy}")
+    print(f"Reynolds number: {Re_tau}")
+    print(f"Max iterations: {num_steps}")
+    print("\n" + "=" * 50 + "\n")
+
+    simulation = TurbulentChannel3D(channel_half_width, Re_tau, u_tau, grid_shape, velocity_set, backend, precision_policy)
+    simulation.run(num_steps, print_interval, post_process_interval=100000)
diff --git a/examples/cfd/windtunnel_3d.py b/examples/cfd/windtunnel_3d.py
index 3522122..92c94aa 100644
--- a/examples/cfd/windtunnel_3d.py
+++ b/examples/cfd/windtunnel_3d.py
@@ -36,7 +36,7 @@ def __init__(self, omega, wind_speed, grid_shape, velocity_set, backend, precisi
         self.velocity_set = velocity_set
         self.backend = backend
         self.precision_policy = precision_policy
-        self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_map = create_nse_fields(grid_shape)
+        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 = []
 
@@ -118,8 +118,8 @@ def setup_boundary_masker(self):
         bc_mesh = self.boundary_conditions[-1]
         dx = self.grid_spacing
         origin, spacing = (0, 0, 0), (dx, dx, dx)
-        self.boundary_map, self.missing_mask = indices_boundary_masker(bclist_other, self.boundary_map, self.missing_mask)
-        self.boundary_map, self.missing_mask = mesh_boundary_masker(bc_mesh, origin, spacing, self.boundary_map, self.missing_mask)
+        self.bc_mask, self.missing_mask = indices_boundary_masker(bclist_other, self.bc_mask, self.missing_mask)
+        self.bc_mask, self.missing_mask = mesh_boundary_masker(bc_mesh, origin, spacing, self.bc_mask, self.missing_mask)
 
     def initialize_fields(self):
         self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend)
@@ -134,7 +134,7 @@ def run(self, num_steps, print_interval, post_process_interval=100):
 
         start_time = time.time()
         for i in range(num_steps):
-            self.f_1 = self.stepper(self.f_0, self.f_1, self.boundary_map, self.missing_mask, i)
+            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 + 1) % print_interval == 0:
@@ -169,7 +169,7 @@ def post_process(self, i):
         save_image(fields["u_magnitude"][:, grid_size_y // 2, :], timestep=i)
 
         # Compute lift and drag
-        boundary_force = self.momentum_transfer(self.f_0, self.boundary_map, self.missing_mask)
+        boundary_force = self.momentum_transfer(self.f_0, self.bc_mask, self.missing_mask)
         drag = np.sqrt(boundary_force[0] ** 2 + boundary_force[1] ** 2)  # xy-plane
         lift = boundary_force[2]
         c_d = 2.0 * drag / (self.wind_speed**2 * self.car_cross_section)
diff --git a/examples/cfd_old_to_be_migrated/flow_past_sphere.py b/examples/cfd_old_to_be_migrated/flow_past_sphere.py
index 7214130..1684266 100644
--- a/examples/cfd_old_to_be_migrated/flow_past_sphere.py
+++ b/examples/cfd_old_to_be_migrated/flow_past_sphere.py
@@ -75,7 +75,7 @@ def warp_implementation(self, rho, u, vel):
     u = grid.create_field(cardinality=velocity_set.d, dtype=wp.float32)
     f0 = grid.create_field(cardinality=velocity_set.q, dtype=wp.float32)
     f1 = grid.create_field(cardinality=velocity_set.q, dtype=wp.float32)
-    boundary_map = grid.create_field(cardinality=1, dtype=wp.uint8)
+    bc_mask = grid.create_field(cardinality=1, dtype=wp.uint8)
     missing_mask = grid.create_field(cardinality=velocity_set.q, dtype=wp.bool)
 
     # Make operators
@@ -154,19 +154,19 @@ def warp_implementation(self, rho, u, vel):
     indices = wp.from_numpy(indices, dtype=wp.int32)
 
     # Set boundary conditions on the indices
-    boundary_map, missing_mask = indices_boundary_masker(indices, half_way_bc.id, boundary_map, missing_mask, (0, 0, 0))
+    bc_mask, missing_mask = indices_boundary_masker(indices, half_way_bc.id, bc_mask, missing_mask, (0, 0, 0))
 
     # Set inlet bc
     lower_bound = (0, 0, 0)
     upper_bound = (0, nr, nr)
     direction = (1, 0, 0)
-    boundary_map, missing_mask = planar_boundary_masker(lower_bound, upper_bound, direction, equilibrium_bc.id, boundary_map, missing_mask, (0, 0, 0))
+    bc_mask, missing_mask = planar_boundary_masker(lower_bound, upper_bound, direction, equilibrium_bc.id, bc_mask, missing_mask, (0, 0, 0))
 
     # Set outlet bc
     lower_bound = (nr - 1, 0, 0)
     upper_bound = (nr - 1, nr, nr)
     direction = (-1, 0, 0)
-    boundary_map, missing_mask = planar_boundary_masker(lower_bound, upper_bound, direction, do_nothing_bc.id, boundary_map, missing_mask, (0, 0, 0))
+    bc_mask, missing_mask = planar_boundary_masker(lower_bound, upper_bound, direction, do_nothing_bc.id, bc_mask, missing_mask, (0, 0, 0))
 
     # Set initial conditions
     rho, u = initializer(rho, u, vel)
@@ -181,7 +181,7 @@ def warp_implementation(self, rho, u, vel):
     num_steps = 1024 * 8
     start = time.time()
     for _ in tqdm(range(num_steps)):
-        f1 = stepper(f0, f1, boundary_map, missing_mask, _)
+        f1 = stepper(f0, f1, bc_mask, missing_mask, _)
         f1, f0 = f0, f1
         if (_ % plot_freq == 0) and (not compute_mlup):
             rho, u = macroscopic(f0, rho, u)
@@ -191,7 +191,7 @@ def warp_implementation(self, rho, u, vel):
             plt.imshow(u[0, :, nr // 2, :].numpy())
             plt.colorbar()
             plt.subplot(1, 2, 2)
-            plt.imshow(boundary_map[0, :, nr // 2, :].numpy())
+            plt.imshow(bc_mask[0, :, nr // 2, :].numpy())
             plt.colorbar()
             plt.savefig(f"{save_dir}/{str(_).zfill(6)}.png")
             plt.close()
diff --git a/examples/cfd_old_to_be_migrated/taylor_green.py b/examples/cfd_old_to_be_migrated/taylor_green.py
index c5b40b7..846ba30 100644
--- a/examples/cfd_old_to_be_migrated/taylor_green.py
+++ b/examples/cfd_old_to_be_migrated/taylor_green.py
@@ -113,7 +113,7 @@ def run_taylor_green(backend, compute_mlup=True):
     u = grid.create_field(cardinality=velocity_set.d, precision=xlb.Precision.FP32)
     f0 = grid.create_field(cardinality=velocity_set.q, precision=xlb.Precision.FP32)
     f1 = grid.create_field(cardinality=velocity_set.q, precision=xlb.Precision.FP32)
-    boundary_map = grid.create_field(cardinality=1, precision=xlb.Precision.UINT8)
+    bc_mask = grid.create_field(cardinality=1, precision=xlb.Precision.UINT8)
     missing_mask = grid.create_field(cardinality=velocity_set.q, precision=xlb.Precision.BOOL)
 
     # Make operators
@@ -149,10 +149,10 @@ def run_taylor_green(backend, compute_mlup=True):
     for _ in tqdm(range(num_steps)):
         # Time step
         if backend == "warp":
-            f1 = stepper(f0, f1, boundary_map, missing_mask, _)
+            f1 = stepper(f0, f1, bc_mask, missing_mask, _)
             f1, f0 = f0, f1
         elif backend == "jax":
-            f0 = stepper(f0, boundary_map, missing_mask, _)
+            f0 = stepper(f0, bc_mask, missing_mask, _)
 
         # Plot if needed
         if (_ % plot_freq == 0) and (not compute_mlup):
diff --git a/examples/performance/mlups_3d.py b/examples/performance/mlups_3d.py
index 602e741..32c3d6d 100644
--- a/examples/performance/mlups_3d.py
+++ b/examples/performance/mlups_3d.py
@@ -32,7 +32,7 @@ def setup_simulation(args):
         raise ValueError("Invalid precision")
 
     xlb.init(
-        velocity_set=xlb.velocity_set.D3Q19(),
+        velocity_set=xlb.velocity_set.D3Q19(precision_policy=precision_policy, backend=backend),
         default_backend=backend,
         default_precision_policy=precision_policy,
     )
@@ -42,9 +42,9 @@ def setup_simulation(args):
 
 def create_grid_and_fields(cube_edge):
     grid_shape = (cube_edge, cube_edge, cube_edge)
-    grid, f_0, f_1, missing_mask, boundary_map = create_nse_fields(grid_shape)
+    grid, f_0, f_1, missing_mask, bc_mask = create_nse_fields(grid_shape)
 
-    return grid, f_0, f_1, missing_mask, boundary_map
+    return grid, f_0, f_1, missing_mask, bc_mask
 
 
 def define_boundary_indices(grid):
@@ -55,7 +55,7 @@ def define_boundary_indices(grid):
         + grid.boundingBoxIndices["right"][i]
         + grid.boundingBoxIndices["front"][i]
         + grid.boundingBoxIndices["back"][i]
-        for i in range(xlb.velocity_set.D3Q19().d)
+        for i in range(len(grid.shape))
     ]
     return lid, walls
 
@@ -67,7 +67,7 @@ def setup_boundary_conditions(grid):
     return [bc_top, bc_walls]
 
 
-def run(f_0, f_1, backend, grid, boundary_map, missing_mask, num_steps):
+def run(f_0, f_1, backend, precision_policy, grid, bc_mask, missing_mask, num_steps):
     omega = 1.0
     stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=setup_boundary_conditions(grid))
 
@@ -75,13 +75,13 @@ def run(f_0, f_1, backend, grid, boundary_map, missing_mask, num_steps):
         stepper = distribute(
             stepper,
             grid,
-            xlb.velocity_set.D3Q19(),
+            xlb.velocity_set.D3Q19(precision_policy=precision_policy, backend=backend),
         )
 
     start_time = time.time()
 
     for i in range(num_steps):
-        f_1 = stepper(f_0, f_1, boundary_map, missing_mask, i)
+        f_1 = stepper(f_0, f_1, bc_mask, missing_mask, i)
         f_0, f_1 = f_1, f_0
     wp.synchronize()
 
@@ -98,10 +98,11 @@ def calculate_mlups(cube_edge, num_steps, elapsed_time):
 def main():
     args = parse_arguments()
     backend, precision_policy = setup_simulation(args)
-    grid, f_0, f_1, missing_mask, boundary_map = create_grid_and_fields(args.cube_edge)
-    f_0 = initialize_eq(f_0, grid, xlb.velocity_set.D3Q19(), backend)
+    velocity_set = xlb.velocity_set.D3Q19(precision_policy=precision_policy, backend=backend)
+    grid, f_0, f_1, missing_mask, bc_mask = create_grid_and_fields(args.cube_edge)
+    f_0 = initialize_eq(f_0, grid, velocity_set, backend)
 
-    elapsed_time = run(f_0, f_1, backend, grid, boundary_map, missing_mask, args.num_steps)
+    elapsed_time = run(f_0, f_1, backend, precision_policy, grid, bc_mask, missing_mask, args.num_steps)
     mlups = calculate_mlups(args.cube_edge, args.num_steps, elapsed_time)
 
     print(f"Simulation completed in {elapsed_time:.2f} seconds")
diff --git a/xlb/grid/grid.py b/xlb/grid/grid.py
index 7d8a678..7494c3e 100644
--- a/xlb/grid/grid.py
+++ b/xlb/grid/grid.py
@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from typing import Tuple
+import numpy as np
 
 from xlb import DefaultConfig
 from xlb.compute_backend import ComputeBackend
@@ -24,52 +25,63 @@ def __init__(self, shape: Tuple[int, ...], compute_backend: ComputeBackend):
         self.shape = shape
         self.dim = len(shape)
         self.compute_backend = compute_backend
-        self._bounding_box_indices()
+        self.boundingBoxIndices = self.bounding_box_indices()
         self._initialize_backend()
 
     @abstractmethod
     def _initialize_backend(self):
         pass
 
-    def _bounding_box_indices(self):
+    def bounding_box_indices(self, remove_edges=False):
         """
         This function calculates the indices of the bounding box of a 2D or 3D grid.
         The bounding box is defined as the set of grid points on the outer edge of the grid.
 
+        Parameters
+        ----------
+        remove_edges : bool, optional
+            If True, the nodes along the edges (not just the corners) are removed from the bounding box indices.
+            Default is False.
+
         Returns
         -------
-            boundingBox (dict): A dictionary where keys are the names of the bounding box faces
-            ("bottom", "top", "left", "right" for 2D; additional "front", "back" for 3D), and values
-            are numpy arrays of indices corresponding to each face.
+        boundingBox (dict): A dictionary where keys are the names of the bounding box faces
+        ("bottom", "top", "left", "right" for 2D; additional "front", "back" for 3D), and values
+        are numpy arrays of indices corresponding to each face.
         """
 
-        def to_tuple(lst):
-            d = len(lst[0])
-            return [tuple([sublist[i] for sublist in lst]) for i in range(d)]
+        # Get the shape of the grid
+        origin = np.array([0, 0, 0])
+        bounds = np.array(self.shape)
+        if remove_edges:
+            origin += 1
+            bounds -= 1
+        slice_x = slice(origin[0], bounds[0])
+        slice_y = slice(origin[1], bounds[1])
+        dim = len(bounds)
+
+        # Generate bounding box indices for each face
+        grid = np.indices(self.shape)
+        boundingBoxIndices = {}
 
-        if self.dim == 2:
-            # For a 2D grid, the bounding box consists of four edges: bottom, top, left, and right.
-            # Each edge is represented as an array of indices. For example, the bottom edge includes
-            # all points where the y-coordinate is 0, so its indices are [[i, 0] for i in range(nx)].
+        if dim == 2:
             nx, ny = self.shape
-            self.boundingBoxIndices = {
-                "bottom": to_tuple([[i, 0] for i in range(nx)]),
-                "top": to_tuple([[i, ny - 1] for i in range(nx)]),
-                "left": to_tuple([[0, i] for i in range(ny)]),
-                "right": to_tuple([[nx - 1, i] for i in range(ny)]),
+            boundingBoxIndices = {
+                "bottom": grid[:, slice_x, 0],
+                "top": grid[:, slice_x, ny - 1],
+                "left": grid[:, 0, slice_y],
+                "right": grid[:, nx - 1, slice_y],
             }
-
-        elif self.dim == 3:
-            # For a 3D grid, the bounding box consists of six faces: bottom, top, left, right, front, and back.
-            # Each face is represented as an array of indices. For example, the bottom face includes all points
-            # where the z-coordinate is 0, so its indices are [[i, j, 0] for i in range(nx) for j in range(ny)].
+        elif dim == 3:
             nx, ny, nz = self.shape
-            self.boundingBoxIndices = {
-                "bottom": to_tuple([[i, j, 0] for i in range(nx) for j in range(ny)]),
-                "top": to_tuple([[i, j, nz - 1] for i in range(nx) for j in range(ny)]),
-                "left": to_tuple([[0, j, k] for j in range(ny) for k in range(nz)]),
-                "right": to_tuple([[nx - 1, j, k] for j in range(ny) for k in range(nz)]),
-                "front": to_tuple([[i, 0, k] for i in range(nx) for k in range(nz)]),
-                "back": to_tuple([[i, ny - 1, k] for i in range(nx) for k in range(nz)]),
+            slice_z = slice(origin[2], bounds[2])
+            boundingBoxIndices = {
+                "bottom": grid[:, slice_x, slice_y, 0].reshape(3, -1),
+                "top": grid[:, slice_x, slice_y, nz - 1].reshape(3, -1),
+                "left": grid[:, 0, slice_y, slice_z].reshape(3, -1),
+                "right": grid[:, nx - 1, slice_y, slice_z].reshape(3, -1),
+                "front": grid[:, slice_x, 0, slice_z].reshape(3, -1),
+                "back": grid[:, slice_x, ny - 1, slice_z].reshape(3, -1),
             }
-        return
+
+        return {k: v.tolist() for k, v in boundingBoxIndices.items()}
diff --git a/xlb/helper/initializers.py b/xlb/helper/initializers.py
index aff3ee8..c8439d9 100644
--- a/xlb/helper/initializers.py
+++ b/xlb/helper/initializers.py
@@ -3,8 +3,10 @@
 
 
 def initialize_eq(f, grid, velocity_set, backend, rho=None, u=None):
-    rho = rho or grid.create_field(cardinality=1, fill_value=1.0)
-    u = u or grid.create_field(cardinality=velocity_set.d, fill_value=0.0)
+    if rho is None:
+        rho = grid.create_field(cardinality=1, fill_value=1.0)
+    if u is None:
+        u = grid.create_field(cardinality=velocity_set.d, fill_value=0.0)
     equilibrium = QuadraticEquilibrium()
 
     if backend == ComputeBackend.JAX:
diff --git a/xlb/helper/nse_solver.py b/xlb/helper/nse_solver.py
index 96befa6..3d0ef3e 100644
--- a/xlb/helper/nse_solver.py
+++ b/xlb/helper/nse_solver.py
@@ -14,6 +14,6 @@ def create_nse_fields(grid_shape: Tuple[int, int, int], velocity_set=None, compu
     f_0 = grid.create_field(cardinality=velocity_set.q, dtype=precision_policy.store_precision)
     f_1 = grid.create_field(cardinality=velocity_set.q, dtype=precision_policy.store_precision)
     missing_mask = grid.create_field(cardinality=velocity_set.q, dtype=Precision.BOOL)
-    boundary_map = grid.create_field(cardinality=1, dtype=Precision.UINT8)
+    bc_mask = grid.create_field(cardinality=1, dtype=Precision.UINT8)
 
-    return grid, f_0, f_1, missing_mask, boundary_map
+    return grid, f_0, f_1, missing_mask, bc_mask
diff --git a/xlb/operator/boundary_condition/bc_do_nothing.py b/xlb/operator/boundary_condition/bc_do_nothing.py
index 6e8d317..dcdc8fd 100644
--- a/xlb/operator/boundary_condition/bc_do_nothing.py
+++ b/xlb/operator/boundary_condition/bc_do_nothing.py
@@ -48,8 +48,8 @@ def __init__(
 
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def jax_implementation(self, f_pre, f_post, boundary_map, missing_mask):
-        boundary = boundary_map == self.id
+    def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask):
+        boundary = bc_mask == self.id
         return jnp.where(boundary, f_pre, f_post)
 
     def _construct_warp(self):
@@ -67,7 +67,7 @@ def functional(
         def kernel2d(
             f_pre: wp.array3d(dtype=Any),
             f_post: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.uint8),
         ):
             # Get the global index
@@ -75,10 +75,10 @@ def kernel2d(
             index = wp.vec2i(i, j)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_2d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_2d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(DoNothingBC.id):
+            if _boundary_id == wp.uint8(DoNothingBC.id):
                 _f_aux = _f_post
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -93,7 +93,7 @@ def kernel2d(
         def kernel3d(
             f_pre: wp.array4d(dtype=Any),
             f_post: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
         ):
             # Get the global index
@@ -101,10 +101,10 @@ def kernel3d(
             index = wp.vec3i(i, j, k)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_3d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_3d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(DoNothingBC.id):
+            if _boundary_id == wp.uint8(DoNothingBC.id):
                 _f_aux = _f_post
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -119,11 +119,11 @@ def kernel3d(
         return functional, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f_pre, f_post, boundary_map, missing_mask):
+    def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
         # Launch the warp kernel
         wp.launch(
             self.warp_kernel,
-            inputs=[f_pre, f_post, boundary_map, missing_mask],
+            inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
         return f_post
diff --git a/xlb/operator/boundary_condition/bc_equilibrium.py b/xlb/operator/boundary_condition/bc_equilibrium.py
index c018a60..716fd8e 100644
--- a/xlb/operator/boundary_condition/bc_equilibrium.py
+++ b/xlb/operator/boundary_condition/bc_equilibrium.py
@@ -62,11 +62,11 @@ def __init__(
 
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def jax_implementation(self, f_pre, f_post, boundary_map, missing_mask):
+    def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask):
         feq = self.equilibrium_operator(jnp.array([self.rho]), jnp.array(self.u))
         new_shape = feq.shape + (1,) * self.velocity_set.d
         feq = lax.broadcast_in_dim(feq, new_shape, [0])
-        boundary = boundary_map == self.id
+        boundary = bc_mask == self.id
 
         return jnp.where(boundary, feq, f_post)
 
@@ -92,7 +92,7 @@ def functional(
         def kernel2d(
             f_pre: wp.array3d(dtype=Any),
             f_post: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
         ):
             # Get the global index
@@ -100,10 +100,10 @@ def kernel2d(
             index = wp.vec2i(i, j)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_2d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_2d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(EquilibriumBC.id):
+            if _boundary_id == wp.uint8(EquilibriumBC.id):
                 _f_aux = _f_post
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -118,7 +118,7 @@ def kernel2d(
         def kernel3d(
             f_pre: wp.array4d(dtype=Any),
             f_post: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
         ):
             # Get the global index
@@ -126,10 +126,10 @@ def kernel3d(
             index = wp.vec3i(i, j, k)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_3d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_3d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(EquilibriumBC.id):
+            if _boundary_id == wp.uint8(EquilibriumBC.id):
                 _f_aux = _f_post
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -144,11 +144,11 @@ def kernel3d(
         return functional, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f_pre, f_post, boundary_map, missing_mask):
+    def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
         # Launch the warp kernel
         wp.launch(
             self.warp_kernel,
-            inputs=[f_pre, f_post, boundary_map, missing_mask],
+            inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
         return f_post
diff --git a/xlb/operator/boundary_condition/bc_extrapolation_outflow.py b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py
index 7bd447e..87c6850 100644
--- a/xlb/operator/boundary_condition/bc_extrapolation_outflow.py
+++ b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py
@@ -94,13 +94,13 @@ def _roll(self, fld, vec):
             return jnp.roll(fld, (vec[0], vec[1], vec[2]), axis=(1, 2, 3))
 
     @partial(jit, static_argnums=(0,), inline=True)
-    def prepare_bc_auxilary_data(self, f_pre, f_post, boundary_map, missing_mask):
+    def prepare_bc_auxilary_data(self, f_pre, f_post, bc_mask, missing_mask):
         """
         Prepare the auxilary distribution functions for the boundary condition.
         Since this function is called post-collisiotn: f_pre = f_post_stream and f_post = f_post_collision
         """
         sound_speed = 1.0 / jnp.sqrt(3.0)
-        boundary = boundary_map == self.id
+        boundary = bc_mask == self.id
         new_shape = (self.velocity_set.q,) + boundary.shape[1:]
         boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1)))
 
@@ -123,8 +123,8 @@ def prepare_bc_auxilary_data(self, f_pre, f_post, boundary_map, missing_mask):
 
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def apply_jax(self, f_pre, f_post, boundary_map, missing_mask):
-        boundary = boundary_map == self.id
+    def apply_jax(self, f_pre, f_post, bc_mask, missing_mask):
+        boundary = bc_mask == self.id
         new_shape = (self.velocity_set.q,) + boundary.shape[1:]
         boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1)))
         return jnp.where(
@@ -135,7 +135,7 @@ def apply_jax(self, f_pre, f_post, boundary_map, missing_mask):
 
     def _construct_warp(self):
         # Set local constants
-        sound_speed = 1.0 / wp.sqrt(3.0)
+        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
@@ -187,7 +187,7 @@ def prepare_bc_auxilary_data(
             _f = f_post
             for l in range(self.velocity_set.q):
                 if missing_mask[l] == wp.uint8(1):
-                    _f[_opp_indices[l]] = (1.0 - sound_speed) * f_pre[l] + sound_speed * f_aux[l]
+                    _f[_opp_indices[l]] = (self.compute_dtype(1.0) - sound_speed) * f_pre[l] + sound_speed * f_aux[l]
             return _f
 
         # Construct the warp kernel
@@ -195,7 +195,7 @@ def prepare_bc_auxilary_data(
         def kernel2d(
             f_pre: wp.array3d(dtype=Any),
             f_post: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
         ):
             # Get the global index
@@ -203,11 +203,11 @@ def kernel2d(
             index = wp.vec2i(i, j)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_2d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_2d(f_pre, f_post, bc_mask, missing_mask, index)
             _f_aux = _f_vec()
 
             # special preparation of auxiliary data
-            if _boundary_map == wp.uint8(ExtrapolationOutflowBC.id):
+            if _boundary_id == wp.uint8(ExtrapolationOutflowBC.id):
                 nv = get_normal_vectors_2d(_missing_mask)
                 for l in range(self.velocity_set.q):
                     if _missing_mask[l] == wp.uint8(1):
@@ -220,7 +220,7 @@ def kernel2d(
                         _f_aux[l] = _f_pre[l, pull_index[0], pull_index[1]]
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(ExtrapolationOutflowBC.id):
+            if _boundary_id == wp.uint8(ExtrapolationOutflowBC.id):
                 # TODO: is there any way for this BC to have a meaningful kernel given that it has two steps after both
                 # collision and streaming?
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
@@ -236,7 +236,7 @@ def kernel2d(
         def kernel3d(
             f_pre: wp.array4d(dtype=Any),
             f_post: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
         ):
             # Get the global index
@@ -244,11 +244,11 @@ def kernel3d(
             index = wp.vec3i(i, j, k)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_3d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_3d(f_pre, f_post, bc_mask, missing_mask, index)
             _f_aux = _f_vec()
 
             # special preparation of auxiliary data
-            if _boundary_map == wp.uint8(ExtrapolationOutflowBC.id):
+            if _boundary_id == wp.uint8(ExtrapolationOutflowBC.id):
                 nv = get_normal_vectors_3d(_missing_mask)
                 for l in range(self.velocity_set.q):
                     if _missing_mask[l] == wp.uint8(1):
@@ -261,7 +261,7 @@ def kernel3d(
                         _f_aux[l] = _f_pre[l, pull_index[0], pull_index[1], pull_index[2]]
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(ExtrapolationOutflowBC.id):
+            if _boundary_id == wp.uint8(ExtrapolationOutflowBC.id):
                 # TODO: is there any way for this BC to have a meaninful kernel given that it has two steps after both
                 # collision and streaming?
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
@@ -277,11 +277,11 @@ def kernel3d(
         return (functional, prepare_bc_auxilary_data), kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f_pre, f_post, boundary_map, missing_mask):
+    def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
         # Launch the warp kernel
         wp.launch(
             self.warp_kernel,
-            inputs=[f_pre, f_post, boundary_map, missing_mask],
+            inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
         return f_post
diff --git a/xlb/operator/boundary_condition/bc_fullway_bounce_back.py b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
index 729dbb4..85bffb1 100644
--- a/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
+++ b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
@@ -48,8 +48,8 @@ def __init__(
 
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def apply_jax(self, f_pre, f_post, boundary_map, missing_mask):
-        boundary = boundary_map == self.id
+    def apply_jax(self, f_pre, f_post, bc_mask, missing_mask):
+        boundary = bc_mask == self.id
         new_shape = (self.velocity_set.q,) + boundary.shape[1:]
         boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1)))
         return jnp.where(boundary, f_pre[self.velocity_set.opp_indices, ...], f_post)
@@ -77,17 +77,17 @@ def functional(
         def kernel2d(
             f_pre: wp.array3d(dtype=Any),
             f_post: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
         ):  # Get the global index
             i, j = wp.tid()
             index = wp.vec2i(i, j)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_2d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_2d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Check if the boundary is active
-            if _boundary_map == wp.uint8(FullwayBounceBackBC.id):
+            if _boundary_id == wp.uint8(FullwayBounceBackBC.id):
                 _f_aux = _f_vec()
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -102,7 +102,7 @@ def kernel2d(
         def kernel3d(
             f_pre: wp.array4d(dtype=Any),
             f_post: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
         ):
             # Get the global index
@@ -110,10 +110,10 @@ def kernel3d(
             index = wp.vec3i(i, j, k)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_3d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_3d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Check if the boundary is active
-            if _boundary_map == wp.uint8(FullwayBounceBackBC.id):
+            if _boundary_id == wp.uint8(FullwayBounceBackBC.id):
                 _f_aux = _f_vec()
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -128,11 +128,11 @@ def kernel3d(
         return functional, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f_pre, f_post, boundary_map, missing_mask):
+    def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
         # Launch the warp kernel
         wp.launch(
             self.warp_kernel,
-            inputs=[f_pre, f_post, boundary_map, missing_mask],
+            inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
         return f_post
diff --git a/xlb/operator/boundary_condition/bc_halfway_bounce_back.py b/xlb/operator/boundary_condition/bc_halfway_bounce_back.py
index 2cb60fa..c2483c9 100644
--- a/xlb/operator/boundary_condition/bc_halfway_bounce_back.py
+++ b/xlb/operator/boundary_condition/bc_halfway_bounce_back.py
@@ -49,10 +49,13 @@ def __init__(
             mesh_vertices,
         )
 
+        # This BC needs padding for finding missing directions when imposed on a geometry that is in the domain interior
+        self.needs_padding = True
+
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def apply_jax(self, f_pre, f_post, boundary_map, missing_mask):
-        boundary = boundary_map == self.id
+    def apply_jax(self, f_pre, f_post, bc_mask, missing_mask):
+        boundary = bc_mask == self.id
         new_shape = (self.velocity_set.q,) + boundary.shape[1:]
         boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1)))
         return jnp.where(
@@ -88,7 +91,7 @@ def functional(
         def kernel2d(
             f_pre: wp.array3d(dtype=Any),
             f_post: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
         ):
             # Get the global index
@@ -96,10 +99,10 @@ def kernel2d(
             index = wp.vec2i(i, j)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_2d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_2d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(HalfwayBounceBackBC.id):
+            if _boundary_id == wp.uint8(HalfwayBounceBackBC.id):
                 _f_aux = _f_post
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -114,7 +117,7 @@ def kernel2d(
         def kernel3d(
             f_pre: wp.array4d(dtype=Any),
             f_post: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
         ):
             # Get the global index
@@ -122,10 +125,10 @@ def kernel3d(
             index = wp.vec3i(i, j, k)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_3d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_3d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(HalfwayBounceBackBC.id):
+            if _boundary_id == wp.uint8(HalfwayBounceBackBC.id):
                 _f_aux = _f_post
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -140,11 +143,11 @@ def kernel3d(
         return functional, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f_pre, f_post, boundary_map, missing_mask):
+    def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
         # Launch the warp kernel
         wp.launch(
             self.warp_kernel,
-            inputs=[f_pre, f_post, boundary_map, missing_mask],
+            inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
         return f_post
diff --git a/xlb/operator/boundary_condition/bc_regularized.py b/xlb/operator/boundary_condition/bc_regularized.py
index 9734368..7879137 100644
--- a/xlb/operator/boundary_condition/bc_regularized.py
+++ b/xlb/operator/boundary_condition/bc_regularized.py
@@ -105,9 +105,9 @@ def regularize_fpop(self, fpop, feq):
 
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def apply_jax(self, f_pre, f_post, boundary_map, missing_mask):
+    def apply_jax(self, f_pre, f_post, bc_mask, missing_mask):
         # creat a mask to slice boundary cells
-        boundary = boundary_map == self.id
+        boundary = bc_mask == self.id
         new_shape = (self.velocity_set.q,) + boundary.shape[1:]
         boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1)))
 
@@ -154,7 +154,7 @@ def _get_fsum(
             fsum_middle = self.compute_dtype(0.0)
             for l in range(_q):
                 if missing_mask[_opp_indices[l]] == wp.uint8(1):
-                    fsum_known += 2.0 * fpop[l]
+                    fsum_known += self.compute_dtype(2.0) * fpop[l]
                 elif missing_mask[l] != wp.uint8(1):
                     fsum_middle += fpop[l]
             return fsum_known + fsum_middle
@@ -202,13 +202,13 @@ def regularize_fpop(
             nt = _d * (_d + 1) // 2
             QiPi1 = _f_vec()
             for l in range(_q):
-                QiPi1[l] = 0.0
+                QiPi1[l] = self.compute_dtype(0.0)
                 for t in range(nt):
                     QiPi1[l] += _qi[l, t] * PiNeq[t]
 
                 # assign all populations based on eq 45 of Latt et al (2008)
                 # fneq ~ f^1
-                fpop1 = 9.0 / 2.0 * _w[l] * QiPi1[l]
+                fpop1 = self.compute_dtype(4.5) * _w[l] * QiPi1[l]
                 fpop[l] = feq[l] + fpop1
             return fpop
 
@@ -230,7 +230,7 @@ def functional3d_velocity(
             unormal = self.compute_dtype(0.0)
             for d in range(_d):
                 unormal += _u[d] * normals[d]
-            _rho = fsum / (1.0 + unormal)
+            _rho = fsum / (self.compute_dtype(1.0) + unormal)
 
             # impose non-equilibrium bounceback
             feq = self.equilibrium_operator.warp_functional(_rho, _u)
@@ -255,7 +255,7 @@ def functional3d_pressure(
 
             # calculate velocity
             fsum = _get_fsum(_f, missing_mask)
-            unormal = -1.0 + fsum / _rho
+            unormal = -self.compute_dtype(1.0) + fsum / _rho
             _u = unormal * normals
 
             # impose non-equilibrium bounceback
@@ -284,7 +284,7 @@ def functional2d_velocity(
             unormal = self.compute_dtype(0.0)
             for d in range(_d):
                 unormal += _u[d] * normals[d]
-            _rho = fsum / (1.0 + unormal)
+            _rho = fsum / (self.compute_dtype(1.0) + unormal)
 
             # impose non-equilibrium bounceback
             feq = self.equilibrium_operator.warp_functional(_rho, _u)
@@ -309,7 +309,7 @@ def functional2d_pressure(
 
             # calculate velocity
             fsum = _get_fsum(_f, missing_mask)
-            unormal = -1.0 + fsum / _rho
+            unormal = -self.compute_dtype(1.0) + fsum / _rho
             _u = unormal * normals
 
             # impose non-equilibrium bounceback
@@ -325,7 +325,7 @@ def functional2d_pressure(
         def kernel2d(
             f_pre: wp.array3d(dtype=Any),
             f_post: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
         ):
             # Get the global index
@@ -333,10 +333,10 @@ def kernel2d(
             index = wp.vec2i(i, j)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_2d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_2d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(self.id):
+            if _boundary_id == wp.uint8(self.id):
                 _f_aux = _f_vec()
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -351,7 +351,7 @@ def kernel2d(
         def kernel3d(
             f_pre: wp.array4d(dtype=Any),
             f_post: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
         ):
             # Get the global index
@@ -359,10 +359,10 @@ def kernel3d(
             index = wp.vec3i(i, j, k)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_3d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_3d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(self.id):
+            if _boundary_id == wp.uint8(self.id):
                 _f_aux = _f_vec()
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -385,11 +385,11 @@ def kernel3d(
         return functional, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f_pre, f_post, boundary_map, missing_mask):
+    def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
         # Launch the warp kernel
         wp.launch(
             self.warp_kernel,
-            inputs=[f_pre, f_post, boundary_map, missing_mask],
+            inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
         return f_post
diff --git a/xlb/operator/boundary_condition/bc_zouhe.py b/xlb/operator/boundary_condition/bc_zouhe.py
index 782eb4c..0bf68b8 100644
--- a/xlb/operator/boundary_condition/bc_zouhe.py
+++ b/xlb/operator/boundary_condition/bc_zouhe.py
@@ -69,6 +69,9 @@ def __init__(
             self.prescribed_value = jnp.atleast_1d(prescribed_value)[(slice(None),) + (None,) * dim]
             # TODO: this won't work if the prescribed values are a profile with the length of bdry indices!
 
+        # This BC needs padding for finding missing directions when imposed on a geometry that is in the domain interior
+        self.needs_padding = True
+
     @partial(jit, static_argnums=(0,), inline=True)
     def _get_known_middle_mask(self, missing_mask):
         known_mask = missing_mask[self.velocity_set.opp_indices]
@@ -156,9 +159,9 @@ def bounceback_nonequilibrium(self, fpop, feq, missing_mask):
 
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def apply_jax(self, f_pre, f_post, boundary_map, missing_mask):
+    def apply_jax(self, f_pre, f_post, bc_mask, missing_mask):
         # creat a mask to slice boundary cells
-        boundary = boundary_map == self.id
+        boundary = bc_mask == self.id
         new_shape = (self.velocity_set.q,) + boundary.shape[1:]
         boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1)))
 
@@ -205,7 +208,7 @@ def _get_fsum(
             fsum_middle = self.compute_dtype(0.0)
             for l in range(_q):
                 if missing_mask[_opp_indices[l]] == wp.uint8(1):
-                    fsum_known += 2.0 * fpop[l]
+                    fsum_known += self.compute_dtype(2.0) * fpop[l]
                 elif missing_mask[l] != wp.uint8(1):
                     fsum_middle += fpop[l]
             return fsum_known + fsum_middle
@@ -247,7 +250,7 @@ def functional3d_velocity(
             unormal = self.compute_dtype(0.0)
             for d in range(_d):
                 unormal += _u[d] * normals[d]
-            _rho = fsum / (1.0 + unormal)
+            _rho = fsum / (self.compute_dtype(1.0) + unormal)
 
             # impose non-equilibrium bounceback
             feq = self.equilibrium_operator.warp_functional(_rho, _u)
@@ -269,7 +272,7 @@ def functional3d_pressure(
 
             # calculate velocity
             fsum = _get_fsum(_f, missing_mask)
-            unormal = -1.0 + fsum / _rho
+            unormal = -self.compute_dtype(1.0) + fsum / _rho
             _u = unormal * normals
 
             # impose non-equilibrium bounceback
@@ -295,7 +298,7 @@ def functional2d_velocity(
             unormal = self.compute_dtype(0.0)
             for d in range(_d):
                 unormal += _u[d] * normals[d]
-            _rho = fsum / (1.0 + unormal)
+            _rho = fsum / (self.compute_dtype(1.0) + unormal)
 
             # impose non-equilibrium bounceback
             feq = self.equilibrium_operator.warp_functional(_rho, _u)
@@ -317,7 +320,7 @@ def functional2d_pressure(
 
             # calculate velocity
             fsum = _get_fsum(_f, missing_mask)
-            unormal = -1.0 + fsum / _rho
+            unormal = -self.compute_dtype(1.0) + fsum / _rho
             _u = unormal * normals
 
             # impose non-equilibrium bounceback
@@ -330,7 +333,7 @@ def functional2d_pressure(
         def kernel2d(
             f_pre: wp.array3d(dtype=Any),
             f_post: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
         ):
             # Get the global index
@@ -338,10 +341,10 @@ def kernel2d(
             index = wp.vec2i(i, j)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_2d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_2d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(self.id):
+            if _boundary_id == wp.uint8(self.id):
                 _f_aux = _f_post
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -356,7 +359,7 @@ def kernel2d(
         def kernel3d(
             f_pre: wp.array4d(dtype=Any),
             f_post: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
         ):
             # Get the global index
@@ -364,10 +367,10 @@ def kernel3d(
             index = wp.vec3i(i, j, k)
 
             # read tid data
-            _f_pre, _f_post, _boundary_map, _missing_mask = self._get_thread_data_3d(f_pre, f_post, boundary_map, missing_mask, index)
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data_3d(f_pre, f_post, bc_mask, missing_mask, index)
 
             # Apply the boundary condition
-            if _boundary_map == wp.uint8(self.id):
+            if _boundary_id == wp.uint8(self.id):
                 _f_aux = _f_post
                 _f = functional(_f_pre, _f_post, _f_aux, _missing_mask)
             else:
@@ -390,11 +393,11 @@ def kernel3d(
         return functional, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f_pre, f_post, boundary_map, missing_mask):
+    def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
         # Launch the warp kernel
         wp.launch(
             self.warp_kernel,
-            inputs=[f_pre, f_post, boundary_map, missing_mask],
+            inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
         return f_post
diff --git a/xlb/operator/boundary_condition/boundary_condition.py b/xlb/operator/boundary_condition/boundary_condition.py
index 2cf6d67..8d47929 100644
--- a/xlb/operator/boundary_condition/boundary_condition.py
+++ b/xlb/operator/boundary_condition/boundary_condition.py
@@ -48,6 +48,10 @@ def __init__(
         # Set the implementation step
         self.implementation_step = implementation_step
 
+        # A flag to indicate whether bc indices need to be padded in both normal directions to identify missing directions
+        # when inside/outside of the geoemtry is not known
+        self.needs_padding = False
+
         if self.compute_backend == ComputeBackend.WARP:
             # Set local constants TODO: This is a hack and should be fixed with warp update
             _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
@@ -66,14 +70,14 @@ def prepare_bc_auxilary_data(
         def _get_thread_data_2d(
             f_pre: wp.array3d(dtype=Any),
             f_post: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
             index: wp.vec2i,
         ):
             # Get the boundary id and missing mask
             _f_pre = _f_vec()
             _f_post = _f_vec()
-            _boundary_map = boundary_map[0, index[0], index[1]]
+            _boundary_id = bc_mask[0, index[0], index[1]]
             _missing_mask = _missing_mask_vec()
             for l in range(self.velocity_set.q):
                 # q-sized vector of populations
@@ -85,20 +89,20 @@ def _get_thread_data_2d(
                     _missing_mask[l] = wp.uint8(1)
                 else:
                     _missing_mask[l] = wp.uint8(0)
-            return _f_pre, _f_post, _boundary_map, _missing_mask
+            return _f_pre, _f_post, _boundary_id, _missing_mask
 
         @wp.func
         def _get_thread_data_3d(
             f_pre: wp.array4d(dtype=Any),
             f_post: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
             index: wp.vec3i,
         ):
             # Get the boundary id and missing mask
             _f_pre = _f_vec()
             _f_post = _f_vec()
-            _boundary_map = boundary_map[0, index[0], index[1], index[2]]
+            _boundary_id = bc_mask[0, index[0], index[1], index[2]]
             _missing_mask = _missing_mask_vec()
             for l in range(self.velocity_set.q):
                 # q-sized vector of populations
@@ -110,7 +114,7 @@ def _get_thread_data_3d(
                     _missing_mask[l] = wp.uint8(1)
                 else:
                     _missing_mask[l] = wp.uint8(0)
-            return _f_pre, _f_post, _boundary_map, _missing_mask
+            return _f_pre, _f_post, _boundary_id, _missing_mask
 
         # Construct some helper warp functions for getting tid data
         if self.compute_backend == ComputeBackend.WARP:
@@ -119,7 +123,7 @@ def _get_thread_data_3d(
             self.prepare_bc_auxilary_data = prepare_bc_auxilary_data
 
     @partial(jit, static_argnums=(0,), inline=True)
-    def prepare_bc_auxilary_data(self, f_pre, f_post, boundary_map, missing_mask):
+    def prepare_bc_auxilary_data(self, f_pre, f_post, bc_mask, missing_mask):
         """
         A placeholder function for prepare the auxilary distribution functions for the boundary condition.
         currently being called after collision only.
diff --git a/xlb/operator/boundary_masker/indices_boundary_masker.py b/xlb/operator/boundary_masker/indices_boundary_masker.py
index 54cb7aa..fdc4331 100644
--- a/xlb/operator/boundary_masker/indices_boundary_masker.py
+++ b/xlb/operator/boundary_masker/indices_boundary_masker.py
@@ -28,25 +28,26 @@ def __init__(
 
     def are_indices_in_interior(self, indices, shape):
         """
-        Check if all 2D or 3D indices are inside the bounds of the domain with the given shape and not
+        Check if each 2D or 3D index is inside the bounds of the domain with the given shape and not
         at its boundary.
 
         :param indices: List of tuples, where each tuple contains indices for each dimension.
         :param shape: Tuple representing the shape of the domain (nx, ny) for 2D or (nx, ny, nz) for 3D.
-        :return: Boolean flag is_inside indicating whether all indices are inside the bounds.
+        :return: List of boolean flags where each flag indicates whether the corresponding index is inside the bounds.
         """
         # Ensure that the number of dimensions in indices matches the domain shape
         dim = len(shape)
         if len(indices) != dim:
             raise ValueError(f"Indices tuple must have {dim} dimensions to match the domain shape.")
 
-        # Check if all indices are within the bounds
-        return all(0 < idx < shape[d] - 1 for d, idx_list in enumerate(indices) for idx in idx_list)
+        # Check each index tuple and return a list of boolean flags
+        flags = [all(0 < idx[d] < shape[d] - 1 for d in range(dim)) for idx in np.array(indices).T]
+        return flags
 
     @Operator.register_backend(ComputeBackend.JAX)
     # TODO HS: figure out why uncommenting the line below fails unlike other operators!
     # @partial(jit, static_argnums=(0))
-    def jax_implementation(self, bclist, boundary_map, missing_mask, start_index=None):
+    def jax_implementation(self, bclist, bc_mask, missing_mask, start_index=None):
         # Pad the missing mask to create a grid mask to identify out of bound boundaries
         # Set padded regin to True (i.e. boundary)
         dim = missing_mask.ndim - 1
@@ -54,24 +55,25 @@ def jax_implementation(self, bclist, boundary_map, missing_mask, start_index=Non
         pad_x, pad_y, pad_z = nDevices, 1, 1
         if dim == 2:
             grid_mask = jnp.pad(missing_mask, ((0, 0), (pad_x, pad_x), (pad_y, pad_y)), constant_values=True)
+            bmap = jnp.pad(bc_mask[0], ((pad_x, pad_x), (pad_y, pad_y)), constant_values=0)
         if dim == 3:
             grid_mask = jnp.pad(missing_mask, ((0, 0), (pad_x, pad_x), (pad_y, pad_y), (pad_z, pad_z)), constant_values=True)
+            bmap = jnp.pad(bc_mask[0], ((pad_x, pad_x), (pad_y, pad_y), (pad_z, pad_z)), constant_values=0)
 
         # shift indices
         shift_tup = (pad_x, pad_y) if dim == 2 else (pad_x, pad_y, pad_z)
         if start_index is None:
             start_index = (0,) * dim
 
-        bmap = boundary_map[0]
-        domain_shape = bmap.shape
+        domain_shape = bc_mask[0].shape
         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
             local_indices = np.array(bc.indices) - np.array(start_index)[:, np.newaxis]
             padded_indices = local_indices + np.array(shift_tup)[:, np.newaxis]
-            bmap = bmap.at[tuple(local_indices)].set(id_number)
-            if self.are_indices_in_interior(bc.indices, domain_shape):
+            bmap = bmap.at[tuple(padded_indices)].set(id_number)
+            if any(self.are_indices_in_interior(bc.indices, domain_shape)) and bc.needs_padding:
                 # checking if all indices associated with this BC are in the interior of the domain (not at the boundary).
                 # This flag is needed e.g. if the no-slip geometry is anywhere but at the boundaries of the computational domain.
                 if dim == 2:
@@ -80,20 +82,21 @@ def jax_implementation(self, bclist, boundary_map, missing_mask, start_index=Non
                     grid_mask = grid_mask.at[:, padded_indices[0], padded_indices[1], padded_indices[2]].set(True)
 
                 # Assign the boundary id to the push indices
-                push_indices = local_indices[:, :, None] + self.velocity_set.c[:, None, :]
-                push_indices = push_indices.reshape(3, -1)
+                push_indices = padded_indices[:, :, None] + self.velocity_set.c[:, None, :]
+                push_indices = push_indices.reshape(dim, -1)
                 bmap = bmap.at[tuple(push_indices)].set(id_number)
 
             # We are done with bc.indices. Remove them from BC objects
             bc.__dict__.pop("indices", None)
 
-        boundary_map = boundary_map.at[0].set(bmap)
         grid_mask = self.stream(grid_mask)
         if dim == 2:
             missing_mask = grid_mask[:, pad_x:-pad_x, pad_y:-pad_y]
+            bc_mask = bc_mask.at[0].set(bmap[pad_x:-pad_x, pad_y:-pad_y])
         if dim == 3:
             missing_mask = grid_mask[:, pad_x:-pad_x, pad_y:-pad_y, pad_z:-pad_z]
-        return boundary_map, missing_mask
+            bc_mask = bc_mask.at[0].set(bmap[pad_x:-pad_x, pad_y:-pad_y, pad_z:-pad_z])
+        return bc_mask, missing_mask
 
     def _construct_warp(self):
         # Make constants for warp
@@ -106,7 +109,7 @@ def kernel2d(
             indices: wp.array2d(dtype=wp.int32),
             id_number: wp.array1d(dtype=wp.uint8),
             is_interior: wp.array1d(dtype=wp.bool),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
             start_index: wp.vec2i,
         ):
@@ -129,12 +132,14 @@ def kernel2d(
                         pull_index[d] = index[d] - _c[d, l]
                         push_index[d] = index[d] + _c[d, l]
 
+                    # set bc_mask for all bc indices
+                    bc_mask[0, index[0], index[1]] = id_number[ii]
+
                     # 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]:
                         # Set the missing mask
                         missing_mask[l, index[0], index[1]] = True
-                        boundary_map[0, index[0], index[1]] = id_number[ii]
 
                     # handling geometries in the interior of the computational domain
                     elif (
@@ -146,7 +151,7 @@ def kernel2d(
                     ):
                         # Set the missing mask
                         missing_mask[l, push_index[0], push_index[1]] = True
-                        boundary_map[0, push_index[0], push_index[1]] = id_number[ii]
+                        bc_mask[0, push_index[0], push_index[1]] = id_number[ii]
 
         # Construct the warp 3D kernel
         @wp.kernel
@@ -154,7 +159,7 @@ def kernel3d(
             indices: wp.array2d(dtype=wp.int32),
             id_number: wp.array1d(dtype=wp.uint8),
             is_interior: wp.array1d(dtype=wp.bool),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
             start_index: wp.vec3i,
         ):
@@ -185,6 +190,9 @@ def kernel3d(
                         pull_index[d] = index[d] - _c[d, l]
                         push_index[d] = index[d] + _c[d, l]
 
+                    # set bc_mask for all bc indices
+                    bc_mask[0, index[0], index[1], index[2]] = id_number[ii]
+
                     # check if pull index is out of bound
                     # These directions will have missing information after streaming
                     if (
@@ -197,7 +205,6 @@ def kernel3d(
                     ):
                         # Set the missing mask
                         missing_mask[l, index[0], index[1], index[2]] = True
-                        boundary_map[0, index[0], index[1], index[2]] = id_number[ii]
 
                     # handling geometries in the interior of the computational domain
                     elif (
@@ -211,14 +218,14 @@ def kernel3d(
                     ):
                         # Set the missing mask
                         missing_mask[l, push_index[0], push_index[1], push_index[2]] = True
-                        boundary_map[0, push_index[0], push_index[1], push_index[2]] = id_number[ii]
+                        bc_mask[0, push_index[0], push_index[1], push_index[2]] = id_number[ii]
 
         kernel = kernel3d if self.velocity_set.d == 3 else kernel2d
 
         return None, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, bclist, boundary_map, missing_mask, start_index=None):
+    def warp_implementation(self, bclist, bc_mask, missing_mask, start_index=None):
         dim = self.velocity_set.d
         index_list = [[] for _ in range(dim)]
         id_list = []
@@ -229,7 +236,7 @@ def warp_implementation(self, bclist, boundary_map, missing_mask, start_index=No
             for d in range(dim):
                 index_list[d] += bc.indices[d]
             id_list += [bc.id] * len(bc.indices[0])
-            is_interior += [self.are_indices_in_interior(bc.indices, boundary_map[0].shape)] * len(bc.indices[0])
+            is_interior += self.are_indices_in_interior(bc.indices, bc_mask[0].shape) if bc.needs_padding else [False] * len(bc.indices[0])
 
             # We are done with bc.indices. Remove them from BC objects
             bc.__dict__.pop("indices", None)
@@ -248,11 +255,11 @@ def warp_implementation(self, bclist, boundary_map, missing_mask, start_index=No
                 indices,
                 id_number,
                 is_interior,
-                boundary_map,
+                bc_mask,
                 missing_mask,
                 start_index,
             ],
             dim=indices.shape[1],
         )
 
-        return boundary_map, missing_mask
+        return bc_mask, missing_mask
diff --git a/xlb/operator/boundary_masker/mesh_boundary_masker.py b/xlb/operator/boundary_masker/mesh_boundary_masker.py
index c43ea02..b88d251 100644
--- a/xlb/operator/boundary_masker/mesh_boundary_masker.py
+++ b/xlb/operator/boundary_masker/mesh_boundary_masker.py
@@ -33,17 +33,17 @@ def jax_implementation(
         bc,
         origin,
         spacing,
-        boundary_map,
+        bc_mask,
         missing_mask,
         start_index=(0, 0, 0),
     ):
         raise NotImplementedError(f"Operation {self.__class__.__name} not implemented in JAX!")
         # Use Warp backend even for this particular operation.
         wp.init()
-        boundary_map = wp.from_jax(boundary_map)
+        bc_mask = wp.from_jax(bc_mask)
         missing_mask = wp.from_jax(missing_mask)
-        boundary_map, missing_mask = self.warp_implementation(bc, origin, spacing, boundary_map, missing_mask, start_index)
-        return wp.to_jax(boundary_map), wp.to_jax(missing_mask)
+        bc_mask, missing_mask = self.warp_implementation(bc, origin, spacing, bc_mask, missing_mask, start_index)
+        return wp.to_jax(bc_mask), wp.to_jax(missing_mask)
 
     def _construct_warp(self):
         # Make constants for warp
@@ -57,7 +57,7 @@ def kernel(
             origin: wp.vec3,
             spacing: wp.vec3,
             id_number: wp.int32,
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
             start_index: wp.vec3i,
         ):
@@ -77,9 +77,9 @@ def kernel(
 
             # Compute the maximum length
             max_length = wp.sqrt(
-                (spacing[0] * wp.float32(boundary_map.shape[1])) ** 2.0
-                + (spacing[1] * wp.float32(boundary_map.shape[2])) ** 2.0
-                + (spacing[2] * wp.float32(boundary_map.shape[3])) ** 2.0
+                (spacing[0] * wp.float32(bc_mask.shape[1])) ** 2.0
+                + (spacing[1] * wp.float32(bc_mask.shape[2])) ** 2.0
+                + (spacing[2] * wp.float32(bc_mask.shape[3])) ** 2.0
             )
 
             # evaluate if point is inside mesh
@@ -98,7 +98,7 @@ def kernel(
                             push_index[d] = index[d] + _c[d, l]
 
                         # Set the boundary id and missing_mask
-                        boundary_map[0, push_index[0], push_index[1], push_index[2]] = wp.uint8(id_number)
+                        bc_mask[0, push_index[0], push_index[1], push_index[2]] = wp.uint8(id_number)
                         missing_mask[l, push_index[0], push_index[1], push_index[2]] = True
 
         return None, kernel
@@ -109,7 +109,7 @@ def warp_implementation(
         bc,
         origin,
         spacing,
-        boundary_map,
+        bc_mask,
         missing_mask,
         start_index=(0, 0, 0),
     ):
@@ -138,11 +138,11 @@ def warp_implementation(
                 origin,
                 spacing,
                 id_number,
-                boundary_map,
+                bc_mask,
                 missing_mask,
                 start_index,
             ],
-            dim=boundary_map.shape[1:],
+            dim=bc_mask.shape[1:],
         )
 
-        return boundary_map, missing_mask
+        return bc_mask, missing_mask
diff --git a/xlb/operator/collision/__init__.py b/xlb/operator/collision/__init__.py
index b48d0ce..0526c8a 100644
--- a/xlb/operator/collision/__init__.py
+++ b/xlb/operator/collision/__init__.py
@@ -1,3 +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
diff --git a/xlb/operator/collision/forced_collision.py b/xlb/operator/collision/forced_collision.py
new file mode 100644
index 0000000..31ef392
--- /dev/null
+++ b/xlb/operator/collision/forced_collision.py
@@ -0,0 +1,135 @@
+import jax.numpy as jnp
+from jax import jit
+import warp as wp
+from typing import Any
+
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.collision.collision import Collision
+from xlb.operator import Operator
+from functools import partial
+from xlb.operator.force import ExactDifference
+
+
+class ForcedCollision(Collision):
+    """
+    A collision operator for LBM with external force.
+    """
+
+    def __init__(
+        self,
+        collision_operator: Operator,
+        forcing_scheme="exact_difference",
+        force_vector=None,
+    ):
+        assert collision_operator is not None
+        self.collision_operator = collision_operator
+        super().__init__(self.collision_operator.omega)
+
+        assert forcing_scheme == "exact_difference", NotImplementedError(f"Force model {forcing_scheme} not implemented!")
+        assert force_vector.shape[0] == self.velocity_set.d, "Check the dimensions of the input force!"
+        self.force_vector = force_vector
+        if forcing_scheme == "exact_difference":
+            self.forcing_operator = ExactDifference(force_vector)
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0,))
+    def jax_implementation(self, f: jnp.ndarray, feq: jnp.ndarray, rho, u):
+        fout = self.collision_operator(f, feq, rho, u)
+        fout = self.forcing_operator(fout, feq, rho, u)
+        return fout
+
+    def _construct_warp(self):
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+
+        # Construct the functional
+        @wp.func
+        def functional(f: Any, feq: Any, rho: Any, u: Any):
+            fout = self.collision_operator.warp_functional(f, feq, rho, u)
+            fout = self.forcing_operator.warp_functional(fout, feq, rho, u)
+            return fout
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel2d(
+            f: wp.array3d(dtype=Any),
+            feq: wp.array3d(dtype=Any),
+            fout: wp.array3d(dtype=Any),
+            rho: wp.array3d(dtype=Any),
+            u: wp.array3d(dtype=Any),
+        ):
+            # Get the global index
+            i, j = wp.tid()
+            index = wp.vec2i(i, j)  # TODO: Warp needs to fix this
+
+            # Load needed values
+            _f = _f_vec()
+            _feq = _f_vec()
+            _d = self.velocity_set.d
+            for l in range(self.velocity_set.q):
+                _f[l] = f[l, index[0], index[1]]
+                _feq[l] = feq[l, index[0], index[1]]
+            _u = _u_vec()
+            for l in range(_d):
+                _u[l] = u[l, index[0], index[1]]
+            _rho = rho[0, index[0], index[1]]
+
+            # Compute the collision
+            _fout = functional(_f, _feq, _rho, _u)
+
+            # Write the result
+            for l in range(self.velocity_set.q):
+                fout[l, index[0], index[1]] = _fout[l]
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel3d(
+            f: wp.array4d(dtype=Any),
+            feq: wp.array4d(dtype=Any),
+            fout: wp.array4d(dtype=Any),
+            rho: wp.array4d(dtype=Any),
+            u: wp.array4d(dtype=Any),
+        ):
+            # Get the global index
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)  # TODO: Warp needs to fix this
+
+            # Load needed values
+            _f = _f_vec()
+            _feq = _f_vec()
+            _d = self.velocity_set.d
+            for l in range(self.velocity_set.q):
+                _f[l] = f[l, index[0], index[1], index[2]]
+                _feq[l] = feq[l, index[0], index[1], index[2]]
+            _u = _u_vec()
+            for l in range(_d):
+                _u[l] = u[l, index[0], index[1], index[2]]
+            _rho = rho[0, index[0], index[1], index[2]]
+
+            # Compute the collision
+            _fout = functional(_f, _feq, _rho, _u)
+
+            # Write the result
+            for l in range(self.velocity_set.q):
+                fout[l, index[0], index[1], index[2]] = _fout[l]
+
+        kernel = kernel3d if self.velocity_set.d == 3 else kernel2d
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f, feq, fout, rho, u):
+        # Launch the warp kernel
+        wp.launch(
+            self.warp_kernel,
+            inputs=[
+                f,
+                feq,
+                fout,
+                rho,
+                u,
+            ],
+            dim=f.shape[1:],
+        )
+        return fout
diff --git a/xlb/operator/collision/kbc.py b/xlb/operator/collision/kbc.py
index 748ebea..7984829 100644
--- a/xlb/operator/collision/kbc.py
+++ b/xlb/operator/collision/kbc.py
@@ -182,6 +182,7 @@ def _construct_warp(self):
             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
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
         _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
         _epsilon = wp.constant(self.compute_dtype(self.epsilon))
         _beta = wp.constant(self.compute_dtype(self.beta))
@@ -216,30 +217,34 @@ def decompose_shear_d3q27(
             s = _f_vec()
 
             # For c = (i, 0, 0), c = (0, j, 0) and c = (0, 0, k)
-            s[9] = (2.0 * nxz - nyz) / 6.0
-            s[18] = (2.0 * nxz - nyz) / 6.0
-            s[3] = (-nxz + 2.0 * nyz) / 6.0
-            s[6] = (-nxz + 2.0 * nyz) / 6.0
-            s[1] = (-nxz - nyz) / 6.0
-            s[2] = (-nxz - nyz) / 6.0
+            two = self.self.compute_dtype(2.0)
+            four = self.self.compute_dtype(4.0)
+            six = self.self.compute_dtype(6.0)
+
+            s[9] = (two * nxz - nyz) / six
+            s[18] = (two * nxz - nyz) / six
+            s[3] = (-nxz + two * nyz) / six
+            s[6] = (-nxz + two * nyz) / six
+            s[1] = (-nxz - nyz) / six
+            s[2] = (-nxz - nyz) / six
 
             # For c = (i, j, 0)
-            s[12] = pi[1] / 4.0
-            s[24] = pi[1] / 4.0
-            s[21] = -pi[1] / 4.0
-            s[15] = -pi[1] / 4.0
+            s[12] = pi[1] / four
+            s[24] = pi[1] / four
+            s[21] = -pi[1] / four
+            s[15] = -pi[1] / four
 
             # For c = (i, 0, k)
-            s[10] = pi[2] / 4.0
-            s[20] = pi[2] / 4.0
-            s[19] = -pi[2] / 4.0
-            s[11] = -pi[2] / 4.0
+            s[10] = pi[2] / four
+            s[20] = pi[2] / four
+            s[19] = -pi[2] / four
+            s[11] = -pi[2] / four
 
             # For c = (0, j, k)
-            s[8] = pi[4] / 4.0
-            s[4] = pi[4] / 4.0
-            s[7] = -pi[4] / 4.0
-            s[5] = -pi[4] / 4.0
+            s[8] = pi[4] / four
+            s[4] = pi[4] / four
+            s[7] = -pi[4] / four
+            s[5] = -pi[4] / four
 
             return s
 
@@ -271,10 +276,11 @@ def functional2d(
 
             # Perform collision
             delta_h = fneq - delta_s
-            gamma = _inv_beta - (2.0 - _inv_beta) * entropic_scalar_product(delta_s, delta_h, feq) / (
+            two = self.compute_dtype(2.0)
+            gamma = _inv_beta - (two - _inv_beta) * entropic_scalar_product(delta_s, delta_h, feq) / (
                 _epsilon + entropic_scalar_product(delta_h, delta_h, feq)
             )
-            fout = f - _beta * (2.0 * delta_s + gamma * delta_h)
+            fout = f - _beta * (two * delta_s + gamma * delta_h)
 
             return fout
 
@@ -293,10 +299,11 @@ def functional3d(
 
             # Perform collision
             delta_h = fneq - delta_s
-            gamma = _inv_beta - (2.0 - _inv_beta) * entropic_scalar_product(delta_s, delta_h, feq) / (
+            two = self.compute_dtype(2.0)
+            gamma = _inv_beta - (two - _inv_beta) * entropic_scalar_product(delta_s, delta_h, feq) / (
                 _epsilon + entropic_scalar_product(delta_h, delta_h, feq)
             )
-            fout = f - _beta * (2.0 * delta_s + gamma * delta_h)
+            fout = f - _beta * (two * delta_s + gamma * delta_h)
 
             return fout
 
@@ -305,9 +312,9 @@ def functional3d(
         def kernel2d(
             f: wp.array3d(dtype=Any),
             feq: wp.array3d(dtype=Any),
+            fout: wp.array3d(dtype=Any),
             rho: wp.array3d(dtype=Any),
             u: wp.array3d(dtype=Any),
-            fout: wp.array3d(dtype=Any),
         ):
             # Get the global index
             i, j = wp.tid()
@@ -320,7 +327,7 @@ def kernel2d(
             for l in range(self.velocity_set.q):
                 _f[l] = f[l, index[0], index[1]]
                 _feq[l] = feq[l, index[0], index[1]]
-            _u = self.warp_u_vec()
+            _u = _u_vec()
             for l in range(_d):
                 _u[l] = u[l, index[0], index[1]]
             _rho = rho[0, index[0], index[1]]
@@ -337,9 +344,9 @@ def kernel2d(
         def kernel3d(
             f: wp.array4d(dtype=Any),
             feq: wp.array4d(dtype=Any),
+            fout: wp.array4d(dtype=Any),
             rho: wp.array4d(dtype=Any),
             u: wp.array4d(dtype=Any),
-            fout: wp.array4d(dtype=Any),
         ):
             # Get the global index
             i, j, k = wp.tid()
@@ -352,7 +359,7 @@ def kernel3d(
             for l in range(self.velocity_set.q):
                 _f[l] = f[l, index[0], index[1], index[2]]
                 _feq[l] = feq[l, index[0], index[1], index[2]]
-            _u = self.warp_u_vec()
+            _u = _u_vec()
             for l in range(_d):
                 _u[l] = u[l, index[0], index[1], index[2]]
             _rho = rho[0, index[0], index[1], index[2]]
diff --git a/xlb/operator/force/__init__.py b/xlb/operator/force/__init__.py
index 6a991ce..2f3e3da 100644
--- a/xlb/operator/force/__init__.py
+++ b/xlb/operator/force/__init__.py
@@ -1,1 +1,2 @@
 from xlb.operator.force.momentum_transfer import MomentumTransfer as MomentumTransfer
+from xlb.operator.force.exact_difference_force import ExactDifference as ExactDifference
diff --git a/xlb/operator/force/exact_difference_force.py b/xlb/operator/force/exact_difference_force.py
new file mode 100644
index 0000000..f148e12
--- /dev/null
+++ b/xlb/operator/force/exact_difference_force.py
@@ -0,0 +1,150 @@
+from functools import partial
+from jax import jit, lax
+import warp as wp
+from typing import Any
+
+from xlb.velocity_set.velocity_set import VelocitySet
+from xlb.precision_policy import PrecisionPolicy
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.operator import Operator
+from xlb.operator.equilibrium import QuadraticEquilibrium
+
+
+class ExactDifference(Operator):
+    """
+    Add external body force based on the exact-difference method due to Kupershtokh (2004)
+
+    References
+    ----------
+    Kupershtokh, A. (2004). New method of incorporating a body force term into the lattice Boltzmann equation. In
+    Proceedings of the 5th International EHD Workshop (pp. 241-246). University of Poitiers, Poitiers, France.
+    Chikatamarla, S. S., & Karlin, I. V. (2013). Entropic lattice Boltzmann method for turbulent flow simulations:
+    Boundary conditions. Physica A, 392, 1925-1930.
+    Krüger, T., et al. (2017). The lattice Boltzmann method. Springer International Publishing, 10.978-3, 4-15.
+    """
+
+    def __init__(
+        self,
+        force_vector,
+        equilibrium: Operator = None,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+    ):
+        # TODO: currently we are limited to a single force vector not a spatially dependent forcing field
+        self.force_vector = force_vector
+        self.equilibrium = QuadraticEquilibrium() if equilibrium is None else equilibrium
+
+        # Call the parent constructor
+        super().__init__(
+            velocity_set,
+            precision_policy,
+            compute_backend,
+        )
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(self, f_postcollision, feq, rho, u):
+        """
+        Parameters
+        ----------
+        f_postcollision: jax.numpy.ndarray
+            The post-collision distribution functions.
+        feq: jax.numpy.ndarray
+            The equilibrium distribution functions.
+        rho: jax.numpy.ndarray
+            The density field.
+
+        u: jax.numpy.ndarray
+            The velocity field.
+
+        Returns
+        -------
+        f_postcollision: jax.numpy.ndarray
+        The post-collision distribution functions with the force applied.
+        """
+        delta_u = lax.broadcast_in_dim(self.force_vector, u.shape, (0,))
+        feq_force = self.equilibrium(rho, u + delta_u)
+        f_postcollision += feq_force - feq
+        return f_postcollision
+
+    def _construct_warp(self):
+        _d = self.velocity_set.d
+        _u_vec = wp.vec(_d, dtype=self.compute_dtype)
+        if _d == 2:
+            _force = _u_vec(self.force_vector[0], self.force_vector[1])
+        else:
+            _force = _u_vec(self.force_vector[0], self.force_vector[1], self.force_vector[2])
+
+        # Construct the functional
+        @wp.func
+        def functional(f_postcollision: Any, feq: Any, rho: Any, u: Any):
+            delta_u = _force
+            feq_force = self.equilibrium.warp_functional(rho, u + delta_u)
+            f_postcollision += feq_force - feq
+            return f_postcollision
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel2d(
+            f_postcollision: Any,
+            feq: Any,
+            fout: wp.array3d(dtype=Any),
+            rho: wp.array3d(dtype=Any),
+            u: wp.array3d(dtype=Any),
+        ):
+            # Get the global index
+            i, j = wp.tid()
+            index = wp.vec2i(i, j)
+
+            # Load needed values
+            _u = _u_vec()
+            for l in range(_d):
+                _u[l] = u[l, index[0], index[1]]
+            _rho = rho[0, index[0], index[1]]
+
+            # Compute the collision
+            _fout = functional(f_postcollision, feq, _rho, _u)
+
+            # Write the result
+            for l in range(self.velocity_set.q):
+                fout[l, index[0], index[1]] = _fout[l]
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel3d(
+            f_postcollision: Any,
+            feq: Any,
+            fout: wp.array4d(dtype=Any),
+            rho: wp.array4d(dtype=Any),
+            u: wp.array4d(dtype=Any),
+        ):
+            # Get the global index
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)  # TODO: Warp needs to fix this
+
+            # Load needed values
+            _u = _u_vec()
+            for l in range(_d):
+                _u[l] = u[l, index[0], index[1], index[2]]
+            _rho = rho[0, index[0], index[1], index[2]]
+
+            # Compute the collision
+            _fout = functional(f_postcollision, feq, _rho, _u)
+
+            # Write the result
+            for l in range(self.velocity_set.q):
+                fout[l, index[0], index[1], index[2]] = _fout[l]
+
+        kernel = kernel3d if self.velocity_set.d == 3 else kernel2d
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f_postcollision, feq, fout, rho, u):
+        # Launch the warp kernel
+        wp.launch(
+            self.warp_kernel,
+            inputs=[f_postcollision, feq, fout, rho, u],
+            dim=f_postcollision.shape[1:],
+        )
+        return fout
diff --git a/xlb/operator/force/momentum_transfer.py b/xlb/operator/force/momentum_transfer.py
index d25baf7..da64e67 100644
--- a/xlb/operator/force/momentum_transfer.py
+++ b/xlb/operator/force/momentum_transfer.py
@@ -50,13 +50,13 @@ def __init__(
 
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def jax_implementation(self, f, boundary_map, missing_mask):
+    def jax_implementation(self, f, bc_mask, missing_mask):
         """
         Parameters
         ----------
         f : jax.numpy.ndarray
             The post-collision distribution function at each node in the grid.
-        boundary_map : jax.numpy.ndarray
+        bc_mask : jax.numpy.ndarray
             A grid field with 0 everywhere except for boundary nodes which are designated
             by their respective boundary id's.
         missing_mask : jax.numpy.ndarray
@@ -71,10 +71,10 @@ def jax_implementation(self, f, boundary_map, missing_mask):
         # Give the input post-collision populations, streaming once and apply the BC the find post-stream values.
         f_post_collision = f
         f_post_stream = self.stream(f_post_collision)
-        f_post_stream = self.no_slip_bc_instance(f_post_collision, f_post_stream, boundary_map, missing_mask)
+        f_post_stream = self.no_slip_bc_instance(f_post_collision, f_post_stream, bc_mask, missing_mask)
 
         # Compute momentum transfer
-        boundary = boundary_map == self.no_slip_bc_instance.id
+        boundary = bc_mask == self.no_slip_bc_instance.id
         new_shape = (self.velocity_set.q,) + boundary.shape[1:]
         boundary = lax.broadcast_in_dim(boundary, new_shape, tuple(range(self.velocity_set.d + 1)))
 
@@ -103,7 +103,7 @@ def _construct_warp(self):
         @wp.kernel
         def kernel2d(
             f: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=wp.uint8),
+            bc_mask: wp.array3d(dtype=wp.uint8),
             missing_mask: wp.array3d(dtype=wp.bool),
             force: wp.array(dtype=Any),
         ):
@@ -112,7 +112,7 @@ def kernel2d(
             index = wp.vec2i(i, j)
 
             # Get the boundary id
-            _boundary_map = boundary_map[0, index[0], index[1]]
+            _boundary_id = bc_mask[0, index[0], index[1]]
             _missing_mask = _missing_mask_vec()
             for l in range(self.velocity_set.q):
                 # TODO fix vec bool
@@ -123,7 +123,7 @@ def kernel2d(
 
             # Determin if boundary is an edge by checking if center is missing
             is_edge = wp.bool(False)
-            if _boundary_map == wp.uint8(_no_slip_id):
+            if _boundary_id == wp.uint8(_no_slip_id):
                 if _missing_mask[_zero_index] == wp.uint8(0):
                     is_edge = wp.bool(True)
 
@@ -140,10 +140,11 @@ def kernel2d(
                 f_post_stream = self.no_slip_bc_instance.warp_functional(f_post_collision, f_post_stream, _f_vec(), _missing_mask)
 
                 # Compute the momentum transfer
-                for l in range(self.velocity_set.q):
-                    if _missing_mask[l] == wp.uint8(1):
-                        phi = f_post_collision[_opp_indices[l]] + f_post_stream[l]
-                        for d in range(self.velocity_set.d):
+                for d in range(self.velocity_set.d):
+                    m[d] = self.compute_dtype(0.0)
+                    for l in range(self.velocity_set.q):
+                        if _missing_mask[l] == wp.uint8(1):
+                            phi = f_post_collision[_opp_indices[l]] + f_post_stream[l]
                             m[d] += phi * self.compute_dtype(_c[d, _opp_indices[l]])
 
             wp.atomic_add(force, 0, m)
@@ -152,7 +153,7 @@ def kernel2d(
         @wp.kernel
         def kernel3d(
             f: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=wp.uint8),
+            bc_mask: wp.array4d(dtype=wp.uint8),
             missing_mask: wp.array4d(dtype=wp.bool),
             force: wp.array(dtype=Any),
         ):
@@ -161,7 +162,7 @@ def kernel3d(
             index = wp.vec3i(i, j, k)
 
             # Get the boundary id
-            _boundary_map = boundary_map[0, index[0], index[1], index[2]]
+            _boundary_id = bc_mask[0, index[0], index[1], index[2]]
             _missing_mask = _missing_mask_vec()
             for l in range(self.velocity_set.q):
                 # TODO fix vec bool
@@ -172,7 +173,7 @@ def kernel3d(
 
             # Determin if boundary is an edge by checking if center is missing
             is_edge = wp.bool(False)
-            if _boundary_map == wp.uint8(_no_slip_id):
+            if _boundary_id == wp.uint8(_no_slip_id):
                 if _missing_mask[_zero_index] == wp.uint8(0):
                     is_edge = wp.bool(True)
 
@@ -189,10 +190,11 @@ def kernel3d(
                 f_post_stream = self.no_slip_bc_instance.warp_functional(f_post_collision, f_post_stream, _f_vec(), _missing_mask)
 
                 # Compute the momentum transfer
-                for l in range(self.velocity_set.q):
-                    if _missing_mask[l] == wp.uint8(1):
-                        phi = f_post_collision[_opp_indices[l]] + f_post_stream[l]
-                        for d in range(self.velocity_set.d):
+                for d in range(self.velocity_set.d):
+                    m[d] = self.compute_dtype(0.0)
+                    for l in range(self.velocity_set.q):
+                        if _missing_mask[l] == wp.uint8(1):
+                            phi = f_post_collision[_opp_indices[l]] + f_post_stream[l]
                             m[d] += phi * self.compute_dtype(_c[d, _opp_indices[l]])
 
             wp.atomic_add(force, 0, m)
@@ -203,14 +205,14 @@ def kernel3d(
         return None, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f, boundary_map, missing_mask):
+    def warp_implementation(self, f, bc_mask, missing_mask):
         # Allocate the force vector (the total integral value will be computed)
         force = wp.zeros((1), dtype=wp.vec3) if self.velocity_set.d == 3 else wp.zeros((1), dtype=wp.vec2)
 
         # Launch the warp kernel
         wp.launch(
             self.warp_kernel,
-            inputs=[f, boundary_map, missing_mask, force],
+            inputs=[f, bc_mask, missing_mask, force],
             dim=f.shape[1:],
         )
         return force.numpy()[0]
diff --git a/xlb/operator/macroscopic/second_moment.py b/xlb/operator/macroscopic/second_moment.py
index 5209d69..bda2a25 100644
--- a/xlb/operator/macroscopic/second_moment.py
+++ b/xlb/operator/macroscopic/second_moment.py
@@ -72,7 +72,7 @@ def functional(
             # Get second order moment (a symmetric tensore shaped into a vector)
             pi = _pi_vec()
             for d in range(_pi_dim):
-                pi[d] = 0.0
+                pi[d] = self.compute_dtype(0.0)
                 for q in range(self.velocity_set.q):
                     pi[d] += _cc[q, d] * fneq[q]
             return pi
diff --git a/xlb/operator/operator.py b/xlb/operator/operator.py
index 83c6538..6e8bbbb 100644
--- a/xlb/operator/operator.py
+++ b/xlb/operator/operator.py
@@ -1,8 +1,10 @@
 import inspect
 import traceback
+import jax
 
 from xlb.compute_backend import ComputeBackend
 from xlb import DefaultConfig
+from xlb.precision_policy import PrecisionPolicy
 
 
 class Operator:
@@ -28,6 +30,12 @@ def __init__(self, velocity_set=None, precision_policy=None, compute_backend=Non
         if self.compute_backend == ComputeBackend.WARP:
             self.warp_functional, self.warp_kernel = self._construct_warp()
 
+        # Updating JAX config in case fp64 is requested
+        if self.compute_backend == ComputeBackend.JAX and (
+            precision_policy == PrecisionPolicy.FP64FP64 or precision_policy == PrecisionPolicy.FP64FP32
+        ):
+            jax.config.update("jax_enable_x64", True)
+
     @classmethod
     def register_backend(cls, backend_name):
         """
diff --git a/xlb/operator/stepper/nse_stepper.py b/xlb/operator/stepper/nse_stepper.py
index 60f610a..18b2167 100644
--- a/xlb/operator/stepper/nse_stepper.py
+++ b/xlb/operator/stepper/nse_stepper.py
@@ -15,10 +15,11 @@
 from xlb.operator.stepper import Stepper
 from xlb.operator.boundary_condition.boundary_condition import ImplementationStep
 from xlb.operator.boundary_condition import DoNothingBC as DummyBC
+from xlb.operator.collision import ForcedCollision
 
 
 class IncompressibleNavierStokesStepper(Stepper):
-    def __init__(self, omega, boundary_conditions=[], collision_type="BGK"):
+    def __init__(self, omega, boundary_conditions=[], collision_type="BGK", forcing_scheme="exact_difference", force_vector=None):
         velocity_set = DefaultConfig.velocity_set
         precision_policy = DefaultConfig.default_precision_policy
         compute_backend = DefaultConfig.default_backend
@@ -29,6 +30,9 @@ def __init__(self, omega, boundary_conditions=[], collision_type="BGK"):
         elif collision_type == "KBC":
             self.collision = KBC(omega, velocity_set, precision_policy, compute_backend)
 
+        if force_vector is not None:
+            self.collision = ForcedCollision(collision_operator=self.collision, forcing_scheme=forcing_scheme, force_vector=force_vector)
+
         # Construct the operators
         self.stream = Stream(velocity_set, precision_policy, compute_backend)
         self.equilibrium = QuadraticEquilibrium(velocity_set, precision_policy, compute_backend)
@@ -40,7 +44,7 @@ def __init__(self, omega, boundary_conditions=[], collision_type="BGK"):
 
     @Operator.register_backend(ComputeBackend.JAX)
     @partial(jit, static_argnums=(0))
-    def jax_implementation(self, f_0, f_1, boundary_map, missing_mask, timestep):
+    def jax_implementation(self, f_0, f_1, bc_mask, missing_mask, timestep):
         """
         Perform a single step of the lattice boltzmann method
         """
@@ -57,7 +61,7 @@ def jax_implementation(self, f_0, f_1, boundary_map, missing_mask, timestep):
                 f_post_stream = bc(
                     f_0,
                     f_post_stream,
-                    boundary_map,
+                    bc_mask,
                     missing_mask,
                 )
 
@@ -72,12 +76,12 @@ def jax_implementation(self, f_0, f_1, boundary_map, missing_mask, timestep):
 
         # Apply collision type boundary conditions
         for bc in self.boundary_conditions:
-            f_post_collision = bc.prepare_bc_auxilary_data(f_post_stream, f_post_collision, boundary_map, missing_mask)
+            f_post_collision = bc.prepare_bc_auxilary_data(f_post_stream, f_post_collision, bc_mask, missing_mask)
             if bc.implementation_step == ImplementationStep.COLLISION:
                 f_post_collision = bc(
                     f_post_stream,
                     f_post_collision,
-                    boundary_map,
+                    bc_mask,
                     missing_mask,
                 )
 
@@ -92,8 +96,6 @@ def _construct_warp(self):
         _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)  # TODO fix vec bool
         _c = self.velocity_set.c
         _q = self.velocity_set.q
-        _opp_indices = self.velocity_set.opp_indices
-        sound_speed = 1.0 / wp.sqrt(3.0)
 
         @wp.struct
         class BoundaryConditionIDStruct:
@@ -116,32 +118,32 @@ def apply_post_streaming_bc(
             f_post: Any,
             f_aux: Any,
             missing_mask: Any,
-            _boundary_map: Any,
+            _boundary_id: Any,
             bc_struct: Any,
         ):
             # Apply post-streaming type boundary conditions
-            if _boundary_map == bc_struct.id_EquilibriumBC:
+            if _boundary_id == bc_struct.id_EquilibriumBC:
                 # Equilibrium boundary condition
                 f_post = self.EquilibriumBC.warp_functional(f_pre, f_post, f_aux, missing_mask)
-            elif _boundary_map == bc_struct.id_DoNothingBC:
+            elif _boundary_id == bc_struct.id_DoNothingBC:
                 # Do nothing boundary condition
                 f_post = self.DoNothingBC.warp_functional(f_pre, f_post, f_aux, missing_mask)
-            elif _boundary_map == bc_struct.id_HalfwayBounceBackBC:
+            elif _boundary_id == bc_struct.id_HalfwayBounceBackBC:
                 # Half way boundary condition
                 f_post = self.HalfwayBounceBackBC.warp_functional(f_pre, f_post, f_aux, missing_mask)
-            elif _boundary_map == bc_struct.id_ZouHeBC_velocity:
+            elif _boundary_id == bc_struct.id_ZouHeBC_velocity:
                 # Zouhe boundary condition (bc type = velocity)
                 f_post = self.ZouHeBC_velocity.warp_functional(f_pre, f_post, f_aux, missing_mask)
-            elif _boundary_map == bc_struct.id_ZouHeBC_pressure:
+            elif _boundary_id == bc_struct.id_ZouHeBC_pressure:
                 # Zouhe boundary condition (bc type = pressure)
                 f_post = self.ZouHeBC_pressure.warp_functional(f_pre, f_post, f_aux, missing_mask)
-            elif _boundary_map == bc_struct.id_RegularizedBC_velocity:
+            elif _boundary_id == bc_struct.id_RegularizedBC_velocity:
                 # Regularized boundary condition (bc type = velocity)
                 f_post = self.RegularizedBC_velocity.warp_functional(f_pre, f_post, f_aux, missing_mask)
-            elif _boundary_map == bc_struct.id_RegularizedBC_pressure:
+            elif _boundary_id == bc_struct.id_RegularizedBC_pressure:
                 # Regularized boundary condition (bc type = velocity)
                 f_post = self.RegularizedBC_pressure.warp_functional(f_pre, f_post, f_aux, missing_mask)
-            elif _boundary_map == bc_struct.id_ExtrapolationOutflowBC:
+            elif _boundary_id == bc_struct.id_ExtrapolationOutflowBC:
                 # Regularized boundary condition (bc type = velocity)
                 f_post = self.ExtrapolationOutflowBC.warp_functional(f_pre, f_post, f_aux, missing_mask)
             return f_post
@@ -152,13 +154,13 @@ def apply_post_collision_bc(
             f_post: Any,
             f_aux: Any,
             missing_mask: Any,
-            _boundary_map: Any,
+            _boundary_id: Any,
             bc_struct: Any,
         ):
-            if _boundary_map == bc_struct.id_FullwayBounceBackBC:
+            if _boundary_id == bc_struct.id_FullwayBounceBackBC:
                 # Full way boundary condition
                 f_post = self.FullwayBounceBackBC.warp_functional(f_pre, f_post, f_aux, missing_mask)
-            elif _boundary_map == bc_struct.id_ExtrapolationOutflowBC:
+            elif _boundary_id == bc_struct.id_ExtrapolationOutflowBC:
                 # f_aux is the neighbour's post-streaming values
                 # Storing post-streaming data in directions that leave the domain
                 f_post = self.ExtrapolationOutflowBC.prepare_bc_auxilary_data(f_pre, f_post, f_aux, missing_mask)
@@ -225,13 +227,13 @@ def get_thread_data_3d(
         def get_bc_auxilary_data_2d(
             f_0: wp.array3d(dtype=Any),
             index: Any,
-            _boundary_map: Any,
+            _boundary_id: Any,
             _missing_mask: Any,
             bc_struct: Any,
         ):
             # special preparation of auxiliary data
             f_auxiliary = _f_vec()
-            if _boundary_map == bc_struct.id_ExtrapolationOutflowBC:
+            if _boundary_id == bc_struct.id_ExtrapolationOutflowBC:
                 nv = get_normal_vectors_2d(_missing_mask)
                 for l in range(self.velocity_set.q):
                     if _missing_mask[l] == wp.uint8(1):
@@ -248,13 +250,13 @@ def get_bc_auxilary_data_2d(
         def get_bc_auxilary_data_3d(
             f_0: wp.array4d(dtype=Any),
             index: Any,
-            _boundary_map: Any,
+            _boundary_id: Any,
             _missing_mask: Any,
             bc_struct: Any,
         ):
             # special preparation of auxiliary data
             f_auxiliary = _f_vec()
-            if _boundary_map == bc_struct.id_ExtrapolationOutflowBC:
+            if _boundary_id == bc_struct.id_ExtrapolationOutflowBC:
                 nv = get_normal_vectors_3d(_missing_mask)
                 for l in range(self.velocity_set.q):
                     if _missing_mask[l] == wp.uint8(1):
@@ -271,7 +273,7 @@ def get_bc_auxilary_data_3d(
         def kernel2d(
             f_0: wp.array3d(dtype=Any),
             f_1: wp.array3d(dtype=Any),
-            boundary_map: wp.array3d(dtype=Any),
+            bc_mask: wp.array3d(dtype=Any),
             missing_mask: wp.array3d(dtype=Any),
             bc_struct: Any,
             timestep: int,
@@ -287,11 +289,11 @@ def kernel2d(
             f_post_stream = self.stream.warp_functional(f_0, index)
 
             # Prepare auxilary data for BC (if applicable)
-            _boundary_map = boundary_map[0, index[0], index[1]]
-            f_auxiliary = get_bc_auxilary_data_2d(f_0, index, _boundary_map, _missing_mask, bc_struct)
+            _boundary_id = bc_mask[0, index[0], index[1]]
+            f_auxiliary = get_bc_auxilary_data_2d(f_0, index, _boundary_id, _missing_mask, bc_struct)
 
             # Apply post-streaming type boundary conditions
-            f_post_stream = apply_post_streaming_bc(f_post_collision, f_post_stream, f_auxiliary, _missing_mask, _boundary_map, bc_struct)
+            f_post_stream = apply_post_streaming_bc(f_post_collision, f_post_stream, f_auxiliary, _missing_mask, _boundary_id, bc_struct)
 
             # Compute rho and u
             rho, u = self.macroscopic.warp_functional(f_post_stream)
@@ -308,7 +310,7 @@ def kernel2d(
             )
 
             # Apply post-collision type boundary conditions
-            f_post_collision = apply_post_collision_bc(f_post_stream, f_post_collision, f_auxiliary, _missing_mask, _boundary_map, bc_struct)
+            f_post_collision = apply_post_collision_bc(f_post_stream, f_post_collision, f_auxiliary, _missing_mask, _boundary_id, bc_struct)
 
             # Set the output
             for l in range(self.velocity_set.q):
@@ -319,7 +321,7 @@ def kernel2d(
         def kernel3d(
             f_0: wp.array4d(dtype=Any),
             f_1: wp.array4d(dtype=Any),
-            boundary_map: wp.array4d(dtype=Any),
+            bc_mask: wp.array4d(dtype=Any),
             missing_mask: wp.array4d(dtype=Any),
             bc_struct: Any,
             timestep: int,
@@ -335,11 +337,11 @@ def kernel3d(
             f_post_stream = self.stream.warp_functional(f_0, index)
 
             # Prepare auxilary data for BC (if applicable)
-            _boundary_map = boundary_map[0, index[0], index[1], index[2]]
-            f_auxiliary = get_bc_auxilary_data_3d(f_0, index, _boundary_map, _missing_mask, bc_struct)
+            _boundary_id = bc_mask[0, index[0], index[1], index[2]]
+            f_auxiliary = get_bc_auxilary_data_3d(f_0, index, _boundary_id, _missing_mask, bc_struct)
 
             # Apply post-streaming type boundary conditions
-            f_post_stream = apply_post_streaming_bc(f_post_collision, f_post_stream, f_auxiliary, _missing_mask, _boundary_map, bc_struct)
+            f_post_stream = apply_post_streaming_bc(f_post_collision, f_post_stream, f_auxiliary, _missing_mask, _boundary_id, bc_struct)
 
             # Compute rho and u
             rho, u = self.macroscopic.warp_functional(f_post_stream)
@@ -351,7 +353,7 @@ def kernel3d(
             f_post_collision = self.collision.warp_functional(f_post_stream, feq, rho, u)
 
             # Apply post-collision type boundary conditions
-            f_post_collision = apply_post_collision_bc(f_post_stream, f_post_collision, f_auxiliary, _missing_mask, _boundary_map, bc_struct)
+            f_post_collision = apply_post_collision_bc(f_post_stream, f_post_collision, f_auxiliary, _missing_mask, _boundary_id, bc_struct)
 
             # Set the output
             for l in range(self.velocity_set.q):
@@ -363,7 +365,7 @@ def kernel3d(
         return BoundaryConditionIDStruct, kernel
 
     @Operator.register_backend(ComputeBackend.WARP)
-    def warp_implementation(self, f_0, f_1, boundary_map, missing_mask, timestep):
+    def warp_implementation(self, f_0, f_1, bc_mask, missing_mask, timestep):
         # Get the boundary condition ids
         from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry
 
@@ -404,7 +406,7 @@ def warp_implementation(self, f_0, f_1, boundary_map, missing_mask, timestep):
             inputs=[
                 f_0,
                 f_1,
-                boundary_map,
+                bc_mask,
                 missing_mask,
                 bc_struct,
                 timestep,
diff --git a/xlb/velocity_set/velocity_set.py b/xlb/velocity_set/velocity_set.py
index 2405b36..41db991 100644
--- a/xlb/velocity_set/velocity_set.py
+++ b/xlb/velocity_set/velocity_set.py
@@ -4,9 +4,12 @@
 import numpy as np
 import warp as wp
 import jax.numpy as jnp
+import jax
 
 from xlb import DefaultConfig
 from xlb.compute_backend import ComputeBackend
+from xlb.precision_policy import PrecisionPolicy
+
 
 class VelocitySet(object):
     """
@@ -31,6 +34,10 @@ def __init__(self, d, q, c, w, precision_policy, backend):
         self.precision_policy = precision_policy
         self.backend = backend
 
+        # Updating JAX config in case fp64 is requested
+        if backend == ComputeBackend.JAX and (precision_policy == PrecisionPolicy.FP64FP64 or precision_policy == PrecisionPolicy.FP64FP32):
+            jax.config.update("jax_enable_x64", True)
+
         # Create all properties in NumPy first
         self._init_numpy_properties(c, w)
 
@@ -83,7 +90,7 @@ def _init_jax_properties(self):
         Convert NumPy properties to JAX-specific properties.
         """
         dtype = self.precision_policy.compute_precision.jax_dtype
-        self.c = jnp.array(self._c, dtype=dtype)
+        self.c = jnp.array(self._c, dtype=jnp.int32)
         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)
@@ -224,4 +231,3 @@ def __repr__(self):
         This function returns the name of the lattice in the format of DxQy.
         """
         return "D{}Q{}".format(self.d, self.q)
-


Test patch:
diff --git a/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py
index bd40dfb..1025e3b 100644
--- a/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py
+++ b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py
@@ -33,7 +33,7 @@ def test_bc_equilibrium_jax(dim, velocity_set, grid_shape):
 
     missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
 
-    boundary_map = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
 
     indices_boundary_masker = IndicesBoundaryMasker()
 
@@ -59,7 +59,7 @@ def test_bc_equilibrium_jax(dim, velocity_set, grid_shape):
         indices=indices,
     )
 
-    boundary_map, missing_mask = indices_boundary_masker([equilibrium_bc], boundary_map, missing_mask, start_index=None)
+    bc_mask, missing_mask = indices_boundary_masker([equilibrium_bc], bc_mask, missing_mask, start_index=None)
 
     f_pre = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.FP32)
 
@@ -67,7 +67,7 @@ def test_bc_equilibrium_jax(dim, velocity_set, grid_shape):
         cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=2.0
     )  # Arbitrary value so that we can check if the values are changed outside the boundary
 
-    f = equilibrium_bc(f_pre, f_post, boundary_map, missing_mask)
+    f = equilibrium_bc(f_pre, f_post, bc_mask, missing_mask)
 
     assert f.shape == (velocity_set.q,) + grid_shape
 
diff --git a/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py
index 9f0cd68..5eb0c10 100644
--- a/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py
+++ b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py
@@ -11,7 +11,7 @@ def init_xlb_env(velocity_set):
     vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
-        default_backend=ComputeBackend.JAX,
+        default_backend=ComputeBackend.WARP,
         velocity_set=vel_set,
     )
 
@@ -32,7 +32,7 @@ def test_bc_equilibrium_warp(dim, velocity_set, grid_shape):
 
     missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
 
-    boundary_map = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
 
     indices_boundary_masker = IndicesBoundaryMasker()
 
@@ -59,7 +59,7 @@ def test_bc_equilibrium_warp(dim, velocity_set, grid_shape):
         indices=indices,
     )
 
-    boundary_map, missing_mask = indices_boundary_masker([equilibrium_bc], boundary_map, missing_mask, start_index=None)
+    bc_mask, missing_mask = indices_boundary_masker([equilibrium_bc], bc_mask, missing_mask, start_index=None)
 
     f = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.FP32)
     f_pre = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.FP32)
@@ -67,7 +67,7 @@ def test_bc_equilibrium_warp(dim, velocity_set, grid_shape):
         cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=2.0
     )  # Arbitrary value so that we can check if the values are changed outside the boundary
 
-    f = equilibrium_bc(f_pre, f_post, boundary_map, missing_mask)
+    f = equilibrium_bc(f_pre, f_post, bc_mask, missing_mask)
 
     f = f.numpy()
     f_post = f_post.numpy()
diff --git a/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py
index dde05bb..cd18975 100644
--- a/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py
+++ b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py
@@ -6,6 +6,7 @@
 from xlb.compute_backend import ComputeBackend
 from xlb.grid import grid_factory
 from xlb import DefaultConfig
+from xlb.operator.boundary_masker import IndicesBoundaryMasker
 
 
 def init_xlb_env(velocity_set):
@@ -35,9 +36,9 @@ def test_fullway_bounce_back_jax(dim, velocity_set, grid_shape):
 
     missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
 
-    boundary_map = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
 
-    indices_boundary_masker = xlb.operator.boundary_masker.IndicesBoundaryMasker()
+    indices_boundary_masker = IndicesBoundaryMasker()
 
     # Make indices for boundary conditions (sphere)
     sphere_radius = grid_shape[0] // 4
@@ -55,7 +56,7 @@ def test_fullway_bounce_back_jax(dim, velocity_set, grid_shape):
     indices = [tuple(indices[i]) for i in range(velocity_set.d)]
     fullway_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices)
 
-    boundary_map, missing_mask = indices_boundary_masker([fullway_bc], boundary_map, missing_mask, start_index=None)
+    bc_mask, missing_mask = indices_boundary_masker([fullway_bc], bc_mask, missing_mask, start_index=None)
 
     f_pre = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=0.0)
     # Generate a random field with the same shape
@@ -68,13 +69,13 @@ def test_fullway_bounce_back_jax(dim, velocity_set, grid_shape):
         cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=2.0
     )  # Arbitrary value so that we can check if the values are changed outside the boundary
 
-    f = fullway_bc(f_pre, f_post, boundary_map, missing_mask)
+    f = fullway_bc(f_pre, f_post, bc_mask, missing_mask)
 
     assert f.shape == (velocity_set.q,) + grid_shape
 
     for i in range(velocity_set.q):
         jnp.allclose(
-            f[velocity_set.get_opp_index(i)][tuple(indices)],
+            f[velocity_set.opp_indices[i]][tuple(indices)],
             f_pre[i][tuple(indices)],
         )
 
diff --git a/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py
index 96e2f21..10b9244 100644
--- a/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py
+++ b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py
@@ -12,7 +12,7 @@ def init_xlb_env(velocity_set):
     vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
-        default_backend=ComputeBackend.JAX,
+        default_backend=ComputeBackend.WARP,
         velocity_set=vel_set,
     )
 
@@ -35,7 +35,7 @@ def test_fullway_bounce_back_warp(dim, velocity_set, grid_shape):
 
     missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
 
-    boundary_map = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
 
     indices_boundary_masker = IndicesBoundaryMasker()
 
@@ -55,7 +55,7 @@ def test_fullway_bounce_back_warp(dim, velocity_set, grid_shape):
     indices = [tuple(indices[i]) for i in range(velocity_set.d)]
     fullway_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices)
 
-    boundary_map, missing_mask = indices_boundary_masker([fullway_bc], boundary_map, missing_mask, start_index=None)
+    bc_mask, missing_mask = indices_boundary_masker([fullway_bc], bc_mask, missing_mask, start_index=None)
 
     # Generate a random field with the same shape
     random_field = np.random.rand(velocity_set.q, *grid_shape).astype(np.float32)
@@ -66,7 +66,7 @@ def test_fullway_bounce_back_warp(dim, velocity_set, grid_shape):
         cardinality=velocity_set.q, dtype=xlb.Precision.FP32, fill_value=2.0
     )  # Arbitrary value so that we can check if the values are changed outside the boundary
 
-    f_pre = fullway_bc(f_pre, f_post, boundary_map, missing_mask)
+    f_pre = fullway_bc(f_pre, f_post, bc_mask, missing_mask)
 
     f = f_pre.numpy()
     f_post = f_post.numpy()
@@ -75,7 +75,7 @@ def test_fullway_bounce_back_warp(dim, velocity_set, grid_shape):
 
     for i in range(velocity_set.q):
         np.allclose(
-            f[velocity_set.get_opp_index(i)][tuple(indices)],
+            f[velocity_set.opp_indices[i]][tuple(indices)],
             f_post[i][tuple(indices)],
         )
 
diff --git a/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py b/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py
index 9325890..79d56d8 100644
--- a/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py
+++ b/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py
@@ -15,6 +15,7 @@ def init_xlb_env(velocity_set):
         velocity_set=vel_set,
     )
 
+
 @pytest.mark.parametrize(
     "dim,velocity_set,grid_shape",
     [
@@ -34,7 +35,7 @@ def test_indices_masker_jax(dim, velocity_set, grid_shape):
 
     missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
 
-    boundary_map = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
 
     indices_boundary_masker = xlb.operator.boundary_masker.IndicesBoundaryMasker()
 
@@ -56,26 +57,26 @@ def test_indices_masker_jax(dim, velocity_set, grid_shape):
     assert len(indices) == dim
     test_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices)
     test_bc.id = 5
-    boundary_map, missing_mask = indices_boundary_masker([test_bc], boundary_map, missing_mask, start_index=None)
+    bc_mask, missing_mask = indices_boundary_masker([test_bc], bc_mask, missing_mask, start_index=None)
 
     assert missing_mask.dtype == xlb.Precision.BOOL.jax_dtype
 
-    assert boundary_map.dtype == xlb.Precision.UINT8.jax_dtype
+    assert bc_mask.dtype == xlb.Precision.UINT8.jax_dtype
 
-    assert boundary_map.shape == (1,) + grid_shape
+    assert bc_mask.shape == (1,) + grid_shape
 
     assert missing_mask.shape == (velocity_set.q,) + grid_shape
 
     if dim == 2:
-        assert jnp.all(boundary_map[0, indices[0], indices[1]] == test_bc.id)
-        # assert that the rest of the boundary_map is zero
-        boundary_map = boundary_map.at[0, indices[0], indices[1]].set(0)
-        assert jnp.all(boundary_map == 0)
+        assert jnp.all(bc_mask[0, indices[0], indices[1]] == test_bc.id)
+        # assert that the rest of the bc_mask is zero
+        bc_mask = bc_mask.at[0, indices[0], indices[1]].set(0)
+        assert jnp.all(bc_mask == 0)
     if dim == 3:
-        assert jnp.all(boundary_map[0, indices[0], indices[1], indices[2]] == test_bc.id)
-        # assert that the rest of the boundary_map is zero
-        boundary_map = boundary_map.at[0, indices[0], indices[1], indices[2]].set(0)
-        assert jnp.all(boundary_map == 0)
+        assert jnp.all(bc_mask[0, indices[0], indices[1], indices[2]] == test_bc.id)
+        # assert that the rest of the bc_mask is zero
+        bc_mask = bc_mask.at[0, indices[0], indices[1], indices[2]].set(0)
+        assert jnp.all(bc_mask == 0)
 
 
 if __name__ == "__main__":
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 43ad052..f56c9ce 100644
--- a/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py
+++ b/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py
@@ -4,13 +4,14 @@
 from xlb.compute_backend import ComputeBackend
 from xlb import DefaultConfig
 from xlb.grid import grid_factory
+from xlb.operator.boundary_masker import IndicesBoundaryMasker
 
 
 def init_xlb_env(velocity_set):
     vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
-        default_backend=ComputeBackend.JAX,
+        default_backend=ComputeBackend.WARP,
         velocity_set=vel_set,
     )
 
@@ -33,9 +34,9 @@ def test_indices_masker_warp(dim, velocity_set, grid_shape):
 
     missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
 
-    boundary_map = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
 
-    indices_boundary_masker = xlb.operator.boundary_masker.IndicesBoundaryMasker()
+    indices_boundary_masker = IndicesBoundaryMasker()
 
     # Make indices for boundary conditions (sphere)
     sphere_radius = grid_shape[0] // 4
@@ -55,33 +56,33 @@ def test_indices_masker_warp(dim, velocity_set, grid_shape):
     assert len(indices) == dim
     test_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices)
     test_bc.id = 5
-    boundary_map, missing_mask = indices_boundary_masker(
+    bc_mask, missing_mask = indices_boundary_masker(
         [test_bc],
-        boundary_map,
+        bc_mask,
         missing_mask,
         start_index=(0, 0, 0) if dim == 3 else (0, 0),
     )
     assert missing_mask.dtype == xlb.Precision.BOOL.wp_dtype
 
-    assert boundary_map.dtype == xlb.Precision.UINT8.wp_dtype
+    assert bc_mask.dtype == xlb.Precision.UINT8.wp_dtype
 
-    boundary_map = boundary_map.numpy()
+    bc_mask = bc_mask.numpy()
     missing_mask = missing_mask.numpy()
 
-    assert boundary_map.shape == (1,) + grid_shape
+    assert bc_mask.shape == (1,) + grid_shape
 
     assert missing_mask.shape == (velocity_set.q,) + grid_shape
 
     if dim == 2:
-        assert np.all(boundary_map[0, indices[0], indices[1]] == test_bc.id)
-        # assert that the rest of the boundary_map is zero
-        boundary_map[0, indices[0], indices[1]] = 0
-        assert np.all(boundary_map == 0)
+        assert np.all(bc_mask[0, indices[0], indices[1]] == test_bc.id)
+        # assert that the rest of the bc_mask is zero
+        bc_mask[0, indices[0], indices[1]] = 0
+        assert np.all(bc_mask == 0)
     if dim == 3:
-        assert np.all(boundary_map[0, indices[0], indices[1], indices[2]] == test_bc.id)
-        # assert that the rest of the boundary_map is zero
-        boundary_map[0, indices[0], indices[1], indices[2]] = 0
-        assert np.all(boundary_map == 0)
+        assert np.all(bc_mask[0, indices[0], indices[1], indices[2]] == test_bc.id)
+        # assert that the rest of the bc_mask is zero
+        bc_mask[0, indices[0], indices[1], indices[2]] = 0
+        assert np.all(bc_mask == 0)
 
 
 if __name__ == "__main__":
diff --git a/tests/grids/test_grid_warp.py b/tests/grids/test_grid_warp.py
index 11c8b2a..782434d 100644
--- a/tests/grids/test_grid_warp.py
+++ b/tests/grids/test_grid_warp.py
@@ -11,10 +11,11 @@ def init_xlb_env(velocity_set):
     vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
-        default_backend=ComputeBackend.JAX,
+        default_backend=ComputeBackend.WARP,
         velocity_set=vel_set,
     )
 
+
 @pytest.mark.parametrize("grid_size", [50, 100, 150])
 def test_warp_grid_create_field(grid_size):
     for grid_shape in [(grid_size, grid_size), (grid_size, grid_size, grid_size)]:
diff --git a/tests/kernels/collision/test_bgk_collision_jax.py b/tests/kernels/collision/test_bgk_collision_jax.py
index aebc726..1672cd5 100644
--- a/tests/kernels/collision/test_bgk_collision_jax.py
+++ b/tests/kernels/collision/test_bgk_collision_jax.py
@@ -16,6 +16,7 @@ def init_xlb_env(velocity_set):
         velocity_set=vel_set,
     )
 
+
 @pytest.mark.parametrize(
     "dim,velocity_set,grid_shape,omega",
     [
diff --git a/tests/kernels/collision/test_bgk_collision_warp.py b/tests/kernels/collision/test_bgk_collision_warp.py
index 2743050..382e368 100644
--- a/tests/kernels/collision/test_bgk_collision_warp.py
+++ b/tests/kernels/collision/test_bgk_collision_warp.py
@@ -12,7 +12,7 @@ def init_xlb_env(velocity_set):
     vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
-        default_backend=ComputeBackend.JAX,
+        default_backend=ComputeBackend.WARP,
         velocity_set=vel_set,
     )
 
diff --git a/tests/kernels/equilibrium/test_equilibrium_jax.py b/tests/kernels/equilibrium/test_equilibrium_jax.py
index 50418bc..aa4f051 100644
--- a/tests/kernels/equilibrium/test_equilibrium_jax.py
+++ b/tests/kernels/equilibrium/test_equilibrium_jax.py
@@ -8,13 +8,14 @@
 
 
 def init_xlb_env(velocity_set):
-    vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
+    vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.JAX)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
         default_backend=ComputeBackend.JAX,
         velocity_set=vel_set,
     )
 
+
 @pytest.mark.parametrize(
     "dim,velocity_set,grid_shape",
     [
diff --git a/tests/kernels/equilibrium/test_equilibrium_warp.py b/tests/kernels/equilibrium/test_equilibrium_warp.py
index 9759fb2..fdd796a 100644
--- a/tests/kernels/equilibrium/test_equilibrium_warp.py
+++ b/tests/kernels/equilibrium/test_equilibrium_warp.py
@@ -11,9 +11,11 @@ def init_xlb_env(velocity_set):
     vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
-        default_backend=ComputeBackend.JAX,
+        default_backend=ComputeBackend.WARP,
         velocity_set=vel_set,
     )
+
+
 @pytest.mark.parametrize(
     "dim,velocity_set,grid_shape",
     [
diff --git a/tests/kernels/macroscopic/test_macroscopic_warp.py b/tests/kernels/macroscopic/test_macroscopic_warp.py
index 6a97927..4f33bc2 100644
--- a/tests/kernels/macroscopic/test_macroscopic_warp.py
+++ b/tests/kernels/macroscopic/test_macroscopic_warp.py
@@ -12,7 +12,7 @@ def init_xlb_env(velocity_set):
     vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
-        default_backend=ComputeBackend.JAX,
+        default_backend=ComputeBackend.WARP,
         velocity_set=vel_set,
     )
 
diff --git a/tests/kernels/stream/test_stream_jax.py b/tests/kernels/stream/test_stream_jax.py
index ef635ea..c1cae52 100644
--- a/tests/kernels/stream/test_stream_jax.py
+++ b/tests/kernels/stream/test_stream_jax.py
@@ -8,10 +8,11 @@
 
 
 def init_xlb_env(velocity_set):
+    vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.JAX)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
         default_backend=ComputeBackend.JAX,
-        velocity_set=velocity_set(),
+        velocity_set=vel_set,
     )
 
 
diff --git a/tests/kernels/stream/test_stream_warp.py b/tests/kernels/stream/test_stream_warp.py
index b83368d..0d100cf 100644
--- a/tests/kernels/stream/test_stream_warp.py
+++ b/tests/kernels/stream/test_stream_warp.py
@@ -10,10 +10,11 @@
 
 
 def init_xlb_env(velocity_set):
+    vel_set = velocity_set(precision_policy=xlb.PrecisionPolicy.FP32FP32, backend=ComputeBackend.WARP)
     xlb.init(
         default_precision_policy=xlb.PrecisionPolicy.FP32FP32,
         default_backend=ComputeBackend.WARP,
-        velocity_set=velocity_set(),
+        velocity_set=vel_set,
     )