From 0a5b0f189eaf54e804e4d699c5a6d62dc04471a4 Mon Sep 17 00:00:00 2001 From: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> Date: Thu, 13 Jun 2024 10:08:44 +0200 Subject: [PATCH] changes following review --- .../dycore/nh_solve/solve_nonhydro.py | 139 +++++++++++------- .../src/icon4py/model/common/constants.py | 8 - .../tests/metric_tests/test_metric_fields.py | 5 +- 3 files changed, 87 insertions(+), 65 deletions(-) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py index 61bb8177df..d5f89b94fe 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py @@ -169,50 +169,70 @@ from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.settings import backend -from enum import IntEnum +import enum # flake8: noqa log = logging.getLogger(__name__) class TimeSteppingScheme(enum.IntEnum): -""" Parameter called `itime_scheme` in ICON namelist.""" + """Parameter called `itime_scheme` in ICON namelist.""" + #: Contravariant vertical velocity is computed in the predictor step only, velocity tendencies are computed in the corrector step only - OPTIMAL = 4 - VERTICAL = 5 # Contravariant vertical velocity is computed in both substeps - VELOCITY = 6 # As itime_vertical, but velocity tendencies are also computed in both substeps + MOST_EFFICIENT = 4 + #: Contravariant vertical velocity is computed in both substeps (beneficial for numerical stability in very-high resolution setups with extremely steep slopes) + STABLE = 5 + #: As STABLE, but velocity tendencies are also computed in both substeps (no benefit, but more expensive) + EXPENSIVE = 6 class RayleighType(enum.IntEnum): - CLASSIC = 1 # classical Rayleigh damping, which makes use of a reference state. - KLEMP = 2 # Klemp (2008) type Rayleigh damping + #: classical Rayleigh damping, which makes use of a reference state. + CLASSIC = 1 + #: Klemp (2008) type Rayleigh damping + KLEMP = 2 class DivergenceDampingType(enum.IntEnum): - TWO_DIMENSIONAL = 2 # divergence damping acting on 2D divergence - THREE_DIMENSIONAL = 3 # divergence damping acting on 3D divergence - COMBINED = 32 # combination of 3D div.damping in the troposphere with transition to 2D div. damping in the stratosphere - - -class DivergenceDamping_Order(IntEnum): - SECOND_ORDER = 2 # 2nd order divergence damping - FOURTH_ORDER = 4 # 4th order divergence damping - COMBINED = 24 # combined 2nd and 4th orders divergence damping and enhanced vertical wind off - centering during initial spinup phase - - -class HorizontalPressureDiscretizationType(IntEnum): -""" Parameter called igradp_method in ICON namelist. """ - NORMAL = 1 # conventional discretization with metric correction term - TAYLOR = 2 # Taylor-expansion-based reconstruction of pressure - TAYLOR_HYDRO = 3 # Similar discretization as igradp_method_taylor, but uses hydrostatic approximation for downward extrapolation over steep slopes - POLYNOMIAL = 4 # Cubic / quadratic polynomial interpolation for pressure reconstruction - POLYNOMIAL_HYDRO = 5 # Same as igradp_method_polynomial, but hydrostatic approximation for downward extrapolation over steep slopes + #: divergence damping acting on 2D divergence + TWO_DIMENSIONAL = 2 + #: divergence damping acting on 3D divergence + THREE_DIMENSIONAL = 3 + #: combination of 3D div.damping in the troposphere with transition to 2D div. damping in the stratosphere + COMBINED = 32 + + +class DivergenceDampingOrder(enum.IntEnum): + #: 2nd order divergence damping + SECOND_ORDER = 2 + #: 4th order divergence damping + FOURTH_ORDER = 4 + #: combined 2nd and 4th orders divergence damping and enhanced vertical wind off - centering during initial spinup phase + COMBINED = 24 + + +class HorizontalPressureDiscretizationType(enum.IntEnum): + """Parameter called igradp_method in ICON namelist.""" + + #: conventional discretization with metric correction term + CONVENTIONAL = 1 + #: Taylor-expansion-based reconstruction of pressure + TAYLOR = 2 + #: Similar discretization as igradp_method_taylor, but uses hydrostatic approximation for downward extrapolation over steep slopes + TAYLOR_HYDRO = 3 + #: Cubic / quadratic polynomial interpolation for pressure reconstruction + POLYNOMIAL = 4 + #: Same as igradp_method_polynomial, but hydrostatic approximation for downward extrapolation over steep slopes + POLYNOMIAL_HYDRO = 5 class RhoThetaAdvectionType(enum.IntEnum): -""" Parameter called iadv_rhotheta in ICON namelist.""" - SIMPLE = 1 # simple 2nd order upwind-biased scheme - MIURA = 2 # 2nd order Miura horizontal + """Parameter called iadv_rhotheta in ICON namelist.""" + + #: simple 2nd order upwind-biased scheme + SIMPLE = 1 + #: 2nd order Miura horizontal + MIURA = 2 @dataclass @@ -278,16 +298,16 @@ class NonHydrostaticConfig: def __init__( self, - itime_scheme: ItimeScheme = ItimeScheme.OPTIMAL, - iadv_rhotheta: Iadv_RhoTheta = Iadv_RhoTheta.MIURA, - igradp_method: IGradp_Method = IGradp_Method.TAYLOR_HYDRO, + itime_scheme: TimeSteppingScheme = TimeSteppingScheme.MOST_EFFICIENT, + iadv_rhotheta: RhoThetaAdvectionType = RhoThetaAdvectionType.MIURA, + igradp_method: HorizontalPressureDiscretizationType = HorizontalPressureDiscretizationType.TAYLOR_HYDRO, ndyn_substeps_var: float = 5.0, - rayleigh_type: Rayleigh_Type = Rayleigh_Type.KLEMP, + rayleigh_type: RayleighType = RayleighType.KLEMP, rayleigh_coeff: float = 0.05, - divdamp_order: DivergenceDamping_Order = DivergenceDamping_Order.COMBINED, # the ICON default is 4, + divdamp_order: DivergenceDampingOrder = DivergenceDampingOrder.COMBINED, # the ICON default is 4, is_iau_active: bool = False, iau_wgt_dyn: float = 0.0, - divdamp_type: DivergenceDamping_Type = DivergenceDamping_Type.THREE_DIMENSIONAL, + divdamp_type: DivergenceDampingType = DivergenceDampingType.THREE_DIMENSIONAL, divdamp_trans_start: float = 12500.0, divdamp_trans_end: float = 17500.0, l_vert_nested: bool = False, @@ -374,16 +394,16 @@ def _validate(self): if self.l_vert_nested: raise NotImplementedError("Vertical nesting support not implemented") - if self.igradp_method != 3: + if self.igradp_method != HorizontalPressureDiscretizationType.TAYLOR_HYDRO: raise NotImplementedError("igradp_method can only be 3") - if self.itime_scheme != 4: + if self.itime_scheme != TimeSteppingScheme.MOST_EFFICIENT: raise NotImplementedError("itime_scheme can only be 4") - if self.divdamp_order != 24: + if self.divdamp_order != DivergenceDampingOrder.COMBINED: raise NotImplementedError("divdamp_order can only be 24") - if self.divdamp_type == 32: + if self.divdamp_type == DivergenceDampingType.COMBINED: raise NotImplementedError("divdamp_type with value 32 not yet implemented") @@ -531,7 +551,7 @@ def _allocate_local_fields(self): def set_timelevels(self, nnow, nnew): # Set time levels of ddt_adv fields for call to velocity_tendencies - if self.config.itime_scheme == 4: + if self.config.itime_scheme == TimeSteppingScheme.MOST_EFFICIENT: self.ntl1 = nnow self.ntl2 = nnew else: @@ -677,7 +697,7 @@ def run_predictor_step( f"running predictor step: dtime = {dtime}, init = {l_init}, recompute = {l_recompute} " ) if l_init or l_recompute: - if self.config.itime_scheme == 4 and not l_init: + if self.config.itime_scheme == TimeSteppingScheme.MOST_EFFICIENT and not l_init: lvn_only = True # Recompute only vn tendency else: lvn_only = False @@ -793,7 +813,7 @@ def run_predictor_step( offset_provider={}, ) - if self.config.igradp_method == 3: + if self.config.igradp_method == HorizontalPressureDiscretizationType.TAYLOR_HYDRO: nhsolve_prog.predictor_stencils_4_5_6( wgtfacq_c_dsl=self.metric_state_nonhydro.wgtfacq_c, z_exner_ex_pr=self.z_exner_ex_pr, @@ -855,7 +875,7 @@ def run_predictor_step( offset_provider=self.grid.offset_providers, ) - if self.config.igradp_method == 3: + if self.config.igradp_method == HorizontalPressureDiscretizationType.TAYLOR_HYDRO: # Second vertical derivative of perturbation Exner pressure (hydrostatic approximation) compute_approx_of_2nd_vertical_derivative_of_exner( z_theta_v_pr_ic=self.z_theta_v_pr_ic, @@ -888,7 +908,7 @@ def run_predictor_step( ) # Compute rho and theta at edges for horizontal flux divergence term - if self.config.iadv_rhotheta == 1: + if self.config.iadv_rhotheta == RhoThetaAdvectionType.SIMPLE: mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl( p_cell_in=prognostic_state[nnow].rho, c_intp=self.interpolation_state.c_intp, @@ -909,7 +929,7 @@ def run_predictor_step( vertical_end=self.grid.num_levels, offset_provider=self.grid.offset_providers, ) - elif self.config.iadv_rhotheta == 2: + elif self.config.iadv_rhotheta == RhoThetaAdvectionType.MIURA: # Compute Green-Gauss gradients for rho and theta mo_math_gradients_grad_green_gauss_cell_dsl( p_grad_1_u=self.z_grad_rth_1, @@ -947,7 +967,7 @@ def run_predictor_step( vertical_end=self.grid.num_levels, offset_provider={}, ) - if self.config.iadv_rhotheta == 2: + if self.config.iadv_rhotheta == RhoThetaAdvectionType.MIURA: # Compute upwind-biased values for rho and theta starting from centered differences # Note: the length of the backward trajectory should be 0.5*dtime*(vn,vt) in order to arrive # at a second-order accurate FV discretization, but twice the length is needed for numerical stability @@ -991,7 +1011,7 @@ def run_predictor_step( offset_provider=self.grid.offset_providers, ) - if self.config.igradp_method == 3: + if self.config.igradp_method == HorizontalPressureDiscretizationType.TAYLOR_HYDRO: # horizontal gradient of Exner pressure, including metric correction # horizontal gradient of Exner pressure, Taylor-expansion-based reconstruction @@ -1024,7 +1044,7 @@ def run_predictor_step( offset_provider=self.grid.offset_providers, ) # compute hydrostatically approximated correction term that replaces downward extrapolation - if self.config.igradp_method == 3: + if self.config.igradp_method == HorizontalPressureDiscretizationType.TAYLOR_HYDRO: compute_hydrostatic_correction_term( theta_v=prognostic_state[nnow].theta_v, ikoffset=self.metric_state_nonhydro.vertoffset_gradp, @@ -1044,7 +1064,7 @@ def run_predictor_step( lowest_level = self.grid.num_levels - 1 hydro_corr_horizontal = as_field((EdgeDim,), self.z_hydro_corr.asnumpy()[:, lowest_level]) - if self.config.igradp_method == 3: + if self.config.igradp_method == HorizontalPressureDiscretizationType.TAYLOR_HYDRO: apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure( ipeidx_dsl=self.metric_state_nonhydro.ipeidx_dsl, pg_exdist=self.metric_state_nonhydro.pg_exdist, @@ -1300,7 +1320,7 @@ def run_predictor_step( offset_provider={}, ) - if self.config.rayleigh_type == constants.RayleighType.RAYLEIGH_KLEMP: + if self.config.rayleigh_type == RayleighType.KLEMP: apply_rayleigh_damping_mechanism( z_raylfac=self.z_raylfac, w_1=prognostic_state[nnew].w_1, @@ -1541,7 +1561,7 @@ def run_corrector_step( offset_provider=self.grid.offset_providers, ) - if self.config.itime_scheme == 4: + if self.config.itime_scheme == TimeSteppingScheme.MOST_EFFICIENT: log.debug(f"corrector: start stencil 23") add_temporal_tendencies_to_vn_by_interpolating_between_time_levels( vn_nnow=prognostic_state[nnow].vn, @@ -1562,7 +1582,10 @@ def run_corrector_step( offset_provider={}, ) - if self.config.divdamp_order == 24 or self.config.divdamp_order == 4: + if ( + self.config.divdamp_order == DivergenceDampingOrder.COMBINED + or self.config.divdamp_order == DivergenceDampingOrder.FOURTH_ORDER + ): # verified for e-10 log.debug(f"corrector start stencil 25") compute_graddiv2_of_vn( @@ -1576,7 +1599,10 @@ def run_corrector_step( offset_provider=self.grid.offset_providers, ) - if self.config.divdamp_order == 24 and scal_divdamp_o2 > 1.0e-6: + if ( + self.config.divdamp_order == DivergenceDampingOrder.COMBINED + and scal_divdamp_o2 > 1.0e-6 + ): log.debug(f"corrector: start stencil 26") apply_2nd_order_divergence_damping( z_graddiv_vn=z_fields.z_graddiv_vn, @@ -1590,7 +1616,10 @@ def run_corrector_step( ) # TODO: this does not get accessed in FORTRAN - if self.config.divdamp_order == 24 and divdamp_fac_o2 <= 4 * self.config.divdamp_fac: + if ( + self.config.divdamp_order == DivergenceDampingOrder.COMBINED + and divdamp_fac_o2 <= 4 * self.config.divdamp_fac + ): if self.grid.limited_area: log.debug("corrector: start stencil 27") apply_weighted_2nd_and_4th_order_divergence_damping( @@ -1702,7 +1731,7 @@ def run_corrector_step( offset_provider=self.grid.offset_providers, ) - if self.config.itime_scheme == 4: + if self.config.itime_scheme == TimeSteppingScheme.MOST_EFFICIENT: log.debug(f"corrector start stencil 42 44 45 45b") nhsolve_prog.stencils_42_44_45_45b( z_w_expl=z_fields.z_w_expl, @@ -1850,7 +1879,7 @@ def run_corrector_step( offset_provider={}, ) - if self.config.rayleigh_type == constants.RayleighType.RAYLEIGH_KLEMP: + if self.config.rayleigh_type == RayleighType.KLEMP: log.debug(f"corrector start stencil 54") apply_rayleigh_damping_mechanism( z_raylfac=self.z_raylfac, diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py index 76bc9adf57..7688760a25 100644 --- a/model/common/src/icon4py/model/common/constants.py +++ b/model/common/src/icon4py/model/common/constants.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later import sys -from enum import IntEnum from typing import Final from icon4py.model.common.type_alias import wpfloat @@ -75,10 +74,3 @@ #: default physics to dynamics time step ratio # TODO (magdalena) not a constant, this is a default config parameter DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Final[float] = 5.0 - - -class RayleighType(IntEnum): - RAYLEIGH_CLASSIC: Final[ - int - ] = 1 # classical Rayleigh damping, which makes use of a reference state. - RAYLEIGH_KLEMP: Final[int] = 2 # Klemp (2008) type Rayleigh damping diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index ec93c1c297..ac667456f3 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -18,6 +18,7 @@ from gt4py.next import as_field from gt4py.next.ffront.fbuiltins import int32 +from icon4py.model.atmosphere.dycore.nh_solve.solve_nonhydro import RayleighType from icon4py.model.common import constants from icon4py.model.common.dimension import ( CellDim, @@ -205,8 +206,8 @@ def test_compute_rayleigh_w(icon_grid, metrics_savepoint, grid_savepoint, backen vct_a=grid_savepoint.vct_a(), damping_height=damping_height, rayleigh_type=rayleigh_type, - rayleigh_classic=constants.RayleighType.RAYLEIGH_CLASSIC, - rayleigh_klemp=constants.RayleighType.RAYLEIGH_KLEMP, + rayleigh_classic=RayleighType.CLASSIC, + rayleigh_klemp=RayleighType.KLEMP, rayleigh_coeff=rayleigh_coeff, vct_a_1=vct_a_1, pi_const=math.pi,