Examples of using kolmogorov_forcing in semi_implicit_navier_stokes

[chenwydj]

Dear researchers:

Do you happen to have any example of using kolmogorov_forcing in semi_implicit_navier_stokes ?

Naively passing kolmogorov_forcing into semi_implicit_navier_stokes will not work.
Mimic the usage in spectral will give errors like "TypeError: unsupported operand type(s) for /: 'function' and 'float'" or "ValueError: arguments have different tree structures"

Examples in https://github.com/google/jax-cfd/blob/main/notebooks/spectral_forced_turbulence.ipynb will work, but one has to manually convert vorticity into velocity.

Thank you very much!