Couple the HBModel and SWModel for use in the Split-Explicit Stepper

[blallen]

Description

A continuation of #1257, adding the necessary communication between the two models needed to have a true split-explicit stepper. The work in this PR enables us to run with up to a 90-minute timestep for the hydrostatic spindown test case.

The broadcasts of reshaped MPIStateArrays depends on #1071 to run on the GPU. I'll pull over the commits from that PR for running bors, but we need to merge that before we merge this one.

There's quite a bit of partial work in this PR that was necessary to get things running, but that either isn't fully implemented everywhere it should be or isn't implemented super elegantly (mainly regarding collecting all the Ocean models into a single module and how multiple dispatch is used for the couple). I'll revisit that once we have the full timestepping machinery running.

I have

• Written and run all necessary tests with CLIMA by including tests/runtests.jl
• Followed all necessary style guidelines and run julia .dev/climaformat.jl .
• Updated the documentation to reflect changes from this PR.

For review by CLIMA developers

• There are no open pull requests for this already
• CLIMA developers with relevant expertise have been assigned to review this submission
• The code conforms to the style guidelines and has consistent naming conventions. julia .dev/format.jl has been run in a separate commit.
• This code does what it is technically intended to do (all numerics make sense physically and/or computationally)

[blallen]

So there's a nasty bug hidden somewhere in the code I used to couple the timesteppers: when using the Statecheck module, the velocity fields can differ after approximately the 4th decimal place after each run. I've tried comparing runs within the same REPL session and across different ones and it happens either way. The weird thing is that the solutions are still within 0.2% of the analytical solution and the VTK output looks reasonable.

As far as I can tell, I've made sure to properly initialize all the variables we are using. I'd appreciate someone taking another look at this before I merge everything in.

[christophernhill]

@blallen there is a lot in this PR that is all lumped together. Is it too late to do over and try and do this in stages? That would be my general instinct (I know thats logistically painful given current CI and monorepo foo, but it is still probably the right approach).

One funky thing I did notice is

boxy_ΔGᵘ = reshape(ΔGᵘ, Nq^2, Nqk, 2, nelemv, nelemh)
         boxy_ΔGᵘ .=
             -reshape(flat_∫du, Nq^2, 1, 2, 1, nelemh) / baroclinic.problem.H

but I think thats probably just left over debugging? or possibly it should be

         boxy_ΔGᵘ .-=
             reshape(flat_∫du, Nq^2, 1, 2, 1, nelemh) / baroclinic.problem.H

in which case spelling it out properly might be better. My general take is Unary ops are an easy way to blow yourself up and don't add anything except potential to make really dumb subtle errors.

• separate (but related)
  @leios does .= broadcast work all the time and in all ways for all things GPU, or are we better off using copy!() ? I tried some .= with the MPI_wrapper PR and got various Dat error messages.

[blallen]

@blallen there is a lot in this PR that is all lumped together. Is it too late to do over and try and do this in stages? That would be my general instinct (I know thats logistically painful given current CI and monorepo foo, but it is still probably the right approach).

I can spin-off 22cb1ff (and related changes) into a separate PR. I was originally planning on doing this afterwards in order to reuse the SimpleBoxProblem setup between both SWModel and HBModel, but realized that having a Base module would also be useful for the coupling traits.

c25791f can also be split into a separate PR, but it will depend on #1071 being merged. I just put it here since all the other broadcast reshapes also depend on #1071 so this PR is CPU only for now.

The rest of the stuff in the PR is pretty hard to disentangle. Basically everything in the src/Ocean/SplitExplicit directory is needed before coupling the models makes sense. This includes the TendencyIntegralModel, the FlowIntegralModel, the FullyCoupled dispatch functions in HydrostaticBoussinesqCoupling.jl and ShallowWaterCoupling.jl, and finally the actual communication between the models in Communication.jl.

We might be able to add some unit tests for the two *IntegralModel where we just provide an input flow/tendency and check that the resulting integrals make sense. Then we could also split those off into a separate PR. I'm not sure if there's anything we could do to split the *Coupling.jl and Communication.jl files into smaller chunks.

One funky thing I did notice is

boxy_ΔGᵘ = reshape(ΔGᵘ, Nq^2, Nqk, 2, nelemv, nelemh)
         boxy_ΔGᵘ .=
             -reshape(flat_∫du, Nq^2, 1, 2, 1, nelemh) / baroclinic.problem.H

but I think thats probably just left over debugging? or possibly it should be

         boxy_ΔGᵘ .-=
             reshape(flat_∫du, Nq^2, 1, 2, 1, nelemh) / baroclinic.problem.H

in which case spelling it out properly might be better. My general take is Unary ops are an easy way to blow yourself up and don't add anything except potential to make really dumb subtle errors.

Ah yeah, that's probably a better way to write it. I copied this straight from @jm-c's branch, but I think your way of writing it is functionally the same (and definitely clearer).

• separate (but related)
  @leios does .= broadcast work all the time and in all ways for all things GPU, or are we better off using copy!() ? I tried some .= with the MPI_wrapper PR and got various Dat error messages.

@leios is going to squash #1071 soon, and then I can cherry-pick that commit and test all of this on the GPU.

[christophernhill]

@blallen doing some splitting would be great. Some is better than none.

Just to check regarding

boxy_ΔGᵘ = reshape(ΔGᵘ, Nq^2, Nqk, 2, nelemv, nelemh)
         boxy_ΔGᵘ .=
             -reshape(flat_∫du, Nq^2, 1, 2, 1, nelemh) / baroclinic.problem.H

As this is written I think the second line is completely overwriting the first assignment (i.e. it is NOT subtracting the integral). The second line is setting every level to minus of the integral. If you meant to be subtracting the integral then I think the operator needs to be

.-=

NOT

.=-

If you are simply assigning then I am not sure why the

boxy_ΔGᵘ = reshape(ΔGᵘ, Nq^2, Nqk, 2, nelemv, nelemh)

is there.

Adding some all-caps just to make sure you didn't miss what I was saying!!!

[blallen]

The function below zeros out the contents of ΔGᵘ at the beginning of each RK stage, so I don't think there is any accumulation anyways. The .-= syntax is definitely clearer than the .=- syntax, but I think the function makes them equivalent in this case. I will change to .-= after rebasing and splitting things off into separate PRs.

@inline function initialize_states!(
    baroclinic::HydrostaticBoussinesqModel,
    barotropic::ShallowWaterModel,
    model_bc,
    model_bt,
    state_bc,
    state_bt,
)
    model_bc.state_auxiliary.ΔGᵘ .= -0

    return nothing
end

[blallen]

Okay, just rebased and rewrote history. Plan of attack is:

59959e6 - gets it own PR #1302 , no new tests needed
3795358 - gets it own PR, need to think up some simple tests for the integrations
46b958c - gets it own PR #1301, depends on #1071
070e655 and e887683 probably squashed to a single commit in the future and remain in this PR, depends on #1071

[blallen]

Rebased on top of #1302, dropped the commit that got sent to #1301 and did some more history rewriting as I realized we could replace the FlowIntegralModel and TendencyIntegralModel with a single VerticalIntegralModel.jl with some minor modifications. This way we only need one test for the integrals instead of one for the flow and one for the tendency.

I'll split off 55d879b into its own PR tomorrow and add some basic tests based on the hydrostatic spindown initial state (and we can add more complicated tests as we get them).

Just to note here, #1301 is a bit more complicated than expected. We'll need to make the corresponding changes here once everything is figured out there.

[blallen]

bors r+

[blallen]

bors r+

[bors]

Build failed:

• format

[blallen]

bors r+

[bors]

Merge conflict.

[blallen]

bors r+

[blallen]

bors try

[blallen]

bors try-

[blallen]

bors try

[bors]

try

Build failed:

• ci/slurmci

[blallen]

bors r+

[bors]

Build succeeded:

• test-os (ubuntu-latest)
• ci/slurmci
• test-os (macos-latest)
• format
• docs-build
• test-os (windows-latest)