autodesk/xlb

XLB: Accelerated Lattice Boltzmann (XLB) based on JAX for Physics-based ML

xlb's Issues

Significant slowdown jaxlib 0.4.16 vs 0.4.14

In XLB library the MLUPs (Millions of Lattice Updates per Second) has dropped by about 15% after updating jaxlib from 0.4.14 to 0.4.16.

Results on a single RTX 6000 Ada:

Version 0.4.14:

omega =  0.4918839153959666
XLA backend: gpu
Number of XLA devices available: 1
WARNING: Checkpointing is disabled for this simulation.
Time to create the grid connectivity bitmask: 0.17581534385681152
Time to create the local bitmasks and normal arrays: 6.792882680892944
WARNING: Default initial conditions assumed: density = 1, velocity = 0
         To set explicit initial density and velocity, use self.initialize_macroscopic_fields.
Domain: 512 x 512 x 512
Number of voxels: 134217728
MLUPS: 827.4825740398888

Version 0.4.14:

WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1695236429.981993  393900 tfrt_cpu_pjrt_client.cc:349] TfrtCpuClient created.
omega =  0.4918839153959666
XLA backend: gpu
Number of XLA devices available: 1
WARNING: Checkpointing is disabled for this simulation.
Time to create the grid connectivity bitmask: 0.17301297187805176
Time to create the local bitmasks and normal arrays: 6.891621112823486
WARNING: Default initial conditions assumed: density = 1, velocity = 0
         To set explicit initial density and velocity, use self.initialize_macroscopic_fields.
W0000 00:00:1695236443.334162  393900 hlo_rematerialization.cc:2946] Can't reduce memory use below 27.08GiB (29077237923 bytes) by rematerialization; only reduced to 29.54GiB (31714181580 bytes), down from 29.54GiB (31714181580 bytes) originally
Domain: 512 x 512 x 512
Number of voxels: 134217728
MLUPS: 726.4662730520944
I0000 00:00:1695236480.893240  393900 tfrt_cpu_pjrt_client.cc:352] TfrtCpuClient destroyed.

Notice the new warnings tfrt_cpu_pjrt_client and hlo_rematerialization generated in version 0.4.16.

Step to reproduce:
Follow the instructions in the library and run:

python examples/performance/MLUPS3d.py 256 200

(the first input in the number of voxels in each dimension and the second number is the number of iterations)

"AttributeError: Unrecognized config option: jax_array" in windtunnel3d.py example

Line 33 jax.config.update('jax_array', True) in windtunnel3d.py throws an exception. Looks it should be removed.

Executing the windtunnel3d.py example in a Google Colab instance throws this exception:

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

[/content/XLB/examples/CFD/windtunnel3d.py](https://localhost:8080/#) in <module>
     31 # disable JIt compilation
     32 
---> 33 jax.config.update('jax_array', True)
     34 
     35 class Car(KBCSim):

[/usr/local/lib/python3.10/dist-packages/jax/_src/config.py](https://localhost:8080/#) in update(self, name, val)
     85   def update(self, name, val):
     86     if name not in self._value_holders:
---> 87       raise AttributeError(f"Unrecognized config option: {name}")
     88     self._value_holders[name]._set(val)
     89 

AttributeError: Unrecognized config option: jax_array

This happens on an L4 instance. I install dependencies by executing, in Colab notebook cells:

!pip install pyvista numpy matplotlib Rtree trimesh jmp orbax-checkpoint termcolor

And then cloning the repository and executing the example:

!git clone https://github.com/Autodesk/XLB

%run XLB/examples/CFD/windtunnel3d.py

Looks like the operation in line 33 was removed after the array migration in JAX 0.4.1, per: https://jax.readthedocs.io/en/latest/jax_array_migration.html and this line should just be removed. Thanks!

The demo occurs in paper

The demo in companying paper (6.2. Differentiable flow control with deep learning) very interests me, but I don't find in example/ folder, could you please give me relevant file?

About the Physics-based machine learning demos

Hello there!

Firstly, I would like to express my gratitude for developing such an invaluable tool.

Currently I'm trying to code the LBM with machine learning by XLB, but I do not know how. And I read form the paper about the Physics-based machine learning demos: Enhancing coarse-grained simulations using deep learning correctors and Differentiable flow control with deep learning, but I do not find any code about them in the repo.

Could you possibly provide some guidance on how to implement these concepts? Alternatively, another simple demo illustrating the integration of machine learning with LBM would be immensely helpful.

Problem applying pressure based boundary condition using ZouHe

If I want to implement inlet or outlet boundary condition to a specific pressure using ZouHe, it shows a TypeError : mul got incompatible shapes for broadcasting: (998,), (9,). I used nx= 1400, ny= 1000. I tried to apply a specific pressure in left vertical inlet wall ( x=0, y=0:1000 ) . Same problem arrives if I want to apply a specific pressure in right outlet boundary.

Outflow BC

outlet = self.boundingBoxIndices['right']
rho_outlet = np.ones(outlet.shape[0], dtype=self.precisionPolicy.compute_dtype)
self.BCs.append(ZouHe(tuple(outlet.T), self.gridInfo, self.precisionPolicy, 'pressure', rho_outlet))

**** Simulation Parameters for Cylinder ****
Parameter | Value

            Omega | 0.5518611517342236
 Grid Points in X | 1400
 Grid Points in Y | 1000
 Grid Points in Z | 0
   Dimensionality | 2
 Precision Policy | f64/f64
     Lattice Type | D2Q9
  Checkpoint Rate | 0

Checkpoint Directory | ./checkpoints
Downsampling Factor | 1
Print Info Rate | 200
I/O Rate | 100
Compute MLUPS | False
Restore Checkpoint | False
Backend | gpu
Number of Devices | 1
Volumetric Ratio of Solid : 0.502656
Time to create the grid mask: 0.3005838394165039
Time to create the local masks and normal arrays: 11.985322952270508
WARNING: Default initial conditions assumed: density = 1, velocity = 0
To set explicit initial density and velocity, use self.initialize_macroscopic_fields.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/mnt/c/RunCodeHere/Projects/XLB/Permeability02_benchmark.py", line 297, in
sim.run(t_max)
File "/mnt/c/RunCodeHere/Projects/XLB/src/base.py", line 905, in run
f, fstar = self.step(f, timestep, return_fpost=self.returnFpost)
File "/mnt/c/RunCodeHere/Projects/XLB/src/base.py", line 841, in step
f_poststreaming = self.apply_bc(f_poststreaming, f_postcollision, timestep, "PostStreaming")
File "/mnt/c/RunCodeHere/Projects/XLB/src/base.py", line 804, in apply_bc
fout = fout.at[bc.indices].set(bc.apply(fout, fin))
File "/mnt/c/RunCodeHere/Projects/XLB/src/boundary_conditions.py", line 781, in apply
feq = self.calculate_equilibrium(fout)
File "/mnt/c/RunCodeHere/Projects/XLB/src/boundary_conditions.py", line 738, in calculate_equilibrium
feq = self.equilibrium(rho, vel)
File "/mnt/c/RunCodeHere/Projects/XLB/src/boundary_conditions.py", line 285, in equilibrium
feq = rho * self.lattice.w * (1.0 + 1.0 * cu + 0.5 * cu**2 - usqr)
File "/home/shouvik/anaconda3/lib/python3.9/site-packages/jax/src/numpy/array_methods.py", line 743, in op
return getattr(self.aval, f"{name}")(self, *args)
File "/home/shouvik/anaconda3/lib/python3.9/site-packages/jax/_src/numpy/array_methods.py", line 271, in deferring_binary_op
return binary_op(*args)
File "/home/shouvik/anaconda3/lib/python3.9/site-packages/jax/src/numpy/ufuncs.py", line 99, in fn
return lax_fn(x1, x2) if x1.dtype != np.bool else bool_lax_fn(x1, x2)
TypeError: mul got incompatible shapes for broadcasting: (998,), (9,).

example program running error

python examples/CFD/cavity2d.py
'rm' 不是内部或外部命令，也不是可运行的程序
或批处理文件。
WARNING:absl:Configured CheckpointManager using deprecated legacy API. Please follow the instructions at https://orbax.readthedocs.io/en/latest/api_refactor.html to migrate by August 1st, 2024.
**** Simulation Parameters for Cavity ****
Parameter | Value

            Omega | 1.2523481527864746
 Grid Points in X | 200
 Grid Points in Y | 200
 Grid Points in Z | 0
   Dimensionality | 2
 Precision Policy | f32/f32
     Lattice Type | D2Q9
  Checkpoint Rate | 1000

Checkpoint Directory | C:\Users\ZZL\Desktop\temp\XLB\checkpoints
Downsampling Factor | 1
Print Info Rate | 100
I/O Rate | 100
Compute MLUPS | False
Restore Checkpoint | False
Backend | cpu
Number of Devices | 1
Time to create the grid mask: 0.11803746223449707
Time to create the local masks and normal arrays: 0.3295261859893799
WARNING: Default initial conditions assumed: density = 1, velocity = 0
To set explicit initial density and velocity, use self.initialize_macroscopic_fields.
Timestep 0 of 5000 completed
Saving data at timestep 0/5000
Saved .\fields_0000000.vtk in 0.003017 seconds.
Saved .\BCs_0000000.vtk in 0.000519 seconds.
Saving checkpoint at timestep 0/5000
WARNING:absl:Attempted to create temporary directory C:\Users\ZZL\Desktop\temp\XLB\checkpoints\0.orbax-checkpoint-tmp-0 which already exists. Removing existing directory since it is not finalized.
Traceback (most recent call last):
File "C:\Users\ZZL\Desktop\temp\XLB\examples\CFD\cavity2d.py", line 96, in
sim.run(5000)
File "C:\Users\ZZL\Desktop\temp\XLB\src\base.py", line 928, in run
self.mngr.save(timestep, state)
File "D:\aruanjian\Python311\Lib\site-packages\orbax\checkpoint\checkpoint_manager.py", line 1110, in save
self._checkpointer.save(save_directory, args=args)
File "D:\aruanjian\Python311\Lib\site-packages\orbax\checkpoint\checkpointer.py", line 186, in save
self._handler.finalize(tmpdir)
File "D:\aruanjian\Python311\Lib\site-packages\orbax\checkpoint\composite_checkpoint_handler.py", line 514, in finalize
handler.finalize(self._get_item_directory(directory, item_name))
File "D:\aruanjian\Python311\Lib\site-packages\orbax\checkpoint\pytree_checkpoint_handler.py", line 759, in finalize
self._handler_impl.finalize(directory)
File "D:\aruanjian\Python311\Lib\site-packages\orbax\checkpoint\base_pytree_checkpoint_handler.py", line 1064, in finalize
type_handlers.merge_ocdbt_per_process_files(directory)
File "D:\aruanjian\Python311\Lib\site-packages\orbax\checkpoint\type_handlers.py", line 670, in merge_ocdbt_per_process_files
asyncio.run(open_and_copy())
File "D:\aruanjian\Python311\Lib\site-packages\nest_asyncio.py", line 30, in run
return loop.run_until_complete(task)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\aruanjian\Python311\Lib\site-packages\nest_asyncio.py", line 98, in run_until_complete
return f.result()
^^^^^^^^^^
File "D:\aruanjian\Python311\Lib\asyncio\futures.py", line 203, in result
raise self._exception.with_traceback(self._exception_tb)
File "D:\aruanjian\Python311\Lib\asyncio\tasks.py", line 279, in __step
result = coro.throw(exc)
^^^^^^^^^^^^^^^
File "D:\aruanjian\Python311\Lib\site-packages\orbax\checkpoint\type_handlers.py", line 667, in open_and_copy
await asyncio.gather(*copy_ops)
File "D:\aruanjian\Python311\Lib\asyncio\tasks.py", line 349, in __wakeup
future.result()
File "D:\aruanjian\Python311\Lib\asyncio\tasks.py", line 279, in __step
result = coro.throw(exc)
^^^^^^^^^^^^^^^
File "D:\aruanjian\Python311\Lib\asyncio\tasks.py", line 694, in _wrap_awaitable
return (yield from awaitable.await())
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "D:\aruanjian\Python311\Lib\asyncio\futures.py", line 287, in await
yield self # This tells Task to wait for completion.
^^^^^^^^^^
File "D:\aruanjian\Python311\Lib\asyncio\tasks.py", line 349, in __wakeup
future.result()
File "D:\aruanjian\Python311\Lib\asyncio\futures.py", line 203, in result
raise self._exception.with_traceback(self._exception_tb)
NotImplementedError: CopyRange not supported

Delete the pull request that I accidentally created

Hi,
I accidentally created a pull request when working on my fork branch. Can you guys help me delete that because I'm working on behave of my company and the code I wrote might involve copyright issues.
Thanks!

folder stl-files missing

When executing windtunnel3d.py, get the following issue:

ValueError: string is not a file: stl-files/DrivAer-Notchback.stl

there is no folder stl-files, or any stl file elsewhere.

assign_fields_sharded with pre-defined initialization

The assign_fields_sharded doesn't work correctly if we define initial rho0 and u0 on multi-GPU systems. It is a simple fix and I'll submit a PR for it. Adding it here so that I don't forget about it.

Performance numbers

Any numbers we can see for performance. Been interested to see if a Jax LBM implementation can get close to optimal performance. Lettuce LBM is around 20x slower for example, https://github.com/lettucecfd/lettuce

cavity2d.py NotImplementedError: CopyRange not supported

please tell how to fix it,thanks in advance

Traceback (most recent call last):
File "e:\code\jupyter\XLB\setup.py", line 99, in
sim.run(5000)
File "e:\code\jupyter\XLB\src\base.py", line 928, in run
self.mngr.save(timestep, state)
File "D:\SomeApps\Miniconda3\envs\python39\lib\site-packages\orbax\checkpoint\checkpoint_manager.py", line 515, in save
self._checkpointers[k].save(item_dir, item, **kwargs)
File "D:\SomeApps\Miniconda3\envs\python39\lib\site-packages\orbax\checkpoint\checkpointer.py", line 150, in save
self._handler.finalize(tmpdir)
File "D:\SomeApps\Miniconda3\envs\python39\lib\site-packages\orbax\checkpoint\pytree_checkpoint_handler.py", line 1411, in finalize
type_handlers.merge_ocdbt_per_process_files(directory)
File "D:\SomeApps\Miniconda3\envs\python39\lib\site-packages\orbax\checkpoint\type_handlers.py", line 744, in merge_ocdbt_per_process_files
asyncio.run(open_and_copy())
File "D:\SomeApps\Miniconda3\envs\python39\lib\asyncio\runners.py", line 44, in run
return loop.run_until_complete(main)
File "D:\SomeApps\Miniconda3\envs\python39\lib\asyncio\base_events.py", line 647, in run_until_complete
return future.result()
File "D:\SomeApps\Miniconda3\envs\python39\lib\site-packages\orbax\checkpoint\type_handlers.py", line 742, in open_and_copy
await asyncio.gather(*copy_ops)
File "D:\SomeApps\Miniconda3\envs\python39\lib\asyncio\tasks.py", line 688, in _wrap_awaitable
return (yield from awaitable.await())
NotImplementedError: CopyRange not supported