Skip to content

Commit e9422c3

Browse files
nefrathenricinefrathenrici
committed
structure atmosmodel
1 parent 1b253e1 commit e9422c3

File tree

2 files changed

+219
-74
lines changed

2 files changed

+219
-74
lines changed

‎src/solver/type_getters.jl

Lines changed: 57 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -85,24 +85,47 @@ function get_atmos(config::AtmosConfig, params)
8585
FT,
8686
)
8787

88-
atmos = AtmosModel(;
88+
# Create grouped structs
89+
moisture = AtmosMoistureModel(;
8990
moisture_model,
90-
ozone,
91-
co2,
92-
radiation_mode,
93-
subsidence = get_subsidence_model(parsed_args, radiation_mode, FT),
94-
ls_adv = get_large_scale_advection_model(parsed_args, FT),
95-
external_forcing = get_external_forcing_model(parsed_args, FT),
96-
scm_coriolis = get_scm_coriolis(parsed_args, FT),
97-
advection_test,
98-
edmfx_model,
9991
precip_model,
10092
cloud_model,
10193
noneq_cloud_formation_mode = implicit_noneq_cloud_formation ?
10294
Implicit() : Explicit(),
103-
forcing_type,
10495
call_cloud_diagnostics_per_stage,
96+
)
97+
98+
forcing = AtmosForcing(;
99+
forcing_type,
100+
subsidence = get_subsidence_model(parsed_args, radiation_mode, FT),
101+
external_forcing = get_external_forcing_model(parsed_args, FT),
102+
)
103+
104+
radiation = AtmosRadiation(;
105+
radiation_mode,
106+
ozone,
107+
co2,
108+
insolation = get_insolation_form(parsed_args),
109+
)
110+
111+
advection = AtmosAdvection(;
112+
ls_adv = get_large_scale_advection_model(parsed_args, FT),
113+
advection_test,
114+
)
115+
116+
turbconv = AtmosTurbconv(;
117+
scm_coriolis = get_scm_coriolis(parsed_args, FT),
118+
edmfx_model,
105119
turbconv_model = get_turbconv_model(FT, parsed_args, turbconv_params),
120+
sgs_adv_mode = implicit_sgs_advection ? Implicit() : Explicit(),
121+
sgs_entr_detr_mode = implicit_sgs_entr_detr ? Implicit() : Explicit(),
122+
sgs_nh_pressure_mode = implicit_sgs_nh_pressure ? Implicit() :
123+
Explicit(),
124+
sgs_mf_mode = implicit_sgs_mass_flux ? Implicit() : Explicit(),
125+
smagorinsky_lilly = get_smagorinsky_lilly_model(parsed_args),
126+
)
127+
128+
gravity_wave = AtmosGravityWave(;
106129
non_orographic_gravity_wave = get_non_orographic_gravity_wave_model(
107130
parsed_args,
108131
FT,
@@ -111,23 +134,37 @@ function get_atmos(config::AtmosConfig, params)
111134
parsed_args,
112135
FT,
113136
),
114-
hyperdiff = get_hyperdiffusion_model(parsed_args, FT),
137+
)
138+
139+
vert_diff_grouped = AtmosVertDiff(;
115140
vert_diff,
116141
diff_mode = implicit_diffusion ? Implicit() : Explicit(),
117-
sgs_adv_mode = implicit_sgs_advection ? Implicit() : Explicit(),
118-
sgs_entr_detr_mode = implicit_sgs_entr_detr ? Implicit() : Explicit(),
119-
sgs_nh_pressure_mode = implicit_sgs_nh_pressure ? Implicit() :
120-
Explicit(),
121-
sgs_mf_mode = implicit_sgs_mass_flux ? Implicit() : Explicit(),
142+
)
143+
144+
sponge = AtmosSponge(;
122145
viscous_sponge = get_viscous_sponge_model(parsed_args, params, FT),
123-
smagorinsky_lilly = get_smagorinsky_lilly_model(parsed_args),
124146
rayleigh_sponge = get_rayleigh_sponge_model(parsed_args, params, FT),
147+
)
148+
149+
surface = AtmosSurface(;
125150
sfc_temperature = get_sfc_temperature_form(parsed_args),
126-
insolation = get_insolation_form(parsed_args),
127-
disable_surface_flux_tendency = parsed_args["disable_surface_flux_tendency"],
128151
surface_model = get_surface_model(parsed_args),
129152
surface_albedo = get_surface_albedo_model(parsed_args, params, FT),
153+
)
154+
155+
atmos = AtmosModel(;
156+
moisture,
157+
forcing,
158+
radiation,
159+
advection,
160+
turbconv,
161+
gravity_wave,
162+
hyperdiff = get_hyperdiffusion_model(parsed_args, FT),
163+
vert_diff = vert_diff_grouped,
164+
sponge,
165+
surface,
130166
numerics = get_numerics(parsed_args),
167+
disable_surface_flux_tendency = parsed_args["disable_surface_flux_tendency"],
131168
)
132169
@assert !@any_reltype(atmos, (UnionAll, DataType))
133170

‎src/solver/types.jl

Lines changed: 162 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -514,85 +514,193 @@ Base.@kwdef struct EDMFXModel{
514514
scale_blending_method::SBM
515515
end
516516

517-
Base.@kwdef struct AtmosModel{
518-
MM,
519-
PM,
520-
CM,
521-
NCFM,
522-
CCDPS,
523-
F,
524-
S,
525-
OZ,
526-
CO2,
527-
RM,
528-
LA,
529-
EXTFORCING,
530-
EC,
531-
AT,
532-
EDMFX,
533-
TCM,
534-
NOGW,
535-
OGW,
536-
HD,
537-
VD,
538-
DM,
539-
SAM,
540-
SEDM,
541-
SNPM,
542-
SMM,
543-
VS,
544-
SL,
545-
RS,
546-
ST,
547-
IN,
548-
SM,
549-
SA,
550-
NUM,
551-
}
517+
# Grouped structures to reduce AtmosModel type parameters
518+
519+
"""
520+
AtmosMoistureModel
521+
522+
Groups moisture-related models and parameters.
523+
"""
524+
Base.@kwdef struct AtmosMoistureModel{MM, PM, CM, NCFM, CCDPS}
552525
moisture_model::MM = nothing
553526
precip_model::PM = nothing
554527
cloud_model::CM = nothing
555528
noneq_cloud_formation_mode::NCFM = nothing
556529
call_cloud_diagnostics_per_stage::CCDPS = nothing
530+
end
531+
532+
"""
533+
AtmosForcing
534+
535+
Groups forcing-related models and parameters.
536+
"""
537+
Base.@kwdef struct AtmosForcing{F, S, EXTFORCING}
557538
forcing_type::F = nothing
558539
subsidence::S = nothing
540+
external_forcing::EXTFORCING = nothing
541+
end
542+
543+
"""
544+
AtmosRadiation
559545
560-
# Currently only relevant for RRTMGP, but will hopefully become standalone
561-
# in the future
562-
"""What to do with ozone for radiation (when using RRTMGP)"""
546+
Groups radiation-related models and parameters.
547+
"""
548+
Base.@kwdef struct AtmosRadiation{RM, OZ, CO2, IN}
549+
radiation_mode::RM = nothing
563550
ozone::OZ = nothing
564-
"""What to do with co2 for radiation (when using RRTMGP)"""
565551
co2::CO2 = nothing
552+
insolation::IN = nothing
553+
end
566554

567-
radiation_mode::RM = nothing
555+
"""
556+
AtmosAdvection
557+
558+
Groups advection-related models and parameters.
559+
"""
560+
Base.@kwdef struct AtmosAdvection{LA, AT}
568561
ls_adv::LA = nothing
569-
external_forcing::EXTFORCING = nothing
570-
scm_coriolis::EC = nothing
571562
advection_test::AT = nothing
563+
end
564+
565+
"""
566+
AtmosTurbconv
567+
568+
Groups turbulence convection-related models and parameters.
569+
"""
570+
Base.@kwdef struct AtmosTurbconv{EC, EDMFX, TCM, SAM, SEDM, SNPM, SMM, SL}
571+
scm_coriolis::EC = nothing
572572
edmfx_model::EDMFX = nothing
573573
turbconv_model::TCM = nothing
574-
non_orographic_gravity_wave::NOGW = nothing
575-
orographic_gravity_wave::OGW = nothing
576-
hyperdiff::HD = nothing
577-
vert_diff::VD = nothing
578-
diff_mode::DM = nothing
579574
sgs_adv_mode::SAM = nothing
580-
"""sgs_entr_detr_mode == Implicit() only works if sgs_adv_mode == Implicit()"""
581575
sgs_entr_detr_mode::SEDM = nothing
582-
"""sgs_nh_pressure_mode == Implicit() only works if sgs_adv_mode == Implicit()"""
583576
sgs_nh_pressure_mode::SNPM = nothing
584-
"""sgs_mf_mode == Implicit() only works if sgs_adv_mode == Implicit() and diff_mode == Implicit()"""
585577
sgs_mf_mode::SMM = nothing
586-
viscous_sponge::VS = nothing
587578
smagorinsky_lilly::SL = nothing
579+
end
580+
581+
"""
582+
AtmosGravityWave
583+
584+
Groups gravity wave-related models and parameters.
585+
"""
586+
Base.@kwdef struct AtmosGravityWave{NOGW, OGW}
587+
non_orographic_gravity_wave::NOGW = nothing
588+
orographic_gravity_wave::OGW = nothing
589+
end
590+
591+
"""
592+
AtmosVertDiff
593+
594+
Groups vertical diffusion-related models and parameters.
595+
"""
596+
Base.@kwdef struct AtmosVertDiff{VD, DM}
597+
vert_diff::VD = nothing
598+
diff_mode::DM = nothing
599+
end
600+
601+
"""
602+
AtmosSponge
603+
604+
Groups sponge-related models and parameters.
605+
"""
606+
Base.@kwdef struct AtmosSponge{VS, RS}
607+
viscous_sponge::VS = nothing
588608
rayleigh_sponge::RS = nothing
609+
end
610+
611+
"""
612+
AtmosSurface
613+
614+
Groups surface-related models and parameters.
615+
"""
616+
Base.@kwdef struct AtmosSurface{ST, SM, SA}
589617
sfc_temperature::ST = nothing
590-
insolation::IN = nothing
591-
"""Whether to apply surface flux tendency (independent of surface conditions)"""
592-
disable_surface_flux_tendency::Bool = false
593618
surface_model::SM = nothing
594619
surface_albedo::SA = nothing
620+
end
621+
622+
# Add broadcastable for the new grouped types
623+
Base.broadcastable(x::AtmosMoistureModel) = tuple(x)
624+
Base.broadcastable(x::AtmosForcing) = tuple(x)
625+
Base.broadcastable(x::AtmosRadiation) = tuple(x)
626+
Base.broadcastable(x::AtmosAdvection) = tuple(x)
627+
Base.broadcastable(x::AtmosTurbconv) = tuple(x)
628+
Base.broadcastable(x::AtmosGravityWave) = tuple(x)
629+
Base.broadcastable(x::AtmosVertDiff) = tuple(x)
630+
Base.broadcastable(x::AtmosSponge) = tuple(x)
631+
Base.broadcastable(x::AtmosSurface) = tuple(x)
632+
633+
Base.@kwdef struct AtmosModel{
634+
MOISTURE,
635+
FORCING,
636+
RADIATION,
637+
ADVECTION,
638+
TURBCONV,
639+
GRAVITY_WAVE,
640+
HD,
641+
VERT_DIFF,
642+
SPONGE,
643+
SURFACE,
644+
NUM,
645+
}
646+
moisture::MOISTURE = AtmosMoistureModel()
647+
forcing::FORCING = AtmosForcing()
648+
radiation::RADIATION = AtmosRadiation()
649+
advection::ADVECTION = AtmosAdvection()
650+
turbconv::TURBCONV = AtmosTurbconv()
651+
gravity_wave::GRAVITY_WAVE = AtmosGravityWave()
652+
hyperdiff::HD = nothing
653+
vert_diff::VERT_DIFF = AtmosVertDiff()
654+
sponge::SPONGE = AtmosSponge()
655+
surface::SURFACE = AtmosSurface()
595656
numerics::NUM = nothing
657+
658+
"""Whether to apply surface flux tendency (independent of surface conditions)"""
659+
disable_surface_flux_tendency::Bool = false
660+
end
661+
662+
# ============================================================================
663+
# BACKWARD COMPATIBILITY FOR AtmosModel PROPERTY ACCESS
664+
# ============================================================================
665+
# We reduced AtmosModel type parameters from ~33 to ~11 by grouping related fields.
666+
# Old: atmos.moisture_model → New: atmos.moisture.moisture_model
667+
#
668+
# This automatically forwards old property access to the new grouped structure,
669+
# so existing code continues to work without changes.
670+
671+
# Map grouped struct types to their AtmosModel field names
672+
const ATMOS_MODEL_GROUPS = (
673+
(AtmosMoistureModel, :moisture),
674+
(AtmosForcing, :forcing),
675+
(AtmosRadiation, :radiation),
676+
(AtmosAdvection, :advection),
677+
(AtmosTurbconv, :turbconv),
678+
(AtmosGravityWave, :gravity_wave),
679+
(AtmosVertDiff, :vert_diff),
680+
(AtmosSponge, :sponge),
681+
(AtmosSurface, :surface),
682+
)
683+
684+
# Auto-generate lookup: property_name => group_field
685+
const GROUPED_PROPERTY_MAP = let
686+
property_map = Dict{Symbol, Symbol}()
687+
for (group_type, group_field) in ATMOS_MODEL_GROUPS
688+
for property in fieldnames(group_type)
689+
property_map[property] = group_field
690+
end
691+
end
692+
property_map
693+
end
694+
695+
# Forward property access: atmos.moisture_model → atmos.moisture.moisture_model
696+
function Base.getproperty(atmos::AtmosModel, property_name::Symbol)
697+
if haskey(GROUPED_PROPERTY_MAP, property_name)
698+
group_field = GROUPED_PROPERTY_MAP[property_name]
699+
group = getfield(atmos, group_field)
700+
return getfield(group, property_name)
701+
else
702+
return getfield(atmos, property_name)
703+
end
596704
end
597705

598706
Base.broadcastable(x::AtmosModel) = tuple(x)

0 commit comments

Comments
 (0)