Fixed Smagorinsky implementation and added JAX backend version

[mehdiataei]

Contributing Guidelines

Fixes #149

• Fixed Warp implementation: the kernel called self._warp_u_vec(), but that method doesn’t exist anywhere in the codebase.
  It was also hard-coded to 3D. the strain logic accessed _c[2, l] and assumed d == 3, so it can’t work for 2D lattices.
  The strain classification was also incorrect. It used (c_x + c_y + c_z) to decide “diagonal vs off-diagonal”. With signed lattice velocities (−1/0/1), that condition does not reliably represent |c| so the test was not passing. The method wrote _fout[l] directly to fout instead of casting to self.store_dtype, unlike BGK/KBC Warp kernels.

• Added JAX implementation for Smagorinsky

• Added tests for both implementations

examples run with both backends with no issues (e.g. run lid driven cavity).

• I have read and understood the CONTRIBUTING.md guidelines

Description

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

How Has This Been Tested?

• All pytest tests pass

Linting and Code Formatting

Make sure the code follows the project's linting and formatting standards. This project uses Ruff for linting.

To run Ruff, execute the following command from the root of the repository:

ruff check .

• Ruff passes