Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Move SG_quad out of cache #3462

Merged
merged 2 commits into from
Dec 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 1 addition & 7 deletions src/cache/cache.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@ struct AtmosCache{
COR,
SFC,
GHOST,
SGQ,
PREC,
SCRA,
HYPE,
Expand Down Expand Up @@ -44,9 +43,6 @@ struct AtmosCache{
"""Center and face ghost buffers used by DSS"""
ghost_buffer::GHOST

"""Struct with sub-grid sampling quadrature"""
SG_quad::SGQ

"""Quantities that are updated with set_precomputed_quantities!"""
precomputed::PREC

Expand Down Expand Up @@ -144,11 +140,10 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)

sfc_setup = surface_setup(params)
scratch = temporary_quantities(Y, atmos)
SG_quad = SGSQuadrature(FT)

precomputed = precomputed_quantities(Y, atmos)
precomputing_arguments =
(; atmos, core, params, sfc_setup, precomputed, scratch, dt, SG_quad)
(; atmos, core, params, sfc_setup, precomputed, scratch, dt)

# Coupler compatibility
isnothing(precomputing_arguments.sfc_setup) &&
Expand Down Expand Up @@ -178,7 +173,6 @@ function build_cache(Y, atmos, params, surface_setup, sim_info, aerosol_names)
core,
sfc_setup,
ghost_buffer,
SG_quad,
precomputed,
scratch,
hyperdiff,
Expand Down
19 changes: 11 additions & 8 deletions src/cache/cloud_fraction.jl
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,10 @@ NVTX.@annotate function set_cloud_fraction!(
Y,
p,
::Union{EquilMoistModel, NonEquilMoistModel},
::QuadratureCloud,
qc::QuadratureCloud,
)
(; SG_quad, params) = p
SG_quad = qc.SG_quad
(; params) = p

FT = eltype(params)
thermo_params = CAP.thermodynamics_params(params)
Expand Down Expand Up @@ -112,10 +113,11 @@ NVTX.@annotate function set_cloud_fraction!(
Y,
p,
::Union{EquilMoistModel, NonEquilMoistModel},
::SGSQuadratureCloud,
qc::SGSQuadratureCloud,
::DiagnosticEDMFX,
)
(; SG_quad, params) = p
SG_quad = qc.SG_quad
(; params) = p

FT = eltype(params)
thermo_params = CAP.thermodynamics_params(params)
Expand Down Expand Up @@ -162,10 +164,11 @@ NVTX.@annotate function set_cloud_fraction!(
Y,
p,
::Union{EquilMoistModel, NonEquilMoistModel},
::SGSQuadratureCloud,
qc::SGSQuadratureCloud,
::PrognosticEDMFX,
)
(; SG_quad, params) = p
SG_quad = qc.SG_quad
(; params) = p

FT = eltype(params)
thermo_params = CAP.thermodynamics_params(params)
Expand Down Expand Up @@ -220,7 +223,8 @@ end
coeff, ᶜlength_scale, thermo_params)

where:
- SG_quad is a struct containing information about quadrature type and order
- SG_quad is a struct containing information about Gaussian quadrature order,
sampling point values and weights
- ts is the thermodynamic state
- ᶜ∇q, ᶜ∇θ are the gradients of q_tot and liquid ice potential temperature
- coeff - a free parameter (to be moved into params)
Expand Down Expand Up @@ -248,7 +252,6 @@ function quad_loop(
# and limited covarainces
function get_x_hat(χ1, χ2)

@assert SG_quad.quadrature_type isa GaussianQuad
FT = eltype(χ1)

q′q′ = covariance_from_grad(coeff, ᶜlength_scale, ᶜ∇q, ᶜ∇q)
Expand Down
5 changes: 3 additions & 2 deletions src/solver/model_getters.jl
Original file line number Diff line number Diff line change
Expand Up @@ -313,12 +313,13 @@ end

function get_cloud_model(parsed_args)
cloud_model = parsed_args["cloud_model"]
FT = parsed_args["FLOAT_TYPE"] == "Float64" ? Float64 : Float32
return if cloud_model == "grid_scale"
GridScaleCloud()
elseif cloud_model == "quadrature"
QuadratureCloud()
QuadratureCloud(SGSQuadrature(FT))
elseif cloud_model == "quadrature_sgs"
SGSQuadratureCloud()
SGSQuadratureCloud(SGSQuadrature(FT))
else
error("Invalid cloud_model $(cloud_model)")
end
Expand Down
102 changes: 66 additions & 36 deletions src/solver/types.jl
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,74 @@ struct NoPrecipitation <: AbstractPrecipitationModel end
struct Microphysics0Moment <: AbstractPrecipitationModel end
struct Microphysics1Moment <: AbstractPrecipitationModel end

"""

AbstractSGSamplingType

How sub-grid scale diagnostic should be sampled in computing cloud fraction.
"""
abstract type AbstractSGSamplingType end

"""
SGSMean

Use the mean value.
"""
struct SGSMean <: AbstractSGSamplingType end

"""
SGSQuadrature

Compute the mean as a weighted sum of the Gauss-Hermite quadrature points.
"""
struct SGSQuadrature{N, A, W} <: AbstractSGSamplingType
a::A # values
w::W # weights
function SGSQuadrature(::Type{FT}; quadrature_order = 3) where {FT}
N = quadrature_order
# TODO: double check this python-> julia translation
# a, w = np.polynomial.hermite.hermgauss(N)
a, w = FastGaussQuadrature.gausshermite(N)
a, w = SA.SVector{N, FT}(a), SA.SVector{N, FT}(w)
return new{N, typeof(a), typeof(w)}(a, w)
end
end
quadrature_order(::SGSQuadrature{N}) where {N} = N

"""
AbstractCloudModel

How to compute the cloud fraction.
"""
abstract type AbstractCloudModel end

"""
GridScaleCloud

Compute the cloud fraction based on grid mean conditions.
"""
struct GridScaleCloud <: AbstractCloudModel end
struct QuadratureCloud <: AbstractCloudModel end
struct SGSQuadratureCloud <: AbstractCloudModel end

"""
QuadratureCloud

Compute the cloud fraction by sampling over the quadrature points, but without
the EDMF sub-grid scale model.
"""
struct QuadratureCloud{SGQ <: AbstractSGSamplingType} <: AbstractCloudModel
SG_quad::SGQ
end

"""
SGSQuadratureCloud

Compute the cloud fraction as a sum of the EDMF environment and updraft
contributions. The EDMF environment cloud fraction is computed by sampling over
the quadrature points.
"""
struct SGSQuadratureCloud{SGQ <: AbstractSGSamplingType} <: AbstractCloudModel
SG_quad::SGQ
end

abstract type AbstractModelConfig end
struct SingleColumnModel <: AbstractModelConfig end
Expand Down Expand Up @@ -274,40 +338,6 @@ struct GeneralizedDetrainment <: AbstractDetrainmentModel end
struct GeneralizedHarmonicsDetrainment <: AbstractDetrainmentModel end
struct SmoothAreaDetrainment <: AbstractDetrainmentModel end

abstract type AbstractQuadratureType end
struct LogNormalQuad <: AbstractQuadratureType end
struct GaussianQuad <: AbstractQuadratureType end

abstract type AbstractSGSamplingType end
struct SGSMean <: AbstractSGSamplingType end
struct SGSQuadrature{N, QT, A, W} <: AbstractSGSamplingType
quadrature_type::QT
a::A
w::W
function SGSQuadrature(
::Type{FT};
quadrature_name = "gaussian",
quadrature_order = 3,
) where {FT}
quadrature_type = if quadrature_name == "log-normal"
LogNormalQuad()
elseif quadrature_name == "gaussian"
GaussianQuad()
else
error("Invalid thermodynamics quadrature $(quadrature_name)")
end
N = quadrature_order
# TODO: double check this python-> julia translation
# a, w = np.polynomial.hermite.hermgauss(N)
a, w = FastGaussQuadrature.gausshermite(N)
a, w = SA.SVector{N, FT}(a), SA.SVector{N, FT}(w)
QT = typeof(quadrature_type)
return new{N, QT, typeof(a), typeof(w)}(quadrature_type, a, w)
end
end
quadrature_order(::SGSQuadrature{N}) where {N} = N
quad_type(::SGSQuadrature{N}) where {N} = N #TODO - this seems wrong?

abstract type AbstractSurfaceThermoState end
struct GCMSurfaceThermoState <: AbstractSurfaceThermoState end

Expand Down
1 change: 0 additions & 1 deletion test/coupler_compatibility.jl
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,6 @@ const T2 = 290
p.core,
sfc_setup,
p.ghost_buffer,
p.SG_quad,
p.precomputed,
p.scratch,
p.hyperdiff,
Expand Down
Loading