Fix up `ShallowWaterModel.jl` about oceananigans.jl HOT 15 OPEN

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

from oceananigans.jl.

@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.

from oceananigans.jl.

@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

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

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

from oceananigans.jl.

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

from oceananigans.jl.

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.

from oceananigans.jl.

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.

from oceananigans.jl.

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.

from oceananigans.jl.

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?

from oceananigans.jl.

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

from oceananigans.jl.

Nice!

from oceananigans.jl.

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

from oceananigans.jl.

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.

from oceananigans.jl.

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.

from oceananigans.jl.

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

from oceananigans.jl.

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

from oceananigans.jl.