Fixed the issues and better handle the lattice cases

examples/CFD/windtunnel3d.py

@@ -110,7 +110,7 @@ def output_data(self, **kwargs):
 
 if __name__ == '__main__':
     precision = 'f32/f32'
-    lattice = LatticeD3Q19(precision)
+    lattice = LatticeD3Q27(precision)
 
     nx = 601
     ny = 351

src/base.py

@@ -23,7 +23,6 @@
 # Local/Custom Libraries
 import src.models
 from src.utils import downsample_field
-from src.lattice import LatticeD2Q9, LatticeD3Q19, LatticeD3Q27
 
 jax.config.update("jax_spmd_mode", 'allow_all')
 # Disables annoying TF warnings
@@ -160,10 +159,10 @@ def lattice(self):
     def lattice(self, value):
         if value is None:
             raise ValueError("Lattice type must be provided.")
-        if self.nz == 0 and not isinstance(value, LatticeD2Q9):
+        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 isinstance(self, src.models.KBCSim) and not isinstance(value, LatticeD3Q27):
-            raise ValueError("For 3D KBC simulations, lattice type must be LatticeD3Q19,")
+        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
 

src/models.py

@@ -39,8 +39,9 @@ class KBCSim(LBMBase):
 
     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,))