From 3dba22232c0e35469b5097c6a8f7cbd2733158f5 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Thu, 21 Sep 2023 13:21:17 -0400
Subject: [PATCH 1/8] added interpolating bounce-back (my own single equation
 and local formulation) where wall proximity is determined through SDF.

---
 examples/CFD/cylinder2d.py |   5 +-
 src/boundary_conditions.py | 120 +++++++++++++++++++++++++++++++++++--
 2 files changed, 120 insertions(+), 5 deletions(-)

diff --git a/examples/CFD/cylinder2d.py b/examples/CFD/cylinder2d.py
index 8196378..5120668 100644
--- a/examples/CFD/cylinder2d.py
+++ b/examples/CFD/cylinder2d.py
@@ -44,7 +44,10 @@ def set_boundary_conditions(self):
         cx, cy = 2.*diam, 2.*diam
         cylinder = (xx - cx)**2 + (yy-cy)**2 <= (diam/2.)**2
         cylinder = coord[cylinder]
-        self.BCs.append(BounceBackHalfway(tuple(cylinder.T), self.gridInfo, self.precisionPolicy))
+        implicit_distance = np.reshape((xx - cx)**2 + (yy-cy)**2 - (diam/2.)**2, (self.nx, self.ny))
+        self.BCs.append(InterpolatedBounceBackLocal(tuple(cylinder.T), implicit_distance,
+                                               self.gridInfo, self.precisionPolicy))
+
         # wall = np.concatenate([cylinder, self.boundingBoxIndices['top'], self.boundingBoxIndices['bottom']])
         # self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
 
diff --git a/src/boundary_conditions.py b/src/boundary_conditions.py
index c33c556..757a6ca 100644
--- a/src/boundary_conditions.py
+++ b/src/boundary_conditions.py
@@ -495,11 +495,11 @@ def configure(self, boundaryBitmask):
         the boundary nodes to be the indices of fluid nodes adjacent of the solid nodes.
         """
         # Perform index shift for halfway BB.
-        shiftDir = ~boundaryBitmask[:, self.lattice.opp_indices]
+        hasFluidNeighbour = ~boundaryBitmask[:, self.lattice.opp_indices]
         idx = np.array(self.indices).T
         idx_trg = []
         for i in range(self.lattice.q):
-            idx_trg.append(idx[shiftDir[:, i], :] + self.lattice.c[:, i])
+            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)
         return
@@ -935,11 +935,11 @@ def configure(self, boundaryBitmask):
         boundaryBitmask : np.ndarray
             The connectivity bitmask for the boundary voxels.
         """        
-        shiftDir = ~boundaryBitmask[:, self.lattice.opp_indices]
+        hasFluidNeighbour = ~boundaryBitmask[:, self.lattice.opp_indices]
         idx = np.array(self.indices).T
         idx_trg = []
         for i in range(self.lattice.q):
-            idx_trg.append(idx[shiftDir[:, i], :] + self.lattice.c[:, i])
+            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)
 
@@ -1008,3 +1008,115 @@ def apply(self, fout, fin):
         fbd = fout[self.indices]
         fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
         return fbd
+
+
+class InterpolatedBounceBackLocal(BoundaryCondition):
+    """
+    A local single-node interpolated bounce-back boundary condition 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 "InterpolatedBounceBack".
+    implementationStep : str
+        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
+        it is "PostStreaming".
+    isSolid : bool
+        Whether the boundary condition represents a solid boundary. For this class, it is True.
+    needsExtraConfiguration : bool
+        Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True.
+    implicit_distances : array-like
+        An array of shape (nx,ny,nz) indicating the signed-distance field from the solid walls
+    """
+
+    def __init__(self, indices, implicit_distances, grid_info, precision_policy):
+
+        super().__init__(indices, grid_info, precision_policy)
+        self.name = "InterpolatedBounceBackLocal"
+        self.implementationStep = "PostStreaming"
+        self.implicit_distances = implicit_distances
+        self.weights = None
+        self.isSolid = True
+        self.needsExtraConfiguration = True
+
+    def configure(self, boundaryBitmask):
+        """
+        Configures the boundary condition.
+
+        Parameters
+        ----------
+        boundaryBitmask : array-like
+            The connectivity bitmask 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 = ~boundaryBitmask[:, 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_new = np.unique(np.vstack(idx_trg), axis=0)
+        self.indices = tuple(indices_new.T)
+        return
+
+    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.
+        """
+        idx = np.array(self.indices).T
+        self.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]
+            mask = self.iknownBitmask[:, q]
+            self.weights[mask, q] = sdf_f[mask] / (sdf_f[mask] - sdf_s[mask])
+        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]
+        fmissing = ((1. - self.weights) * f_postcollision_iknown +
+                    self.weights * (f_postcollision_imissing + f_poststreaming_iknown)) / (1.0 + self.weights)
+        fbd = fbd.at[bindex, self.imissing].set(fmissing)
+
+        return fbd
\ No newline at end of file

From 047919d4fc8854b8173449bcc23c828e7fda2fb6 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Mon, 2 Oct 2023 16:01:54 -0400
Subject: [PATCH 2/8] added the Bouzidi BC scheme

---
 requirements.txt           |   4 +-
 src/boundary_conditions.py | 130 +++++++++++++++++++++++++++++++++++--
 2 files changed, 128 insertions(+), 6 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 2c912ed..38e4d2d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
-jax==0.4.11
-jaxlib==0.4.11
+jax==0.4.14
+jaxlib==0.4.14
 jmp==0.0.4
 matplotlib==3.7.1
 numpy==1.24.2
diff --git a/src/boundary_conditions.py b/src/boundary_conditions.py
index 757a6ca..de02151 100644
--- a/src/boundary_conditions.py
+++ b/src/boundary_conditions.py
@@ -1010,9 +1010,10 @@ def apply(self, fout, fin):
         return fbd
 
 
-class InterpolatedBounceBackLocal(BoundaryCondition):
+class InterpolatedBounceBackDifferentiable(BoundaryCondition):
     """
-    A local single-node interpolated bounce-back boundary condition for a lattice Boltzmann method simulation.
+    A local and differentiable single-node interpolated bounce-back boundary condition for a lattice Boltzmann method
+    simulation.
 
     This class implements a interpolated bounce-back boundary condition. The boundary condition is applied after
     the streaming step.
@@ -1020,7 +1021,7 @@ class InterpolatedBounceBackLocal(BoundaryCondition):
     Attributes
     ----------
     name : str
-        The name of the boundary condition. For this class, it is "InterpolatedBounceBack".
+        The name of the boundary condition. For this class, it is "InterpolatedBounceBackDifferentiable".
     implementationStep : str
         The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
         it is "PostStreaming".
@@ -1035,7 +1036,7 @@ class InterpolatedBounceBackLocal(BoundaryCondition):
     def __init__(self, indices, implicit_distances, grid_info, precision_policy):
 
         super().__init__(indices, grid_info, precision_policy)
-        self.name = "InterpolatedBounceBackLocal"
+        self.name = "InterpolatedBounceBackDifferentiable"
         self.implementationStep = "PostStreaming"
         self.implicit_distances = implicit_distances
         self.weights = None
@@ -1119,4 +1120,125 @@ def apply(self, fout, fin):
                     self.weights * (f_postcollision_imissing + f_poststreaming_iknown)) / (1.0 + self.weights)
         fbd = fbd.at[bindex, self.imissing].set(fmissing)
 
+        return fbd
+
+
+class InterpolatedBounceBackBouzidi(BoundaryCondition):
+    """
+    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".
+    implementationStep : str
+        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
+        it is "PostStreaming".
+    isSolid : bool
+        Whether the boundary condition represents a solid boundary. For this class, it is True.
+    needsExtraConfiguration : bool
+        Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True.
+    implicit_distances : array-like
+        An array of shape (nx,ny,nz) indicating the signed-distance field from the solid walls
+    """
+
+    def __init__(self, indices, implicit_distances, grid_info, precision_policy):
+
+        super().__init__(indices, grid_info, precision_policy)
+        self.name = "InterpolatedBounceBackBouzidi"
+        self.implementationStep = "PostStreaming"
+        self.implicit_distances = implicit_distances
+        self.weights = None
+        self.isSolid = True
+        self.needsExtraConfiguration = True
+
+    def configure(self, boundaryBitmask):
+        """
+        Configures the boundary condition.
+
+        Parameters
+        ----------
+        boundaryBitmask : array-like
+            The connectivity bitmask 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 = ~boundaryBitmask[:, 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_new = np.unique(np.vstack(idx_trg), axis=0)
+        self.indices = tuple(indices_new.T)
+        return
+
+    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.
+        """
+        idx = np.array(self.indices).T
+        self.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]
+            mask = self.iknownBitmask[:, q]
+            self.weights[mask, q] = sdf_f[mask] / (sdf_f[mask] - sdf_s[mask])
+        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)
         return fbd
\ No newline at end of file

From 1cc98ec7eda2a9e05e3815b59a15ba7fca236e6d Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Mon, 2 Oct 2023 16:13:03 -0400
Subject: [PATCH 3/8] adding the modified exampels that are being used for
 benchmarking XLB

---
 examples/CFD/cavity3d.py   |  58 +++++++++++--------
 examples/CFD/cylinder2d.py | 114 ++++++++++++++++++++++---------------
 src/models.py              |  34 +++++++++++
 3 files changed, 134 insertions(+), 72 deletions(-)

diff --git a/examples/CFD/cavity3d.py b/examples/CFD/cavity3d.py
index d912786..f7e35c8 100644
--- a/examples/CFD/cavity3d.py
+++ b/examples/CFD/cavity3d.py
@@ -4,7 +4,7 @@
 
 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.
+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:
 
@@ -14,21 +14,19 @@
 4. Visualization: The simulation outputs data in VTK format for visualization. The data can be visualized using software like Paraview.
 
 """
-
-import os
-
 # Use 8 CPU devices
 # os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
 from src.models import BGKSim, KBCSim
 from src.lattice import LatticeD3Q27
-import numpy as np
 from src.utils import *
 from jax.config import config
 from src.boundary_conditions import *
+import json, codecs
 
-precision = 'f32/f32'
+precision = 'f64/f64'
+config.update('jax_enable_x64', True)
 
-class Cavity(KBCSim):
+class Cavity(BGKSim):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
@@ -39,7 +37,7 @@ def set_boundary_conditions(self):
              self.boundingBoxIndices['front'], self.boundingBoxIndices['back'],
              self.boundingBoxIndices['bottom']))
         # apply bounce back boundary condition to the walls
-        self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy))
+        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']
@@ -47,14 +45,14 @@ def set_boundary_conditions(self):
         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))
+        self.BCs.append(BounceBackHalfway(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, 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, 1:-1])
-        u = np.array(kwargs['u'][1:-1, 1:-1, 1:-1, :])
+        rho = np.array(kwargs['rho'])
+        u = np.array(kwargs['u'])
         timestep = kwargs['timestep']
-        u_prev = kwargs['u_prev'][1:-1, 1:-1, 1:-1, :]
+        u_prev = kwargs['u_prev']
 
         u_old = np.linalg.norm(u_prev, axis=2)
         u_new = np.linalg.norm(u, axis=2)
@@ -62,26 +60,36 @@ def output_data(self, **kwargs):
         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)
+        # save_fields_vtk(timestep, fields)
+
+        # output profiles of velocity at mid-plane for benchmarking
+        output_filename = "./profiles_" + f"{timestep:07d}.json"
+        ldc_ref_result = {'ux(x=y=0)': list(u[nx//2, ny//2, :, 0]/prescribed_vel),
+                          'uz(z=y=0)': list(u[:, ny//2, nz//2, 2]/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)
+        # u_mag = np.linalg.norm(u, axis=2)
         # live_volume_randering(timestep, u_mag)
 
 if __name__ == '__main__':
     lattice = LatticeD3Q27(precision)
 
-    nx = 101
-    ny = 101
-    nz = 101
+    nx = 256
+    ny = 256
+    nz = 256
 
-    Re = 50000.0
-    prescribed_vel = 0.1
-    clength = nx - 1
+    Re = 1000.0
+    prescribed_vel = 0.06
+    clength = nx - 2
 
     visc = prescribed_vel * clength / Re
     omega = 1.0 / (3. * visc + 0.5)
     print('omega = ', omega)
-    
+
     os.system("rm -rf ./*.vtk && rm -rf ./*.png")
 
     kwargs = {
@@ -91,9 +99,9 @@ def output_data(self, **kwargs):
         'ny': ny,
         'nz': nz,
         'precision': precision,
-        'io_rate': 100,
-        'print_info_rate': 100,
-        'downsampling_factor': 2
+        'io_rate': 10000,
+        'print_info_rate': 10000,
+        'downsampling_factor': 1
     }
     sim = Cavity(**kwargs)
-    sim.run(2000)
\ No newline at end of file
+    sim.run(1000000)
\ No newline at end of file
diff --git a/examples/CFD/cylinder2d.py b/examples/CFD/cylinder2d.py
index 5120668..deade62 100644
--- a/examples/CFD/cylinder2d.py
+++ b/examples/CFD/cylinder2d.py
@@ -27,13 +27,14 @@
 from src.models import BGKSim, KBCSim
 import jax.numpy as jnp
 import os
+import json
 
 # Use 8 CPU devices
 # os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
 import jax
 jax.config.update('jax_enable_x64', True)
 
-class Cylinder(KBCSim):
+class Cylinder(BGKSim):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
@@ -45,8 +46,8 @@ def set_boundary_conditions(self):
         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(InterpolatedBounceBackLocal(tuple(cylinder.T), implicit_distance,
-                                               self.gridInfo, self.precisionPolicy))
+        self.BCs.append(InterpolatedBounceBackBouzidi(tuple(cylinder.T), implicit_distance, self.gridInfo, self.precisionPolicy))
+        # self.BCs.append(BounceBackHalfway(tuple(cylinder.T), self.gridInfo, self.precisionPolicy))
 
         # wall = np.concatenate([cylinder, self.boundingBoxIndices['top'], self.boundingBoxIndices['bottom']])
         # self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
@@ -76,53 +77,72 @@ def output_data(self, **kwargs):
         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)
+        if timestep == 0:
+            self.CL_max = 0.0
+            self.CD_max = 0.0
+        if timestep > 0.8 * t_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'
-    prescribed_vel = 0.005
-    diam = 80
-    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)
-
-    print('omega = ', omega)
-    print("Mesh size: ", nx, ny)
-    print("Number of voxels: ", nx * ny)
-    
-    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)
+    diam_list = [10, 20, 30, 40, 60, 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)
+
+        print('omega = ', omega)
+        print("Mesh size: ", nx, ny)
+        print("Number of voxels: ", nx * ny)
+        
+        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
+        }
+        sim = Cylinder(**kwargs)
+        t_max = int(1000000 / scale_factor)
+        sim.run(t_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/src/models.py b/src/models.py
index 7e5e825..7d0424c 100644
--- a/src/models.py
+++ b/src/models.py
@@ -68,6 +68,40 @@ def collision(self, f):
         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):
+        """
+        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, 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):

From 2871c87785f911fff54f79aa0e28b3aed33b4227 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Tue, 3 Oct 2023 09:40:14 -0400
Subject: [PATCH 4/8] the order of BC needed to change to get the absolute
 correct data

---
 examples/CFD/cavity3d.py | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/examples/CFD/cavity3d.py b/examples/CFD/cavity3d.py
index f7e35c8..0afa648 100644
--- a/examples/CFD/cavity3d.py
+++ b/examples/CFD/cavity3d.py
@@ -17,7 +17,7 @@
 # Use 8 CPU devices
 # os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
 from src.models import BGKSim, KBCSim
-from src.lattice import LatticeD3Q27
+from src.lattice import LatticeD3Q19
 from src.utils import *
 from jax.config import config
 from src.boundary_conditions import *
@@ -31,6 +31,13 @@ def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     def set_boundary_conditions(self):
+
+        # apply inlet equilibrium 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))
+
         # concatenate the indices of the left, right, and bottom walls
         walls = np.concatenate(
             (self.boundingBoxIndices['left'], self.boundingBoxIndices['right'],
@@ -38,14 +45,7 @@ def set_boundary_conditions(self):
              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(BounceBackHalfway(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, vel_wall))
+        return
 
     def output_data(self, **kwargs):
         # 1: -1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
@@ -64,8 +64,10 @@ def output_data(self, **kwargs):
 
         # output profiles of velocity at mid-plane for benchmarking
         output_filename = "./profiles_" + f"{timestep:07d}.json"
-        ldc_ref_result = {'ux(x=y=0)': list(u[nx//2, ny//2, :, 0]/prescribed_vel),
-                          'uz(z=y=0)': list(u[:, ny//2, nz//2, 2]/prescribed_vel)}
+        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,
@@ -76,7 +78,7 @@ def output_data(self, **kwargs):
         # live_volume_randering(timestep, u_mag)
 
 if __name__ == '__main__':
-    lattice = LatticeD3Q27(precision)
+    lattice = LatticeD3Q19(precision)
 
     nx = 256
     ny = 256

From 0ddb3a4a3ba6cd2228a04147631ed9916c0b1bd6 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Wed, 11 Oct 2023 17:26:03 -0400
Subject: [PATCH 5/8] making necessary changes to handle LDC at Re=10000

---
 examples/CFD/cavity3d.py | 32 ++++++++++++++++++++++----------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/examples/CFD/cavity3d.py b/examples/CFD/cavity3d.py
index 0afa648..ee6cd12 100644
--- a/examples/CFD/cavity3d.py
+++ b/examples/CFD/cavity3d.py
@@ -17,7 +17,7 @@
 # Use 8 CPU devices
 # os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
 from src.models import BGKSim, KBCSim
-from src.lattice import LatticeD3Q19
+from src.lattice import LatticeD3Q19, LatticeD3Q27
 from src.utils import *
 from jax.config import config
 from src.boundary_conditions import *
@@ -26,17 +26,21 @@
 precision = 'f64/f64'
 config.update('jax_enable_x64', True)
 
-class Cavity(BGKSim):
+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 equilibrium boundary condition to the top wall
+        # 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(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(
@@ -44,7 +48,9 @@ def set_boundary_conditions(self):
              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))
+        # 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):
@@ -78,16 +84,22 @@ def output_data(self, **kwargs):
         # live_volume_randering(timestep, u_mag)
 
 if __name__ == '__main__':
-    lattice = LatticeD3Q19(precision)
+    # Note: 
+    # We have used BGK with D3Q19 (or D3Q27) for Re=(1000, 3200) and KBC with D3Q27 for Re=10,000
+    lattice = LatticeD3Q27(precision)
 
     nx = 256
     ny = 256
     nz = 256
 
-    Re = 1000.0
+    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)
     print('omega = ', omega)
@@ -101,9 +113,9 @@ def output_data(self, **kwargs):
         'ny': ny,
         'nz': nz,
         'precision': precision,
-        'io_rate': 10000,
-        'print_info_rate': 10000,
+        'io_rate': int(10//tc),
+        'print_info_rate': int(10//tc),
         'downsampling_factor': 1
     }
     sim = Cavity(**kwargs)
-    sim.run(1000000)
\ No newline at end of file
+    sim.run(niter_max)
\ No newline at end of file

From 25fcd4cb0c20314bac10abd6004e162b430e6e5e Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Thu, 19 Oct 2023 10:52:41 -0400
Subject: [PATCH 6/8] fixed a bug in setting up IC for TG problem

---
 examples/CFD/taylor_green_vortex.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/CFD/taylor_green_vortex.py b/examples/CFD/taylor_green_vortex.py
index 1071025..f7019ee 100644
--- a/examples/CFD/taylor_green_vortex.py
+++ b/examples/CFD/taylor_green_vortex.py
@@ -48,7 +48,7 @@ def initialize_populations(self, rho, 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(int(20000*32/nx))
+        f = ADE.run(int(20000*nx/32))
         return f
 
     def output_data(self, **kwargs):

From 8dced2b8b293b886bbe398b27a15238efc503ff3 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Tue, 24 Oct 2023 13:20:58 -0400
Subject: [PATCH 7/8] refactored the interpolating BC to follow related class
 inheritance. Also added more comments.

---
 src/base.py                |   3 +-
 src/boundary_conditions.py | 173 ++++++++++++-------------------------
 2 files changed, 58 insertions(+), 118 deletions(-)

diff --git a/src/base.py b/src/base.py
index b359b6a..0ca82a7 100644
--- a/src/base.py
+++ b/src/base.py
@@ -11,17 +11,16 @@
 
 # JAX-related imports
 from jax import jit, lax, vmap
-from jax.config import config
 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
-import src.models
 from src.utils import downsample_field
 
 jax.config.update("jax_spmd_mode", 'allow_all')
diff --git a/src/boundary_conditions.py b/src/boundary_conditions.py
index 3b01ccb..2bb6119 100644
--- a/src/boundary_conditions.py
+++ b/src/boundary_conditions.py
@@ -498,14 +498,28 @@ def configure(self, boundaryBitmask):
         """
         # Perform index shift for halfway BB.
         hasFluidNeighbour = ~boundaryBitmask[:, 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):
         """
@@ -526,13 +540,10 @@ def apply(self, fout, fin):
         nbd = len(self.indices[0])
         bindex = np.arange(nbd)[:, None]
         fbd = fout[self.indices]
-        if self.vel is not None:
-            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].set(fin[self.indices][bindex, self.iknown] - cu[bindex, self.iknown])
-        else:
-            fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
 
+        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):
@@ -1012,10 +1023,10 @@ def apply(self, fout, fin):
         return fbd
 
 
-class InterpolatedBounceBackDifferentiable(BoundaryCondition):
+class InterpolatedBounceBackBouzidi(BounceBackHalfway):
     """
-    A local and differentiable single-node interpolated bounce-back boundary condition for a lattice Boltzmann method
-    simulation.
+    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.
@@ -1023,7 +1034,7 @@ class InterpolatedBounceBackDifferentiable(BoundaryCondition):
     Attributes
     ----------
     name : str
-        The name of the boundary condition. For this class, it is "InterpolatedBounceBackDifferentiable".
+        The name of the boundary condition. For this class, it is "InterpolatedBounceBackBouzidi".
     implementationStep : str
         The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
         it is "PostStreaming".
@@ -1035,43 +1046,12 @@ class InterpolatedBounceBackDifferentiable(BoundaryCondition):
         An array of shape (nx,ny,nz) indicating the signed-distance field from the solid walls
     """
 
-    def __init__(self, indices, implicit_distances, grid_info, precision_policy):
+    def __init__(self, indices, implicit_distances, grid_info, precision_policy, vel=None):
 
-        super().__init__(indices, grid_info, precision_policy)
-        self.name = "InterpolatedBounceBackDifferentiable"
-        self.implementationStep = "PostStreaming"
+        super().__init__(indices, grid_info, precision_policy, vel=vel)
+        self.name = "InterpolatedBounceBackBouzidi"
         self.implicit_distances = implicit_distances
         self.weights = None
-        self.isSolid = True
-        self.needsExtraConfiguration = True
-
-    def configure(self, boundaryBitmask):
-        """
-        Configures the boundary condition.
-
-        Parameters
-        ----------
-        boundaryBitmask : array-like
-            The connectivity bitmask 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 = ~boundaryBitmask[:, 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_new = np.unique(np.vstack(idx_trg), axis=0)
-        self.indices = tuple(indices_new.T)
-        return
 
     def set_proximity_ratio(self):
         """
@@ -1118,17 +1098,36 @@ def apply(self, fout, fin):
         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_postcollision_iknown +
-                    self.weights * (f_postcollision_imissing + f_poststreaming_iknown)) / (1.0 + self.weights)
+
+        # 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 InterpolatedBounceBackBouzidi(BoundaryCondition):
+class InterpolatedBounceBackDifferentiable(InterpolatedBounceBackBouzidi):
     """
-    A local single-node version of the interpolated bounce-back boundary condition due to Bouzidi for a lattice
-    Boltzmann method simulation.
+    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.
@@ -1136,7 +1135,7 @@ class InterpolatedBounceBackBouzidi(BoundaryCondition):
     Attributes
     ----------
     name : str
-        The name of the boundary condition. For this class, it is "InterpolatedBounceBackBouzidi".
+        The name of the boundary condition. For this class, it is "InterpolatedBounceBackDifferentiable".
     implementationStep : str
         The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
         it is "PostStreaming".
@@ -1148,63 +1147,11 @@ class InterpolatedBounceBackBouzidi(BoundaryCondition):
         An array of shape (nx,ny,nz) indicating the signed-distance field from the solid walls
     """
 
-    def __init__(self, indices, implicit_distances, grid_info, precision_policy):
+    def __init__(self, indices, implicit_distances, grid_info, precision_policy, vel=None):
 
-        super().__init__(indices, grid_info, precision_policy)
-        self.name = "InterpolatedBounceBackBouzidi"
-        self.implementationStep = "PostStreaming"
-        self.implicit_distances = implicit_distances
-        self.weights = None
-        self.isSolid = True
-        self.needsExtraConfiguration = True
-
-    def configure(self, boundaryBitmask):
-        """
-        Configures the boundary condition.
-
-        Parameters
-        ----------
-        boundaryBitmask : array-like
-            The connectivity bitmask 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 = ~boundaryBitmask[:, 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_new = np.unique(np.vstack(idx_trg), axis=0)
-        self.indices = tuple(indices_new.T)
-        return
-
-    def set_proximity_ratio(self):
-        """
-        Creates the interpolation data needed for the boundary condition.
+        super().__init__(indices, implicit_distances, grid_info, precision_policy, vel=vel)
+        self.name = "InterpolatedBounceBackDifferentiable"
 
-        Returns
-        -------
-        None. The function updates the object's weights attribute in place.
-        """
-        idx = np.array(self.indices).T
-        self.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]
-            mask = self.iknownBitmask[:, q]
-            self.weights[mask, q] = sdf_f[mask] / (sdf_f[mask] - sdf_s[mask])
-        return
 
     @partial(jit, static_argnums=(0,))
     def apply(self, fout, fin):
@@ -1231,16 +1178,10 @@ def apply(self, fout, fin):
         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)
+        fmissing = ((1. - self.weights) * f_postcollision_iknown +
+                    self.weights * (f_postcollision_imissing + f_poststreaming_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

From b101e16353b4d3f0976664ed9662aa351c122ef3 Mon Sep 17 00:00:00 2001
From: Hesam Salehipour <hesam.salehipour@autodesk.com>
Date: Thu, 23 Nov 2023 11:52:42 -0500
Subject: [PATCH 8/8] fixed a bug in InterpDiffBB. Added more comments.

---
 src/boundary_conditions.py | 24 ++++++------------------
 src/models.py              |  9 ++++++++-
 2 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/src/boundary_conditions.py b/src/boundary_conditions.py
index 2bb6119..8d0ad10 100644
--- a/src/boundary_conditions.py
+++ b/src/boundary_conditions.py
@@ -1035,15 +1035,12 @@ class InterpolatedBounceBackBouzidi(BounceBackHalfway):
     ----------
     name : str
         The name of the boundary condition. For this class, it is "InterpolatedBounceBackBouzidi".
-    implementationStep : str
-        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
-        it is "PostStreaming".
-    isSolid : bool
-        Whether the boundary condition represents a solid boundary. For this class, it is True.
-    needsExtraConfiguration : bool
-        Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True.
     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):
@@ -1136,15 +1133,6 @@ class InterpolatedBounceBackDifferentiable(InterpolatedBounceBackBouzidi):
     ----------
     name : str
         The name of the boundary condition. For this class, it is "InterpolatedBounceBackDifferentiable".
-    implementationStep : str
-        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
-        it is "PostStreaming".
-    isSolid : bool
-        Whether the boundary condition represents a solid boundary. For this class, it is True.
-    needsExtraConfiguration : bool
-        Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True.
-    implicit_distances : array-like
-        An array of shape (nx,ny,nz) indicating the signed-distance field from the solid walls
     """
 
     def __init__(self, indices, implicit_distances, grid_info, precision_policy, vel=None):
@@ -1178,8 +1166,8 @@ def apply(self, fout, fin):
         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_postcollision_iknown +
-                    self.weights * (f_postcollision_imissing + f_poststreaming_iknown)) / (1.0 + self.weights)
+        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:
diff --git a/src/models.py b/src/models.py
index be58185..a0500c8 100644
--- a/src/models.py
+++ b/src/models.py
@@ -73,7 +73,14 @@ def collision(self, f):
     @partial(jit, static_argnums=(0,), donate_argnums=(1,))
     def collision_modified(self, f):
         """
-        KBC collision step for lattice.
+        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