Repository: Autodesk/XLB (language: python)

Issue:
Marged 2D and 3D kernels in Warp
129 additions and 1,265 deletions

![image](https://github.com/user-attachments/assets/3cc1086a-435b-4380-9c95-581dd7eb3683)



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
new file mode 100644
index 0000000..1b44c5a
--- /dev/null
+++ b/.github/workflows/lint.yml
@@ -0,0 +1,27 @@
+name: Lint
+
+on:
+  pull_request:
+    branches:
+      - major-refactoring  # Remember to add main branch later
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install ruff
+
+    - name: Run Ruff
+      run: ruff check .
diff --git a/.gitignore b/.gitignore
index d774451..2fad3f7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -39,7 +39,7 @@ Thumbs.db
 __pycache__/
 *.py[cod]
 *$py.class
-
+**pyc
 # C extensions
 *.so
 
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..6a2bd2f
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,6 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.5.6
+    hooks:
+      - id: ruff
+        args: [--fix]
diff --git a/examples/CFD/airfoil3d.py b/examples/CFD/airfoil3d.py
deleted file mode 100644
index 77e4b58..0000000
--- a/examples/CFD/airfoil3d.py
+++ /dev/null
@@ -1,195 +0,0 @@
-"""
-This is a example for simulating fluid flow around a NACA airfoil using the lattice Boltzmann method (LBM). 
-The LBM is a computational fluid dynamics method for simulating fluid flow and is particularly effective 
-for complex geometries and multiphase flow. 
-
-In this example you'll be introduced to the following concepts:
-
-1. Lattice: The example uses a D3Q27 lattice, which is a three-dimensional lattice model that considers 
-    27 discrete velocity directions. This allows for a more accurate representation of the fluid flow 
-    in three dimensions.
-
-2. NACA Airfoil Generation: The example includes a function to generate a NACA airfoil shape, which is 
-    common in aerodynamics. The function allows for customization of the length, thickness, and angle 
-    of the airfoil.
-
-3. Boundary Conditions: The example includes several boundary conditions. These include a "bounce back" 
-    condition on the airfoil surface and the top and bottom of the domain, a "do nothing" condition 
-    at the outlet (right side of the domain), and an "equilibrium" condition at the inlet 
-    (left side of the domain) to simulate a uniform flow.
-
-4. Simulation Parameters: The example allows for the setting of various simulation parameters, 
-    including the Reynolds number, inlet velocity, and characteristic length. 
-
-5. In-situ visualization: The example outputs rendering images of the q-criterion using
-     PhantomGaze library (https://github.com/loliverhennigh/PhantomGaze) without any I/O overhead 
-     while the data is still on the GPU. 
-"""
-
-
-import numpy as np
-# from IPython import display
-import matplotlib.pylab as plt
-from src.models import BGKSim, KBCSim
-from src.lattice import LatticeD3Q19, LatticeD3Q27
-from src.boundary_conditions import *
-import numpy as np
-from src.utils import *
-from jax import config
-import os
-#os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-import jax
-import scipy
-
-# PhantomGaze for in-situ rendering
-import phantomgaze as pg
-
-def makeNacaAirfoil(length, thickness=30, angle=0):
-    def nacaAirfoil(x, thickness, chordLength):
-        coeffs = [0.2969, -0.1260, -0.3516, 0.2843, -0.1015]
-        exponents = [0.5, 1, 2, 3, 4]
-        yt = [coeff * (x / chordLength) ** exp for coeff, exp in zip(coeffs, exponents)]
-        yt = 5. * thickness / 100 * chordLength * np.sum(yt)
-
-        return yt
-
-    x = np.linspace(0, length, num=length)
-    yt = np.array([nacaAirfoil(xi, thickness, length) for xi in x])
-
-    y_max = int(np.max(yt)) + 1
-    domain = np.zeros((2 * y_max, len(x)), dtype=int)
-
-    for i, xi in enumerate(x):
-        upper_bound = int(y_max + yt[i])
-        lower_bound = int(y_max - yt[i])
-        domain[lower_bound:upper_bound, i] = 1
-
-    domain = scipy.ndimage.rotate(domain, angle, reshape=True)
-    domain = np.where(domain > 0.5, 1, 0)
-
-    return domain
-
-class Airfoil(KBCSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        tx, ty = np.array([self.nx, self.ny], dtype=int) - airfoil.shape
-
-        airfoil_mask = np.pad(airfoil, ((tx // 3, tx - tx // 3), (ty // 2, ty - ty // 2)), 'constant', constant_values=False)
-        airfoil_mask = np.repeat(airfoil_mask[:, :, np.newaxis], self.nz, axis=2)
-        
-        airfoil_indices = np.argwhere(airfoil_mask)
-        wall = np.concatenate((airfoil_indices,
-                               self.boundingBoxIndices['bottom'], self.boundingBoxIndices['top']))
-        self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
-
-        # Store airfoil boundary for visualization
-        self.visualization_bc = jnp.zeros((self.nx, self.ny, self.nz), dtype=jnp.float32)
-        self.visualization_bc = self.visualization_bc.at[tuple(airfoil_indices.T)].set(1.0)
-
-        doNothing = self.boundingBoxIndices['right']
-        self.BCs.append(DoNothing(tuple(doNothing.T), self.gridInfo, self.precisionPolicy))
-
-        inlet = self.boundingBoxIndices['left']
-        rho_inlet = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_inlet = np.zeros((inlet.shape), dtype=self.precisionPolicy.compute_dtype)
-
-        vel_inlet[:, 0] = prescribed_vel
-        self.BCs.append(EquilibriumBC(tuple(inlet.T), self.gridInfo, self.precisionPolicy, rho_inlet, vel_inlet))
-
-    def output_data(self, **kwargs):
-        # Compute q-criterion and vorticity using finite differences
-        # Get velocity field  
-        u = kwargs['u'][..., 1:-1, :]
-        # vorticity and q-criterion
-        norm_mu, q = q_criterion(u)
-
-        # Make phantomgaze volume
-        dx = 0.01
-        origin = (0.0, 0.0, 0.0)
-        upper_bound = (self.visualization_bc.shape[0] * dx, self.visualization_bc.shape[1] * dx, self.visualization_bc.shape[2] * dx)
-        q_volume = pg.objects.Volume(
-            q,
-            spacing=(dx, dx, dx),
-            origin=origin,
-        )
-        norm_mu_volume = pg.objects.Volume(
-            norm_mu,
-            spacing=(dx, dx, dx),
-            origin=origin,
-        )
-        boundary_volume = pg.objects.Volume(
-            self.visualization_bc,
-            spacing=(dx, dx, dx),
-            origin=origin,
-        )
-
-        # Make colormap for norm_mu
-        colormap = pg.Colormap("jet", vmin=0.0, vmax=0.05)
-
-        # Get camera parameters
-        focal_point = (self.visualization_bc.shape[0] * dx / 2, self.visualization_bc.shape[1] * dx / 2, self.visualization_bc.shape[2] * dx / 2)
-        radius = 5.0
-        angle = kwargs['timestep'] * 0.0001
-        camera_position = (focal_point[0] + radius * np.sin(angle), focal_point[1], focal_point[2] + radius * np.cos(angle))
-
-        # Rotate camera 
-        camera = pg.Camera(position=camera_position, focal_point=focal_point, view_up=(0.0, 1.0, 0.0), max_depth=30.0, height=1080, width=1920, background=pg.SolidBackground(color=(0.0, 0.0, 0.0)))
-
-        # Make wireframe
-        screen_buffer = pg.render.wireframe(lower_bound=origin, upper_bound=upper_bound, thickness=0.01, camera=camera)
-
-        # Render axes
-        screen_buffer = pg.render.axes(size=0.1, center=(0.0, 0.0, 1.1), camera=camera, screen_buffer=screen_buffer)
-
-        # Render q-criterion
-        screen_buffer = pg.render.contour(q_volume, threshold=0.00003, color=norm_mu_volume, colormap=colormap, camera=camera, screen_buffer=screen_buffer)
-
-        # Render boundary
-        boundary_colormap = pg.Colormap("bone_r", vmin=0.0, vmax=3.0, opacity=np.linspace(0.0, 6.0, 256))
-        screen_buffer = pg.render.volume(boundary_volume, camera=camera, colormap=boundary_colormap, screen_buffer=screen_buffer)
-
-        # Show the rendered image
-        plt.imsave('q_criterion_' + str(kwargs['timestep']).zfill(7) + '.png', np.minimum(screen_buffer.image.get(), 1.0))
-
-
-if __name__ == '__main__':
-    airfoil_length = 101
-    airfoil_thickness = 30
-    airfoil_angle = 20
-    airfoil = makeNacaAirfoil(length=airfoil_length, thickness=airfoil_thickness, angle=airfoil_angle).T
-    precision = 'f32/f32'
-
-    lattice = LatticeD3Q27(precision)
-
-    nx = airfoil.shape[0]
-    ny = airfoil.shape[1]
-
-    ny = 3 * ny
-    nx = 5 * nx
-    nz = 101
-
-    Re = 30000.0
-    prescribed_vel = 0.1
-    clength = airfoil_length
-
-    visc = prescribed_vel * clength / Re
-    omega = 1.0 / (3. * visc + 0.5)
-    
-    os.system('rm -rf ./*.vtk && rm -rf ./*.png')
-
-    # Set the parameters for the simulation
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': nx,
-        'ny': ny,
-        'nz': nz,
-        'precision': precision,
-        'io_rate': 100,
-        'print_info_rate': 100,
-    }
-
-    sim = Airfoil(**kwargs)
-    sim.run(20000)
diff --git a/examples/CFD/cavity2d.py b/examples/CFD/cavity2d.py
deleted file mode 100644
index 28536e4..0000000
--- a/examples/CFD/cavity2d.py
+++ /dev/null
@@ -1,96 +0,0 @@
-"""
-This example implements a 2D Lid-Driven Cavity Flow simulation using the lattice Boltzmann method (LBM). 
-The Lid-Driven Cavity Flow is a standard test case for numerical schemes applied to fluid dynamics, which involves fluid in a square cavity with a moving lid (top boundary).
-
-In this example you'll be introduced to the following concepts:
-
-1. Lattice: The simulation employs a D2Q9 lattice. It's a 2D lattice model with nine discrete velocity directions, which is typically used for 2D simulations.
-
-2. Boundary Conditions: The code implements two types of boundary conditions:
-
-    BounceBackHalfway: This condition is applied to the stationary walls (left, right, and bottom). It models a no-slip boundary where the velocity of fluid at the wall is zero.
-    EquilibriumBC: This condition is used for the moving lid (top boundary). It defines a boundary with a set velocity, simulating the "driving" of the cavity by the lid.
-
-3. Checkpointing: The simulation supports checkpointing. Checkpoints are saved periodically (determined by the 'checkpoint_rate'), allowing the simulation to be stopped and restarted from the last checkpoint. This can be beneficial for long simulations or in case of unexpected interruptions.
-
-4. Visualization: The simulation outputs data in VTK format for visualization. It also provides images of the velocity field and saves the boundary conditions at each time step. The data can be visualized using software like Paraview.
-
-"""
-from jax import config
-import numpy as np
-import jax.numpy as jnp
-import os
-
-from src.boundary_conditions import *
-from src.models import BGKSim, KBCSim
-from src.lattice import LatticeD2Q9
-from src.utils import *
-
-# Use 8 CPU devices
-# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-
-class Cavity(KBCSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-
-        # concatenate the indices of the left, right, and bottom walls
-        walls = np.concatenate((self.boundingBoxIndices["left"], self.boundingBoxIndices["right"], self.boundingBoxIndices["bottom"]))
-        # apply bounce back boundary condition to the walls
-        self.BCs.append(BounceBackHalfway(tuple(walls.T), self.gridInfo, self.precisionPolicy))
-
-        # apply inlet equilibrium boundary condition to the top wall
-        moving_wall = self.boundingBoxIndices["top"]
-
-        rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype)
-        vel_wall[:, 0] = prescribed_vel
-        self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall))
-
-    def output_data(self, **kwargs):
-        # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
-        rho = np.array(kwargs["rho"][1:-1, 1:-1])
-        u = np.array(kwargs["u"][1:-1, 1:-1, :])
-        timestep = kwargs["timestep"]
-
-        save_image(timestep, u)
-        fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1]}
-        save_fields_vtk(timestep, fields)
-        save_BCs_vtk(timestep, self.BCs, self.gridInfo)
-
-if __name__ == "__main__":
-    precision = "f32/f32"
-    lattice = LatticeD2Q9(precision)
-
-    nx = 200
-    ny = 200
-
-    Re = 200.0
-    prescribed_vel = 0.1
-    clength = nx - 1
-
-    checkpoint_rate = 1000
-    checkpoint_dir = os.path.abspath("./checkpoints")
-
-    visc = prescribed_vel * clength / Re
-    omega = 1.0 / (3.0 * visc + 0.5)
-    
-    os.system("rm -rf ./*.vtk && rm -rf ./*.png")
-
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': nx,
-        'ny': ny,
-        'nz': 0,
-        'precision': precision,
-        'io_rate': 100,
-        'print_info_rate': 100,
-        'checkpoint_rate': checkpoint_rate,
-        'checkpoint_dir': checkpoint_dir,
-        'restore_checkpoint': False,
-    }
-
-    sim = Cavity(**kwargs)
-    sim.run(5000)
diff --git a/examples/CFD/cavity3d.py b/examples/CFD/cavity3d.py
deleted file mode 100644
index 2c30d28..0000000
--- a/examples/CFD/cavity3d.py
+++ /dev/null
@@ -1,123 +0,0 @@
-"""
-This example implements a 3D Lid-Driven Cavity Flow simulation using the lattice Boltzmann method (LBM). 
-The Lid-Driven Cavity Flow is a standard test case for numerical schemes applied to fluid dynamics, which involves fluid in a square cavity with a moving lid (top boundary).
-
-In this example you'll be introduced to the following concepts:
-
-1. Lattice: The simulation employs a D3Q27 lattice. It's a 3D lattice model with 27 discrete velocity directions.
-
-2. Boundary Conditions: The code implements two types of boundary conditions:
-
-    BounceBack: This condition is applied to the stationary walls, except the top wall. It models a no-slip boundary where the velocity of fluid at the wall is zero.
-    EquilibriumBC: This condition is used for the moving lid (top boundary). It defines a boundary with a set velocity, simulating the "driving" of the cavity by the lid.
-
-4. Visualization: The simulation outputs data in VTK format for visualization. The data can be visualized using software like Paraview.
-
-"""
-# Use 8 CPU devices
-# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-
-import numpy as np
-from src.utils import *
-from jax import config
-import json, codecs
-
-from src.models import BGKSim, KBCSim
-from src.lattice import LatticeD3Q19, LatticeD3Q27
-from src.boundary_conditions import *
-
-
-config.update('jax_enable_x64', True)
-
-class Cavity(KBCSim):
-    # Note: We have used BGK with D3Q19 (or D3Q27) for Re=(1000, 3200) and KBC with D3Q27 for Re=10,000
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        # Note:
-        # We have used halfway BB for Re=(1000, 3200) and regularized BC for Re=10,000
-
-        # apply inlet boundary condition to the top wall
-        moving_wall = self.boundingBoxIndices['top']
-        vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype)
-        vel_wall[:, 0] = prescribed_vel
-        # self.BCs.append(BounceBackHalfway(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, vel_wall))
-        self.BCs.append(Regularized(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall))
-
-        # concatenate the indices of the left, right, and bottom walls
-        walls = np.concatenate(
-            (self.boundingBoxIndices['left'], self.boundingBoxIndices['right'],
-             self.boundingBoxIndices['front'], self.boundingBoxIndices['back'],
-             self.boundingBoxIndices['bottom']))
-        # apply bounce back boundary condition to the walls
-        # self.BCs.append(BounceBackHalfway(tuple(walls.T), self.gridInfo, self.precisionPolicy))
-        vel_wall = np.zeros(walls.shape, dtype=self.precisionPolicy.compute_dtype)
-        self.BCs.append(Regularized(tuple(walls.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall))
-        return
-
-    def output_data(self, **kwargs):
-        # 1: -1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
-        rho = np.array(kwargs['rho'])
-        u = np.array(kwargs['u'])
-        timestep = kwargs['timestep']
-        u_prev = kwargs['u_prev']
-
-        u_old = np.linalg.norm(u_prev, axis=2)
-        u_new = np.linalg.norm(u, axis=2)
-
-        err = np.sum(np.abs(u_old - u_new))
-        print('error= {:07.6f}'.format(err))
-        fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]}
-        # save_fields_vtk(timestep, fields)
-
-        # output profiles of velocity at mid-plane for benchmarking
-        output_filename = "./profiles_" + f"{timestep:07d}.json"
-        ux_mid = 0.5*(u[nx//2, ny//2, :, 0] + u[nx//2+1, ny//2+1, :, 0])
-        uz_mid = 0.5*(u[:, ny//2, nz//2, 2] + u[:, ny//2+1, nz//2+1, 2])
-        ldc_ref_result = {'ux(x=y=0)': list(ux_mid/prescribed_vel),
-                          'uz(z=y=0)': list(uz_mid/prescribed_vel)}
-        json.dump(ldc_ref_result, codecs.open(output_filename, 'w', encoding='utf-8'),
-                separators=(',', ':'),
-                sort_keys=True,
-                indent=4)
-
-        # Calculate the velocity magnitude
-        # u_mag = np.linalg.norm(u, axis=2)
-        # live_volume_randering(timestep, u_mag)
-
-if __name__ == '__main__':
-    # Note: 
-    # We have used BGK with D3Q19 (or D3Q27) for Re=(1000, 3200) and KBC with D3Q27 for Re=10,000
-    precision = 'f64/f64'
-    lattice = LatticeD3Q27(precision)
-
-    nx = 256
-    ny = 256
-    nz = 256
-
-    Re = 10000.0
-    prescribed_vel = 0.06
-    clength = nx - 2
-
-    # characteristic time
-    tc = prescribed_vel/clength
-    niter_max = int(500//tc)
-
-    visc = prescribed_vel * clength / Re
-    omega = 1.0 / (3. * visc + 0.5)
-    os.system("rm -rf ./*.vtk && rm -rf ./*.png")
-
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': nx,
-        'ny': ny,
-        'nz': nz,
-        'precision': precision,
-        'io_rate': int(10//tc),
-        'print_info_rate': int(10//tc),
-        'downsampling_factor': 1
-    }
-    sim = Cavity(**kwargs)
-    sim.run(niter_max)
\ No newline at end of file
diff --git a/examples/CFD/channel3d.py b/examples/CFD/channel3d.py
deleted file mode 100644
index 2c7ab73..0000000
--- a/examples/CFD/channel3d.py
+++ /dev/null
@@ -1,157 +0,0 @@
-"""
-This script performs a 3D simulation of turbulent channel flow using the lattice Boltzmann method (LBM). 
-Turbulent channel flow, also known as plane Couette flow, is a fundamental case in the study of wall-bounded turbulent flows.
-
-In this example you'll be introduced to the following concepts:
-
-1. Lattice: A D3Q27 lattice is used, which is a three-dimensional lattice model with 27 discrete velocity directions. This type of lattice allows for a more precise representation of fluid flow in three dimensions.
-
-2. Initial Conditions: The initial conditions for the flow are randomly generated, and the populations are initialized to be the solution of an advection-diffusion equation.
-
-3. Boundary Conditions: Bounce back boundary conditions are applied at the top and bottom walls, simulating a no-slip condition typical for wall-bounded flows.
-
-4. External Force: An external force is applied to drive the flow.
-
-"""
-
-from src.boundary_conditions import *
-from jax import config
-from src.utils import *
-import numpy as np
-from src.lattice import LatticeD3Q27
-from src.models import KBCSim, AdvectionDiffusionBGK
-import jax.numpy as jnp
-import os
-import matplotlib.pyplot as plt
-
-# Use 8 CPU devices
-# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-import jax
-
-# disable JIt compilation
-
-jax.config.update('jax_enable_x64', True)
-
-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
-    """
-    dns_dic = {
-        "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,0.0000707,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.2170e-13,-7.3193e-10,-1.5832e-07,-3.7598e-06,-3.3837e-05,-1.7683e-04,-6.5008e-04,-1.8650e-03,-4.4488e-03,-9.2047e-03,-1.7023e-02,-2.8777e-02,-4.5228e-02,-6.6952e-02,-9.4281e-02,-1.2724e-01,-1.6551e-01,-2.0842e-01,-2.5498e-01,-3.0396e-01,-3.5398e-01,-4.0362e-01,-4.5163e-01,-4.9698e-01,-5.3880e-01,-5.7639e-01,-6.0919e-01,-6.3686e-01,-6.5930e-01,-6.7652e-01,-6.8867e-01,-6.9613e-01,-6.9928e-01,-6.9854e-01,-6.9444e-01,-6.8744e-01,-6.7802e-01,-6.6675e-01,-6.5429e-01,-6.4131e-01,-6.2817e-01,-6.1487e-01,-6.0122e-01,-5.8703e-01,-5.7221e-01,-5.5678e-01,-5.4090e-01,-5.2493e-01,-5.0917e-01,-4.9371e-01,-4.7867e-01,-4.6421e-01,-4.5050e-01,-4.3759e-01,-4.2550e-01,-4.1436e-01,-4.0444e-01,-3.9595e-01,-3.8900e-01,-3.8360e-01,-3.7966e-01,-3.7702e-01,-3.7542e-01,-3.7460e-01,-3.7436e-01]
-        }
-    return dns_dic
-
-class TurbulentChannel(KBCSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        # top and bottom sides of the channel are no-slip and the other directions are periodic
-        wall = np.concatenate((self.boundingBoxIndices['bottom'], self.boundingBoxIndices['top']))
-        self.BCs.append(Regularized(tuple(wall.T), self.gridInfo, self.precisionPolicy, 'velocity', np.zeros((wall.shape[0], 3))))
-        return
-
-    def initialize_macroscopic_fields(self):
-        rho = self.precisionPolicy.cast_to_output(1.0)
-        u = self.distributed_array_init((self.nx, self.ny, self.nz, self.dim),
-                                         self.precisionPolicy.compute_dtype, init_val=1e-2 * np.random.random((self.nx, self.ny, self.nz, self.dim)))
-        u = self.precisionPolicy.cast_to_output(u)
-        return rho, u
-
-    def initialize_populations(self, rho, u):
-        omegaADE = 1.0
-        lattice = LatticeD3Q27(precision)
-
-        kwargs = {'lattice': lattice, 'nx': self.nx, 'ny': self.ny, 'nz': self.nz,  'precision': precision, 'omega': omegaADE, 'vel': u}
-        ADE = AdvectionDiffusionBGK(**kwargs)
-        ADE.initialize_macroscopic_fields = self.initialize_macroscopic_fields
-        print("Initializing the distribution functions using the specified macroscopic fields....")
-        f = ADE.run(50000)
-        return f
-
-    def get_force(self):
-        # define the external force
-        force = np.zeros((self.nx, self.ny, self.nz, 3))
-        force[..., 0] = Re_tau**2 * visc**2 / h**3
-        return self.precisionPolicy.cast_to_output(force)
-
-    def output_data(self, **kwargs):
-        rho = np.array(kwargs["rho"])
-        u = np.array(kwargs["u"])
-        timestep = kwargs["timestep"]
-        u_prev = kwargs['u_prev']
-
-        u_old = np.linalg.norm(u_prev, axis=2)
-        u_new = np.linalg.norm(u, axis=2)
-
-        err = np.sum(np.abs(u_old - u_new))
-        print("error= {:07.6f}".format(err))
-
-        # mean streamwise velocity in wall units u^+(z)
-        uplus = np.mean(u[..., 0], axis=(0,1))/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) + '.pdf'
-        plt.savefig(fname, format='pdf')
-        fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]}
-        save_fields_vtk(timestep, fields)
-
-
-
-if __name__ == "__main__":
-    precision = "f64/f64"
-    lattice = LatticeD3Q27(precision)
-
-    # h: channel half-width
-    h = 50
-
-    # Define channel geometry based on h
-    nx = 6*h
-    ny = 3*h
-    nz = 2*h
-
-    # Define flow regime
-    Re_tau = 180
-    u_tau = 0.001
-    DeltaPlus = Re_tau/h    # DeltaPlus = u_tau / nu * Delta where u_tau / nu = Re_tau/h
-    visc = u_tau * h / Re_tau
-    omega = 1.0 / (3.0 * visc + 0.5)
-
-    # Wall distance in wall units to be used inside output_data
-    zz = np.arange(nz)
-    zz = np.minimum(zz, zz.max() - zz)
-    yplus = zz * u_tau / visc
-
-    os.system("rm -rf ./*.vtk && rm -rf ./*.png")
-
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': nx,
-        'ny': ny,
-        'nz': nz,
-        'precision': precision,
-        'io_rate': 500000,
-        'print_info_rate': 100000
-    }
-    sim = turbulentChannel(**kwargs)
-    sim.run(10000000)
diff --git a/examples/CFD/couette2d.py b/examples/CFD/couette2d.py
deleted file mode 100644
index 1c15a6a..0000000
--- a/examples/CFD/couette2d.py
+++ /dev/null
@@ -1,79 +0,0 @@
-"""
-This script performs a 2D simulation of Couette flow using the lattice Boltzmann method (LBM). 
-"""
-
-import os
-import jax.numpy as jnp
-import numpy as np
-from src.utils import *
-from jax import config
-
-
-from src.models import BGKSim
-from src.boundary_conditions import *
-from src.lattice import LatticeD2Q9
-
-# config.update('jax_disable_jit', True)
-# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=4'
-
-class Couette(BGKSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        walls = np.concatenate((self.boundingBoxIndices["top"], self.boundingBoxIndices["bottom"]))
-        self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy))
-
-        outlet = self.boundingBoxIndices["right"]
-        inlet = self.boundingBoxIndices["left"]
-
-        rho_wall = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_wall = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype)
-        vel_wall[:, 0] = prescribed_vel
-        self.BCs.append(EquilibriumBC(tuple(inlet.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall))
-
-        self.BCs.append(DoNothing(tuple(outlet.T), self.gridInfo, self.precisionPolicy))
-
-    def output_data(self, **kwargs):
-        # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
-        rho = np.array(kwargs["rho"][..., 1:-1, :])
-        u = np.array(kwargs["u"][..., 1:-1, :])
-        timestep = kwargs["timestep"]
-        u_prev = kwargs["u_prev"][..., 1:-1, :]
-
-        u_old = np.linalg.norm(u_prev, axis=2)
-        u_new = np.linalg.norm(u, axis=2)
-        err = np.sum(np.abs(u_old - u_new))
-        print("error= {:07.6f}".format(err))
-        save_image(timestep, u)
-        fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1]}
-        save_fields_vtk(timestep, fields)
-
-if __name__ == "__main__":
-    precision = "f32/f32"
-    lattice = LatticeD2Q9(precision)
-    nx = 501
-    ny = 101
-
-    Re = 100.0
-    prescribed_vel = 0.1
-    clength = nx - 1
-
-    visc = prescribed_vel * clength / Re
-
-    omega = 1.0 / (3.0 * visc + 0.5)
-    assert omega < 1.98, "omega must be less than 2.0"
-    os.system("rm -rf ./*.vtk && rm -rf ./*.png")
-
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': nx,
-        'ny': ny,
-        'nz': 0,
-        'precision': precision,
-        'io_rate': 100,
-        'print_info_rate': 100
-        }
-    sim = Couette(**kwargs)
-    sim.run(20000)
diff --git a/examples/CFD/cylinder2d.py b/examples/CFD/cylinder2d.py
deleted file mode 100644
index 2c9887d..0000000
--- a/examples/CFD/cylinder2d.py
+++ /dev/null
@@ -1,148 +0,0 @@
-"""
-This script conducts a 2D simulation of flow around a cylinder using the lattice Boltzmann method (LBM). This is a classic problem in fluid dynamics and is often used to examine the behavior of fluid flow over a bluff body.
-
-In this example you'll be introduced to the following concepts:
-
-1. Lattice: A D2Q9 lattice is used, which is a two-dimensional lattice model with nine discrete velocity directions. This type of lattice allows for a precise representation of fluid flow in two dimensions.
-
-2. Boundary Conditions: The script implements several types of boundary conditions:
-
-    BounceBackHalfway: This condition is applied to the cylinder surface, simulating a no-slip condition where the fluid at the cylinder surface has zero velocity.
-    ExtrapolationOutflow: This condition is applied at the outlet (right boundary), where the fluid is allowed to exit the simulation domain freely.
-    Regularized: This condition is applied at the inlet (left boundary) and models the inflow of fluid into the domain with a specified velocity profile. Another Regularized condition is used for the stationary top and bottom walls.
-3. Velocity Profile: The script uses a Poiseuille flow profile for the inlet velocity. This is a parabolic profile commonly seen in pipe flow.
-
-4. Drag and lift calculation: The script computes the lift and drag on the cylinder, which are important quantities in fluid dynamics and aerodynamics.
-
-5. Visualization: The simulation outputs data in VTK format for visualization. It also generates images of the velocity field. The data can be visualized using software like ParaView.
-
-# To run type:
-nohup python3 examples/CFD/cylinder2d.py > logfile.log &
-"""
-import os
-import json
-import jax
-from time import time
-from jax import config
-import numpy as np
-import jax.numpy as jnp
-
-from src.utils import *
-from src.boundary_conditions import *
-from src.models import BGKSim, KBCSim
-from src.lattice import LatticeD2Q9
-
-# Use 8 CPU devices
-# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-jax.config.update('jax_enable_x64', True)
-
-class Cylinder(BGKSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        # Define the cylinder surface
-        coord = np.array([(i, j) for i in range(self.nx) for j in range(self.ny)])
-        xx, yy = coord[:, 0], coord[:, 1]
-        cx, cy = 2.*diam, 2.*diam
-        cylinder = (xx - cx)**2 + (yy-cy)**2 <= (diam/2.)**2
-        cylinder = coord[cylinder]
-        implicit_distance = np.reshape((xx - cx)**2 + (yy-cy)**2 - (diam/2.)**2, (self.nx, self.ny))
-        self.BCs.append(InterpolatedBounceBackBouzidi(tuple(cylinder.T), implicit_distance, self.gridInfo, self.precisionPolicy))
-
-        # Outflow BC
-        outlet = self.boundingBoxIndices['right']
-        rho_outlet = np.ones((outlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        self.BCs.append(ExtrapolationOutflow(tuple(outlet.T), self.gridInfo, self.precisionPolicy))
-        # self.BCs.append(ZouHe(tuple(outlet.T), self.gridInfo, self.precisionPolicy, 'pressure', rho_outlet))
-
-        # Inlet BC
-        inlet = self.boundingBoxIndices['left']
-        rho_inlet = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_inlet = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype)
-        yy_inlet = yy.reshape(self.nx, self.ny)[tuple(inlet.T)]
-        vel_inlet[:, 0] = poiseuille_profile(yy_inlet,
-                                             yy_inlet.min(),
-                                             yy_inlet.max()-yy_inlet.min(), 3.0 / 2.0 * prescribed_vel)
-        self.BCs.append(Regularized(tuple(inlet.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_inlet))
-
-        # No-slip BC for top and bottom
-        wall = np.concatenate([self.boundingBoxIndices['top'], self.boundingBoxIndices['bottom']])
-        vel_wall = np.zeros(wall.shape, dtype=self.precisionPolicy.compute_dtype)
-        self.BCs.append(Regularized(tuple(wall.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall))
-
-    def output_data(self, **kwargs):
-        # 1:-1 to remove boundary voxels (not needed for visualization when using bounce-back)
-        rho = np.array(kwargs["rho"][..., 1:-1, :])
-        u = np.array(kwargs["u"][..., 1:-1, :])
-        timestep = kwargs["timestep"]
-        u_prev = kwargs["u_prev"][..., 1:-1, :]
-
-        if timestep == 0:
-            self.CL_max = 0.0
-            self.CD_max = 0.0
-        if timestep > 0.5 * niter_max:
-            # compute lift and drag over the cyliner
-            cylinder = self.BCs[0]
-            boundary_force = cylinder.momentum_exchange_force(kwargs['f_poststreaming'], kwargs['f_postcollision'])
-            boundary_force = np.sum(np.array(boundary_force), axis=0)
-            drag = boundary_force[0]
-            lift = boundary_force[1]
-            cd = 2. * drag / (prescribed_vel ** 2 * diam)
-            cl = 2. * lift / (prescribed_vel ** 2 * diam)
-
-            u_old = np.linalg.norm(u_prev, axis=2)
-            u_new = np.linalg.norm(u, axis=2)
-            err = np.sum(np.abs(u_old - u_new))
-            self.CL_max = max(self.CL_max, cl)
-            self.CD_max = max(self.CD_max, cd)
-            print('error= {:07.6f}, CL = {:07.6f}, CD = {:07.6f}'.format(err, cl, cd))
-            # save_image(timestep, u)
-
-# Helper function to specify a parabolic poiseuille profile
-poiseuille_profile  = lambda x,x0,d,umax: np.maximum(0.,4.*umax/(d**2)*((x-x0)*d-(x-x0)**2))
-
-if __name__ == '__main__':
-    precision = 'f64/f64'
-    # diam_list = [10, 20, 30, 40, 60, 80]
-    diam_list = [80]
-    CL_list, CD_list = [], []
-    result_dict = {}
-    result_dict['resolution_list'] = diam_list
-    for diam in diam_list:
-        scale_factor = 80 / diam
-        prescribed_vel = 0.003 * scale_factor
-        lattice = LatticeD2Q9(precision)
-
-        nx = int(22*diam)
-        ny = int(4.1*diam)
-
-        Re = 100.0
-        visc = prescribed_vel * diam / Re
-        omega = 1.0 / (3. * visc + 0.5)
-        
-        os.system('rm -rf ./*.vtk && rm -rf ./*.png')
-
-        kwargs = {
-            'lattice': lattice,
-            'omega': omega,
-            'nx': nx,
-            'ny': ny,
-            'nz': 0,
-            'precision': precision,
-            'io_rate': int(500 / scale_factor),
-            'print_info_rate': int(10000 / scale_factor),
-            'return_fpost': True    # Need to retain fpost-collision for computation of lift and drag
-        }
-        # characteristic time
-        tc = prescribed_vel/diam
-        niter_max = int(100//tc)
-        sim = Cylinder(**kwargs)
-        sim.run(niter_max)
-        CL_list.append(sim.CL_max)
-        CD_list.append(sim.CD_max)
-
-    result_dict['CL'] = CL_list
-    result_dict['CD'] = CD_list
-    with open('data.json', 'w') as fp:
-        json.dump(result_dict, fp)
diff --git a/examples/CFD/oscilating_cylinder2d.py b/examples/CFD/oscilating_cylinder2d.py
deleted file mode 100644
index 97d6746..0000000
--- a/examples/CFD/oscilating_cylinder2d.py
+++ /dev/null
@@ -1,146 +0,0 @@
-"""
-This script conducts a 2D simulation of flow around a cylinder using the lattice Boltzmann method (LBM). This is a classic problem in fluid dynamics and is often used to examine the behavior of fluid flow over a bluff body.
-
-In this example you'll be introduced to the following concepts:
-
-1. Lattice: A D2Q9 lattice is used, which is a two-dimensional lattice model with nine discrete velocity directions. This type of lattice allows for a precise representation of fluid flow in two dimensions.
-
-2. Boundary Conditions: The script implements several types of boundary conditions:
-
-    BounceBackMoving: This condition is applied to the cylinder surface. Unlike the usual BounceBack condition, this one takes into account the motion of the cylinder.
-    ExtrapolationOutflow: This condition is applied at the outlet (right boundary), where the fluid is allowed to exit the simulation domain freely.
-    Regularized: This condition is applied at the inlet (left boundary) and models the inflow of fluid into the domain with a specified velocity profile. Another Regularized condition is used for the stationary top and bottom walls.
-3. Velocity Profile: The script uses a Poiseuille flow profile for the inlet velocity. This is a parabolic profile commonly seen in pipe flow.
-
-4. Drag and lift calculation: The script computes the lift and drag on the cylinder, which are important quantities in fluid dynamics and aerodynamics.
-
-5. Visualization: The simulation outputs data in VTK format for visualization. It also generates images of the velocity field. The data can be visualized using software like ParaView.
-
-"""
-
-
-import os
-import jax
-from time import time
-from jax import config
-import numpy as np
-import jax.numpy as jnp
-
-from src.utils import *
-from src.boundary_conditions import *
-from src.models import BGKSim, KBCSim
-from src.lattice import LatticeD2Q9
-
-# Use 8 CPU devices
-# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-jax.config.update('jax_enable_x64', True)
-
-class Cylinder(KBCSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        wall = np.concatenate([self.boundingBoxIndices['top'], self.boundingBoxIndices['bottom']])
-        self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
-
-        coord = np.array([np.unravel_index(i, (self.nx, self.ny)) for i in range(self.nx*self.ny)]) 
-        xx, yy = coord[:, 0], coord[:, 1]
-        cx, cy = 2.*diam, 2.*diam
-        cyl = ((xx) - cx)**2 + (yy-cy)**2 <= (diam/2.)**2
-        cyl = jnp.array(coord[cyl])
-    
-        # Define update rules for boundary conditions
-        def update_function(time: int):
-            # Move the cylinder up and down sinusoidally with time
-            # Define the scale for the sinusoidal motion
-            scale = 10000
-            
-            # Amplitude of the motion, a quarter of the y-dimension of the grid
-            A = ny // 4
-            
-            # Calculate the new y-coordinates of the cylinder. The cylinder moves up and down, 
-            # its motion dictated by the sinusoidal function. We use `astype(int)` to ensure 
-            # the indices are integers, as they will be used for array indexing.
-            new_y_coords = cyl[:, 1] + jnp.array((jnp.sin(time/scale)*A).astype(int))
-            
-            # Define the indices of the grid points occupied by the cylinder
-            indices = (cyl[:, 0], new_y_coords)
-            
-            # Calculate the velocity of the cylinder. The x-component is always 0 (the cylinder 
-            # doesn't move horizontally), and the y-component is the derivative of the sinusoidal 
-            # function governing the cylinder's motion, scaled by the amplitude and the scale factor.
-            velocity = jnp.array([0., jnp.cos(time/scale)* A / scale], dtype=self.precisionPolicy.compute_dtype)
-            
-            return indices, velocity
-
-        self.BCs.append(BounceBackMoving(self.gridInfo, self.precisionPolicy, update_function=update_function))
-
-
-        outlet = self.boundingBoxIndices['right']
-        self.BCs.append(ExtrapolationOutflow(tuple(outlet.T), self.gridInfo, self.precisionPolicy))
-
-        inlet = self.boundingBoxIndices['left']
-        vel_inlet = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype)
-        yy_inlet = yy.reshape(self.nx, self.ny)[tuple(inlet.T)]
-        vel_inlet[:, 0] = poiseuille_profile(yy_inlet,
-                                             yy_inlet.min(),
-                                             yy_inlet.max()-yy_inlet.min(), 3.0 / 2.0 * prescribed_vel)
-        self.BCs.append(Regularized(tuple(inlet.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_inlet))
-
-
-    def output_data(self, **kwargs):
-        # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
-        rho = np.array(kwargs["rho"][..., 1:-1, :])
-        u = np.array(kwargs["u"][..., 1:-1, :])
-        timestep = kwargs["timestep"]
-        u_prev = kwargs["u_prev"][..., 1:-1, :]
-
-        # compute lift and drag over the cyliner
-        cylinder = self.BCs[0]
-        boundary_force = cylinder.momentum_exchange_force(kwargs['f_poststreaming'], kwargs['f_postcollision'])
-        boundary_force = np.sum(boundary_force, axis=0)
-        drag = boundary_force[0]
-        lift = boundary_force[1]
-        cd = 2. * drag / (prescribed_vel ** 2 * diam)
-        cl = 2. * lift / (prescribed_vel ** 2 * diam)
-
-        u_old = np.linalg.norm(u_prev, axis=2)
-        u_new = np.linalg.norm(u, axis=2)
-        err = np.sum(np.abs(u_old - u_new))
-        print('error= {:07.6f}, CL = {:07.6f}, CD = {:07.6f}'.format(err, cl, cd))
-        save_image(timestep, u)
-        # u magnitude
-        fields = {'rho': rho[..., 0], 'u': np.linalg.norm(u, axis=2)}
-        save_fields_vtk(timestep, fields)
-        save_BCs_vtk(timestep, self.BCs, self.gridInfo)
-
-# Helper function to specify a parabolic poiseuille profile
-poiseuille_profile  = lambda x,x0,d,umax: np.maximum(0.,4.*umax/(d**2)*((x-x0)*d-(x-x0)**2))
-
-if __name__ == '__main__':
-    precision = 'f64/f64'
-    lattice = LatticeD2Q9(precision)
-    prescribed_vel = 0.005
-    diam = 20
-    nx = int(22*diam)
-    ny = int(4.1*diam)
-
-    Re = 10.0
-    visc = prescribed_vel * diam / Re
-    omega = 1.0 / (3. * visc + 0.5)
-
-    os.system('rm -rf ./*.vtk && rm -rf ./*.png')
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': nx,
-        'ny': ny,
-        'nz': 0,
-        'precision': precision,
-        'io_rate': 500,
-        'print_info_rate': 500,
-        'return_fpost': True,    # Need to retain fpost-collision for computation of lift and drag
-    }
-    sim = Cylinder(**kwargs)
-
-    sim.run(1000000)
diff --git a/examples/CFD/taylor_green_vortex.py b/examples/CFD/taylor_green_vortex.py
deleted file mode 100644
index 374c499..0000000
--- a/examples/CFD/taylor_green_vortex.py
+++ /dev/null
@@ -1,127 +0,0 @@
-"""
-The given script sets up a simulation for the Taylor-Green vortex flow.
-The Taylor-Green vortex is a type of two-dimensional, incompressible fluid flow with a known analytical solution, making it an ideal test case for fluid dynamics simulations.
-The flow is characterized by a pair of counter-rotating vortices. In this script, the initial fields for the Taylor-Green vortex are set using a known function.
-"""
-
-
-import os
-import json
-import jax
-import numpy as np
-import matplotlib.pyplot as plt
-
-from src.utils import *
-from src.boundary_conditions import *
-from src.models import BGKSim, KBCSim, AdvectionDiffusionBGK
-from src.lattice import LatticeD2Q9
-
-
-# Use 8 CPU devices
-# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-# disable JIt compilation
-
-jax.config.update('jax_enable_x64', True)
-
-def taylor_green_initial_fields(xx, yy, u0, rho0, nu, time):
-    ux = u0 * np.sin(xx) * np.cos(yy) * np.exp(-2 * nu * time)
-    uy = -u0 * np.cos(xx) * np.sin(yy) * np.exp(-2 * nu * time)
-    rho = 1.0 - rho0 * u0 ** 2 / 12. * (np.cos(2. * xx) + np.cos(2. * yy)) * np.exp(-4 * nu * time)
-    return ux, uy, np.expand_dims(rho, axis=-1)
-
-class TaylorGreenVortex(KBCSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        # no boundary conditions implying periodic BC in all directions
-        return
-
-    def initialize_macroscopic_fields(self):
-        ux, uy, rho = taylor_green_initial_fields(xx, yy, vel_ref, 1, 0., 0.)
-        rho = self.distributed_array_init(rho.shape, self.precisionPolicy.output_dtype, init_val=1.0, sharding=self.sharding)
-        u = np.stack([ux, uy], axis=-1)
-        u = self.distributed_array_init(u.shape, self.precisionPolicy.output_dtype, init_val=u, sharding=self.sharding)
-        return rho, u
-
-    def initialize_populations(self, rho, u):
-        omegaADE = 1.0
-        kwargs = {'lattice': lattice, 'nx': self.nx, 'ny': self.ny, 'nz': self.nz,  'precision': precision, 'omega': omegaADE, 'vel': u, 'print_info_rate': 0, 'io_rate': 0}
-        ADE = AdvectionDiffusionBGK(**kwargs)
-        ADE.initialize_macroscopic_fields = self.initialize_macroscopic_fields
-        print("Initializing the distribution functions using the specified macroscopic fields....")
-        f = ADE.run(int(20000*nx/32))
-        return f
-
-    def output_data(self, **kwargs):
-        # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
-        rho = np.array(kwargs["rho"])
-        u = np.array(kwargs["u"])
-        timestep = kwargs["timestep"]
-
-        # theoretical results
-        time = timestep * (kx**2 + ky**2)/2.
-        ux_th, uy_th, rho_th = taylor_green_initial_fields(xx, yy, vel_ref, 1, visc, time)
-        vel_err_L2 = np.sqrt(np.sum((u[..., 0]-ux_th)**2 + (u[..., 1]-uy_th)**2) / np.sum(ux_th**2 + uy_th**2))
-        rho_err_L2 = np.sqrt(np.sum((rho - rho_th)**2) / np.sum(rho_th**2))
-        print("Vel error= {:07.6f}, Pressure error= {:07.6f}".format(vel_err_L2, rho_err_L2))
-        if timestep == endTime:
-            ErrL2ResList.append(vel_err_L2)
-            ErrL2ResListRho.append(rho_err_L2)
-        # save_image(timestep, u)
-
-
-if __name__ == "__main__":
-    precision_list = ["f32/f32", "f64/f32", "f64/f64"]
-    resList = [32, 64, 128, 256, 512, 1024]
-    result_dict = dict.fromkeys(precision_list)
-    result_dict['resolution_list'] = resList
-
-    for precision in precision_list:
-        lattice = LatticeD2Q9(precision)
-        ErrL2ResList = []
-        ErrL2ResListRho = []
-        result_dict[precision] = dict.fromkeys(['vel_error', 'rho_error'])
-        for nx in resList:
-            ny = nx
-            twopi = 2.0 * np.pi
-            coord = np.array([(i, j) for i in range(nx) for j in range(ny)])
-            xx, yy = coord[:, 0], coord[:, 1]
-            kx, ky = twopi / nx, twopi / ny
-            xx = xx.reshape((nx, ny)) * kx
-            yy = yy.reshape((nx, ny)) * ky
-
-            Re = 1600.0
-            vel_ref = 0.04*32/nx
-
-            visc = vel_ref * nx / Re
-            omega = 1.0 / (3.0 * visc + 0.5)
-            os.system("rm -rf ./*.vtk && rm -rf ./*.png")
-            kwargs = {
-                'lattice': lattice,
-                'omega': omega,
-                'nx': nx,
-                'ny': ny,
-                'nz': 0,
-                'precision': precision,
-                'io_rate': 5000,
-                'print_info_rate': 1000
-            }
-            sim = TaylorGreenVortex(**kwargs)
-            tc = 2.0/(2. * visc * (kx**2 + ky**2))
-            endTime = int(0.05*tc)
-            sim.run(endTime)
-        result_dict[precision]['vel_error'] = ErrL2ResList
-        result_dict[precision]['rho_error'] = ErrL2ResListRho
-
-    with open('data.json', 'w') as fp:
-        json.dump(result_dict, fp)
-
-    # plt.loglog(resList, ErrL2ResList, '-o')
-    # plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--')
-    # plt.savefig('ErrorVel.png'); plt.savefig('ErrorVel.pdf', format='pdf')
-
-    # plt.figure()
-    # plt.loglog(resList, ErrL2ResListRho, '-o')
-    # plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--')
-    # plt.savefig('ErrorRho.png'); plt.savefig('ErrorRho.pdf', format='pdf')
diff --git a/examples/CFD/windtunnel3d.py b/examples/CFD/windtunnel3d.py
deleted file mode 100644
index 081f9b5..0000000
--- a/examples/CFD/windtunnel3d.py
+++ /dev/null
@@ -1,140 +0,0 @@
-"""
-This script performs a Lattice Boltzmann Method (LBM) simulation of fluid flow over a car model. Here are the main concepts and steps in the simulation:
-
-Here are the main concepts introduced simulation:
-
-1. Lattice: Given the usually high Reynolds number required for these simulations, a D3Q27 lattice is used, which is a three-dimensional lattice model with 27 discrete velocity directions.
-
-2. Loading geometry and voxelization: The geometry of the car is loaded from a STL file. 
-This is a file format commonly used for 3D models. The model is then voxelized to a binary matrix which represents the presence or absence of the object in the lattice. We use the DrivAer model, which is a common car model used for aerodynamic simulations.
-
-3. Output: After each specified number of iterations, the script outputs the state of the simulation. This includes the error (difference between consecutive velocity fields), lift and drag coefficients, and visualization files in the VTK format.
-"""
-
-
-import os
-import jax
-import trimesh
-from time import time
-import numpy as np
-import jax.numpy as jnp
-from jax import config
-
-from src.utils import *
-from src.models import BGKSim, KBCSim
-from src.lattice import LatticeD3Q19, LatticeD3Q27
-from src.boundary_conditions import *
-
-# Use 8 CPU devices
-# os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-
-# disable JIt compilation
-
-jax.config.update('jax_array', True)
-
-class Car(KBCSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def voxelize_stl(self, stl_filename, length_lbm_unit):
-        mesh = trimesh.load_mesh(stl_filename, process=False)
-        length_phys_unit = mesh.extents.max()
-        pitch = length_phys_unit/length_lbm_unit
-        mesh_voxelized = mesh.voxelized(pitch=pitch)
-        mesh_matrix = mesh_voxelized.matrix
-        return mesh_matrix, pitch
-
-    def set_boundary_conditions(self):
-        print('Voxelizing mesh...')
-        time_start = time()
-        stl_filename = 'stl-files/DrivAer-Notchback.stl'
-        car_length_lbm_unit = self.nx / 4
-        car_voxelized, pitch = voxelize_stl(stl_filename, car_length_lbm_unit)
-        car_matrix = car_voxelized.matrix
-        print('Voxelization time for pitch={}: {} seconds'.format(pitch, time() - time_start))
-        print("Car matrix shape: ", car_matrix.shape)
-
-        self.car_area = np.prod(car_matrix.shape[1:])
-        tx, ty, tz = np.array([nx, ny, nz]) - car_matrix.shape
-        shift = [tx//4, ty//2, 0]
-        car_indices = np.argwhere(car_matrix) + shift
-        self.BCs.append(BounceBackHalfway(tuple(car_indices.T), self.gridInfo, self.precisionPolicy))
-
-        wall = np.concatenate((self.boundingBoxIndices['bottom'], self.boundingBoxIndices['top'],
-                               self.boundingBoxIndices['front'], self.boundingBoxIndices['back']))
-        self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
-
-        doNothing = self.boundingBoxIndices['right']
-        self.BCs.append(DoNothing(tuple(doNothing.T), self.gridInfo, self.precisionPolicy))
-        self.BCs[-1].implementationStep = 'PostCollision'
-        # rho_outlet = np.ones(doNothing.shape[0], dtype=self.precisionPolicy.compute_dtype)
-        # self.BCs.append(ZouHe(tuple(doNothing.T),
-        #                                          self.gridInfo,
-        #                                          self.precisionPolicy,
-        #                                          'pressure', rho_outlet))
-
-        inlet = self.boundingBoxIndices['left']
-        rho_inlet = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_inlet = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype)
-
-        vel_inlet[:, 0] = prescribed_vel
-        self.BCs.append(EquilibriumBC(tuple(inlet.T), self.gridInfo, self.precisionPolicy, rho_inlet, vel_inlet))
-        # self.BCs.append(ZouHe(tuple(inlet.T),
-        #                                          self.gridInfo,
-        #                                          self.precisionPolicy,
-        #                                          'velocity', vel_inlet))
-
-    def output_data(self, **kwargs):
-        # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
-        rho = np.array(kwargs['rho'][..., 1:-1, 1:-1, :])
-        u = np.array(kwargs['u'][..., 1:-1, 1:-1, :])
-        timestep = kwargs['timestep']
-        u_prev = kwargs['u_prev'][..., 1:-1, 1:-1, :]
-
-        # compute lift and drag over the car
-        car = self.BCs[0]
-        boundary_force = car.momentum_exchange_force(kwargs['f_poststreaming'], kwargs['f_postcollision'])
-        boundary_force = np.sum(boundary_force, axis=0)
-        drag = np.sqrt(boundary_force[0]**2 + boundary_force[1]**2)     #xy-plane
-        lift = boundary_force[2]                                        #z-direction
-        cd = 2. * drag / (prescribed_vel ** 2 * self.car_area)
-        cl = 2. * lift / (prescribed_vel ** 2 * self.car_area)
-
-        u_old = np.linalg.norm(u_prev, axis=2)
-        u_new = np.linalg.norm(u, axis=2)
-
-        err = np.sum(np.abs(u_old - u_new))
-        print('error= {:07.6f}, CL = {:07.6f}, CD = {:07.6f}'.format(err, cl, cd))
-        fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]}
-        save_fields_vtk(timestep, fields)
-
-if __name__ == '__main__':
-    precision = 'f32/f32'
-    lattice = LatticeD3Q27(precision)
-
-    nx = 601
-    ny = 351
-    nz = 251
-
-    Re = 50000.0
-    prescribed_vel = 0.05
-    clength = nx - 1
-
-    visc = prescribed_vel * clength / Re
-    omega = 1.0 / (3. * visc + 0.5)
-
-    os.system('rm -rf ./*.vtk && rm -rf ./*.png')
-
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': nx,
-        'ny': ny,
-        'nz': nz,
-        'precision': precision,
-        'io_rate': 100,
-        'print_info_rate': 100,
-        'return_fpost': True  # Need to retain fpost-collision for computation of lift and drag
-    }
-    sim = Car(**kwargs)
-    sim.run(200000)
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
new file mode 100644
index 0000000..1b5905e
--- /dev/null
+++ b/examples/cfd/flow_past_sphere_3d.py
@@ -0,0 +1,145 @@
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.precision_policy import PrecisionPolicy
+from xlb.helper import create_nse_fields, initialize_eq, check_bc_overlaps
+from xlb.operator.stepper import IncompressibleNavierStokesStepper
+from xlb.operator.boundary_condition import (
+    FullwayBounceBackBC,
+    HalfwayBounceBackBC,
+    ZouHeBC,
+    RegularizedBC,
+    EquilibriumBC,
+    DoNothingBC,
+    ExtrapolationOutflowBC,
+)
+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 time
+
+
+class FlowOverSphere:
+    def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
+        # initialize backend
+        xlb.init(
+            velocity_set=velocity_set,
+            default_backend=backend,
+            default_precision_policy=precision_policy,
+        )
+
+        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(omega)
+
+    def _setup(self, omega):
+        self.setup_boundary_conditions()
+        self.setup_boundary_masker()
+        self.initialize_fields()
+        self.setup_stepper(omega)
+
+    def define_boundary_indices(self):
+        box = self.grid.bounding_box_indices()
+        box_no_edge = self.grid.bounding_box_indices(remove_edges=True)
+        inlet = box_no_edge["left"]
+        outlet = box_no_edge["right"]
+        walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(self.velocity_set.d)]
+        walls = np.unique(np.array(walls), axis=-1).tolist()
+
+        sphere_radius = self.grid_shape[1] // 12
+        x = np.arange(self.grid_shape[0])
+        y = np.arange(self.grid_shape[1])
+        z = np.arange(self.grid_shape[2])
+        X, Y, Z = np.meshgrid(x, y, z, indexing="ij")
+        indices = np.where(
+            (X - self.grid_shape[0] // 6) ** 2 + (Y - self.grid_shape[1] // 2) ** 2 + (Z - self.grid_shape[2] // 2) ** 2 < sphere_radius**2
+        )
+        sphere = [tuple(indices[i]) for i in range(self.velocity_set.d)]
+
+        return inlet, outlet, walls, sphere
+
+    def setup_boundary_conditions(self):
+        inlet, outlet, walls, sphere = self.define_boundary_indices()
+        bc_left = RegularizedBC("velocity", (0.04, 0.0, 0.0), indices=inlet)
+        # bc_left = EquilibriumBC(rho = 1, u=(0.04, 0.0, 0.0), indices=inlet)
+        bc_walls = FullwayBounceBackBC(indices=walls)
+        # bc_outlet = RegularizedBC("pressure", 1.0, indices=outlet)
+        # bc_outlet = DoNothingBC(indices=outlet)
+        bc_outlet = ExtrapolationOutflowBC(indices=outlet)
+        bc_sphere = HalfwayBounceBackBC(indices=sphere)
+        self.boundary_conditions = [bc_walls, bc_left, bc_outlet, bc_sphere]
+
+    def setup_boundary_masker(self):
+        # check boundary condition list for duplicate indices before creating bc mask
+        check_bc_overlaps(self.boundary_conditions, self.velocity_set.d, self.backend)
+
+        indices_boundary_masker = IndicesBoundaryMasker(
+            velocity_set=self.velocity_set,
+            precision_policy=self.precision_policy,
+            compute_backend=self.backend,
+        )
+        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.precision_policy, self.backend)
+
+    def setup_stepper(self, omega):
+        self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions, collision_type="BGK")
+
+    def run(self, num_steps, post_process_interval=100):
+        start_time = time.time()
+        for i in range(num_steps):
+            self.f_0, self.f_1 = self.stepper(self.f_0, self.f_1, self.bc_mask, self.missing_mask, i)
+            self.f_0, self.f_1 = self.f_1, self.f_0
+
+            if i % post_process_interval == 0 or i == num_steps - 1:
+                self.post_process(i)
+                end_time = time.time()
+                print(f"Completing {i} iterations. Time elapsed for 1000 LBM steps in {end_time - start_time:.6f} seconds.")
+                start_time = time.time()
+
+    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.D3Q19(precision_policy=self.precision_policy, backend=ComputeBackend.JAX),
+        )
+        rho, u = macro(f_0)
+
+        # remove boundary cells
+        u = u[:, 1:-1, 1:-1, 1:-1]
+        rho = rho[:, 1:-1, 1:-1, 1:-1]
+        u_magnitude = (u[0] ** 2 + u[1] ** 2 + u[2] ** 2) ** 0.5
+
+        fields = {"u_magnitude": u_magnitude, "u_x": u[0], "u_y": u[1], "u_z": u[2], "rho": rho[0]}
+
+        # save_fields_vtk(fields, timestep=i)
+        save_image(fields["u_magnitude"][:, self.grid_shape[1] // 2, :], timestep=i)
+
+
+if __name__ == "__main__":
+    # Running the simulation
+    grid_shape = (512 // 2, 128 // 2, 128 // 2)
+    backend = ComputeBackend.WARP
+    precision_policy = PrecisionPolicy.FP32FP32
+
+    velocity_set = xlb.velocity_set.D3Q19(precision_policy=precision_policy, backend=backend)
+    omega = 1.6
+
+    simulation = FlowOverSphere(omega, grid_shape, velocity_set, backend, precision_policy)
+    simulation.run(num_steps=10000, post_process_interval=1000)
diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py
new file mode 100644
index 0000000..d790c2c
--- /dev/null
+++ b/examples/cfd/lid_driven_cavity_2d.py
@@ -0,0 +1,118 @@
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.precision_policy import PrecisionPolicy
+from xlb.helper import create_nse_fields, initialize_eq, check_bc_overlaps
+from xlb.operator.boundary_masker import IndicesBoundaryMasker
+from xlb.operator.stepper import IncompressibleNavierStokesStepper
+from xlb.operator.boundary_condition import HalfwayBounceBackBC, EquilibriumBC
+from xlb.operator.macroscopic import Macroscopic
+from xlb.utils import save_fields_vtk, save_image
+import xlb.velocity_set
+import warp as wp
+import jax.numpy as jnp
+import numpy as np
+
+
+class LidDrivenCavity2D:
+    def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
+        # initialize backend
+        xlb.init(
+            velocity_set=velocity_set,
+            default_backend=backend,
+            default_precision_policy=precision_policy,
+        )
+
+        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(omega)
+
+    def _setup(self, omega):
+        self.setup_boundary_conditions()
+        self.setup_boundary_masker()
+        self.initialize_fields()
+        self.setup_stepper(omega)
+
+    def define_boundary_indices(self):
+        box = self.grid.bounding_box_indices()
+        box_no_edge = self.grid.bounding_box_indices(remove_edges=True)
+        lid = box_no_edge["top"]
+        walls = [box["bottom"][i] + box["left"][i] + box["right"][i] for i in range(self.velocity_set.d)]
+        walls = np.unique(np.array(walls), axis=-1).tolist()
+        return lid, walls
+
+    def setup_boundary_conditions(self):
+        lid, walls = self.define_boundary_indices()
+        bc_top = EquilibriumBC(rho=1.0, u=(0.02, 0.0), indices=lid)
+        bc_walls = HalfwayBounceBackBC(indices=walls)
+        self.boundary_conditions = [bc_walls, bc_top]
+
+    def setup_boundary_masker(self):
+        # check boundary condition list for duplicate indices before creating bc mask
+        check_bc_overlaps(self.boundary_conditions, self.velocity_set.d, self.backend)
+        indices_boundary_masker = IndicesBoundaryMasker(
+            velocity_set=self.velocity_set,
+            precision_policy=self.precision_policy,
+            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):
+        self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.precision_policy, self.backend)
+
+    def setup_stepper(self, omega):
+        self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions)
+
+    def run(self, num_steps, post_process_interval=100):
+        for i in range(num_steps):
+            self.f_0, self.f_1 = self.stepper(self.f_0, self.f_1, self.bc_mask, self.missing_mask, i)
+            self.f_0, self.f_1 = self.f_1, self.f_0
+
+            if i % 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):
+            # If the backend is warp, we need to drop the last dimension added by warp for 2D simulations
+            f_0 = wp.to_jax(self.f_0)[..., 0]
+        else:
+            f_0 = self.f_0
+
+        macro = Macroscopic(
+            compute_backend=ComputeBackend.JAX,
+            precision_policy=self.precision_policy,
+            velocity_set=xlb.velocity_set.D2Q9(precision_policy=self.precision_policy, backend=ComputeBackend.JAX),
+        )
+
+        rho, u = macro(f_0)
+
+        # remove boundary cells
+        rho = rho[:, 1:-1, 1:-1]
+        u = u[:, 1:-1, 1:-1]
+        u_magnitude = (u[0] ** 2 + u[1] ** 2) ** 0.5
+
+        fields = {"rho": rho[0], "u_x": u[0], "u_y": u[1], "u_magnitude": u_magnitude}
+
+        save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity")
+        save_image(fields["u_magnitude"], timestep=i, prefix="lid_driven_cavity")
+
+
+if __name__ == "__main__":
+    # Running the simulation
+    grid_size = 500
+    grid_shape = (grid_size, grid_size)
+    backend = ComputeBackend.WARP
+    precision_policy = PrecisionPolicy.FP32FP32
+
+    velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend)
+    omega = 1.6
+
+    simulation = LidDrivenCavity2D(omega, grid_shape, velocity_set, backend, precision_policy)
+    simulation.run(num_steps=5000, post_process_interval=1000)
diff --git a/examples/cfd/lid_driven_cavity_2d_distributed.py b/examples/cfd/lid_driven_cavity_2d_distributed.py
new file mode 100644
index 0000000..a7d43b4
--- /dev/null
+++ b/examples/cfd/lid_driven_cavity_2d_distributed.py
@@ -0,0 +1,35 @@
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.precision_policy import PrecisionPolicy
+from xlb.operator.stepper import IncompressibleNavierStokesStepper
+from xlb.distribute import distribute
+from lid_driven_cavity_2d import LidDrivenCavity2D
+
+
+class LidDrivenCavity2D_distributed(LidDrivenCavity2D):
+    def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
+        super().__init__(omega, grid_shape, velocity_set, backend, precision_policy)
+
+    def setup_stepper(self, omega):
+        stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions)
+        distributed_stepper = distribute(
+            stepper,
+            self.grid,
+            self.velocity_set,
+        )
+        self.stepper = distributed_stepper
+        return
+
+
+if __name__ == "__main__":
+    # Running the simulation
+    grid_size = 512
+    grid_shape = (grid_size, grid_size)
+    backend = ComputeBackend.JAX  # Must be JAX for distributed multi-GPU computations. Distributed computations on WARP are not supported yet!
+    precision_policy = PrecisionPolicy.FP32FP32
+
+    velocity_set = xlb.velocity_set.D2Q9(precision_policy=precision_policy, backend=backend)
+    omega = 1.6
+
+    simulation = LidDrivenCavity2D_distributed(omega, grid_shape, velocity_set, backend, precision_policy)
+    simulation.run(num_steps=5000, post_process_interval=1000)
diff --git a/examples/cfd/turbulent_channel_3d.py b/examples/cfd/turbulent_channel_3d.py
new file mode 100644
index 0000000..eb73fdc
--- /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
+        box = self.grid.bounding_box_indices(remove_edges=True)
+        walls = [box["bottom"][i] + box["top"][i] for i in range(self.velocity_set.d)]
+        return walls
+
+    def setup_boundary_conditions(self):
+        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.precision_policy, 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_0, self.f_1 = self.stepper(self.f_0, self.f_1, self.bc_mask, self.missing_mask, i)
+            self.f_0, self.f_1 = self.f_1, self.f_0
+
+            if (i + 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
new file mode 100644
index 0000000..96c79f7
--- /dev/null
+++ b/examples/cfd/windtunnel_3d.py
@@ -0,0 +1,258 @@
+import xlb
+import trimesh
+import time
+from xlb.compute_backend import ComputeBackend
+from xlb.precision_policy import PrecisionPolicy
+from xlb.helper import create_nse_fields, initialize_eq, check_bc_overlaps
+from xlb.operator.stepper import IncompressibleNavierStokesStepper
+from xlb.operator.boundary_condition import (
+    FullwayBounceBackBC,
+    EquilibriumBC,
+    DoNothingBC,
+    RegularizedBC,
+    HalfwayBounceBackBC,
+    ExtrapolationOutflowBC,
+    GradsApproximationBC,
+)
+from xlb.operator.force.momentum_transfer import MomentumTransfer
+from xlb.operator.macroscopic import Macroscopic
+from xlb.operator.boundary_masker import IndicesBoundaryMasker, MeshBoundaryMasker, MeshDistanceBoundaryMasker
+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
+
+
+class WindTunnel3D:
+    def __init__(self, omega, wind_speed, grid_shape, velocity_set, backend, precision_policy):
+        # initialize backend
+        xlb.init(
+            velocity_set=velocity_set,
+            default_backend=backend,
+            default_precision_policy=precision_policy,
+        )
+
+        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.wind_speed = wind_speed
+        self.omega = omega
+        self._setup()
+
+        # Make list to store drag coefficients
+        self.time_steps = []
+        self.drag_coefficients = []
+        self.lift_coefficients = []
+
+    def _setup(self):
+        # NOTE: it is important to initialize fields before setup_boundary_masker is called because f_0 or f_1 might be used to store BC information
+        self.initialize_fields()
+        self.setup_boundary_conditions()
+        self.setup_boundary_masker()
+        self.setup_stepper()
+
+    def voxelize_stl(self, stl_filename, length_lbm_unit):
+        mesh = trimesh.load_mesh(stl_filename, process=False)
+        length_phys_unit = mesh.extents.max()
+        pitch = length_phys_unit / length_lbm_unit
+        mesh_voxelized = mesh.voxelized(pitch=pitch)
+        mesh_matrix = mesh_voxelized.matrix
+        return mesh_matrix, pitch
+
+    def define_boundary_indices(self):
+        box = self.grid.bounding_box_indices()
+        box_no_edge = self.grid.bounding_box_indices(remove_edges=True)
+        inlet = box_no_edge["left"]
+        outlet = box_no_edge["right"]
+        walls = [box["bottom"][i] + box["top"][i] + box["front"][i] + box["back"][i] for i in range(self.velocity_set.d)]
+        walls = np.unique(np.array(walls), axis=-1).tolist()
+
+        # Load the mesh (replace with your own mesh)
+        stl_filename = "../stl-files/DrivAer-Notchback.stl"
+        mesh = trimesh.load_mesh(stl_filename, process=False)
+        mesh_vertices = mesh.vertices
+
+        # Transform the mesh points to be located in the right position in the wind tunnel
+        mesh_vertices -= mesh_vertices.min(axis=0)
+        mesh_extents = mesh_vertices.max(axis=0)
+        length_phys_unit = mesh_extents.max()
+        length_lbm_unit = self.grid_shape[0] / 4
+        dx = length_phys_unit / length_lbm_unit
+        shift = np.array([self.grid_shape[0] * dx / 4, (self.grid_shape[1] * dx - mesh_extents[1]) / 2, 0.0])
+        car = mesh_vertices + shift
+        self.grid_spacing = dx
+        self.car_cross_section = np.prod(mesh_extents[1:]) / dx**2
+
+        return inlet, outlet, walls, car
+
+    def setup_boundary_conditions(self):
+        inlet, outlet, walls, car = self.define_boundary_indices()
+        bc_left = EquilibriumBC(rho=1.0, u=(self.wind_speed, 0.0, 0.0), indices=inlet)
+        # bc_left = RegularizedBC('velocity', (self.wind_speed, 0.0, 0.0), indices=inlet)
+        bc_walls = FullwayBounceBackBC(indices=walls)
+        bc_do_nothing = ExtrapolationOutflowBC(indices=outlet)
+        # bc_car = HalfwayBounceBackBC(mesh_vertices=car)
+        bc_car = GradsApproximationBC(mesh_vertices=car)
+        # bc_car = FullwayBounceBackBC(mesh_vertices=car)
+        self.boundary_conditions = [bc_walls, bc_left, bc_do_nothing, bc_car]
+
+    def setup_boundary_masker(self):
+        # check boundary condition list for duplicate indices before creating bc mask
+        check_bc_overlaps(self.boundary_conditions, self.velocity_set.d, self.backend)
+
+        indices_boundary_masker = IndicesBoundaryMasker(
+            velocity_set=self.velocity_set,
+            precision_policy=self.precision_policy,
+            compute_backend=self.backend,
+        )
+        # mesh_boundary_masker = MeshBoundaryMasker(
+        #     velocity_set=self.velocity_set,
+        #     precision_policy=self.precision_policy,
+        #     compute_backend=self.backend,
+        # )
+        mesh_distance_boundary_masker = MeshDistanceBoundaryMasker(
+            velocity_set=self.velocity_set,
+            precision_policy=self.precision_policy,
+            compute_backend=self.backend,
+        )
+        bclist_other = self.boundary_conditions[:-1]
+        bc_mesh = self.boundary_conditions[-1]
+        dx = self.grid_spacing
+        origin, spacing = (0, 0, 0), (dx, dx, dx)
+        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)
+        self.bc_mask, self.missing_mask, self.f_1 = mesh_distance_boundary_masker(bc_mesh, origin, spacing, self.bc_mask, self.missing_mask, self.f_1)
+
+    def initialize_fields(self):
+        self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.precision_policy, self.backend)
+        self.f_1 = initialize_eq(self.f_1, self.grid, self.velocity_set, self.precision_policy, self.backend)
+
+    def setup_stepper(self):
+        self.stepper = IncompressibleNavierStokesStepper(self.omega, boundary_conditions=self.boundary_conditions, collision_type="KBC")
+
+    def run(self, num_steps, print_interval, post_process_interval=100):
+        # Setup the operator for computing surface forces at the interface of the specified BC
+        bc_car = self.boundary_conditions[-1]
+        self.momentum_transfer = MomentumTransfer(bc_car)
+
+        start_time = time.time()
+        for i in range(num_steps):
+            self.f_0, self.f_1 = self.stepper(self.f_0, self.f_1, self.bc_mask, self.missing_mask, i)
+            self.f_0, self.f_1 = self.f_1, self.f_0
+
+            if (i + 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)
+
+        # remove boundary cells
+        u = u[:, 1:-1, 1:-1, 1:-1]
+        u_magnitude = (u[0] ** 2 + u[1] ** 2 + u[2] ** 2) ** 0.5
+
+        fields = {"u_magnitude": u_magnitude}
+
+        save_fields_vtk(fields, timestep=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.f_1, 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)
+        c_l = 2.0 * lift / (self.wind_speed**2 * self.car_cross_section)
+        self.drag_coefficients.append(c_d)
+        self.lift_coefficients.append(c_l)
+        self.time_steps.append(i)
+
+        # Save monitor plot
+        self.plot_drag_coefficient()
+        return
+
+    def plot_drag_coefficient(self):
+        # Compute moving average of drag coefficient, 100, 1000, 10000
+        drag_coefficients = np.array(self.drag_coefficients)
+        self.drag_coefficients_ma_10 = np.convolve(drag_coefficients, np.ones(10) / 10, mode="valid")
+        self.drag_coefficients_ma_100 = np.convolve(drag_coefficients, np.ones(100) / 100, mode="valid")
+        self.drag_coefficients_ma_1000 = np.convolve(drag_coefficients, np.ones(1000) / 1000, mode="valid")
+        self.drag_coefficients_ma_10000 = np.convolve(drag_coefficients, np.ones(10000) / 10000, mode="valid")
+        self.drag_coefficients_ma_100000 = np.convolve(drag_coefficients, np.ones(100000) / 100000, mode="valid")
+
+        # Plot drag coefficient
+        plt.plot(self.time_steps, drag_coefficients, label="Raw")
+        if len(self.time_steps) > 10:
+            plt.plot(self.time_steps[9:], self.drag_coefficients_ma_10, label="MA 10")
+        if len(self.time_steps) > 100:
+            plt.plot(self.time_steps[99:], self.drag_coefficients_ma_100, label="MA 100")
+        if len(self.time_steps) > 1000:
+            plt.plot(self.time_steps[999:], self.drag_coefficients_ma_1000, label="MA 1,000")
+        if len(self.time_steps) > 10000:
+            plt.plot(self.time_steps[9999:], self.drag_coefficients_ma_10000, label="MA 10,000")
+        if len(self.time_steps) > 100000:
+            plt.plot(self.time_steps[99999:], self.drag_coefficients_ma_100000, label="MA 100,000")
+
+        plt.ylim(-1.0, 1.0)
+        plt.legend()
+        plt.xlabel("Time step")
+        plt.ylabel("Drag coefficient")
+        plt.savefig("drag_coefficient_ma.png")
+        plt.close()
+
+
+if __name__ == "__main__":
+    # Grid parameters
+    grid_size_x, grid_size_y, grid_size_z = 512, 128, 128
+    grid_shape = (grid_size_x, grid_size_y, grid_size_z)
+
+    # Configuration
+    backend = ComputeBackend.WARP
+    precision_policy = PrecisionPolicy.FP32FP32
+    velocity_set = xlb.velocity_set.D3Q27(precision_policy=precision_policy, backend=backend)
+    wind_speed = 0.02
+    num_steps = 100000
+    print_interval = 1000
+
+    # Set up Reynolds number and deduce relaxation time (omega)
+    # Re = 50000.0
+    Re = 500000000000.0
+    clength = grid_size_x - 1
+    visc = wind_speed * clength / Re
+    omega = 1.0 / (3.0 * visc + 0.5)
+
+    # 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"Prescribed velocity: {wind_speed}")
+    print(f"Reynolds number: {Re}")
+    print(f"Max iterations: {num_steps}")
+    print("\n" + "=" * 50 + "\n")
+
+    simulation = WindTunnel3D(omega, wind_speed, grid_shape, velocity_set, backend, precision_policy)
+    simulation.run(num_steps, print_interval, post_process_interval=1000)
diff --git a/examples/performance/MLUPS2d.py b/examples/performance/MLUPS2d.py
deleted file mode 100644
index 77d32a4..0000000
--- a/examples/performance/MLUPS2d.py
+++ /dev/null
@@ -1,67 +0,0 @@
-"""
-This script computes the MLUPS (Million Lattice Updates per Second) in 2D by simulating fluid flow inside a 2D cavity.
-"""
-
-import os
-import argparse
-import jax.numpy as jnp
-import numpy as np
-from jax import config
-from time import time
-
-from src.utils import *
-from src.boundary_conditions import *
-from src.lattice import LatticeD2Q9
-from src.models import BGKSim
-
-class Cavity(BGKSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        # concatenate the indices of the left, right, and bottom walls
-        walls = np.concatenate((self.boundingBoxIndices['left'], self.boundingBoxIndices['right'], self.boundingBoxIndices['bottom']))
-        # apply bounce back boundary condition to the walls
-        self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy))
-
-        # apply inlet equilibrium boundary condition to the top wall
-        moving_wall = self.boundingBoxIndices['top']
-
-        rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype)
-        vel_wall[:, 0] = u_wall
-        self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall))
-
-
-if __name__ == '__main__':
-    precision = 'f32/f32'
-    lattice = LatticeD2Q9(precision)
-
-    parser = argparse.ArgumentParser("simple_example")
-    parser.add_argument("N", help="The total number of voxels will be NxN", type=int)
-    parser.add_argument("timestep", help="Number of timesteps", type=int)
-    args = parser.parse_args()
-
-    n = args.N
-    max_iter = args.timestep
-    Re = 100.0
-    u_wall = 0.1
-    clength = n - 1
-
-    visc = u_wall * clength / Re
-    omega = 1.0 / (3. * visc + 0.5)
-    print('omega = ', omega)
-
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': n,
-        'ny': n,
-        'nz': 0,
-        'precision': precision,
-        'compute_MLUPS': True
-    }
-    
-    os.system('rm -rf ./*.vtk && rm -rf ./*.png')
-    sim = Cavity(**kwargs)
-    sim.run(max_iter)
diff --git a/examples/performance/MLUPS3d.py b/examples/performance/MLUPS3d.py
deleted file mode 100644
index 8a9f9e3..0000000
--- a/examples/performance/MLUPS3d.py
+++ /dev/null
@@ -1,76 +0,0 @@
-"""
-This script computes the MLUPS (Million Lattice Updates per Second) in 3D by simulating fluid flow inside a 2D cavity.
-"""
-
-import os
-import argparse
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-from jax import config
-from time import time
-#config.update('jax_disable_jit', True)
-# Use 8 CPU devices
-#os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-#config.update("jax_enable_x64", True)
-from src.utils import *
-from src.boundary_conditions import *
-from src.models import BGKSim
-from src.lattice import LatticeD3Q19
-class Cavity(BGKSim):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        # concatenate the indices of the left, right, and bottom walls
-        walls = np.concatenate((self.boundingBoxIndices['left'], self.boundingBoxIndices['right'], self.boundingBoxIndices['bottom'], self.boundingBoxIndices['front'], self.boundingBoxIndices['back']))
-        # apply bounce back boundary condition to the walls
-        self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy))
-
-        # apply inlet equilibrium boundary condition to the top wall
-        moving_wall = self.boundingBoxIndices['top']
-    
-        rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype)
-        vel_wall[:, 0] = u_wall
-        self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall))
-
-if __name__ == '__main__':
-    precision = 'f32/f32'
-    lattice = LatticeD3Q19(precision)
-    # Create a parser that will read the command line arguments
-    parser = argparse.ArgumentParser("Calculate MLUPS for a 3D cavity flow simulation")
-    parser.add_argument("N", help="The total number of voxels all directions. The final dimension will be N*NxN", default=100, type=int)
-    parser.add_argument("N_ITERS", help="Number of timesteps", default=10000, type=int)    
-
-    args = parser.parse_args()
-    n = args.N
-    n_iters = args.N_ITERS
-
-    # Store the Reynolds number in the variable Re
-    Re = 100.0
-    # Store the velocity of the lid in the variable u_wall
-    u_wall = 0.1
-    # Store the length of the cavity in the variable clength
-    clength = n - 1
-
-    # Compute the viscosity from the Reynolds number, the lid velocity, and the length of the cavity
-    visc = u_wall * clength / Re
-    # Compute the relaxation parameter from the viscosity
-    omega = 1.0 / (3. * visc + 0.5)
-
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': n,
-        'ny': n,
-        'nz': n,
-        'precision': precision,
-        'compute_MLUPS': True
-    }
-
-    sim = Cavity(**kwargs)
-    # Run the simulation
-    sim.run(n_iters)
-    
\ No newline at end of file
diff --git a/examples/performance/MLUPS3d_distributed.py b/examples/performance/MLUPS3d_distributed.py
deleted file mode 100644
index 70d5328..0000000
--- a/examples/performance/MLUPS3d_distributed.py
+++ /dev/null
@@ -1,108 +0,0 @@
-"""
-This script computes the MLUPS (Million Lattice Updates per Second) in 3D by simulating fluid flow inside a 2D cavity. 
-This script is equivalent to MLUPS3d.py, but uses JAX distributed to run the simulation on distributed systems (multi-host, multi-GPUs).
-Please refer to https://jax.readthedocs.io/en/latest/multi_process.html for more information on JAX distributed.
-"""
-
-
-# Standard Libraries
-import argparse
-import os
-import jax
-
-import jax.numpy as jnp
-import numpy as np
-
-from jax import config
-
-from src.boundary_conditions import *
-from src.models import BGKSim
-from src.lattice import LatticeD3Q19
-from src.utils import *
-
-#config.update('jax_disable_jit', True)
-# Use 8 CPU devices
-#os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-#config.update("jax_enable_x64", True)
-
-class Cavity(BGKSim):
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def set_boundary_conditions(self):
-        # concatenate the indices of the left, right, and bottom walls
-        walls = np.concatenate((self.boundingBoxIndices['left'], self.boundingBoxIndices['right'], self.boundingBoxIndices['bottom'], self.boundingBoxIndices['front'], self.boundingBoxIndices['back']))
-        # apply bounce back boundary condition to the walls
-        self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy))
-
-        # apply inlet equilibrium boundary condition to the top wall
-        moving_wall = self.boundingBoxIndices['top']
-    
-        rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype)
-        vel_wall[:, 0] = u_wall
-        self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall))
-
-if __name__ == '__main__':
-    # Create a parser that will read the command line arguments
-    parser = argparse.ArgumentParser("Calculate MLUPS for a 3D cavity flow simulation")
-    parser.add_argument("N", help="The total number of voxels in one direction. The final dimension will be N*NxN",
-                        default=100, type=int)
-    parser.add_argument("N_ITERS", help="Number of iterations", default=10000, type=int)
-    parser.add_argument("N_PROCESSES", help="Number of processes. If >1, call jax.distributed.initialize with that number of process. If -1 will call jax.distributed.initialize without any arsgument. So it should pick up the values from SLURM env variable.",
-                        default=1, type=int)
-    parser.add_argument("IP", help="IP of the master node for multi-node. Useless if using SLURM.",
-                        default='127.0.0.1', type=str, nargs='?')
-    parser.add_argument("PROCESS_ID_INCREMENT", help="For multi-node only. Useless if using SLURM.",
-                        default=0, type=int, nargs='?')
-
-    args = parser.parse_args()
-    n = args.N
-    n_iters = args.N_ITERS
-    n_processes = args.N_PROCESSES
-    # Initialize JAX distributed. The IP, number of processes and process id must be set correctly.
-    print("N processes, ", n_processes)
-    print("N iter, ", n_iters)
-    if n_processes > 1:
-        process_id = int(os.environ.get('CUDA_VISIBLE_DEVICES', 0)) + args.PROCESS_ID_INCREMENT
-        print("ip, num_processes, process_id, ", args.IP, n_processes, process_id)
-        jax.distributed.initialize(args.IP, num_processes=n_processes,
-                                   process_id=process_id)
-    elif n_processes == -1:
-        print("Will call jax.distributed.initialize()")
-        jax.distributed.initialize()
-        print("jax.distributed.initialize() ended")
-    else:
-        print("No call to jax.distributed.initialize")
-    print("JAX local devices: ", jax.local_devices())
-
-    precision = 'f32/f32'
-    # Create a 3D lattice with the D3Q19 scheme
-    lattice = LatticeD3Q19(precision)
-
-    # Store the Reynolds number in the variable Re
-    Re = 100.0
-    # Store the velocity of the lid in the variable u_wall
-    u_wall = 0.1
-    # Store the length of the cavity in the variable clength
-    clength = n - 1
-
-    # Compute the viscosity from the Reynolds number, the lid velocity, and the length of the cavity
-    visc = u_wall * clength / Re
-    # Compute the relaxation parameter from the viscosity
-    omega = 1.0 / (3. * visc + 0.5)
-    
-    # Create a new instance of the Cavity class
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': n,
-        'ny': n,
-        'nz': n,
-        'precision': precision,
-        'compute_MLUPS': True
-    }
-
-    sim = Cavity(**kwargs)    # Run the simulation
-    sim.run(n_iters)
diff --git a/examples/performance/mlups_3d.py b/examples/performance/mlups_3d.py
new file mode 100644
index 0000000..2001fb2
--- /dev/null
+++ b/examples/performance/mlups_3d.py
@@ -0,0 +1,108 @@
+import xlb
+import argparse
+import time
+import warp as wp
+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 FullwayBounceBackBC, EquilibriumBC
+from xlb.distribute import distribute
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description="MLUPS for 3D Lattice Boltzmann Method Simulation (BGK)")
+    parser.add_argument("cube_edge", type=int, help="Length of the edge of the cubic grid")
+    parser.add_argument("num_steps", type=int, help="Timestep for the simulation")
+    parser.add_argument("backend", type=str, help="Backend for the simulation (jax or warp)")
+    parser.add_argument("precision", type=str, help="Precision for the simulation (e.g., fp32/fp32)")
+    return parser.parse_args()
+
+
+def setup_simulation(args):
+    backend = ComputeBackend.JAX if args.backend == "jax" else ComputeBackend.WARP
+    precision_policy_map = {
+        "fp32/fp32": PrecisionPolicy.FP32FP32,
+        "fp64/fp64": PrecisionPolicy.FP64FP64,
+        "fp64/fp32": PrecisionPolicy.FP64FP32,
+        "fp32/fp16": PrecisionPolicy.FP32FP16,
+    }
+    precision_policy = precision_policy_map.get(args.precision)
+    if precision_policy is None:
+        raise ValueError("Invalid precision")
+
+    xlb.init(
+        velocity_set=xlb.velocity_set.D3Q19(precision_policy=precision_policy, backend=backend),
+        default_backend=backend,
+        default_precision_policy=precision_policy,
+    )
+
+    return backend, precision_policy
+
+
+def create_grid_and_fields(cube_edge):
+    grid_shape = (cube_edge, cube_edge, cube_edge)
+    grid, f_0, f_1, missing_mask, bc_mask = create_nse_fields(grid_shape)
+
+    return grid, f_0, f_1, missing_mask, bc_mask
+
+
+def define_boundary_indices(grid):
+    box = grid.bounding_box_indices()
+    box_no_edge = grid.bounding_box_indices(remove_edges=True)
+    lid = box_no_edge["top"]
+    walls = [box["bottom"][i] + box["left"][i] + box["right"][i] + box["front"][i] + box["back"][i] for i in range(len(grid.shape))]
+    return lid, walls
+
+
+def setup_boundary_conditions(grid):
+    lid, walls = define_boundary_indices(grid)
+    bc_top = EquilibriumBC(rho=1.0, u=(0.02, 0.0, 0.0), indices=lid)
+    bc_walls = FullwayBounceBackBC(indices=walls)
+    return [bc_top, bc_walls]
+
+
+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))
+
+    if backend == ComputeBackend.JAX:
+        stepper = distribute(
+            stepper,
+            grid,
+            xlb.velocity_set.D3Q19(precision_policy=precision_policy, backend=backend),
+        )
+
+    start_time = time.time()
+
+    for i in range(num_steps):
+        f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, i)
+        f_0, f_1 = f_1, f_0
+    wp.synchronize()
+
+    end_time = time.time()
+    return end_time - start_time
+
+
+def calculate_mlups(cube_edge, num_steps, elapsed_time):
+    total_lattice_updates = cube_edge**3 * num_steps
+    mlups = (total_lattice_updates / elapsed_time) / 1e6
+    return mlups
+
+
+def main():
+    args = parse_arguments()
+    backend, precision_policy = setup_simulation(args)
+    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, precision_policy, backend)
+
+    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")
+    print(f"MLUPs: {mlups:.2f}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/requirements.txt b/requirements.txt
index 11ee0fd..4d0cd2c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,11 +1,10 @@
-jax==0.4.20
-jaxlib==0.4.20
-jmp==0.0.4
-matplotlib==3.8.0
-numpy==1.26.1
-pyvista==0.42.3
-Rtree==1.0.1
-trimesh==4.0.0
-orbax-checkpoint==0.4.1
-termcolor==2.3.0
-PhantomGaze @ git+https://github.com/loliverhennigh/PhantomGaze.git
\ No newline at end of file
+jax[cuda]
+matplotlib
+numpy
+pyvista
+Rtree
+trimesh
+warp-lang
+numpy-stl
+pydantic
+ruff
\ No newline at end of file
diff --git a/ruff.toml b/ruff.toml
new file mode 100644
index 0000000..83c6370
--- /dev/null
+++ b/ruff.toml
@@ -0,0 +1,42 @@
+# Adopted from tinygrad's ruff.toml thanks @geohot
+indent-width = 4
+preview = true
+target-version = "py38"
+
+lint.select = [
+  "F",  # Pyflakes
+  "W6",
+  "E71",
+  "E72",
+  "E112",   # no-indented-block
+  "E113",   # unexpected-indentation
+  # "E124",
+  "E203",   # whitespace-before-punctuation
+  "E272",   # multiple-spaces-before-keyword
+  "E303",   # too-many-blank-lines
+  "E304",   # blank-line-after-decorator
+  "E501",   # line-too-long
+  # "E502",
+  "E702",   # multiple-statements-on-one-line-semicolon
+  "E703",   # useless-semicolon
+  "E731",   # lambda-assignment
+  "W191",   # tab-indentation
+  "W291",   # trailing-whitespace
+  "W293",   # blank-line-with-whitespace
+  "UP039",  # unnecessary-class-parentheses
+  "C416",   # unnecessary-comprehension
+  "RET506", # superfluous-else-raise
+  "RET507", # superfluous-else-continue
+  "A",      # builtin-variable-shadowing, builtin-argument-shadowing, builtin-attribute-shadowing
+  "SIM105", # suppressible-exception
+  "FURB110",# if-exp-instead-of-or-operator
+]
+
+# unused-variable, shadowing a Python builtin module, Module imported but unused
+lint.ignore = ["F841", "A005", "F401"]
+line-length = 150
+
+exclude = [
+  "docs/",
+  "xlb/experimental/",
+]
\ No newline at end of file
diff --git a/setup.py b/setup.py
index e69de29..6f9780a 100644
--- a/setup.py
+++ b/setup.py
@@ -0,0 +1,10 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="XLB",
+    version="0.0.1",
+    author="",
+    packages=find_packages(),
+    install_requires=[],
+    include_package_data=True,
+)
diff --git a/src/base.py b/src/base.py
deleted file mode 100644
index 7fb4f57..0000000
--- a/src/base.py
+++ /dev/null
@@ -1,1086 +0,0 @@
-# Standard Libraries
-import os
-import time
-
-# Third-Party Libraries
-import jax
-import jax.numpy as jnp
-import jmp
-import numpy as np
-from termcolor import colored
-
-# JAX-related imports
-from jax import jit, lax, vmap
-from jax.experimental import mesh_utils
-from jax.experimental.multihost_utils import process_allgather
-from jax.experimental.shard_map import shard_map
-from jax.sharding import NamedSharding, PartitionSpec, PositionalSharding, Mesh
-import orbax.checkpoint as orb
-
-# functools imports
-from functools import partial
-
-# Local/Custom Libraries
-from src.utils import downsample_field
-
-jax.config.update("jax_spmd_mode", 'allow_all')
-# Disables annoying TF warnings
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 
-
-class LBMBase(object):
-    """
-    LBMBase: A class that represents a base for Lattice Boltzmann Method simulation.
-    
-    Parameters
-    ----------
-        lattice (object): The lattice object that contains the lattice structure and weights.
-        omega (float): The relaxation parameter for the LBM simulation.
-        nx (int): Number of grid points in the x-direction.
-        ny (int): Number of grid points in the y-direction.
-        nz (int, optional): Number of grid points in the z-direction. Defaults to 0.
-        precision (str, optional): A string specifying the precision used for the simulation. Defaults to "f32/f32".
-    """
-    
-    def __init__(self, **kwargs):
-        self.omega = kwargs.get("omega")
-        self.nx = kwargs.get("nx")
-        self.ny = kwargs.get("ny")
-        self.nz = kwargs.get("nz")
-
-        self.precision = kwargs.get("precision")
-        computedType, storedType = self.set_precisions(self.precision)
-        self.precisionPolicy = jmp.Policy(compute_dtype=computedType,
-                                            param_dtype=computedType, output_dtype=storedType)
-        
-        self.lattice = kwargs.get("lattice")
-        self.checkpointRate = kwargs.get("checkpoint_rate", 0)
-        self.checkpointDir = kwargs.get("checkpoint_dir", './checkpoints')
-        self.downsamplingFactor = kwargs.get("downsampling_factor", 1)
-        self.printInfoRate = kwargs.get("print_info_rate", 100)
-        self.ioRate = kwargs.get("io_rate", 0)
-        self.returnFpost = kwargs.get("return_fpost", False)
-        self.computeMLUPS = kwargs.get("compute_MLUPS", False)
-        self.restore_checkpoint = kwargs.get("restore_checkpoint", False)
-        self.nDevices = jax.device_count()
-        self.backend = jax.default_backend()
-
-        if self.computeMLUPS:
-            self.restore_checkpoint = False
-            self.ioRate = 0
-            self.checkpointRate = 0
-            self.printInfoRate = 0
-
-        # Check for distributed mode
-        if self.nDevices > jax.local_device_count():
-            print("WARNING: Running in distributed mode. Make sure that jax.distributed.initialize is called before performing any JAX computations.")
-                    
-        self.c = self.lattice.c
-        self.q = self.lattice.q
-        self.w = self.lattice.w
-        self.dim = self.lattice.d
-
-        # Set the checkpoint manager
-        if self.checkpointRate > 0:
-            mngr_options = orb.CheckpointManagerOptions(save_interval_steps=self.checkpointRate, max_to_keep=1)
-            self.mngr = orb.CheckpointManager(self.checkpointDir, orb.PyTreeCheckpointer(), options=mngr_options)
-        else:
-            self.mngr = None
-        
-        # Adjust the number of grid points in the x direction, if necessary.
-        # If the number of grid points is not divisible by the number of devices
-        # it increases the number of grid points to the next multiple of the number of devices.
-        # This is done in order to accommodate the domain sharding per XLA device
-        nx, ny, nz = kwargs.get("nx"), kwargs.get("ny"), kwargs.get("nz")
-        if None in {nx, ny, nz}:
-            raise ValueError("nx, ny, and nz must be provided. For 2D examples, nz must be set to 0.")
-        self.nx = nx
-        if nx % self.nDevices:
-            self.nx = nx + (self.nDevices - nx % self.nDevices)
-            print("WARNING: nx increased from {} to {} in order to accommodate domain sharding per XLA device.".format(nx, self.nx))
-        self.ny = ny
-        self.nz = nz
-
-        self.show_simulation_parameters()
-    
-        # Store grid information
-        self.gridInfo = {
-            "nx": self.nx,
-            "ny": self.ny,
-            "nz": self.nz,
-            "dim": self.lattice.d,
-            "lattice": self.lattice
-        }
-
-        P = PartitionSpec
-
-        # Define the right permutation
-        self.rightPerm = [(i, (i + 1) % self.nDevices) for i in range(self.nDevices)]
-        # Define the left permutation
-        self.leftPerm = [((i + 1) % self.nDevices, i) for i in range(self.nDevices)]
-
-        # Set up the sharding and streaming for 2D and 3D simulations
-        if self.dim == 2:
-            self.devices = mesh_utils.create_device_mesh((self.nDevices, 1, 1))
-            self.mesh = Mesh(self.devices, axis_names=("x", "y", "value"))
-            self.sharding = NamedSharding(self.mesh, P("x", "y", "value"))
-
-            self.streaming = jit(shard_map(self.streaming_m, mesh=self.mesh,
-                                                      in_specs=P("x", None, None), out_specs=P("x", None, None), check_rep=False))
-
-        # Set up the sharding and streaming for 2D and 3D simulations
-        elif self.dim == 3:
-            self.devices = mesh_utils.create_device_mesh((self.nDevices, 1, 1, 1))
-            self.mesh = Mesh(self.devices, axis_names=("x", "y", "z", "value"))
-            self.sharding = NamedSharding(self.mesh, P("x", "y", "z", "value"))
-
-            self.streaming = jit(shard_map(self.streaming_m, mesh=self.mesh,
-                                                      in_specs=P("x", None, None, None), out_specs=P("x", None, None, None), check_rep=False))
-
-        else:
-            raise ValueError(f"dim = {self.dim} not supported")
-        
-        # Compute the bounding box indices for boundary conditions
-        self.boundingBoxIndices= self.bounding_box_indices()
-        # Create boundary data for the simulation
-        self._create_boundary_data()
-        self.force = self.get_force()
-
-    @property
-    def lattice(self):
-        return self._lattice
-
-    @lattice.setter
-    def lattice(self, value):
-        if value is None:
-            raise ValueError("Lattice type must be provided.")
-        if self.nz == 0 and value.name not in ['D2Q9']:
-            raise ValueError("For 2D simulations, lattice type must be LatticeD2Q9.")
-        if self.nz != 0 and value.name not in ['D3Q19', 'D3Q27']:
-            raise ValueError("For 3D simulations, lattice type must be LatticeD3Q19, or LatticeD3Q27.")
-                            
-        self._lattice = value
-
-    @property
-    def omega(self):
-        return self._omega
-
-    @omega.setter
-    def omega(self, value):
-        if value is None:
-            raise ValueError("omega must be provided")
-        if not isinstance(value, float):
-            raise TypeError("omega must be a float")
-        self._omega = value
-
-    @property
-    def nx(self):
-        return self._nx
-
-    @nx.setter
-    def nx(self, value):
-        if value is None:
-            raise ValueError("nx must be provided")
-        if not isinstance(value, int):
-            raise TypeError("nx must be an integer")
-        self._nx = value
-
-    @property
-    def ny(self):
-        return self._ny
-
-    @ny.setter
-    def ny(self, value):
-        if value is None:
-            raise ValueError("ny must be provided")
-        if not isinstance(value, int):
-            raise TypeError("ny must be an integer")
-        self._ny = value
-
-    @property
-    def nz(self):
-        return self._nz
-
-    @nz.setter
-    def nz(self, value):
-        if value is None:
-            raise ValueError("nz must be provided")
-        if not isinstance(value, int):
-            raise TypeError("nz must be an integer")
-        self._nz = value
-
-    @property
-    def precision(self):
-        return self._precision
-
-    @precision.setter
-    def precision(self, value):
-        if not isinstance(value, str):
-            raise TypeError("precision must be a string")
-        self._precision = value
-
-    @property
-    def checkpointRate(self):
-        return self._checkpointRate
-
-    @checkpointRate.setter
-    def checkpointRate(self, value):
-        if not isinstance(value, int):
-            raise TypeError("checkpointRate must be an integer")
-        self._checkpointRate = value
-
-    @property
-    def checkpointDir(self):
-        return self._checkpointDir
-
-    @checkpointDir.setter
-    def checkpointDir(self, value):
-        if not isinstance(value, str):
-            raise TypeError("checkpointDir must be a string")
-        self._checkpointDir = value
-
-    @property
-    def downsamplingFactor(self):
-        return self._downsamplingFactor
-
-    @downsamplingFactor.setter
-    def downsamplingFactor(self, value):
-        if not isinstance(value, int):
-            raise TypeError("downsamplingFactor must be an integer")
-        self._downsamplingFactor = value
-
-    @property
-    def printInfoRate(self):
-        return self._printInfoRate
-
-    @printInfoRate.setter
-    def printInfoRate(self, value):
-        if not isinstance(value, int):
-            raise TypeError("printInfoRate must be an integer")
-        self._printInfoRate = value
-
-    @property
-    def ioRate(self):
-        return self._ioRate
-
-    @ioRate.setter
-    def ioRate(self, value):
-        if not isinstance(value, int):
-            raise TypeError("ioRate must be an integer")
-        self._ioRate = value
-
-    @property
-    def returnFpost(self):
-        return self._returnFpost
-
-    @returnFpost.setter
-    def returnFpost(self, value):
-        if not isinstance(value, bool):
-            raise TypeError("returnFpost must be a boolean")
-        self._returnFpost = value
-
-    @property
-    def computeMLUPS(self):
-        return self._computeMLUPS
-
-    @computeMLUPS.setter
-    def computeMLUPS(self, value):
-        if not isinstance(value, bool):
-            raise TypeError("computeMLUPS must be a boolean")
-        self._computeMLUPS = value
-
-    @property
-    def restore_checkpoint(self):
-        return self._restore_checkpoint
-
-    @restore_checkpoint.setter
-    def restore_checkpoint(self, value):
-        if not isinstance(value, bool):
-            raise TypeError("restore_checkpoint must be a boolean")
-        self._restore_checkpoint = value
-
-    @property
-    def nDevices(self):
-        return self._nDevices
-
-    @nDevices.setter
-    def nDevices(self, value):
-        if not isinstance(value, int):
-            raise TypeError("nDevices must be an integer")
-        self._nDevices = value
-
-    def show_simulation_parameters(self):
-        attributes_to_show = [
-            'omega', 'nx', 'ny', 'nz', 'dim', 'precision', 'lattice', 
-            'checkpointRate', 'checkpointDir', 'downsamplingFactor', 
-            'printInfoRate', 'ioRate', 'computeMLUPS', 
-            'restore_checkpoint', 'backend', 'nDevices'
-        ]
-
-        descriptive_names = {
-            'omega': 'Omega',
-            'nx': 'Grid Points in X',
-            'ny': 'Grid Points in Y',
-            'nz': 'Grid Points in Z',
-            'dim': 'Dimensionality',
-            'precision': 'Precision Policy',
-            'lattice': 'Lattice Type',
-            'checkpointRate': 'Checkpoint Rate',
-            'checkpointDir': 'Checkpoint Directory',
-            'downsamplingFactor': 'Downsampling Factor',
-            'printInfoRate': 'Print Info Rate',
-            'ioRate': 'I/O Rate',
-            'computeMLUPS': 'Compute MLUPS',
-            'restore_checkpoint': 'Restore Checkpoint',
-            'backend': 'Backend',
-            'nDevices': 'Number of Devices'
-        }
-        simulation_name = self.__class__.__name__
-        
-        print(colored(f'**** Simulation Parameters for {simulation_name} ****', 'green'))
-                
-        header = f"{colored('Parameter', 'blue'):>30} | {colored('Value', 'yellow')}"
-        print(header)
-        print('-' * 50)
-        
-        for attr in attributes_to_show:
-            value = getattr(self, attr, 'Attribute not set')
-            descriptive_name = descriptive_names.get(attr, attr)  # Use the attribute name as a fallback
-            row = f"{colored(descriptive_name, 'blue'):>30} | {colored(value, 'yellow')}"
-            print(row)
-
-    def _create_boundary_data(self):
-        """
-        Create boundary data for the Lattice Boltzmann simulation by setting boundary conditions,
-        creating grid mask, and preparing local masks and normal arrays.
-        """
-        self.BCs = []
-        self.set_boundary_conditions()
-        # Accumulate the indices of all BCs to create the grid mask with FALSE along directions that
-        # stream into a boundary voxel.
-        solid_halo_list = [np.array(bc.indices).T for bc in self.BCs if bc.isSolid]
-        solid_halo_voxels = np.unique(np.vstack(solid_halo_list), axis=0) if solid_halo_list else None
-
-        # Create the grid mask on each process
-        start = time.time()
-        grid_mask = self.create_grid_mask(solid_halo_voxels)
-        print("Time to create the grid mask:", time.time() - start)
-
-        start = time.time()
-        for bc in self.BCs:
-            assert bc.implementationStep in ['PostStreaming', 'PostCollision']
-            bc.create_local_mask_and_normal_arrays(grid_mask)
-        print("Time to create the local masks and normal arrays:", time.time() - start)
-
-    # This is another non-JITed way of creating the distributed arrays. It is not used at the moment.
-    # def distributed_array_init(self, shape, type, init_val=None):
-    #     sharding_dim = shape[0] // self.nDevices
-    #     sharded_shape = (self.nDevices, sharding_dim,  *shape[1:])
-    #     device_shape = sharded_shape[1:]
-    #     arrays = []
-
-    #     for d, index in self.sharding.addressable_devices_indices_map(sharded_shape).items():
-    #         jax.default_device = d
-    #         if init_val is None:
-    #             x = jnp.zeros(shape=device_shape, dtype=type)
-    #         else:
-    #             x = jnp.full(shape=device_shape, fill_value=init_val, dtype=type)  
-    #         arrays += [jax.device_put(x, d)] 
-    #     jax.default_device = jax.devices()[0]
-    #     return jax.make_array_from_single_device_arrays(shape, self.sharding, arrays)
-
-    @partial(jit, static_argnums=(0, 1, 2, 4))
-    def distributed_array_init(self, shape, type, init_val=0, sharding=None):
-        """
-        Initialize a distributed array using JAX, with a specified shape, data type, and initial value.
-        Optionally, provide a custom sharding strategy.
-
-        Parameters
-        ----------
-            shape (tuple): The shape of the array to be created.
-            type (dtype): The data type of the array to be created.
-            init_val (scalar, optional): The initial value to fill the array with. Defaults to 0.
-            sharding (Sharding, optional): The sharding strategy to use. Defaults to `self.sharding`.
-
-        Returns
-        -------
-            jax.numpy.ndarray: A JAX array with the specified shape, data type, initial value, and sharding strategy.
-        """
-        if sharding is None:
-            sharding = self.sharding
-        x = jnp.full(shape=shape, fill_value=init_val, dtype=type)        
-        return jax.lax.with_sharding_constraint(x, sharding)
-    
-    @partial(jit, static_argnums=(0,))
-    def create_grid_mask(self, solid_halo_voxels):
-        """
-        This function creates a mask for the background grid that accounts for the location of the boundaries.
-        
-        Parameters
-        ----------
-            solid_halo_voxels: A numpy array representing the voxels in the halo of the solid object.
-            
-        Returns
-        -------
-            A JAX array representing the grid mask of the grid.
-        """
-        # Halo width (hw_x is different to accommodate the domain sharding per XLA device)
-        hw_x = self.nDevices
-        hw_y = hw_z = 1
-        if self.dim == 2:
-            grid_mask = self.distributed_array_init((self.nx + 2 * hw_x, self.ny + 2 * hw_y, self.lattice.q), jnp.bool_, init_val=True)
-            grid_mask = grid_mask.at[(slice(hw_x, -hw_x), slice(hw_y, -hw_y), slice(None))].set(False)
-            if solid_halo_voxels is not None:
-                solid_halo_voxels = solid_halo_voxels.at[:, 0].add(hw_x)
-                solid_halo_voxels = solid_halo_voxels.at[:, 1].add(hw_y)
-                grid_mask = grid_mask.at[tuple(solid_halo_voxels.T)].set(True)  
-
-            grid_mask = self.streaming(grid_mask)
-            return lax.with_sharding_constraint(grid_mask, self.sharding)
-
-        elif self.dim == 3:
-            grid_mask = self.distributed_array_init((self.nx + 2 * hw_x, self.ny + 2 * hw_y, self.nz + 2 * hw_z, self.lattice.q), jnp.bool_, init_val=True)
-            grid_mask = grid_mask.at[(slice(hw_x, -hw_x), slice(hw_y, -hw_y), slice(hw_z, -hw_z), slice(None))].set(False)
-            if solid_halo_voxels is not None:
-                solid_halo_voxels = solid_halo_voxels.at[:, 0].add(hw_x)
-                solid_halo_voxels = solid_halo_voxels.at[:, 1].add(hw_y)
-                solid_halo_voxels = solid_halo_voxels.at[:, 2].add(hw_z)
-                grid_mask = grid_mask.at[tuple(solid_halo_voxels.T)].set(True)
-            grid_mask = self.streaming(grid_mask)
-            return lax.with_sharding_constraint(grid_mask, self.sharding)
-
-    def bounding_box_indices(self):
-        """
-        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.
-
-        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.
-        """
-        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(self.nx)].
-            bounding_box = {"bottom": np.array([[i, 0] for i in range(self.nx)], dtype=int),
-                           "top": np.array([[i, self.ny - 1] for i in range(self.nx)], dtype=int),
-                           "left": np.array([[0, i] for i in range(self.ny)], dtype=int),
-                           "right": np.array([[self.nx - 1, i] for i in range(self.ny)], dtype=int)}
-                            
-            return bounding_box
-
-        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(self.nx) for j in range(self.ny)].
-            bounding_box = {
-                "bottom": np.array([[i, j, 0] for i in range(self.nx) for j in range(self.ny)], dtype=int),
-                "top": np.array([[i, j, self.nz - 1] for i in range(self.nx) for j in range(self.ny)],dtype=int),
-                "left": np.array([[0, j, k] for j in range(self.ny) for k in range(self.nz)], dtype=int),
-                "right": np.array([[self.nx - 1, j, k] for j in range(self.ny) for k in range(self.nz)], dtype=int),
-                "front": np.array([[i, 0, k] for i in range(self.nx) for k in range(self.nz)], dtype=int),
-                "back": np.array([[i, self.ny - 1, k] for i in range(self.nx) for k in range(self.nz)], dtype=int)}
-
-            return bounding_box
-
-    def set_precisions(self, precision):
-        """
-        This function sets the precision of the computations. The precision is defined by a pair of values,
-        representing the precision of the computation and the precision of the storage, respectively.
-
-        Parameters
-        ----------
-            precision (str): A string representing the desired precision. The string should be in the format
-            "computation/storage", where "computation" and "storage" are either "f64", "f32", or "f16",
-            representing 64-bit, 32-bit, or 16-bit floating point numbers, respectively.
-
-        Returns
-        -------
-            tuple: A pair of jax.numpy data types representing the computation and storage precisions, respectively.
-            If the input string does not match any of the predefined options, the function defaults to (jnp.float32, jnp.float32).
-        """
-        return {
-            "f64/f64": (jnp.float64, jnp.float64),
-            "f32/f32": (jnp.float32, jnp.float32),
-            "f32/f16": (jnp.float32, jnp.float16),
-            "f16/f16": (jnp.float16, jnp.float16),
-            "f64/f32": (jnp.float64, jnp.float32),
-            "f64/f16": (jnp.float64, jnp.float16),
-        }.get(precision, (jnp.float32, jnp.float32))
-
-    def initialize_macroscopic_fields(self):
-        """
-        This function initializes the macroscopic fields (density and velocity) to their default values.
-        The default density is 1 and the default velocity is 0.
-
-        Note: This function is a placeholder and should be overridden in a subclass or in an instance of the class
-        to provide specific initial conditions.
-
-        Returns
-        -------
-            None, None: The default density and velocity, both None. This indicates that the actual values should be set elsewhere.
-        """
-        print("WARNING: Default initial conditions assumed: density = 1, velocity = 0")
-        print("         To set explicit initial density and velocity, use self.initialize_macroscopic_fields.")
-        return None, None
-
-    def assign_fields_sharded(self):
-        """
-        This function is used to initialize the simulation by assigning the macroscopic fields and populations.
-
-        The function first initializes the macroscopic fields, which are the density (rho0) and velocity (u0).
-        Depending on the dimension of the simulation (2D or 3D), it then sets the shape of the array that will hold the 
-        distribution functions (f).
-
-        If the density or velocity are not provided, the function initializes the distribution functions with a default 
-        value (self.w), representing density=1 and velocity=0. Otherwise, it uses the provided density and velocity to initialize the populations.
-
-        Parameters
-        ----------
-        None
-
-        Returns
-        -------
-        f: a distributed JAX array of shape (nx, ny, nz, q) or (nx, ny, q) holding the distribution functions for the simulation.
-        """
-        rho0, u0 = self.initialize_macroscopic_fields()
-
-        if self.dim == 2:
-            shape = (self.nx, self.ny, self.lattice.q)
-        if self.dim == 3:
-            shape = (self.nx, self.ny, self.nz, self.lattice.q)
-    
-        if rho0 is None or u0 is None:
-            f = self.distributed_array_init(shape, self.precisionPolicy.output_dtype, init_val=self.w)
-        else:
-            f = self.initialize_populations(rho0, u0)
-
-        return f
-    
-    def initialize_populations(self, rho0, u0):
-        """
-        This function initializes the populations (distribution functions) for the simulation.
-        It uses the equilibrium distribution function, which is a function of the macroscopic 
-        density and velocity.
-
-        Parameters
-        ----------
-        rho0: jax.numpy.ndarray
-            The initial density field.
-        u0: jax.numpy.ndarray
-            The initial velocity field.
-
-        Returns
-        -------
-        f: jax.numpy.ndarray
-            The array holding the initialized distribution functions for the simulation.
-        """
-        return self.equilibrium(rho0, u0)
-
-    def send_right(self, x, axis_name):
-        """
-        This function sends the data to the right neighboring process in a parallel computing environment.
-        It uses a permutation operation provided by the LAX library.
-
-        Parameters
-        ----------
-        x: jax.numpy.ndarray
-            The data to be sent.
-        axis_name: str
-            The name of the axis along which the data is sent.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The data after being sent to the right neighboring process.
-        """
-        return lax.ppermute(x, perm=self.rightPerm, axis_name=axis_name)
-   
-    def send_left(self, x, axis_name):
-        """
-        This function sends the data to the left neighboring process in a parallel computing environment.
-        It uses a permutation operation provided by the LAX library.
-
-        Parameters
-        ----------
-        x: jax.numpy.ndarray
-            The data to be sent.
-        axis_name: str
-            The name of the axis along which the data is sent.
-
-        Returns
-        -------
-            The data after being sent to the left neighboring process.
-        """
-        return lax.ppermute(x, perm=self.leftPerm, axis_name=axis_name)
-    
-    def streaming_m(self, f):
-        """
-        This function performs the streaming step in the Lattice Boltzmann Method, which is 
-        the propagation of the distribution functions in the lattice.
-
-        To enable multi-GPU/TPU functionality, it extracts the left and right boundary slices of the
-        distribution functions that need to be communicated to the neighboring processes.
-
-        The function then sends the left boundary slice to the right neighboring process and the right 
-        boundary slice to the left neighboring process. The received data is then set to the 
-        corresponding indices in the receiving domain.
-
-        Parameters
-        ----------
-        f: jax.numpy.ndarray
-            The array holding the distribution functions for the simulation.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The distribution functions after the streaming operation.
-        """
-        f = self.streaming_p(f)
-        left_comm, right_comm = f[:1, ..., self.lattice.right_indices], f[-1:, ..., self.lattice.left_indices]
-
-        left_comm, right_comm = self.send_right(left_comm, 'x'), self.send_left(right_comm, 'x')
-        f = f.at[:1, ..., self.lattice.right_indices].set(left_comm)
-        f = f.at[-1:, ..., self.lattice.left_indices].set(right_comm)
-        return f
-
-    @partial(jit, static_argnums=(0,))
-    def streaming_p(self, f):
-        """
-        Perform streaming operation on a partitioned (in the x-direction) distribution function.
-        
-        The function uses the vmap operation provided by the JAX library to vectorize the computation 
-        over all lattice directions.
-
-        Parameters
-        ----------
-            f: The distribution function.
-
-        Returns
-        -------
-            The updated distribution function after streaming.
-        """
-        def streaming_i(f, c):
-            """
-            Perform individual streaming operation in a direction.
-
-            Parameters
-            ----------
-                f: The distribution function.
-                c: The streaming direction vector.
-
-            Returns
-            -------
-                jax.numpy.ndarray
-                The updated distribution function after streaming.
-            """
-            if self.dim == 2:
-                return jnp.roll(f, (c[0], c[1]), axis=(0, 1))
-            elif self.dim == 3:
-                return jnp.roll(f, (c[0], c[1], c[2]), axis=(0, 1, 2))
-
-        return vmap(streaming_i, in_axes=(-1, 0), out_axes=-1)(f, self.c.T)
-
-    @partial(jit, static_argnums=(0, 3), inline=True)
-    def equilibrium(self, rho, u, cast_output=True):
-        """
-        This function computes the equilibrium distribution function in the Lattice Boltzmann Method.
-        The equilibrium distribution function is a function of the macroscopic density and velocity.
-
-        The function first casts the density and velocity to the compute precision if the cast_output flag is True.
-        The function finally casts the equilibrium distribution function to the output precision if the cast_output 
-        flag is True.
-
-        Parameters
-        ----------
-        rho: jax.numpy.ndarray
-            The macroscopic density.
-        u: jax.numpy.ndarray
-            The macroscopic velocity.
-        cast_output: bool, optional
-            A flag indicating whether to cast the density, velocity, and equilibrium distribution function to the 
-            compute and output precisions. Default is True.
-
-        Returns
-        -------
-        feq: ja.numpy.ndarray
-            The equilibrium distribution function.
-        """
-        # Cast the density and velocity to the compute precision if the cast_output flag is True
-        if cast_output:
-            rho, u = self.precisionPolicy.cast_to_compute((rho, u))
-
-        # Cast c to compute precision so that XLA call FXX matmul, 
-        # which is faster (it is faster in some older versions of JAX, newer versions are smart enough to do this automatically)
-        c = jnp.array(self.c, dtype=self.precisionPolicy.compute_dtype)
-        cu = 3.0 * jnp.dot(u, c)
-        usqr = 1.5 * jnp.sum(jnp.square(u), axis=-1, keepdims=True)
-        feq = rho * self.w * (1.0 + cu * (1.0 + 0.5 * cu) - usqr)
-
-        if cast_output:
-            return self.precisionPolicy.cast_to_output(feq)
-        else:
-            return feq
-
-    @partial(jit, static_argnums=(0,))
-    def momentum_flux(self, fneq):
-        """
-        This function computes the momentum flux, which is the product of the non-equilibrium 
-        distribution functions (fneq) and the lattice moments (cc).
-
-        The momentum flux is used in the computation of the stress tensor in the Lattice Boltzmann 
-        Method (LBM).
-
-        Parameters
-        ----------
-        fneq: jax.numpy.ndarray
-            The non-equilibrium distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The computed momentum flux.
-        """
-        return jnp.dot(fneq, self.lattice.cc)
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def update_macroscopic(self, f):
-        """
-        This function computes the macroscopic variables (density and velocity) based on the 
-        distribution functions (f).
-
-        The density is computed as the sum of the distribution functions over all lattice directions. 
-        The velocity is computed as the dot product of the distribution functions and the lattice 
-        velocities, divided by the density.
-
-        Parameters
-        ----------
-        f: jax.numpy.ndarray
-            The distribution functions.
-
-        Returns
-        -------
-        rho: jax.numpy.ndarray
-            Computed density.
-        u: jax.numpy.ndarray
-            Computed velocity.
-        """
-        rho =jnp.sum(f, axis=-1, keepdims=True)
-        c = jnp.array(self.c, dtype=self.precisionPolicy.compute_dtype).T
-        u = jnp.dot(f, c) / rho
-
-        return rho, u
-    
-    @partial(jit, static_argnums=(0, 4), inline=True)
-    def apply_bc(self, fout, fin, timestep, implementation_step):
-        """
-        This function applies the boundary conditions to the distribution functions.
-
-        It iterates over all boundary conditions (BCs) and checks if the implementation step of the 
-        boundary condition matches the provided implementation step. If it does, it applies the 
-        boundary condition to the post-streaming distribution functions (fout).
-
-        Parameters
-        ----------
-        fout: jax.numpy.ndarray
-            The post-collision distribution functions.
-        fin: jax.numpy.ndarray
-            The post-streaming distribution functions.
-        implementation_step: str
-            The implementation step at which the boundary conditions should be applied.
-
-        Returns
-        -------
-        ja.numpy.ndarray
-            The output distribution functions after applying the boundary conditions.
-        """
-        for bc in self.BCs:
-            fout = bc.prepare_populations(fout, fin, implementation_step)
-            if bc.implementationStep == implementation_step:
-                if bc.isDynamic:
-                    fout = bc.apply(fout, fin, timestep)
-                else:
-                    fout = fout.at[bc.indices].set(bc.apply(fout, fin))
-                    
-        return fout
-
-    @partial(jit, static_argnums=(0, 3), donate_argnums=(1,))
-    def step(self, f_poststreaming, timestep, return_fpost=False):
-        """
-        This function performs a single step of the LBM simulation.
-
-        It first performs the collision step, which is the relaxation of the distribution functions 
-        towards their equilibrium values. It then applies the respective boundary conditions to the 
-        post-collision distribution functions.
-
-        The function then performs the streaming step, which is the propagation of the distribution 
-        functions in the lattice. It then applies the respective boundary conditions to the post-streaming 
-        distribution functions.
-
-        Parameters
-        ----------
-        f_poststreaming: jax.numpy.ndarray
-            The post-streaming distribution functions.
-        timestep: int
-            The current timestep of the simulation.
-        return_fpost: bool, optional
-            If True, the function also returns the post-collision distribution functions.
-
-        Returns
-        -------
-        f_poststreaming: jax.numpy.ndarray
-            The post-streaming distribution functions after the simulation step.
-        f_postcollision: jax.numpy.ndarray or None
-            The post-collision distribution functions after the simulation step, or None if 
-            return_fpost is False.
-        """
-        f_postcollision = self.collision(f_poststreaming)
-        f_postcollision = self.apply_bc(f_postcollision, f_poststreaming, timestep, "PostCollision")
-        f_poststreaming = self.streaming(f_postcollision)
-        f_poststreaming = self.apply_bc(f_poststreaming, f_postcollision, timestep, "PostStreaming")
-
-        if return_fpost:
-            return f_poststreaming, f_postcollision
-        else:
-            return f_poststreaming, None
-
-    def run(self, t_max):
-        """
-        This function runs the LBM simulation for a specified number of time steps.
-
-        It first initializes the distribution functions and then enters a loop where it performs the 
-        simulation steps (collision, streaming, and boundary conditions) for each time step.
-
-        The function can also print the progress of the simulation, save the simulation data, and 
-        compute the performance of the simulation in million lattice updates per second (MLUPS).
-
-        Parameters
-        ----------
-        t_max: int
-            The total number of time steps to run the simulation.
-        Returns
-        -------
-        f: jax.numpy.ndarray
-            The distribution functions after the simulation.
-        """
-        f = self.assign_fields_sharded()
-        start_step = 0
-        if self.restore_checkpoint:
-            latest_step = self.mngr.latest_step()
-            if latest_step is not None:  # existing checkpoint present
-                # Assert that the checkpoint manager is not None
-                assert self.mngr is not None, "Checkpoint manager does not exist."
-                state = {'f': f}
-                shardings = jax.tree_map(lambda x: x.sharding, state)
-                restore_args = orb.checkpoint_utils.construct_restore_args(state, shardings)
-                try:
-                    f = self.mngr.restore(latest_step, restore_kwargs={'restore_args': restore_args})['f']
-                    print(f"Restored checkpoint at step {latest_step}.")
-                except ValueError:
-                    raise ValueError(f"Failed to restore checkpoint at step {latest_step}.")
-                
-                start_step = latest_step + 1
-                if not (t_max > start_step):
-                    raise ValueError(f"Simulation already exceeded maximum allowable steps (t_max = {t_max}). Consider increasing t_max.")
-        if self.computeMLUPS:
-            start = time.time()
-        # Loop over all time steps
-        for timestep in range(start_step, t_max + 1):
-            io_flag = self.ioRate > 0 and (timestep % self.ioRate == 0 or timestep == t_max)
-            print_iter_flag = self.printInfoRate> 0 and timestep % self.printInfoRate== 0
-            checkpoint_flag = self.checkpointRate > 0 and timestep % self.checkpointRate == 0
-
-            if io_flag:
-                # Update the macroscopic variables and save the previous values (for error computation)
-                rho_prev, u_prev = self.update_macroscopic(f)
-                rho_prev = downsample_field(rho_prev, self.downsamplingFactor)
-                u_prev = downsample_field(u_prev, self.downsamplingFactor)
-                # Gather the data from all processes and convert it to numpy arrays (move to host memory)
-                rho_prev = process_allgather(rho_prev)
-                u_prev = process_allgather(u_prev)
-
-
-            # Perform one time-step (collision, streaming, and boundary conditions)
-            f, fstar = self.step(f, timestep, return_fpost=self.returnFpost)
-            # Print the progress of the simulation
-            if print_iter_flag:
-                print(colored("Timestep ", 'blue') + colored(f"{timestep}", 'green') + colored(" of ", 'blue') + colored(f"{t_max}", 'green') + colored(" completed", 'blue'))
-
-            if io_flag:
-                # Save the simulation data
-                print(f"Saving data at timestep {timestep}/{t_max}")
-                rho, u = self.update_macroscopic(f)
-                rho = downsample_field(rho, self.downsamplingFactor)
-                u = downsample_field(u, self.downsamplingFactor)
-                
-                # Gather the data from all processes and convert it to numpy arrays (move to host memory)
-                rho = process_allgather(rho)
-                u = process_allgather(u)
-
-                # Save the data
-                self.handle_io_timestep(timestep, f, fstar, rho, u, rho_prev, u_prev)
-            
-            if checkpoint_flag:
-                # Save the checkpoint
-                print(f"Saving checkpoint at timestep {timestep}/{t_max}")
-                state = {'f': f}
-                self.mngr.save(timestep, state)
-            
-            # Start the timer for the MLUPS computation after the first timestep (to remove compilation overhead)
-            if self.computeMLUPS and timestep == 1:
-                jax.block_until_ready(f)
-                start = time.time()
-
-        if self.computeMLUPS:
-            # Compute and print the performance of the simulation in MLUPS
-            jax.block_until_ready(f)
-            end = time.time()
-            if self.dim == 2:
-                print(colored("Domain: ", 'blue') + colored(f"{self.nx} x {self.ny}", 'green') if self.dim == 2 else colored(f"{self.nx} x {self.ny} x {self.nz}", 'green'))
-                print(colored("Number of voxels: ", 'blue') + colored(f"{self.nx * self.ny}", 'green') if self.dim == 2 else colored(f"{self.nx * self.ny * self.nz}", 'green'))
-                print(colored("MLUPS: ", 'blue') + colored(f"{self.nx * self.ny * t_max / (end - start) / 1e6}", 'red'))
-
-            elif self.dim == 3:
-                print(colored("Domain: ", 'blue') + colored(f"{self.nx} x {self.ny} x {self.nz}", 'green'))
-                print(colored("Number of voxels: ", 'blue') + colored(f"{self.nx * self.ny * self.nz}", 'green'))
-                print(colored("MLUPS: ", 'blue') + colored(f"{self.nx * self.ny * self.nz * t_max / (end - start) / 1e6}", 'red'))
-
-        return f
-
-    def handle_io_timestep(self, timestep, f, fstar, rho, u, rho_prev, u_prev):
-        """
-        This function handles the input/output (I/O) operations at each time step of the simulation.
-
-        It prepares the data to be saved and calls the output_data function, which can be overwritten 
-        by the user to customize the I/O operations.
-
-        Parameters
-        ----------
-        timestep: int
-            The current time step of the simulation.
-        f: jax.numpy.ndarray
-            The post-streaming distribution functions at the current time step.
-        fstar: jax.numpy.ndarray
-            The post-collision distribution functions at the current time step.
-        rho: jax.numpy.ndarray
-            The density field at the current time step.
-        u: jax.numpy.ndarray
-            The velocity field at the current time step.
-        """
-        kwargs = {
-            "timestep": timestep,
-            "rho": rho,
-            "rho_prev": rho_prev,
-            "u": u,
-            "u_prev": u_prev,
-            "f_poststreaming": f,
-            "f_postcollision": fstar
-        }
-        self.output_data(**kwargs)
-
-    def output_data(self, **kwargs):
-        """
-        This function is intended to be overwritten by the user to customize the input/output (I/O) 
-        operations of the simulation.
-
-        By default, it does nothing. When overwritten, it could save the simulation data to files, 
-        display the simulation results in real time, send the data to another process for analysis, etc.
-
-        Parameters
-        ----------
-        **kwargs: dict
-            A dictionary containing the simulation data to be outputted. The keys are the names of the 
-            data fields, and the values are the data fields themselves.
-        """
-        pass
-
-    def set_boundary_conditions(self):
-        """
-        This function sets the boundary conditions for the simulation.
-
-        It is intended to be overwritten by the user to specify the boundary conditions according to 
-        the specific problem being solved.
-
-        By default, it does nothing. When overwritten, it could set periodic boundaries, no-slip 
-        boundaries, inflow/outflow boundaries, etc.
-        """
-        pass
-
-    @partial(jit, static_argnums=(0,), donate_argnums=(1,))
-    def collision(self, fin):
-        """
-        This function performs the collision step in the Lattice Boltzmann Method.
-
-        It is intended to be overwritten by the user to specify the collision operator according to 
-        the specific LBM model being used.
-
-        By default, it does nothing. When overwritten, it could implement the BGK collision operator,
-        the MRT collision operator, etc.
-
-        Parameters
-        ----------
-        fin: jax.numpy.ndarray
-            The pre-collision distribution functions.
-
-        Returns
-        -------
-        fin: jax.numpy.ndarray
-            The post-collision distribution functions.
-        """
-        pass
-
-    def get_force(self):
-        """
-        This function computes the force to be applied to the fluid in the Lattice Boltzmann Method.
-
-        It is intended to be overwritten by the user to specify the force according to the specific 
-        problem being solved.
-
-        By default, it does nothing and returns None. When overwritten, it could implement a constant 
-        force term.
-
-        Returns
-        -------
-        force: jax.numpy.ndarray
-            The force to be applied to the fluid.
-        """
-        pass
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def apply_force(self, f_postcollision, feq, rho, u):
-        """
-        add force based on exact-difference method due to Kupershtokh
-
-        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.
-        
-        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.
-        """
-        delta_u = self.get_force()
-        feq_force = self.equilibrium(rho, u + delta_u, cast_output=False)
-        f_postcollision = f_postcollision + feq_force - feq
-        return f_postcollision
-    
-
-    
diff --git a/src/boundary_conditions.py b/src/boundary_conditions.py
deleted file mode 100644
index 6d3352a..0000000
--- a/src/boundary_conditions.py
+++ /dev/null
@@ -1,1178 +0,0 @@
-import jax.numpy as jnp
-from jax import jit, device_count
-from functools import partial
-import numpy as np
-class BoundaryCondition(object):
-    """
-    Base class for boundary conditions in a LBM simulation.
-
-    This class provides a general structure for implementing boundary conditions. It includes methods for preparing the
-    boundary attributes and for applying the boundary condition. Specific boundary conditions should be implemented as
-    subclasses of this class, with the `apply` method overridden as necessary.
-
-    Attributes
-    ----------
-    lattice : Lattice
-        The lattice used in the simulation.
-    nx:
-        The number of nodes in the x direction.
-    ny:
-        The number of nodes in the y direction.
-    nz:
-        The number of nodes in the z direction.
-    dim : int
-        The number of dimensions in the simulation (2 or 3).
-    precision_policy : PrecisionPolicy
-        The precision policy used in the simulation.
-    indices : array-like
-        The indices of the boundary nodes.
-    name : str or None
-        The name of the boundary condition. This should be set in subclasses.
-    isSolid : bool
-        Whether the boundary condition is for a solid boundary. This should be set in subclasses.
-    isDynamic : bool
-        Whether the boundary condition is dynamic (changes over time). This should be set in subclasses.
-    needsExtraConfiguration : bool
-        Whether the boundary condition requires extra configuration. This should be set in subclasses.
-    implementationStep : str
-        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. This should be set in subclasses.
-    """
-
-    def __init__(self, indices, gridInfo, precision_policy):
-        self.lattice = gridInfo["lattice"]
-        self.nx = gridInfo["nx"]
-        self.ny = gridInfo["ny"]
-        self.nz = gridInfo["nz"]
-        self.dim = gridInfo["dim"]
-        self.precisionPolicy = precision_policy
-        self.indices = indices
-        self.name = None
-        self.isSolid = False
-        self.isDynamic = False
-        self.needsExtraConfiguration = False
-        self.implementationStep = "PostStreaming"
-
-    def create_local_mask_and_normal_arrays(self, grid_mask):
-
-        """
-        Creates local mask and normal arrays for the boundary condition.
-
-        Parameters
-        ----------
-        grid_mask : array-like
-            The grid mask for the lattice.
-
-        Returns
-        -------
-        None
-
-        Notes
-        -----
-        This method creates local mask and normal arrays for the boundary condition based on the grid mask.
-        If the boundary condition requires extra configuration, the `configure` method is called.
-        """
-
-        if self.needsExtraConfiguration:
-            boundaryMask = self.get_boundary_mask(grid_mask)
-            self.configure(boundaryMask)
-            self.needsExtraConfiguration = False
-
-        boundaryMask = self.get_boundary_mask(grid_mask)
-        self.normals = self.get_normals(boundaryMask)
-        self.imissing, self.iknown = self.get_missing_indices(boundaryMask)
-        self.imissingMask, self.iknownMask, self.imiddleMask = self.get_missing_mask(boundaryMask)
-
-        return
-
-    def get_boundary_mask(self, grid_mask):  
-        """
-        Add jax.device_count() to the self.indices in x-direction, and 1 to the self.indices other directions
-        This is to make sure the boundary condition is applied to the correct nodes as grid_mask is
-        expanded by (jax.device_count(), 1, 1)
-
-        Parameters
-        ----------
-        grid_mask : array-like
-            The grid mask for the lattice.
-        
-        Returns
-        -------
-        boundaryMask : array-like
-        """   
-        shifted_indices = np.array(self.indices)
-        shifted_indices[0] += device_count()
-        shifted_indices[1:] += 1
-        # Convert back to tuple
-        shifted_indices = tuple(shifted_indices)
-        boundaryMask = np.array(grid_mask[shifted_indices])
-
-        return boundaryMask
-
-    def configure(self, boundaryMask):
-        """
-        Configures the boundary condition.
-
-        Parameters
-        ----------
-        boundaryMask : array-like
-            The grid mask for the boundary voxels.
-
-        Returns
-        -------
-        None
-
-        Notes
-        -----
-        This method should be overridden in subclasses if the boundary condition requires extra configuration.
-        """
-        return
-
-    @partial(jit, static_argnums=(0, 3), inline=True)
-    def prepare_populations(self, fout, fin, implementation_step):
-        """
-        Prepares the distribution functions for the boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The incoming distribution functions.
-        fin : jax.numpy.ndarray
-            The outgoing distribution functions.
-        implementation_step : str
-            The step in the lattice Boltzmann method algorithm at which the preparation is applied.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The prepared distribution functions.
-
-        Notes
-        -----
-        This method should be overridden in subclasses if the boundary condition requires preparation of the distribution functions during post-collision or post-streaming. See ExtrapolationBoundaryCondition for an example.
-        """   
-        return fout
-
-    def get_normals(self, boundaryMask):
-        """
-        Calculates the normal vectors at the boundary nodes.
-
-        Parameters
-        ----------
-        boundaryMask : array-like
-            The boundary mask for the lattice.
-
-        Returns
-        -------
-        array-like
-            The normal vectors at the boundary nodes.
-
-        Notes
-        -----
-        This method calculates the normal vectors by dotting the boundary mask with the main lattice directions.
-        """
-        main_c = self.lattice.c.T[self.lattice.main_indices]
-        m = boundaryMask[..., self.lattice.main_indices]
-        normals = -np.dot(m, main_c)
-        return normals
-
-    def get_missing_indices(self, boundaryMask):
-        """
-        Returns two int8 arrays the same shape as boundaryMask. The non-zero entries of these arrays indicate missing
-        directions that require BCs (imissing) as well as their corresponding opposite directions (iknown).
-
-        Parameters
-        ----------
-        boundaryMask : array-like
-            The boundary mask for the lattice.
-
-        Returns
-        -------
-        tuple of array-like
-            The missing and known indices for the boundary condition.
-
-        Notes
-        -----
-        This method calculates the missing and known indices based on the boundary mask. The missing indices are the
-        non-zero entries of the boundary mask, and the known indices are their corresponding opposite directions.
-        """
-
-        # Find imissing, iknown 1-to-1 corresponding indices
-        # Note: the "zero" index is used as default value here and won't affect BC computations
-        nbd = len(self.indices[0])
-        imissing = np.vstack([np.arange(self.lattice.q, dtype='uint8')] * nbd)
-        iknown = np.vstack([self.lattice.opp_indices] * nbd)
-        imissing[~boundaryMask] = 0
-        iknown[~boundaryMask] = 0
-        return imissing, iknown
-
-    def get_missing_mask(self, boundaryMask):
-        """
-        Returns three boolean arrays the same shape as boundaryMask.
-        Note: these boundary masks are useful for reduction (eg. summation) operators of selected q-directions.
-
-        Parameters
-        ----------
-        boundaryMask : array-like
-            The boundary mask for the lattice.
-
-        Returns
-        -------
-        tuple of array-like
-            The missing, known, and middle masks for the boundary condition.
-
-        Notes
-        -----
-        This method calculates the missing, known, and middle masks based on the boundary mask. The missing mask
-        is the boundary mask, the known mask is the opposite directions of the missing mask, and the middle mask
-        is the directions that are neither missing nor known.
-        """
-        # Find masks for imissing, iknown and imiddle
-        imissingMask = boundaryMask
-        iknownMask = imissingMask[:, self.lattice.opp_indices]
-        imiddleMask = ~(imissingMask | iknownMask)
-        return imissingMask, iknownMask, imiddleMask
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin):
-        """
-        Applies the boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-
-        Returns
-        -------
-        None
-
-        Notes
-        -----
-        This method should be overridden in subclasses to implement the specific boundary condition. The method should
-        modify the output distribution functions in place to apply the boundary condition.
-        """
-        pass
-
-    @partial(jit, static_argnums=(0,))
-    def equilibrium(self, rho, u):
-        """
-        Compute equilibrium distribution function.
-
-        Parameters
-        ----------
-        rho : jax.numpy.ndarray
-            The density at each node in the lattice.
-        u : jax.numpy.ndarray
-            The velocity at each node in the lattice.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The equilibrium distribution function at each node in the lattice.
-
-        Notes
-        -----
-        This method computes the equilibrium distribution function based on the density and velocity. The computation is
-        performed in the compute precision specified by the precision policy. The result is not cast to the output precision as
-        this is function is used inside other functions that require the compute precision.
-        """
-        rho, u = self.precisionPolicy.cast_to_compute((rho, u))
-        c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
-        cu = 3.0 * jnp.dot(u, c)
-        usqr = 1.5 * jnp.sum(u**2, axis=-1, keepdims=True)
-        feq = rho * self.lattice.w * (1.0 + 1.0 * cu + 0.5 * cu**2 - usqr)
-
-        return feq
-
-    @partial(jit, static_argnums=(0,))
-    def momentum_flux(self, fneq):
-        """
-        Compute the momentum flux.
-
-        Parameters
-        ----------
-        fneq : jax.numpy.ndarray
-            The non-equilibrium distribution function at each node in the lattice.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The momentum flux at each node in the lattice.
-
-        Notes
-        -----
-        This method computes the momentum flux by dotting the non-equilibrium distribution function with the lattice
-        direction vectors.
-        """
-        return jnp.dot(fneq, self.lattice.cc)
-
-    @partial(jit, static_argnums=(0,))
-    def momentum_exchange_force(self, f_poststreaming, f_postcollision):
-        """
-        Using the momentum exchange method to compute the boundary force vector exerted on the solid geometry
-        based on [1] as described in [3]. Ref [2] shows how [1] is applicable to curved geometries only by using a
-        bounce-back method (e.g. Bouzidi) that accounts for curved boundaries.
-        NOTE: this function should be called after BC's are imposed.
-        [1] A.J.C. Ladd, Numerical simulations of particular suspensions via a discretized Boltzmann equation.
-            Part 2 (numerical results), J. Fluid Mech. 271 (1994) 311-339.
-        [2] R. Mei, D. Yu, W. Shyy, L.-S. Luo, Force evaluation in the lattice Boltzmann method involving
-            curved geometry, Phys. Rev. E 65 (2002) 041203.
-        [3] Caiazzo, A., & Junk, M. (2008). Boundary forces in lattice Boltzmann: Analysis of momentum exchange
-            algorithm. Computers & Mathematics with Applications, 55(7), 1415-1423.
-
-        Parameters
-        ----------
-        f_poststreaming : jax.numpy.ndarray
-            The post-streaming distribution function at each node in the lattice.
-        f_postcollision : jax.numpy.ndarray
-            The post-collision distribution function at each node in the lattice.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The force exerted on the solid geometry at each boundary node.
-
-        Notes
-        -----
-        This method computes the force exerted on the solid geometry at each boundary node using the momentum exchange method. 
-        The force is computed based on the post-streaming and post-collision distribution functions. This method
-        should be called after the boundary conditions are imposed.
-        """
-        c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
-        nbd = len(self.indices[0])
-        bindex = np.arange(nbd)[:, None]
-        phi = f_postcollision[self.indices][bindex, self.iknown] + \
-              f_poststreaming[self.indices][bindex, self.imissing]
-        force = jnp.sum(c[:, self.iknown] * phi, axis=-1).T
-        return force
-
-class BounceBack(BoundaryCondition):
-    """
-    Bounce-back boundary condition for a lattice Boltzmann method simulation.
-
-    This class implements a full-way bounce-back boundary condition, where particles hitting the boundary are reflected
-    back in the direction they came from. The boundary condition is applied after the collision step.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "BounceBackFullway".
-    implementationStep : str
-        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
-        it is "PostCollision".
-    """
-    def __init__(self, indices, gridInfo, precision_policy):
-        super().__init__(indices, gridInfo, precision_policy)
-        self.name = "BounceBackFullway"
-        self.implementationStep = "PostCollision"
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin):
-        """
-        Applies the bounce-back boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-
-        Notes
-        -----
-        This method applies the bounce-back boundary condition by reflecting the input distribution functions at the
-        boundary nodes in the opposite direction.
-        """
-
-        return fin[self.indices][..., self.lattice.opp_indices]
-
-class BounceBackMoving(BoundaryCondition):
-    """
-    Moving bounce-back boundary condition for a lattice Boltzmann method simulation.
-
-    This class implements a moving bounce-back boundary condition, where particles hitting the boundary are reflected
-    back in the direction they came from, with an additional velocity due to the movement of the boundary. The boundary
-    condition is applied after the collision step.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "BounceBackFullwayMoving".
-    implementationStep : str
-        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
-        it is "PostCollision".
-    isDynamic : bool
-        Whether the boundary condition is dynamic (changes over time). For this class, it is True.
-    update_function : function
-        A function that updates the boundary condition. For this class, it is a function that updates the boundary
-        condition based on the current time step. The signature of the function is `update_function(time) -> (indices, vel)`,
-
-    """
-    def __init__(self, gridInfo, precision_policy, update_function=None):
-        # We get the indices at time zero to pass to the parent class for initialization
-        indices, _ = update_function(0)
-        super().__init__(indices, gridInfo, precision_policy)
-        self.name = "BounceBackFullwayMoving"
-        self.implementationStep = "PostCollision"
-        self.isDynamic = True
-        self.update_function = jit(update_function)
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin, time):
-        """
-        Applies the moving bounce-back boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-        time : int
-            The current time step.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-        """
-        indices, vel = self.update_function(time)
-        c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
-        cu = 6.0 * self.lattice.w * jnp.dot(vel, c)
-        return fout.at[indices].set(fin[indices][..., self.lattice.opp_indices] - cu)
-
-
-class BounceBackHalfway(BoundaryCondition):
-    """
-    Halfway bounce-back boundary condition for a lattice Boltzmann method simulation.
-
-    This class implements a halfway bounce-back boundary condition. The boundary condition is applied after
-    the streaming step.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "BounceBackHalfway".
-    implementationStep : str
-        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
-        it is "PostStreaming".
-    needsExtraConfiguration : bool
-        Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True.
-    isSolid : bool
-        Whether the boundary condition represents a solid boundary. For this class, it is True.
-    vel : array-like
-        The prescribed value of velocity vector for the boundary condition. No-slip BC is assumed if vel=None (default).
-    """
-    def __init__(self, indices, gridInfo, precision_policy, vel=None):
-        super().__init__(indices, gridInfo, precision_policy)
-        self.name = "BounceBackHalfway"
-        self.implementationStep = "PostStreaming"
-        self.needsExtraConfiguration = True
-        self.isSolid = True
-        self.vel = vel
-
-    def configure(self, boundaryMask):
-        """
-        Configures the boundary condition.
-
-        Parameters
-        ----------
-        boundaryMask : array-like
-            The grid mask for the boundary voxels.
-
-        Returns
-        -------
-        None
-
-        Notes
-        -----
-        This method performs an index shift for the halfway bounce-back boundary condition. It updates the indices of
-        the boundary nodes to be the indices of fluid nodes adjacent of the solid nodes.
-        """
-        # Perform index shift for halfway BB.
-        hasFluidNeighbour = ~boundaryMask[:, self.lattice.opp_indices]
-        nbd_orig = len(self.indices[0])
-        idx = np.array(self.indices).T
-        idx_trg = []
-        for i in range(self.lattice.q):
-            idx_trg.append(idx[hasFluidNeighbour[:, i], :] + self.lattice.c[:, i])
-        indices_new = np.unique(np.vstack(idx_trg), axis=0)
-        self.indices = tuple(indices_new.T)
-        nbd_modified = len(self.indices[0])
-        if (nbd_orig != nbd_modified) and self.vel is not None:
-            vel_avg = np.mean(self.vel, axis=0)
-            self.vel = jnp.zeros(indices_new.shape, dtype=self.precisionPolicy.compute_dtype) + vel_avg
-            print("WARNING: assuming a constant averaged velocity vector is imposed at all BC cells!")
-
-        return
-
-    @partial(jit, static_argnums=(0,))
-    def impose_boundary_vel(self, fbd, bindex):
-        c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
-        cu = 6.0 * self.lattice.w * jnp.dot(self.vel, c)
-        fbd = fbd.at[bindex, self.imissing].add(-cu[bindex, self.iknown])
-        return fbd
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin):
-        """
-        Applies the halfway bounce-back boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-        """
-        nbd = len(self.indices[0])
-        bindex = np.arange(nbd)[:, None]
-        fbd = fout[self.indices]
-
-        fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
-        if self.vel is not None:
-            fbd = self.impose_boundary_vel(fbd, bindex)
-        return fbd
-    
-class EquilibriumBC(BoundaryCondition):
-    """
-    Equilibrium boundary condition for a lattice Boltzmann method simulation.
-
-    This class implements an equilibrium boundary condition, where the distribution function at the boundary nodes is
-    set to the equilibrium distribution function. The boundary condition is applied after the streaming step.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "EquilibriumBC".
-    implementationStep : str
-        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
-        it is "PostStreaming".
-    out : jax.numpy.ndarray
-        The equilibrium distribution function at the boundary nodes.
-    """
-
-    def __init__(self, indices, gridInfo, precision_policy, rho, u):
-        super().__init__(indices, gridInfo, precision_policy)
-        self.out = self.precisionPolicy.cast_to_output(self.equilibrium(rho, u))
-        self.name = "EquilibriumBC"
-        self.implementationStep = "PostStreaming"
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin):
-        """
-        Applies the equilibrium boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-
-        Notes
-        -----
-        This method applies the equilibrium boundary condition by setting the output distribution functions at the
-        boundary nodes to the equilibrium distribution function.
-        """
-        return self.out
-
-class DoNothing(BoundaryCondition):
-    def __init__(self, indices, gridInfo, precision_policy):
-        """
-        Do-nothing boundary condition for a lattice Boltzmann method simulation.
-
-        This class implements a do-nothing boundary condition, where no action is taken at the boundary nodes. The boundary
-        condition is applied after the streaming step.
-
-        Attributes
-        ----------
-        name : str
-            The name of the boundary condition. For this class, it is "DoNothing".
-        implementationStep : str
-            The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
-            it is "PostStreaming".
-
-        Notes
-        -----
-        This boundary condition enforces skipping of streaming altogether as it sets post-streaming equal to post-collision
-        populations (so no streaming at this BC voxels). The problem with returning post-streaming values or "fout[self.indices]
-        is that the information that exit the domain on the opposite side of this boundary, would "re-enter". This is because
-        we roll the entire array and so the boundary condition acts like a one-way periodic BC. If EquilibriumBC is used as
-        the BC for that opposite boundary, then the rolled-in values are taken from the initial condition at equilibrium.
-        Otherwise if ZouHe is used for example the simulation looks like a run-down simulation at low-Re. The opposite boundary
-        may be even a wall (consider pipebend example). If we correct imissing directions and assign "fin", this method becomes
-        much less stable and also one needs to correctly take care of corner cases.
-        """
-        super().__init__(indices, gridInfo, precision_policy)
-        self.name = "DoNothing"
-        self.implementationStep = "PostStreaming"
-
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin):
-        """
-        Applies the do-nothing boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-
-        Notes
-        -----
-        This method applies the do-nothing boundary condition by simply returning the input distribution functions at the
-        boundary nodes.
-        """
-        return fin[self.indices]
-
-
-class ZouHe(BoundaryCondition):
-    """
-    Zou-He boundary condition for a lattice Boltzmann method simulation.
-
-    This class implements the Zou-He boundary condition, which is a non-equilibrium bounce-back boundary condition.
-    It can be used to set inflow and outflow boundary conditions with prescribed pressure or velocity.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "ZouHe".
-    implementationStep : str
-        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
-        it is "PostStreaming".
-    type : str
-        The type of the boundary condition. It can be either 'velocity' for a prescribed velocity boundary condition,
-        or 'pressure' for a prescribed pressure boundary condition.
-    prescribed : float or array-like
-        The prescribed values for the boundary condition. It can be either the prescribed velocities for a 'velocity'
-        boundary condition, or the prescribed pressures for a 'pressure' boundary condition.
-
-    References
-    ----------
-    Zou, Q., & He, X. (1997). On pressure and velocity boundary conditions for the lattice Boltzmann BGK model.
-    Physics of Fluids, 9(6), 1591-1598. doi:10.1063/1.869307
-    """
-    def __init__(self, indices, gridInfo, precision_policy, type, prescribed):
-        super().__init__(indices, gridInfo, precision_policy)
-        self.name = "ZouHe"
-        self.implementationStep = "PostStreaming"
-        self.type = type
-        self.prescribed = prescribed
-        self.needsExtraConfiguration = True
-
-    def configure(self, boundaryMask):
-        """
-        Correct boundary indices to ensure that only voxelized surfaces with normal vectors along main cartesian axes
-        are assigned this type of BC.
-        """
-        nv = np.dot(self.lattice.c, ~boundaryMask.T)
-        corner_voxels = np.count_nonzero(nv, axis=0) > 1
-        # removed_voxels = np.array(self.indices)[:, corner_voxels]
-        self.indices = tuple(np.array(self.indices)[:, ~corner_voxels])
-        self.prescribed = self.prescribed[~corner_voxels]
-        return
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def calculate_vel(self, fpop, rho):
-        """
-        Calculate velocity based on the prescribed pressure/density (Zou/He BC)
-        """
-        unormal = -1. + 1. / rho * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) +
-                               2. * jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True))
-
-        # Return the above unormal as a normal vector which sets the tangential velocities to zero
-        vel = unormal * self.normals
-        return vel
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def calculate_rho(self, fpop, vel):
-        """
-        Calculate density based on the prescribed velocity (Zou/He BC)
-        """
-        unormal = np.sum(self.normals*vel, axis=1)
-
-        rho = (1.0/(1.0 + unormal))[..., None] * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) +
-                                  2.*jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True))
-        return rho
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def calculate_equilibrium(self, fpop):
-        """
-        This is the ZouHe method of calculating the missing macroscopic variables at the boundary.
-        """
-        if self.type == 'velocity':
-            vel = self.prescribed
-            rho = self.calculate_rho(fpop, vel)
-        elif self.type == 'pressure':
-            rho = self.prescribed
-            vel = self.calculate_vel(fpop, rho)
-        else:
-            raise ValueError(f"type = {self.type} not supported! Use \'pressure\' or \'velocity\'.")
-
-        # compute feq at the boundary
-        feq = self.equilibrium(rho, vel)
-        return feq
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def bounceback_nonequilibrium(self, fpop, feq):
-        """
-        Calculate unknown populations using bounce-back of non-equilibrium populations
-        a la original Zou & He formulation
-        """
-        nbd = len(self.indices[0])
-        bindex = np.arange(nbd)[:, None]
-        fbd = fpop[self.indices]
-        fknown = fpop[self.indices][bindex, self.iknown] + feq[bindex, self.imissing] - feq[bindex, self.iknown]
-        fbd = fbd.at[bindex, self.imissing].set(fknown)
-        return fbd
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, _):
-        """
-        Applies the Zou-He boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        _ : jax.numpy.ndarray
-            The input distribution functions. This is not used in this method.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-
-        Notes
-        -----
-        This method applies the Zou-He boundary condition by first computing the equilibrium distribution functions based
-        on the prescribed values and the type of boundary condition, and then setting the unknown distribution functions
-        based on the non-equilibrium bounce-back method. 
-        Tangential velocity is not ensured to be zero by adding transverse contributions based on
-        Hecth & Harting (2010) (doi:10.1088/1742-5468/2010/01/P01018) as it caused numerical instabilities at higher
-        Reynolds numbers. One needs to use "Regularized" BC at higher Reynolds.
-        """
-        # compute the equilibrium based on prescribed values and the type of BC
-        feq = self.calculate_equilibrium(fout)
-
-        # set the unknown f populations based on the non-equilibrium bounce-back method
-        fbd = self.bounceback_nonequilibrium(fout, feq)
-
-
-        return fbd
-
-class Regularized(ZouHe):
-    """
-    Regularized boundary condition for a lattice Boltzmann method simulation.
-
-    This class implements the regularized boundary condition, which is a non-equilibrium bounce-back boundary condition
-    with additional regularization. It can be used to set inflow and outflow boundary conditions with prescribed pressure
-    or velocity.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "Regularized".
-    Qi : numpy.ndarray
-        The Qi tensor, which is used in the regularization of the distribution functions.
-
-    References
-    ----------
-    Latt, J. (2007). Hydrodynamic limit of lattice Boltzmann equations. PhD thesis, University of Geneva.
-    Latt, J., Chopard, B., Malaspinas, O., Deville, M., & Michler, A. (2008). Straight velocity boundaries in the
-    lattice Boltzmann method. Physical Review E, 77(5), 056703. doi:10.1103/PhysRevE.77.056703
-    """
-
-    def __init__(self, indices, gridInfo, precision_policy, type, prescribed):
-        super().__init__(indices, gridInfo, precision_policy, type, prescribed)
-        self.name = "Regularized"
-        #TODO for Hesam: check to understand why corner cases cause instability here.
-        # self.needsExtraConfiguration = False
-        self.construct_symmetric_lattice_moment()
-
-    def construct_symmetric_lattice_moment(self):
-        """
-        Construct the symmetric lattice moment Qi.
-
-        The Qi tensor is used in the regularization of the distribution functions. It is defined as Qi = cc - cs^2*I,
-        where cc is the tensor of lattice velocities, cs is the speed of sound, and I is the identity tensor.
-        """
-        Qi = self.lattice.cc
-        if self.dim == 3:
-            diagonal = (0, 3, 5)
-            offdiagonal = (1, 2, 4)
-        elif self.dim == 2:
-            diagonal = (0, 2)
-            offdiagonal = (1,)
-        else:
-            raise ValueError(f"dim = {self.dim} not supported")
-
-        # Qi = cc - cs^2*I
-        Qi = Qi.at[:, diagonal].set(self.lattice.cc[:, diagonal] - 1./3.)
-
-        # multiply off-diagonal elements by 2 because the Q tensor is symmetric
-        Qi = Qi.at[:, offdiagonal].set(self.lattice.cc[:, offdiagonal] * 2.0)
-
-        self.Qi = Qi.T
-        return
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def regularize_fpop(self, fpop, feq):
-        """
-        Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop.
-
-        Parameters
-        ----------
-        fpop : jax.numpy.ndarray
-            The distribution functions.
-        feq : jax.numpy.ndarray
-            The equilibrium distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The regularized distribution functions.
-        """
-
-        # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq}
-        f_neq = fpop - feq
-        PiNeq = self.momentum_flux(f_neq)
-        # PiNeq = self.momentum_flux(fpop) - self.momentum_flux(feq)
-
-        # Compute double dot product Qi:Pi1
-        # QiPi1 = np.zeros_like(fpop)
-        # Pi1 = PiNeq
-        # QiPi1 = jnp.dot(Qi, Pi1)
-        QiPi1 = jnp.dot(PiNeq, self.Qi)
-
-        # assign all populations based on eq 45 of Latt et al (2008)
-        # fneq ~ f^1
-        fpop1 = 9. / 2. * self.lattice.w[None, :] * QiPi1
-        fpop_regularized = feq + fpop1
-
-        return fpop_regularized
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, _):
-        """
-        Applies the regularized boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        _ : jax.numpy.ndarray
-            The input distribution functions. This is not used in this method.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-
-        Notes
-        -----
-        This method applies the regularized boundary condition by first computing the equilibrium distribution functions based
-        on the prescribed values and the type of boundary condition, then setting the unknown distribution functions
-        based on the non-equilibrium bounce-back method, and finally regularizing the distribution functions.
-        """
-
-        # compute the equilibrium based on prescribed values and the type of BC
-        feq = self.calculate_equilibrium(fout)
-
-        # set the unknown f populations based on the non-equilibrium bounce-back method
-        fbd = self.bounceback_nonequilibrium(fout, feq)
-
-        # Regularize the boundary fpop
-        fbd = self.regularize_fpop(fbd, feq)
-        return fbd
-
-
-class ExtrapolationOutflow(BoundaryCondition):
-    """
-    Extrapolation outflow boundary condition for a lattice Boltzmann method simulation.
-
-    This class implements the extrapolation outflow boundary condition, which is a type of outflow boundary condition
-    that uses extrapolation to avoid strong wave reflections.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "ExtrapolationOutflow".
-    sound_speed : float
-        The speed of sound in the simulation.
-
-    References
-    ----------
-    Geier, M., Schönherr, M., Pasquali, A., & Krafczyk, M. (2015). The cumulant lattice Boltzmann equation in three
-    dimensions: Theory and validation. Computers & Mathematics with Applications, 70(4), 507–547.
-    doi:10.1016/j.camwa.2015.05.001.
-    """
-
-    def __init__(self, indices, gridInfo, precision_policy):
-        super().__init__(indices, gridInfo, precision_policy)
-        self.name = "ExtrapolationOutflow"
-        self.needsExtraConfiguration = True
-        self.sound_speed = 1./jnp.sqrt(3.)
-
-    def configure(self, boundaryMask):
-        """
-        Configure the boundary condition by finding neighbouring voxel indices.
-
-        Parameters
-        ----------
-        boundaryMask : np.ndarray
-            The grid mask for the boundary voxels.
-        """        
-        hasFluidNeighbour = ~boundaryMask[:, self.lattice.opp_indices]
-        idx = np.array(self.indices).T
-        idx_trg = []
-        for i in range(self.lattice.q):
-            idx_trg.append(idx[hasFluidNeighbour[:, i], :] + self.lattice.c[:, i])
-        indices_nbr = np.unique(np.vstack(idx_trg), axis=0)
-        self.indices_nbr = tuple(indices_nbr.T)
-
-        return
-
-    @partial(jit, static_argnums=(0, 3), inline=True)
-    def prepare_populations(self, fout, fin, implementation_step):
-        """
-        Prepares the distribution functions for the boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The incoming distribution functions.
-        fin : jax.numpy.ndarray
-            The outgoing distribution functions.
-        implementation_step : str
-            The step in the lattice Boltzmann method algorithm at which the preparation is applied.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The prepared distribution functions.
-
-        Notes
-        -----
-        Because this function is called "PostCollision", f_poststreaming refers to previous time step or t-1
-        """
-        f_postcollision = fout
-        f_poststreaming = fin
-        if implementation_step == 'PostStreaming':
-            return f_postcollision
-        nbd = len(self.indices[0])
-        bindex = np.arange(nbd)[:, None]
-        fps_bdr = f_poststreaming[self.indices]
-        fps_nbr = f_poststreaming[self.indices_nbr]
-        fpc_bdr = f_postcollision[self.indices]
-        fpop = fps_bdr[bindex, self.imissing]
-        fpop_neighbour = fps_nbr[bindex, self.imissing]
-        fpop_extrapolated = self.sound_speed * fpop_neighbour + (1. - self.sound_speed) * fpop
-
-        # Use the iknown directions of f_postcollision that leave the domain during streaming to store the BC data
-        fpc_bdr = fpc_bdr.at[bindex, self.iknown].set(fpop_extrapolated)
-        f_postcollision = f_postcollision.at[self.indices].set(fpc_bdr)
-        return f_postcollision
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin):
-        """
-        Applies the extrapolation outflow boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-        """
-        nbd = len(self.indices[0])
-        bindex = np.arange(nbd)[:, None]
-        fbd = fout[self.indices]
-        fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
-        return fbd
-
-
-class InterpolatedBounceBackBouzidi(BounceBackHalfway):
-    """
-    A local single-node version of the interpolated bounce-back boundary condition due to Bouzidi for a lattice
-    Boltzmann method simulation.
-
-    This class implements a interpolated bounce-back boundary condition. The boundary condition is applied after
-    the streaming step.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "InterpolatedBounceBackBouzidi".
-    implicit_distances : array-like
-        An array of shape (nx,ny,nz) indicating the signed-distance field from the solid walls
-    weights : array-like
-        An array of shape (number_of_bc_cells, q) initialized as None and constructed using implicit_distances array
-        during runtime. These "weights" are associated with the fractional distance of fluid cell to the boundary 
-        position defined as: weights(dir_i) = |x_fluid - x_boundary(dir_i)| / |x_fluid - x_solid(dir_i)|.
-    """
-
-    def __init__(self, indices, implicit_distances, grid_info, precision_policy, vel=None):
-
-        super().__init__(indices, grid_info, precision_policy, vel=vel)
-        self.name = "InterpolatedBounceBackBouzidi"
-        self.implicit_distances = implicit_distances
-        self.weights = None
-
-    def set_proximity_ratio(self):
-        """
-        Creates the interpolation data needed for the boundary condition.
-
-        Returns
-        -------
-        None. The function updates the object's weights attribute in place.
-        """
-        epsilon = 1e-12
-        nbd = len(self.indices[0])
-        idx = np.array(self.indices).T
-        bindex = np.arange(nbd)[:, None]
-        weights = np.full((idx.shape[0], self.lattice.q), 0.5)
-        c = np.array(self.lattice.c)
-        sdf_f = self.implicit_distances[self.indices]
-        for q in range(1, self.lattice.q):
-            solid_indices = idx + c[:, q]
-            solid_indices_tuple = tuple(map(tuple, solid_indices.T))
-            sdf_s = self.implicit_distances[solid_indices_tuple]
-            weights[:, q] = sdf_f / (sdf_f - sdf_s + epsilon)
-        self.weights = weights[bindex, self.iknown]
-        return
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin):
-        """
-        Applies the halfway bounce-back boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-        """
-        if self.weights is None:
-            self.set_proximity_ratio()
-        nbd = len(self.indices[0])
-        bindex = np.arange(nbd)[:, None]
-        fbd = fout[self.indices]
-        f_postcollision_iknown = fin[self.indices][bindex, self.iknown]
-        f_postcollision_imissing = fin[self.indices][bindex, self.imissing]
-        f_poststreaming_iknown = fout[self.indices][bindex, self.iknown]
-
-        # if weights<0.5
-        fs_near = 2. * self.weights * f_postcollision_iknown + \
-                  (1.0 - 2.0 * self.weights) * f_poststreaming_iknown
-
-        # if weights>=0.5
-        fs_far = 1.0 / (2. * self.weights) * f_postcollision_iknown + \
-                 (2.0 * self.weights - 1.0) / (2. * self.weights) * f_postcollision_imissing
-
-        # combine near and far contributions
-        fmissing = jnp.where(self.weights < 0.5, fs_near, fs_far)
-        fbd = fbd.at[bindex, self.imissing].set(fmissing)
-
-        if self.vel is not None:
-            fbd = self.impose_boundary_vel(fbd, bindex)
-        return fbd
-
-
-class InterpolatedBounceBackDifferentiable(InterpolatedBounceBackBouzidi):
-    """
-    A differentiable variant of the "InterpolatedBounceBackBouzidi" BC scheme. This BC is now differentiable at
-    self.weight = 0.5 unlike the original Bouzidi scheme which switches between 2 equations at weight=0.5. Refer to
-    [1] (their Appendix E) for more information.
-
-    References
-    ----------
-    [1] Geier, M., Schönherr, M., Pasquali, A., & Krafczyk, M. (2015). The cumulant lattice Boltzmann equation in three
-    dimensions: Theory and validation. Computers & Mathematics with Applications, 70(4), 507–547.
-    doi:10.1016/j.camwa.2015.05.001.
-
-
-    This class implements a interpolated bounce-back boundary condition. The boundary condition is applied after
-    the streaming step.
-
-    Attributes
-    ----------
-    name : str
-        The name of the boundary condition. For this class, it is "InterpolatedBounceBackDifferentiable".
-    """
-
-    def __init__(self, indices, implicit_distances, grid_info, precision_policy, vel=None):
-
-        super().__init__(indices, implicit_distances, grid_info, precision_policy, vel=vel)
-        self.name = "InterpolatedBounceBackDifferentiable"
-
-
-    @partial(jit, static_argnums=(0,))
-    def apply(self, fout, fin):
-        """
-        Applies the halfway bounce-back boundary condition.
-
-        Parameters
-        ----------
-        fout : jax.numpy.ndarray
-            The output distribution functions.
-        fin : jax.numpy.ndarray
-            The input distribution functions.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            The modified output distribution functions after applying the boundary condition.
-        """
-        if self.weights is None:
-            self.set_proximity_ratio()
-        nbd = len(self.indices[0])
-        bindex = np.arange(nbd)[:, None]
-        fbd = fout[self.indices]
-        f_postcollision_iknown = fin[self.indices][bindex, self.iknown]
-        f_postcollision_imissing = fin[self.indices][bindex, self.imissing]
-        f_poststreaming_iknown = fout[self.indices][bindex, self.iknown]
-        fmissing = ((1. - self.weights) * f_poststreaming_iknown +
-                    self.weights * (f_postcollision_imissing + f_postcollision_iknown)) / (1.0 + self.weights)
-        fbd = fbd.at[bindex, self.imissing].set(fmissing)
-
-        if self.vel is not None:
-            fbd = self.impose_boundary_vel(fbd, bindex)
-        return fbd
\ No newline at end of file
diff --git a/src/lattice.py b/src/lattice.py
deleted file mode 100644
index 788796b..0000000
--- a/src/lattice.py
+++ /dev/null
@@ -1,281 +0,0 @@
-import re
-import numpy as np
-import jax.numpy as jnp
-
-
-class Lattice(object):
-    """
-    This class represents a lattice in the Lattice Boltzmann Method.
-
-    It stores the properties of the lattice, including the dimensions, the number of 
-    velocities, the velocity vectors, the weights, the moments, and the indices of the 
-    opposite, main, right, and left velocities.
-
-    The class also provides methods to construct these properties based on the name of the 
-    lattice.
-
-    Parameters
-    ----------
-    name: str
-        The name of the lattice, which specifies the dimensions and the number of velocities.
-        For example, "D2Q9" represents a 2D lattice with 9 velocities.
-    precision: str, optional
-        The precision of the computations. It can be "f32/f32", "f32/f16", "f64/f64", 
-        "f64/f32", or "f64/f16". The first part before the slash is the precision of the 
-        computations, and the second part after the slash is the precision of the outputs.
-    """
-    def __init__(self, name, precision="f32/f32") -> None:
-        self.name = name
-        dq = re.findall(r"\d+", name)
-        self.precision = precision
-        self.d = int(dq[0])
-        self.q = int(dq[1])
-        if precision == "f32/f32" or precision == "f32/f16":
-            self.precisionPolicy = jnp.float32
-        elif precision == "f64/f64" or precision == "f64/f32" or precision == "f64/f16":
-            self.precisionPolicy = jnp.float64
-        elif precision == "f16/f16":
-            self.precisionPolicy = jnp.float16
-        else:
-            raise ValueError("precision not supported")
-
-        # Construct the properties of the lattice
-        self.c = jnp.array(self.construct_lattice_velocity(), dtype=jnp.int8)
-        self.w = jnp.array(self.construct_lattice_weight(), dtype=self.precisionPolicy)
-        self.cc = jnp.array(self.construct_lattice_moment(), dtype=self.precisionPolicy)
-        self.opp_indices = jnp.array(self.construct_opposite_indices(), dtype=jnp.int8)
-        self.main_indices = jnp.array(self.construct_main_indices(), dtype=jnp.int8)
-        self.right_indices = np.array(self.construct_right_indices(), dtype=jnp.int8)
-        self.left_indices = np.array(self.construct_left_indices(), dtype=jnp.int8)
-
-    def construct_opposite_indices(self):
-        """
-        This function constructs the indices of the opposite velocities for each velocity.
-
-        The opposite velocity of a velocity is the velocity that has the same magnitude but the 
-        opposite direction.
-
-        Returns
-        -------
-        opposite: numpy.ndarray
-            The indices of the opposite velocities.
-        """
-        c = self.c.T
-        opposite = np.array([c.tolist().index((-c[i]).tolist()) for i in range(self.q)])
-        return opposite
-    
-    def construct_right_indices(self):
-        """
-        This function constructs the indices of the velocities that point in the positive 
-        x-direction.
-
-        Returns
-        -------
-        numpy.ndarray
-            The indices of the right velocities.
-        """
-        c = self.c.T
-        return np.nonzero(c[:, 0] == 1)[0]
-    
-    def construct_left_indices(self):
-        """
-        This function constructs the indices of the velocities that point in the negative 
-        x-direction.
-
-        Returns
-        -------
-        numpy.ndarray
-            The indices of the left velocities.
-        """
-        c = self.c.T
-        return np.nonzero(c[:, 0] == -1)[0]
-    
-    def construct_main_indices(self):
-        """
-        This function constructs the indices of the main velocities.
-
-        The main velocities are the velocities that have a magnitude of 1 in lattice units.
-
-        Returns
-        -------
-        numpy.ndarray
-            The indices of the main velocities.
-        """
-        c = self.c.T
-        if self.d == 2:
-            return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) == 1))[0]
-
-        elif self.d == 3:
-            return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) + np.abs(c[:, 2]) == 1))[0]
-
-    def construct_lattice_velocity(self):
-        """
-        This function constructs the velocity vectors of the lattice.
-
-        The velocity vectors are defined based on the name of the lattice. For example, for a D2Q9 
-        lattice, there are 9 velocities: (0,0), (1,0), (-1,0), (0,1), (0,-1), (1,1), (-1,-1), 
-        (1,-1), and (-1,1).
-
-        Returns
-        -------
-        c.T: numpy.ndarray
-            The velocity vectors of the lattice.
-        """
-        if self.name == "D2Q9":  # D2Q9
-            cx = [0, 0, 0, 1, -1, 1, -1, 1, -1]
-            cy = [0, 1, -1, 0, 1, -1, 0, 1, -1]
-            c = np.array(tuple(zip(cx, cy)))
-        elif self.name == "D3Q19":  # D3Q19
-            c = [(x, y, z) for x in [0, -1, 1] for y in [0, -1, 1] for z in [0, -1, 1]]
-            c = np.array([ci for ci in c if np.linalg.norm(ci) < 1.5])
-        elif self.name == "D3Q27":  # D3Q27
-            c = [(x, y, z) for x in [0, -1, 1] for y in [0, -1, 1] for z in [0, -1, 1]]
-            # c = np.array([ci for ci in c if np.linalg.norm(ci) < 1.5])
-            c = np.array(c)
-        else:
-            raise ValueError("Supported Lattice types are D2Q9, D3Q19 and D3Q27")
-
-        return c.T
-
-    def construct_lattice_weight(self):
-        """
-        This function constructs the weights of the lattice.
-
-        The weights are defined based on the name of the lattice. For example, for a D2Q9 lattice, 
-        the weights are 4/9 for the rest velocity, 1/9 for the main velocities, and 1/36 for the 
-        diagonal velocities.
-
-        Returns
-        -------
-        w: numpy.ndarray
-            The weights of the lattice.
-        """
-        # Get the transpose of the lattice vector
-        c = self.c.T
-
-        # Initialize the weights to be 1/36
-        w = 1.0 / 36.0 * np.ones(self.q)
-
-        # Update the weights for 2D and 3D lattices
-        if self.name == "D2Q9":
-            w[np.linalg.norm(c, axis=1) < 1.1] = 1.0 / 9.0
-            w[0] = 4.0 / 9.0
-        elif self.name == "D3Q19":
-            w[np.linalg.norm(c, axis=1) < 1.1] = 2.0 / 36.0
-            w[0] = 1.0 / 3.0
-        elif self.name == "D3Q27":
-            cl = np.linalg.norm(c, axis=1)
-            w[np.isclose(cl, 1.0, atol=1e-8)] = 2.0 / 27.0
-            w[(cl > 1) & (cl <= np.sqrt(2))] = 1.0 / 54.0
-            w[(cl > np.sqrt(2)) & (cl <= np.sqrt(3))] = 1.0 / 216.0
-            w[0] = 8.0 / 27.0
-        else:
-            raise ValueError("Supported Lattice types are D2Q9, D3Q19 and D3Q27")
-
-        # Return the weights
-        return w
-
-    def construct_lattice_moment(self):
-        """
-        This function constructs the moments of the lattice.
-
-        The moments are the products of the velocity vectors, which are used in the computation of 
-        the equilibrium distribution functions and the collision operator in the Lattice Boltzmann 
-        Method (LBM).
-
-        Returns
-        -------
-        cc: numpy.ndarray
-            The moments of the lattice.
-        """
-        c = self.c.T
-        # Counter for the loop
-        cntr = 0
-
-        # nt: number of independent elements of a symmetric tensor
-        nt = self.d * (self.d + 1) // 2
-
-        cc = np.zeros((self.q, nt))
-        for a in range(0, self.d):
-            for b in range(a, self.d):
-                cc[:, cntr] = c[:, a] * c[:, b]
-                cntr += 1
-
-        return cc
-    
-    def __str__(self):
-        return self.name
-
-class LatticeD2Q9(Lattice):
-    """
-    Lattice class for 2D D2Q9 lattice.
-
-    D2Q9 stands for two-dimensional nine-velocity model. It is a common model used in the 
-    Lat tice Boltzmann Method for simulating fluid flows in two dimensions.
-
-    Parameters
-    ----------
-    precision: str, optional
-        The precision of the lattice. The default is "f32/f32"
-    """
-    def __init__(self, precision="f32/f32"):
-        super().__init__("D2Q9", precision)
-        self._set_constants()
-
-    def _set_constants(self):
-        self.cs = jnp.sqrt(3) / 3.0
-        self.cs2 = 1.0 / 3.0
-        self.inv_cs2 = 3.0
-        self.i_s = jnp.asarray(list(range(9)))
-        self.im = 3  # Number of imiddles (includes center)
-        self.ik = 3  # Number of iknowns or iunknowns
-
-
-class LatticeD3Q19(Lattice):
-    """
-    Lattice class for 3D D3Q19 lattice.
-
-    D3Q19 stands for three-dimensional nineteen-velocity model. It is a common model used in the 
-    Lattice Boltzmann Method for simulating fluid flows in three dimensions.
-
-    Parameters
-    ----------
-    precision: str, optional
-        The precision of the lattice. The default is "f32/f32"
-    """
-    def __init__(self, precision="f32/f32"):
-        super().__init__("D3Q19", precision)
-        self._set_constants()
-
-    def _set_constants(self):
-        self.cs = jnp.sqrt(3) / 3.0
-        self.cs2 = 1.0 / 3.0
-        self.inv_cs2 = 3.0
-        self.i_s = jnp.asarray(list(range(19)), dtype=jnp.int8)
-
-        self.im = 9  # Number of imiddles (includes center)
-        self.ik = 5  # Number of iknowns or iunknowns
-
-
-class LatticeD3Q27(Lattice):
-    """
-    Lattice class for 3D D3Q27 lattice.
-
-    D3Q27 stands for three-dimensional twenty-seven-velocity model. It is a common model used in the 
-    Lattice Boltzmann Method for simulating fluid flows in three dimensions.
-
-    Parameters
-    ----------
-    precision: str, optional
-        The precision of the lattice. The default is "f32/f32"
-    """
-
-    def __init__(self, precision="f32/f32"):
-        super().__init__("D3Q27", precision)
-        self._set_constants()
-
-    def _set_constants(self):
-        self.cs = jnp.sqrt(3) / 3.0
-        self.cs2 = 1.0 / 3.0
-        self.inv_cs2 = 3.0
-        self.i_s = jnp.asarray(list(range(27)), dtype=jnp.int8)
\ No newline at end of file
diff --git a/src/models.py b/src/models.py
deleted file mode 100644
index a0500c8..0000000
--- a/src/models.py
+++ /dev/null
@@ -1,260 +0,0 @@
-import jax.numpy as jnp
-from jax import jit
-from functools import partial
-from src.base import LBMBase
-"""
-Collision operators are defined in this file for different models.
-"""
-
-class BGKSim(LBMBase):
-    """
-    BGK simulation class.
-
-    This class implements the Bhatnagar-Gross-Krook (BGK) approximation for the collision step in the Lattice Boltzmann Method.
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    @partial(jit, static_argnums=(0,), donate_argnums=(1,))
-    def collision(self, f):
-        """
-        BGK collision step for lattice.
-
-        The collision step is where the main physics of the LBM is applied. In the BGK approximation, 
-        the distribution function is relaxed towards the equilibrium distribution function.
-        """
-        f = self.precisionPolicy.cast_to_compute(f)
-        rho, u = self.update_macroscopic(f)
-        feq = self.equilibrium(rho, u, cast_output=False)
-        fneq = f - feq
-        fout = f - self.omega * fneq
-        if self.force is not None:
-            fout = self.apply_force(fout, feq, rho, u)
-        return self.precisionPolicy.cast_to_output(fout)
-
-class KBCSim(LBMBase):
-    """
-    KBC simulation class.
-
-    This class implements the Karlin-Bösch-Chikatamarla (KBC) model for the collision step in the Lattice Boltzmann Method.
-    """
-    def __init__(self, **kwargs):
-        if kwargs.get('lattice').name != 'D3Q27' and kwargs.get('nz') > 0:
-            raise ValueError("KBC collision operator in 3D must only be used with D3Q27 lattice.")
-        super().__init__(**kwargs)
-
-    @partial(jit, static_argnums=(0,), donate_argnums=(1,))
-    def collision(self, f):
-        """
-        KBC collision step for lattice.
-        """
-        f = self.precisionPolicy.cast_to_compute(f)
-        tiny = 1e-32
-        beta = self.omega * 0.5
-        rho, u = self.update_macroscopic(f)
-        feq = self.equilibrium(rho, u, cast_output=False)
-        fneq = f - feq
-        if self.dim == 2:
-            deltaS = self.fdecompose_shear_d2q9(fneq) * rho / 4.0
-        else:
-            deltaS = self.fdecompose_shear_d3q27(fneq) * rho
-        deltaH = fneq - deltaS
-        invBeta = 1.0 / beta
-        gamma = invBeta - (2.0 - invBeta) * self.entropic_scalar_product(deltaS, deltaH, feq) / (tiny + self.entropic_scalar_product(deltaH, deltaH, feq))
-
-        fout = f - beta * (2.0 * deltaS + gamma[..., None] * deltaH)
-
-        # add external force
-        if self.force is not None:
-            fout = self.apply_force(fout, feq, rho, u)
-        return self.precisionPolicy.cast_to_output(fout)
-    
-    @partial(jit, static_argnums=(0,), donate_argnums=(1,))
-    def collision_modified(self, f):
-        """
-        Alternative KBC collision step for lattice.
-        Note: 
-        At low Reynolds number the orignal KBC collision above produces inaccurate results because
-        it does not check for the entropy increase/decrease. The KBC stabalizations should only be 
-        applied in principle to cells whose entropy decrease after a regular BGK collision. This is 
-        the case in most cells at higher Reynolds numbers and hence a check may not be needed. 
-        Overall the following alternative collision is more reliable and may replace the original 
-        implementation. The issue at the moment is that it is about 60-80% slower than the above method.
-        """
-        f = self.precisionPolicy.cast_to_compute(f)
-        tiny = 1e-32
-        beta = self.omega * 0.5
-        rho, u = self.update_macroscopic(f)
-        feq = self.equilibrium(rho, u, castOutput=False)
-
-        # Alternative KBC: only stabalizes for voxels whose entropy decreases after BGK collision.
-        f_bgk = f - self.omega * (f - feq)
-        H_fin = jnp.sum(f * jnp.log(f / self.w), axis=-1, keepdims=True)
-        H_fout = jnp.sum(f_bgk * jnp.log(f_bgk / self.w), axis=-1, keepdims=True)
-
-        # the rest is identical to collision_deprecated
-        fneq = f - feq
-        if self.dim == 2:
-            deltaS = self.fdecompose_shear_d2q9(fneq) * rho / 4.0
-        else:
-            deltaS = self.fdecompose_shear_d3q27(fneq) * rho
-        deltaH = fneq - deltaS
-        invBeta = 1.0 / beta
-        gamma = invBeta - (2.0 - invBeta) * self.entropic_scalar_product(deltaS, deltaH, feq) / (tiny + self.entropic_scalar_product(deltaH, deltaH, feq))
-
-        f_kbc = f - beta * (2.0 * deltaS + gamma[..., None] * deltaH)
-        fout = jnp.where(H_fout > H_fin, f_kbc, f_bgk)
-
-        # add external force
-        if self.force is not None:
-            fout = self.apply_force(fout, feq, rho, u)
-        return self.precisionPolicy.cast_to_output(fout)
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def entropic_scalar_product(self, x, y, feq):
-        """
-        Compute the entropic scalar product of x and y to approximate gamma in KBC.
-
-        Returns
-        -------
-        jax.numpy.array
-            Entropic scalar product of x, y, and feq.
-        """
-        return jnp.sum(x * y / feq, axis=-1)
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def fdecompose_shear_d2q9(self, fneq):
-        """
-        Decompose fneq into shear components for D2Q9 lattice.
-
-        Parameters
-        ----------
-        fneq : jax.numpy.array
-            Non-equilibrium distribution function.
-
-        Returns
-        -------
-        jax.numpy.array
-            Shear components of fneq.
-        """
-        Pi = self.momentum_flux(fneq)
-        N = Pi[..., 0] - Pi[..., 2]
-        s = jnp.zeros_like(fneq)
-        s = s.at[..., 6].set(N)
-        s = s.at[..., 3].set(N)
-        s = s.at[..., 2].set(-N)
-        s = s.at[..., 1].set(-N)
-        s = s.at[..., 8].set(Pi[..., 1])
-        s = s.at[..., 4].set(-Pi[..., 1])
-        s = s.at[..., 5].set(-Pi[..., 1])
-        s = s.at[..., 7].set(Pi[..., 1])
-
-        return s
-
-    @partial(jit, static_argnums=(0,), inline=True)
-    def fdecompose_shear_d3q27(self, fneq):
-        """
-        Decompose fneq into shear components for D3Q27 lattice.
-
-        Parameters
-        ----------
-        fneq : jax.numpy.ndarray
-            Non-equilibrium distribution function.
-
-        Returns
-        -------
-        jax.numpy.ndarray
-            Shear components of fneq.
-        """
-        # if self.grid.dim == 3:
-        #     diagonal    = (0, 3, 5)
-        #     offdiagonal = (1, 2, 4)
-        # elif self.grid.dim == 2:
-        #     diagonal    = (0, 2)
-        #     offdiagonal = (1,)
-
-        # c=
-        # array([[0, 0, 0],-----0
-        #        [0, 0, -1],----1
-        #        [0, 0, 1],-----2
-        #        [0, -1, 0],----3
-        #        [0, -1, -1],---4
-        #        [0, -1, 1],----5
-        #        [0, 1, 0],-----6
-        #        [0, 1, -1],----7
-        #        [0, 1, 1],-----8
-        #        [-1, 0, 0],----9
-        #        [-1, 0, -1],--10
-        #        [-1, 0, 1],---11
-        #        [-1, -1, 0],--12
-        #        [-1, -1, -1],-13
-        #        [-1, -1, 1],--14
-        #        [-1, 1, 0],---15
-        #        [-1, 1, -1],--16
-        #        [-1, 1, 1],---17
-        #        [1, 0, 0],----18
-        #        [1, 0, -1],---19
-        #        [1, 0, 1],----20
-        #        [1, -1, 0],---21
-        #        [1, -1, -1],--22
-        #        [1, -1, 1],---23
-        #        [1, 1, 0],----24
-        #        [1, 1, -1],---25
-        #        [1, 1, 1]])---26
-        Pi = self.momentum_flux(fneq)
-        Nxz = Pi[..., 0] - Pi[..., 5]
-        Nyz = Pi[..., 3] - Pi[..., 5]
-
-        # For c = (i, 0, 0), c = (0, j, 0) and c = (0, 0, k)
-        s = jnp.zeros_like(fneq)
-        s = s.at[..., 9].set((2.0 * Nxz - Nyz) / 6.0)
-        s = s.at[..., 18].set((2.0 * Nxz - Nyz) / 6.0)
-        s = s.at[..., 3].set((-Nxz + 2.0 * Nyz) / 6.0)
-        s = s.at[..., 6].set((-Nxz + 2.0 * Nyz) / 6.0)
-        s = s.at[..., 1].set((-Nxz - Nyz) / 6.0)
-        s = s.at[..., 2].set((-Nxz - Nyz) / 6.0)
-
-        # For c = (i, j, 0)
-        s = s.at[..., 12].set(Pi[..., 1] / 4.0)
-        s = s.at[..., 24].set(Pi[..., 1] / 4.0)
-        s = s.at[..., 21].set(-Pi[..., 1] / 4.0)
-        s = s.at[..., 15].set(-Pi[..., 1] / 4.0)
-
-        # For c = (i, 0, k)
-        s = s.at[..., 10].set(Pi[..., 2] / 4.0)
-        s = s.at[..., 20].set(Pi[..., 2] / 4.0)
-        s = s.at[..., 19].set(-Pi[..., 2] / 4.0)
-        s = s.at[..., 11].set(-Pi[..., 2] / 4.0)
-
-        # For c = (0, j, k)
-        s = s.at[..., 8].set(Pi[..., 4] / 4.0)
-        s = s.at[..., 4].set(Pi[..., 4] / 4.0)
-        s = s.at[..., 7].set(-Pi[..., 4] / 4.0)
-        s = s.at[..., 5].set(-Pi[..., 4] / 4.0)
-
-        return s
-
-
-class AdvectionDiffusionBGK(LBMBase):
-    """
-    Advection Diffusion Model based on the BGK model.
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.vel = kwargs.get("vel", None)
-        if self.vel is None:
-            raise ValueError("Velocity must be specified for AdvectionDiffusionBGK.")
-
-    @partial(jit, static_argnums=(0,), donate_argnums=(1,))
-    def collision(self, f):
-        """
-        BGK collision step for lattice.
-        """
-        f = self.precisionPolicy.cast_to_compute(f)
-        rho =jnp.sum(f, axis=-1, keepdims=True)
-        feq = self.equilibrium(rho, self.vel, cast_output=False)
-        fneq = f - feq
-        fout = f - self.omega * fneq
-        return self.precisionPolicy.cast_to_output(fout)
\ No newline at end of file
diff --git a/xlb/__init__.py b/xlb/__init__.py
new file mode 100644
index 0000000..b58db3b
--- /dev/null
+++ b/xlb/__init__.py
@@ -0,0 +1,29 @@
+# Enum classes
+from xlb.compute_backend import ComputeBackend as ComputeBackend
+from xlb.precision_policy import PrecisionPolicy as PrecisionPolicy, Precision as Precision
+from xlb.physics_type import PhysicsType as PhysicsType
+
+# Config
+from .default_config import init as init, DefaultConfig as DefaultConfig
+
+# Velocity Set
+import xlb.velocity_set
+
+# Operators
+import xlb.operator.equilibrium
+import xlb.operator.collision
+import xlb.operator.stream
+import xlb.operator.boundary_condition
+import xlb.operator.macroscopic
+
+# Grids
+import xlb.grid
+
+# Solvers
+import xlb.helper
+
+# Utils
+import xlb.utils
+
+# Distributed computing
+import xlb.distribute
diff --git a/xlb/compute_backend.py b/xlb/compute_backend.py
new file mode 100644
index 0000000..60da291
--- /dev/null
+++ b/xlb/compute_backend.py
@@ -0,0 +1,8 @@
+# Enum used to keep track of the compute backends
+
+from enum import Enum, auto
+
+
+class ComputeBackend(Enum):
+    JAX = auto()
+    WARP = auto()
diff --git a/xlb/default_config.py b/xlb/default_config.py
new file mode 100644
index 0000000..20eac44
--- /dev/null
+++ b/xlb/default_config.py
@@ -0,0 +1,40 @@
+import jax
+from xlb.compute_backend import ComputeBackend
+from dataclasses import dataclass
+from xlb.precision_policy import PrecisionPolicy
+
+
+@dataclass
+class DefaultConfig:
+    default_precision_policy = None
+    velocity_set = None
+    default_backend = None
+
+
+def init(velocity_set, default_backend, default_precision_policy):
+    DefaultConfig.velocity_set = velocity_set
+    DefaultConfig.default_backend = default_backend
+    DefaultConfig.default_precision_policy = default_precision_policy
+
+    if default_backend == ComputeBackend.WARP:
+        import warp as wp
+
+        wp.init()  # TODO: Must be removed in the future versions of WARP
+    elif default_backend == ComputeBackend.JAX:
+        check_multi_gpu_support()
+    else:
+        raise ValueError(f"Unsupported compute backend: {default_backend}")
+
+
+def default_backend() -> ComputeBackend:
+    return DefaultConfig.default_backend
+
+
+def check_multi_gpu_support():
+    gpus = jax.devices("gpu")
+    if len(gpus) > 1:
+        print("Multi-GPU support is available: {} GPUs detected.".format(len(gpus)))
+    elif len(gpus) == 1:
+        print("Single-GPU support is available: 1 GPU detected.")
+    else:
+        print("No GPU support is available; CPU fallback will be used.")
diff --git a/xlb/distribute/__init__.py b/xlb/distribute/__init__.py
new file mode 100644
index 0000000..dd9f33d
--- /dev/null
+++ b/xlb/distribute/__init__.py
@@ -0,0 +1,1 @@
+from .distribute import distribute
diff --git a/xlb/distribute/distribute.py b/xlb/distribute/distribute.py
new file mode 100644
index 0000000..c62b915
--- /dev/null
+++ b/xlb/distribute/distribute.py
@@ -0,0 +1,105 @@
+from jax.sharding import PartitionSpec as P
+from xlb.operator import Operator
+from xlb.operator.stepper import IncompressibleNavierStokesStepper
+from xlb.operator.boundary_condition.boundary_condition import ImplementationStep
+from jax import lax
+from jax.experimental.shard_map import shard_map
+from jax import jit
+
+
+def distribute_operator(
+    operator: Operator,
+    grid,
+    velocity_set,
+    num_results=1,
+    ops="permute",
+) -> Operator:
+    # Define the sharded operator
+    def _sharded_operator(*args):
+        result = operator(*args)
+
+        if ops == "permute":
+            # Define permutation rules for right and left communication
+            rightPerm = [(i, (i + 1) % grid.nDevices) for i in range(grid.nDevices)]
+            leftPerm = [((i + 1) % grid.nDevices, i) for i in range(grid.nDevices)]
+
+            left_comm, right_comm = (
+                result[velocity_set.right_indices, :1, ...],
+                result[velocity_set.left_indices, -1:, ...],
+            )
+
+            left_comm = lax.ppermute(
+                left_comm,
+                perm=rightPerm,
+                axis_name="x",
+            )
+
+            right_comm = lax.ppermute(
+                right_comm,
+                perm=leftPerm,
+                axis_name="x",
+            )
+
+            result = result.at[velocity_set.right_indices, :1, ...].set(left_comm)
+            result = result.at[velocity_set.left_indices, -1:, ...].set(right_comm)
+
+            return result
+        else:
+            raise NotImplementedError(f"Operation {ops} not implemented")
+
+    # Build sharding_flags and in_specs based on args
+    def build_specs(grid, *args):
+        sharding_flags = []
+        in_specs = []
+        for arg in args:
+            if arg.shape[1:] == grid.shape:
+                sharding_flags.append(True)
+            else:
+                sharding_flags.append(False)
+
+        in_specs = tuple(P(*((None, "x") + (grid.dim - 1) * (None,))) if flag else P() for flag in sharding_flags)
+        out_specs = tuple(P(*((None, "x") + (grid.dim - 1) * (None,))) for _ in range(num_results))
+        return tuple(sharding_flags), in_specs, out_specs
+
+    def _wrapped_operator(*args):
+        sharding_flags, in_specs, out_specs = build_specs(grid, *args)
+
+        if len(out_specs) == 1:
+            out_specs = out_specs[0]
+
+        distributed_operator = shard_map(
+            _sharded_operator,
+            mesh=grid.global_mesh,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            check_rep=False,
+        )
+        return distributed_operator(*args)
+
+    return jit(_wrapped_operator)
+
+
+def distribute(operator, grid, velocity_set, num_results=1, ops="permute"):
+    """
+    Distribute an operator or a stepper.
+    If the operator is a stepper, check for post-streaming boundary conditions
+    before deciding how to distribute.
+    """
+    if isinstance(operator, IncompressibleNavierStokesStepper):
+        # Check for post-streaming boundary conditions
+        has_post_streaming_bc = any(bc.implementation_step == ImplementationStep.STREAMING for bc in operator.boundary_conditions)
+
+        if has_post_streaming_bc:
+            # If there are post-streaming BCs, only distribute the stream operator
+            distributed_stream = distribute_operator(operator.stream, grid, velocity_set)
+            operator.stream = distributed_stream
+        else:
+            # If no post-streaming BCs, distribute the whole operator
+            distributed_op = distribute_operator(operator, grid, velocity_set, num_results=num_results, ops=ops)
+            return distributed_op
+
+        return operator
+    else:
+        # For other operators, apply the original distribution logic
+        distributed_op = distribute_operator(operator, grid, velocity_set, num_results=num_results, ops=ops)
+        return distributed_op
diff --git a/xlb/experimental/__init__.py b/xlb/experimental/__init__.py
new file mode 100644
index 0000000..8b13789
--- /dev/null
+++ b/xlb/experimental/__init__.py
@@ -0,0 +1,1 @@
+
diff --git a/xlb/experimental/ooc/__init__.py b/xlb/experimental/ooc/__init__.py
new file mode 100644
index 0000000..5206cc1
--- /dev/null
+++ b/xlb/experimental/ooc/__init__.py
@@ -0,0 +1,2 @@
+from xlb.experimental.ooc.out_of_core import OOCmap
+from xlb.experimental.ooc.ooc_array import OOCArray
diff --git a/xlb/experimental/ooc/ooc_array.py b/xlb/experimental/ooc/ooc_array.py
new file mode 100644
index 0000000..11aedc3
--- /dev/null
+++ b/xlb/experimental/ooc/ooc_array.py
@@ -0,0 +1,440 @@
+import numpy as np
+import cupy as cp
+
+# from mpi4py import MPI
+import itertools
+
+from xlb.experimental.ooc.tiles.dense_tile import DenseTile, DenseGPUTile, DenseCPUTile
+from xlb.experimental.ooc.tiles.compressed_tile import (
+    CompressedTile,
+    CompressedGPUTile,
+    CompressedCPUTile,
+)
+
+
+class OOCArray:
+    """An out-of-core distributed array class.
+
+    Parameters
+    ----------
+    shape : tuple
+        The shape of the array.
+    dtype : cp.dtype
+        The data type of the array.
+    tile_shape : tuple
+        The shape of the tiles. Should be a factor of the shape.
+    padding : int or tuple
+        The padding of the tiles.
+    comm : MPI communicator
+        The MPI communicator.
+    devices : list of cp.cuda.Device
+        The list of GPU devices to use.
+    codec : Codec
+        The codec to use for compression. None for no compression (Dense tiles).
+    nr_compute_tiles : int
+        The number of compute tiles used for asynchronous copies.
+        TODO currently only 1 is supported when using JAX.
+    """
+
+    def __init__(
+        self,
+        shape,
+        dtype,
+        tile_shape,
+        padding=1,
+        comm=None,
+        devices=[cp.cuda.Device(0)],
+        codec=None,
+        nr_compute_tiles=1,
+    ):
+        self.shape = shape
+        self.tile_shape = tile_shape
+        self.dtype = dtype
+        if isinstance(padding, int):
+            padding = (padding,) * len(shape)
+        self.padding = padding
+        self.comm = comm
+        self.devices = devices
+        self.codec = codec
+        self.nr_compute_tiles = nr_compute_tiles
+
+        # Set tile class
+        if self.codec is None:
+            self.Tile = DenseTile
+            self.DeviceTile = DenseGPUTile
+            self.HostTile = DenseCPUTile  # TODO: Possibly make HardDiskTile or something
+
+        else:
+            self.Tile = CompressedTile
+            self.DeviceTile = CompressedGPUTile
+            self.HostTile = CompressedCPUTile
+
+        # Get process id and number of processes
+        self.pid = self.comm.Get_rank()
+        self.nr_proc = self.comm.Get_size()
+
+        # Check that the tile shape divides the array shape.
+        if any([shape[i] % tile_shape[i] != 0 for i in range(len(shape))]):
+            raise ValueError(f"Tile shape {tile_shape} does not divide shape {shape}.")
+        self.tile_dims = tuple([shape[i] // tile_shape[i] for i in range(len(shape))])
+        self.nr_tiles = np.prod(self.tile_dims)
+
+        # Get number of tiles per process
+        if self.nr_tiles % self.nr_proc != 0:
+            raise ValueError(f"Number of tiles {self.nr_tiles} does not divide number of processes {self.nr_proc}.")
+        self.nr_tiles_per_proc = self.nr_tiles // self.nr_proc
+
+        # Make the tile mapppings
+        self.tile_process_map = {}
+        self.tile_device_map = {}
+        for i, tile_index in enumerate(itertools.product(*[range(n) for n in self.tile_dims])):
+            self.tile_process_map[tile_index] = i % self.nr_proc
+            self.tile_device_map[tile_index] = devices[i % len(devices)]  # Checkoboard pattern, TODO: may not be optimal
+
+        # Get my device
+        if self.nr_proc != len(self.devices):
+            raise ValueError(f"Number of processes {self.nr_proc} does not equal number of devices {len(self.devices)}.")
+        self.device = self.devices[self.pid]
+
+        # Make the tiles
+        self.tiles = {}
+        for tile_index in self.tile_process_map.keys():
+            if self.pid == self.tile_process_map[tile_index]:
+                self.tiles[tile_index] = self.HostTile(self.tile_shape, self.dtype, self.padding, self.codec)
+
+        # Make GPU tiles for copying data between CPU and GPU
+        if self.nr_tiles % self.nr_compute_tiles != 0:
+            raise ValueError(
+                f"Number of tiles {self.nr_tiles} does not divide number of compute tiles {self.nr_compute_tiles}. This is used for asynchronous copies."
+            )
+        compute_array_shape = [s + 2 * p for (s, p) in zip(self.tile_shape, self.padding)]
+        self.compute_tiles_htd = []
+        self.compute_tiles_dth = []
+        self.compute_streams_htd = []
+        self.compute_streams_dth = []
+        self.compute_arrays = []
+        self.current_compute_index = 0
+        with cp.cuda.Device(self.device):
+            for i in range(self.nr_compute_tiles):
+                # Make compute tiles for copying data
+                compute_tile = self.DeviceTile(self.tile_shape, self.dtype, self.padding, self.codec)
+                self.compute_tiles_htd.append(compute_tile)
+                compute_tile = self.DeviceTile(self.tile_shape, self.dtype, self.padding, self.codec)
+                self.compute_tiles_dth.append(compute_tile)
+
+                # Make cupy stream
+                self.compute_streams_htd.append(cp.cuda.Stream(non_blocking=True))
+                self.compute_streams_dth.append(cp.cuda.Stream(non_blocking=True))
+
+                # Make compute array
+
+                self.compute_arrays.append(cp.empty(compute_array_shape, self.dtype))
+
+        # Make compute tile mappings
+        self.compute_tile_mapping_htd = {}
+        self.compute_tile_mapping_dth = {}
+        self.compute_stream_mapping_htd = {}
+
+    def size(self):
+        """Return number of allocated bytes for all host tiles."""
+        return sum([tile.size() for tile in self.tiles.values()])
+
+    def nbytes(self):
+        """Return number of bytes for all host tiles."""
+        return sum([tile.nbytes for tile in self.tiles.values()])
+
+    def compression_ratio(self):
+        """Return the compression ratio for all host tiles."""
+        return self.nbytes() / self.size()
+
+    def compression_ratio(self):
+        """Return the compression ratio aggregated over all tiles."""
+
+        if self.codec is None:
+            return 1.0
+        else:
+            total_bytes = 0
+            total_uncompressed_bytes = 0
+            for tile in self.tiles.values():
+                (
+                    tile_total_bytes_uncompressed,
+                    tile_total_bytes_compressed,
+                ) = tile.compression_ratio()
+                total_bytes += tile_total_bytes_compressed
+                total_uncompressed_bytes += tile_total_bytes_uncompressed
+            return total_uncompressed_bytes / total_bytes
+
+    def update_compute_index(self):
+        """Update the current compute index."""
+        self.current_compute_index = (self.current_compute_index + 1) % self.nr_compute_tiles
+
+    def _guess_next_tile_index(self, tile_index):
+        """Guess the next tile index to use for the compute array."""
+        # TODO: This assumes access is sequential
+        tile_indices = list(self.tiles.keys())
+        current_ind = tile_indices.index(tile_index)
+        next_ind = current_ind + 1
+        if next_ind >= len(tile_indices):
+            return None
+        else:
+            return tile_indices[next_ind]
+
+    def reset_queue_htd(self):
+        """Reset the queue for host to device copies."""
+
+        self.compute_tile_mapping_htd = {}
+        self.compute_stream_mapping_htd = {}
+        self.current_compute_index = 0
+
+    def managed_compute_tiles_htd(self, tile_index):
+        """Get the compute tiles needed for computation.
+
+        Parameters
+        ----------
+        tile_index : tuple
+            The tile index.
+
+        Returns
+        -------
+        compute_tile : ComputeTile
+            The compute tile needed for computation.
+        """
+
+        ###################################################
+        # TODO: This assumes access is sequential for tiles
+        ###################################################
+
+        # Que up the next tiles
+        cur_tile_index = tile_index
+        cur_compute_index = self.current_compute_index
+        for i in range(self.nr_compute_tiles):
+            # Check if already in compute tile map and if not que it
+            if cur_tile_index not in self.compute_tile_mapping_htd.keys():
+                # Get the store tile
+                tile = self.tiles[cur_tile_index]
+
+                # Get the compute tile
+                compute_tile = self.compute_tiles_htd[cur_compute_index]
+
+                # Get the compute stream
+                compute_stream = self.compute_streams_htd[cur_compute_index]
+
+                # Copy the tile to the compute tile using the compute stream
+                with compute_stream:
+                    tile.to_gpu_tile(compute_tile)
+                tile.to_gpu_tile(compute_tile)
+
+                # Set the compute tile mapping
+                self.compute_tile_mapping_htd[cur_tile_index] = compute_tile
+                self.compute_stream_mapping_htd[cur_tile_index] = compute_stream
+
+            # Update the tile index and compute index
+            cur_tile_index = self._guess_next_tile_index(cur_tile_index)
+            if cur_tile_index is None:
+                break
+            cur_compute_index = (cur_compute_index + 1) % self.nr_compute_tiles
+
+        # Get the compute tile
+        self.compute_stream_mapping_htd[tile_index].synchronize()
+        compute_tile = self.compute_tile_mapping_htd[tile_index]
+
+        # Pop the tile from the compute tile map
+        self.compute_tile_mapping_htd.pop(tile_index)
+        self.compute_stream_mapping_htd.pop(tile_index)
+
+        # Return the compute tile
+        return compute_tile
+
+    def get_compute_array(self, tile_index):
+        """Given a tile index, copy the tile to the compute array.
+
+        Parameters
+        ----------
+        tile_index : tuple
+            The tile index.
+
+        Returns
+        -------
+        compute_array : array
+            The compute array.
+        global_index : tuple
+            The lower bound index that the compute array corresponds to in the global array.
+            For example, if the compute array is the 0th tile and has padding 1, then the
+            global index will be (-1, -1, ..., -1).
+        """
+
+        # Get the compute tile
+        compute_tile = self.managed_compute_tiles_htd(tile_index)
+
+        # Concatenate the sub-arrays to make the compute array
+        compute_tile.to_array(self.compute_arrays[self.current_compute_index])
+
+        # Return the compute array index in global array
+        global_index = tuple([i * s - p for (i, s, p) in zip(tile_index, self.tile_shape, self.padding)])
+
+        return self.compute_arrays[self.current_compute_index], global_index
+
+    def set_tile(self, compute_array, tile_index):
+        """Given a tile index, copy the compute array to the tile.
+
+        Parameters
+        ----------
+        compute_array : array
+            The compute array.
+        tile_index : tuple
+            The tile index.
+        """
+
+        # Syncronize the current stream dth stream
+        stream = self.compute_streams_dth[self.current_compute_index]
+        stream.synchronize()
+        cp.cuda.get_current_stream().synchronize()
+
+        # Set the compute tile to the correct one
+        compute_tile = self.compute_tiles_dth[self.current_compute_index]
+
+        # Split the compute array into a tile
+        compute_tile.from_array(compute_array)
+
+        # Syncronize the current stream and the compute stream
+        cp.cuda.get_current_stream().synchronize()
+
+        # Copy the tile from the compute tile to the store tile
+        with stream:
+            compute_tile.to_cpu_tile(self.tiles[tile_index])
+        compute_tile.to_cpu_tile(self.tiles[tile_index])
+
+    def update_padding(self):
+        """Perform a padding swap between neighboring tiles."""
+
+        # Get padding indices
+        pad_ind = self.compute_tiles_htd[0].pad_ind
+
+        # Loop over tiles
+        comm_tag = 0
+        for tile_index in self.tile_process_map.keys():
+            # Loop over all padding
+            for pad_index in pad_ind:
+                # Get neighboring tile index
+                neigh_tile_index = tuple([(i + p) % s for (i, p, s) in zip(tile_index, pad_index, self.tile_dims)])
+                neigh_pad_index = tuple([-p for p in pad_index])  # flip
+
+                # 4 cases:
+                # 1. the tile and neighboring tile are on the same process
+                # 2. the tile is on this process and the neighboring tile is on another process
+                # 3. the tile is on another process and the neighboring tile is on this process
+                # 4. the tile and neighboring tile are on different processes
+
+                # Case 1: the tile and neighboring tile are on the same process
+                if self.pid == self.tile_process_map[tile_index] and self.pid == self.tile_process_map[neigh_tile_index]:
+                    # Get the tile and neighboring tile
+                    tile = self.tiles[tile_index]
+                    neigh_tile = self.tiles[neigh_tile_index]
+
+                    # Get pointer to padding and neighboring padding
+                    padding = tile._padding[pad_index]
+                    neigh_padding = neigh_tile._buf_padding[neigh_pad_index]
+
+                    # Swap padding
+                    tile._padding[pad_index] = neigh_padding
+                    neigh_tile._buf_padding[neigh_pad_index] = padding
+
+                # Case 2: the tile is on this process and the neighboring tile is on another process
+                if self.pid == self.tile_process_map[tile_index] and self.pid != self.tile_process_map[neigh_tile_index]:
+                    # Get the tile and padding
+                    tile = self.tiles[tile_index]
+                    padding = tile._padding[pad_index]
+
+                    # Send padding to neighboring process
+                    self.comm.Send(
+                        padding,
+                        dest=self.tile_process_map[neigh_tile_index],
+                        tag=comm_tag,
+                    )
+
+                # Case 3: the tile is on another process and the neighboring tile is on this process
+                if self.pid != self.tile_process_map[tile_index] and self.pid == self.tile_process_map[neigh_tile_index]:
+                    # Get the neighboring tile and padding
+                    neigh_tile = self.tiles[neigh_tile_index]
+                    neigh_padding = neigh_tile._buf_padding[neigh_pad_index]
+
+                    # Receive padding from neighboring process
+                    self.comm.Recv(
+                        neigh_padding,
+                        source=self.tile_process_map[tile_index],
+                        tag=comm_tag,
+                    )
+
+                # Case 4: the tile and neighboring tile are on different processes
+                if self.pid != self.tile_process_map[tile_index] and self.pid != self.tile_process_map[neigh_tile_index]:
+                    pass
+
+                # Increment the communication tag
+                comm_tag += 1
+
+        # Shuffle padding with buffers
+        for tile in self.tiles.values():
+            tile.swap_buf_padding()
+
+    def get_array(self):
+        """Get the full array out from all the sub-arrays. This should only be used for testing."""
+
+        # Get the full array
+        if self.comm.rank == 0:
+            array = np.ones(self.shape, dtype=self.dtype)
+        else:
+            array = None
+
+        # Loop over tiles
+        comm_tag = 0
+        for tile_index in self.tile_process_map.keys():
+            # Set the center array in the full array
+            slice_index = tuple([slice(i * s, (i + 1) * s) for (i, s) in zip(tile_index, self.tile_shape)])
+
+            # if tile on this process compute the center array
+            if self.comm.rank == self.tile_process_map[tile_index]:
+                # Get the tile
+                tile = self.tiles[tile_index]
+
+                # Copy the tile to the compute tile
+                tile.to_gpu_tile(self.compute_tiles_htd[0])
+
+                # Get the compute array
+                self.compute_tiles_htd[0].to_array(self.compute_arrays[0])
+
+                # Get the center array
+                center_array = self.compute_arrays[0][tile._slice_center].get()
+
+            # 4 cases:
+            # 1. the tile is on rank 0 and this process is rank 0
+            # 2. the tile is on another rank and this process is rank 0
+            # 3. the tile is on this rank and this process is not rank 0
+            # 4. the tile is not on rank 0 and this process is not rank 0
+
+            # Case 1: the tile is on rank 0
+            if self.comm.rank == 0 and self.tile_process_map[tile_index] == 0:
+                # Set the center array in the full array
+                array[slice_index] = center_array
+
+            # Case 2: the tile is on another rank and this process is rank 0
+            if self.comm.rank == 0 and self.tile_process_map[tile_index] != 0:
+                # Get the data from the other rank
+                center_array = np.empty(self.tile_shape, dtype=self.dtype)
+                self.comm.Recv(center_array, source=self.tile_process_map[tile_index], tag=comm_tag)
+
+                # Set the center array in the full array
+                array[slice_index] = center_array
+
+            # Case 3: the tile is on this rank and this process is not rank 0
+            if self.comm.rank != 0 and self.tile_process_map[tile_index] == self.comm.rank:
+                # Send the data to rank 0
+                self.comm.Send(center_array, dest=0, tag=comm_tag)
+
+            # Case 4: the tile is not on rank 0 and this process is not rank 0
+            if self.comm.rank != 0 and self.tile_process_map[tile_index] != 0:
+                pass
+
+            # Update the communication tag
+            comm_tag += 1
+
+        return array
diff --git a/xlb/experimental/ooc/out_of_core.py b/xlb/experimental/ooc/out_of_core.py
new file mode 100644
index 0000000..bc42fab
--- /dev/null
+++ b/xlb/experimental/ooc/out_of_core.py
@@ -0,0 +1,104 @@
+# Out-of-core decorator for functions that take a lot of memory
+
+import cupy as cp
+
+from xlb.experimental.ooc.ooc_array import OOCArray
+from xlb.experimental.ooc.utils import (
+    _cupy_to_backend,
+    _backend_to_cupy,
+)
+
+
+def OOCmap(comm, ref_args, add_index=False, backend="jax"):
+    """Decorator for out-of-core functions.
+
+    Parameters
+    ----------
+    comm : MPI communicator
+        The MPI communicator. (TODO add functionality)
+    ref_args : List[int]
+        The indices of the arguments that are OOC arrays to be written to by outputs of the function.
+    add_index : bool, optional
+        Whether to add the index of the global array to the arguments of the function. Default is False.
+        If true the function will take in a tuple of (array, index) instead of just the array.
+    backend : str, optional
+        The backend to use for the function. Default is 'jax'.
+        Options are 'jax' and 'warp'.
+    give_stream : bool, optional
+        Whether to give the function a stream to run on. Default is False.
+        If true the function will take in a last argument of the stream to run on.
+    """
+
+    def decorator(func):
+        def wrapper(*args):
+            # Get list of OOC arrays
+            ooc_array_args = []
+            for arg in args:
+                if isinstance(arg, OOCArray):
+                    ooc_array_args.append(arg)
+
+            # Check that all ooc arrays are compatible
+            # TODO: Add better checks
+            for ooc_array in ooc_array_args:
+                if ooc_array_args[0].tile_dims != ooc_array.tile_dims:
+                    raise ValueError(f"Tile dimensions of ooc arrays do not match. {ooc_array_args[0].tile_dims} != {ooc_array.tile_dims}")
+
+            # Apply the function to each of the ooc arrays
+            for tile_index in ooc_array_args[0].tiles.keys():
+                # Run through args and kwargs and replace ooc arrays with their compute arrays
+                new_args = []
+                for arg in args:
+                    if isinstance(arg, OOCArray):
+                        # Get the compute array (this performs all the memory copies)
+                        compute_array, global_index = arg.get_compute_array(tile_index)
+
+                        # Convert to backend array
+                        compute_array = _cupy_to_backend(compute_array, backend)
+
+                        # Add index to the arguments if requested
+                        if add_index:
+                            compute_array = (compute_array, global_index)
+
+                        new_args.append(compute_array)
+                    else:
+                        new_args.append(arg)
+
+                # Run the function
+                results = func(*new_args)
+
+                # Convert the results to a tuple if not already
+                if not isinstance(results, tuple):
+                    results = (results,)
+
+                # Convert the results back to cupy arrays
+                results = tuple([_backend_to_cupy(result, backend) for result in results])
+
+                # Write the results back to the ooc array
+                for arg_index, result in zip(ref_args, results):
+                    args[arg_index].set_tile(result, tile_index)
+
+                # Update the ooc arrays compute tile index
+                for ooc_array in ooc_array_args:
+                    ooc_array.update_compute_index()
+
+            # Syncronize all processes
+            cp.cuda.Device().synchronize()
+            comm.Barrier()
+
+            # Update the ooc arrays padding
+            for i, ooc_array in enumerate(ooc_array_args):
+                if i in ref_args:
+                    ooc_array.update_padding()
+
+                # Reset que
+                ooc_array.reset_queue_htd()
+
+            # Return OOC arrays
+            if len(ref_args) == 1:
+                return ooc_array_args[ref_args[0]]
+            else:
+                return tuple([args[arg_index] for arg_index in ref_args])
+
+        return wrapper
+
+    return decorator
diff --git a/xlb/experimental/ooc/tiles/__init__.py b/xlb/experimental/ooc/tiles/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/xlb/experimental/ooc/tiles/compressed_tile.py b/xlb/experimental/ooc/tiles/compressed_tile.py
new file mode 100644
index 0000000..ccdd2bb
--- /dev/null
+++ b/xlb/experimental/ooc/tiles/compressed_tile.py
@@ -0,0 +1,248 @@
+import numpy as np
+import cupy as cp
+import warnings
+
+try:
+    from kvikio._lib.arr import asarray
+except ImportError:
+    warnings.warn("kvikio not installed. Compression will not work.")
+
+from xlb.experimental.ooc.tiles.tile import Tile
+from xlb.experimental.ooc.tiles.dense_tile import DenseGPUTile
+from xlb.experimental.ooc.tiles.dynamic_array import DynamicPinnedArray
+
+
+def _decode(comp_array, dest_array, codec):
+    """
+    Decompresses comp_array into dest_array.
+
+    Parameters
+    ----------
+    comp_array : cupy array
+        The compressed array to be decompressed. Data type uint8.
+    dest_array : cupy array
+        The storage array for the decompressed data.
+    codec : Codec
+        The codec to use for decompression. For example, `kvikio.nvcomp.CascadedManager`.
+
+    """
+
+    # Store needed information
+    dtype = dest_array.dtype
+    shape = dest_array.shape
+
+    # Reshape dest_array to match to make into buffer
+    dest_array = dest_array.view(cp.uint8).reshape(-1)
+
+    # Decompress
+    codec._manager.decompress(asarray(dest_array), asarray(comp_array))
+
+    return dest_array.view(dtype).reshape(shape)
+
+
+def _encode(array, dest_array, codec):
+    """
+    Compresses array into dest_array.
+
+    Parameters
+    ----------
+    array : cupy array
+        The array to be compressed.
+    dest_array : cupy array
+        The storage array for the compressed data. Data type uint8.
+    codec : Codec
+        The codec to use for compression. For example, `kvikio.nvcomp.CascadedManager`.
+    """
+
+    # Make sure array is contiguous
+    array = cp.ascontiguousarray(array)
+
+    # Configure compression
+    codec._manager.configure_compression(array.nbytes)
+
+    # Compress
+    size = codec._manager.compress(asarray(array), asarray(dest_array))
+    return size
+
+
+class CompressedTile(Tile):
+    """A Tile where the data is stored in compressed form."""
+
+    def __init__(self, shape, dtype, padding, codec):
+        super().__init__(shape, dtype, padding, codec)
+
+    def allocate_array(self, shape):
+        """Returns a cupy array with the given shape."""
+        raise NotImplementedError
+
+    def to_array(self, array):
+        """Copy a tile to a full array."""
+        # Only implemented for GPU tiles
+        raise NotImplementedError
+
+    def from_array(self, array):
+        """Copy a full array to tile."""
+        # Only implemented for GPU tiles
+        raise NotImplementedError
+
+    def compression_ratio(self):
+        """Returns the compression ratio of the tile."""
+        # Get total number of bytes in tile
+        total_bytes = self._array.size()
+        for pad_ind in self.pad_ind:
+            total_bytes += self._padding[pad_ind].size()
+
+        # Get total number of bytes in uncompressed tile
+        total_bytes_uncompressed = np.prod(self.shape) * self.dtype_itemsize
+        for pad_ind in self.pad_ind:
+            total_bytes_uncompressed += np.prod(self._padding_shape[pad_ind]) * self.dtype_itemsize
+
+        # Return compression ratio
+        return total_bytes_uncompressed, total_bytes
+
+
+class CompressedCPUTile(CompressedTile):
+    """A tile with cells on the CPU."""
+
+    def __init__(self, shape, dtype, padding, codec):
+        super().__init__(shape, dtype, padding, codec)
+
+    def size(self):
+        """Returns the size of the tile in bytes."""
+        size = self._array.size()
+        for pad_ind in self.pad_ind:
+            size += self._padding[pad_ind].size()
+        return size
+
+    def allocate_array(self, shape):
+        """Returns a cupy array with the given shape."""
+        # Make zero array
+        cp_array = cp.zeros(shape, dtype=self.dtype)
+
+        # compress array
+        codec = self.codec()
+        compressed_cp_array = codec.compress(cp_array)
+
+        # Allocate array on CPU
+        array = DynamicPinnedArray(compressed_cp_array.nbytes)
+
+        # Copy array
+        compressed_cp_array.get(out=array.array)
+
+        # add nbytes
+        self.nbytes += cp_array.nbytes
+
+        # delete GPU arrays
+        del compressed_cp_array
+        del cp_array
+
+        return array
+
+    def to_gpu_tile(self, dst_gpu_tile):
+        """Copy tile to a GPU tile."""
+
+        # Check tile is Compressed
+        assert isinstance(dst_gpu_tile, CompressedGPUTile), "Destination tile must be a CompressedGPUTile"
+
+        # Copy array
+        dst_gpu_tile._array[: len(self._array.array)].set(self._array.array)
+        dst_gpu_tile._array_bytes = self._array.nbytes
+
+        # Copy padding
+        for pad_ind in self.pad_ind:
+            dst_gpu_tile._padding[pad_ind][: len(self._padding[pad_ind].array)].set(self._padding[pad_ind].array)
+            dst_gpu_tile._padding_bytes[pad_ind] = self._padding[pad_ind].nbytes
+
+
+class CompressedGPUTile(CompressedTile):
+    """A sub-array with ghost cells on the GPU."""
+
+    def __init__(self, shape, dtype, padding, codec):
+        super().__init__(shape, dtype, padding, codec)
+
+        # Allocate dense GPU tile
+        self.dense_gpu_tile = DenseGPUTile(shape, dtype, padding)
+
+        # Set bytes for each array and padding
+        self._array_bytes = -1
+        self._padding_bytes = {}
+        for pad_ind in self.pad_ind:
+            self._padding_bytes[pad_ind] = -1
+
+        # Set codec for each array and padding
+        self._array_codec = None
+        self._padding_codec = {}
+
+    def allocate_array(self, shape):
+        """Returns a cupy array with the given shape."""
+        nbytes = np.prod(shape) * self.dtype_itemsize
+        codec = self.codec()
+        max_compressed_buffer = codec._manager.configure_compression(nbytes)["max_compressed_buffer_size"]
+        array = cp.zeros((max_compressed_buffer,), dtype=np.uint8)
+        return array
+
+    def to_array(self, array):
+        """Copy a tile to a full array."""
+
+        # Copy center array
+        if self._array_codec is None:
+            self._array_codec = self.codec()
+            self._array_codec._manager.configure_decompression_with_compressed_buffer(asarray(self._array[: self._array_bytes]))
+            self._array_codec.decompression_config = self._array_codec._manager.configure_decompression_with_compressed_buffer(
+                asarray(self._array[: self._array_bytes])
+            )
+        self.dense_gpu_tile._array = _decode(
+            self._array[: self._array_bytes],
+            self.dense_gpu_tile._array,
+            self._array_codec,
+        )
+        array[self._slice_center] = self.dense_gpu_tile._array
+
+        # Copy padding
+        for pad_ind in self.pad_ind:
+            if pad_ind not in self._padding_codec:
+                self._padding_codec[pad_ind] = self.codec()
+                self._padding_codec[pad_ind].decompression_config = self._padding_codec[
+                    pad_ind
+                ]._manager.configure_decompression_with_compressed_buffer(asarray(self._padding[pad_ind][: self._padding_bytes[pad_ind]]))
+            self.dense_gpu_tile._padding[pad_ind] = _decode(
+                self._padding[pad_ind][: self._padding_bytes[pad_ind]],
+                self.dense_gpu_tile._padding[pad_ind],
+                self._padding_codec[pad_ind],
+            )
+            array[self._slice_padding_to_array[pad_ind]] = self.dense_gpu_tile._padding[pad_ind]
+
+    def from_array(self, array):
+        """Copy a full array to tile."""
+
+        # Copy center array
+        if self._array_codec is None:
+            self._array_codec = self.codec()
+            self._array_codec.configure_compression(self._array.nbytes)
+        self._array_bytes = _encode(array[self._slice_center], self._array, self._array_codec)
+
+        # Copy padding
+        for pad_ind in self.pad_ind:
+            if pad_ind not in self._padding_codec:
+                self._padding_codec[pad_ind] = self.codec()
+                self._padding_codec[pad_ind].configure_compression(self._padding[pad_ind].nbytes)
+            self._padding_bytes[pad_ind] = _encode(
+                array[self._slice_array_to_padding[pad_ind]],
+                self._padding[pad_ind],
+                self._padding_codec[pad_ind],
+            )
+
+    def to_cpu_tile(self, dst_cpu_tile):
+        """Copy tile to a CPU tile."""
+
+        # Check tile is Compressed
+        assert isinstance(dst_cpu_tile, CompressedCPUTile), "Destination tile must be a CompressedCPUTile"
+
+        # Copy array
+        dst_cpu_tile._array.resize(self._array_bytes)
+        self._array[: self._array_bytes].get(out=dst_cpu_tile._array.array)
+
+        # Copy padding
+        for pad_ind in self.pad_ind:
+            dst_cpu_tile._padding[pad_ind].resize(self._padding_bytes[pad_ind])
+            self._padding[pad_ind][: self._padding_bytes[pad_ind]].get(out=dst_cpu_tile._padding[pad_ind].array)
diff --git a/xlb/experimental/ooc/tiles/dense_tile.py b/xlb/experimental/ooc/tiles/dense_tile.py
new file mode 100644
index 0000000..41fc129
--- /dev/null
+++ b/xlb/experimental/ooc/tiles/dense_tile.py
@@ -0,0 +1,88 @@
+import numpy as np
+import cupy as cp
+
+from xlb.experimental.ooc.tiles.tile import Tile
+
+
+class DenseTile(Tile):
+    """A Tile where the data is stored in a dense array of the requested dtype."""
+
+    def allocate_array(self, shape):
+        """Returns a cupy array with the given shape."""
+        raise NotImplementedError
+
+    def to_array(self, array):
+        """Copy a tile to a full array."""
+        # TODO: This can be done with a single kernel call, profile to see if it is faster and needs to be done.
+
+        # Copy center array
+        array[self._slice_center] = self._array
+
+        # Copy padding
+        for pad_ind in self.pad_ind:
+            array[self._slice_padding_to_array[pad_ind]] = self._padding[pad_ind]
+
+    def from_array(self, array):
+        """Copy a full array to tile."""
+        # TODO: This can be done with a single kernel call, profile to see if it is faster and needs to be done.
+
+        # Copy center array
+        self._array[...] = array[self._slice_center]
+
+        # Copy padding
+        for pad_ind in self.pad_ind:
+            self._padding[pad_ind][...] = array[self._slice_array_to_padding[pad_ind]]
+
+
+class DenseCPUTile(DenseTile):
+    """A dense tile with cells on the CPU."""
+
+    def __init__(self, shape, dtype, padding, codec=None):
+        super().__init__(shape, dtype, padding, None)
+
+    def allocate_array(self, shape):
+        """Returns a cupy array with the given shape."""
+        # TODO: Seems hacky, but it works. Is there a better way?
+        mem = cp.cuda.alloc_pinned_memory(np.prod(shape) * self.dtype_itemsize)
+        array = np.frombuffer(mem, dtype=self.dtype, count=np.prod(shape)).reshape(shape)
+        self.nbytes += mem.size()
+        return array
+
+    def to_gpu_tile(self, dst_gpu_tile):
+        """Copy tile to a GPU tile."""
+
+        # Check that the destination tile is on the GPU
+        assert isinstance(dst_gpu_tile, DenseGPUTile), "Destination tile must be on GPU"
+
+        # Copy array
+        dst_gpu_tile._array.set(self._array)
+
+        # Copy padding
+        for src_array, dst_gpu_array in zip(self._padding.values(), dst_gpu_tile._padding.values()):
+            dst_gpu_array.set(src_array)
+
+
+class DenseGPUTile(DenseTile):
+    """A sub-array with ghost cells on the GPU."""
+
+    def __init__(self, shape, dtype, padding, codec=None):
+        super().__init__(shape, dtype, padding, None)
+
+    def allocate_array(self, shape):
+        """Returns a cupy array with the given shape."""
+        array = cp.zeros(shape, dtype=self.dtype)
+        self.nbytes += array.nbytes
+        return array
+
+    def to_cpu_tile(self, dst_cpu_tile):
+        """Copy tile to a CPU tile."""
+
+        # Check that the destination tile is on the CPU
+        assert isinstance(dst_cpu_tile, DenseCPUTile), "Destination tile must be on CPU"
+
+        # Copy arra
+        self._array.get(out=dst_cpu_tile._array)
+
+        # Copy padding
+        for src_array, dst_array in zip(self._padding.values(), dst_cpu_tile._padding.values()):
+            src_array.get(out=dst_array)
diff --git a/xlb/experimental/ooc/tiles/dynamic_array.py b/xlb/experimental/ooc/tiles/dynamic_array.py
new file mode 100644
index 0000000..403d164
--- /dev/null
+++ b/xlb/experimental/ooc/tiles/dynamic_array.py
@@ -0,0 +1,66 @@
+# Dynamic array class for pinned memory allocation
+
+import math
+import cupy as cp
+import numpy as np
+
+
+class DynamicArray:
+    """
+    Dynamic pinned memory array class.
+
+    Attributes
+    ----------
+    nbytes : int
+        The number of bytes in the array.
+    bytes_resize : int
+        The number of bytes to resize the array by if the number of bytes requested exceeds the allocated number of bytes.
+    """
+
+    def __init__(self, nbytes, bytes_resize_factor=0.025):
+        # Set the number of bytes
+        self.nbytes = nbytes
+        self.bytes_resize_factor = bytes_resize_factor
+        self.bytes_resize = math.ceil(bytes_resize_factor * nbytes)
+
+        # Set the number of bytes
+        self.allocated_bytes = math.ceil(nbytes / self.bytes_resize) * self.bytes_resize
+
+
+class DynamicPinnedArray(DynamicArray):
+    def __init__(self, nbytes, bytes_resize_factor=0.05):
+        super().__init__(nbytes, bytes_resize_factor)
+
+        # Allocate the memory
+        self.mem = cp.cuda.alloc_pinned_memory(self.allocated_bytes)
+
+        # Make np array that points to the pinned memory
+        self.array = np.frombuffer(self.mem, dtype=np.uint8, count=int(self.nbytes))
+
+    def size(self):
+        return self.mem.size()
+
+    def resize(self, nbytes):
+        # Set the new number of bytes
+        self.nbytes = nbytes
+
+        # Check if the number of bytes requested is less than 2xbytes_resize or if the number of bytes requested exceeds the allocated number of bytes
+        if nbytes < (self.allocated_bytes - 2 * self.bytes_resize) or nbytes > self.allocated_bytes:
+            ## Free the memory
+            # del self.mem
+
+            # Set the new number of allocated bytes
+            self.allocated_bytes = math.ceil(nbytes / self.bytes_resize) * self.bytes_resize
+
+            # Allocate the memory
+            self.mem = cp.cuda.alloc_pinned_memory(self.allocated_bytes)
+
+            # Make np array that points to the pinned memory
+            self.array = np.frombuffer(self.mem, dtype=np.uint8, count=int(self.nbytes))
+
+            # Set new resize number of bytes
+            self.bytes_resize = math.ceil(self.bytes_resize_factor * nbytes)
+
+        # Otherwise change numpy array size
+        else:
+            self.array = np.frombuffer(self.mem, dtype=np.uint8, count=int(self.nbytes))
diff --git a/xlb/experimental/ooc/tiles/tile.py b/xlb/experimental/ooc/tiles/tile.py
new file mode 100644
index 0000000..9bb347b
--- /dev/null
+++ b/xlb/experimental/ooc/tiles/tile.py
@@ -0,0 +1,109 @@
+import cupy as cp
+import itertools
+
+
+class Tile:
+    """Base class for Tile with ghost cells. This tile is used to build a distributed array.
+
+    Attributes
+    ----------
+    shape : tuple
+        Shape of the tile. This will be the shape of the array without padding/ghost cells.
+    dtype : cp.dtype
+        Data type the tile represents. Note that the data data may be stored in a different
+        data type. For example, if it is stored in compressed form.
+    padding : tuple
+        Number of padding/ghost cells in each dimension.
+    """
+
+    def __init__(self, shape, dtype, padding, codec=None):
+        # Store parameters
+        self.shape = shape
+        self.dtype = dtype
+        self.padding = padding
+        self.dtype_itemsize = cp.dtype(self.dtype).itemsize
+        self.nbytes = 0  # Updated when array is allocated
+        self.codec = codec  # Codec to use for compression TODO: Find better abstraction for this
+
+        # Make center array
+        self._array = self.allocate_array(self.shape)
+
+        # Make padding indices
+        pad_dir = []
+        for i in range(len(self.shape)):
+            if self.padding[i] == 0:
+                pad_dir.append((0,))
+            else:
+                pad_dir.append((-1, 0, 1))
+        self.pad_ind = list(itertools.product(*pad_dir))
+        self.pad_ind.remove((0,) * len(self.shape))
+
+        # Make padding and padding buffer arrays
+        self._padding = {}
+        self._buf_padding = {}
+        for ind in self.pad_ind:
+            # determine array shape
+            shape = []
+            for i in range(len(self.shape)):
+                if ind[i] == -1 or ind[i] == 1:
+                    shape.append(self.padding[i])
+                else:
+                    shape.append(self.shape[i])
+
+            # Make padding and padding buffer
+            self._padding[ind] = self.allocate_array(shape)
+            self._buf_padding[ind] = self.allocate_array(shape)
+
+        # Get slicing for array copies
+        self._slice_center = tuple([slice(pad, pad + shape) for (pad, shape) in zip(self.padding, self.shape)])
+        self._slice_padding_to_array = {}
+        self._slice_array_to_padding = {}
+        self._padding_shape = {}
+        for pad_ind in self.pad_ind:
+            slice_padding_to_array = []
+            slice_array_to_padding = []
+            padding_shape = []
+            for pad, ind, s in zip(self.padding, pad_ind, self.shape):
+                if ind == -1:
+                    slice_padding_to_array.append(slice(0, pad))
+                    slice_array_to_padding.append(slice(pad, 2 * pad))
+                    padding_shape.append(pad)
+                elif ind == 0:
+                    slice_padding_to_array.append(slice(pad, s + pad))
+                    slice_array_to_padding.append(slice(pad, s + pad))
+                    padding_shape.append(s)
+                else:
+                    slice_padding_to_array.append(slice(s + pad, s + 2 * pad))
+                    slice_array_to_padding.append(slice(s, s + pad))
+                    padding_shape.append(pad)
+            self._slice_padding_to_array[pad_ind] = tuple(slice_padding_to_array)
+            self._slice_array_to_padding[pad_ind] = tuple(slice_array_to_padding)
+            self._padding_shape[pad_ind] = tuple(padding_shape)
+
+    def size(self):
+        """Returns the number of bytes allocated for the tile."""
+        raise NotImplementedError
+
+    def allocate_array(self, shape):
+        """Returns a cupy array with the given shape."""
+        raise NotImplementedError
+
+    def copy_tile(self, dst_tile):
+        """Copy a tile from one tile to another."""
+        raise NotImplementedError
+
+    def to_array(self, array):
+        """Copy a tile to a full array."""
+        raise NotImplementedError
+
+    def from_array(self, array):
+        """Copy a full array to a tile."""
+        raise NotImplementedError
+
+    def swap_buf_padding(self):
+        """Swap the padding buffer pointer with the padding pointer."""
+        for index in self.pad_ind:
+            (self._buf_padding[index], self._padding[index]) = (
+                self._padding[index],
+                self._buf_padding[index],
+            )
diff --git a/xlb/experimental/ooc/utils.py b/xlb/experimental/ooc/utils.py
new file mode 100644
index 0000000..1179c76
--- /dev/null
+++ b/xlb/experimental/ooc/utils.py
@@ -0,0 +1,79 @@
+import warp as wp
+import cupy as cp
+import jax.dlpack as jdlpack
+import jax
+
+
+def _cupy_to_backend(cupy_array, backend):
+    """
+    Convert cupy array to backend array
+
+    Parameters
+    ----------
+    cupy_array : cupy.ndarray
+        Input cupy array
+    backend : str
+        Backend to convert to. Options are "jax", "warp", or "cupy"
+    """
+
+    # Convert cupy array to backend array
+    dl_array = cupy_array.toDlpack()
+    if backend == "jax":
+        backend_array = jdlpack.from_dlpack(dl_array)
+    elif backend == "warp":
+        backend_array = wp.from_dlpack(dl_array)
+    elif backend == "cupy":
+        backend_array = cupy_array
+    else:
+        raise ValueError(f"Backend {backend} not supported")
+    return backend_array
+
+
+def _backend_to_cupy(backend_array, backend):
+    """
+    Convert backend array to cupy array
+
+    Parameters
+    ----------
+    backend_array : backend.ndarray
+        Input backend array
+    backend : str
+        Backend to convert from. Options are "jax", "warp", or "cupy"
+    """
+
+    # Convert backend array to cupy array
+    if backend == "jax":
+        (jax.device_put(0.0) + 0).block_until_ready()
+        dl_array = jdlpack.to_dlpack(backend_array)
+    elif backend == "warp":
+        dl_array = wp.to_dlpack(backend_array)
+    elif backend == "cupy":
+        return backend_array
+    else:
+        raise ValueError(f"Backend {backend} not supported")
+    cupy_array = cp.fromDlpack(dl_array)
+    return cupy_array
+
+
+def _stream_to_backend(stream, backend):
+    """
+    Convert cupy stream to backend stream
+
+    Parameters
+    ----------
+    stream : cupy.cuda.Stream
+        Input cupy stream
+    backend : str
+        Backend to convert to. Options are "jax", "warp", or "cupy"
+    """
+
+    # Convert stream to backend stream
+    if backend == "jax":
+        raise ValueError("Jax currently does not support streams")
+    if backend == "warp":
+        backend_stream = wp.Stream(cuda_stream=stream.ptr)
+    elif backend == "cupy":
+        backend_stream = stream
+    else:
+        raise ValueError(f"Backend {backend} not supported")
+    return backend_stream
diff --git a/xlb/grid/__init__.py b/xlb/grid/__init__.py
new file mode 100644
index 0000000..7d9ec24
--- /dev/null
+++ b/xlb/grid/__init__.py
@@ -0,0 +1,1 @@
+from xlb.grid.grid import grid_factory as grid_factory
diff --git a/xlb/grid/grid.py b/xlb/grid/grid.py
new file mode 100644
index 0000000..53139fc
--- /dev/null
+++ b/xlb/grid/grid.py
@@ -0,0 +1,86 @@
+from abc import ABC, abstractmethod
+from typing import Tuple
+import numpy as np
+
+from xlb import DefaultConfig
+from xlb.compute_backend import ComputeBackend
+
+
+def grid_factory(shape: Tuple[int, ...], compute_backend: ComputeBackend = None):
+    compute_backend = compute_backend or DefaultConfig.default_backend
+    if compute_backend == ComputeBackend.WARP:
+        from xlb.grid.warp_grid import WarpGrid
+
+        return WarpGrid(shape)
+    elif compute_backend == ComputeBackend.JAX:
+        from xlb.grid.jax_grid import JaxGrid
+
+        return JaxGrid(shape)
+
+    raise ValueError(f"Compute backend {compute_backend} is not supported")
+
+
+class Grid(ABC):
+    def __init__(self, shape: Tuple[int, ...], compute_backend: ComputeBackend):
+        self.shape = shape
+        self.dim = len(shape)
+        self.compute_backend = compute_backend
+        self._initialize_backend()
+
+    @abstractmethod
+    def _initialize_backend(self):
+        pass
+
+    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.
+        """
+
+        # 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 dim == 2:
+            nx, ny = self.shape
+            boundingBoxIndices = {
+                "bottom": grid[:, slice_x, 0],
+                "top": grid[:, slice_x, ny - 1],
+                "left": grid[:, 0, slice_y],
+                "right": grid[:, nx - 1, slice_y],
+            }
+        elif dim == 3:
+            nx, ny, nz = self.shape
+            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 {k: v.tolist() for k, v in boundingBoxIndices.items()}
diff --git a/xlb/grid/jax_grid.py b/xlb/grid/jax_grid.py
new file mode 100644
index 0000000..24eeb03
--- /dev/null
+++ b/xlb/grid/jax_grid.py
@@ -0,0 +1,59 @@
+from typing import Literal
+from jax.sharding import PartitionSpec as P
+from jax.sharding import NamedSharding, Mesh
+from jax.experimental import mesh_utils
+from xlb.compute_backend import ComputeBackend
+
+import jax.numpy as jnp
+import jax
+
+from xlb import DefaultConfig
+
+
+from .grid import Grid
+from xlb.precision_policy import Precision
+
+
+class JaxGrid(Grid):
+    def __init__(self, shape):
+        super().__init__(shape, ComputeBackend.JAX)
+
+    def _initialize_backend(self):
+        self.nDevices = jax.device_count()
+        self.backend = jax.default_backend()
+        self.device_mesh = (
+            mesh_utils.create_device_mesh((1, self.nDevices, 1)) if self.dim == 2 else mesh_utils.create_device_mesh((1, self.nDevices, 1, 1))
+        )
+        self.global_mesh = (
+            Mesh(self.device_mesh, axis_names=("cardinality", "x", "y"))
+            if self.dim == 2
+            else Mesh(self.device_mesh, axis_names=("cardinality", "x", "y", "z"))
+        )
+        self.sharding = (
+            NamedSharding(self.global_mesh, P("cardinality", "x", "y"))
+            if self.dim == 2
+            else NamedSharding(self.global_mesh, P("cardinality", "x", "y", "z"))
+        )
+
+    def create_field(
+        self,
+        cardinality: int,
+        dtype: Literal[Precision.FP32, Precision.FP64, Precision.FP16, Precision.BOOL] = None,
+        fill_value=None,
+    ):
+        sharding_dim = self.shape[0] // self.nDevices
+        device_shape = (cardinality, sharding_dim, *self.shape[1:])
+        full_shape = (cardinality, *self.shape)
+        arrays = []
+
+        dtype = dtype.jax_dtype if dtype else DefaultConfig.default_precision_policy.store_precision.jax_dtype
+
+        for d, index in self.sharding.addressable_devices_indices_map(full_shape).items():
+            jax.default_device = d
+            if fill_value:
+                x = jnp.full(device_shape, fill_value, dtype=dtype)
+            else:
+                x = jnp.zeros(shape=device_shape, dtype=dtype)
+            arrays += [jax.device_put(x, d)]
+        jax.default_device = jax.devices()[0]
+        return jax.make_array_from_single_device_arrays(full_shape, self.sharding, arrays)
diff --git a/xlb/grid/warp_grid.py b/xlb/grid/warp_grid.py
new file mode 100644
index 0000000..c74fc2f
--- /dev/null
+++ b/xlb/grid/warp_grid.py
@@ -0,0 +1,32 @@
+import warp as wp
+
+from .grid import Grid
+from xlb.precision_policy import Precision
+from xlb.compute_backend import ComputeBackend
+from typing import Literal
+from xlb import DefaultConfig
+
+
+class WarpGrid(Grid):
+    def __init__(self, shape):
+        super().__init__(shape, ComputeBackend.WARP)
+
+    def _initialize_backend(self):
+        pass
+
+    def create_field(
+        self,
+        cardinality: int,
+        dtype: Literal[Precision.FP32, Precision.FP64, Precision.FP16] = None,
+        fill_value=None,
+    ):
+        dtype = dtype.wp_dtype if dtype else DefaultConfig.default_precision_policy.store_precision.wp_dtype
+
+        # Check if shape is 2D, and if so, append a singleton dimension to the shape
+        shape = (cardinality,) + (self.shape if len(self.shape) != 2 else self.shape + (1,))
+
+        if fill_value is None:
+            f = wp.zeros(shape, dtype=dtype)
+        else:
+            f = wp.full(shape, fill_value, dtype=dtype)
+        return f
diff --git a/xlb/grid_backend.py b/xlb/grid_backend.py
new file mode 100644
index 0000000..42cd022
--- /dev/null
+++ b/xlb/grid_backend.py
@@ -0,0 +1,9 @@
+# Enum used to keep track of the compute backends
+
+from enum import Enum, auto
+
+
+class GridBackend(Enum):
+    JAX = auto()
+    WARP = auto()
+    OOC = auto()
diff --git a/xlb/helper/__init__.py b/xlb/helper/__init__.py
new file mode 100644
index 0000000..d52f206
--- /dev/null
+++ b/xlb/helper/__init__.py
@@ -0,0 +1,3 @@
+from xlb.helper.nse_solver import create_nse_fields
+from xlb.helper.initializers import initialize_eq
+from xlb.helper.check_boundary_overlaps import check_bc_overlaps
diff --git a/xlb/helper/check_boundary_overlaps.py b/xlb/helper/check_boundary_overlaps.py
new file mode 100644
index 0000000..1adceb2
--- /dev/null
+++ b/xlb/helper/check_boundary_overlaps.py
@@ -0,0 +1,24 @@
+import numpy as np
+from xlb.compute_backend import ComputeBackend
+
+
+def check_bc_overlaps(bclist, dim, backend):
+    index_list = [[] for _ in range(dim)]
+    for bc in bclist:
+        if bc.indices is None:
+            continue
+        # Detect duplicates within bc.indices
+        index_arr = np.unique(bc.indices, axis=-1)
+        if index_arr.shape[-1] != len(bc.indices[0]):
+            if backend == ComputeBackend.WARP:
+                raise ValueError(f"Boundary condition {bc.__class__.__name__} has duplicate indices!")
+            print(f"WARNING: there are duplicate indices in {bc.__class__.__name__} and hence the order in bc list matters!")
+        for d in range(dim):
+            index_list[d] += bc.indices[d]
+
+    # Detect duplicates within bclist
+    index_arr = np.unique(index_list, axis=-1)
+    if index_arr.shape[-1] != len(index_list[0]):
+        if backend == ComputeBackend.WARP:
+            raise ValueError("Boundary condition list containes duplicate indices!")
+        print("WARNING: there are duplicate indices in the boundary condition list and hence the order in this list matters!")
diff --git a/xlb/helper/initializers.py b/xlb/helper/initializers.py
new file mode 100644
index 0000000..ccb4a82
--- /dev/null
+++ b/xlb/helper/initializers.py
@@ -0,0 +1,20 @@
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.equilibrium import QuadraticEquilibrium
+
+
+def initialize_eq(f, grid, velocity_set, precision_policy, backend, rho=None, u=None):
+    if rho is None:
+        rho = grid.create_field(cardinality=1, fill_value=1.0, dtype=precision_policy.compute_precision)
+    if u is None:
+        u = grid.create_field(cardinality=velocity_set.d, fill_value=0.0, dtype=precision_policy.compute_precision)
+    equilibrium = QuadraticEquilibrium()
+
+    if backend == ComputeBackend.JAX:
+        f = equilibrium(rho, u)
+
+    elif backend == ComputeBackend.WARP:
+        f = equilibrium(rho, u, f)
+
+    del rho, u
+
+    return f
diff --git a/xlb/helper/nse_solver.py b/xlb/helper/nse_solver.py
new file mode 100644
index 0000000..3d0ef3e
--- /dev/null
+++ b/xlb/helper/nse_solver.py
@@ -0,0 +1,19 @@
+from xlb import DefaultConfig
+from xlb.grid import grid_factory
+from xlb.precision_policy import Precision
+from typing import Tuple
+
+
+def create_nse_fields(grid_shape: Tuple[int, int, int], velocity_set=None, compute_backend=None, precision_policy=None):
+    velocity_set = velocity_set or DefaultConfig.velocity_set
+    compute_backend = compute_backend or DefaultConfig.default_backend
+    precision_policy = precision_policy or DefaultConfig.default_precision_policy
+    grid = grid_factory(grid_shape, compute_backend=compute_backend)
+
+    # Create fields
+    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)
+    bc_mask = grid.create_field(cardinality=1, dtype=Precision.UINT8)
+
+    return grid, f_0, f_1, missing_mask, bc_mask
diff --git a/xlb/operator/__init__.py b/xlb/operator/__init__.py
new file mode 100644
index 0000000..02b8a59
--- /dev/null
+++ b/xlb/operator/__init__.py
@@ -0,0 +1,2 @@
+from xlb.operator.operator import Operator
+from xlb.operator.parallel_operator import ParallelOperator
diff --git a/xlb/operator/boundary_condition/__init__.py b/xlb/operator/boundary_condition/__init__.py
new file mode 100644
index 0000000..4782ea0
--- /dev/null
+++ b/xlb/operator/boundary_condition/__init__.py
@@ -0,0 +1,10 @@
+from xlb.operator.boundary_condition.boundary_condition import BoundaryCondition
+from xlb.operator.boundary_condition.boundary_condition_registry import BoundaryConditionRegistry
+from xlb.operator.boundary_condition.bc_equilibrium import EquilibriumBC
+from xlb.operator.boundary_condition.bc_do_nothing import DoNothingBC
+from xlb.operator.boundary_condition.bc_halfway_bounce_back import HalfwayBounceBackBC
+from xlb.operator.boundary_condition.bc_fullway_bounce_back import FullwayBounceBackBC
+from xlb.operator.boundary_condition.bc_zouhe import ZouHeBC
+from xlb.operator.boundary_condition.bc_regularized import RegularizedBC
+from xlb.operator.boundary_condition.bc_extrapolation_outflow import ExtrapolationOutflowBC
+from xlb.operator.boundary_condition.bc_grads_approximation import GradsApproximationBC
diff --git a/xlb/operator/boundary_condition/bc_do_nothing.py b/xlb/operator/boundary_condition/bc_do_nothing.py
new file mode 100644
index 0000000..56a332f
--- /dev/null
+++ b/xlb/operator/boundary_condition/bc_do_nothing.py
@@ -0,0 +1,79 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+from functools import partial
+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.boundary_condition.boundary_condition import (
+    ImplementationStep,
+    BoundaryCondition,
+)
+from xlb.operator.boundary_condition.boundary_condition_registry import (
+    boundary_condition_registry,
+)
+
+
+class DoNothingBC(BoundaryCondition):
+    """
+    Do nothing boundary condition. This boundary condition skips the streaming step for the
+    boundary nodes.
+    """
+
+    def __init__(
+        self,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        super().__init__(
+            ImplementationStep.STREAMING,
+            velocity_set,
+            precision_policy,
+            compute_backend,
+            indices,
+            mesh_vertices,
+        )
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    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):
+        # Construct the functional for this BC
+        @wp.func
+        def functional(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            return f_pre
+
+        kernel = self._construct_kernel(functional)
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    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, bc_mask, missing_mask],
+            dim=f_pre.shape[1:],
+        )
+        return f_post
\ No newline at end of file
diff --git a/xlb/operator/boundary_condition/bc_equilibrium.py b/xlb/operator/boundary_condition/bc_equilibrium.py
new file mode 100644
index 0000000..77f408f
--- /dev/null
+++ b/xlb/operator/boundary_condition/bc_equilibrium.py
@@ -0,0 +1,104 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import jax.lax as lax
+from functools import partial
+import warp as wp
+from typing import Tuple, Any
+
+from xlb.velocity_set.velocity_set import VelocitySet
+from xlb.precision_policy import PrecisionPolicy
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.equilibrium.equilibrium import Equilibrium
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.operator.operator import Operator
+from xlb.operator.boundary_condition.boundary_condition import (
+    ImplementationStep,
+    BoundaryCondition,
+)
+from xlb.operator.boundary_condition.boundary_condition_registry import (
+    boundary_condition_registry,
+)
+
+
+class EquilibriumBC(BoundaryCondition):
+    """
+    Full Bounce-back boundary condition for a lattice Boltzmann method simulation.
+    """
+
+    def __init__(
+        self,
+        rho: float,
+        u: Tuple[float, float, float],
+        equilibrium_operator: Operator = None,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        # Store the equilibrium information
+        self.rho = rho
+        self.u = u
+        self.equilibrium_operator = QuadraticEquilibrium() if equilibrium_operator is None else equilibrium_operator
+        # Raise error if equilibrium operator is not a subclass of Equilibrium
+        if not issubclass(type(self.equilibrium_operator), Equilibrium):
+            raise ValueError("Equilibrium operator must be a subclass of Equilibrium")
+
+        # Call the parent constructor
+        super().__init__(
+            ImplementationStep.STREAMING,
+            velocity_set,
+            precision_policy,
+            compute_backend,
+            indices,
+            mesh_vertices,
+        )
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    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 = bc_mask == self.id
+
+        return jnp.where(boundary, feq, f_post)
+
+    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)
+        _rho = self.compute_dtype(self.rho)
+        _u = _u_vec(self.u[0], self.u[1], self.u[2]) if self.velocity_set.d == 3 else _u_vec(self.u[0], self.u[1])
+
+        # Construct the functional for this BC
+        @wp.func
+        def functional(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            _f = self.equilibrium_operator.warp_functional(_rho, _u)
+            return _f
+
+        # Use the parent class's kernel and pass the functional
+        kernel = self._construct_kernel(functional)
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    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, bc_mask, missing_mask],
+            dim=f_pre.shape[1:],
+        )
+        return f_post
\ No newline at end of file
diff --git a/xlb/operator/boundary_condition/bc_extrapolation_outflow.py b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py
new file mode 100644
index 0000000..38657e5
--- /dev/null
+++ b/xlb/operator/boundary_condition/bc_extrapolation_outflow.py
@@ -0,0 +1,208 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import jax.lax as lax
+from functools import partial
+import warp as wp
+from typing import Any
+from collections import Counter
+import numpy as np
+
+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.boundary_condition.boundary_condition import (
+    ImplementationStep,
+    BoundaryCondition,
+)
+from xlb.operator.boundary_condition.boundary_condition_registry import (
+    boundary_condition_registry,
+)
+
+
+class ExtrapolationOutflowBC(BoundaryCondition):
+    """
+    Extrapolation outflow boundary condition for a lattice Boltzmann method simulation.
+
+    This class implements the extrapolation outflow boundary condition, which is a type of outflow boundary condition
+    that uses extrapolation to avoid strong wave reflections.
+
+    References
+    ----------
+    Geier, M., Schönherr, M., Pasquali, A., & Krafczyk, M. (2015). The cumulant lattice Boltzmann equation in three
+    dimensions: Theory and validation. Computers & Mathematics with Applications, 70(4), 507-547.
+    doi:10.1016/j.camwa.2015.05.001.
+    """
+
+    def __init__(
+        self,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        # Call the parent constructor
+        super().__init__(
+            ImplementationStep.STREAMING,
+            velocity_set,
+            precision_policy,
+            compute_backend,
+            indices,
+            mesh_vertices,
+        )
+
+        # find and store the normal vector using indices
+        self._get_normal_vec(indices)
+
+        # Unpack the two warp functionals needed for this BC!
+        if self.compute_backend == ComputeBackend.WARP:
+            self.warp_functional, self.prepare_bc_auxilary_data = self.warp_functional
+
+    def _get_normal_vec(self, indices):
+        # Get the frequency count and most common element directly
+        freq_counts = [Counter(coord).most_common(1)[0] for coord in indices]
+
+        # Extract counts and elements
+        counts = np.array([count for _, count in freq_counts])
+        elements = np.array([element for element, _ in freq_counts])
+
+        # Normalize the counts
+        self.normal = counts // counts.max()
+
+        # Reverse the normal vector if the most frequent element is 0
+        if elements[np.argmax(counts)] == 0:
+            self.normal *= -1
+        return
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def _roll(self, fld, vec):
+        """
+        Perform rolling operation of a field with dimentions [q, nx, ny, nz] in a direction
+        given by vec. All q-directions are rolled at the same time.
+        # TODO: how to improve this for multi-gpu runs?
+        """
+        if self.velocity_set.d == 2:
+            return jnp.roll(fld, (vec[0], vec[1]), axis=(1, 2))
+        elif self.velocity_set.d == 3:
+            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, 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 = 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)))
+
+        # Roll boundary mask in the opposite of the normal vector to mask its next immediate neighbour
+        neighbour = self._roll(boundary, -self.normal)
+
+        # gather post-streaming values associated with previous time-step to construct the auxilary data for BC
+        fpop = jnp.where(boundary, f_pre, f_post)
+        fpop_neighbour = jnp.where(neighbour, f_pre, f_post)
+
+        # With fpop_neighbour isolated, now roll it back to be positioned at the boundary for subsequent operations
+        fpop_neighbour = self._roll(fpop_neighbour, self.normal)
+        fpop_extrapolated = sound_speed * fpop_neighbour + (1.0 - sound_speed) * fpop
+
+        # Use the iknown directions of f_postcollision that leave the domain during streaming to store the BC data
+        opp = self.velocity_set.opp_indices
+        known_mask = missing_mask[opp]
+        f_post = jnp.where(jnp.logical_and(boundary, known_mask), fpop_extrapolated[opp], f_post)
+        return f_post
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(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(
+            jnp.logical_and(missing_mask, boundary),
+            f_pre[self.velocity_set.opp_indices],
+            f_post,
+        )
+
+    def _construct_warp(self):
+        # Set local constants
+        sound_speed = self.compute_dtype(1.0 / wp.sqrt(3.0))
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _c = self.velocity_set.c
+        _q = self.velocity_set.q
+        _opp_indices = self.velocity_set.opp_indices
+
+        @wp.func
+        def get_normal_vectors(
+            missing_mask: Any,
+        ):
+            for l in range(_q):
+                if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
+                    return -wp.vec3i(_c[0, l], _c[1, l], _c[2, l])
+
+        # Construct the functionals for this BC
+        @wp.func
+        def functional(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # Post-streaming values are only modified at missing direction
+            _f = f_post
+            for l in range(self.velocity_set.q):
+                # If the mask is missing then take the opposite index
+                if missing_mask[l] == wp.uint8(1):
+                    _f[l] = f_pre[_opp_indices[l]]
+            return _f
+
+        @wp.func
+        def prepare_bc_auxilary_data(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # Preparing the formulation for this BC using the neighbour's populations stored in f_aux and
+            # f_pre (post-streaming values of the current voxel). We use directions that leave the domain
+            # for storing this prepared data.
+            _f = f_post
+            nv = get_normal_vectors(missing_mask)
+            for l in range(self.velocity_set.q):
+                if missing_mask[l] == wp.uint8(1):
+                    # f_0 is the post-collision values of the current time-step
+                    # Get pull index associated with the "neighbours" pull_index
+                    pull_index = type(index)()
+                    for d in range(self.velocity_set.d):
+                        pull_index[d] = index[d] - (_c[d, l] + nv[d])
+                    # The following is the post-streaming values of the neighbor cell
+                    f_aux = self.compute_dtype(f_0[l, pull_index[0], pull_index[1], pull_index[2]])
+                    _f[_opp_indices[l]] = (self.compute_dtype(1.0) - sound_speed) * f_pre[l] + sound_speed * f_aux
+            return _f
+
+        kernel = self._construct_kernel(functional)
+
+        return (functional, prepare_bc_auxilary_data), kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    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, bc_mask, missing_mask],
+            dim=f_pre.shape[1:],
+        )
+        return f_post
\ No newline at end of file
diff --git a/xlb/operator/boundary_condition/bc_fullway_bounce_back.py b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
new file mode 100644
index 0000000..19a3013
--- /dev/null
+++ b/xlb/operator/boundary_condition/bc_fullway_bounce_back.py
@@ -0,0 +1,89 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import jax.lax as lax
+from functools import partial
+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 import Operator
+from xlb.operator.boundary_condition.boundary_condition import (
+    BoundaryCondition,
+    ImplementationStep,
+)
+from xlb.operator.boundary_condition.boundary_condition_registry import (
+    boundary_condition_registry,
+)
+
+
+class FullwayBounceBackBC(BoundaryCondition):
+    """
+    Full Bounce-back boundary condition for a lattice Boltzmann method simulation.
+    """
+
+    def __init__(
+        self,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        super().__init__(
+            ImplementationStep.COLLISION,
+            velocity_set,
+            precision_policy,
+            compute_backend,
+            indices,
+            mesh_vertices,
+        )
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(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)
+
+    def _construct_warp(self):
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        _opp_indices = self.velocity_set.opp_indices
+        _q = wp.constant(self.velocity_set.q)
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+
+        # Construct the functional for this BC
+        @wp.func
+        def functional(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            fliped_f = _f_vec()
+            for l in range(_q):
+                fliped_f[l] = f_pre[_opp_indices[l]]
+            return fliped_f
+
+        kernel = self._construct_kernel(functional)
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    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, bc_mask, missing_mask],
+            dim=f_pre.shape[1:],
+        )
+        return f_post
\ No newline at end of file
diff --git a/xlb/operator/boundary_condition/bc_grads_approximation.py b/xlb/operator/boundary_condition/bc_grads_approximation.py
new file mode 100644
index 0000000..94ddba3
--- /dev/null
+++ b/xlb/operator/boundary_condition/bc_grads_approximation.py
@@ -0,0 +1,324 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import jax.lax as lax
+from functools import partial
+import warp as wp
+from typing import Any
+from collections import Counter
+import numpy as np
+
+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.macroscopic import Macroscopic
+from xlb.operator.macroscopic.zero_moment import ZeroMoment
+from xlb.operator.macroscopic.second_moment import SecondMoment as MomentumFlux
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.operator.boundary_condition.boundary_condition import (
+    ImplementationStep,
+    BoundaryCondition,
+)
+from xlb.operator.boundary_condition.boundary_condition_registry import (
+    boundary_condition_registry,
+)
+
+
+class GradsApproximationBC(BoundaryCondition):
+    """
+    Purpose: Using Grad's approximation to represent fpop based on macroscopic inputs used for outflow [1] and
+    Dirichlet BCs [2]
+    [1] S. Chikatamarla, S. Ansumali, and I. Karlin, "Grad's approximation for missing data in lattice Boltzmann
+        simulations", Europhys. Lett. 74, 215 (2006).
+    [2] Dorschner, B., Chikatamarla, S. S., Bösch, F., & Karlin, I. V. (2015). Grad's approximation for moving and
+        stationary walls in entropic lattice Boltzmann simulations. Journal of Computational Physics, 295, 340-354.
+
+    """
+
+    def __init__(
+        self,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        # TODO: the input velocity must be suitably stored elesewhere when mesh is moving.
+        self.u = (0, 0, 0)
+
+        # Call the parent constructor
+        super().__init__(
+            ImplementationStep.STREAMING,
+            velocity_set,
+            precision_policy,
+            compute_backend,
+            indices,
+            mesh_vertices,
+        )
+
+        # Instantiate the operator for computing macroscopic values
+        self.macroscopic = Macroscopic()
+        self.zero_moment = ZeroMoment()
+        self.equilibrium = QuadraticEquilibrium()
+        self.momentum_flux = MomentumFlux()
+
+        # This BC needs implicit distance to the mesh
+        self.needs_mesh_distance = True
+
+        # If this BC is defined using indices, it would need padding in order to find missing directions
+        # when imposed on a geometry that is in the domain interior
+        if self.mesh_vertices is None:
+            assert self.indices is not None
+            self.needs_padding = True
+
+        # Raise error if used for 2d examples:
+        if self.velocity_set.d == 2:
+            raise NotImplementedError("This BC is not implemented in 2D!")
+
+        # if indices is not None:
+        #     # this BC would be limited to stationary boundaries
+        #     # assert mesh_vertices is None
+        # if mesh_vertices is not None:
+        #     # this BC would be applicable for stationary and moving boundaries
+        #     assert indices is None
+        #     if mesh_velocity_function is not None:
+        #         # mesh is moving and/or deforming
+
+        assert self.compute_backend == ComputeBackend.WARP, "This BC is currently only implemented with the Warp backend!"
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask):
+        # TODO
+        raise NotImplementedError(f"Operation {self.__class__.__name} not implemented in JAX!")
+        return
+
+    def _construct_warp(self):
+        # Set local variables and constants
+        _c = self.velocity_set.c
+        _q = self.velocity_set.q
+        _d = self.velocity_set.d
+        _w = self.velocity_set.w
+        _qi = self.velocity_set.qi
+        _opp_indices = self.velocity_set.opp_indices
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _u_wall = _u_vec(self.u[0], self.u[1], self.u[2]) if _d == 3 else _u_vec(self.u[0], self.u[1])
+        # diagonal = wp.vec3i(0, 3, 5) if _d == 3 else wp.vec2i(0, 2)
+
+        @wp.func
+        def regularize_fpop(
+            missing_mask: Any,
+            rho: Any,
+            u: Any,
+            fpop: Any,
+        ):
+            """
+            Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop.
+            """
+            # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq}
+            feq = self.equilibrium.warp_functional(rho, u)
+            f_neq = fpop - feq
+            PiNeq = self.momentum_flux.warp_functional(f_neq)
+
+            # Compute double dot product Qi:Pi1 (where Pi1 = PiNeq)
+            nt = _d * (_d + 1) // 2
+            for l in range(_q):
+                QiPi1 = self.compute_dtype(0.0)
+                for t in range(nt):
+                    QiPi1 += _qi[l, t] * PiNeq[t]
+
+                # assign all populations based on eq 45 of Latt et al (2008)
+                # fneq ~ f^1
+                fpop1 = self.compute_dtype(4.5) * _w[l] * QiPi1
+                fpop[l] = feq[l] + fpop1
+            return fpop
+
+        @wp.func
+        def grads_approximate_fpop(
+            missing_mask: Any,
+            rho: Any,
+            u: Any,
+            f_post: Any,
+        ):
+            # Purpose: Using Grad's approximation to represent fpop based on macroscopic inputs used for outflow [1] and
+            # Dirichlet BCs [2]
+            # [1] S. Chikatax`marla, S. Ansumali, and I. Karlin, "Grad's approximation for missing data in lattice Boltzmann
+            #   simulations", Europhys. Lett. 74, 215 (2006).
+            # [2] Dorschner, B., Chikatamarla, S. S., Bösch, F., & Karlin, I. V. (2015). Grad's approximation for moving and
+            #    stationary walls in entropic lattice Boltzmann simulations. Journal of Computational Physics, 295, 340-354.
+
+            # Note: See also self.regularize_fpop function which is somewhat similar.
+
+            # Compute pressure tensor Pi using all f_post-streaming values
+            Pi = self.momentum_flux.warp_functional(f_post)
+
+            # Compute double dot product Qi:Pi1 (where Pi1 = PiNeq)
+            nt = _d * (_d + 1) // 2
+            for l in range(_q):
+                # if missing_mask[l] == wp.uint8(1):
+                QiPi = self.compute_dtype(0.0)
+                for t in range(nt):
+                    if t == 0 or t == 3 or t == 5:
+                        QiPi += _qi[l, t] * (Pi[t] - rho / self.compute_dtype(3.0))
+                    else:
+                        QiPi += _qi[l, t] * Pi[t]
+
+                # Compute c.u
+                cu = self.compute_dtype(0.0)
+                for d in range(self.velocity_set.d):
+                    if _c[d, l] == 1:
+                        cu += u[d]
+                    elif _c[d, l] == -1:
+                        cu -= u[d]
+                cu *= self.compute_dtype(3.0)
+
+                # change f_post using the Grad's approximation
+                f_post[l] = rho * _w[l] * (self.compute_dtype(1.0) + cu) + _w[l] * self.compute_dtype(4.5) * QiPi
+
+            return f_post
+
+        # Construct the functionals for this BC
+        @wp.func
+        def functional_method1(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # NOTE: this BC has been reformulated to become entirely local and so has differences compared to the original paper.
+            #       Here we use the current time-step populations (f_pre = f_post_collision and f_post = f_post_streaming).
+            one = self.compute_dtype(1.0)
+            for l in range(_q):
+                # If the mask is missing then take the opposite index
+                if missing_mask[l] == wp.uint8(1):
+                    # The implicit distance to the boundary or "weights" have been stored in known directions of f_1
+                    # weight = f_1[_opp_indices[l], index[0], index[1], index[2]]
+                    weight = self.compute_dtype(0.5)
+
+                    # Use differentiable interpolated BB to find f_missing:
+                    f_post[l] = ((one - weight) * f_post[_opp_indices[l]] + weight * (f_pre[l] + f_pre[_opp_indices[l]])) / (one + weight)
+
+                    # # Add contribution due to moving_wall to f_missing as is usual in regular Bouzidi BC
+                    # cu = self.compute_dtype(0.0)
+                    # for d in range(_d):
+                    #     if _c[d, l] == 1:
+                    #         cu += _u_wall[d]
+                    #     elif _c[d, l] == -1:
+                    #         cu -= _u_wall[d]
+                    # cu *= self.compute_dtype(-6.0) * _w[l]
+                    # f_post[l] += cu
+
+            # Compute density, velocity using all f_post-streaming values
+            rho, u = self.macroscopic.warp_functional(f_post)
+
+            # Compute Grad's appriximation using full equation as in Eq (10) of Dorschner et al.
+            f_post = regularize_fpop(missing_mask, rho, u, f_post)
+            # f_post = grads_approximate_fpop(missing_mask, rho, u, f_post)
+            return f_post
+
+        # Construct the functionals for this BC
+        @wp.func
+        def functional_method2(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # NOTE: this BC has been reformulated to become entirely local and so has differences compared to the original paper.
+            #       Here we use the current time-step populations (f_pre = f_post_collision and f_post = f_post_streaming).
+            # NOTE: f_aux should contain populations at "x_f" (see their fig 1) in the missign direction of the BC which amounts
+            #       to post-collision values being pulled from appropriate cells like ExtrapolationBC
+            #
+            # here I need to compute all terms in Eq (10)
+            # Strategy:
+            # 1) "weights" should have been stored somewhere to be used here.
+            # 2) Given "weights", "u_w" (input to the BC) and "u_f" (computed from f_aux), compute "u_target" as per Eq (14)
+            #    NOTE: in the original paper "u_target" is associated with the previous time step not current time.
+            # 3) Given "weights" use differentiable interpolated BB to find f_missing as I had before:
+            # fmissing = ((1. - weights) * f_poststreaming_iknown + weights * (f_postcollision_imissing + f_postcollision_iknown)) / (1.0 + weights)
+            # 4) Add contribution due to u_w to f_missing as is usual in regular Bouzidi BC (ie. -6.0 * self.lattice.w * jnp.dot(self.vel, c)
+            # 5) Compute rho_target = \sum(f_ibb) based on these values
+            # 6) Compute feq using feq = self.equilibrium(rho_target, u_target)
+            # 7) Compute Pi_neq and Pi_eq using all f_post-streaming values as per:
+            #       Pi_neq = self.momentum_flux(fneq) and Pi_eq = self.momentum_flux(feq)
+            # 8) Compute Grad's appriximation using full equation as in Eq (10)
+            #    NOTE: this is very similar to the regularization procedure.
+
+            _f_nbr = _f_vec()
+            u_target = _u_vec(0.0, 0.0, 0.0) if _d == 3 else _u_vec(0.0, 0.0)
+            num_missing = 0
+            one = self.compute_dtype(1.0)
+            for l in range(_q):
+                # If the mask is missing then take the opposite index
+                if missing_mask[l] == wp.uint8(1):
+                    # Find the neighbour and its velocity value
+                    for ll in range(_q):
+                        # f_0 is the post-collision values of the current time-step
+                        # Get index associated with the fluid neighbours
+                        fluid_nbr_index = type(index)()
+                        for d in range(_d):
+                            fluid_nbr_index[d] = index[d] + _c[d, l]
+                        # The following is the post-collision values of the fluid neighbor cell
+                        _f_nbr[ll] = self.compute_dtype(f_0[ll, fluid_nbr_index[0], fluid_nbr_index[1], fluid_nbr_index[2]])
+
+                    # Compute the velocity vector at the fluid neighbouring cells
+                    _, u_f = self.macroscopic.warp_functional(_f_nbr)
+
+                    # Record the number of missing directions
+                    num_missing += 1
+
+                    # The implicit distance to the boundary or "weights" have been stored in known directions of f_1
+                    weight = f_1[_opp_indices[l], index[0], index[1], index[2]]
+
+                    # Given "weights", "u_w" (input to the BC) and "u_f" (computed from f_aux), compute "u_target" as per Eq (14)
+                    for d in range(_d):
+                        u_target[d] += (weight * u_f[d] + _u_wall[d]) / (one + weight)
+
+                    # Use differentiable interpolated BB to find f_missing:
+                    f_post[l] = ((one - weight) * f_post[_opp_indices[l]] + weight * (f_pre[l] + f_pre[_opp_indices[l]])) / (one + weight)
+
+                    # Add contribution due to moving_wall to f_missing as is usual in regular Bouzidi BC
+                    cu = self.compute_dtype(0.0)
+                    for d in range(_d):
+                        if _c[d, l] == 1:
+                            cu += _u_wall[d]
+                        elif _c[d, l] == -1:
+                            cu -= _u_wall[d]
+                    cu *= self.compute_dtype(-6.0) * _w[l]
+                    f_post[l] += cu
+
+            # Compute rho_target = \sum(f_ibb) based on these values
+            rho_target = self.zero_moment.warp_functional(f_post)
+            for d in range(_d):
+                u_target[d] /= num_missing
+
+            # Compute Grad's appriximation using full equation as in Eq (10) of Dorschner et al.
+            f_post = grads_approximate_fpop(missing_mask, rho_target, u_target, f_post)
+            return f_post
+
+        functional = functional_method1
+
+        kernel = self._construct_kernel(functional)
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    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, bc_mask, missing_mask],
+            dim=f_pre.shape[1:],
+        )
+        return f_post
\ No newline at end of file
diff --git a/xlb/operator/boundary_condition/bc_halfway_bounce_back.py b/xlb/operator/boundary_condition/bc_halfway_bounce_back.py
new file mode 100644
index 0000000..bf04af0
--- /dev/null
+++ b/xlb/operator/boundary_condition/bc_halfway_bounce_back.py
@@ -0,0 +1,102 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import jax.lax as lax
+from functools import partial
+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.boundary_condition.boundary_condition import (
+    ImplementationStep,
+    BoundaryCondition,
+)
+from xlb.operator.boundary_condition.boundary_condition_registry import (
+    boundary_condition_registry,
+)
+
+
+class HalfwayBounceBackBC(BoundaryCondition):
+    """
+    Halfway Bounce-back boundary condition for a lattice Boltzmann method simulation.
+
+    TODO: Implement moving boundary conditions for this
+    """
+
+    def __init__(
+        self,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        # Call the parent constructor
+        super().__init__(
+            ImplementationStep.STREAMING,
+            velocity_set,
+            precision_policy,
+            compute_backend,
+            indices,
+            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 jax_implementation(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(
+            jnp.logical_and(missing_mask, boundary),
+            f_pre[self.velocity_set.opp_indices],
+            f_post,
+        )
+
+    def _construct_warp(self):
+        # Set local constants
+        _opp_indices = self.velocity_set.opp_indices
+
+        # Construct the functional for this BC
+        @wp.func
+        def functional(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # Post-streaming values are only modified at missing direction
+            _f = f_post
+            for l in range(self.velocity_set.q):
+                # If the mask is missing then take the opposite index
+                if missing_mask[l] == wp.uint8(1):
+                    # Get the pre-streaming distribution function in oppisite direction
+                    _f[l] = f_pre[_opp_indices[l]]
+
+            return _f
+
+        kernel = self._construct_kernel(functional)
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    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, 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
new file mode 100644
index 0000000..af4c783
--- /dev/null
+++ b/xlb/operator/boundary_condition/bc_regularized.py
@@ -0,0 +1,281 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import jax.lax as lax
+from functools import partial
+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.boundary_condition.bc_zouhe import ZouHeBC
+from xlb.operator.boundary_condition.boundary_condition import ImplementationStep
+from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry
+from xlb.operator.macroscopic.second_moment import SecondMoment as MomentumFlux
+
+
+class RegularizedBC(ZouHeBC):
+    """
+    Regularized boundary condition for a lattice Boltzmann method simulation.
+
+    This class implements the regularized boundary condition, which is a non-equilibrium bounce-back boundary condition
+    with additional regularization. It can be used to set inflow and outflow boundary conditions with prescribed pressure
+    or velocity.
+
+    Attributes
+    ----------
+    name : str
+        The name of the boundary condition. For this class, it is "Regularized".
+    Qi : numpy.ndarray
+        The Qi tensor, which is used in the regularization of the distribution functions.
+
+    References
+    ----------
+    Latt, J. (2007). Hydrodynamic limit of lattice Boltzmann equations. PhD thesis, University of Geneva.
+    Latt, J., Chopard, B., Malaspinas, O., Deville, M., & Michler, A. (2008). Straight velocity boundaries in the
+    lattice Boltzmann method. Physical Review E, 77(5), 056703. doi:10.1103/PhysRevE.77.056703
+    """
+
+    def __init__(
+        self,
+        bc_type,
+        prescribed_value,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        # Call the parent constructor
+        super().__init__(
+            bc_type,
+            prescribed_value,
+            velocity_set,
+            precision_policy,
+            compute_backend,
+            indices,
+            mesh_vertices,
+        )
+        # Overwrite the boundary condition registry id with the bc_type in the name
+        self.momentum_flux = MomentumFlux()
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def regularize_fpop(self, fpop, feq):
+        """
+        Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop.
+
+        Parameters
+        ----------
+        fpop : jax.numpy.ndarray
+            The distribution functions.
+        feq : jax.numpy.ndarray
+            The equilibrium distribution functions.
+
+        Returns
+        -------
+        jax.numpy.ndarray
+            The regularized distribution functions.
+        """
+        # Qi = cc - cs^2*I
+        dim = self.velocity_set.d
+        weights = self.velocity_set.w[(slice(None),) + (None,) * dim]
+        Qi = jnp.array(self.velocity_set.qi, dtype=self.compute_dtype)
+
+        # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq}
+        f_neq = fpop - feq
+        PiNeq = self.momentum_flux(f_neq)
+        # PiNeq = self.momentum_flux(fpop) - self.momentum_flux(feq)
+
+        # Compute double dot product Qi:Pi1
+        # QiPi1 = np.zeros_like(fpop)
+        # Pi1 = PiNeq
+        QiPi1 = jnp.tensordot(Qi, PiNeq, axes=(1, 0))
+
+        # assign all populations based on eq 45 of Latt et al (2008)
+        # fneq ~ f^1
+        fpop1 = 9.0 / 2.0 * weights * QiPi1
+        fpop_regularized = feq + fpop1
+        return fpop_regularized
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask):
+        # creat a mask to slice boundary cells
+        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)))
+
+        # compute the equilibrium based on prescribed values and the type of BC
+        feq = self.calculate_equilibrium(f_post, missing_mask)
+
+        # set the unknown f populations based on the non-equilibrium bounce-back method
+        f_post_bd = self.bounceback_nonequilibrium(f_post, feq, missing_mask)
+
+        # Regularize the boundary fpop
+        f_post_bd = self.regularize_fpop(f_post_bd, feq)
+
+        # apply bc
+        f_post = jnp.where(boundary, f_post_bd, f_post)
+        return f_post
+
+    def _construct_warp(self):
+        # assign placeholders for both u and rho based on prescribed_value
+        _d = self.velocity_set.d
+        _q = self.velocity_set.q
+        u = self.prescribed_value if self.bc_type == "velocity" else (0,) * _d
+        rho = self.prescribed_value if self.bc_type == "pressure" else 0.0
+
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        # _u_vec = wp.vec(_d, dtype=self.compute_dtype)
+        # compute Qi tensor and store it in self
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _rho = self.compute_dtype(rho)
+        _u = _u_vec(u[0], u[1], u[2]) if _d == 3 else _u_vec(u[0], u[1])
+        _opp_indices = self.velocity_set.opp_indices
+        _w = self.velocity_set.w
+        _c = self.velocity_set.c
+        _c_float = self.velocity_set.c_float
+        _qi = self.velocity_set.qi
+        # TODO: related to _c_float: this is way less than ideal. we should not be making new types
+
+        @wp.func
+        def _get_fsum(
+            fpop: Any,
+            missing_mask: Any,
+        ):
+            fsum_known = self.compute_dtype(0.0)
+            fsum_middle = self.compute_dtype(0.0)
+            for l in range(_q):
+                if missing_mask[_opp_indices[l]] == wp.uint8(1):
+                    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
+
+        @wp.func
+        def get_normal_vectors(
+            missing_mask: Any,
+        ):
+            for l in range(_q):
+                if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
+                    return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l])
+
+        @wp.func
+        def bounceback_nonequilibrium(
+            fpop: Any,
+            feq: Any,
+            missing_mask: Any,
+        ):
+            for l in range(_q):
+                if missing_mask[l] == wp.uint8(1):
+                    fpop[l] = fpop[_opp_indices[l]] + feq[l] - feq[_opp_indices[l]]
+            return fpop
+
+        @wp.func
+        def regularize_fpop(
+            fpop: Any,
+            feq: Any,
+        ):
+            """
+            Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop.
+            """
+            # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq}
+            f_neq = fpop - feq
+            PiNeq = self.momentum_flux.warp_functional(f_neq)
+
+            # Compute double dot product Qi:Pi1 (where Pi1 = PiNeq)
+            nt = _d * (_d + 1) // 2
+            for l in range(_q):
+                QiPi1 = self.compute_dtype(0.0)
+                for t in range(nt):
+                    QiPi1 += _qi[l, t] * PiNeq[t]
+
+                # assign all populations based on eq 45 of Latt et al (2008)
+                # fneq ~ f^1
+                fpop1 = self.compute_dtype(4.5) * _w[l] * QiPi1
+                fpop[l] = feq[l] + fpop1
+            return fpop
+
+        @wp.func
+        def functional_velocity(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # Post-streaming values are only modified at missing direction
+            _f = f_post
+
+            # Find normal vector
+            normals = get_normal_vectors(missing_mask)
+
+            # calculate rho
+            fsum = _get_fsum(_f, missing_mask)
+            unormal = self.compute_dtype(0.0)
+            for d in range(_d):
+                unormal += _u[d] * normals[d]
+            _rho = fsum / (self.compute_dtype(1.0) + unormal)
+
+            # impose non-equilibrium bounceback
+            feq = self.equilibrium_operator.warp_functional(_rho, _u)
+            _f = bounceback_nonequilibrium(_f, feq, missing_mask)
+
+            # Regularize the boundary fpop
+            _f = regularize_fpop(_f, feq)
+            return _f
+
+        @wp.func
+        def functional_pressure(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # Post-streaming values are only modified at missing direction
+            _f = f_post
+
+            # Find normal vector
+            normals = get_normal_vectors(missing_mask)
+
+            # calculate velocity
+            fsum = _get_fsum(_f, missing_mask)
+            unormal = -self.compute_dtype(1.0) + fsum / _rho
+            _u = unormal * normals
+
+            # impose non-equilibrium bounceback
+            feq = self.equilibrium_operator.warp_functional(_rho, _u)
+            _f = bounceback_nonequilibrium(_f, feq, missing_mask)
+
+            # Regularize the boundary fpop
+            _f = regularize_fpop(_f, feq)
+            return _f
+
+        if self.bc_type == "velocity":
+            functional = functional_velocity
+        elif self.bc_type == "pressure":
+            functional = functional_pressure
+        kernel = self._construct_kernel(functional)
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    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, 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
new file mode 100644
index 0000000..a92d909
--- /dev/null
+++ b/xlb/operator/boundary_condition/bc_zouhe.py
@@ -0,0 +1,298 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import jax.lax as lax
+from functools import partial
+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.boundary_condition.boundary_condition import (
+    ImplementationStep,
+    BoundaryCondition,
+)
+from xlb.operator.boundary_condition.boundary_condition_registry import (
+    boundary_condition_registry,
+)
+from xlb.operator.equilibrium import QuadraticEquilibrium
+
+
+class ZouHeBC(BoundaryCondition):
+    """
+    Zou-He boundary condition for a lattice Boltzmann method simulation.
+
+    This method applies the Zou-He boundary condition by first computing the equilibrium distribution functions based
+    on the prescribed values and the type of boundary condition, and then setting the unknown distribution functions
+    based on the non-equilibrium bounce-back method.
+    Tangential velocity is not ensured to be zero by adding transverse contributions based on
+    Hecth & Harting (2010) (doi:10.1088/1742-5468/2010/01/P01018) as it caused numerical instabilities at higher
+    Reynolds numbers. One needs to use "Regularized" BC at higher Reynolds.
+    """
+
+    def __init__(
+        self,
+        bc_type,
+        prescribed_value,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        # Important Note: it is critical to add id inside __init__ for this BC because different instantiations of this BC
+        # may have different types (velocity or pressure).
+        assert bc_type in ["velocity", "pressure"], f"type = {bc_type} not supported! Use 'pressure' or 'velocity'."
+        self.bc_type = bc_type
+        self.equilibrium_operator = QuadraticEquilibrium()
+        self.prescribed_value = prescribed_value
+
+        # Call the parent constructor
+        super().__init__(
+            ImplementationStep.STREAMING,
+            velocity_set,
+            precision_policy,
+            compute_backend,
+            indices,
+            mesh_vertices,
+        )
+
+        # Set the prescribed value for pressure or velocity
+        dim = self.velocity_set.d
+        if self.compute_backend == ComputeBackend.JAX:
+            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]
+        middle_mask = ~(missing_mask | known_mask)
+        return known_mask, middle_mask
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def _get_normal_vec(self, missing_mask):
+        main_c = self.velocity_set.c[:, self.velocity_set.main_indices]
+        m = missing_mask[self.velocity_set.main_indices]
+        normals = -jnp.tensordot(main_c, m, axes=(-1, 0))
+        return normals
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def get_rho(self, fpop, missing_mask):
+        if self.bc_type == "velocity":
+            vel = self.prescribed_value
+            rho = self.calculate_rho(fpop, vel, missing_mask)
+        elif self.bc_type == "pressure":
+            rho = self.prescribed_value
+        else:
+            raise ValueError(f"type = {self.bc_type} not supported! Use 'pressure' or 'velocity'.")
+        return rho
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def get_vel(self, fpop, missing_mask):
+        if self.bc_type == "velocity":
+            vel = self.prescribed_value
+        elif self.bc_type == "pressure":
+            rho = self.prescribed_value
+            vel = self.calculate_vel(fpop, rho, missing_mask)
+        else:
+            raise ValueError(f"type = {self.bc_type} not supported! Use 'pressure' or 'velocity'.")
+        return vel
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def calculate_vel(self, fpop, rho, missing_mask):
+        """
+        Calculate velocity based on the prescribed pressure/density (Zou/He BC)
+        """
+
+        normals = self._get_normal_vec(missing_mask)
+        known_mask, middle_mask = self._get_known_middle_mask(missing_mask)
+        fsum = jnp.sum(fpop * middle_mask, axis=0, keepdims=True) + 2.0 * jnp.sum(fpop * known_mask, axis=0, keepdims=True)
+        unormal = -1.0 + fsum / rho
+
+        # Return the above unormal as a normal vector which sets the tangential velocities to zero
+        vel = unormal * normals
+        return vel
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def calculate_rho(self, fpop, vel, missing_mask):
+        """
+        Calculate density based on the prescribed velocity (Zou/He BC)
+        """
+        normals = self._get_normal_vec(missing_mask)
+        known_mask, middle_mask = self._get_known_middle_mask(missing_mask)
+        unormal = jnp.sum(normals * vel, keepdims=True, axis=0)
+        fsum = jnp.sum(fpop * middle_mask, axis=0, keepdims=True) + 2.0 * jnp.sum(fpop * known_mask, axis=0, keepdims=True)
+        rho = fsum / (1.0 + unormal)
+        return rho
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def calculate_equilibrium(self, fpop, missing_mask):
+        """
+        This is the ZouHe method of calculating the missing macroscopic variables at the boundary.
+        """
+        rho = self.get_rho(fpop, missing_mask)
+        vel = self.get_vel(fpop, missing_mask)
+
+        # compute feq at the boundary
+        feq = self.equilibrium_operator(rho, vel)
+        return feq
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def bounceback_nonequilibrium(self, fpop, feq, missing_mask):
+        """
+        Calculate unknown populations using bounce-back of non-equilibrium populations
+        a la original Zou & He formulation
+        """
+        opp = self.velocity_set.opp_indices
+        fknown = fpop[opp] + feq - feq[opp]
+        fpop = jnp.where(missing_mask, fknown, fpop)
+        return fpop
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(self, f_pre, f_post, bc_mask, missing_mask):
+        # creat a mask to slice boundary cells
+        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)))
+
+        # compute the equilibrium based on prescribed values and the type of BC
+        feq = self.calculate_equilibrium(f_post, missing_mask)
+
+        # set the unknown f populations based on the non-equilibrium bounce-back method
+        f_post_bd = self.bounceback_nonequilibrium(f_post, feq, missing_mask)
+        f_post = jnp.where(boundary, f_post_bd, f_post)
+        return f_post
+
+    def _construct_warp(self):
+        # assign placeholders for both u and rho based on prescribed_value
+        _d = self.velocity_set.d
+        _q = self.velocity_set.q
+        u = self.prescribed_value if self.bc_type == "velocity" else (0,) * _d
+        rho = self.prescribed_value if self.bc_type == "pressure" else 0.0
+
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        # _u_vec = wp.vec(_d, dtype=self.compute_dtype)
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _rho = self.compute_dtype(rho)
+        _u = _u_vec(u[0], u[1], u[2]) if _d == 3 else _u_vec(u[0], u[1])
+        _opp_indices = self.velocity_set.opp_indices
+        _c = self.velocity_set.c
+        _c_float = self.velocity_set.c_float
+        # TODO: this is way less than ideal. we should not be making new types
+
+        @wp.func
+        def _get_fsum(
+            fpop: Any,
+            missing_mask: Any,
+        ):
+            fsum_known = self.compute_dtype(0.0)
+            fsum_middle = self.compute_dtype(0.0)
+            for l in range(_q):
+                if missing_mask[_opp_indices[l]] == wp.uint8(1):
+                    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
+
+        @wp.func
+        def get_normal_vectors(
+            missing_mask: Any,
+        ):
+            for l in range(_q):
+                if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) + wp.abs(_c[2, l]) == 1:
+                    return -_u_vec(_c_float[0, l], _c_float[1, l], _c_float[2, l])
+
+        @wp.func
+        def bounceback_nonequilibrium(
+            fpop: Any,
+            feq: Any,
+            missing_mask: Any,
+        ):
+            for l in range(_q):
+                if missing_mask[l] == wp.uint8(1):
+                    fpop[l] = fpop[_opp_indices[l]] + feq[l] - feq[_opp_indices[l]]
+            return fpop
+
+        @wp.func
+        def functional_velocity(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # Post-streaming values are only modified at missing direction
+            _f = f_post
+
+            # Find normal vector
+            normals = get_normal_vectors(missing_mask)
+
+            # calculate rho
+            fsum = _get_fsum(_f, missing_mask)
+            unormal = self.compute_dtype(0.0)
+            for d in range(_d):
+                unormal += _u[d] * normals[d]
+            _rho = fsum / (self.compute_dtype(1.0) + unormal)
+
+            # impose non-equilibrium bounceback
+            feq = self.equilibrium_operator.warp_functional(_rho, _u)
+            _f = bounceback_nonequilibrium(_f, feq, missing_mask)
+            return _f
+
+        @wp.func
+        def functional_pressure(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            # Post-streaming values are only modified at missing direction
+            _f = f_post
+
+            # Find normal vector
+            normals = get_normal_vectors(missing_mask)
+
+            # calculate velocity
+            fsum = _get_fsum(_f, missing_mask)
+            unormal = -self.compute_dtype(1.0) + fsum / _rho
+            _u = unormal * normals
+
+            # impose non-equilibrium bounceback
+            feq = self.equilibrium_operator.warp_functional(_rho, _u)
+            _f = bounceback_nonequilibrium(_f, feq, missing_mask)
+            return _f
+
+        if self.bc_type == "velocity":
+            functional = functional_velocity
+        elif self.bc_type == "pressure":
+            functional = functional_pressure
+        elif self.bc_type == "velocity":
+            functional = functional_pressure
+
+        kernel = self._construct_kernel(functional)
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    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, 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
new file mode 100644
index 0000000..bf1eef2
--- /dev/null
+++ b/xlb/operator/boundary_condition/boundary_condition.py
@@ -0,0 +1,148 @@
+"""
+Base class for boundary conditions in a LBM simulation.
+"""
+
+from enum import Enum, auto
+import warp as wp
+from typing import Any
+from jax import jit
+from functools import partial
+
+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 import DefaultConfig
+from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry
+
+
+# Enum for implementation step
+class ImplementationStep(Enum):
+    COLLISION = auto()
+    STREAMING = auto()
+
+
+class BoundaryCondition(Operator):
+    """
+    Base class for boundary conditions in a LBM simulation.
+    """
+
+    def __init__(
+        self,
+        implementation_step: ImplementationStep,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+        indices=None,
+        mesh_vertices=None,
+    ):
+        self.id = boundary_condition_registry.register_boundary_condition(self.__class__.__name__ + "_" + str(hash(self)))
+        velocity_set = velocity_set or DefaultConfig.velocity_set
+        precision_policy = precision_policy or DefaultConfig.default_precision_policy
+        compute_backend = compute_backend or DefaultConfig.default_backend
+
+        super().__init__(velocity_set, precision_policy, compute_backend)
+
+        # Set the BC indices
+        self.indices = indices
+        self.mesh_vertices = mesh_vertices
+
+        # 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
+
+        # A flag for BCs that need implicit boundary distance between the grid and a mesh (to be set to True if applicable inside each BC)
+        self.needs_mesh_distance = 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)
+            _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)  # TODO fix vec bool
+
+        @wp.func
+        def prepare_bc_auxilary_data(
+            index: Any,
+            timestep: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+        ):
+            return f_post
+
+        @wp.func
+        def _get_thread_data(
+            f_pre: wp.array4d(dtype=Any),
+            f_post: wp.array4d(dtype=Any),
+            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_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
+                _f_pre[l] = self.compute_dtype(f_pre[l, index[0], index[1], index[2]])
+                _f_post[l] = self.compute_dtype(f_post[l, index[0], index[1], index[2]])
+
+                # TODO fix vec bool
+                if missing_mask[l, index[0], index[1], index[2]]:
+                    _missing_mask[l] = wp.uint8(1)
+                else:
+                    _missing_mask[l] = wp.uint8(0)
+            return _f_pre, _f_post, _boundary_id, _missing_mask
+
+        # Construct some helper warp functions for getting tid data
+        if self.compute_backend == ComputeBackend.WARP:
+            self._get_thread_data = _get_thread_data
+            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, bc_mask, missing_mask):
+        """
+        A placeholder function for prepare the auxilary distribution functions for the boundary condition.
+        currently being called after collision only.
+        """
+        return f_post
+
+    def _construct_kernel(self, functional):
+        """
+        Constructs the warp kernel for the boundary condition.
+        The functional is specific to each boundary condition and should be passed as an argument.
+        """
+        _id = wp.uint8(self.id)
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            f_pre: wp.array4d(dtype=Any),
+            f_post: wp.array4d(dtype=Any),
+            bc_mask: wp.array4d(dtype=wp.uint8),
+            missing_mask: wp.array4d(dtype=wp.bool),
+        ):
+            # Get the global index
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            # read tid data
+            _f_pre, _f_post, _boundary_id, _missing_mask = self._get_thread_data(f_pre, f_post, bc_mask, missing_mask, index)
+
+            # Apply the boundary condition
+            if _boundary_id == _id:
+                timestep = 0
+                _f = functional(index, timestep, _missing_mask, f_pre, f_post, _f_pre, _f_post)
+            else:
+                _f = _f_post
+
+            # Write the result
+            for l in range(self.velocity_set.q):
+                f_post[l, index[0], index[1], index[2]] = self.store_dtype(_f[l])
+
+        return kernel
diff --git a/xlb/operator/boundary_condition/boundary_condition_registry.py b/xlb/operator/boundary_condition/boundary_condition_registry.py
new file mode 100644
index 0000000..6238fc5
--- /dev/null
+++ b/xlb/operator/boundary_condition/boundary_condition_registry.py
@@ -0,0 +1,30 @@
+"""
+Registry for boundary conditions in a LBM simulation.
+"""
+
+
+class BoundaryConditionRegistry:
+    """
+    Registry for boundary conditions in a LBM simulation.
+    """
+
+    def __init__(
+        self,
+    ):
+        self.id_to_bc = {}  # Maps id number to boundary condition
+        self.bc_to_id = {}  # Maps boundary condition to id number
+        self.next_id = 1  # 0 is reserved for no boundary condition
+
+    def register_boundary_condition(self, boundary_condition):
+        """
+        Register a boundary condition.
+        """
+        _id = self.next_id
+        self.next_id += 1
+        self.id_to_bc[_id] = boundary_condition
+        self.bc_to_id[boundary_condition] = _id
+        print(f"registered bc {boundary_condition} with id {_id}")
+        return _id
+
+
+boundary_condition_registry = BoundaryConditionRegistry()
diff --git a/xlb/operator/boundary_masker/__init__.py b/xlb/operator/boundary_masker/__init__.py
new file mode 100644
index 0000000..d03636a
--- /dev/null
+++ b/xlb/operator/boundary_masker/__init__.py
@@ -0,0 +1,3 @@
+from xlb.operator.boundary_masker.indices_boundary_masker import IndicesBoundaryMasker
+from xlb.operator.boundary_masker.mesh_boundary_masker import MeshBoundaryMasker
+from xlb.operator.boundary_masker.mesh_distance_boundary_masker import MeshDistanceBoundaryMasker
diff --git a/xlb/operator/boundary_masker/indices_boundary_masker.py b/xlb/operator/boundary_masker/indices_boundary_masker.py
new file mode 100644
index 0000000..0c1f7e1
--- /dev/null
+++ b/xlb/operator/boundary_masker/indices_boundary_masker.py
@@ -0,0 +1,220 @@
+import numpy as np
+import warp as wp
+import jax
+import jax.numpy as jnp
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.operator import Operator
+from xlb.operator.stream.stream import Stream
+from xlb.grid import grid_factory
+from xlb.precision_policy import Precision
+
+
+class IndicesBoundaryMasker(Operator):
+    """
+    Operator for creating a boundary mask
+    """
+
+    def __init__(
+        self,
+        velocity_set=None,
+        precision_policy=None,
+        compute_backend=None,
+    ):
+        # Make stream operator
+        self.stream = Stream(velocity_set, precision_policy, compute_backend)
+
+        # Call super
+        super().__init__(velocity_set, precision_policy, compute_backend)
+
+    def are_indices_in_interior(self, indices, shape):
+        """
+        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: Array of indices, where each column 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: Array of boolean flags where each flag indicates whether the corresponding index is inside the bounds.
+        """
+        shape_array = np.array(shape)
+        return np.all((indices > 0) & (indices < shape_array[:, np.newaxis] - 1), axis=0)
+
+    @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, 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
+        nDevices = jax.device_count()
+        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
+
+        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
+            bc_indices = np.array(bc.indices)
+            local_indices = bc_indices - np.array(start_index)[:, np.newaxis]
+            padded_indices = local_indices + np.array(shift_tup)[:, np.newaxis]
+            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.
+                # This flag is needed e.g. if the no-slip geometry is anywhere but at the boundaries of the computational domain.
+                if dim == 2:
+                    grid_mask = grid_mask.at[:, padded_indices[0], padded_indices[1]].set(True)
+                if dim == 3:
+                    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 = 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)
+
+        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]
+            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
+        _c = self.velocity_set.c
+        _q = wp.constant(self.velocity_set.q)
+
+        @wp.func
+        def check_index_bounds(index: wp.vec3i, shape: wp.vec3i):
+            is_in_bounds = index[0] >= 0 and index[0] < shape[0] and index[1] >= 0 and index[1] < shape[1] and index[2] >= 0 and index[2] < shape[2]
+            return is_in_bounds
+
+        # Construct the warp 3D kernel
+        @wp.kernel
+        def kernel(
+            indices: wp.array2d(dtype=wp.int32),
+            id_number: wp.array1d(dtype=wp.uint8),
+            is_interior: wp.array1d(dtype=wp.bool),
+            bc_mask: wp.array4d(dtype=wp.uint8),
+            missing_mask: wp.array4d(dtype=wp.bool),
+            start_index: wp.vec3i,
+        ):
+            # Get the index of indices
+            ii = wp.tid()
+
+            # Get local indices
+            index = wp.vec3i()
+            index[0] = indices[0, ii] - start_index[0]
+            index[1] = indices[1, ii] - start_index[1]
+            index[2] = indices[2, ii] - start_index[2]
+
+            # Check if index is in bounds
+            shape = wp.vec3i(missing_mask.shape[1], missing_mask.shape[2], missing_mask.shape[3])
+            if check_index_bounds(index, shape):
+                # Stream indices
+                for l in range(_q):
+                    # Get the index of the streaming direction
+                    pull_index = wp.vec3i()
+                    push_index = wp.vec3i()
+                    for d in range(self.velocity_set.d):
+                        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 not check_index_bounds(pull_index, shape):
+                        # Set the missing mask
+                        missing_mask[l, index[0], index[1], index[2]] = True
+
+                    # handling geometries in the interior of the computational domain
+                    elif check_index_bounds(pull_index, shape) and is_interior[ii]:
+                        # Set the missing mask
+                        missing_mask[l, push_index[0], push_index[1], push_index[2]] = True
+                        bc_mask[0, push_index[0], push_index[1], push_index[2]] = id_number[ii]
+
+        return None, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, bclist, bc_mask, missing_mask, start_index=None):
+        # Pre-allocate arrays with maximum possible size
+        max_size = sum(len(bc.indices[0]) if isinstance(bc.indices, list) else bc.indices.shape[1] for bc in bclist if bc.indices is not None)
+        indices = np.zeros((3, max_size), dtype=np.int32)
+        id_numbers = np.zeros(max_size, dtype=np.uint8)
+        is_interior = np.zeros(max_size, dtype=bool)
+
+        current_index = 0
+        for bc in bclist:
+            assert bc.indices is not None, f'Please specify indices associated with the {bc.__class__.__name__} BC using keyword "indices"!'
+            assert bc.mesh_vertices is None, f"Please use MeshBoundaryMasker operator if {bc.__class__.__name__} is imposed on a mesh (e.g. STL)!"
+
+            bc_indices = np.asarray(bc.indices)
+            num_indices = bc_indices.shape[1]
+
+            # Ensure indices are 3D
+            if bc_indices.shape[0] == 2:
+                bc_indices = np.vstack([bc_indices, np.zeros(num_indices, dtype=int)])
+
+            # Add indices to the pre-allocated array
+            indices[:, current_index : current_index + num_indices] = bc_indices
+
+            # Set id numbers
+            id_numbers[current_index : current_index + num_indices] = bc.id
+
+            # Set is_interior flags
+            if bc.needs_padding:
+                is_interior[current_index : current_index + num_indices] = self.are_indices_in_interior(bc_indices, bc_mask[0].shape)
+            else:
+                is_interior[current_index : current_index + num_indices] = False
+
+            current_index += num_indices
+
+            # Remove indices from BC objects
+            bc.__dict__.pop("indices", None)
+
+        # Trim arrays to actual size
+        indices = indices[:, :current_index]
+        id_numbers = id_numbers[:current_index]
+        is_interior = is_interior[:current_index]
+
+        # Convert to Warp arrays
+        wp_indices = wp.array(indices, dtype=wp.int32)
+        wp_id_numbers = wp.array(id_numbers, dtype=wp.uint8)
+        wp_is_interior = wp.array(is_interior, dtype=wp.bool)
+
+        if start_index is None:
+            start_index = wp.vec3i(0, 0, 0)
+        else:
+            start_index = wp.vec3i(*start_index)
+
+        # Launch the warp kernel
+        wp.launch(
+            self.warp_kernel,
+            dim=current_index,
+            inputs=[
+                wp_indices,
+                wp_id_numbers,
+                wp_is_interior,
+                bc_mask,
+                missing_mask,
+                start_index,
+            ],
+        )
+
+        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
new file mode 100644
index 0000000..edfdb42
--- /dev/null
+++ b/xlb/operator/boundary_masker/mesh_boundary_masker.py
@@ -0,0 +1,153 @@
+# Base class for all equilibriums
+
+import numpy as np
+import warp as wp
+import jax
+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
+
+
+class MeshBoundaryMasker(Operator):
+    """
+    Operator for creating a boundary missing_mask from an STL file
+    """
+
+    def __init__(
+        self,
+        velocity_set: VelocitySet,
+        precision_policy: PrecisionPolicy,
+        compute_backend: ComputeBackend.WARP,
+    ):
+        # Call super
+        super().__init__(velocity_set, precision_policy, compute_backend)
+
+        # Raise error if used for 2d examples:
+        if self.velocity_set.d == 2:
+            raise NotImplementedError("This Operator is not implemented in 2D!")
+
+        # Also using Warp kernels for JAX implementation
+        if self.compute_backend == ComputeBackend.JAX:
+            self.warp_functional, self.warp_kernel = self._construct_warp()
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    def jax_implementation(
+        self,
+        bc,
+        origin,
+        spacing,
+        id_number,
+        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()
+        bc_mask = wp.from_jax(bc_mask)
+        missing_mask = wp.from_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
+        _c = self.velocity_set.c
+        _q = wp.constant(self.velocity_set.q)
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            mesh_id: wp.uint64,
+            origin: wp.vec3,
+            spacing: wp.vec3,
+            id_number: wp.int32,
+            bc_mask: wp.array4d(dtype=wp.uint8),
+            missing_mask: wp.array4d(dtype=wp.bool),
+            start_index: wp.vec3i,
+        ):
+            # get index
+            i, j, k = wp.tid()
+
+            # Get local indices
+            index = wp.vec3i()
+            index[0] = i - start_index[0]
+            index[1] = j - start_index[1]
+            index[2] = k - start_index[2]
+
+            # position of the point
+            ijk = wp.vec3(wp.float32(index[0]), wp.float32(index[1]), wp.float32(index[2]))
+            ijk = ijk + wp.vec3(0.5, 0.5, 0.5)  # cell center
+            pos = wp.cw_mul(ijk, spacing) + origin
+
+            # Compute the maximum length
+            max_length = wp.sqrt(
+                (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
+            query = wp.mesh_query_point_sign_winding_number(mesh_id, pos, max_length)
+            if query.result:
+                # set point to be solid
+                if query.sign <= 0:  # TODO: fix this
+                    # Stream indices
+                    for l in range(1, _q):
+                        # Get the index of the streaming direction
+                        push_index = wp.vec3i()
+                        for d in range(self.velocity_set.d):
+                            push_index[d] = index[d] + _c[d, l]
+
+                        # Set the boundary id and missing_mask
+                        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
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(
+        self,
+        bc,
+        origin,
+        spacing,
+        bc_mask,
+        missing_mask,
+        start_index=(0, 0, 0),
+    ):
+        assert bc.mesh_vertices is not None, f'Please provide the mesh vertices for {bc.__class__.__name__} BC using keyword "mesh_vertices"!'
+        assert bc.indices is None, f"Please use IndicesBoundaryMasker operator if {bc.__class__.__name__} is imposed on known indices of the grid!"
+        assert (
+            bc.mesh_vertices.shape[1] == self.velocity_set.d
+        ), "Mesh points must be reshaped into an array (N, 3) where N indicates number of points!"
+        mesh_vertices = bc.mesh_vertices
+        id_number = bc.id
+
+        # We are done with bc.mesh_vertices. Remove them from BC objects
+        bc.__dict__.pop("mesh_vertices", None)
+
+        # Ensure this masker is called only for BCs that need implicit distance to the mesh
+        assert not bc.needs_mesh_distance, 'Please use "MeshDistanceBoundaryMasker" if this BC needs mesh distance!'
+
+        mesh_indices = np.arange(mesh_vertices.shape[0])
+        mesh = wp.Mesh(
+            points=wp.array(mesh_vertices, dtype=wp.vec3),
+            indices=wp.array(mesh_indices, dtype=int),
+        )
+
+        # Launch the warp kernel
+        wp.launch(
+            self.warp_kernel,
+            inputs=[
+                mesh.id,
+                origin,
+                spacing,
+                id_number,
+                bc_mask,
+                missing_mask,
+                start_index,
+            ],
+            dim=bc_mask.shape[1:],
+        )
+
+        return bc_mask, missing_mask
diff --git a/xlb/operator/boundary_masker/mesh_distance_boundary_masker.py b/xlb/operator/boundary_masker/mesh_distance_boundary_masker.py
new file mode 100644
index 0000000..87af94c
--- /dev/null
+++ b/xlb/operator/boundary_masker/mesh_distance_boundary_masker.py
@@ -0,0 +1,184 @@
+# Base class for all equilibriums
+
+import numpy as np
+import warp as wp
+import jax
+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
+
+
+class MeshDistanceBoundaryMasker(Operator):
+    """
+    Operator for creating a boundary missing_mask from an STL file
+    """
+
+    def __init__(
+        self,
+        velocity_set: VelocitySet,
+        precision_policy: PrecisionPolicy,
+        compute_backend: ComputeBackend.WARP,
+    ):
+        # Call super
+        super().__init__(velocity_set, precision_policy, compute_backend)
+
+        # Raise error if used for 2d examples:
+        if self.velocity_set.d == 2:
+            raise NotImplementedError("This Operator is not implemented in 2D!")
+
+        # Also using Warp kernels for JAX implementation
+        if self.compute_backend == ComputeBackend.JAX:
+            self.warp_functional, self.warp_kernel = self._construct_warp()
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    def jax_implementation(
+        self,
+        bc,
+        origin,
+        spacing,
+        id_number,
+        bc_mask,
+        missing_mask,
+        f_field,
+        start_index=(0, 0, 0),
+    ):
+        raise NotImplementedError(f"Operation {self.__class__.__name} not implemented in JAX!")
+
+    def _construct_warp(self):
+        # Make constants for warp
+        _c = self.velocity_set.c
+        _q = wp.constant(self.velocity_set.q)
+        _opp_indices = self.velocity_set.opp_indices
+
+        @wp.func
+        def check_index_bounds(index: wp.vec3i, shape: wp.vec3i):
+            is_in_bounds = index[0] >= 0 and index[0] < shape[0] and index[1] >= 0 and index[1] < shape[1] and index[2] >= 0 and index[2] < shape[2]
+            return is_in_bounds
+
+        @wp.func
+        def index_to_position(index: wp.vec3i, origin: wp.vec3, spacing: wp.vec3):
+            # position of the point
+            ijk = wp.vec3(wp.float32(index[0]), wp.float32(index[1]), wp.float32(index[2]))
+            ijk = ijk + wp.vec3(0.5, 0.5, 0.5)  # cell center
+            pos = wp.cw_mul(ijk, spacing) + origin
+            return pos
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            mesh_id: wp.uint64,
+            origin: wp.vec3,
+            spacing: wp.vec3,
+            id_number: wp.int32,
+            bc_mask: wp.array4d(dtype=wp.uint8),
+            missing_mask: wp.array4d(dtype=wp.bool),
+            f_field: wp.array4d(dtype=Any),
+            start_index: wp.vec3i,
+        ):
+            # get index
+            i, j, k = wp.tid()
+
+            # Get local indices
+            index = wp.vec3i()
+            index[0] = i - start_index[0]
+            index[1] = j - start_index[1]
+            index[2] = k - start_index[2]
+
+            # position of the point
+            pos_solid_cell = index_to_position(index, origin, spacing)
+
+            # Compute the maximum length
+            max_length = wp.sqrt(
+                (spacing[0] * wp.float32(missing_mask.shape[1])) ** 2.0
+                + (spacing[1] * wp.float32(missing_mask.shape[2])) ** 2.0
+                + (spacing[2] * wp.float32(missing_mask.shape[3])) ** 2.0
+            )
+
+            # evaluate if point is inside mesh
+            query = wp.mesh_query_point_sign_winding_number(mesh_id, pos_solid_cell, max_length)
+            if query.result and query.sign <= 0:  # TODO: fix this
+                # Set bc_mask of solid to a large number to enable skipping LBM operations
+                bc_mask[0, index[0], index[1], index[2]] = wp.uint8(255)
+
+                # Find neighboring fluid cells along each lattice direction and the their fractional distance to the mesh
+                for l in range(1, _q):
+                    # Get the index of the streaming direction
+                    push_index = wp.vec3i()
+                    for d in range(self.velocity_set.d):
+                        push_index[d] = index[d] + _c[d, l]
+                    shape = wp.vec3i(missing_mask.shape[1], missing_mask.shape[2], missing_mask.shape[3])
+                    if check_index_bounds(push_index, shape):
+                        # find neighbouring fluid cell
+                        pos_fluid_cell = index_to_position(push_index, origin, spacing)
+                        query = wp.mesh_query_point_sign_winding_number(mesh_id, pos_fluid_cell, max_length)
+                        if query.result and query.sign > 0:
+                            # Set the boundary id and missing_mask
+                            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
+
+                            # get position of the mesh triangle that intersects with the solid cell
+                            pos_mesh = wp.mesh_eval_position(mesh_id, query.face, query.u, query.v)
+                            weight = wp.length(pos_fluid_cell - pos_mesh) / wp.length(pos_fluid_cell - pos_solid_cell)
+                            f_field[_opp_indices[l], push_index[0], push_index[1], push_index[2]] = self.store_dtype(weight)
+
+        return None, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(
+        self,
+        bc,
+        origin,
+        spacing,
+        bc_mask,
+        missing_mask,
+        f_field,
+        start_index=(0, 0, 0),
+    ):
+        assert bc.mesh_vertices is not None, f'Please provide the mesh vertices for {bc.__class__.__name__} BC using keyword "mesh_vertices"!'
+        assert bc.indices is None, f"Cannot find the implicit distance to the boundary for {bc.__class__.__name__} without a mesh!"
+        assert (
+            bc.mesh_vertices.shape[1] == self.velocity_set.d
+        ), "Mesh points must be reshaped into an array (N, 3) where N indicates number of points!"
+        assert (
+            f_field is not None and f_field.shape == missing_mask.shape
+        ), "To compute and store the implicit distance to the boundary for this BC, use a population field!"
+        mesh_vertices = bc.mesh_vertices
+        id_number = bc.id
+
+        # We are done with bc.mesh_vertices. Remove them from BC objects
+        bc.__dict__.pop("mesh_vertices", None)
+
+        # Ensure this masker is called only for BCs that need implicit distance to the mesh
+        assert bc.needs_mesh_distance, 'Please use "MeshBoundaryMasker" if this BC does NOT need mesh distance!'
+
+        mesh_indices = np.arange(mesh_vertices.shape[0])
+        mesh = wp.Mesh(
+            points=wp.array(mesh_vertices, dtype=wp.vec3),
+            indices=wp.array(mesh_indices, dtype=int),
+        )
+
+        # Convert input tuples to warp vectors
+        origin = wp.vec3(origin[0], origin[1], origin[2])
+        spacing = wp.vec3(spacing[0], spacing[1], spacing[2])
+        start_index = wp.vec3i(start_index[0], start_index[1], start_index[2])
+        mesh_id = wp.uint64(mesh.id)
+
+        # Launch the warp kernel
+        wp.launch(
+            self.warp_kernel,
+            inputs=[
+                mesh_id,
+                origin,
+                spacing,
+                id_number,
+                bc_mask,
+                missing_mask,
+                f_field,
+                start_index,
+            ],
+            dim=missing_mask.shape[1:],
+        )
+
+        return bc_mask, missing_mask, f_field
diff --git a/xlb/operator/collision/__init__.py b/xlb/operator/collision/__init__.py
new file mode 100644
index 0000000..2f92bb4
--- /dev/null
+++ b/xlb/operator/collision/__init__.py
@@ -0,0 +1,4 @@
+from xlb.operator.collision.collision import Collision
+from xlb.operator.collision.bgk import BGK
+from xlb.operator.collision.kbc import KBC
+from xlb.operator.collision.forced_collision import ForcedCollision
diff --git a/xlb/operator/collision/bgk.py b/xlb/operator/collision/bgk.py
new file mode 100644
index 0000000..115ed9a
--- /dev/null
+++ b/xlb/operator/collision/bgk.py
@@ -0,0 +1,80 @@
+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
+
+
+class BGK(Collision):
+    """
+    BGK collision operator for LBM.
+    """
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0,))
+    def jax_implementation(self, f: jnp.ndarray, feq: jnp.ndarray, rho, u):
+        fneq = f - feq
+        fout = f - self.compute_dtype(self.omega) * fneq
+        return fout
+
+    def _construct_warp(self):
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        _w = self.velocity_set.w
+        _omega = wp.constant(self.compute_dtype(self.omega))
+        _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):
+            fneq = f - feq
+            fout = f - _omega * fneq
+            return fout
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            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()
+            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]]
+
+            # 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]] = self.store_dtype(_fout[l])
+
+        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/collision.py b/xlb/operator/collision/collision.py
new file mode 100644
index 0000000..00a8dfd
--- /dev/null
+++ b/xlb/operator/collision/collision.py
@@ -0,0 +1,31 @@
+"""
+Base class for Collision operators
+"""
+
+from xlb.velocity_set import VelocitySet
+from xlb.operator import Operator
+
+
+class Collision(Operator):
+    """
+    Base class for collision operators.
+
+    This class defines the collision step for the Lattice Boltzmann Method.
+
+    Parameters
+    ----------
+    omega : float
+        Relaxation parameter for collision step. Default value is 0.6.
+    shear : bool
+        Flag to indicate whether the collision step requires the shear stress.
+    """
+
+    def __init__(
+        self,
+        omega: float,
+        velocity_set: VelocitySet = None,
+        precision_policy=None,
+        compute_backend=None,
+    ):
+        self.omega = omega
+        super().__init__(velocity_set, precision_policy, compute_backend)
diff --git a/xlb/operator/collision/forced_collision.py b/xlb/operator/collision/forced_collision.py
new file mode 100644
index 0000000..2036bab
--- /dev/null
+++ b/xlb/operator/collision/forced_collision.py
@@ -0,0 +1,101 @@
+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 kernel(
+            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]
+
+        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
new file mode 100644
index 0000000..5c5c29e
--- /dev/null
+++ b/xlb/operator/collision/kbc.py
@@ -0,0 +1,335 @@
+"""
+KBC collision operator for LBM.
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import warp as wp
+from typing import Any
+from functools import partial
+
+from xlb.velocity_set import VelocitySet, D2Q9, D3Q27
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.collision.collision import Collision
+from xlb.operator import Operator
+from xlb.operator.macroscopic import SecondMoment as MomentumFlux
+
+
+class KBC(Collision):
+    """
+    KBC collision operator for LBM.
+
+    This class implements the Karlin-Bösch-Chikatamarla (KBC) model for the collision step in the Lattice Boltzmann Method.
+    """
+
+    def __init__(
+        self,
+        omega: float,
+        velocity_set: VelocitySet = None,
+        precision_policy=None,
+        compute_backend=None,
+    ):
+        self.momentum_flux = MomentumFlux()
+        self.epsilon = 1e-32
+        self.beta = omega * 0.5
+        self.inv_beta = 1.0 / self.beta
+
+        super().__init__(
+            omega=omega,
+            velocity_set=velocity_set,
+            precision_policy=precision_policy,
+            compute_backend=compute_backend,
+        )
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0,), donate_argnums=(1, 2, 3))
+    def jax_implementation(
+        self,
+        f: jnp.ndarray,
+        feq: jnp.ndarray,
+        rho: jnp.ndarray,
+        u: jnp.ndarray,
+    ):
+        """
+        KBC collision step for lattice.
+
+        Parameters
+        ----------
+        f : jax.numpy.array
+            Distribution function.
+        feq : jax.numpy.array
+            Equilibrium distribution function.
+        rho : jax.numpy.array
+            Density.
+        u : jax.numpy.array
+            Velocity.
+        """
+        fneq = f - feq
+        if isinstance(self.velocity_set, D2Q9):
+            shear = self.decompose_shear_d2q9_jax(fneq)
+            delta_s = shear * rho / 4.0  # TODO: Check this
+        elif isinstance(self.velocity_set, D3Q27):
+            shear = self.decompose_shear_d3q27_jax(fneq)
+            delta_s = shear * rho
+        else:
+            raise NotImplementedError("Velocity set not supported: {}".format(type(self.velocity_set)))
+
+        # Perform collision
+        delta_h = fneq - delta_s
+        gamma = self.inv_beta - (2.0 - self.inv_beta) * self.entropic_scalar_product(delta_s, delta_h, feq) / (
+            self.epsilon + self.entropic_scalar_product(delta_h, delta_h, feq)
+        )
+
+        fout = f - self.beta * (2.0 * delta_s + gamma[None, ...] * delta_h)
+
+        return fout
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def entropic_scalar_product(self, x: jnp.ndarray, y: jnp.ndarray, feq: jnp.ndarray):
+        """
+        Compute the entropic scalar product of x and y to approximate gamma in KBC.
+
+        Returns
+        -------
+        jax.numpy.array
+            Entropic scalar product of x, y, and feq.
+        """
+        return jnp.sum(x * y / feq, axis=0)
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def decompose_shear_d3q27_jax(self, fneq):
+        """
+        Decompose fneq into shear components for D3Q27 lattice.
+
+        Parameters
+        ----------
+        fneq : jax.numpy.ndarray
+            Non-equilibrium distribution function.
+
+        Returns
+        -------
+        jax.numpy.ndarray
+            Shear components of fneq.
+        """
+
+        # Calculate the momentum flux
+        Pi = self.momentum_flux(fneq)
+        # Calculating Nxz and Nyz with indices moved to the first dimension
+        Nxz = Pi[0, ...] - Pi[5, ...]
+        Nyz = Pi[3, ...] - Pi[5, ...]
+
+        # For c = (i, 0, 0), c = (0, j, 0) and c = (0, 0, k)
+        s = jnp.zeros_like(fneq)
+        s = s.at[9, ...].set((2.0 * Nxz - Nyz) / 6.0)
+        s = s.at[18, ...].set((2.0 * Nxz - Nyz) / 6.0)
+        s = s.at[3, ...].set((-Nxz + 2.0 * Nyz) / 6.0)
+        s = s.at[6, ...].set((-Nxz + 2.0 * Nyz) / 6.0)
+        s = s.at[1, ...].set((-Nxz - Nyz) / 6.0)
+        s = s.at[2, ...].set((-Nxz - Nyz) / 6.0)
+
+        # For c = (i, j, 0)
+        s = s.at[12, ...].set(Pi[1, ...] / 4.0)
+        s = s.at[24, ...].set(Pi[1, ...] / 4.0)
+        s = s.at[21, ...].set(-Pi[1, ...] / 4.0)
+        s = s.at[15, ...].set(-Pi[1, ...] / 4.0)
+
+        # For c = (i, 0, k)
+        s = s.at[10, ...].set(Pi[2, ...] / 4.0)
+        s = s.at[20, ...].set(Pi[2, ...] / 4.0)
+        s = s.at[19, ...].set(-Pi[2, ...] / 4.0)
+        s = s.at[11, ...].set(-Pi[2, ...] / 4.0)
+
+        # For c = (0, j, k)
+        s = s.at[8, ...].set(Pi[4, ...] / 4.0)
+        s = s.at[4, ...].set(Pi[4, ...] / 4.0)
+        s = s.at[7, ...].set(-Pi[4, ...] / 4.0)
+        s = s.at[5, ...].set(-Pi[4, ...] / 4.0)
+
+        return s
+
+    @partial(jit, static_argnums=(0,), inline=True)
+    def decompose_shear_d2q9_jax(self, fneq):
+        """
+        Decompose fneq into shear components for D2Q9 lattice.
+
+        Parameters
+        ----------
+        fneq : jax.numpy.array
+            Non-equilibrium distribution function.
+
+        Returns
+        -------
+        jax.numpy.array
+            Shear components of fneq.
+        """
+        Pi = self.momentum_flux(fneq)
+        N = Pi[0, ...] - Pi[2, ...]
+        s = jnp.zeros_like(fneq)
+        s = s.at[3, ...].set(N)
+        s = s.at[6, ...].set(N)
+        s = s.at[2, ...].set(-N)
+        s = s.at[1, ...].set(-N)
+        s = s.at[8, ...].set(Pi[1, ...])
+        s = s.at[4, ...].set(-Pi[1, ...])
+        s = s.at[5, ...].set(-Pi[1, ...])
+        s = s.at[7, ...].set(Pi[1, ...])
+
+        return s
+
+    def _construct_warp(self):
+        # Raise error if velocity set is not supported
+        if not (isinstance(self.velocity_set, D3Q27) or isinstance(self.velocity_set, D2Q9)):
+            raise NotImplementedError("Velocity set not supported for warp backend: {}".format(type(self.velocity_set)))
+
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        _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))
+        _inv_beta = wp.constant(self.compute_dtype(1.0 / self.beta))
+
+        @wp.func
+        def decompose_shear_d2q9(fneq: Any):
+            pi = self.momentum_flux.warp_functional(fneq)
+            N = pi[0] - pi[1]
+            s = _f_vec()
+            s[3] = N
+            s[6] = N
+            s[2] = -N
+            s[1] = -N
+            s[8] = pi[2]
+            s[4] = -pi[2]
+            s[5] = -pi[2]
+            s[7] = pi[2]
+            return s
+
+        # Construct functional for decomposing shear
+        @wp.func
+        def decompose_shear_d3q27(
+            fneq: Any,
+        ):
+            # Get momentum flux
+            pi = self.momentum_flux.warp_functional(fneq)
+            nxz = pi[0] - pi[5]
+            nyz = pi[3] - pi[5]
+
+            # set shear components
+            s = _f_vec()
+
+            # For c = (i, 0, 0), c = (0, j, 0) and c = (0, 0, k)
+            two = self.compute_dtype(2.0)
+            four = self.compute_dtype(4.0)
+            six = 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] / 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] / 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] / four
+            s[4] = pi[4] / four
+            s[7] = -pi[4] / four
+            s[5] = -pi[4] / four
+
+            return s
+
+        # Construct functional for computing entropic scalar product
+        @wp.func
+        def entropic_scalar_product(
+            x: Any,
+            y: Any,
+            feq: Any,
+        ):
+            e = wp.cw_div(wp.cw_mul(x, y), feq)
+            e_sum = self.compute_dtype(0.0)
+            for i in range(self.velocity_set.q):
+                e_sum += e[i]
+            return e_sum
+
+        # Construct the functional
+        @wp.func
+        def functional(
+            f: Any,
+            feq: Any,
+            rho: Any,
+            u: Any,
+        ):
+            # Compute shear and delta_s
+            fneq = f - feq
+            shear = decompose_shear_d3q27(fneq)
+            delta_s = shear * rho  # TODO: Check this
+
+            # Perform collision
+            delta_h = fneq - delta_s
+            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 * (two * delta_s + gamma * delta_h)
+
+            return fout
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            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]] = self.store_dtype(_fout[l])
+
+        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/equilibrium/__init__.py b/xlb/operator/equilibrium/__init__.py
new file mode 100644
index 0000000..987aa74
--- /dev/null
+++ b/xlb/operator/equilibrium/__init__.py
@@ -0,0 +1,1 @@
+from xlb.operator.equilibrium.quadratic_equilibrium import Equilibrium, QuadraticEquilibrium
diff --git a/xlb/operator/equilibrium/equilibrium.py b/xlb/operator/equilibrium/equilibrium.py
new file mode 100644
index 0000000..11f4155
--- /dev/null
+++ b/xlb/operator/equilibrium/equilibrium.py
@@ -0,0 +1,17 @@
+# Base class for all equilibriums
+from xlb.velocity_set.velocity_set import VelocitySet
+from xlb.operator.operator import Operator
+
+
+class Equilibrium(Operator):
+    """
+    Base class for all equilibriums
+    """
+
+    def __init__(
+        self,
+        velocity_set: VelocitySet = None,
+        precision_policy=None,
+        compute_backend=None,
+    ):
+        super().__init__(velocity_set, precision_policy, compute_backend)
diff --git a/xlb/operator/equilibrium/quadratic_equilibrium.py b/xlb/operator/equilibrium/quadratic_equilibrium.py
new file mode 100644
index 0000000..62cc041
--- /dev/null
+++ b/xlb/operator/equilibrium/quadratic_equilibrium.py
@@ -0,0 +1,98 @@
+from functools import partial
+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.equilibrium.equilibrium import Equilibrium
+from xlb.operator import Operator
+
+
+class QuadraticEquilibrium(Equilibrium):
+    """
+    Quadratic equilibrium of Boltzmann equation using hermite polynomials.
+    Standard equilibrium model for LBM.
+    """
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(self, rho, u):
+        cu = 3.0 * jnp.tensordot(self.velocity_set.c, u, axes=(0, 0))
+        usqr = 1.5 * jnp.sum(jnp.square(u), axis=0, keepdims=True)
+        w = self.velocity_set.w.reshape((-1,) + (1,) * (len(rho.shape) - 1))
+        feq = rho * w * (1.0 + cu * (1.0 + 0.5 * cu) - usqr)
+        return feq
+
+    def _construct_warp(self):
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        _c = self.velocity_set.c
+        _w = self.velocity_set.w
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+
+        # Construct the equilibrium functional
+        @wp.func
+        def functional(
+            rho: Any,
+            u: Any,
+        ):
+            # Allocate the equilibrium
+            feq = _f_vec()
+
+            # Compute the equilibrium
+            for l in range(self.velocity_set.q):
+                # Compute cu
+                cu = self.compute_dtype(0.0)
+                for d in range(self.velocity_set.d):
+                    if _c[d, l] == 1:
+                        cu += u[d]
+                    elif _c[d, l] == -1:
+                        cu -= u[d]
+                cu *= self.compute_dtype(3.0)
+
+                # Compute usqr
+                usqr = self.compute_dtype(1.5) * wp.dot(u, u)
+
+                # Compute feq
+                feq[l] = rho * _w[l] * (self.compute_dtype(1.0) + cu * (self.compute_dtype(1.0) + self.compute_dtype(0.5) * cu) - usqr)
+
+            return feq
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            rho: wp.array4d(dtype=Any),
+            u: wp.array4d(dtype=Any),
+            f: wp.array4d(dtype=Any),
+        ):
+            # Get the global index
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            # Get the equilibrium
+            _u = _u_vec()
+            for d in range(self.velocity_set.d):
+                _u[d] = u[d, index[0], index[1], index[2]]
+            _rho = rho[0, index[0], index[1], index[2]]
+            feq = functional(_rho, _u)
+
+            # Set the output
+            for l in range(self.velocity_set.q):
+                f[l, index[0], index[1], index[2]] = self.store_dtype(feq[l])
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, rho, u, f):
+        # Launch the warp kernel
+        wp.launch(
+            self.warp_kernel,
+            inputs=[
+                rho,
+                u,
+                f,
+            ],
+            dim=rho.shape[1:],
+        )
+        return f
diff --git a/xlb/operator/force/__init__.py b/xlb/operator/force/__init__.py
new file mode 100644
index 0000000..ba8a13c
--- /dev/null
+++ b/xlb/operator/force/__init__.py
@@ -0,0 +1,2 @@
+from xlb.operator.force.momentum_transfer import MomentumTransfer
+from xlb.operator.force.exact_difference_force import 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..ec1ef5b
--- /dev/null
+++ b/xlb/operator/force/exact_difference_force.py
@@ -0,0 +1,123 @@
+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 kernel(
+            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]] = self.store_dtype(_fout[l])
+
+        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
new file mode 100644
index 0000000..dbd5307
--- /dev/null
+++ b/xlb/operator/force/momentum_transfer.py
@@ -0,0 +1,172 @@
+from functools import partial
+import jax.numpy as jnp
+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.stream import Stream
+
+
+class MomentumTransfer(Operator):
+    """
+    An opertor for the momentum exchange method to compute the boundary force vector exerted on the solid geometry
+    based on [1] as described in [3]. Ref [2] shows how [1] is applicable to curved geometries only by using a
+    bounce-back method (e.g. Bouzidi) that accounts for curved boundaries.
+    NOTE: this function should be called after BC's are imposed.
+    [1] A.J.C. Ladd, Numerical simulations of particular suspensions via a discretized Boltzmann equation.
+        Part 2 (numerical results), J. Fluid Mech. 271 (1994) 311-339.
+    [2] R. Mei, D. Yu, W. Shyy, L.-S. Luo, Force evaluation in the lattice Boltzmann method involving
+        curved geometry, Phys. Rev. E 65 (2002) 041203.
+    [3] Caiazzo, A., & Junk, M. (2008). Boundary forces in lattice Boltzmann: Analysis of momentum exchange
+        algorithm. Computers & Mathematics with Applications, 55(7), 1415-1423.
+
+    Notes
+    -----
+    This method computes the force exerted on the solid geometry at each boundary node using the momentum exchange method.
+    The force is computed based on the post-streaming and post-collision distribution functions. This method
+    should be called after the boundary conditions are imposed.
+    """
+
+    def __init__(
+        self,
+        no_slip_bc_instance,
+        velocity_set: VelocitySet = None,
+        precision_policy: PrecisionPolicy = None,
+        compute_backend: ComputeBackend = None,
+    ):
+        self.no_slip_bc_instance = no_slip_bc_instance
+        self.stream = Stream(velocity_set, precision_policy, compute_backend)
+
+        # 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, bc_mask, missing_mask):
+        """
+        Parameters
+        ----------
+        f : jax.numpy.ndarray
+            The post-collision distribution function at each node in the grid.
+        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
+            A grid field with lattice cardinality that specifies missing lattice directions
+            for each boundary node.
+
+        Returns
+        -------
+        jax.numpy.ndarray
+            The force exerted on the solid geometry at each boundary node.
+        """
+        # 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, bc_mask, missing_mask)
+
+        # Compute momentum transfer
+        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)))
+
+        # the following will return force as a grid-based field with zero everywhere except for boundary nodes.
+        opp = self.velocity_set.opp_indices
+        phi = f_post_collision[opp] + f_post_stream
+        phi = jnp.where(jnp.logical_and(boundary, missing_mask), phi, 0.0)
+        force = jnp.tensordot(self.velocity_set.c[:, opp], phi, axes=(-1, 0))
+        return force
+
+    def _construct_warp(self):
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        _c = self.velocity_set.c
+        _opp_indices = self.velocity_set.opp_indices
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)  # TODO fix vec bool
+        _no_slip_id = self.no_slip_bc_instance.id
+
+        # Find velocity index for 0, 0, 0
+        for l in range(self.velocity_set.q):
+            if _c[0, l] == 0 and _c[1, l] == 0 and _c[2, l] == 0:
+                zero_index = l
+        _zero_index = wp.int32(zero_index)
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            f_0: wp.array4d(dtype=Any),
+            f_1: wp.array4d(dtype=Any),
+            bc_mask: wp.array4d(dtype=wp.uint8),
+            missing_mask: wp.array4d(dtype=wp.bool),
+            force: wp.array(dtype=Any),
+        ):
+            # Get the global index
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            # Get the boundary id
+            _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
+                if missing_mask[l, index[0], index[1], index[2]]:
+                    _missing_mask[l] = wp.uint8(1)
+                else:
+                    _missing_mask[l] = wp.uint8(0)
+
+            # Determin if boundary is an edge by checking if center is missing
+            is_edge = wp.bool(False)
+            if _boundary_id == wp.uint8(_no_slip_id):
+                if _missing_mask[_zero_index] == wp.uint8(0):
+                    is_edge = wp.bool(True)
+
+            # If the boundary is an edge then add the momentum transfer
+            m = _u_vec()
+            if is_edge:
+                # Get the distribution function
+                f_post_collision = _f_vec()
+                for l in range(self.velocity_set.q):
+                    f_post_collision[l] = f_0[l, index[0], index[1], index[2]]
+
+                # Apply streaming (pull method)
+                timestep = 0
+                f_post_stream = self.stream.warp_functional(f_0, index)
+                f_post_stream = self.no_slip_bc_instance.warp_functional(index, timestep, _missing_mask, f_0, f_1, f_post_collision, f_post_stream)
+
+                # Compute the momentum transfer
+                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]
+                            if _c[d, _opp_indices[l]] == 1:
+                                m[d] += phi
+                            elif _c[d, _opp_indices[l]] == -1:
+                                m[d] -= phi
+
+            wp.atomic_add(force, 0, m)
+
+        return None, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f_0, f_1, bc_mask, missing_mask):
+        # Allocate the force vector (the total integral value will be computed)
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+        force = wp.zeros((1), dtype=_u_vec)
+
+        # Launch the warp kernel
+        wp.launch(
+            self.warp_kernel,
+            inputs=[f_0, f_1, bc_mask, missing_mask, force],
+            dim=f_0.shape[1:],
+        )
+        return force.numpy()[0]
diff --git a/xlb/operator/macroscopic/__init__.py b/xlb/operator/macroscopic/__init__.py
new file mode 100644
index 0000000..75dec9e
--- /dev/null
+++ b/xlb/operator/macroscopic/__init__.py
@@ -0,0 +1,4 @@
+from xlb.operator.macroscopic.macroscopic import Macroscopic
+from xlb.operator.macroscopic.second_moment import SecondMoment
+from xlb.operator.macroscopic.zero_moment import ZeroMoment
+from xlb.operator.macroscopic.first_moment import FirstMoment
diff --git a/xlb/operator/macroscopic/first_moment.py b/xlb/operator/macroscopic/first_moment.py
new file mode 100644
index 0000000..cb99a9f
--- /dev/null
+++ b/xlb/operator/macroscopic/first_moment.py
@@ -0,0 +1,67 @@
+from functools import partial
+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.operator import Operator
+
+
+class FirstMoment(Operator):
+    """A class to compute the first moment (velocity) of distribution functions."""
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0), inline=True)
+    def jax_implementation(self, f, rho):
+        u = jnp.tensordot(self.velocity_set.c, f, axes=(-1, 0)) / rho
+        return u
+
+    def _construct_warp(self):
+        _c = self.velocity_set.c
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _u_vec = wp.vec(self.velocity_set.d, dtype=self.compute_dtype)
+
+        @wp.func
+        def functional(
+            f: _f_vec,
+            rho: Any,
+        ):
+            u = _u_vec()
+            for l in range(self.velocity_set.q):
+                for d in range(self.velocity_set.d):
+                    if _c[d, l] == 1:
+                        u[d] += f[l]
+                    elif _c[d, l] == -1:
+                        u[d] -= f[l]
+            u /= rho
+            return u
+
+        @wp.kernel
+        def kernel(
+            f: wp.array4d(dtype=Any),
+            rho: wp.array4d(dtype=Any),
+            u: wp.array4d(dtype=Any),
+        ):
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            _f = _f_vec()
+            for l in range(self.velocity_set.q):
+                _f[l] = f[l, index[0], index[1], index[2]]
+            _rho = rho[0, index[0], index[1], index[2]]
+            _u = functional(_f, _rho)
+
+            for d in range(self.velocity_set.d):
+                u[d, index[0], index[1], index[2]] = self.store_dtype(_u[d])
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f, rho, u):
+        wp.launch(
+            self.warp_kernel,
+            inputs=[f, rho, u],
+            dim=u.shape[1:],
+        )
+        return u
diff --git a/xlb/operator/macroscopic/macroscopic.py b/xlb/operator/macroscopic/macroscopic.py
new file mode 100644
index 0000000..ab1193b
--- /dev/null
+++ b/xlb/operator/macroscopic/macroscopic.py
@@ -0,0 +1,66 @@
+from functools import partial
+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.operator import Operator
+from xlb.operator.macroscopic.zero_moment import ZeroMoment
+from xlb.operator.macroscopic.first_moment import FirstMoment
+
+
+class Macroscopic(Operator):
+    """A class to compute both zero and first moments of distribution functions (rho, u)."""
+
+    def __init__(self, *args, **kwargs):
+        self.zero_moment = ZeroMoment(*args, **kwargs)
+        self.first_moment = FirstMoment(*args, **kwargs)
+        super().__init__(*args, **kwargs)
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0), inline=True)
+    def jax_implementation(self, f):
+        rho = self.zero_moment(f)
+        u = self.first_moment(f, rho)
+        return rho, u
+
+    def _construct_warp(self):
+        zero_moment_func = self.zero_moment.warp_functional
+        first_moment_func = self.first_moment.warp_functional
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+
+        @wp.func
+        def functional(f: _f_vec):
+            rho = zero_moment_func(f)
+            u = first_moment_func(f, rho)
+            return rho, u
+
+        @wp.kernel
+        def kernel(
+            f: wp.array4d(dtype=Any),
+            rho: wp.array4d(dtype=Any),
+            u: wp.array4d(dtype=Any),
+        ):
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            _f = _f_vec()
+            for l in range(self.velocity_set.q):
+                _f[l] = f[l, index[0], index[1], index[2]]
+            _rho, _u = functional(_f)
+
+            rho[0, index[0], index[1], index[2]] = self.store_dtype(_rho)
+            for d in range(self.velocity_set.d):
+                u[d, index[0], index[1], index[2]] = self.store_dtype(_u[d])
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f, rho, u):
+        wp.launch(
+            self.warp_kernel,
+            inputs=[f, rho, u],
+            dim=rho.shape[1:],
+        )
+        return rho, u
diff --git a/xlb/operator/macroscopic/second_moment.py b/xlb/operator/macroscopic/second_moment.py
new file mode 100644
index 0000000..6c7e70e
--- /dev/null
+++ b/xlb/operator/macroscopic/second_moment.py
@@ -0,0 +1,106 @@
+# Base class for all equilibriums
+
+from functools import partial
+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.operator import Operator
+
+
+class SecondMoment(Operator):
+    """
+    Operator to calculate the second moment of distribution functions.
+
+    The second moment may be used to compute the momentum flux in the computation of
+    the stress tensor in the Lattice Boltzmann Method (LBM).
+
+    Important Note:
+    Note that this rank 2 symmetric tensor (dim*dim) has been converted into a rank one
+    vector where the diagonal and off-diagonal components correspond to the following elements of
+    the vector:
+    if self.grid.dim == 3:
+        diagonal    = (0, 3, 5)
+        offdiagonal = (1, 2, 4)
+    elif self.grid.dim == 2:
+        diagonal    = (0, 2)
+        offdiagonal = (1,)
+
+    ** For any reduction operation on the full tensor it is crucial to account for the full tensor by
+    considering all diagonal and off-diagonal components.
+    """
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0,), donate_argnums=(1,))
+    def jax_implementation(
+        self,
+        fneq: jnp.ndarray,
+    ):
+        """
+        This function computes the second order moment, which is the product of the
+        distribution functions (f) and the lattice moments (cc).
+
+        Parameters
+        ----------
+        fneq: jax.numpy.ndarray
+            The distribution functions.
+
+        Returns
+        -------
+        jax.numpy.ndarray
+            The computed second moment.
+        """
+        return jnp.tensordot(self.velocity_set.cc, fneq, axes=(0, 0))
+
+    def _construct_warp(self):
+        # Make constants for warp
+        _cc = self.velocity_set.cc
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _pi_dim = self.velocity_set.d * (self.velocity_set.d + 1) // 2
+        _pi_vec = wp.vec(
+            _pi_dim,
+            dtype=self.compute_dtype,
+        )
+
+        # Construct functional for computing second moment
+        @wp.func
+        def functional(
+            fneq: Any,
+        ):
+            # Get second order moment (a symmetric tensore shaped into a vector)
+            pi = _pi_vec()
+            for d in range(_pi_dim):
+                pi[d] = self.compute_dtype(0.0)
+                for q in range(self.velocity_set.q):
+                    pi[d] += _cc[q, d] * fneq[q]
+            return pi
+
+        # Construct the kernel
+        @wp.kernel
+        def kernel(
+            f: wp.array4d(dtype=Any),
+            pi: wp.array4d(dtype=Any),
+        ):
+            # Get the global index
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            # Get the equilibrium
+            _f = _f_vec()
+            for l in range(self.velocity_set.q):
+                _f[l] = f[l, index[0], index[1], index[2]]
+            _pi = functional(_f)
+
+            # Set the output
+            for d in range(_pi_dim):
+                pi[d, index[0], index[1], index[2]] = self.store_dtype(_pi[d])
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f, pi):
+        # Launch the warp kernel
+        wp.launch(self.warp_kernel, inputs=[f, pi], dim=pi.shape[1:])
+        return pi
diff --git a/xlb/operator/macroscopic/zero_moment.py b/xlb/operator/macroscopic/zero_moment.py
new file mode 100644
index 0000000..8abb4de
--- /dev/null
+++ b/xlb/operator/macroscopic/zero_moment.py
@@ -0,0 +1,49 @@
+from functools import partial
+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.operator import Operator
+
+
+class ZeroMoment(Operator):
+    """A class to compute the zeroth moment (density) of distribution functions."""
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0), inline=True)
+    def jax_implementation(self, f):
+        return jnp.sum(f, axis=0, keepdims=True)
+
+    def _construct_warp(self):
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+
+        @wp.func
+        def functional(f: _f_vec):
+            rho = self.compute_dtype(0.0)
+            for l in range(self.velocity_set.q):
+                rho += f[l]
+            return rho
+
+        @wp.kernel
+        def kernel(
+            f: wp.array4d(dtype=Any),
+            rho: wp.array4d(dtype=Any),
+        ):
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            _f = _f_vec()
+            for l in range(self.velocity_set.q):
+                _f[l] = f[l, index[0], index[1], index[2]]
+            _rho = functional(_f)
+
+            rho[0, index[0], index[1], index[2]] = _rho
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f, rho):
+        wp.launch(self.warp_kernel, inputs=[f, rho], dim=rho.shape[1:])
+        return rho
diff --git a/xlb/operator/operator.py b/xlb/operator/operator.py
new file mode 100644
index 0000000..6e8bbbb
--- /dev/null
+++ b/xlb/operator/operator.py
@@ -0,0 +1,134 @@
+import inspect
+import traceback
+import jax
+
+from xlb.compute_backend import ComputeBackend
+from xlb import DefaultConfig
+from xlb.precision_policy import PrecisionPolicy
+
+
+class Operator:
+    """
+    Base class for all operators, collision, streaming, equilibrium, etc.
+
+    This class is responsible for handling compute backends.
+    """
+
+    _backends = {}
+
+    def __init__(self, velocity_set=None, precision_policy=None, compute_backend=None):
+        # Set the default values from the global config
+        self.velocity_set = velocity_set or DefaultConfig.velocity_set
+        self.precision_policy = precision_policy or DefaultConfig.default_precision_policy
+        self.compute_backend = compute_backend or DefaultConfig.default_backend
+
+        # Check if the compute backend is supported
+        if self.compute_backend not in ComputeBackend:
+            raise ValueError(f"Compute backend {compute_backend} is not supported")
+
+        # Construct the kernel based backend functions TODO: Maybe move this to the register or something
+        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):
+        """
+        Decorator to register a backend for the operator.
+        """
+
+        def decorator(func):
+            subclass_name = func.__qualname__.split(".")[0]
+            signature = inspect.signature(func)
+            key = (subclass_name, backend_name, str(signature))
+            cls._backends[key] = func
+            return func
+
+        return decorator
+
+    def __call__(self, *args, callback=None, **kwargs):
+        method_candidates = [
+            (key, method) for key, method in self._backends.items() if key[0] == self.__class__.__name__ and key[1] == self.compute_backend
+        ]
+        bound_arguments = None
+        for key, backend_method in method_candidates:
+            try:
+                # This attempts to bind the provided args and kwargs to the backend method's signature
+                bound_arguments = inspect.signature(backend_method).bind(self, *args, **kwargs)
+                bound_arguments.apply_defaults()  # This fills in any default values
+                result = backend_method(self, *args, **kwargs)
+                callback_arg = result if result is not None else (args, kwargs)
+                if callback and callable(callback):
+                    callback(callback_arg)
+                return result
+            except Exception as e:
+                error = e
+                traceback_str = traceback.format_exc()
+                continue  # This skips to the next candidate if binding fails
+
+        raise Exception(f"Error captured for backend with key {key} for operator {self.__class__.__name__}: {error}\n {traceback_str}")
+
+    @property
+    def supported_compute_backend(self):
+        """
+        Returns the supported compute backend for the operator
+        """
+        return list(self._backends.keys())
+
+    def _is_method_overridden(self, method_name):
+        """
+        Helper method to check if a method is overridden in a subclass.
+        """
+        method = getattr(self, method_name, None)
+        if method is None:
+            return False
+        return method.__func__ is not getattr(Operator, method_name, None).__func__
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}()"
+
+    @property
+    def backend(self):
+        """
+        Returns the compute backend object for the operator (e.g. jax, warp)
+        This should be used with caution as all backends may not have the same API.
+        """
+        if self.compute_backend == ComputeBackend.JAX:
+            import jax.numpy as backend
+        elif self.compute_backend == ComputeBackend.WARP:
+            import warp as backend
+        return backend
+
+    @property
+    def compute_dtype(self):
+        """
+        Returns the compute dtype
+        """
+        if self.compute_backend == ComputeBackend.JAX:
+            return self.precision_policy.compute_precision.jax_dtype
+        elif self.compute_backend == ComputeBackend.WARP:
+            return self.precision_policy.compute_precision.wp_dtype
+
+    @property
+    def store_dtype(self):
+        """
+        Returns the store dtype
+        """
+        if self.compute_backend == ComputeBackend.JAX:
+            return self.precision_policy.store_precision.jax_dtype
+        elif self.compute_backend == ComputeBackend.WARP:
+            return self.precision_policy.store_precision.wp_dtype
+
+    def _construct_warp(self):
+        """
+        Construct the warp functional and kernel of the operator
+        TODO: Maybe a better way to do this?
+        Maybe add this to the backend decorator?
+        Leave it for now, as it is not clear how the warp backend will evolve
+        """
+        return None, None
diff --git a/xlb/operator/parallel_operator.py b/xlb/operator/parallel_operator.py
new file mode 100644
index 0000000..9f9b5c5
--- /dev/null
+++ b/xlb/operator/parallel_operator.py
@@ -0,0 +1,82 @@
+from jax.experimental.shard_map import shard_map
+from jax.sharding import PartitionSpec as P
+from jax import lax
+
+
+class ParallelOperator:
+    """
+    A generic class for parallelizing operations across multiple GPUs/TPUs.
+    """
+
+    def __init__(self, grid, func, velocity_set):
+        """
+        Initialize the ParallelOperator.
+
+        Parameters
+        ----------
+        grid : Grid object
+            The computational grid.
+        func : function
+            The function to be parallelized.
+        velocity_set : VelocitySet object
+            The velocity set used in the Lattice Boltzmann Method.
+        """
+        self.grid = grid
+        self.func = func
+        self.velocity_set = velocity_set
+
+    def __call__(self, f):
+        """
+        Execute the parallel operation.
+
+        Parameters
+        ----------
+        f : jax.numpy.ndarray
+            The input data for the operation.
+
+        Returns
+        -------
+        jax.numpy.ndarray
+            The result after applying the parallel operation.
+        """
+        in_specs = P(*((None, "x") + (self.grid.dim - 1) * (None,)))
+        out_specs = in_specs
+
+        f = shard_map(
+            self._parallel_func,
+            mesh=self.grid.global_mesh,
+            in_specs=in_specs,
+            out_specs=out_specs,
+            check_rep=False,
+        )(f)
+        return f
+
+    def _parallel_func(self, f):
+        """
+        Internal function to handle data communication and apply the given function in parallel.
+
+        Parameters
+        ----------
+        f : jax.numpy.ndarray
+            The input data.
+
+        Returns
+        -------
+        jax.numpy.ndarray
+            The processed data.
+        """
+        rightPerm = [(i, (i + 1) % self.grid.nDevices) for i in range(self.grid.nDevices)]
+        leftPerm = [((i + 1) % self.grid.nDevices, i) for i in range(self.grid.nDevices)]
+        f = self.func(f)
+        left_comm, right_comm = (
+            f[self.velocity_set.right_indices, :1, ...],
+            f[self.velocity_set.left_indices, -1:, ...],
+        )
+        left_comm, right_comm = (
+            lax.ppermute(left_comm, perm=rightPerm, axis_name="x"),
+            lax.ppermute(right_comm, perm=leftPerm, axis_name="x"),
+        )
+        f = f.at[self.velocity_set.right_indices, :1, ...].set(left_comm)
+        f = f.at[self.velocity_set.left_indices, -1:, ...].set(right_comm)
+
+        return f
diff --git a/xlb/operator/precision_caster/__init__.py b/xlb/operator/precision_caster/__init__.py
new file mode 100644
index 0000000..a027c52
--- /dev/null
+++ b/xlb/operator/precision_caster/__init__.py
@@ -0,0 +1,1 @@
+from xlb.operator.precision_caster.precision_caster import PrecisionCaster
diff --git a/xlb/operator/precision_caster/precision_caster.py b/xlb/operator/precision_caster/precision_caster.py
new file mode 100644
index 0000000..5427cba
--- /dev/null
+++ b/xlb/operator/precision_caster/precision_caster.py
@@ -0,0 +1,96 @@
+"""
+Base class for casting precision of the input data to the desired precision
+"""
+
+import jax.numpy as jnp
+from jax import jit
+import warp as wp
+from functools import partial
+
+from xlb.operator.operator import Operator
+from xlb.velocity_set import VelocitySet
+from xlb.precision_policy import Precision, PrecisionPolicy
+from xlb.compute_backend import ComputeBackend
+
+
+class PrecisionCaster(Operator):
+    """
+    Class that handles the construction of lattice boltzmann precision casting operator.
+    """
+
+    def __init__(
+        self,
+        input_precision: Precision,
+        output_precision: Precision,
+        velocity_set: VelocitySet,
+        precision_policy: PrecisionPolicy,
+        compute_backend: ComputeBackend,
+    ):
+        super().__init__(
+            velocity_set=velocity_set,
+            precision_policy=precision_policy,
+            compute_backend=compute_backend,
+        )
+
+        # Set the input and output precision based on the backend
+        self.input_precision = self._precision_to_dtype(input_precision)
+        self.output_precision = self._precision_to_dtype(output_precision)
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0,))
+    def jax_implementation(self, f: jnp.ndarray) -> jnp.ndarray:
+        return self.output_precision(f)
+
+    def _construct_warp(self):
+        # Construct needed types and constants
+        from_lattice_vec = wp.vec(self.velocity_set.q, dtype=self.input_precision)
+        to_lattice_vec = wp.vec(self.velocity_set.q, dtype=self.output_precision)
+        from_array_type = wp.array4d(dtype=self.input_precision)
+        to_array_type = wp.array4d(dtype=self.output_precision)
+        _q = wp.constant(self.velocity_set.q)
+
+        # Construct the functional
+        @wp.func
+        def functional(
+            from_f: from_lattice_vec,
+        ) -> to_lattice_vec:
+            to_f = to_lattice_vec()
+            for i in range(self.velocity_set.q):
+                to_f[i] = self.output_precision(from_f[i])
+            return to_f
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            from_f: from_array_type,
+            to_f: to_array_type,
+        ):
+            # Get the global index
+            i, j, k = wp.tid()
+
+            # Get f
+            _from_f = from_lattice_vec()
+            for l in range(_q):
+                _from_f[l] = from_f[l, i, j, k]
+
+            # Cast the precision
+            _to_f = functional(_from_f)
+
+            # Set f
+            for l in range(_q):
+                to_f[l, i, j, k] = _to_f[l]
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, from_f, to_f):
+        # Launch the warp kernel
+        wp.launch(
+            self._kernel,
+            inputs=[
+                from_f,
+                to_f,
+            ],
+            dim=from_f.shape[1:],
+        )
+        return to_f
diff --git a/xlb/operator/stepper/__init__.py b/xlb/operator/stepper/__init__.py
new file mode 100644
index 0000000..e5d159c
--- /dev/null
+++ b/xlb/operator/stepper/__init__.py
@@ -0,0 +1,2 @@
+from xlb.operator.stepper.stepper import Stepper
+from xlb.operator.stepper.nse_stepper import IncompressibleNavierStokesStepper
diff --git a/xlb/operator/stepper/nse_stepper.py b/xlb/operator/stepper/nse_stepper.py
new file mode 100644
index 0000000..e08e95c
--- /dev/null
+++ b/xlb/operator/stepper/nse_stepper.py
@@ -0,0 +1,213 @@
+# Base class for all stepper operators
+
+from functools import partial
+from jax import jit
+import warp as wp
+from typing import Any
+
+from xlb import DefaultConfig
+from xlb.compute_backend import ComputeBackend
+from xlb.operator import Operator
+from xlb.operator.stream import Stream
+from xlb.operator.collision import BGK, KBC
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.operator.macroscopic import Macroscopic
+from xlb.operator.stepper import Stepper
+from xlb.operator.boundary_condition.boundary_condition import ImplementationStep
+from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry
+from xlb.operator.collision import ForcedCollision
+
+
+class IncompressibleNavierStokesStepper(Stepper):
+    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
+
+        # Construct the collision operator
+        if collision_type == "BGK":
+            self.collision = BGK(omega, velocity_set, precision_policy, compute_backend)
+        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)
+        self.macroscopic = Macroscopic(velocity_set, precision_policy, compute_backend)
+
+        operators = [self.macroscopic, self.equilibrium, self.collision, self.stream]
+
+        super().__init__(operators, boundary_conditions)
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(self, f_0, f_1, bc_mask, missing_mask, timestep):
+        """
+        Perform a single step of the lattice boltzmann method
+        """
+        # Cast to compute precisioni
+        f_0 = self.precision_policy.cast_to_compute_jax(f_0)
+        f_1 = self.precision_policy.cast_to_compute_jax(f_1)
+
+        # Apply streaming
+        f_post_stream = self.stream(f_0)
+
+        # Apply boundary conditions
+        for bc in self.boundary_conditions:
+            if bc.implementation_step == ImplementationStep.STREAMING:
+                f_post_stream = bc(
+                    f_0,
+                    f_post_stream,
+                    bc_mask,
+                    missing_mask,
+                )
+
+        # Compute the macroscopic variables
+        rho, u = self.macroscopic(f_post_stream)
+
+        # Compute equilibrium
+        feq = self.equilibrium(rho, u)
+
+        # Apply collision
+        f_post_collision = self.collision(f_post_stream, feq, rho, u)
+
+        # 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, bc_mask, missing_mask)
+            if bc.implementation_step == ImplementationStep.COLLISION:
+                f_post_collision = bc(
+                    f_post_stream,
+                    f_post_collision,
+                    bc_mask,
+                    missing_mask,
+                )
+
+        # Copy back to store precision
+        f_1 = self.precision_policy.cast_to_store_jax(f_post_collision)
+
+        return f_0, f_1
+
+    def _construct_warp(self):
+        # Set local constants
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+        _missing_mask_vec = wp.vec(self.velocity_set.q, dtype=wp.uint8)
+        _opp_indices = self.velocity_set.opp_indices
+
+        # Read the list of bc_to_id created upon instantiation
+        bc_to_id = boundary_condition_registry.bc_to_id
+        id_to_bc = boundary_condition_registry.id_to_bc
+
+        # Gather IDs of ExtrapolationOutflowBC boundary conditions
+        extrapolation_outflow_bc_ids = []
+        for bc_name, bc_id in bc_to_id.items():
+            if bc_name.startswith("ExtrapolationOutflowBC"):
+                extrapolation_outflow_bc_ids.append(bc_id)
+        # Group active boundary conditions
+        active_bcs = set(boundary_condition_registry.id_to_bc[bc.id] for bc in self.boundary_conditions)
+
+        @wp.func
+        def apply_bc(
+            index: Any,
+            timestep: Any,
+            _boundary_id: Any,
+            missing_mask: Any,
+            f_0: Any,
+            f_1: Any,
+            f_pre: Any,
+            f_post: Any,
+            is_post_streaming: bool,
+        ):
+            f_result = f_post
+
+            # Unroll the loop over boundary conditions
+            for i in range(wp.static(len(self.boundary_conditions))):
+                if is_post_streaming:
+                    if wp.static(self.boundary_conditions[i].implementation_step == ImplementationStep.STREAMING):
+                        if _boundary_id == wp.static(self.boundary_conditions[i].id):
+                            f_result = wp.static(self.boundary_conditions[i].warp_functional)(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
+                else:
+                    if wp.static(self.boundary_conditions[i].implementation_step == ImplementationStep.COLLISION):
+                        if _boundary_id == wp.static(self.boundary_conditions[i].id):
+                            f_result = wp.static(self.boundary_conditions[i].warp_functional)(index, timestep, missing_mask, f_0, f_1, f_pre, f_post)
+                    if wp.static(self.boundary_conditions[i].id in extrapolation_outflow_bc_ids):
+                        if _boundary_id == wp.static(self.boundary_conditions[i].id):
+                            f_result = wp.static(self.boundary_conditions[i].prepare_bc_auxilary_data)(
+                                index, timestep, missing_mask, f_0, f_1, f_pre, f_post
+                            )
+            return f_result
+
+        @wp.func
+        def get_thread_data(
+            f0_buffer: wp.array4d(dtype=Any),
+            f1_buffer: wp.array4d(dtype=Any),
+            missing_mask: wp.array4d(dtype=Any),
+            index: Any,
+        ):
+            # Read thread data for populations
+            _f0_thread = _f_vec()
+            _f1_thread = _f_vec()
+            _missing_mask = _missing_mask_vec()
+            for l in range(self.velocity_set.q):
+                # q-sized vector of pre-streaming populations
+                _f0_thread[l] = self.compute_dtype(f0_buffer[l, index[0], index[1], index[2]])
+                _f1_thread[l] = self.compute_dtype(f1_buffer[l, index[0], index[1], index[2]])
+                if missing_mask[l, index[0], index[1], index[2]]:
+                    _missing_mask[l] = wp.uint8(1)
+                else:
+                    _missing_mask[l] = wp.uint8(0)
+
+            return _f0_thread, _f1_thread, _missing_mask
+
+        @wp.kernel
+        def kernel(
+            f_0: wp.array4d(dtype=Any),
+            f_1: wp.array4d(dtype=Any),
+            bc_mask: wp.array4d(dtype=Any),
+            missing_mask: wp.array4d(dtype=Any),
+            timestep: int,
+        ):
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            _boundary_id = bc_mask[0, index[0], index[1], index[2]]
+            if _boundary_id == wp.uint8(255):
+                return
+
+            # Apply streaming
+            _f_post_stream = self.stream.warp_functional(f_0, index)
+
+            _f0_thread, _f1_thread, _missing_mask = get_thread_data(f_0, f_1, missing_mask, index)
+            _f_post_collision = _f0_thread
+
+            # Apply post-streaming boundary conditions
+            _f_post_stream = apply_bc(index, timestep, _boundary_id, _missing_mask, f_0, f_1, _f_post_collision, _f_post_stream, True)
+
+            _rho, _u = self.macroscopic.warp_functional(_f_post_stream)
+            _feq = self.equilibrium.warp_functional(_rho, _u)
+            _f_post_collision = self.collision.warp_functional(_f_post_stream, _feq, _rho, _u)
+
+            # Apply post-collision boundary conditions
+            _f_post_collision = apply_bc(index, timestep, _boundary_id, _missing_mask, f_0, f_1, _f_post_stream, _f_post_collision, False)
+
+            # Store the result in f_1
+            for l in range(self.velocity_set.q):
+                # TODO: Improve this later
+                if wp.static("GradsApproximationBC" in active_bcs):
+                    if _boundary_id == wp.static(boundary_condition_registry.bc_to_id["GradsApproximationBC"]):
+                        if _missing_mask[l] == wp.uint8(1):
+                            f_0[_opp_indices[l], index[0], index[1], index[2]] = self.store_dtype(_f1_thread[_opp_indices[l]])
+                f_1[l, index[0], index[1], index[2]] = self.store_dtype(_f_post_collision[l])
+
+        return None, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f_0, f_1, bc_mask, missing_mask, timestep):
+        wp.launch(
+            self.warp_kernel,
+            inputs=[f_0, f_1, bc_mask, missing_mask, timestep],
+            dim=f_0.shape[1:],
+        )
+        return f_0, f_1
diff --git a/xlb/operator/stepper/stepper.py b/xlb/operator/stepper/stepper.py
new file mode 100644
index 0000000..44aab5f
--- /dev/null
+++ b/xlb/operator/stepper/stepper.py
@@ -0,0 +1,32 @@
+# Base class for all stepper operators
+from xlb.operator import Operator
+from xlb import DefaultConfig
+
+
+class Stepper(Operator):
+    """
+    Class that handles the construction of lattice boltzmann stepping operator
+    """
+
+    def __init__(self, operators, boundary_conditions):
+        # Get the boundary condition ids
+        from xlb.operator.boundary_condition.boundary_condition_registry import boundary_condition_registry
+
+        self.operators = operators
+        self.boundary_conditions = boundary_conditions
+
+        # Get velocity set, precision policy, and compute backend
+        velocity_sets = set([op.velocity_set for op in self.operators if op is not None])
+        assert len(velocity_sets) < 2, "All velocity sets must be the same. Got {}".format(velocity_sets)
+        velocity_set = DefaultConfig.velocity_set if not velocity_sets else velocity_sets.pop()
+
+        precision_policies = set([op.precision_policy for op in self.operators if op is not None])
+        assert len(precision_policies) < 2, "All precision policies must be the same. Got {}".format(precision_policies)
+        precision_policy = DefaultConfig.default_precision_policy if not precision_policies else precision_policies.pop()
+
+        compute_backends = set([op.compute_backend for op in self.operators if op is not None])
+        assert len(compute_backends) < 2, "All compute backends must be the same. Got {}".format(compute_backends)
+        compute_backend = DefaultConfig.default_backend if not compute_backends else compute_backends.pop()
+
+        # Initialize operator
+        super().__init__(velocity_set, precision_policy, compute_backend)
diff --git a/xlb/operator/stream/__init__.py b/xlb/operator/stream/__init__.py
new file mode 100644
index 0000000..9093da7
--- /dev/null
+++ b/xlb/operator/stream/__init__.py
@@ -0,0 +1,1 @@
+from xlb.operator.stream.stream import Stream
diff --git a/xlb/operator/stream/stream.py b/xlb/operator/stream/stream.py
new file mode 100644
index 0000000..247fa5a
--- /dev/null
+++ b/xlb/operator/stream/stream.py
@@ -0,0 +1,114 @@
+# Base class for all streaming operators
+
+from functools import partial
+import jax.numpy as jnp
+from jax import jit, vmap
+import warp as wp
+from typing import Any
+
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.operator import Operator
+
+
+class Stream(Operator):
+    """
+    Base class for all streaming operators. This is used for pulling the distribution
+    """
+
+    @Operator.register_backend(ComputeBackend.JAX)
+    @partial(jit, static_argnums=(0))
+    def jax_implementation(self, f):
+        """
+        JAX implementation of the streaming step.
+
+        TODO: Make sure this works with pull scheme.
+
+        Parameters
+        ----------
+        f: jax.numpy.ndarray
+            The distribution function.
+        """
+
+        def _streaming_jax_i(f, c):
+            """
+            Perform individual streaming operation in a direction.
+
+            Parameters
+            ----------
+                f: The distribution function.
+                c: The streaming direction vector.
+
+            Returns
+            -------
+                jax.numpy.ndarray
+                The updated distribution function after streaming.
+            """
+            if self.velocity_set.d == 2:
+                return jnp.roll(f, (c[0], c[1]), axis=(0, 1))
+            elif self.velocity_set.d == 3:
+                return jnp.roll(f, (c[0], c[1], c[2]), axis=(0, 1, 2))
+
+        return vmap(_streaming_jax_i, in_axes=(0, 0), out_axes=0)(f, jnp.array(self.velocity_set.c).T)
+
+    def _construct_warp(self):
+        # Set local constants TODO: This is a hack and should be fixed with warp update
+        _c = self.velocity_set.c
+        _f_vec = wp.vec(self.velocity_set.q, dtype=self.compute_dtype)
+
+        # Construct the funcional to get streamed indices
+        @wp.func
+        def functional(
+            f: wp.array4d(dtype=Any),
+            index: Any,
+        ):
+            # Pull the distribution function
+            _f = _f_vec()
+            for l in range(self.velocity_set.q):
+                # Get pull index
+                pull_index = type(index)()
+                for d in range(self.velocity_set.d):
+                    pull_index[d] = index[d] - _c[d, l]
+
+                    # impose periodicity for out of bound values
+                    if pull_index[d] < 0:
+                        pull_index[d] = f.shape[d + 1] - 1
+                    elif pull_index[d] >= f.shape[d + 1]:
+                        pull_index[d] = 0
+
+                # Read the distribution function
+                # Unlike other functionals, we need to cast the type here since we read from the buffer
+                _f[l] = self.compute_dtype(f[l, pull_index[0], pull_index[1], pull_index[2]])
+
+            return _f
+
+        # Construct the warp kernel
+        @wp.kernel
+        def kernel(
+            f_0: wp.array4d(dtype=Any),
+            f_1: wp.array4d(dtype=Any),
+        ):
+            # Get the global index
+            i, j, k = wp.tid()
+            index = wp.vec3i(i, j, k)
+
+            # Set the output
+            _f = functional(f_0, index)
+
+            # Write the output
+            for l in range(self.velocity_set.q):
+                f_1[l, index[0], index[1], index[2]] = self.store_dtype(_f[l])
+
+        return functional, kernel
+
+    @Operator.register_backend(ComputeBackend.WARP)
+    def warp_implementation(self, f_0, f_1):
+        # Launch the warp kernel
+        wp.launch(
+            self.warp_kernel,
+            inputs=[
+                f_0,
+                f_1,
+            ],
+            dim=f_0.shape[1:],
+        )
+        return f_1
diff --git a/xlb/physics_type.py b/xlb/physics_type.py
new file mode 100644
index 0000000..39841fc
--- /dev/null
+++ b/xlb/physics_type.py
@@ -0,0 +1,8 @@
+# Enum used to keep track of the physics types supported by different operators
+
+from enum import Enum, auto
+
+
+class PhysicsType(Enum):
+    NSE = auto()  # Navier-Stokes Equations
+    ADE = auto()  # Advection-Diffusion Equations
diff --git a/xlb/precision_policy.py b/xlb/precision_policy.py
new file mode 100644
index 0000000..7d31c8a
--- /dev/null
+++ b/xlb/precision_policy.py
@@ -0,0 +1,89 @@
+# Enum for precision policy
+
+from enum import Enum, auto
+import jax.numpy as jnp
+import warp as wp
+
+
+class Precision(Enum):
+    FP64 = auto()
+    FP32 = auto()
+    FP16 = auto()
+    UINT8 = auto()
+    BOOL = auto()
+
+    @property
+    def wp_dtype(self):
+        if self == Precision.FP64:
+            return wp.float64
+        elif self == Precision.FP32:
+            return wp.float32
+        elif self == Precision.FP16:
+            return wp.float16
+        elif self == Precision.UINT8:
+            return wp.uint8
+        elif self == Precision.BOOL:
+            return wp.bool
+        else:
+            raise ValueError("Invalid precision")
+
+    @property
+    def jax_dtype(self):
+        if self == Precision.FP64:
+            return jnp.float64
+        elif self == Precision.FP32:
+            return jnp.float32
+        elif self == Precision.FP16:
+            return jnp.float16
+        elif self == Precision.UINT8:
+            return jnp.uint8
+        elif self == Precision.BOOL:
+            return jnp.bool_
+        else:
+            raise ValueError("Invalid precision")
+
+
+class PrecisionPolicy(Enum):
+    FP64FP64 = auto()
+    FP64FP32 = auto()
+    FP64FP16 = auto()
+    FP32FP32 = auto()
+    FP32FP16 = auto()
+
+    @property
+    def compute_precision(self):
+        if self == PrecisionPolicy.FP64FP64:
+            return Precision.FP64
+        elif self == PrecisionPolicy.FP64FP32:
+            return Precision.FP64
+        elif self == PrecisionPolicy.FP64FP16:
+            return Precision.FP64
+        elif self == PrecisionPolicy.FP32FP32:
+            return Precision.FP32
+        elif self == PrecisionPolicy.FP32FP16:
+            return Precision.FP32
+        else:
+            raise ValueError("Invalid precision policy")
+
+    @property
+    def store_precision(self):
+        if self == PrecisionPolicy.FP64FP64:
+            return Precision.FP64
+        elif self == PrecisionPolicy.FP64FP32:
+            return Precision.FP32
+        elif self == PrecisionPolicy.FP64FP16:
+            return Precision.FP16
+        elif self == PrecisionPolicy.FP32FP32:
+            return Precision.FP32
+        elif self == PrecisionPolicy.FP32FP16:
+            return Precision.FP16
+        else:
+            raise ValueError("Invalid precision policy")
+
+    def cast_to_compute_jax(self, array):
+        compute_precision = self.compute_precision
+        return jnp.array(array, dtype=compute_precision.jax_dtype)
+
+    def cast_to_store_jax(self, array):
+        store_precision = self.store_precision
+        return jnp.array(array, dtype=store_precision.jax_dtype)
diff --git a/xlb/precision_policy/precision_policy.py b/xlb/precision_policy/precision_policy.py
new file mode 100644
index 0000000..6b400ee
--- /dev/null
+++ b/xlb/precision_policy/precision_policy.py
@@ -0,0 +1,49 @@
+from xlb.compute_backend import ComputeBackend
+from xlb import DefaultConfig
+from xlb.precision_policy.jax_precision_policy import (
+    JaxFp32Fp32,
+    JaxFp32Fp16,
+    JaxFp64Fp64,
+    JaxFp64Fp32,
+    JaxFp64Fp16,
+)
+
+
+class Fp64Fp64:
+    def __new__(cls):
+        if DefaultConfig.compute_backend == ComputeBackend.JAX:
+            return JaxFp64Fp64()
+        else:
+            raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}")
+
+
+class Fp64Fp32:
+    def __new__(cls):
+        if DefaultConfig.compute_backend == ComputeBackend.JAX:
+            return JaxFp64Fp32()
+        else:
+            raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}")
+
+
+class Fp32Fp32:
+    def __new__(cls):
+        if DefaultConfig.compute_backend == ComputeBackend.JAX:
+            return JaxFp32Fp32()
+        else:
+            raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}")
+
+
+class Fp64Fp16:
+    def __new__(cls):
+        if DefaultConfig.compute_backend == ComputeBackend.JAX:
+            return JaxFp64Fp16()
+        else:
+            raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}")
+
+
+class Fp32Fp16:
+    def __new__(cls):
+        if DefaultConfig.compute_backend == ComputeBackend.JAX:
+            return JaxFp32Fp16()
+        else:
+            raise ValueError(f"Unsupported compute backend: {DefaultConfig.compute_backend}")
diff --git a/xlb/utils/__init__.py b/xlb/utils/__init__.py
new file mode 100644
index 0000000..3c8032e
--- /dev/null
+++ b/xlb/utils/__init__.py
@@ -0,0 +1,9 @@
+from .utils import (
+    downsample_field,
+    save_image,
+    save_fields_vtk,
+    save_BCs_vtk,
+    rotate_geometry,
+    voxelize_stl,
+    axangle2mat,
+)
diff --git a/src/utils.py b/xlb/utils/utils.py
similarity index 64%
rename from src/utils.py
rename to xlb/utils/utils.py
index 1720ca6..0a9858a 100644
--- a/src/utils.py
+++ b/xlb/utils/utils.py
@@ -1,7 +1,6 @@
 import numpy as np
 import matplotlib.pylab as plt
 from matplotlib import cm
-import numpy as np
 from time import time
 import pyvista as pv
 from jax.image import resize
@@ -15,7 +14,7 @@
 
 
 @partial(jit, static_argnums=(1, 2))
-def downsample_field(field, factor, method='bicubic'):
+def downsample_field(field, factor, method="bicubic"):
     """
     Downsample a JAX array by a factor of `factor` along each axis.
 
@@ -44,7 +43,8 @@ def downsample_field(field, factor, method='bicubic'):
 
         return jnp.stack(downsampled_components, axis=-1)
 
-def save_image(timestep, fld, prefix=None):
+
+def save_image(fld, timestep=None, prefix=None, **kwargs):
     """
     Save an image of a field at a given timestep.
 
@@ -63,31 +63,38 @@ def save_image(timestep, fld, prefix=None):
 
     Notes
     -----
-    This function saves the field as an image in the PNG format. The filename is based on the name of the main script file, the provided prefix, and the timestep number.
-    If the field is 3D, the magnitude of the field is calculated and saved. The image is saved with the 'nipy_spectral' colormap and the origin set to 'lower'.
+    This function saves the field as an image in the PNG format.
+    The filename is based on the name of the main script file, the provided prefix, and the timestep number.
+    If the field is 3D, the magnitude of the field is calculated and saved.
+    The image is saved with the 'nipy_spectral' colormap and the origin set to 'lower'.
     """
-    fname = os.path.basename(__main__.__file__)
-    fname = os.path.splitext(fname)[0]
-    if prefix is not None:
-        fname = prefix + fname
-    fname = fname + "_" + str(timestep).zfill(4)
+    if prefix is None:
+        fname = os.path.basename(__main__.__file__)
+        fname = os.path.splitext(fname)[0]
+    else:
+        fname = prefix
+
+    if timestep is not None:
+        fname = fname + "_" + str(timestep).zfill(4)
 
     if len(fld.shape) > 3:
         raise ValueError("The input field should be 2D!")
-    elif len(fld.shape) == 3:
-        fld = np.sqrt(fld[..., 0] ** 2 + fld[..., 1] ** 2)
+    if len(fld.shape) == 3:
+        fld = np.sqrt(fld[0, ...] ** 2 + fld[0, ...] ** 2)
 
     plt.clf()
-    plt.imsave(fname + '.png', fld.T, cmap=cm.nipy_spectral, origin='lower')
+    kwargs.pop("cmap", None)
+    plt.imsave(fname + ".png", fld.T, cmap=cm.nipy_spectral, origin="lower", **kwargs)
 
-def save_fields_vtk(timestep, fields, output_dir='.', prefix='fields'):
+
+def save_fields_vtk(fields, timestep, output_dir=".", prefix="fields"):
     """
     Save VTK fields to the specified directory.
 
     Parameters
     ----------
     timestep (int): The timestep number to be associated with the saved fields.
-    fields (Dict[str, np.ndarray]): A dictionary of fields to be saved. Each field must be an array-like object 
+    fields (Dict[str, np.ndarray]): A dictionary of fields to be saved. Each field must be an array-like object
         with dimensions (nx, ny) for 2D fields or (nx, ny, nz) for 3D fields, where:
             - nx : int, number of grid points along the x-axis
             - ny : int, number of grid points along the y-axis
@@ -107,14 +114,14 @@ def save_fields_vtk(timestep, fields, output_dir='.', prefix='fields'):
     will be saved as 'fields_0000010.vtk'in the specified directory.
 
     """
-    # Assert that all fields have the same dimensions except for the last dimension assuming fields is a dictionary
+    # Assert that all fields have the same dimensions
     for key, value in fields.items():
         if key == list(fields.keys())[0]:
             dimensions = value.shape
         else:
             assert value.shape == dimensions, "All fields must have the same dimensions!"
 
-    output_filename = os.path.join(output_dir, prefix +  "_" + f"{timestep:07d}.vtk")
+    output_filename = os.path.join(output_dir, prefix + "_" + f"{timestep:07d}.vtk")
 
     # Add 1 to the dimensions tuple as we store cell values
     dimensions = tuple([dim + 1 for dim in dimensions])
@@ -127,59 +134,15 @@ def save_fields_vtk(timestep, fields, output_dir='.', prefix='fields'):
 
     # Add the fields to the grid
     for key, value in fields.items():
-        grid[key] = value.flatten(order='F')
+        grid[key] = value.flatten(order="F")
 
     # Save the grid to a VTK file
     start = time()
     grid.save(output_filename, binary=True)
     print(f"Saved {output_filename} in {time() - start:.6f} seconds.")
 
-def live_volume_randering(timestep, field):
-    # WORK IN PROGRESS
-    """
-    Live rendering of a 3D volume using pyvista.
-
-    Parameters
-    ----------
-    field (np.ndarray): A 3D array containing the field to be rendered.
-
-    Returns
-    -------
-    None
-
-    Notes
-    -----
-    This function uses pyvista to render a 3D volume. The volume is rendered with a colormap based on the field values.
-    The colormap is updated every 0.1 seconds to reflect changes to the field.
 
-    """
-    # Create a uniform grid (Note that the field must be 3D) otherwise raise error
-    if field.ndim != 3:
-        raise ValueError("The input field must be 3D!")
-    dimensions = field.shape
-    grid = pv.ImageData(dimensions=dimensions)
-
-    # Add the field to the grid
-    grid['field'] = field.flatten(order='F')
-
-    # Create the rendering scene
-    if timestep == 0:
-        plt.ion()
-        plt.figure(figsize=(10, 10))
-        plt.axis('off')
-        plt.title("Live rendering of the field")
-        pl = pv.Plotter(off_screen=True)
-        pl.add_volume(grid, cmap='nipy_spectral', opacity='sigmoid_10', shade=False)
-        plt.imshow(pl.screenshot())
-
-    else:
-        pl = pv.Plotter(off_screen=True)
-        pl.add_volume(grid, cmap='nipy_spectral', opacity='sigmoid_10', shade=False)
-        # Update the rendering scene every 0.1 seconds
-        plt.imshow(pl.screenshot())
-        plt.pause(0.1)
-
-def save_BCs_vtk(timestep, BCs, gridInfo,  output_dir='.'):
+def save_BCs_vtk(timestep, BCs, gridInfo, output_dir="."):
     """
     Save boundary conditions as VTK format to the specified directory.
 
@@ -200,12 +163,12 @@ def save_BCs_vtk(timestep, BCs, gridInfo,  output_dir='.'):
     """
 
     # Create a uniform grid
-    if gridInfo['nz'] == 0:
-        gridDimensions = (gridInfo['nx'] + 1, gridInfo['ny'] + 1, 1)
-        fieldDimensions = (gridInfo['nx'], gridInfo['ny'], 1)
+    if gridInfo["nz"] == 0:
+        gridDimensions = (gridInfo["nx"] + 1, gridInfo["ny"] + 1, 1)
+        fieldDimensions = (gridInfo["nx"], gridInfo["ny"], 1)
     else:
-        gridDimensions = (gridInfo['nx'] + 1, gridInfo['ny'] + 1, gridInfo['nz'] + 1)
-        fieldDimensions = (gridInfo['nx'], gridInfo['ny'], gridInfo['nz'])
+        gridDimensions = (gridInfo["nx"] + 1, gridInfo["ny"] + 1, gridInfo["nz"] + 1)
+        fieldDimensions = (gridInfo["nx"], gridInfo["ny"], gridInfo["nz"])
 
     grid = pv.ImageData(dimensions=gridDimensions)
 
@@ -226,16 +189,16 @@ def save_BCs_vtk(timestep, BCs, gridInfo,  output_dir='.'):
             bcIndices = bc.indices
 
         # Convert indices to 1D indices
-        if gridInfo['dim'] == 2:
-            bcIndices = np.ravel_multi_index(bcIndices, fieldDimensions[:-1], order='F')
+        if gridInfo["dim"] == 2:
+            bcIndices = np.ravel_multi_index(bcIndices, fieldDimensions[:-1], order="F")
         else:
-            bcIndices = np.ravel_multi_index(bcIndices, fieldDimensions, order='F')
+            bcIndices = np.ravel_multi_index(bcIndices, fieldDimensions, order="F")
 
-        grid[bcName] = np.zeros(fieldDimensions, dtype=bool).flatten(order='F')
+        grid[bcName] = np.zeros(fieldDimensions, dtype=bool).flatten(order="F")
         grid[bcName][bcIndices] = True
 
     # Save the grid to a VTK file
-    output_filename = os.path.join(output_dir,  "BCs_" + f"{timestep:07d}.vtk")
+    output_filename = os.path.join(output_dir, "BCs_" + f"{timestep:07d}.vtk")
 
     start = time()
     grid.save(output_filename, binary=True)
@@ -268,7 +231,8 @@ def rotate_geometry(indices, origin, axis, angle):
     using the axis-angle representation of rotations. The origin of the rotation axis is assumed to be at (0, 0, 0).
     """
     indices_rotated = (jnp.array(indices).T - origin) @ axangle2mat(axis, angle) + origin
-    return tuple(jnp.rint(indices_rotated).astype('int32').T)
+    return tuple(jnp.rint(indices_rotated).astype("int32").T)
+
 
 def voxelize_stl(stl_filename, length_lbm_unit=None, tranformation_matrix=None, pitch=None):
     """
@@ -309,7 +273,7 @@ def voxelize_stl(stl_filename, length_lbm_unit=None, tranformation_matrix=None,
 
 
 def axangle2mat(axis, angle, is_normalized=False):
-    ''' Rotation matrix for rotation angle `angle` around `axis`
+    """Rotation matrix for rotation angle `angle` around `axis`
     Parameters
     ----------
     axis : 3 element sequence
@@ -326,7 +290,7 @@ def axangle2mat(axis, angle, is_normalized=False):
     -----
     From : https://github.com/matthew-brett/transforms3d
     Ref : http://en.wikipedia.org/wiki/Rotation_matrix#Axis_and_angle
-    '''
+    """
     x, y, z = axis
     if not is_normalized:
         n = jnp.sqrt(x * x + y * y + z * z)
@@ -348,67 +312,5 @@ def axangle2mat(axis, angle, is_normalized=False):
     return jnp.array([
         [x * xC + c, xyC - zs, zxC + ys],
         [xyC + zs, y * yC + c, yzC - xs],
-        [zxC - ys, yzC + xs, z * zC + c]])
-
-@partial(jit)
-def q_criterion(u):
-    # Compute derivatives
-    u_x = u[..., 0]
-    u_y = u[..., 1]
-    u_z = u[..., 2]
-
-    # Compute derivatives
-    u_x_dx = (u_x[2:, 1:-1, 1:-1] - u_x[:-2, 1:-1, 1:-1]) / 2
-    u_x_dy = (u_x[1:-1, 2:, 1:-1] - u_x[1:-1, :-2, 1:-1]) / 2
-    u_x_dz = (u_x[1:-1, 1:-1, 2:] - u_x[1:-1, 1:-1, :-2]) / 2
-    u_y_dx = (u_y[2:, 1:-1, 1:-1] - u_y[:-2, 1:-1, 1:-1]) / 2
-    u_y_dy = (u_y[1:-1, 2:, 1:-1] - u_y[1:-1, :-2, 1:-1]) / 2
-    u_y_dz = (u_y[1:-1, 1:-1, 2:] - u_y[1:-1, 1:-1, :-2]) / 2
-    u_z_dx = (u_z[2:, 1:-1, 1:-1] - u_z[:-2, 1:-1, 1:-1]) / 2
-    u_z_dy = (u_z[1:-1, 2:, 1:-1] - u_z[1:-1, :-2, 1:-1]) / 2
-    u_z_dz = (u_z[1:-1, 1:-1, 2:] - u_z[1:-1, 1:-1, :-2]) / 2
-
-    # Compute vorticity
-    mu_x = u_z_dy - u_y_dz
-    mu_y = u_x_dz - u_z_dx
-    mu_z = u_y_dx - u_x_dy
-    norm_mu = jnp.sqrt(mu_x ** 2 + mu_y ** 2 + mu_z ** 2)
-
-    # Compute strain rate
-    s_0_0 = u_x_dx
-    s_0_1 = 0.5 * (u_x_dy + u_y_dx)
-    s_0_2 = 0.5 * (u_x_dz + u_z_dx)
-    s_1_0 = s_0_1
-    s_1_1 = u_y_dy
-    s_1_2 = 0.5 * (u_y_dz + u_z_dy)
-    s_2_0 = s_0_2
-    s_2_1 = s_1_2
-    s_2_2 = u_z_dz
-    s_dot_s = (
-        s_0_0 ** 2 + s_0_1 ** 2 + s_0_2 ** 2 +
-        s_1_0 ** 2 + s_1_1 ** 2 + s_1_2 ** 2 +
-        s_2_0 ** 2 + s_2_1 ** 2 + s_2_2 ** 2
-    )
-
-    # Compute omega
-    omega_0_0 = 0.0
-    omega_0_1 = 0.5 * (u_x_dy - u_y_dx)
-    omega_0_2 = 0.5 * (u_x_dz - u_z_dx)
-    omega_1_0 = -omega_0_1
-    omega_1_1 = 0.0
-    omega_1_2 = 0.5 * (u_y_dz - u_z_dy)
-    omega_2_0 = -omega_0_2
-    omega_2_1 = -omega_1_2
-    omega_2_2 = 0.0
-    omega_dot_omega = (
-        omega_0_0 ** 2 + omega_0_1 ** 2 + omega_0_2 ** 2 +
-        omega_1_0 ** 2 + omega_1_1 ** 2 + omega_1_2 ** 2 +
-        omega_2_0 ** 2 + omega_2_1 ** 2 + omega_2_2 ** 2
-    )
-
-    # Compute q-criterion
-    q = 0.5 * (omega_dot_omega - s_dot_s)
-
-    return norm_mu, q
-
-
+        [zxC - ys, yzC + xs, z * zC + c],
+    ])
diff --git a/xlb/velocity_set/__init__.py b/xlb/velocity_set/__init__.py
new file mode 100644
index 0000000..5b7b737
--- /dev/null
+++ b/xlb/velocity_set/__init__.py
@@ -0,0 +1,4 @@
+from xlb.velocity_set.velocity_set import VelocitySet
+from xlb.velocity_set.d2q9 import D2Q9
+from xlb.velocity_set.d3q19 import D3Q19
+from xlb.velocity_set.d3q27 import D3Q27
diff --git a/xlb/velocity_set/d2q9.py b/xlb/velocity_set/d2q9.py
new file mode 100644
index 0000000..5324618
--- /dev/null
+++ b/xlb/velocity_set/d2q9.py
@@ -0,0 +1,24 @@
+# Description: Lattice class for 2D D2Q9 lattice.
+
+import numpy as np
+
+from xlb.velocity_set.velocity_set import VelocitySet
+
+
+class D2Q9(VelocitySet):
+    """
+    Velocity Set for 2D D2Q9 lattice.
+
+    D2Q9 stands for two-dimensional nine-velocity model. It is a common model used in the
+    Lattice Boltzmann Method for simulating fluid flows in two dimensions.
+    """
+
+    def __init__(self, precision_policy, backend):
+        # Construct the velocity vectors and weights
+        cx = [0, 0, 0, 1, -1, 1, -1, 1, -1]
+        cy = [0, 1, -1, 0, 1, -1, 0, 1, -1]
+        c = np.array(tuple(zip(cx, cy))).T
+        w = np.array([4 / 9, 1 / 9, 1 / 9, 1 / 9, 1 / 36, 1 / 36, 1 / 9, 1 / 36, 1 / 36])
+
+        # Call the parent constructor
+        super().__init__(2, 9, c, w, precision_policy=precision_policy, backend=backend)
diff --git a/xlb/velocity_set/d3q19.py b/xlb/velocity_set/d3q19.py
new file mode 100644
index 0000000..4a48c2f
--- /dev/null
+++ b/xlb/velocity_set/d3q19.py
@@ -0,0 +1,30 @@
+# Description: Lattice class for 3D D3Q19 lattice.
+
+import itertools
+import numpy as np
+
+from xlb.velocity_set.velocity_set import VelocitySet
+
+
+class D3Q19(VelocitySet):
+    """
+    Velocity Set for 3D D3Q19 lattice.
+
+    D3Q19 stands for three-dimensional nineteen-velocity model. It is a common model used in the
+    Lattice Boltzmann Method for simulating fluid flows in three dimensions.
+    """
+
+    def __init__(self, precision_policy, backend):
+        # Construct the velocity vectors and weights
+        c = np.array([ci for ci in itertools.product([-1, 0, 1], repeat=3) if np.sum(np.abs(ci)) <= 2]).T
+        w = np.zeros(19)
+        for i in range(19):
+            if np.sum(np.abs(c[:, i])) == 0:
+                w[i] = 1.0 / 3.0
+            elif np.sum(np.abs(c[:, i])) == 1:
+                w[i] = 1.0 / 18.0
+            elif np.sum(np.abs(c[:, i])) == 2:
+                w[i] = 1.0 / 36.0
+
+        # Initialize the lattice
+        super().__init__(3, 19, c, w, precision_policy=precision_policy, backend=backend)
diff --git a/xlb/velocity_set/d3q27.py b/xlb/velocity_set/d3q27.py
new file mode 100644
index 0000000..b056d53
--- /dev/null
+++ b/xlb/velocity_set/d3q27.py
@@ -0,0 +1,32 @@
+# Description: Lattice class for 3D D3Q27 lattice.
+
+import itertools
+import numpy as np
+
+from xlb.velocity_set.velocity_set import VelocitySet
+
+
+class D3Q27(VelocitySet):
+    """
+    Velocity Set for 3D D3Q27 lattice.
+
+    D3Q27 stands for three-dimensional twenty-seven-velocity model. It is a common model used in the
+    Lattice Boltzmann Method for simulating fluid flows in three dimensions.
+    """
+
+    def __init__(self, precision_policy, backend):
+        # Construct the velocity vectors and weights
+        c = np.array(list(itertools.product([0, -1, 1], repeat=3))).T
+        w = np.zeros(27)
+        for i in range(27):
+            if np.sum(np.abs(c[:, i])) == 0:
+                w[i] = 8.0 / 27.0
+            elif np.sum(np.abs(c[:, i])) == 1:
+                w[i] = 2.0 / 27.0
+            elif np.sum(np.abs(c[:, i])) == 2:
+                w[i] = 1.0 / 54.0
+            elif np.sum(np.abs(c[:, i])) == 3:
+                w[i] = 1.0 / 216.0
+
+        # Initialize the Lattice
+        super().__init__(3, 27, c, w, precision_policy=precision_policy, backend=backend)
diff --git a/xlb/velocity_set/velocity_set.py b/xlb/velocity_set/velocity_set.py
new file mode 100644
index 0000000..41db991
--- /dev/null
+++ b/xlb/velocity_set/velocity_set.py
@@ -0,0 +1,233 @@
+# Base Velocity Set class
+
+import math
+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):
+    """
+    Base class for the velocity set of the Lattice Boltzmann Method (LBM), e.g. D2Q9, D3Q27, etc.
+
+    Parameters
+    ----------
+    d: int
+        The dimension of the lattice.
+    q: int
+        The number of velocities of the lattice.
+    c: numpy.ndarray
+        The velocity vectors of the lattice. Shape: (q, d)
+    w: numpy.ndarray
+        The weights of the lattice. Shape: (q,)
+    """
+
+    def __init__(self, d, q, c, w, precision_policy, backend):
+        # Store the dimension and the number of velocities
+        self.d = d
+        self.q = q
+        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)
+
+        # Convert properties to backend-specific format
+        if self.backend == ComputeBackend.WARP:
+            self._init_warp_properties()
+        elif self.backend == ComputeBackend.JAX:
+            self._init_jax_properties()
+        else:
+            raise ValueError(f"Unsupported compute backend: {self.backend}")
+
+        # Set up backend-specific constants
+        self._init_backend_constants()
+
+    def _init_numpy_properties(self, c, w):
+        """
+        Initialize all properties in NumPy first.
+        """
+        self._c = np.array(c)
+        self._w = np.array(w)
+        self._opp_indices = self._construct_opposite_indices()
+        self._cc = self._construct_lattice_moment()
+        self._c_float = self._c.astype(np.float64)
+        self._qi = self._construct_qi()
+
+        # Constants in NumPy
+        self.cs = np.float64(math.sqrt(3) / 3.0)
+        self.cs2 = np.float64(1.0 / 3.0)
+        self.inv_cs2 = np.float64(3.0)
+
+        # Indices
+        self.main_indices = self._construct_main_indices()
+        self.right_indices = self._construct_right_indices()
+        self.left_indices = self._construct_left_indices()
+
+    def _init_warp_properties(self):
+        """
+        Convert NumPy properties to Warp-specific properties.
+        """
+        dtype = self.precision_policy.compute_precision.wp_dtype
+        self.c = wp.constant(wp.mat((self.d, self.q), dtype=wp.int32)(self._c))
+        self.w = wp.constant(wp.vec(self.q, dtype=dtype)(self._w))
+        self.opp_indices = wp.constant(wp.vec(self.q, dtype=wp.int32)(self._opp_indices))
+        self.cc = wp.constant(wp.mat((self.q, self.d * (self.d + 1) // 2), dtype=dtype)(self._cc))
+        self.c_float = wp.constant(wp.mat((self.d, self.q), dtype=dtype)(self._c_float))
+        self.qi = wp.constant(wp.mat((self.q, self.d * (self.d + 1) // 2), dtype=dtype)(self._qi))
+
+    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=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)
+        self.c_float = jnp.array(self._c_float, dtype=dtype)
+        self.qi = jnp.array(self._qi, dtype=dtype)
+
+    def _init_backend_constants(self):
+        """
+        Initialize the constants for the backend.
+        """
+        if self.backend == ComputeBackend.WARP:
+            dtype = self.precision_policy.compute_precision.wp_dtype
+            self.cs = wp.constant(dtype(self.cs))
+            self.cs2 = wp.constant(dtype(self.cs2))
+            self.inv_cs2 = wp.constant(dtype(self.inv_cs2))
+        elif self.backend == ComputeBackend.JAX:
+            dtype = self.precision_policy.compute_precision.jax_dtype
+            self.cs = jnp.array(self.cs, dtype=dtype)
+            self.cs2 = jnp.array(self.cs2, dtype=dtype)
+            self.inv_cs2 = jnp.array(self.inv_cs2, dtype=dtype)
+
+    def warp_lattice_vec(self, dtype):
+        return wp.vec(len(self.c), dtype=dtype)
+
+    def warp_u_vec(self, dtype):
+        return wp.vec(self.d, dtype=dtype)
+
+    def warp_stream_mat(self, dtype):
+        return wp.mat((self.q, self.d), dtype=dtype)
+
+    def _construct_qi(self):
+        # Qi = cc - cs^2*I
+        dim = self.d
+        Qi = self._cc.copy()
+        if dim == 3:
+            diagonal, offdiagonal = (0, 3, 5), (1, 2, 4)
+        elif dim == 2:
+            diagonal, offdiagonal = (0, 2), (1,)
+        else:
+            raise ValueError(f"dim = {dim} not supported")
+
+        # multiply off-diagonal elements by 2 because the Q tensor is symmetric
+        Qi[:, diagonal] += -1.0 / 3.0
+        Qi[:, offdiagonal] *= 2.0
+        return Qi
+
+    def _construct_lattice_moment(self):
+        """
+        This function constructs the moments of the lattice.
+
+        The moments are the products of the velocity vectors, which are used in the computation of
+        the equilibrium distribution functions and the collision operator in the Lattice Boltzmann
+        Method (LBM).
+
+        Returns
+        -------
+        cc: numpy.ndarray
+            The moments of the lattice.
+        """
+        c = self._c.T
+        # Counter for the loop
+        cntr = 0
+        c = self._c.T
+        # nt: number of independent elements of a symmetric tensor
+        nt = self.d * (self.d + 1) // 2
+        cc = np.zeros((self.q, nt))
+        cntr = 0
+        for a in range(self.d):
+            for b in range(a, self.d):
+                cc[:, cntr] = c[:, a] * c[:, b]
+                cntr += 1
+        return cc
+
+    def _construct_opposite_indices(self):
+        """
+        This function constructs the indices of the opposite velocities for each velocity.
+
+        The opposite velocity of a velocity is the velocity that has the same magnitude but the
+        opposite direction.
+
+        Returns
+        -------
+        opposite: numpy.ndarray
+            The indices of the opposite velocities.
+        """
+        c = self._c.T
+        return np.array([c.tolist().index((-c[i]).tolist()) for i in range(self.q)])
+
+    def _construct_main_indices(self):
+        """
+        This function constructs the indices of the main velocities.
+
+        The main velocities are the velocities that have a magnitude of 1 in lattice units.
+
+        Returns
+        -------
+        numpy.ndarray
+            The indices of the main velocities.
+        """
+        c = self._c.T
+        if self.d == 2:
+            return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) == 1))[0]
+        elif self.d == 3:
+            return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) + np.abs(c[:, 2]) == 1))[0]
+
+    def _construct_right_indices(self):
+        """
+        This function constructs the indices of the velocities that point in the positive
+        x-direction.
+
+        Returns
+        -------
+        numpy.ndarray
+            The indices of the right velocities.
+        """
+        return np.nonzero(self._c.T[:, 0] == 1)[0]
+
+    def _construct_left_indices(self):
+        """
+        This function constructs the indices of the velocities that point in the negative
+        x-direction.
+
+        Returns
+        -------
+        numpy.ndarray
+            The indices of the left velocities.
+        """
+        return np.nonzero(self._c.T[:, 0] == -1)[0]
+
+    def __str__(self):
+        """
+        This function returns the name of the lattice in the format of DxQy.
+        """
+        return self.__repr__()
+
+    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/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py
new file mode 100644
index 0000000..1025e3b
--- /dev/null
+++ b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py
@@ -0,0 +1,94 @@
+import pytest
+import numpy as np
+import jax.numpy as jnp
+import xlb
+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):
+    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",
+    [
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+    ],
+)
+def test_bc_equilibrium_jax(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+    velocity_set = DefaultConfig.velocity_set
+
+    missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
+
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+
+    indices_boundary_masker = IndicesBoundaryMasker()
+
+    # Make indices for boundary conditions (sphere)
+    sphere_radius = grid_shape[0] // 4
+    nr = grid_shape[0]
+    x = np.arange(nr)
+    y = np.arange(nr)
+    z = np.arange(nr)
+    if dim == 2:
+        X, Y = np.meshgrid(x, y)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2)
+    else:
+        X, Y, Z = np.meshgrid(x, y, z)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2)
+
+    indices = [tuple(indices[i]) for i in range(velocity_set.d)]
+
+    equilibrium_bc = xlb.operator.boundary_condition.EquilibriumBC(
+        rho=1.0,
+        u=(0.0, 0.0, 0.0) if dim == 3 else (0.0, 0.0),
+        equilibrium_operator=xlb.operator.equilibrium.QuadraticEquilibrium(),
+        indices=indices,
+    )
+
+    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)
+
+    f_post = my_grid.create_field(
+        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, bc_mask, missing_mask)
+
+    assert f.shape == (velocity_set.q,) + grid_shape
+
+    # Assert that the values are correct in the indices of the sphere
+    weights = velocity_set.w
+    for i, weight in enumerate(weights):
+        if dim == 2:
+            assert jnp.allclose(f[i, indices[0], indices[1]], weight), f"Direction {i} in f does not match the expected weight"
+        else:
+            assert jnp.allclose(f[i, indices[0], indices[1], indices[2]], weight), f"Direction {i} in f does not match the expected weight"
+
+    # Make sure that everywhere else the values are the same as f_post. Note that indices are just int values
+    mask_outside = np.ones(grid_shape, dtype=bool)
+    mask_outside[indices] = False  # Mark boundary as false
+    if dim == 2:
+        for i in range(velocity_set.q):
+            assert jnp.allclose(f[i, mask_outside], f_post[i, mask_outside])
+    else:
+        for i in range(velocity_set.q):
+            assert jnp.allclose(f[i, mask_outside], f_post[i, mask_outside])
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py
new file mode 100644
index 0000000..6bd9311
--- /dev/null
+++ b/tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py
@@ -0,0 +1,97 @@
+import pytest
+import numpy as np
+import xlb
+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):
+    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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape",
+    [
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+    ],
+)
+def test_bc_equilibrium_warp(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+    velocity_set = DefaultConfig.velocity_set
+
+    missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
+
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+
+    indices_boundary_masker = IndicesBoundaryMasker()
+
+    # Make indices for boundary conditions (sphere)
+    sphere_radius = grid_shape[0] // 4
+    nr = grid_shape[0]
+    x = np.arange(nr)
+    y = np.arange(nr)
+    z = np.arange(nr)
+    if dim == 2:
+        X, Y = np.meshgrid(x, y)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2)
+    else:
+        X, Y, Z = np.meshgrid(x, y, z)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2)
+
+    indices = [tuple(indices[i]) for i in range(velocity_set.d)]
+    equilibrium = xlb.operator.equilibrium.QuadraticEquilibrium()
+
+    equilibrium_bc = xlb.operator.boundary_condition.EquilibriumBC(
+        rho=1.0,
+        u=(0.0, 0.0, 0.0) if dim == 3 else (0.0, 0.0),
+        equilibrium_operator=equilibrium,
+        indices=indices,
+    )
+
+    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)
+    f_post = my_grid.create_field(
+        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, bc_mask, missing_mask)
+
+    f = f.numpy()
+    f_post = f_post.numpy()
+
+    assert f.shape == (velocity_set.q,) + grid_shape if dim == 3 else (velocity_set.q, grid_shape[0], grid_shape[1], 1)
+
+    # Assert that the values are correct in the indices of the sphere
+    weights = velocity_set.w
+    for i, weight in enumerate(weights):
+        if dim == 2:
+            assert np.allclose(f[i, indices[0], indices[1]], weight), f"Direction {i} in f does not match the expected weight"
+        else:
+            assert np.allclose(f[i, indices[0], indices[1], indices[2]], weight), f"Direction {i} in f does not match the expected weight"
+
+    # Make sure that everywhere else the values are the same as f_post. Note that indices are just int values
+    mask_outside = np.ones(grid_shape, dtype=bool)
+    mask_outside[indices] = False  # Mark boundary as false
+    if dim == 2:
+        for i in range(velocity_set.q):
+            assert np.allclose(f[i, mask_outside], f_post[i, mask_outside])
+    else:
+        for i in range(velocity_set.q):
+            assert np.allclose(f[i, mask_outside], f_post[i, mask_outside])
+
+
+if __name__ == "__main__":
+    pytest.main()
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
new file mode 100644
index 0000000..cd18975
--- /dev/null
+++ b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py
@@ -0,0 +1,90 @@
+import pytest
+import numpy as np
+import jax.numpy as jnp
+import xlb
+import jax
+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):
+    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",
+    [
+        (2, xlb.velocity_set.D2Q9, (50, 50)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (100, 100, 100)),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q27, (100, 100, 100)),
+    ],
+)
+def test_fullway_bounce_back_jax(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+    velocity_set = DefaultConfig.velocity_set
+
+    missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
+
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+
+    indices_boundary_masker = IndicesBoundaryMasker()
+
+    # Make indices for boundary conditions (sphere)
+    sphere_radius = grid_shape[0] // 4
+    nr = grid_shape[0]
+    x = np.arange(nr)
+    y = np.arange(nr)
+    z = np.arange(nr)
+    if dim == 2:
+        X, Y = np.meshgrid(x, y)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2)
+    else:
+        X, Y, Z = np.meshgrid(x, y, z)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2)
+
+    indices = [tuple(indices[i]) for i in range(velocity_set.d)]
+    fullway_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices)
+
+    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
+    key = jax.random.PRNGKey(0)
+    random_field = jax.random.uniform(key, f_pre.shape)
+    # Add the random field to f_pre
+    f_pre += random_field
+
+    f_post = my_grid.create_field(
+        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, bc_mask, missing_mask)
+
+    assert f.shape == (velocity_set.q,) + grid_shape
+
+    for i in range(velocity_set.q):
+        jnp.allclose(
+            f[velocity_set.opp_indices[i]][tuple(indices)],
+            f_pre[i][tuple(indices)],
+        )
+
+    # Make sure that everywhere else the values are the same as f_post. Note that indices are just int values
+    mask_outside = np.ones(grid_shape, dtype=bool)
+    mask_outside[indices] = False  # Mark boundary as false
+    if dim == 2:
+        for i in range(velocity_set.q):
+            assert jnp.allclose(f[i, mask_outside], f_post[i, mask_outside])
+    else:
+        for i in range(velocity_set.q):
+            assert jnp.allclose(f[i, mask_outside], f_post[i, mask_outside])
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
new file mode 100644
index 0000000..59c6c9d
--- /dev/null
+++ b/tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py
@@ -0,0 +1,93 @@
+import pytest
+import numpy as np
+import warp as wp
+import xlb
+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):
+    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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape",
+    [
+        (2, xlb.velocity_set.D2Q9, (50, 50)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (100, 100, 100)),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q27, (100, 100, 100)),
+    ],
+)
+def test_fullway_bounce_back_warp(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+    velocity_set = DefaultConfig.velocity_set
+
+    missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
+
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+
+    indices_boundary_masker = IndicesBoundaryMasker()
+
+    # Make indices for boundary conditions (sphere)
+    sphere_radius = grid_shape[0] // 4
+    nr = grid_shape[0]
+    x = np.arange(nr)
+    y = np.arange(nr)
+    z = np.arange(nr)
+    if dim == 2:
+        X, Y = np.meshgrid(x, y)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2)
+    else:
+        X, Y, Z = np.meshgrid(x, y, z)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2)
+
+    indices = [tuple(indices[i]) for i in range(velocity_set.d)]
+    fullway_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices)
+
+    bc_mask, missing_mask = indices_boundary_masker([fullway_bc], bc_mask, missing_mask, start_index=None)
+
+    # Generate a random field with the same shape
+    if dim == 2:
+        random_field = np.random.rand(velocity_set.q, grid_shape[0], grid_shape[1], 1).astype(np.float32)
+    else:
+        random_field = np.random.rand(velocity_set.q, grid_shape[0], grid_shape[1], grid_shape[2]).astype(np.float32)
+    # Add the random field to f_pre
+    f_pre = wp.array(random_field)
+
+    f_post = my_grid.create_field(
+        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, bc_mask, missing_mask)
+
+    f = f_pre.numpy()
+    f_post = f_post.numpy()
+
+    assert f.shape == (velocity_set.q,) + grid_shape if dim == 3 else (velocity_set.q, grid_shape[0], grid_shape[1], 1)
+
+    for i in range(velocity_set.q):
+        np.allclose(
+            f[velocity_set.opp_indices[i]][tuple(indices)],
+            f_post[i][tuple(indices)],
+        )
+
+    # Make sure that everywhere else the values are the same as f_post. Note that indices are just int values
+    mask_outside = np.ones(grid_shape, dtype=bool)
+    mask_outside[indices] = False  # Mark boundary as false
+    if dim == 2:
+        for i in range(velocity_set.q):
+            assert np.allclose(f[i, mask_outside], f_post[i, mask_outside])
+    else:
+        for i in range(velocity_set.q):
+            assert np.allclose(f[i, mask_outside], f_post[i, mask_outside])
diff --git a/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py b/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py
new file mode 100644
index 0000000..79d56d8
--- /dev/null
+++ b/tests/boundary_conditions/mask/test_bc_indices_masker_jax.py
@@ -0,0 +1,83 @@
+import pytest
+import jax.numpy as jnp
+import numpy as np
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb import DefaultConfig
+from xlb.grid import grid_factory
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape",
+    [
+        (2, xlb.velocity_set.D2Q9, (4, 4)),
+        (2, xlb.velocity_set.D2Q9, (50, 50)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (100, 100, 100)),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q27, (100, 100, 100)),
+    ],
+)
+def test_indices_masker_jax(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+    velocity_set = DefaultConfig.velocity_set
+
+    missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
+
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+
+    indices_boundary_masker = xlb.operator.boundary_masker.IndicesBoundaryMasker()
+
+    # Make indices for boundary conditions (sphere)
+    sphere_radius = grid_shape[0] // 4
+    nr = grid_shape[0]
+    x = np.arange(nr)
+    y = np.arange(nr)
+    z = np.arange(nr)
+    if dim == 2:
+        X, Y = np.meshgrid(x, y)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2)
+    else:
+        X, Y, Z = np.meshgrid(x, y, z)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2)
+
+    indices = [tuple(indices[i]) for i in range(velocity_set.d)]
+
+    assert len(indices) == dim
+    test_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices)
+    test_bc.id = 5
+    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 bc_mask.dtype == xlb.Precision.UINT8.jax_dtype
+
+    assert bc_mask.shape == (1,) + grid_shape
+
+    assert missing_mask.shape == (velocity_set.q,) + grid_shape
+
+    if dim == 2:
+        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(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__":
+    pytest.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
new file mode 100644
index 0000000..f56c9ce
--- /dev/null
+++ b/tests/boundary_conditions/mask/test_bc_indices_masker_warp.py
@@ -0,0 +1,89 @@
+import pytest
+import numpy as np
+import xlb
+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.WARP,
+        velocity_set=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape",
+    [
+        (2, xlb.velocity_set.D2Q9, (50, 50)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (100, 100, 100)),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q27, (100, 100, 100)),
+    ],
+)
+def test_indices_masker_warp(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+    velocity_set = DefaultConfig.velocity_set
+
+    missing_mask = my_grid.create_field(cardinality=velocity_set.q, dtype=xlb.Precision.BOOL)
+
+    bc_mask = my_grid.create_field(cardinality=1, dtype=xlb.Precision.UINT8)
+
+    indices_boundary_masker = IndicesBoundaryMasker()
+
+    # Make indices for boundary conditions (sphere)
+    sphere_radius = grid_shape[0] // 4
+    nr = grid_shape[0]
+    x = np.arange(nr)
+    y = np.arange(nr)
+    z = np.arange(nr)
+    if dim == 2:
+        X, Y = np.meshgrid(x, y)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 < sphere_radius**2)
+    else:
+        X, Y, Z = np.meshgrid(x, y, z)
+        indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2)
+
+    indices = [tuple(indices[i]) for i in range(velocity_set.d)]
+
+    assert len(indices) == dim
+    test_bc = xlb.operator.boundary_condition.FullwayBounceBackBC(indices=indices)
+    test_bc.id = 5
+    bc_mask, missing_mask = indices_boundary_masker(
+        [test_bc],
+        bc_mask,
+        missing_mask,
+        start_index=(0, 0, 0) if dim == 3 else (0, 0),
+    )
+    assert missing_mask.dtype == xlb.Precision.BOOL.wp_dtype
+
+    assert bc_mask.dtype == xlb.Precision.UINT8.wp_dtype
+
+    bc_mask = bc_mask.numpy()
+    missing_mask = missing_mask.numpy()
+
+    assert bc_mask.shape == (1,) + grid_shape
+
+    assert missing_mask.shape == (velocity_set.q,) + grid_shape
+
+    if dim == 2:
+        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(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__":
+    pytest.main()
diff --git a/tests/grids/__init__.py b/tests/grids/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/grids/test_grid_jax.py b/tests/grids/test_grid_jax.py
new file mode 100644
index 0000000..dd74da6
--- /dev/null
+++ b/tests/grids/test_grid_jax.py
@@ -0,0 +1,74 @@
+import pytest
+import jax
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.grid import grid_factory
+from jax.sharding import Mesh
+from jax.experimental import mesh_utils
+import jax.numpy as jnp
+
+
+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,
+        velocity_set=vel_set,
+    )
+
+
+@pytest.mark.parametrize("grid_size", [50, 100, 150])
+def test_jax_2d_grid_initialization(grid_size):
+    init_xlb_env(xlb.velocity_set.D2Q9)
+    grid_shape = (grid_size, grid_size)
+    my_grid = grid_factory(grid_shape)
+    f = my_grid.create_field(cardinality=9)
+    n_devices = jax.device_count()
+
+    device_mesh = mesh_utils.create_device_mesh((1, n_devices, 1))
+    expected_mesh = Mesh(device_mesh, axis_names=("cardinality", "x", "y"))
+
+    assert f.shape == (9,) + grid_shape, "Field shape is incorrect"
+    assert f.sharding.mesh == expected_mesh, "Field sharding mesh is incorrect"
+    assert f.sharding.spec == ("cardinality", "x", "y"), "PartitionSpec is incorrect"
+
+
+@pytest.mark.parametrize("grid_size", [50, 100, 150])
+def test_jax_3d_grid_initialization(grid_size):
+    init_xlb_env(xlb.velocity_set.D3Q19)
+    grid_shape = (grid_size, grid_size, grid_size)
+    my_grid = grid_factory(grid_shape)
+    f = my_grid.create_field(cardinality=9)
+    n_devices = jax.device_count()
+
+    device_mesh = mesh_utils.create_device_mesh((1, n_devices, 1, 1))
+    expected_mesh = Mesh(device_mesh, axis_names=("cardinality", "x", "y", "z"))
+
+    assert f.shape == (9,) + grid_shape, "Field shape is incorrect"
+    assert f.sharding.mesh == expected_mesh, "Field sharding mesh is incorrect"
+    assert f.sharding.spec == (
+        "cardinality",
+        "x",
+        "y",
+        "z",
+    ), "PartitionSpec is incorrect"
+
+
+def test_jax_grid_create_field_fill_value():
+    init_xlb_env(xlb.velocity_set.D2Q9)
+    grid_shape = (100, 100)
+    fill_value = 3.14
+    my_grid = grid_factory(grid_shape)
+
+    f = my_grid.create_field(cardinality=9, fill_value=fill_value)
+    assert f.shape == (9,) + grid_shape, "Field shape is incorrect"
+    assert jnp.allclose(f, fill_value), "Field not properly initialized with fill_value"
+
+
+@pytest.fixture(autouse=True)
+def setup_xlb_env():
+    init_xlb_env(xlb.velocity_set.D2Q9)
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/grids/test_grid_warp.py b/tests/grids/test_grid_warp.py
new file mode 100644
index 0000000..61b27d4
--- /dev/null
+++ b/tests/grids/test_grid_warp.py
@@ -0,0 +1,51 @@
+import pytest
+import warp as wp
+import numpy as np
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.grid import grid_factory
+from xlb.precision_policy import Precision
+
+
+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=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)]:
+        init_xlb_env(xlb.velocity_set.D3Q19)
+        my_grid = grid_factory(grid_shape)
+        f = my_grid.create_field(cardinality=9, dtype=Precision.FP32)
+        if len(grid_shape) == 2:
+            assert f.shape == (9,) + grid_shape + (1,), "Field shape is incorrect got {}".format(f.shape)
+        else:
+            assert f.shape == (9,) + grid_shape, "Field shape is incorrect got {}".format(f.shape)
+        assert isinstance(f, wp.array), "Field should be a Warp ndarray"
+
+
+def test_warp_grid_create_field_fill_value():
+    init_xlb_env(xlb.velocity_set.D2Q9)
+    grid_shape = (100, 100)
+    fill_value = 3.14
+    my_grid = grid_factory(grid_shape)
+
+    f = my_grid.create_field(cardinality=9, dtype=Precision.FP32, fill_value=fill_value)
+    assert isinstance(f, wp.array), "Field should be a Warp ndarray"
+
+    f = f.numpy()
+    assert np.allclose(f, fill_value), "Field not properly initialized with fill_value"
+
+
+@pytest.fixture(autouse=True)
+def setup_xlb_env():
+    init_xlb_env(xlb.velocity_set.D2Q9)
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/kernels/__init__.py b/tests/kernels/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/kernels/collision/test_bgk_collision_jax.py b/tests/kernels/collision/test_bgk_collision_jax.py
new file mode 100644
index 0000000..1672cd5
--- /dev/null
+++ b/tests/kernels/collision/test_bgk_collision_jax.py
@@ -0,0 +1,54 @@
+import pytest
+import jax.numpy as jnp
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.operator.collision import BGK
+from xlb.grid import grid_factory
+from xlb import DefaultConfig
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape,omega",
+    [
+        (2, xlb.velocity_set.D2Q9, (100, 100), 0.6),
+        (2, xlb.velocity_set.D2Q9, (100, 100), 1.0),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 0.6),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.0),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50), 0.6),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50), 1.0),
+    ],
+)
+def test_bgk_ollision(dim, velocity_set, grid_shape, omega):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+
+    rho = my_grid.create_field(cardinality=1, fill_value=1.0)
+    u = my_grid.create_field(cardinality=dim, fill_value=0.0)
+
+    # Compute equilibrium
+    compute_macro = QuadraticEquilibrium()
+    f_eq = compute_macro(rho, u)
+
+    # Compute collision
+
+    compute_collision = BGK(omega=omega)
+
+    f_orig = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q)
+
+    f_out = compute_collision(f_orig, f_eq, rho, u)
+
+    assert jnp.allclose(f_out, f_orig - omega * (f_orig - f_eq))
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/kernels/collision/test_bgk_collision_warp.py b/tests/kernels/collision/test_bgk_collision_warp.py
new file mode 100644
index 0000000..382e368
--- /dev/null
+++ b/tests/kernels/collision/test_bgk_collision_warp.py
@@ -0,0 +1,57 @@
+import pytest
+import numpy as np
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.operator.collision import BGK
+from xlb.grid import grid_factory
+from xlb import DefaultConfig
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape,omega",
+    [
+        (2, xlb.velocity_set.D2Q9, (100, 100), 0.6),
+        (2, xlb.velocity_set.D2Q9, (100, 100), 1.0),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 0.6),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.0),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50), 0.6),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50), 1.0),
+    ],
+)
+def test_bgk_collision_warp(dim, velocity_set, grid_shape, omega):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+
+    rho = my_grid.create_field(cardinality=1, fill_value=1.0)
+    u = my_grid.create_field(cardinality=dim, fill_value=0.0)
+
+    compute_macro = QuadraticEquilibrium()
+
+    f_eq = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q)
+    f_eq = compute_macro(rho, u, f_eq)
+
+    compute_collision = BGK(omega=omega)
+    f_orig = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q)
+
+    f_out = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q)
+    f_out = compute_collision(f_orig, f_eq, f_out, rho, u)
+
+    f_eq = f_eq.numpy()
+    f_out = f_out.numpy()
+    f_orig = f_orig.numpy()
+
+    assert np.allclose(f_out, f_orig - omega * (f_orig - f_eq), atol=1e-5)
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/kernels/equilibrium/test_equilibrium_jax.py b/tests/kernels/equilibrium/test_equilibrium_jax.py
new file mode 100644
index 0000000..aa4f051
--- /dev/null
+++ b/tests/kernels/equilibrium/test_equilibrium_jax.py
@@ -0,0 +1,52 @@
+import pytest
+import numpy as np
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.grid import grid_factory
+from xlb import DefaultConfig
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape",
+    [
+        (2, xlb.velocity_set.D2Q9, (50, 50)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (100, 100, 100)),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q27, (100, 100, 100)),
+    ],
+)
+def test_quadratic_equilibrium_jax(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+
+    rho = my_grid.create_field(cardinality=1, fill_value=1.0)
+    u = my_grid.create_field(cardinality=dim, fill_value=0.0)
+
+    # Compute equilibrium
+    compute_macro = QuadraticEquilibrium()
+    f_eq = compute_macro(rho, u)
+
+    # Test sum of f_eq across cardinality at each point
+    sum_f_eq = np.sum(f_eq, axis=0)
+    assert np.allclose(sum_f_eq, 1.0), "Sum of f_eq should be 1.0 across all directions at each grid point"
+
+    # Test that each direction matches the expected weights
+    weights = DefaultConfig.velocity_set.w
+    for i, weight in enumerate(weights):
+        assert np.allclose(f_eq[i, ...], weight), f"Direction {i} in f_eq does not match the expected weight"
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/kernels/equilibrium/test_equilibrium_warp.py b/tests/kernels/equilibrium/test_equilibrium_warp.py
new file mode 100644
index 0000000..fdd796a
--- /dev/null
+++ b/tests/kernels/equilibrium/test_equilibrium_warp.py
@@ -0,0 +1,58 @@
+import pytest
+import numpy as np
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.grid import grid_factory
+from xlb import DefaultConfig
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape",
+    [
+        (2, xlb.velocity_set.D2Q9, (50, 50)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (100, 100, 100)),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q27, (100, 100, 100)),
+    ],
+)
+def test_quadratic_equilibrium_warp(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+
+    rho = my_grid.create_field(cardinality=1, fill_value=1.0)
+    u = my_grid.create_field(cardinality=dim, fill_value=0.0)
+
+    f_eq = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q)
+
+    compute_macro = QuadraticEquilibrium()
+    f_eq = compute_macro(rho, u, f_eq)
+
+    f_eq_np = f_eq.numpy()
+
+    sum_f_eq = np.sum(f_eq_np, axis=0)
+    assert np.allclose(sum_f_eq, 1.0), "Sum of f_eq should be 1.0 across all directions at each grid point"
+
+    weights = DefaultConfig.velocity_set.w
+    for i, weight in enumerate(weights):
+        assert np.allclose(f_eq_np[i, ...], weight), f"Direction {i} in f_eq does not match the expected weight"
+
+
+# @pytest.fixture(autouse=True)
+# def setup_xlb_env(request):
+#     dim, velocity_set, grid_shape = request.param
+#     init_xlb_env(velocity_set)
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/kernels/macroscopic/test_macroscopic_jax.py b/tests/kernels/macroscopic/test_macroscopic_jax.py
new file mode 100644
index 0000000..2c2ad55
--- /dev/null
+++ b/tests/kernels/macroscopic/test_macroscopic_jax.py
@@ -0,0 +1,50 @@
+import pytest
+import numpy as np
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.operator.macroscopic import Macroscopic
+from xlb.grid import grid_factory
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape,rho,velocity",
+    [
+        (2, xlb.velocity_set.D2Q9, (100, 100), 1.0, 0.0),
+        (2, xlb.velocity_set.D2Q9, (100, 100), 1.1, 1.0),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.0, 0.0),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.1, 1.0),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50), 1.0, 0.0),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50), 1.1, 1.0),
+    ],
+)
+def test_macroscopic_jax(dim, velocity_set, grid_shape, rho, velocity):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+
+    rho_field = my_grid.create_field(cardinality=1, fill_value=rho)
+    velocity_field = my_grid.create_field(cardinality=dim, fill_value=velocity)
+
+    # Compute equilibrium
+    f_eq = QuadraticEquilibrium()(rho_field, velocity_field)
+
+    compute_macro = Macroscopic()
+
+    rho_calc, u_calc = compute_macro(f_eq)
+
+    # Test sum of f_eq which should be 1.0 for rho and 0.0 for u
+    assert np.allclose(rho_calc, rho), "Sum of f_eq should be {rho} for rho"
+    assert np.allclose(u_calc, velocity), "Sum of f_eq should be {velocity} for u"
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/kernels/macroscopic/test_macroscopic_warp.py b/tests/kernels/macroscopic/test_macroscopic_warp.py
new file mode 100644
index 0000000..4f33bc2
--- /dev/null
+++ b/tests/kernels/macroscopic/test_macroscopic_warp.py
@@ -0,0 +1,53 @@
+import pytest
+import numpy as np
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.equilibrium import QuadraticEquilibrium
+from xlb.operator.macroscopic import Macroscopic
+from xlb.grid import grid_factory
+from xlb import DefaultConfig
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape,rho,velocity",
+    [
+        (2, xlb.velocity_set.D2Q9, (100, 100), 1.0, 0.0),
+        (2, xlb.velocity_set.D2Q9, (100, 100), 1.1, 1.0),
+        (2, xlb.velocity_set.D2Q9, (100, 100), 1.1, 2.0),
+        (2, xlb.velocity_set.D2Q9, (50, 50), 1.1, 2.0),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.0, 0.0),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.1, 1.0),  # TODO: Uncommenting will cause a Warp error. Needs investigation.
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50), 1.1, 2.0),  # TODO: Uncommenting will cause a Warp error. Needs investigation.
+    ],
+)
+def test_macroscopic_warp(dim, velocity_set, grid_shape, rho, velocity):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+
+    rho_field = my_grid.create_field(cardinality=1, fill_value=rho)
+    velocity_field = my_grid.create_field(cardinality=dim, fill_value=velocity)
+
+    f_eq = my_grid.create_field(cardinality=DefaultConfig.velocity_set.q)
+    f_eq = QuadraticEquilibrium()(rho_field, velocity_field, f_eq)
+
+    compute_macro = Macroscopic()
+    rho_calc = my_grid.create_field(cardinality=1)
+    u_calc = my_grid.create_field(cardinality=dim)
+
+    rho_calc, u_calc = compute_macro(f_eq, rho_calc, u_calc)
+
+    assert np.allclose(rho_calc.numpy(), rho), f"Computed density should be close to initialized density {rho}"
+    assert np.allclose(u_calc.numpy(), velocity), f"Computed velocity should be close to initialized velocity {velocity}"
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/kernels/stream/test_stream_jax.py b/tests/kernels/stream/test_stream_jax.py
new file mode 100644
index 0000000..c1cae52
--- /dev/null
+++ b/tests/kernels/stream/test_stream_jax.py
@@ -0,0 +1,69 @@
+import pytest
+import jax.numpy as jnp
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.stream import Stream
+from xlb import DefaultConfig
+from xlb.grid import grid_factory
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape",
+    [
+        (2, xlb.velocity_set.D2Q9, (50, 50)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (100, 100, 100)),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q27, (100, 100, 100)),
+    ],
+)
+def test_stream_operator_jax(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid = grid_factory(grid_shape)
+    velocity_set = DefaultConfig.velocity_set
+
+    stream_op = Stream()
+
+    f_initial = my_grid.create_field(cardinality=velocity_set.q)
+    f_initial = f_initial.at[..., f_initial.shape[-1] // 2].set(1)
+
+    f_streamed = stream_op(f_initial)
+
+    expected = []
+
+    if dim == 2:
+        for i in range(velocity_set.q):
+            expected.append(
+                jnp.roll(
+                    f_initial[i, ...],
+                    (velocity_set.c[0][i], velocity_set.c[1][i]),
+                    axis=(0, 1),
+                )
+            )
+    elif dim == 3:
+        for i in range(velocity_set.q):
+            expected.append(
+                jnp.roll(
+                    f_initial[i, ...],
+                    (velocity_set.c[0][i], velocity_set.c[1][i], velocity_set.c[2][i]),
+                    axis=(0, 1, 2),
+                )
+            )
+
+    expected = jnp.stack(expected, axis=0)
+
+    assert jnp.allclose(f_streamed, expected), "Streaming did not occur as expected"
+
+
+if __name__ == "__main__":
+    pytest.main()
diff --git a/tests/kernels/stream/test_stream_warp.py b/tests/kernels/stream/test_stream_warp.py
new file mode 100644
index 0000000..95fcc05
--- /dev/null
+++ b/tests/kernels/stream/test_stream_warp.py
@@ -0,0 +1,81 @@
+import pytest
+import jax.numpy as jnp
+import numpy as np
+import warp as wp
+import xlb
+from xlb.compute_backend import ComputeBackend
+from xlb.operator.stream import Stream
+from xlb import DefaultConfig
+from xlb.grid import grid_factory
+
+
+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=vel_set,
+    )
+
+
+@pytest.mark.parametrize(
+    "dim,velocity_set,grid_shape",
+    [
+        (2, xlb.velocity_set.D2Q9, (50, 50)),
+        (2, xlb.velocity_set.D2Q9, (100, 100)),
+        (3, xlb.velocity_set.D3Q19, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q19, (100, 100, 100)),
+        (3, xlb.velocity_set.D3Q27, (50, 50, 50)),
+        (3, xlb.velocity_set.D3Q27, (100, 100, 100)),
+    ],
+)
+def test_stream_operator_warp(dim, velocity_set, grid_shape):
+    init_xlb_env(velocity_set)
+    my_grid_jax = grid_factory(grid_shape, compute_backend=ComputeBackend.JAX)
+    velocity_set = DefaultConfig.velocity_set
+
+    f_initial = my_grid_jax.create_field(cardinality=velocity_set.q)
+    f_initial = f_initial.at[..., f_initial.shape[-1] // 2].set(1)
+
+    expected = []
+
+    if dim == 2:
+        for i in range(velocity_set.q):
+            expected.append(
+                jnp.roll(
+                    f_initial[i, ...],
+                    (velocity_set.c[0][i], velocity_set.c[1][i]),
+                    axis=(0, 1),
+                )
+            )
+    elif dim == 3:
+        for i in range(velocity_set.q):
+            expected.append(
+                jnp.roll(
+                    f_initial[i, ...],
+                    (velocity_set.c[0][i], velocity_set.c[1][i], velocity_set.c[2][i]),
+                    axis=(0, 1, 2),
+                )
+            )
+
+    expected = jnp.stack(expected, axis=0)
+
+    if dim == 2:
+        f_initial_warp = wp.array(f_initial[..., np.newaxis])
+
+    elif dim == 3:
+        f_initial_warp = wp.array(f_initial)
+
+    stream_op = Stream()
+    my_grid_warp = grid_factory(grid_shape, compute_backend=ComputeBackend.WARP)
+    f_streamed = my_grid_warp.create_field(cardinality=velocity_set.q)
+    f_streamed = stream_op(f_initial_warp, f_streamed)
+
+    if len(grid_shape) == 2:
+        assert jnp.allclose(f_streamed.numpy()[..., 0], np.array(expected)), "Streaming did not occur as expected"
+    else:
+        assert jnp.allclose(f_streamed.numpy(), np.array(expected)), "Streaming did not occur as expected"
+
+
+if __name__ == "__main__":
+    pytest.main()