Fix up `ShallowWaterModel.jl`

[francispoulin]

This issue will help to combine several outstanding issues involving the ShallowWaterModel.jl. I suggest the following in this order:

• [ ] Fix inconsistent tendencies in the two models (#2928)
• [ ] Fix up shallow water regression tests (#3049)
• [ ] Add viscosity
• [ ] Speed up shallow water example when building docs (#3151)
• [ ] Make shallow water Bickley jet less expensive (#3169)
• [ ] Update immersed bondaries validation scripts (#2985)
• [ ] Introduce multi-layer shallow water model (#2507)
• [ ] Validate positive preserving WENO schemes

For the first part, I put together this document that shows and derives the equations for the two different models. Next I will look at the tendencies.

Any suggestions on what else to do, or what to do differently, are welcome!

ShallowWaterModel_eqns.pdf

[francispoulin]

@simone-silvestri and I created a new branch to try and start tackle these issues, fp-ss/shallow-water-version2.

When ran the shallow_water_Bickley_jet.jl on this branch there was an error. We were using superscripts fcc and cfc instead if faa and cfc on the averaging operators. I have since fixed that.

Now I'm going to try and merge main into this branch.

[francispoulin]

@simone-silvestri, I'm trying to merge main into fp-ss/shallow-water-version2 and having some problems with src/Grids/latitude_longitude-grid.jl.

The version you put together has this in line 189 https://github.com/CliMA/Oceananigans.jl/blob/6bd880053dc865c58a1667cfd2649624cb3e8aa8/src/Grids/latitude_longitude_grid.jl#L189

Main however has a very different syntax on the left hand side of equals.

https://github.com/CliMA/Oceananigans.jl/blob/0391b3a90d0dddd67625cd4373e80012e9806df6/src/Grids/latitude_longitude_grid.jl#L189

I remember that you worked on the sizes when we chatted. Any suggestions as to how we should resolve the conflicts?

[francispoulin]

@simone-silvestri , any advice on how to resolve this merger with main?

[glwagner]

@francispoulin Nλ, Nφ, Nz = size and Hλ, Hφ, Hz = halo

[francispoulin]

Thanks. I remember Simone made a change here on this branch and I'm not sure why but I'll try and return to what is done in main and go from there.

[francispoulin]

The branch fp-ss/shallow-water-version2 has (hopefully) corrected the problems with the tendency calculations brought up by @glwagner in #2928 (and elsewhere).

Does someone what to take a look to confirm that we have done this correctly? If yes the changes start here.

I thought the next step might be to add in the regression tests one by one, which were removed in #3050. I am going to look into test_shallow_water_models.jl next to see what runs and what doesn't with the current code.

[francispoulin]

All the tests seem to pass but I get two errors.

The first error is on a GPU, running test_shallow_water_diffusion_cosine on the field u using the ConservativeFormulation(). When I run this in isolation I get the following error. Any suggestions on what to do here?

[2024/02/27 10:33:43.977] INFO    Testing ShallowWaterModel cosine viscous diffusion [u, ConservativeFormulation()]
[2024/02/27 10:33:43.977] WARN  The ShallowWaterModel is currently unvalidated, subject to change, and should not be used for scientific research without adequate validation. -@-> /home/fpoulin/Software/Oceananigans.jl/src/Models/ShallowWaterModels/shallow_water_model.jl:129
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore are only permitted from the REPL for prototyping purposes.
If you did intend to index this array, annotate the caller with @allowscalar.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/uOYfN/src/GPUArraysCore.jl:103
  [3] getindex(xs::CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}, I::Int64)
    @ GPUArrays ~/.julia/packages/GPUArrays/EZkix/src/host/indexing.jl:9
  [4] getindex
    @ ~/.julia/packages/OffsetArrays/0MOrf/src/OffsetArrays.jl:438 [inlined]
  [5] _getindex
    @ ./abstractarray.jl:1321 [inlined]
  [6] getindex
    @ ./abstractarray.jl:1288 [inlined]
  [7] getindex
    @ ~/Software/Oceananigans.jl/src/Fields/field.jl:399 [inlined]
  [8] _broadcast_getindex
    @ ./broadcast.jl:675 [inlined]
  [9] _getindex
    @ ./broadcast.jl:705 [inlined]
 [10] _broadcast_getindex
    @ ./broadcast.jl:681 [inlined]
 [11] getindex
    @ ./broadcast.jl:636 [inlined]
 [12] macro expansion
    @ ./broadcast.jl:1004 [inlined]
 [13] macro expansion
    @ ./simdloop.jl:77 [inlined]
 [14] copyto!
    @ ./broadcast.jl:1003 [inlined]
 [15] copyto!
    @ ./broadcast.jl:956 [inlined]
 [16] copy
    @ ./broadcast.jl:928 [inlined]
 [17] materialize
    @ ./broadcast.jl:903 [inlined]
 [18] isapprox(a::Field{…}, b::Field{…}; kw::@Kwargs{…})
    @ Oceananigans.Fields ~/Software/Oceananigans.jl/src/Fields/field.jl:712
 [19] test_shallow_water_diffusion_cosine(grid::RectilinearGrid{…}, formulation::ConservativeFormulation, fieldname::Symbol, ξ::Base.ReshapedArray{…})
    @ Main ~/Software/Oceananigans.jl/test/test_shallow_water_models.jl:84
 [20] top-level scope
    @ ./REPL[25]:4
 [21] top-level scope
    @ ~/.julia/packages/CUDA/nbRJk/src/initialization.jl:205
Some type information was truncated. Use `show(err)` to see complete types.

[francispoulin]

I was getting two errors but when I added CUDA.allowscalar(true) into test_shallow_water_models.jl the tests all passed on a CPU and GPU on my laptop.

Note that in dependencies_for_runtests.jl this line here only found a GPU on my computer. When I changed it temporarily to test both CPU and GPU and all 80 tests passed.

I remember last year in #3050 @navidcy found that it ran on some computers but not others. I presume that is still a concern. But lots has changed since then, for example we are no longer using julia 1.8.

@navidcy, might you be able to try the shallow water tests on the same computer you found the failers in last year to see if the problem persists?

[francispoulin]

I also ran the tests on a server and all the tests passed on both CPUs and GPUs.

[simone-silvestri]

Nice!

[simone-silvestri]

All the tests seem to pass but I get two errors.

The first error is on a GPU, running test_shallow_water_diffusion_cosine on the field u using the ConservativeFormulation(). When I run this in isolation I get the following error. Any suggestions on what to do here?

[2024/02/27 10:33:43.977] INFO    Testing ShallowWaterModel cosine viscous diffusion [u, ConservativeFormulation()]
[2024/02/27 10:33:43.977] WARN  The ShallowWaterModel is currently unvalidated, subject to change, and should not be used for scientific research without adequate validation. -@-> /home/fpoulin/Software/Oceananigans.jl/src/Models/ShallowWaterModels/shallow_water_model.jl:129
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore are only permitted from the REPL for prototyping purposes.
If you did intend to index this array, annotate the caller with @allowscalar.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/uOYfN/src/GPUArraysCore.jl:103
  [3] getindex(xs::CuArray{Float64, 3, CUDA.Mem.DeviceBuffer}, I::Int64)
    @ GPUArrays ~/.julia/packages/GPUArrays/EZkix/src/host/indexing.jl:9
  [4] getindex
    @ ~/.julia/packages/OffsetArrays/0MOrf/src/OffsetArrays.jl:438 [inlined]
  [5] _getindex
    @ ./abstractarray.jl:1321 [inlined]
  [6] getindex
    @ ./abstractarray.jl:1288 [inlined]
  [7] getindex
    @ ~/Software/Oceananigans.jl/src/Fields/field.jl:399 [inlined]
  [8] _broadcast_getindex
    @ ./broadcast.jl:675 [inlined]
  [9] _getindex
    @ ./broadcast.jl:705 [inlined]
 [10] _broadcast_getindex
    @ ./broadcast.jl:681 [inlined]
 [11] getindex
    @ ./broadcast.jl:636 [inlined]
 [12] macro expansion
    @ ./broadcast.jl:1004 [inlined]
 [13] macro expansion
    @ ./simdloop.jl:77 [inlined]
 [14] copyto!
    @ ./broadcast.jl:1003 [inlined]
 [15] copyto!
    @ ./broadcast.jl:956 [inlined]
 [16] copy
    @ ./broadcast.jl:928 [inlined]
 [17] materialize
    @ ./broadcast.jl:903 [inlined]
 [18] isapprox(a::Field{…}, b::Field{…}; kw::@Kwargs{…})
    @ Oceananigans.Fields ~/Software/Oceananigans.jl/src/Fields/field.jl:712
 [19] test_shallow_water_diffusion_cosine(grid::RectilinearGrid{…}, formulation::ConservativeFormulation, fieldname::Symbol, ξ::Base.ReshapedArray{…})
    @ Main ~/Software/Oceananigans.jl/test/test_shallow_water_models.jl:84
 [20] top-level scope
    @ ./REPL[25]:4
 [21] top-level scope
    @ ~/.julia/packages/CUDA/nbRJk/src/initialization.jl:205
Some type information was truncated. Use `show(err)` to see complete types.

Looks like it's a problem with the test, not really with the model

[francispoulin]

Thanks @simone-silvestri. The problem has been fixed.

These tests have all passed for me on CPUs and GPUs on my laptop and a server, and I presume they will for everyone else.

But I don't beleive these tests were the regression tests that failed before. I am happy to look into those and revive them one by one, if someone can point me to where I might find them.

[francispoulin]

I updated the regression tests for shallow water and added them to this branch.

I tested it on my laptop and the a cluster and all tested passed on both the CPU and GPU.

Does anyone else want to try out the regression tests to see if they work for them?

I am happy to do testing but I can't test it for me as everything seems to be working. Help on this would be greatly apprecited.

[simone-silvestri]

If you create a draft PR using this branch the CI should kick in.

[francispoulin]

Thanks for the suggestion! I will do that right away.