fixed a bug in bc and kbc (Warp) that emerged after merging 2d and 3d…

setup.py

@@ -1,32 +1,30 @@
 from setuptools import setup, find_packages
 
 setup(
-    name='xlb',
-    version='0.2.1',
-    description='XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML',
-    long_description=open('README.md').read(),
-    long_description_content_type='text/markdown',
-    author='Mehdi Ataei',
-    url='https://github.com/Autodesk/XLB',
-    license='Apache License 2.0',
+    name="xlb",
+    version="0.2.1",
+    description="XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    author="XLB Team",
+    url="https://github.com/Autodesk/XLB",
+    license="Apache License 2.0",
     packages=find_packages(),
     install_requires=[
-        'matplotlib>=3.9.2',
-        'numpy>=2.1.2',
-        'pyvista>=0.44.1',
-        'trimesh>=4.4.9',
-        'warp-lang>=1.4.0',
-        'numpy-stl>=3.1.2',
-        'pydantic>=2.9.1',
-        'ruff>=0.6.5',
-        'jax>=0.4.34'  # Base JAX CPU-only requirement
+        "matplotlib>=3.9.2",
+        "numpy>=2.1.2",
+        "pyvista>=0.44.1",
+        "trimesh>=4.4.9",
+        "warp-lang>=1.4.0",
+        "numpy-stl>=3.1.2",
+        "pydantic>=2.9.1",
+        "ruff>=0.6.5",
+        "jax>=0.4.34",  # Base JAX CPU-only requirement
     ],
     extras_require={
-        'cuda': ['jax[cuda12]>=0.4.34'],  # For CUDA installations
-        'tpu': ['jax[tpu]>=0.4.34'],  # For TPU installations
+        "cuda": ["jax[cuda12]>=0.4.34"],  # For CUDA installations
+        "tpu": ["jax[tpu]>=0.4.34"],  # For TPU installations
     },
-    python_requires='>=3.10',
-    dependency_links=[
-        'https://storage.googleapis.com/jax-releases/libtpu_releases.html'
-    ],
+    python_requires=">=3.10",
+    dependency_links=["https://storage.googleapis.com/jax-releases/libtpu_releases.html"],
 )

tests/boundary_conditions/mask/test_bc_indices_masker_warp.py

@@ -60,7 +60,6 @@ def test_indices_masker_warp(dim, velocity_set, grid_shape):
         [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
 
@@ -69,9 +68,12 @@ def test_indices_masker_warp(dim, velocity_set, grid_shape):
     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 len(grid_shape) == 2:
+        assert bc_mask.shape == (1,) + grid_shape + (1,), "bc_mask shape is incorrect got {}".format(bc_mask.shape)
+        assert missing_mask.shape == (velocity_set.q,) + grid_shape + (1,), "missing_mask shape is incorrect got {}".format(missing_mask.shape)
+    else:
+        assert bc_mask.shape == (1,) + grid_shape, "bc_mask shape is incorrect got {}".format(bc_mask.shape)
+        assert missing_mask.shape == (velocity_set.q,) + grid_shape, "missing_mask shape is incorrect got {}".format(missing_mask.shape)
 
     if dim == 2:
         assert np.all(bc_mask[0, indices[0], indices[1]] == test_bc.id)

xlb/operator/boundary_condition/bc_do_nothing.py

@@ -76,4 +76,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
+        return f_post

xlb/operator/boundary_condition/bc_equilibrium.py

@@ -101,4 +101,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
+        return f_post

xlb/operator/boundary_condition/bc_extrapolation_outflow.py

@@ -205,4 +205,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
+        return f_post

xlb/operator/boundary_condition/bc_fullway_bounce_back.py

@@ -86,4 +86,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
+        return f_post

xlb/operator/boundary_condition/bc_grads_approximation.py

@@ -321,4 +321,4 @@ def warp_implementation(self, f_pre, f_post, bc_mask, missing_mask):
             inputs=[f_pre, f_post, bc_mask, missing_mask],
             dim=f_pre.shape[1:],
         )
-        return f_post
+        return f_post

xlb/operator/boundary_condition/bc_regularized.py

@@ -162,9 +162,14 @@ def _get_fsum(
         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])
+            if wp.static(_d == 3):
+                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])
+            else:
+                for l in range(_q):
+                    if missing_mask[l] == wp.uint8(1) and wp.abs(_c[0, l]) + wp.abs(_c[1, l]) == 1:
+                        return -_u_vec(_c_float[0, l], _c_float[1, l])
 
         @wp.func
         def bounceback_nonequilibrium(

xlb/operator/collision/kbc.py

@@ -271,7 +271,10 @@ def functional(
         ):
             # Compute shear and delta_s
             fneq = f - feq
-            shear = decompose_shear_d3q27(fneq)
+            if wp.static(self.velocity_set.d == 3):
+                shear = decompose_shear_d3q27(fneq)
+            else:
+                shear = decompose_shear_d2q9(fneq)
             delta_s = shear * rho  # TODO: Check this
 
             # Perform collision