Hello!

First, thanks for sharing this amazing project. Second, sorry if this isn’t the right place to ask this question.

I would like to know if there is an ongoing effort to develop a block structured NS-FV LES/WMLES code able to simulate academic relevant flows, i.e turbulent channel flow, NASA Hump etc.

If yes, could someone please point me to the right direction? If no, do you think the current code could be used as a starting point?

Best,
Eduardo

Upon following the instructions to setup a development environment, i.e:

git clone https://github.com/google/jax-cfd; cd jax-cfd
python -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install jaxlib
pip install -e ".[complete]"

and running the test suite, as per:

python -m pytest -n auto jax_cfd --dist=loadfile --ignore=jax_cfd/base/validation_test.py

the test-suite fails with:

11 failed, 447 passed, 242 warnings, 4 errors 

I've attached a full output of the tests here.

Quick note on running tests: the instructions in the README.md don't run if you set the environment up like I did above. The modification with python -m prefixing the command is required.

Are there any plans to fix these to ensure that all tests pass?

I have copied the code from spectral_forced_turbulence.ipynb, but it gives the following error:

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

UnfilteredStackTrace                      Traceback (most recent call last)

<timed exec> in <module>

[/usr/local/lib/python3.8/dist-packages/jax_cfd/base/initial_conditions.py](https://localhost:8080/#) in filtered_velocity_field(rng_key, grid, maximum_velocity, peak_wavenumber, iterations)
    110   # specified maximum velocity.
--> 111   return funcutils.repeated(project_and_normalize, iterations)(velocity)
    112 

28 frames

UnfilteredStackTrace: TypeError: jnp.linalg.norm requires ndarray or scalar arguments, got <class 'list'> at position 0.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

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


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

TypeError                                 Traceback (most recent call last)

<timed exec> in <module>

[/usr/local/lib/python3.8/dist-packages/jax/_src/numpy/util.py](https://localhost:8080/#) in _check_arraylike(fun_name, *args)
    343                     if not _arraylike(arg))
    344     msg = "{} requires ndarray or scalar arguments, got {} at position {}."
--> 345     raise TypeError(msg.format(fun_name, type(arg), pos))
    346 
    347 

TypeError: jnp.linalg.norm requires ndarray or scalar arguments, got <class 'list'> at position 0.

You can reproduce this with this Google Colab. It is also confusing that the Google Colab provided in the README does not install jax-cfd. What is the intended way of running them?

Edit: it seems that the jax-cfd version on PyPi is outdated. Downloading the source from the repository works.

Using the 2D channel flow notebook structure I have tried making some changes, starting by changing the Re of the flow, but it remains laminar. Is there a setting preventing the transition to turbulent that I am missing?

Thank you

When I run ml_model_inference_demo.ipynb, I have the following problems. It is obvious that the results are different from yours. What is the problem? Is there a connection to my environment?

Hi,

Thanks for this amazing library!
I installed the module using the command given in the README pip install jax-cfd[complete] to have access to all submodules.

When I try to import the spectral module via import jax_cfd.spectral as spectral I get the following error message:

import jax_cfd.spectral as spectral
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'jax_cfd.spectral'

Some more information about my system and the environment:

• Python version: Python 3.9.7 (default, Sep 16 2021, 13:09:58)
• JAX version: 0.2.21
• JAX CFD version: 0.1.0
• OS: Linux 5.10.16.3-microsoft-standard-WSL2

Is there any way that I can resolve this issue?

Thanks in advance!

Plot the x-velocity

ds.pipe(lambda ds: ds.u).plot.imshow(
'x', 'y', col='time',cmap=seaborn.cm.rocket, robust=True, col_wrap=4, aspect=2);

Solution:

Plot the x-velocity

ds.pipe(lambda ds: ds.u).plot.imshow(x='x',y='y',col='time',cmap=seaborn.cm.rocket, robust=True, col_wrap=4, aspect=2);

Jax.devicearray is no longer available in recent JAX distributions.

https://github.com/tequilahub/tequila#quantumchemistry

I'm trying to run channel_flow_demo.ipynb, but I get the following error when I call the step_fn :

(massive traceback truncated)

/usr/local/lib/python3.7/dist-packages/jax_cfd/base/advection.py in (.0)
107 for u, target_offset in zip(v, target_offsets))
108 aligned_c = tuple(c_interpolation_fn(c, target_offset, aligned_v, dt)
--> 109 for target_offset in target_offsets)
110 return _advect_aligned(aligned_c, aligned_v)
111

/usr/local/lib/python3.7/dist-packages/jax_cfd/base/interpolation.py in tvd_interpolation(c, offset, v, dt)
282 c_left = c.shift(-1, axis)
283 c_right = c.shift(1, axis)
--> 284 c_next_right = c.shift(2, axis)
285 # Velocities of different sign are evaluated with limiters at different
286 # points. See equations (4.34) -- (4.39) from the reference above.

/usr/local/lib/python3.7/dist-packages/jax_cfd/base/grids.py in shift(self, offset, axis)
269 GridArray has offset u.offset + offset.
270 """
--> 271 return self.bc.shift(self.array, offset, axis)
272
273 def _interior_grid(self) -> Grid:

/usr/local/lib/python3.7/dist-packages/jax_cfd/base/boundaries.py in shift(self, u, offset, axis)
77 u.offset + offset.
78 """
---> 79 padded = self._pad(u, offset, axis)
80 trimmed = self._trim(padded, -offset, axis)
81 return trimmed

/usr/local/lib/python3.7/dist-packages/jax_cfd/base/boundaries.py in _pad(self, u, width, axis)
116 if bc_type != BCType.PERIODIC and abs(width) > 1:
117 raise ValueError(
--> 118 'Padding past 1 ghost cell is not defined in nonperiodic case.')
119
120 if bc_type == BCType.PERIODIC:

ValueError: Padding past 1 ghost cell is not defined in nonperiodic case.

I installed JAX-CFD on a fresh Colab notebook using pip install git+https://github.com/google/jax-cfd.git

Please help me fix this. Let me know if I'm using the wrong version (since the repo is constantly being developed).

Thanks!

Hi,

I tried to define some boundary conditions. Until recently, everything went perfectly fine with

import jax_cfd.base.grids as grids
bc = grids.BoundaryConditions((grids.PERIODIC, grids.PERIODIC))

Traceback (most recent call last):
  File "<string>", line 1, in <module>
AttributeError: module 'jax_cfd.base.grids' has no attribute 'BoundaryConditions'

I saw that you moved the boundary conditions to a new module called boundaries.py. I tried then to access them accordingly from the new module via:

import jax_cfd.base as cfd
bc = cfd.boundaries.periodic_boundary_conditions(ndim=2)

Traceback (most recent call last):
  File "<string>", line 1, in <module>
AttributeError: module 'jax_cfd.base' has no attribute 'boundaries'

But now, neither of these work anymore (see the appended error messages) and I wonder how I can define boundary conditions.
I tried installing JAX-CFD according to the documentation and according to #90, but neither of them gave me a solution to my problem.

My python version is 3.8.

Python 3.8.12 (default, Oct 12 2021, 13:49:34) 
[GCC 7.5.0] :: Anaconda, Inc. on linux

I have run your program through your ipynb, but I can't find where is the data input? Are the nc and kpl files downloaded through gsutils? I would like to try out some other simple examples with your method. In addition, is the method of using machine learning in base/ml? I see that ml_model_inference_demo.ipynb uses ml, I hope to get your help, thank you very much.

Hello,

This is a very nice JAX implementation for CFDs! I have a few questions about your latest released notebook spectral_forced_turbulence.ipynb.

1. What is the exact NS equation that the notebook aims to solve? Would it be possible to provide some detailed expressions of NS equations and boundary conditions? Or would it be possible to provide some references? I think It will be super helpful to people like me who are more familiar with machine learning and know less about CFD.

2. If I understand correctly, I think the code aims to solve the NS equation (vorticity-velocity form) using the pseudo-spectral method. So we need to enforce periodic boundary conditions for the velocity (or vorticity?) However, when I try to run the following piece of code (which should give me a random velocity field and visualize it at the boundary)

v0 = cfd.initial_conditions.filtered_velocity_field(jax.random.PRNGKey(0), grid, max_velocity, 4)
u = v0[0].data
v = v0[1].data

plt.plot(u[:, 0])
plt.plot(u[:, -1])
plt.show()

plt.plot(u[0, :])
plt.plot(u[-1, :])
plt.show()

I got the following figures below. It seems that the periodic boundary condition is not fully imposed. Could you please take a look or correct me if I am wrong?

Once again, this is really a project! Looking forward to your reply.

When running layers_util_test.py with jax[CPU] version later than 0.4.1, it throws a key error 'experimental_xmap_spmd_lowering'.

I haven't read through the paper yet but I figured I'm inquire, is this research limited in scope to just 2d or are tensor based calculations appropriate and practical for 3d modeling as well?

Hi! I have a problem and sincerely hope to get your help!
Problem description: I wrote a python program myself to implement some functions of the project, but I encountered a problem in the first step.
In grids.py, line 66, grid: Grid, I didn't see the definition or import of Grid. In the end, I implemented the def velocity_trajectory_to_xarray() function of xarray_utils.py of the data submodule in the python file I wrote and reported an error " NameError: name 'Grid' is not defined", what should I do about this?

Hi,

I'm trying to understand the logic behind the various functions in ml.train_utils.

in the definition of loss_and_gradient the returned _loss function uses the target trajectory as part of the input of the trajectory_fn. I cannot make sense of it. The description of the function states that the trajectory_fn should accepts params and initial_velocity, which make sense to me, but then I don't understand why would we want to use the target_trajectory as the initial_velocity (which presumably doesn't have the same shape as initial_velocity since one is a trajectory and the other the velocity at a single time step).

I'm pretty much stuck here because I now don't know how to use this in the context of defining a train_step function etc.

Best regards,

Gwen

Hi! This is a great project, and I'm a big fan of both the machine learning applications here and also some of the smaller, helpful structures, in particular base.grids.

Currently, it is possible to add two GridArrays, but it is not possible to add two GridVariables. So this works fine:

import jax_cfd.base.grids as gd
import jax.numpy as jnp

grid = gd.Grid([4,], domain = [(0, 1),])

array_of_values = jnp.array([2.0, 2.0, 3.0, 4.0])

centered_array = grid.center(array_of_values)

print(centered_array + centered_array)

But this throws an exception:

bc = gd.BoundaryConditions((gd.PERIODIC,))

centered_variable = gd.GridVariable(centered, bc)

print(centered_variable + centered_variable)

I'm happy to have a go at implementing this myself, if someone isn't already working on it.

Also, am I correct in thinking that the way to use a JaxNumPy function on a GridArray is to call it via NumPy? For example, this throws an exception:

print(jnp.abs(centered_array))

But this works:

import numpy as np
print(np.abs(centered_array))

I assume it's implemented this way because NumPy has an automatic mix-in that we can employ to funnel things to the appropriate JaxNumPy function, but JaxNumPy does not.

The readme says that an example notebook showing how to train a simple hybrid ML + CFD model is being prepared. Is there an ETA on this?
I'm trying to get it to work using JAX-CFD and Flax but I'm running into problems. An example notebook would help me a lot!

Any plan to move to colab?

Hi, I'm so excited to see your work, it's so exciting! I want to implement your project on my own platform system, but since there is no jaxlib source code on the python official website, the compiled jaxlib is not suitable for my working platform. I would be very grateful if you can provide the source code of jaxlib