Implement directionally-averaged, scale-invariant dynamic Smagorinsky closure

[tomchor]

Closes https://github.com/CliMA/Oceananigans.jl/issues/3637

This PR implements the scale-invariant dynamic Smagorinsky. My main reference for this was Bou-Zeid et al. (2005), just because it's nicely summarized there.

As a refresher, the coefficient in this case is calculated as

c_s^2 = \frac{\langle L_{ij} M_{ij}\rangle}{\langle M_{ij} M_{ij} \rangle},

where $M_{ij}$ and $L_{ij}$ are obtained by applying the Germano identity and minimizing the error. $\langle \cdot \rangle$ is an average which usually is implemented as planar averaging. I'm implementing it here with an arbitrary DirectionalAveraging procedure that can average in whatever direction the user wants. Following a suggestion by @glwagner in #3637, I'm also starting to put the pieces in place for a more general averaging procedure so that eventually we can add Lagrangian averaging as an option for this model.

Note that this PR is related to https://github.com/CliMA/Oceananigans.jl/pull/3638. The rationale for having this PR in addition to #3638 is that for most LES are doubly periodic, in which case there's not much advantage in having the Lagrangian averaging that's implemented in #3638. In these cases I'd argue the model implemented here is more efficient, so having both types of models is probably desirable.

Another note is that, once this PR is merged, it should be very straightforward to implement a scale-dependent version, which works better close to boundaries.

The model seems to be working. Here's an animation of an unforced 3D turbulence simulation on a 32³ grid. Left is vorticity and right is the strain rate modulus squared:

https://github.com/CliMA/Oceananigans.jl/assets/13205162/7131a99d-df6c-4883-850d-d4a87988cdb7

Note that the value of the calculated Smag coefficient $c_s$ (which I'm showing at the top of the plots above) is about 0.17, which is very close to the theoretical value obtained by Lilly of 0.16.

I'm opening this as a draft PR for now because there are some things that need doing:

• [ ] Generalize the filter for stretched grids. For now it assumes a regular grid for simplicity, but it's trivial to generalize.
• [ ] Optimize the calculation of the coefficient. At the moment I'm creating four extra fields in order to calculate the Smag coefficient: LM, MM are 3D fields; and LM_avg and MM_avg are 1D or 2D. What I'm doing is to first calculate $L_{ij} M_{ij}$ and $M_{ij} M_{ij}$ pointwise, and then LM_avg and MM_avg receive their averages. We should be able to calculate everything without needing LM, MM are 3D fields, I just couldn't figure out how yet :)
• [ ] Write docs
• [ ] Write tests
• [ ] Validate that model is working as intended

CC @glwagner @simone-silvestri @xkykai @whitleyv

Feel free to add more things to the to-do list that I may have forgotten.

[glwagner]

Not sure if I am missing something but first I think the user can just write

averaging=(1, 2)

which would mean average over xy. Or if you want to be more verbose then

averaging=DimensionAveraging(dims=(1, 2))

or something.

Next for the average itself it seems you need

LM_op = KernelFunctionOperation{Center, Center, Center}(LᵢⱼMᵢⱼ_ccc, grid, u, v, w)
LM_avg = Field(Average(LM_op, averaging.dims))

In the constructor. Then you can just call compute!(LM_avg) without needing lots of additional dispatch.

[tomchor]

Not sure if I am missing something but first I think the user can just write

averaging=(1, 2)

I considered taking that approach but I felt it made it a bit awkward to implement Lagrangian averaging in the future. That said, I do agree that my currently implementation is probably unnecessarily verbose. I think your suggestion below of DimensionAveraging(dims=(1, 2)) is a good compromise.

which would mean average over xy. Or if you want to be more verbose then

averaging=DimensionAveraging(dims=(1, 2))

or something.

Next for the average itself it seems you need

LM_op = KernelFunctionOperation{Center, Center, Center}(LᵢⱼMᵢⱼ_ccc, grid, u, v, w)
LM_avg = Field(Average(LM_op, averaging.dims))

In the constructor. Then you can just call compute!(LM_avg) without needing lots of additional dispatch.

Ah, that's so simple! I'll implement that.

[glwagner]

Not sure if I am missing something but first I think the user can just write

averaging=(1, 2)

I considered taking that approach but I felt it made it a bit awkward to implement Lagrangian averaging in the future. That said, I do agree that my currently implementation is probably unnecessarily verbose. I think your suggestion below of DimensionAveraging(dims=(1, 2)) is a good compromise.

My suggestion is an API thing. If the user writes averaging = (1, 2) then you can interpret this as DimensionAveraging((1, 2)). Saves some characters.

Implementing this requires writing a function like

averaging(dims::Tuple) = DimenisionAveraging(dims)
averaging(other_averaging) = other_averaging

Also just brainstorming... maybe a constructor

DimensionAveraging(dims...) = DimensionAveraging(tuple(dims))
DimensionAveraging(dims::Colon) = DimensionAveraging(:)

will save some characters because arguably DimensionAveraging(1, 2) is easier to read than DimensionAveraging((1, 2)).

I think we are a little keyword-happy in the API, sometimes its nice to user positional arguments when meanings are obvious...

[tomchor]

Here's some validation. The animation below is from a triply-periodic 3D simulation with an initial reasonably-resolved initial condition without any forcings. Note that the IC looks a bit weird right now (a couple of straight lines in strain) because I'm taking a shortcut in defining a well-resolved random IC, but I can easily change it later if necessary.

Left panel shows the strain rate squared for the SmagorinskyLilly closure, while the one on the right shows results for the ScaleInvarianteSmagorinsky closure, implemented in this PR. The bottom panel shows the evolution of the dynamically-calculated Smagorinsky coefficient for the latter, in comparison with the constant value of 0.16 imposed on the former.

Important here is that the value of 0.16 was analytically derived by Lilly by assuming an isotropic 3D turbulence, a Kolmogorovian energy cascade, and further assuming that the cut-off filter is in the inertial range. I think all these assumptions are valid in this simulation, so we expect the dynamically-calculated value of the Smagorinsky coefficient (the black line) to approach the value 0.16 as time goes on.

https://github.com/CliMA/Oceananigans.jl/assets/13205162/4049e7cf-452e-4883-a709-a675cf12277c

Although the match is not exact (the value it approaches is ~0.17), I think this is close enough. That said, I'm planning on also implementing a boundary layer validation along similar lines, which we can use to validate the model in the same fashion as Bou-Zeid et al. (2005).

One thing to note is that the current implementation appears to be very slow. While the simulation with the SmagorinskyLilly closure runs on my laptop in 10 seconds, it takes 4 minutes for the simulation with the ScaleInvariantSmagorinsky. I know the dynamic model will be slower given the extra computations, but such a difference seems large to me, so I'm hoping something can be changed here to improve performance:

┌ Info: Running
└   closure = SmagorinskyLilly: C=0.16, Cb=1.0, Pr=1.0
[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (2.738 seconds)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (4.852 seconds).
[ Info: Simulation is stopping after running for 11.657 seconds.
[ Info: Simulation time 1.333 minutes equals or exceeds stop time 1.333 minutes.
┌ Info: Running
└   closure = ScaleInvariantSmagorinsky: averaging=Averaging over directions Colon(), Pr=1.0
[ Info: Initializing simulation...
[ Info:     ... simulation initialization complete (2.195 seconds)
[ Info: Executing initial time step...
[ Info:     ... initial time step complete (2.812 seconds).
[ Info: Simulation is stopping after running for 3.735 minutes.
[ Info: Simulation time 1.333 minutes equals or exceeds stop time 1.333 minutes.

[tomchor]

It's also worth noting that right now many calculations are done more than once in each timestep. For example for each component of $M_{ij}$ I'm calculating the whole strain rate tensor modulus in addition to the strain rate tensor component needed:

https://github.com/CliMA/Oceananigans.jl/blob/25cc34e6c395e210e0aecf8181919c25435d7919/src/TurbulenceClosures/turbulence_closure_implementations/scale_invariant_smagorinsky.jl#L241-L257

This is done for legibility of the code, but it may be necessary to forfeit that in favor of doing fewer calculations. (Also note that I'm using a weird way to define function names here, so lmk if you guys think I should change it.)

Another thing to note that it's common to update dynamic Smagorinsky coefficients once every 5 or so time-steps only, since they can be pretty expensive. afaik this is generally done for the scale-dependent versions, which have two test filters instead of the one needed in this PR, but I wouldn't be surprised if it's occasionally necessary for the scale-invariant versions as well.

[tomchor]

One thing to note is that the current implementation appears to be very slow. While the simulation with the SmagorinskyLilly closure runs on my laptop in 10 seconds, it takes 4 minutes for the simulation with the ScaleInvariantSmagorinsky. I know the dynamic model will be slower given the extra computations, but such a difference seems large to me, so I'm hoping something can be changed here to improve performance:

Avoided recomputation of the strain rate at ccc and sped things up a bit more. Now it runs in 2.9 minutes. A lot more to go...

[tomchor]

I have some validation scripts already that work on the CPU and indicate that the new closure is working. However, I can't make this work on the GPU. I keep getting this error:

ERROR: LoadError: GPU compilation of MethodInstance for Oceananigans.TurbulenceClosures.gpu__compute_scale_invariant_smagorinsky_viscosity!(::KernelAbstractions.CompilerMetadata{…}, ::OffsetArrays.OffsetArray{…}, ::Field{…}, ::Field{…}, ::RectilinearGrid{…}, ::ScaleInvariantSmagorinsky{…}, ::Nothing, ::@NamedTuple{…}, ::@NamedTuple{}) failed
KernelError: passing and using non-bitstype argument

Argument 7 to your kernel function is of type ScaleInvariantSmagorinsky{Oceananigans.TurbulenceClosures.ExplicitTimeDiscretization, Oceananigans.TurbulenceClosures.DirectionalAveraging{Tuple{Int64, Int64}}, Float64, @NamedTuple{}, Integer}, which is not isbits:
  .update_frequency is of type Integer which is not isbits.

Reading up on the CUDA.jl docs I think I understand where this error comes from (although I thought update_frequency, which is an Integer, should work, but it throws an error). Still couldn't figure out how to fix it in this case here. I assume it's not hard to fix though, so I was wondering if someone (I'm assuming @simone-silvestri or @glwagner) can please give me a hand or at least point me in the right direction?

[ali-ramadhan]

@tomchor Not sure if this is the issue but have you tried adapting ScaleInvariantSmagorinsky for the GPU?

Looks like ScalarBiharmonicDiffusivity needed it here: https://github.com/CliMA/Oceananigans.jl/blob/b3be950ef6a1f2139c2f65e083f20f1f7958ea55/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl#L109-L113

Not sure why update_frequency is coming up as an Integer (which is an abstract supertype) instead of Int (which is a concrete type and an alias for Int64) but maybe adapting will fix this?

[tomchor]

@tomchor Not sure if this is the issue but have you tried adapting ScaleInvariantSmagorinsky for the GPU?

Looks like ScalarBiharmonicDiffusivity needed it here:

https://github.com/CliMA/Oceananigans.jl/blob/b3be950ef6a1f2139c2f65e083f20f1f7958ea55/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl#L109-L113

Not sure why update_frequency is coming up as an Integer (which is an abstract supertype) instead of Int (which is a concrete type and an alias for Int64) but maybe adapting will fix this?

I tried that and got a different error. But now looking at ScalarBiharmonicDiffusivity I think I probably did something wrong. I'll try again soon!

[tomchor]

I wrote some validation scripts and I think the model is working, but I'd appreciate some feedback. First, this animation shows two 3D decaying turbulence simulations, one with SmagorinskyLilly and another with the ScaleInvSmag. I'm also plotting the value of the Smagorinsky constant from the dynamic model. For similar conditions (3D homogeneous tubrulence), Lilly assumed a sharp spectral filter in the inertial subrange, matched the SGS dissipation rate to the TKE dissipation rate and obtained the theoretical value of the Smag constant of $c_s$ ≈ 0.165. I think these assumptions should be roughly valid in this validation script, and indeed the value that we obtain for $c_s$ is kinda close:

https://github.com/user-attachments/assets/f210d18b-82bf-4e9e-bb26-58bbd4251820

Although to be honest I expected the value to be closer. That said, I believe Lilly also assumed a stationary turbulent flow (although I'd have to double check that), which is not the case here and may be affecting the coefficient value. Btw, I had posted a similar video before but I made some modifications to the model since then.

I also reproduced the first test simulation in the original Bou-Zeid (2005) paper: namely a channel flow forced by a pressure gradient, with a Monin-Obukhov-compliant bottom drag. I couldn't reproduce the simulation perfectly (for example, they use a pseudo-spectral scheme in the horizontal directions, force their drag with a horizontal average, and modify the SmagLilly constant slightly) so the results are a bit different from what's in their paper, but I think that's okay. Let me know if anyone thinks otherwise. Importantly, the ScaleInvSmag shows an improvement over the SmagLilly, with the later being expectedly overly diffusive and pretty much killing any turbulence at this resolution:

https://github.com/user-attachments/assets/6440450d-52f4-43ef-9ad5-0f8fcd9993d5

For reference, this is what similar plots from the paper look like (the equivalent for us here would be SMAG and PASI):

Some quick notes:

• Many tests are failing become I made an ad-hoc modification for now which passes the velocities to DiffusivityFields(). I did that because otherwise it was hard to make the model performant and simulations were taking way too long. We can (and should) review this and either come up with a better way to construct the diffusivity_fields or make this change separately in another PR, which will require changing the other models too.
• I'm updating the dynamic model once every 5 time steps only (this is user-defined). This is common practice for dynamic models since their cost is significantly higher than that of constant Smag or even AMD. With that practice, the dynamic model is taking about 3 times longer to run than the constant Smagorinsky. It does take significantly longer to compile (I haven't timed it). I believe there might be some optimizations still on the table though.
• I also found that the precise value of the Smag coefficient calculated via the dynamic procedure is dependent on the advection scheme, with WENO generally leading to smaller values. In hindsight, I think that's not surprising though.

[tomchor]

With the latest changes the compilation time is much shorter (around 10x shorter on the GPU) and simulations have been running in about twice to four times as long as the simulations with constant Smagorinsky (per time step).

[glwagner]

Some results from the "closure comparison" validation test (rotating wind-driven mixing in a stratified fluid):

dz = 1 meter, dx = dy = 2 meter

dz = 2 meter, dx = dy = 4 meter

Seems like dynamic smag is doing what it should which is great. I'm surprised to see that the Lilly coefficient doesn't do much at coarse resolution at least.

On cost: it is a bit interesting. I think as kernel saturate, WENO(order=9) will start to win. For small kernels however, AMD seems ideal. We'll have to do a proper benchmark.

This PR also contains an experimental implementation of Lagrangian-averaged (scale-invariant) dynamic Smagorinsky.

[tomchor]

I'm surprised at how much SmagLilly kills turbulence at these resolutions. To the point that $w^2$ is pretty much zero. A visualization of that effect for a wall flow can be seen here. Crazy stuff.

[glwagner]

I'm surprised at how much SmagLilly kills turbulence at these resolutions. To the point that w 2 is pretty much zero. A visualization of that effect for a wall flow can be seen here. Crazy stuff.

Right, it requires high resolution to do something remotely physical. I think that's a big limitation.

I can't see the image because it seems to be protected.

[tomchor]

I can't see the image because it seems to be protected.

Weird. It's the second video in this comment in this PR.

[glwagner]

dz = 0.5 meters, dx = dy = 1 meter

[tomchor]

Here are results including the Lagrangian-averaged version of the dynamic Smagorinsky with $\Delta z=2$ m:

As expected, for this simulation the results are pretty similar to the directionally-averaged version.

[glwagner]

Here are results including the Lagrangian-averaged version of the dynamic Smagorinsky with Δ z = 2 m:

As expected, for this simulation the results are pretty similar to the directionally-averaged version.

Lagrangian-averaged should be slightly less diffusive right?

[tomchor]

Lagrangian-averaged should be slightly less diffusive right?

Hmm, I'm not aware of that. afaik it's about the same in a doubly periodic domain. Looking at Bou-Zeid's Figure 1 I also don't see any obvious trend that way. What is generally true is that scale-dependent formulations (not implemented here) tend to be less diffusive then scale-invariant ones (which is the one implemented here).

[glwagner]

Interesting. I disagree that fig 1 does not have an obvious trend, the trend is quite clear away from the wall:

Both Lagrangian schemes have larger coefficients than their planar-averaged counter parts. Also scale dependence increases the Lagrangian coefficient.

But still, I was referring to the eddy viscosity, not the coefficient. It would be interesting to me if the Lagrangian scheme was consistently more diffusive. But that might actually also be consistent with convergence at lower resolution, perhaps.

[tomchor]

I disagree that fig 1 does not have an obvious trend, the trend is quite clear away from the wall:

Both Lagrangian schemes have larger coefficients than their planar-averaged counter parts. Also scale dependence increases the Lagrangian coefficient.

I agree that far from the wall there is a clear trend. But I'm not sure that's a fair way to look at it.

But still, I was referring to the eddy viscosity, not the coefficient. It would be interesting to me if the Lagrangian scheme was consistently more diffusive. But that might actually also be consistent with convergence at lower resolution, perhaps.

That said, good point that the coefficient doesn't necessarily translate into the eddy viscosity. I'll check the eddy viscosity later and see what comes up as I am now curious :)

[tomchor]

~I've been getting the following error right at the end of docs on buildkite that I can't reproduce locally (and therefore have not been able to solve):~

┌ Error: 4 docstrings not included in the manual:
│ 
│     Oceananigans.TurbulenceClosures.Smagorinskys.DynamicCoefficient :: Union{Tuple{}, Tuple{Any}}
│     Oceananigans.TurbulenceClosures.Smagorinskys.SmagorinskyLilly :: Union{Tuple{}, Tuple{Any}, Tuple{Any, Any}}
│     Oceananigans.TurbulenceClosures.Smagorinskys.LillyCoefficient :: Union{Tuple{}, Tuple{Any}}
│     Oceananigans.TurbulenceClosures.Smagorinskys.Smagorinsky :: Union{Tuple{}, Tuple{TD}, Tuple{TD, Any}} where TD
│ 
│ These are docstrings in the checked modules (configured with the modules keyword)
│ that are not included in canonical @docs or @autodocs blocks.
└ @ Documenter /storage5/buildkite-agent/.julia-18496/packages/Documenter/C1XEF/src/utilities/utilities.jl:44

~Anyone know what that's about? Those are the 4 new docstrings that I added, all in the Smagorinskys module.~

Nevermind, it appears I have solved it.

[tomchor]

FYI the tests are failing only because somehow tests are getting a kill signal or the servers are being lost:

Maybe something's up with he servers? As far as I can tell, all tests are passing and this PR is ready for review.

[tomchor]

Hmm, now I'm getting a "no space on device error":

[glwagner]

Hmm, now I'm getting a "no space on device error":

gotta delete somem stuff

[tomchor]

@glwagner (or anyone) could you please give me a hand with the tests? The turbulence tests are getting stuck when building a NonhydrostaticModel with a dynamic Smagorinsky (with both kinds of averaging):

The thing is that I have not been able to reproduce this locally, so it's been hard for me to fix it. Do you have any ideas of how to proceed? All other tests seems to pass.

[tomchor]

It seems that whenever the model has a DynamicSmag, update_state!() gets stuck here: https://github.com/CliMA/Oceananigans.jl/blob/5dd01ddfe73dd03515b52e6499b27c1b61c71364/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl#L43

which ends up calling compute_diffusivities!() here: https://github.com/CliMA/Oceananigans.jl/blob/5dd01ddfe73dd03515b52e6499b27c1b61c71364/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl#L64-L67

In this branch, compute_diffusivities!() looks like https://github.com/CliMA/Oceananigans.jl/blob/d00a7e6802ca73aaeb7e9c247f02775af719648f/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/smagorinsky.jl#L108-L121

and is general for all Smagorinsky models, so I'm guessing the culprit is in compute_coefficient_fields!(), which for the DirectionallyAveragedDynamicSmagorinsky is: https://github.com/CliMA/Oceananigans.jl/blob/d00a7e6802ca73aaeb7e9c247f02775af719648f/src/TurbulenceClosures/turbulence_closure_implementations/Smagorinskys/dynamic_coefficient.jl#L117-L139

And something in there makes the compiler get stuck on buildkite. I have no idea what that something is. I was not able to reproduce this on any local machine that I've tried, which is making this hard for me to pinpoint the cause.

[glwagner]

Is the problem for GPU or CPU? CPU runs on tartarus but GPU runs on sverdrup.

What closure configuration is this a problem for?

[tomchor]

I haven't been able to run the GPU tests for a while, but it's happening on CPUs for sure

[glwagner]

Which test / which file causes the problem?