Skip to content

Conversation

charleskawczynski
Copy link
Member

@charleskawczynski charleskawczynski commented Mar 20, 2025

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 CliMA/ClimaCore.jl#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.

@charleskawczynski charleskawczynski force-pushed the ck/api branch 4 times, most recently from 5242056 to 88fc7c1 Compare March 21, 2025 15:54
@charleskawczynski charleskawczynski force-pushed the ck/api branch 4 times, most recently from 3fbbbee to 7ac3cf0 Compare March 21, 2025 18:53
@charleskawczynski charleskawczynski marked this pull request as ready for review March 21, 2025 19:02
@charleskawczynski charleskawczynski force-pushed the ck/api branch 8 times, most recently from cab6563 to fbe61d0 Compare March 26, 2025 20:11
Copy link

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!

@charleskawczynski
Copy link
Member Author

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
Copy link
Member

szy21 commented Apr 8, 2025

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
Copy link
Member Author

Can @Sbozzolo also take a look at this?

Copy link
Member

@Sbozzolo Sbozzolo left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add please add documentation?

The documentation should provide enough information/instructions so that a technically versed user can use the code without input from experts.

This is currently not possible, and there's only one (incomplete) sentence about using the code:

First, users can construct a settings object using `RRTMGPSolver`

Many docstrings are also incomplete or could be more useful. For example, where does a user learn what AllSkyRadiationWithClearSkyDiagnostics mean?

Or, it's never explained what isothermal_boundary_layer is, or this docstring has an incomplete sentence:

"""
	set_cols!(data, value::Union{Number, AbstractArray})
Sets the columns in data to values in value.
If this
"""

Address comments

Address some comments

Use correct _lookup_tables method

Address comment

Apply formatter

Revise volume_mixing_ratio definitions

Add more tests

Change aero column_mass_density and radius to lookup methods

Fixes

GPU unit test fixes

Combine doc sections

Fix tests for gpu

Apply suggestion
@charleskawczynski
Copy link
Member Author

charleskawczynski commented Apr 19, 2025

Could you add please add documentation?

Many docstrings are also incomplete or could be more useful.

I've fixed the example that you gave.

@charleskawczynski charleskawczynski merged commit 5c1b0b1 into main Apr 19, 2025
6 of 10 checks passed
@Sbozzolo
Copy link
Member

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

@charleskawczynski
Copy link
Member Author

charleskawczynski commented Apr 21, 2025

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
Copy link
Member

sriharshakandala commented Apr 24, 2025

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: #578).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

6 participants