Add API to RRTMGP

[charleskawczynski]

The docs are not complete, but I'd like some early feedback before finishing documenting.

In efforts towards the RRTMGP interface refactoring, this PR adds an API to RRTMGP, which should achieve a few things:

• simplify some constructors through some helper types. In particular, lots of RRTMGP constructs all contain DA, FT, nlay, ncol-- all of which are now packed into RRTMGPGridParams
• improving inference by avoiding symbol comparisons. In particular, source_func_longwave and source_func_shortwave are type-unstable, and they need not be (by simply removing them and having users utilize union-splitting when constructing these source functions).
• provide users with an API to access data, and views into that data, that they passed to RRTMGP

One of the requirements in RRTMGP interface refactoring was to eliminate set_and_save!. set_and_save! does a few things. It:

• determines the range of the domain, by checking if the variable name starts with center_ or face_
• creates a view given the range (excluding the isothermal boundary layer) to the array
• calls set_array!, which checks the sizes of the inputs and outputs, and assigns the view to the initial value
• pushes the view into a large list, and that list is later converted into a NamedTuple and stored in RRTMGPModel

Some issues with this function are:

• It's difficult to extend and reason about, because there is a lot of implicit assumptions about the data that is passed in (e.g., prefixes with center_, face_)
• It requires creating a huge NamedTuple of views into data that we pass to RRTMGP. This gigantic struct impacts our stack traces and may be impacting our compile times. Ultimately, RRTMGP does not yet provide an interface to retrieve data that users pass to it-- that is a big component of what this PR does.

To understand the range of types of data passed to rrtmgp, I used grep for DA{ to find the arrays created in that are passed to set_and_save/set_array in RRTMGPInterface, I've found:

DA{FT}(undef,lu_kwargs.nbnd_lw,ncol)
DA{FT}(undef,ncol)
DA{FT}(undef,lu_kwargs.nbnd_sw,ncol)
DA{FT}(undef,ncol,ngpt_sw)
DA{FT}(undef,nlay+1,ncol)
DA{FT}(undef,nlay,ncol)
DA{FT}(undef,4,nlay,ncol)
DA{FT}(undef,lu_kwargs.ngas_sw)
DA{FT}(undef,lu_kwargs.ngas_sw,nlay,ncol)
DA{Bool}(undef,nlay,ncol)
DA{FT}(undef,n_aerosol_sizes,nlay,ncol)
DA{FT}(undef,n_aerosols,nlay,ncol)

If we generalize all of the non-nlay and non-ncol values to N, this simplifies to 8 cases:

DA{FT}(undef,N,ncol)
DA{FT}(undef,ncol)
DA{FT}(undef,ncol,N)
DA{FT}(undef,nlay+1,ncol)
DA{FT}(undef,nlay,ncol)
DA{FT}(undef,N,nlay,ncol)
DA{FT}(undef,N)
DA{Bool}(undef,nlay,ncol)

If we provide users with a mechanism to specify nlay and the eltype (FT or Bool), then this can further simplify to 6 cases:

DA{FT}(undef,N,ncol)
DA{FT}(undef,ncol)
DA{FT}(undef,ncol,N)
DA{FT}(undef,nlay,ncol)
DA{FT}(undef,N,nlay,ncol)
DA{FT}(undef,N)

It is a bit odd that both DA{FT}(undef,N,ncol) and DA{FT}(undef,ncol,N) exist, and perhaps they too should be unified--but they need not be. Also, technically, DA{FT}(undef,nlay,ncol) and DA{FT}(undef,ncol) are somewhat subsets of DA{FT}(undef,N,nlay,ncol) and DA{FT}(undef,N,ncol). The important point here is that we need to be able to disambiguate ncol,N and, for example, nlay,ncol data, if we want to provide users a safe and generic way to operate on this data, and to be able to write extendable custom kernels, as proposed in https://github.com/CliMA/ClimaCore.jl/pull/2233.

To make this possible, this PR introduces a new array type, RRTMGPData <: AbstractArray, with 4 constructors:

NVCData(gp::RRTMGPGridParams, ::Type{T}; N, nlay)
VCData(gp::RRTMGPGridParams, ::Type{T}; nlay)
NCData(gp::RRTMGPGridParams, ::Type{T}; N)
NData(gp::RRTMGPGridParams, ::Type{T}; N)

These constructors all return a RRTMGPData object with a subtype parameter of AbstractIndexOrder, of which, there are 4:

• NVCOrder Index order with dimensions (N, vertical level, column).
• VCOrder Index order with dimensions (vertical level, column).
• NCOrder Index order with dimensions (N, column).
• NOrder Index order with a single N dimension.

This will allow us to eliminate set_and_save! like the following:

    z_lay = RRTMGP.VCData(grid_params; nlay) # returns a RRTMGPData
    RRTMGP.set_domain!(z_lay, get(kwargs, :center_z, NaN), grid_params) # sets the domain in `z_lay` to `get(kwargs, :center_z, NaN)`

Later, when ClimaAtmos needs z_lay, it can call:

z_lay = center_z(rrtmgp_solver)

and the view can be constructed on the fly with:

z_lay_view = RRTMGP.domain_view(grid_params, z_lay)

This will allow us to eliminate the views that are stored in RRTMGPModel.

Fortunately, this works out of the box, without requiring changes RRTMGP since it is mostly written with generic concrete types, like AtmosphericState:

struct AtmosphericState{FTA1D, FTA1DN, FTA2D, D, VMR, CLD, AER} <: AbstractAtmosphericState
    "longitude, in degrees (`ncol`), optional"
    lon::FTA1DN
    "latitude, in degrees (`ncol`), optional"
    lat::FTA1DN
    "storage for col_dry [`molecules per cm^2 of dry air`], play `[Pa, mb]`, tlay `[K]`, rel_hum; `(4, nlay, ncol)`"
    layerdata::D
    "Level pressures `[Pa, mb]`; `(nlay+1,ncol)`"
    p_lev::FTA2D

and these types simply assume that inputs are arrays.

[gemini-code-assist[bot]]

[!WARNING] You have reached your daily quota limit. As a reminder, free tier users are limited to 5 requests per day. Please wait up to 24 hours and I will start processing your requests again!

[szy21]

I'll try to finish reviewing this PR before the end of tomorrow, although I will likely only skim through it without checking very carefully if everything is correct.

[Sbozzolo]

The google doc seem to be set to private and I cannot open it

[charleskawczynski]

The google doc seem to be set to private and I cannot open it

Invitation sent.

[szy21]

The overall change looks good to me. For the data layout, or the API in general, it looks like an improvement, but it is hard for me to evaluate without seeing corresponding changes in ClimaAtmos. I will leave it to other people who are more qualified to comment on it.

[charleskawczynski]

Suggestions on the following function names (and internal structure names)? I suppose the internals can be cleaned up, but it might be nice to at least solidify the user-facing part.

lw_flux_up(s::RRTMGPSolver)                           = s.lws.flux.flux_up
lw_flux_dn(s::RRTMGPSolver)                           = s.lws.flux.flux_dn
lw_flux(s::RRTMGPSolver)                              = s.lws.flux.flux_net
clear_lw_flux_up(s::RRTMGPSolver)                     = s.clear_flux_lw.flux_up
clear_lw_flux_dn(s::RRTMGPSolver)                     = s.clear_flux_lw.flux_dn
clear_lw_flux(s::RRTMGPSolver)                        = s.clear_flux_lw.flux_net
sw_flux_up(s::RRTMGPSolver)                           = s.sws.flux.flux_up
sw_flux_dn(s::RRTMGPSolver)                           = s.sws.flux.flux_dn
sw_flux(s::RRTMGPSolver)                              = s.sws.flux.flux_net
clear_sw_flux_up(s::RRTMGPSolver)                     = s.clear_flux_sw.flux_up
clear_sw_flux_dn(s::RRTMGPSolver)                     = s.clear_flux_sw.flux_dn
clear_sw_direct_flux_dn(s::RRTMGPSolver)              = s.clear_flux_sw.flux_dn_dir
clear_sw_flux(s::RRTMGPSolver)                        = s.clear_flux_sw.flux_net

[szy21]

Suggestions on the following function names (and internal structure names)? I suppose the internals can be cleaned up, but it might be nice to at least solidify the user-facing part.

lw_flux_up(s::RRTMGPSolver)                           = s.lws.flux.flux_up
lw_flux_dn(s::RRTMGPSolver)                           = s.lws.flux.flux_dn
lw_flux(s::RRTMGPSolver)                              = s.lws.flux.flux_net
clear_lw_flux_up(s::RRTMGPSolver)                     = s.clear_flux_lw.flux_up
clear_lw_flux_dn(s::RRTMGPSolver)                     = s.clear_flux_lw.flux_dn
clear_lw_flux(s::RRTMGPSolver)                        = s.clear_flux_lw.flux_net
sw_flux_up(s::RRTMGPSolver)                           = s.sws.flux.flux_up
sw_flux_dn(s::RRTMGPSolver)                           = s.sws.flux.flux_dn
sw_flux(s::RRTMGPSolver)                              = s.sws.flux.flux_net
clear_sw_flux_up(s::RRTMGPSolver)                     = s.clear_flux_sw.flux_up
clear_sw_flux_dn(s::RRTMGPSolver)                     = s.clear_flux_sw.flux_dn
clear_sw_direct_flux_dn(s::RRTMGPSolver)              = s.clear_flux_sw.flux_dn_dir
clear_sw_flux(s::RRTMGPSolver)                        = s.clear_flux_sw.flux_net

Maybe use lw_flux_net instead of lw_flux (and same for other net fluxes) to be more clear?

[charleskawczynski]

Can @Sbozzolo also take a look at this?

[charleskawczynski]

Could you add please add documentation?

Many docstrings are also incomplete or could be more useful.

I've fixed the example that you gave.

[Sbozzolo]

I just wanted to point out that my comments have not been addressed and this PR was merged against my review.

[charleskawczynski]

I just wanted to point out that my comments have not been addressed and this PR was merged against my review.

Sorry about that. I'll address the remaining pieces in a follow-up PR. I'd like to get started, in parallel, on the PR in atmos where we leverage these new pieces.

[sriharshakandala]

Is there a reason OS Unit Tests / test-os (ubuntu-latest) (pull_request) tests are turned off. These used to be mandatory for merging a PR. The aqua tests are only run on the Github CI to avoid redundancy. The aqua tests seems to be failing due to method ambiguities from this PR ( issue: https://github.com/CliMA/RRTMGP.jl/issues/578).