diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py
index 10c7e6ed15..144e000edf 100644
--- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py
+++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py
@@ -68,6 +68,7 @@
CPD,
DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO,
GAS_CONSTANT_DRY_AIR,
+ dbl_eps,
)
from icon4py.model.common.decomposition.definitions import ExchangeRuntime, SingleNodeExchange
from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim
@@ -771,6 +772,7 @@ def _do_diffusion_step(
theta_v=prognostic_state.theta_v,
theta_ref_mc=self.metric_state.theta_ref_mc,
thresh_tdiff=self.thresh_tdiff,
+ smallest_vpfloat=dbl_eps,
kh_smag_e=self.kh_smag_e,
horizontal_start=edge_start_nudging,
horizontal_end=edge_end_halo,
diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py
index 76b5f7289d..d5db49b331 100644
--- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py
+++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py
@@ -16,9 +16,10 @@
from gt4py.next import as_field
from gt4py.next.common import Dimension, Field
from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import broadcast, int32, maximum, minimum
+from gt4py.next.ffront.fbuiltins import broadcast, int32, minimum
-from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim
+from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim
+from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift
# TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py
@@ -98,35 +99,6 @@ def setup_fields_for_initial_step(
_setup_fields_for_initial_step(k4, hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit))
-@field_operator
-def _en_smag_fac_for_zero_nshift(
- vect_a: Field[[KDim], float],
- hdiff_smag_fac: float,
- hdiff_smag_fac2: float,
- hdiff_smag_fac3: float,
- hdiff_smag_fac4: float,
- hdiff_smag_z: float,
- hdiff_smag_z2: float,
- hdiff_smag_z3: float,
- hdiff_smag_z4: float,
-) -> Field[[KDim], float]:
- dz21 = hdiff_smag_z2 - hdiff_smag_z
- alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21
- df32 = hdiff_smag_fac3 - hdiff_smag_fac2
- df42 = hdiff_smag_fac4 - hdiff_smag_fac2
- dz32 = hdiff_smag_z3 - hdiff_smag_z2
- dz42 = hdiff_smag_z4 - hdiff_smag_z2
-
- bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32))
- aqdr = df32 / dz32 - bqdr * dz32
- zf = 0.5 * (vect_a + vect_a(Koff[1]))
-
- dzlin = minimum(dz21, maximum(0.0, zf - hdiff_smag_z))
- dzqdr = minimum(dz42, maximum(0.0, zf - hdiff_smag_z2))
- enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr)
- return enh_smag_fac
-
-
@field_operator
def _init_diffusion_local_fields_for_regular_timestemp(
k4: float,
@@ -143,7 +115,7 @@ def _init_diffusion_local_fields_for_regular_timestemp(
) -> tuple[Field[[KDim], float], Field[[KDim], float], Field[[KDim], float]]:
diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps)
smag_limit = _setup_smag_limit(diff_multfac_vn)
- enh_smag_fac = _en_smag_fac_for_zero_nshift(
+ enh_smag_fac = en_smag_fac_for_zero_nshift(
vect_a,
hdiff_smag_fac,
hdiff_smag_fac2,
diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py
index a02ce5b67c..6506fd20de 100644
--- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py
+++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion.py
@@ -19,14 +19,13 @@
from icon4py.model.atmosphere.diffusion.diffusion_utils import scale_k
from icon4py.model.common.grid.horizontal import CellParams, EdgeParams
from icon4py.model.common.grid.vertical import VerticalModelParams
+from icon4py.model.common.test_utils.reference_funcs import enhanced_smagorinski_factor_numpy
from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionInitSavepoint
-from .utils import (
- diff_multfac_vn_numpy,
- enhanced_smagorinski_factor_numpy,
- smag_limit_numpy,
- verify_diffusion_fields,
-)
+from .utils import diff_multfac_vn_numpy, smag_limit_numpy, verify_diffusion_fields
+
+
+backend = icon4py.model.atmosphere.diffusion.diffusion.backend
backend = icon4py.model.atmosphere.diffusion.diffusion.backend
diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py
index 46e135d326..dfb4dba5d0 100644
--- a/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py
+++ b/model/atmosphere/diffusion/tests/diffusion_tests/test_diffusion_utils.py
@@ -16,7 +16,6 @@
from icon4py.model.atmosphere.diffusion.diffusion import DiffusionParams
from icon4py.model.atmosphere.diffusion.diffusion_utils import (
- _en_smag_fac_for_zero_nshift,
_setup_runtime_diff_multfac_vn,
_setup_smag_limit,
scale_k,
@@ -27,7 +26,7 @@
from icon4py.model.common.grid.simple import SimpleGrid
from icon4py.model.common.test_utils.helpers import random_field, zero_field
-from .utils import diff_multfac_vn_numpy, enhanced_smagorinski_factor_numpy, smag_limit_numpy
+from .utils import diff_multfac_vn_numpy, smag_limit_numpy
def initial_diff_multfac_vn_numpy(shape, k4, hdiff_efdt_ratio):
@@ -104,25 +103,10 @@ def test_diff_multfac_vn_smag_limit_for_loop_run_with_k4_substeps(backend):
assert np.allclose(expected_smag_limit, smag_limit.asnumpy())
-def test_init_enh_smag_fac(backend):
- grid = SimpleGrid()
- enh_smag_fac = zero_field(grid, KDim)
- a_vec = random_field(grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
- fac = (0.67, 0.5, 1.3, 0.8)
- z = (0.1, 0.2, 0.3, 0.4)
-
- enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, a_vec.asnumpy())
-
- _en_smag_fac_for_zero_nshift.with_backend(backend)(
- a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim}
- )
- assert np.allclose(enhanced_smag_fac_np, enh_smag_fac.asnumpy())
-
-
def test_set_zero_vertex_k(backend):
grid = SimpleGrid()
f = random_field(grid, VertexDim, KDim)
- set_zero_v_k.with_backend(backend)(f, offset_provider={})
+ set_zero_v_k(f, offset_provider={})
assert np.allclose(0.0, f.asnumpy())
diff --git a/model/atmosphere/diffusion/tests/diffusion_tests/utils.py b/model/atmosphere/diffusion/tests/diffusion_tests/utils.py
index b75dd10c85..082a94516f 100644
--- a/model/atmosphere/diffusion/tests/diffusion_tests/utils.py
+++ b/model/atmosphere/diffusion/tests/diffusion_tests/utils.py
@@ -58,19 +58,3 @@ def smag_limit_numpy(func, *args):
def diff_multfac_vn_numpy(shape, k4, substeps):
factor = min(1.0 / 128.0, k4 * substeps / 3.0)
return factor * np.ones(shape)
-
-
-def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec: np.ndarray):
- alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0])
- df32 = factor_in[2] - factor_in[1]
- df42 = factor_in[3] - factor_in[1]
- dz32 = heigths_in[2] - heigths_in[1]
- dz42 = heigths_in[3] - heigths_in[1]
- bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32))
- aqdr = df32 / dz32 - bqdr * dz32
- zf = 0.5 * (a_vec[:-1] + a_vec[1:])
- max0 = np.maximum(0.0, zf - heigths_in[0])
- dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0)
- max1 = np.maximum(0.0, zf - heigths_in[1])
- dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1)
- return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr)
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 1be86b7b99..1b2bfbb318 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
@@ -10,6 +10,7 @@
# distribution for a copy of the license or check .
#
# SPDX-License-Identifier: GPL-3.0-or-later
+import logging
from typing import Final, Optional
import numpy as np
@@ -145,10 +146,8 @@
from icon4py.model.atmosphere.dycore.state_utils.utils import (
_allocate,
_allocate_indices,
- _calculate_bdy_divdamp,
- _en_smag_fac_for_zero_nshift,
+ _calculate_divdamp_fields,
compute_z_raylfac,
- scal_divdamp_calcs,
set_zero_c_k,
set_zero_e_k,
)
@@ -156,13 +155,16 @@
from icon4py.model.atmosphere.dycore.velocity.velocity_advection import VelocityAdvection
from icon4py.model.common.decomposition.definitions import ExchangeRuntime, SingleNodeExchange
from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim
-from icon4py.model.common.grid.horizontal import EdgeParams, HorizontalMarkerIndex
+from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex
from icon4py.model.common.grid.icon import IconGrid
from icon4py.model.common.grid.vertical import VerticalModelParams
+from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift
from icon4py.model.common.states.prognostic_state import PrognosticState
backend = run_gtfn
+# flake8: noqa
+log = logging.getLogger(__name__)
class NonHydrostaticConfig:
@@ -197,10 +199,10 @@ def __init__(
divdamp_fac2: float = 0.004,
divdamp_fac3: float = 0.004,
divdamp_fac4: float = 0.004,
- divdamp_z: float = 32500,
- divdamp_z2: float = 40000,
- divdamp_z3: float = 60000,
- divdamp_z4: float = 80000,
+ divdamp_z: float = 32500.0,
+ divdamp_z2: float = 40000.0,
+ divdamp_z3: float = 60000.0,
+ divdamp_z4: float = 80000.0,
):
# parameters from namelist diffusion_nml
self.itime_scheme: int = itime_scheme
@@ -277,20 +279,6 @@ def __init__(self, config: NonHydrostaticConfig):
if self.kstart_moist != 1:
raise NotImplementedError("kstart_moist can only be 1")
- self.alin = (config.divdamp_fac2 - config.divdamp_fac) / (
- config.divdamp_z2 - config.divdamp_z
- )
-
- self.df32 = config.divdamp_fac3 - config.divdamp_fac2
- self.dz32 = config.divdamp_z3 - config.divdamp_z2
- self.df42 = config.divdamp_fac4 - config.divdamp_fac2
- self.dz42 = config.divdamp_z4 - config.divdamp_z2
-
- self.bqdr = (self.df42 * self.dz32 - self.df32 * self.dz42) / (
- self.dz32 * self.dz42 * (self.dz42 - self.dz32)
- )
- self.aqdr = self.df32 / self.dz32 - self.bqdr * self.dz32
-
class SolveNonhydro:
def __init__(self, exchange: ExchangeRuntime = SingleNodeExchange()):
@@ -303,11 +291,12 @@ def __init__(self, exchange: ExchangeRuntime = SingleNodeExchange()):
self.interpolation_state: Optional[InterpolationState] = None
self.vertical_params: Optional[VerticalModelParams] = None
self.edge_geometry: Optional[EdgeParams] = None
- self.cell_areas: Optional[Field[[CellDim], float]] = None
+ self.cell_params: Optional[CellParams] = None
self.velocity_advection: Optional[VelocityAdvection] = None
self.l_vert_nested: bool = False
- self.enh_divdamp_fac = None
- self.scal_divdamp: Field[[KDim], float] = None
+ self.enh_divdamp_fac: Optional[Field[[KDim], float]] = None
+ self.scal_divdamp: Optional[Field[[KDim], float]] = None
+ self._bdy_divdamp: Optional[Field[[KDim], float]] = None
self.p_test_run = True
self.jk_start = 0 # used in stencil_55
self.ntl1 = 0
@@ -322,12 +311,8 @@ def init(
interpolation_state: InterpolationState,
vertical_params: VerticalModelParams,
edge_geometry: EdgeParams,
- cell_areas: Field[[CellDim], float],
+ cell_geometry: CellParams,
owner_mask: Field[[CellDim], bool],
- a_vec: Field[[KDim], float],
- enh_smag_fac: Field[[KDim], float],
- fac: tuple,
- z: tuple,
):
"""
Initialize NonHydrostatic granule with configuration.
@@ -341,7 +326,7 @@ def init(
self.interpolation_state: InterpolationState = interpolation_state
self.vertical_params = vertical_params
self.edge_geometry = edge_geometry
- self.cell_areas = cell_areas
+ self.cell_params = cell_geometry
self.velocity_advection = VelocityAdvection(
grid,
metric_state_nonhydro,
@@ -360,21 +345,20 @@ def init(
else:
self.jk_start = 0
- out = enh_smag_fac
- _en_smag_fac_for_zero_nshift.with_backend(run_gtfn)(
- a_vec, *fac, *z, out=enh_smag_fac, offset_provider={"Koff": KDim}
- )
- self.enh_divdamp_fac = enh_smag_fac
-
- cell_areas_avg = np.sum(cell_areas.asnumpy()) / float(self.grid.num_cells)
- # TODO @tehrengruber: fix power
- scal_divdamp_calcs.with_backend(backend)(
- enh_smag_fac,
- out,
- cell_areas_avg,
- offset_provider={},
+ en_smag_fac_for_zero_nshift.with_backend(run_gtfn)(
+ self.vertical_params.vct_a,
+ self.config.divdamp_fac,
+ self.config.divdamp_fac2,
+ self.config.divdamp_fac3,
+ self.config.divdamp_fac4,
+ self.config.divdamp_z,
+ self.config.divdamp_z2,
+ self.config.divdamp_z3,
+ self.config.divdamp_z4,
+ out=self.enh_divdamp_fac,
+ offset_provider={"Koff": KDim},
)
- self.scal_divdamp = out
+
self.p_test_run = True
self._initialized = True
@@ -408,6 +392,9 @@ def _allocate_local_fields(self):
self.z_w_concorr_me = _allocate(EdgeDim, KDim, grid=self.grid)
self.z_hydro_corr_horizontal = _allocate(EdgeDim, grid=self.grid)
self.z_raylfac = _allocate(KDim, grid=self.grid)
+ self.enh_divdamp_fac = _allocate(KDim, grid=self.grid)
+ self._bdy_divdamp = _allocate(KDim, grid=self.grid)
+ self.scal_divdamp = _allocate(KDim, grid=self.grid)
def set_timelevels(self, nnow, nnew):
# Set time levels of ddt_adv fields for call to velocity_tendencies
@@ -425,7 +412,7 @@ def time_step(
prep_adv: PrepAdvection,
z_fields: ZFields, # TODO (Magdalena): move local fields to SolveNonHydro
nh_constants: NHConstants,
- bdy_divdamp: Field[[KDim], float],
+ divdamp_fac_o2: float,
dtime: float,
idyn_timestep: float,
l_recompute: bool,
@@ -435,18 +422,9 @@ def time_step(
lclean_mflx: bool,
lprep_adv: bool,
):
- # Coefficient for reduced fourth-order divergence damping along nest boundaries
- _calculate_bdy_divdamp(
- self.scal_divdamp,
- self.config.nudge_max_coeff,
- constants.dbl_eps,
- out=bdy_divdamp,
- offset_provider={},
+ log.info(
+ f"running timestep: dtime = {dtime}, dyn_timestep = {idyn_timestep}, init = {l_init}, recompute = {l_recompute}, prep_adv = {lprep_adv} "
)
-
- # scaling factor for second-order divergence damping: divdamp_fac_o2*delta_x**2
- # delta_x**2 is approximated by the mean cell area
-
start_cell_lb = self.grid.get_start_index(
CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim)
)
@@ -489,15 +467,17 @@ def time_step(
prognostic_state=prognostic_state_ls,
z_fields=z_fields,
prep_adv=prep_adv,
+ divdamp_fac_o2=divdamp_fac_o2,
dtime=dtime,
nnew=nnew,
nnow=nnow,
lclean_mflx=lclean_mflx,
nh_constants=nh_constants,
- bdy_divdamp=bdy_divdamp,
lprep_adv=lprep_adv,
)
-
+ log.info(
+ f"running corrector step: dtime = {dtime}, dyn_timestep = {idyn_timestep}, prep_adv = {lprep_adv}, divdamp_fac_od = {divdamp_fac_o2} "
+ )
start_cell_lb = self.grid.get_start_index(
CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim)
)
@@ -559,6 +539,9 @@ def run_predictor_step(
nnow: int,
nnew: int,
):
+ log.info(
+ f"running predictor step: dtime = {dtime}, dyn_timestep = {idyn_timestep}, init = {l_init}, recompute = {l_recompute} "
+ )
if l_init or l_recompute:
if self.config.itime_scheme == 4 and not l_init:
lvn_only = True # Recompute only vn tendency
@@ -574,7 +557,7 @@ def run_predictor_step(
z_vt_ie=z_fields.z_vt_ie,
dtime=dtime,
ntnd=self.ntl1,
- cell_areas=self.cell_areas,
+ cell_areas=self.cell_params.area,
)
p_dthalf = 0.5 * dtime
@@ -634,10 +617,11 @@ def run_predictor_step(
end_cell_local_minus1 = self.grid.get_end_index(
CellDim, HorizontalMarkerIndex.local(CellDim) - 1
)
-
+ end_cell_halo = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.halo(CellDim))
start_cell_nudging = self.grid.get_start_index(
CellDim, HorizontalMarkerIndex.nudging(CellDim)
)
+
end_cell_local = self.grid.get_end_index(CellDim, HorizontalMarkerIndex.local(CellDim))
# Precompute Rayleigh damping factor
@@ -691,7 +675,6 @@ def run_predictor_step(
vertical_end=self.grid.num_levels + 1,
offset_provider={"Koff": KDim},
)
-
if self.vertical_params.nflatlev == 1:
# Perturbation Exner pressure on top half level
raise NotImplementedError("nflatlev=1 not implemented")
@@ -1025,7 +1008,7 @@ def run_predictor_step(
vertical_end=self.grid.num_levels,
offset_provider={},
)
-
+ log.debug("exchanging prognostic field 'vn' and local field 'z_rho_e'")
self._exchange.exchange_and_wait(EdgeDim, prognostic_state[nnew].vn, z_fields.z_rho_e)
mo_solve_nonhydro_stencil_30.with_backend(backend)(
@@ -1341,9 +1324,10 @@ def run_predictor_step(
vertical_end=self.grid.num_levels,
offset_provider={"Koff": KDim},
)
-
+ log.debug("exchanging prognostic field 'w' and local field 'z_dwdz_dd'")
self._exchange.exchange_and_wait(CellDim, prognostic_state[nnew].w, z_fields.z_dwdz_dd)
else:
+ log.debug("exchanging prognostic field 'w'")
self._exchange.exchange_and_wait(CellDim, prognostic_state[nnew].w)
def run_corrector_step(
@@ -1352,17 +1336,33 @@ def run_corrector_step(
prognostic_state: list[PrognosticState],
z_fields: ZFields,
nh_constants: NHConstants,
+ divdamp_fac_o2: float,
prep_adv: PrepAdvection,
dtime: float,
nnew: int,
nnow: int,
lclean_mflx: bool,
- bdy_divdamp: Field[[KDim], float],
lprep_adv: bool,
):
# Inverse value of ndyn_substeps for tracer advection precomputations
r_nsubsteps = 1.0 / self.config.ndyn_substeps_var
+ # scaling factor for second-order divergence damping: divdamp_fac_o2*delta_x**2
+ # delta_x**2 is approximated by the mean cell area
+ # Coefficient for reduced fourth-order divergence d
+ scal_divdamp_o2 = divdamp_fac_o2 * self.cell_params.mean_cell_area
+
+ _calculate_divdamp_fields.with_backend(run_gtfn)(
+ self.enh_divdamp_fac,
+ int32(self.config.divdamp_order),
+ self.cell_params.mean_cell_area,
+ divdamp_fac_o2,
+ self.config.nudge_max_coeff,
+ constants.dbl_eps,
+ out=(self.scal_divdamp, self._bdy_divdamp),
+ offset_provider={},
+ )
+
start_cell_lb_plus2 = self.grid.get_start_index(
CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2
)
@@ -1407,7 +1407,7 @@ def run_corrector_step(
z_vt_ie=z_fields.z_vt_ie,
dtime=dtime,
ntnd=self.ntl2,
- cell_areas=self.cell_areas,
+ cell_areas=self.cell_params.area,
)
nvar = nnew
@@ -1446,7 +1446,6 @@ def run_corrector_step(
vertical_end=self.grid.num_levels,
offset_provider={"Koff": KDim},
)
-
if self.config.l_open_ubc and not self.l_vert_nested:
raise NotImplementedError("l_open_ubc not implemented")
@@ -1503,11 +1502,11 @@ def run_corrector_step(
)
if self.config.lhdiff_rcf:
- if self.config.divdamp_order == 24 and nh_constants.scal_divdamp_o2 > 1.0e-6:
+ if self.config.divdamp_order == 24 and scal_divdamp_o2 > 1.0e-6:
mo_solve_nonhydro_stencil_26.with_backend(backend)(
z_graddiv_vn=z_fields.z_graddiv_vn,
vn=prognostic_state[nnew].vn,
- scal_divdamp_o2=nh_constants.scal_divdamp_o2,
+ scal_divdamp_o2=scal_divdamp_o2,
horizontal_start=start_edge_nudging_plus1,
horizontal_end=end_edge_local,
vertical_start=0,
@@ -1523,7 +1522,7 @@ def run_corrector_step(
if self.grid.limited_area:
mo_solve_nonhydro_stencil_27.with_backend(backend)(
scal_divdamp=self.scal_divdamp,
- bdy_divdamp=bdy_divdamp,
+ bdy_divdamp=self._bdy_divdamp,
nudgecoeff_e=self.interpolation_state.nudgecoeff_e,
z_graddiv2_vn=self.z_graddiv2_vn,
vn=prognostic_state[nnew].vn,
@@ -1557,7 +1556,7 @@ def run_corrector_step(
vertical_end=self.grid.num_levels,
offset_provider={},
)
-
+ log.debug("exchanging prognostic field 'vn'")
self._exchange.exchange_and_wait(EdgeDim, (prognostic_state[nnew].vn))
mo_solve_nonhydro_stencil_31.with_backend(backend)(
e_flx_avg=self.interpolation_state.e_flx_avg,
@@ -1867,7 +1866,7 @@ def run_corrector_step(
vertical_end=self.grid.num_levels,
offset_provider={},
)
-
+ log.debug("exchange prognostic fields 'rho' , 'exner', 'w'")
self._exchange.exchange_and_wait(
CellDim,
prognostic_state[nnew].rho,
diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/interpolation_state.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/interpolation_state.py
index 2b504f3ea2..440c4a4b9c 100644
--- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/interpolation_state.py
+++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/interpolation_state.py
@@ -13,13 +13,10 @@
from dataclasses import dataclass
-import numpy as np
from gt4py.next.common import Field
-from gt4py.next.iterator.embedded import np_as_located_field
from icon4py.model.common.dimension import (
C2E2CODim,
- CECDim,
CEDim,
CellDim,
E2C2EDim,
@@ -64,17 +61,3 @@ class InterpolationState:
pos_on_tplane_e_1: Field[[ECDim], float]
pos_on_tplane_e_2: Field[[ECDim], float]
e_flx_avg: Field[[EdgeDim, E2C2EODim], float]
-
- @property
- def geofac_n2s_c(self) -> Field[[CellDim], float]:
- return np_as_located_field(CellDim)(np.asarray(self.geofac_n2s)[:, 0])
-
- @property
- def geofac_n2s_nbh(self) -> Field[[CECDim], float]:
- geofac_nbh_ar = np.asarray(self.geofac_n2s)[:, 1:]
- old_shape = geofac_nbh_ar.shape
- return np_as_located_field(CECDim)(
- geofac_nbh_ar.reshape(
- old_shape[0] * old_shape[1],
- )
- )
diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/nh_constants.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/nh_constants.py
index e3cd107be0..9caae3a93f 100644
--- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/nh_constants.py
+++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/nh_constants.py
@@ -20,5 +20,3 @@ class NHConstants:
wgt_nnew_rth: float
wgt_nnow_vel: float
wgt_nnew_vel: float
- scal_divdamp: float
- scal_divdamp_o2: float
diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py
index 62bafcd15e..cb335ba2ae 100644
--- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py
+++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/utils.py
@@ -10,7 +10,6 @@
# distribution for a copy of the license or check .
#
# SPDX-License-Identifier: GPL-3.0-or-later
-from typing import Tuple
import numpy as np
from gt4py.next import as_field
@@ -21,10 +20,9 @@
broadcast,
int32,
maximum,
- minimum,
)
-from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, Koff, VertexDim
+from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim
def indices_field(dim: Dimension, grid, is_halfdim, dtype=int):
@@ -112,35 +110,6 @@ def set_zero_c_k(
)
-@field_operator
-def _setup_smag_limit(diff_multfac_vn: Field[[KDim], float]) -> Field[[KDim], float]:
- return 0.125 - 4.0 * diff_multfac_vn
-
-
-@field_operator
-def _setup_initial_diff_multfac_vn(k4: float, hdiff_efdt_ratio: float) -> Field[[KDim], float]:
- return broadcast(k4 / 3.0 * hdiff_efdt_ratio, (KDim,))
-
-
-@field_operator
-def _setup_fields_for_initial_step(
- k4: float, hdiff_efdt_ratio: float
-) -> Tuple[Field[[KDim], float], Field[[KDim], float]]:
- diff_multfac_vn = _setup_initial_diff_multfac_vn(k4, hdiff_efdt_ratio)
- smag_limit = _setup_smag_limit(diff_multfac_vn)
- return diff_multfac_vn, smag_limit
-
-
-@program
-def setup_fields_for_initial_step(
- k4: float,
- hdiff_efdt_ratio: float,
- diff_multfac_vn: Field[[KDim], float],
- smag_limit: Field[[KDim], float],
-):
- _setup_fields_for_initial_step(k4, hdiff_efdt_ratio, out=(diff_multfac_vn, smag_limit))
-
-
@field_operator
def _calculate_bdy_divdamp(
scal_divdamp: Field[[KDim], float], nudge_max_coeff: float, dbl_eps: float
@@ -148,34 +117,35 @@ def _calculate_bdy_divdamp(
return 0.75 / (nudge_max_coeff + dbl_eps) * abs(scal_divdamp)
-# TODO (magdalena) copied from diffusion/diffusion_utils/_en_smag_fac_for_zero_nshift makes sense?
@field_operator
-def _en_smag_fac_for_zero_nshift(
- vect_a: Field[[KDim], float],
- hdiff_smag_fac: float,
- hdiff_smag_fac2: float,
- hdiff_smag_fac3: float,
- hdiff_smag_fac4: float,
- hdiff_smag_z: float,
- hdiff_smag_z2: float,
- hdiff_smag_z3: float,
- hdiff_smag_z4: float,
+def _calculate_scal_divdamp(
+ enh_divdamp_fac: Field[[KDim], float],
+ divdamp_order: int32,
+ mean_cell_area: float,
+ divdamp_fac_o2: float,
) -> Field[[KDim], float]:
- dz21 = hdiff_smag_z2 - hdiff_smag_z
- alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21
- df32 = hdiff_smag_fac3 - hdiff_smag_fac2
- df42 = hdiff_smag_fac4 - hdiff_smag_fac2
- dz32 = hdiff_smag_z3 - hdiff_smag_z2
- dz42 = hdiff_smag_z4 - hdiff_smag_z2
+ enh_divdamp_fac = (
+ maximum(0.0, enh_divdamp_fac - 0.25 * divdamp_fac_o2)
+ if divdamp_order == 24
+ else enh_divdamp_fac
+ )
+ return -enh_divdamp_fac * mean_cell_area**2
- bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32))
- aqdr = df32 / dz32 - bqdr * dz32
- zf = 0.5 * (vect_a + vect_a(Koff[1]))
- dzlin = minimum(dz21, maximum(0.0, zf - hdiff_smag_z))
- dzqdr = minimum(dz42, maximum(0.0, zf - hdiff_smag_z2))
- enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr)
- return enh_smag_fac
+@field_operator
+def _calculate_divdamp_fields(
+ enh_divdamp_fac: Field[[KDim], float],
+ divdamp_order: int32,
+ mean_cell_area: float,
+ divdamp_fac_o2: float,
+ nudge_max_coeff: float,
+ dbl_eps: float,
+) -> tuple[Field[[KDim], float], Field[[KDim], float]]:
+ scal_divdamp = _calculate_scal_divdamp(
+ enh_divdamp_fac, divdamp_order, mean_cell_area, divdamp_fac_o2
+ )
+ bdy_divdamp = _calculate_bdy_divdamp(scal_divdamp, nudge_max_coeff, dbl_eps)
+ return (scal_divdamp, bdy_divdamp)
@field_operator
@@ -188,19 +158,3 @@ def compute_z_raylfac(
rayleigh_w: Field[[KDim], float], dtime: float, z_raylfac: Field[[KDim], float]
):
_compute_z_raylfac(rayleigh_w, dtime, out=z_raylfac)
-
-
-@field_operator
-def _scal_divdamp_calcs(
- enh_divdamp_fac: Field[[KDim], float], mean_cell_area: float
-) -> Field[[KDim], float]:
- return -enh_divdamp_fac * mean_cell_area**2
-
-
-@program
-def scal_divdamp_calcs(
- enh_divdamp_fac: Field[[KDim], float],
- out: Field[[KDim], float],
- mean_cell_area: float,
-):
- _scal_divdamp_calcs(enh_divdamp_fac, mean_cell_area, out=out)
diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py
index e1a39ed368..e6340459d5 100644
--- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py
+++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/velocity/velocity_advection.py
@@ -395,7 +395,6 @@ def run_predictor_step(
},
)
- # This behaviour needs to change for multiple blocks
self.levelmask = self.levmask
mo_velocity_advection_stencil_19.with_backend(backend)(
diff --git a/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/test_parallel_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/test_parallel_solve_nonhydro.py
index b02cdeae8f..be1704e21a 100644
--- a/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/test_parallel_solve_nonhydro.py
+++ b/model/atmosphere/dycore/tests/dycore_tests/mpi_tests/test_parallel_solve_nonhydro.py
@@ -23,7 +23,7 @@
from icon4py.model.atmosphere.dycore.state_utils.diagnostic_state import DiagnosticStateNonHydro
from icon4py.model.atmosphere.dycore.state_utils.nh_constants import NHConstants
from icon4py.model.atmosphere.dycore.state_utils.prep_adv_state import PrepAdvection
-from icon4py.model.atmosphere.dycore.state_utils.utils import _allocate, zero_field
+from icon4py.model.atmosphere.dycore.state_utils.utils import _allocate
from icon4py.model.atmosphere.dycore.state_utils.z_fields import ZFields
from icon4py.model.common.decomposition import definitions
from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim
@@ -33,7 +33,7 @@
from icon4py.model.common.test_utils.datatest_fixtures import ( # noqa : F401 fixture
decomposition_info,
)
-from icon4py.model.common.test_utils.helpers import dallclose, random_field
+from icon4py.model.common.test_utils.helpers import dallclose
from icon4py.model.common.test_utils.parallel_helpers import ( # noqa : F401 fixture
check_comm_size,
processor_props,
@@ -107,10 +107,6 @@ def test_run_solve_nonhydro_single_step(
vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic()
)
- enh_smag_fac = zero_field(icon_grid, KDim)
- a_vec = random_field(icon_grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
- fac = (0.67, 0.5, 1.3, 0.8)
- z = (0.1, 0.2, 0.3, 0.4)
nnow = 0
nnew = 1
recompute = sp_v.get_metadata("recompute").get("recompute")
@@ -178,8 +174,6 @@ def test_run_solve_nonhydro_single_step(
wgt_nnew_rth=sp.wgt_nnew_rth(),
wgt_nnow_vel=sp.wgt_nnow_vel(),
wgt_nnew_vel=sp.wgt_nnew_vel(),
- scal_divdamp=sp.scal_divdamp(),
- scal_divdamp_o2=sp.scal_divdamp_o2(),
)
interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro()
@@ -198,12 +192,8 @@ def test_run_solve_nonhydro_single_step(
interpolation_state=interpolation_state,
vertical_params=vertical_params,
edge_geometry=edge_geometry,
- cell_areas=cell_geometry.area,
+ cell_geometry=cell_geometry,
owner_mask=grid_savepoint.c_owner_mask(),
- a_vec=a_vec,
- enh_smag_fac=enh_smag_fac,
- fac=fac,
- z=z,
)
prognostic_state_ls = [prognostic_state_nnow, prognostic_state_nnew]
@@ -217,7 +207,7 @@ def test_run_solve_nonhydro_single_step(
prep_adv=prep_adv,
z_fields=z_fields,
nh_constants=nh_constants,
- bdy_divdamp=sp.bdy_divdamp(),
+ divdamp_fac_o2=0.032,
dtime=dtime,
idyn_timestep=dyn_timestep,
l_recompute=recompute,
diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py
index def23664cd..7d56058655 100644
--- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py
+++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py
@@ -10,8 +10,10 @@
# distribution for a copy of the license or check .
#
# SPDX-License-Identifier: GPL-3.0-or-later
+import logging
import pytest
+from gt4py.next.program_processors.runners.gtfn import run_gtfn
from icon4py.model.atmosphere.dycore.nh_solve.solve_nonhydro import (
NonHydrostaticConfig,
@@ -21,42 +23,63 @@
from icon4py.model.atmosphere.dycore.state_utils.diagnostic_state import DiagnosticStateNonHydro
from icon4py.model.atmosphere.dycore.state_utils.nh_constants import NHConstants
from icon4py.model.atmosphere.dycore.state_utils.prep_adv_state import PrepAdvection
-from icon4py.model.atmosphere.dycore.state_utils.utils import _allocate
+from icon4py.model.atmosphere.dycore.state_utils.utils import (
+ _allocate,
+ _calculate_bdy_divdamp,
+ _calculate_scal_divdamp,
+)
from icon4py.model.atmosphere.dycore.state_utils.z_fields import ZFields
+from icon4py.model.common import constants
from icon4py.model.common.dimension import CellDim, EdgeDim, KDim
from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalMarkerIndex
from icon4py.model.common.grid.vertical import VerticalModelParams
+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.test_utils.helpers import dallclose, random_field, zero_field
+from icon4py.model.common.test_utils.helpers import dallclose
-@pytest.mark.datatest
-def test_nonhydro_params():
- config = NonHydrostaticConfig()
- nonhydro_params = NonHydrostaticParams(config)
+backend = run_gtfn
- assert nonhydro_params.df32 == pytest.approx(
- config.divdamp_fac3 - config.divdamp_fac2, abs=1e-12
- )
- assert nonhydro_params.dz32 == pytest.approx(config.divdamp_z3 - config.divdamp_z2, abs=1e-12)
- assert nonhydro_params.df42 == pytest.approx(
- config.divdamp_fac4 - config.divdamp_fac2, abs=1e-12
- )
- assert nonhydro_params.dz42 == pytest.approx(config.divdamp_z4 - config.divdamp_z2, abs=1e-12)
- assert nonhydro_params.bqdr == pytest.approx(
- (nonhydro_params.df42 * nonhydro_params.dz32 - nonhydro_params.df32 * nonhydro_params.dz42)
- / (
- nonhydro_params.dz32
- * nonhydro_params.dz42
- * (nonhydro_params.dz42 - nonhydro_params.dz32)
- ),
- abs=1e-12,
- )
- assert nonhydro_params.aqdr == pytest.approx(
- nonhydro_params.df32 / nonhydro_params.dz32 - nonhydro_params.bqdr * nonhydro_params.dz32,
- abs=1e-12,
- )
+@pytest.mark.datatest
+def test_validate_divdamp_fields_against_savepoint_values(
+ grid_savepoint,
+ savepoint_nonhydro_init,
+ icon_grid,
+):
+ config = NonHydrostaticConfig()
+ divdamp_fac_o2 = 0.032
+ mean_cell_area = grid_savepoint.mean_cell_area()
+ enh_divdamp_fac = _allocate(KDim, is_halfdim=False, dtype=float, grid=icon_grid)
+ scal_divdamp = _allocate(KDim, is_halfdim=False, dtype=float, grid=icon_grid)
+ bdy_divdamp = _allocate(KDim, is_halfdim=False, dtype=float, grid=icon_grid)
+ en_smag_fac_for_zero_nshift.with_backend(backend)(
+ grid_savepoint.vct_a(),
+ config.divdamp_fac,
+ config.divdamp_fac2,
+ config.divdamp_fac3,
+ config.divdamp_fac4,
+ config.divdamp_z,
+ config.divdamp_z2,
+ config.divdamp_z3,
+ config.divdamp_z4,
+ out=enh_divdamp_fac,
+ offset_provider={"Koff": KDim},
+ )
+ _calculate_scal_divdamp.with_backend(backend)(
+ enh_divdamp_fac=enh_divdamp_fac,
+ divdamp_order=config.divdamp_order,
+ mean_cell_area=mean_cell_area,
+ divdamp_fac_o2=divdamp_fac_o2,
+ out=scal_divdamp,
+ offset_provider={},
+ )
+ _calculate_bdy_divdamp.with_backend(backend)(
+ scal_divdamp, config.nudge_max_coeff, constants.dbl_eps, out=bdy_divdamp, offset_provider={}
+ )
+
+ assert dallclose(scal_divdamp.asnumpy(), savepoint_nonhydro_init.scal_divdamp().asnumpy())
+ assert dallclose(bdy_divdamp.asnumpy(), savepoint_nonhydro_init.bdy_divdamp().asnumpy())
@pytest.mark.datatest
@@ -78,27 +101,20 @@ def test_nonhydro_predictor_step(
metrics_savepoint,
interpolation_savepoint,
savepoint_nonhydro_exit,
+ caplog,
):
+ caplog.set_level(logging.DEBUG)
config = NonHydrostaticConfig()
sp = savepoint_nonhydro_init
sp_exit = savepoint_nonhydro_exit
nonhydro_params = NonHydrostaticParams(config)
- vertical_params = VerticalModelParams(
- vct_a=grid_savepoint.vct_a(),
- rayleigh_damping_height=damping_height,
- nflat_gradp=grid_savepoint.nflat_gradp(),
- nflatlev=grid_savepoint.nflatlev(),
- )
+ vertical_params = create_vertical_params(damping_height, grid_savepoint)
sp_v = savepoint_velocity_init
dtime = sp_v.get_metadata("dtime").get("dtime")
recompute = sp_v.get_metadata("recompute").get("recompute")
dyn_timestep = sp.get_metadata("dyn_timestep").get("dyn_timestep")
linit = sp_v.get_metadata("linit").get("linit")
- enh_smag_fac = zero_field(icon_grid, KDim)
- a_vec = random_field(icon_grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
- fac = (0.67, 0.5, 1.3, 0.8)
- z = (0.1, 0.2, 0.3, 0.4)
nnow = 0
nnew = 1
@@ -125,38 +141,7 @@ def test_nonhydro_predictor_step(
exner_incr=None, # sp.exner_incr(),
)
- prognostic_state_nnow = PrognosticState(
- w=sp.w_now(),
- vn=sp.vn_now(),
- theta_v=sp.theta_v_now(),
- rho=sp.rho_now(),
- exner=sp.exner_now(),
- )
-
- prognostic_state_nnew = PrognosticState(
- w=sp.w_new(),
- vn=sp.vn_new(),
- theta_v=sp.theta_v_new(),
- rho=sp.rho_new(),
- exner=sp.exner_new(),
- )
-
- z_fields = ZFields(
- z_gradh_exner=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_alpha=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_beta=_allocate(CellDim, KDim, grid=icon_grid),
- z_w_expl=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_exner_expl=_allocate(CellDim, KDim, grid=icon_grid),
- z_q=_allocate(CellDim, KDim, grid=icon_grid),
- z_contr_w_fl_l=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_rho_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_theta_v_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_graddiv_vn=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_rho_expl=_allocate(CellDim, KDim, grid=icon_grid),
- z_dwdz_dd=_allocate(CellDim, KDim, grid=icon_grid),
- z_kin_hor_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_vt_ie=_allocate(EdgeDim, KDim, grid=icon_grid),
- )
+ z_fields = allocate_z_fields(icon_grid)
interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro()
metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels)
@@ -174,15 +159,11 @@ def test_nonhydro_predictor_step(
interpolation_state=interpolation_state,
vertical_params=vertical_params,
edge_geometry=edge_geometry,
- cell_areas=cell_geometry.area,
+ cell_geometry=cell_geometry,
owner_mask=grid_savepoint.c_owner_mask(),
- a_vec=a_vec,
- enh_smag_fac=enh_smag_fac,
- fac=fac,
- z=z,
)
- prognostic_state_ls = [prognostic_state_nnow, prognostic_state_nnew]
+ prognostic_state_ls = create_prognostic_states(sp)
solve_nonhydro.set_timelevels(nnow, nnew)
solve_nonhydro.run_predictor_step(
diagnostic_state_nh=diagnostic_state_nh,
@@ -205,8 +186,6 @@ def test_nonhydro_predictor_step(
icon_result_w_concorr_c = sp_exit.w_concorr_c().asnumpy()
icon_result_mass_fl_e = sp_exit.mass_fl_e().asnumpy()
- # TODO: @abishekg7 remove bounds from asserts?
- # stencils 2, 3
cell_start_lb_plus2 = icon_grid.get_start_index(
CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2
)
@@ -221,6 +200,7 @@ def test_nonhydro_predictor_step(
EdgeDim, HorizontalMarkerIndex.nudging(EdgeDim) + 1
)
+ # stencils 2, 3
assert dallclose(
sp_exit.exner_pr().asnumpy()[cell_start_lb_plus2:, :],
diagnostic_state_nh.exner_pr.asnumpy()[cell_start_lb_plus2:, :],
@@ -236,15 +216,16 @@ def test_nonhydro_predictor_step(
sp_exit.z_exner_ic().asnumpy()[cell_start_lb_plus2:, nlev - 1],
solve_nonhydro.z_exner_ic.asnumpy()[cell_start_lb_plus2:, nlev - 1],
)
+ nflatlev = vertical_params.nflatlev
assert dallclose(
- sp_exit.z_exner_ic().asnumpy()[cell_start_lb_plus2:, 4 : nlev - 1],
- solve_nonhydro.z_exner_ic.asnumpy()[cell_start_lb_plus2:, 4 : nlev - 1],
+ sp_exit.z_exner_ic().asnumpy()[cell_start_lb_plus2:, nflatlev : nlev - 1],
+ solve_nonhydro.z_exner_ic.asnumpy()[cell_start_lb_plus2:, nflatlev : nlev - 1],
rtol=1.0e-9,
)
# stencil 6
assert dallclose(
- sp_exit.z_dexner_dz_c(1).asnumpy()[cell_start_lb_plus2:, :],
- solve_nonhydro.z_dexner_dz_c_1.asnumpy()[cell_start_lb_plus2:, :],
+ sp_exit.z_dexner_dz_c(1).asnumpy()[cell_start_lb_plus2:, nflatlev:],
+ solve_nonhydro.z_dexner_dz_c_1.asnumpy()[cell_start_lb_plus2:, nflatlev:],
atol=5e-18,
)
@@ -279,9 +260,10 @@ def test_nonhydro_predictor_step(
)
# stencils 12
+ nflat_gradp = vertical_params.nflat_gradp
assert dallclose(
- sp_exit.z_dexner_dz_c(2).asnumpy()[cell_start_lb_plus2:, :],
- solve_nonhydro.z_dexner_dz_c_2.asnumpy()[cell_start_lb_plus2:, :],
+ sp_exit.z_dexner_dz_c(2).asnumpy()[cell_start_lb_plus2:, nflat_gradp:],
+ solve_nonhydro.z_dexner_dz_c_2.asnumpy()[cell_start_lb_plus2:, nflat_gradp:],
atol=1e-22,
)
@@ -329,6 +311,7 @@ def test_nonhydro_predictor_step(
solve_nonhydro.z_hydro_corr.asnumpy()[edge_start_nuding_plus1:, nlev - 1],
atol=1e-20,
)
+ prognostic_state_nnew = prognostic_state_ls[1]
# stencils 24
assert dallclose(
icon_result_vn_new[edge_start_nuding_plus1:, :],
@@ -393,7 +376,6 @@ def test_nonhydro_predictor_step(
z_fields.z_kin_hor_e.asnumpy()[edge_start_lb_plus4:, :],
atol=10e-13,
)
-
# stencil 35
assert dallclose(
sp_exit.z_w_concorr_me().asnumpy()[edge_start_lb_plus4:, vertical_params.nflatlev :],
@@ -403,7 +385,7 @@ def test_nonhydro_predictor_step(
# stencils 39,40
assert dallclose(
- icon_result_w_concorr_c.asnumpy()[cell_start_lb_plus2:, :],
+ icon_result_w_concorr_c[cell_start_lb_plus2:, :],
diagnostic_state_nh.w_concorr_c.asnumpy()[cell_start_lb_plus2:, :],
atol=1e-15,
)
@@ -460,7 +442,7 @@ def test_nonhydro_predictor_step(
z_fields.z_rho_expl.asnumpy()[cell_start_nudging:, :],
atol=2e-15,
)
- # stencil 48, 49 # TODO (magdalena) there is a problem at last 2 klevel=nlev-1,nlev-2
+ # stencil 48, 49
assert dallclose(
sp_exit.z_exner_expl().asnumpy()[cell_start_nudging:, :],
z_fields.z_exner_expl.asnumpy()[cell_start_nudging:, :],
@@ -473,10 +455,62 @@ def test_nonhydro_predictor_step(
# not tested
assert dallclose(icon_result_exner_new, prognostic_state_nnew.exner.asnumpy())
-
assert dallclose(icon_result_theta_v_new, prognostic_state_nnew.theta_v.asnumpy())
+def construct_diagnostics(sp, sp_v):
+ return DiagnosticStateNonHydro(
+ theta_v_ic=sp.theta_v_ic(),
+ exner_pr=sp.exner_pr(),
+ rho_ic=sp.rho_ic(),
+ ddt_exner_phy=sp.ddt_exner_phy(),
+ grf_tend_rho=sp.grf_tend_rho(),
+ grf_tend_thv=sp.grf_tend_thv(),
+ grf_tend_w=sp.grf_tend_w(),
+ mass_fl_e=sp.mass_fl_e(),
+ ddt_vn_phy=sp.ddt_vn_phy(),
+ grf_tend_vn=sp.grf_tend_vn(),
+ ddt_vn_apc_ntl1=sp_v.ddt_vn_apc_pc(1),
+ ddt_vn_apc_ntl2=sp_v.ddt_vn_apc_pc(2),
+ ddt_w_adv_ntl1=sp_v.ddt_w_adv_pc(1),
+ ddt_w_adv_ntl2=sp_v.ddt_w_adv_pc(2),
+ vt=sp_v.vt(),
+ vn_ie=sp_v.vn_ie(),
+ w_concorr_c=sp_v.w_concorr_c(),
+ rho_incr=None, # sp.rho_incr(),
+ vn_incr=None, # sp.vn_incr(),
+ exner_incr=None, # sp.exner_incr(),
+ )
+
+
+def allocate_z_fields(icon_grid):
+ return ZFields(
+ z_gradh_exner=_allocate(EdgeDim, KDim, grid=icon_grid),
+ z_alpha=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
+ z_beta=_allocate(CellDim, KDim, grid=icon_grid),
+ z_w_expl=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
+ z_exner_expl=_allocate(CellDim, KDim, grid=icon_grid),
+ z_q=_allocate(CellDim, KDim, grid=icon_grid),
+ z_contr_w_fl_l=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
+ z_rho_e=_allocate(EdgeDim, KDim, grid=icon_grid),
+ z_theta_v_e=_allocate(EdgeDim, KDim, grid=icon_grid),
+ z_graddiv_vn=_allocate(EdgeDim, KDim, grid=icon_grid),
+ z_rho_expl=_allocate(CellDim, KDim, grid=icon_grid),
+ z_dwdz_dd=_allocate(CellDim, KDim, grid=icon_grid),
+ z_kin_hor_e=_allocate(EdgeDim, KDim, grid=icon_grid),
+ z_vt_ie=_allocate(EdgeDim, KDim, grid=icon_grid),
+ )
+
+
+def create_vertical_params(damping_height, grid_savepoint):
+ return VerticalModelParams(
+ vct_a=grid_savepoint.vct_a(),
+ rayleigh_damping_height=damping_height,
+ nflat_gradp=grid_savepoint.nflat_gradp(),
+ nflatlev=grid_savepoint.nflatlev(),
+ )
+
+
@pytest.mark.datatest
@pytest.mark.parametrize(
"istep_init, istep_exit, step_date_init, step_date_exit",
@@ -495,7 +529,9 @@ def test_nonhydro_corrector_step(
metrics_savepoint,
interpolation_savepoint,
savepoint_nonhydro_exit,
+ caplog,
):
+ caplog.set_level(logging.DEBUG)
config = NonHydrostaticConfig()
sp = savepoint_nonhydro_init
nonhydro_params = NonHydrostaticParams(config)
@@ -507,53 +543,16 @@ def test_nonhydro_corrector_step(
)
sp_v = savepoint_velocity_init
dtime = sp_v.get_metadata("dtime").get("dtime")
+ clean_mflx = sp_v.get_metadata("clean_mflx").get("clean_mflx")
+ lprep_adv = sp_v.get_metadata("prep_adv").get("prep_adv")
prep_adv = PrepAdvection(
vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic()
)
- enh_smag_fac = zero_field(icon_grid, KDim)
- a_vec = random_field(icon_grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
- fac = (0.67, 0.5, 1.3, 0.8)
- z = (0.1, 0.2, 0.3, 0.4)
-
- diagnostic_state_nh = DiagnosticStateNonHydro(
- theta_v_ic=sp.theta_v_ic(),
- exner_pr=sp.exner_pr(),
- rho_ic=sp.rho_ic(),
- ddt_exner_phy=sp.ddt_exner_phy(),
- grf_tend_rho=sp.grf_tend_rho(),
- grf_tend_thv=sp.grf_tend_thv(),
- grf_tend_w=sp.grf_tend_w(),
- mass_fl_e=sp.mass_fl_e(),
- ddt_vn_phy=sp.ddt_vn_phy(),
- grf_tend_vn=sp.grf_tend_vn(),
- ddt_vn_apc_ntl1=sp_v.ddt_vn_apc_pc(1),
- ddt_vn_apc_ntl2=sp_v.ddt_vn_apc_pc(2),
- ddt_w_adv_ntl1=sp_v.ddt_w_adv_pc(1),
- ddt_w_adv_ntl2=sp_v.ddt_w_adv_pc(2),
- vt=sp_v.vt(), # sp_v.vt(), #TODO: @abishekg7 change back to sp_v
- vn_ie=sp_v.vn_ie(),
- w_concorr_c=sp_v.w_concorr_c(),
- rho_incr=None, # sp.rho_incr(),
- vn_incr=None, # sp.vn_incr(),
- exner_incr=None, # sp.exner_incr(),
- )
-
- prognostic_state_nnow = PrognosticState(
- w=sp.w_now(),
- vn=sp.vn_now(),
- theta_v=sp.theta_v_now(),
- rho=sp.rho_now(),
- exner=sp.exner_now(),
- )
+ nnow = 0 # TODO: @abishekg7 read from serialized data?
+ nnew = 1
- prognostic_state_nnew = PrognosticState(
- w=sp.w_new(),
- vn=sp.vn_new(),
- theta_v=sp.theta_v_new(),
- rho=sp.rho_new(),
- exner=sp.exner_new(),
- )
+ diagnostic_state_nh = construct_diagnostics(sp, sp_v)
z_fields = ZFields(
z_gradh_exner=sp.z_gradh_exner(),
@@ -572,14 +571,8 @@ def test_nonhydro_corrector_step(
z_vt_ie=sp_v.z_vt_ie(),
)
- nh_constants = NHConstants(
- wgt_nnow_rth=sp.wgt_nnow_rth(),
- wgt_nnew_rth=sp.wgt_nnew_rth(),
- wgt_nnow_vel=sp.wgt_nnow_vel(),
- wgt_nnew_vel=sp.wgt_nnew_vel(),
- scal_divdamp=sp.scal_divdamp(),
- scal_divdamp_o2=sp.scal_divdamp_o2(),
- )
+ nh_constants = create_nh_constants(sp)
+ divdamp_fac_o2 = sp.divdamp_fac_o2()
interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro()
metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels)
@@ -596,32 +589,24 @@ def test_nonhydro_corrector_step(
interpolation_state=interpolation_state,
vertical_params=vertical_params,
edge_geometry=edge_geometry,
- cell_areas=cell_geometry.area,
+ cell_geometry=cell_geometry,
owner_mask=grid_savepoint.c_owner_mask(),
- a_vec=a_vec,
- enh_smag_fac=enh_smag_fac,
- fac=fac,
- z=z,
)
- nnow = 0
- nnew = 1
- lprep_adv = sp_v.get_metadata("prep_adv").get("prep_adv")
- clean_mflx = sp_v.get_metadata("clean_mflx").get("clean_mflx")
-
- prognostic_state_ls = [prognostic_state_nnow, prognostic_state_nnew]
+ prognostic_state_ls = create_prognostic_states(sp)
solve_nonhydro.set_timelevels(nnow, nnew)
+ solve_nonhydro._bdy_divdamp = sp.bdy_divdamp()
solve_nonhydro.run_corrector_step(
diagnostic_state_nh=diagnostic_state_nh,
prognostic_state=prognostic_state_ls,
z_fields=z_fields,
prep_adv=prep_adv,
+ divdamp_fac_o2=divdamp_fac_o2,
dtime=dtime,
nnew=nnew,
nnow=nnow,
lclean_mflx=clean_mflx,
nh_constants=nh_constants,
- bdy_divdamp=sp.bdy_divdamp(),
lprep_adv=lprep_adv,
)
@@ -709,18 +694,15 @@ def test_run_solve_nonhydro_single_step(
interpolation_savepoint,
savepoint_nonhydro_exit,
savepoint_nonhydro_step_exit,
+ caplog,
):
+ caplog.set_level(logging.DEBUG)
config = NonHydrostaticConfig()
sp = savepoint_nonhydro_init
sp_exit = savepoint_nonhydro_exit
sp_step_exit = savepoint_nonhydro_step_exit
nonhydro_params = NonHydrostaticParams(config)
- vertical_params = VerticalModelParams(
- vct_a=grid_savepoint.vct_a(),
- rayleigh_damping_height=damping_height,
- nflat_gradp=grid_savepoint.nflat_gradp(),
- nflatlev=grid_savepoint.nflatlev(),
- )
+ vertical_params = create_vertical_params(damping_height, grid_savepoint)
sp_v = savepoint_velocity_init
dtime = sp_v.get_metadata("dtime").get("dtime")
lprep_adv = sp_v.get_metadata("prep_adv").get("prep_adv")
@@ -729,80 +711,17 @@ def test_run_solve_nonhydro_single_step(
vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic()
)
- enh_smag_fac = zero_field(icon_grid, KDim)
- a_vec = random_field(icon_grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
- fac = (0.67, 0.5, 1.3, 0.8)
- z = (0.1, 0.2, 0.3, 0.4)
nnow = 0
nnew = 1
recompute = sp_v.get_metadata("recompute").get("recompute")
linit = sp_v.get_metadata("linit").get("linit")
dyn_timestep = sp_v.get_metadata("dyn_timestep").get("dyn_timestep")
- diagnostic_state_nh = DiagnosticStateNonHydro(
- theta_v_ic=sp.theta_v_ic(),
- exner_pr=sp.exner_pr(),
- rho_ic=sp.rho_ic(),
- ddt_exner_phy=sp.ddt_exner_phy(),
- grf_tend_rho=sp.grf_tend_rho(),
- grf_tend_thv=sp.grf_tend_thv(),
- grf_tend_w=sp.grf_tend_w(),
- mass_fl_e=sp.mass_fl_e(),
- ddt_vn_phy=sp.ddt_vn_phy(),
- grf_tend_vn=sp.grf_tend_vn(),
- ddt_vn_apc_ntl1=sp_v.ddt_vn_apc_pc(1),
- ddt_vn_apc_ntl2=sp_v.ddt_vn_apc_pc(2),
- ddt_w_adv_ntl1=sp_v.ddt_w_adv_pc(1),
- ddt_w_adv_ntl2=sp_v.ddt_w_adv_pc(2),
- vt=sp_v.vt(),
- vn_ie=sp_v.vn_ie(),
- w_concorr_c=sp_v.w_concorr_c(),
- rho_incr=None, # sp.rho_incr(),
- vn_incr=None, # sp.vn_incr(),
- exner_incr=None, # sp.exner_incr(),
- )
-
- prognostic_state_nnow = PrognosticState(
- w=sp.w_now(),
- vn=sp.vn_now(),
- theta_v=sp.theta_v_now(),
- rho=sp.rho_now(),
- exner=sp.exner_now(),
- )
+ diagnostic_state_nh = construct_diagnostics(sp, sp_v)
- prognostic_state_nnew = PrognosticState(
- w=sp.w_new(),
- vn=sp.vn_new(),
- theta_v=sp.theta_v_new(),
- rho=sp.rho_new(),
- exner=sp.exner_new(),
- )
-
- z_fields = ZFields(
- z_gradh_exner=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_alpha=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_beta=_allocate(CellDim, KDim, grid=icon_grid),
- z_w_expl=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_exner_expl=_allocate(CellDim, KDim, grid=icon_grid),
- z_q=_allocate(CellDim, KDim, grid=icon_grid),
- z_contr_w_fl_l=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_rho_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_theta_v_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_graddiv_vn=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_rho_expl=_allocate(CellDim, KDim, grid=icon_grid),
- z_dwdz_dd=_allocate(CellDim, KDim, grid=icon_grid),
- z_kin_hor_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_vt_ie=_allocate(EdgeDim, KDim, grid=icon_grid),
- )
+ z_fields = allocate_z_fields(icon_grid)
- nh_constants = NHConstants(
- wgt_nnow_rth=sp.wgt_nnow_rth(),
- wgt_nnew_rth=sp.wgt_nnew_rth(),
- wgt_nnow_vel=sp.wgt_nnow_vel(),
- wgt_nnew_vel=sp.wgt_nnew_vel(),
- scal_divdamp=sp.scal_divdamp(),
- scal_divdamp_o2=sp.scal_divdamp_o2(),
- )
+ nh_constants = create_nh_constants(sp)
interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro()
metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels)
@@ -819,23 +738,20 @@ def test_run_solve_nonhydro_single_step(
interpolation_state=interpolation_state,
vertical_params=vertical_params,
edge_geometry=edge_geometry,
- cell_areas=cell_geometry.area,
+ cell_geometry=cell_geometry,
owner_mask=grid_savepoint.c_owner_mask(),
- a_vec=a_vec,
- enh_smag_fac=enh_smag_fac,
- fac=fac,
- z=z,
)
- prognostic_state_ls = [prognostic_state_nnow, prognostic_state_nnew]
+ prognostic_state_ls = create_prognostic_states(sp)
+ initial_divdamp_fac = sp.divdamp_fac_o2()
solve_nonhydro.time_step(
diagnostic_state_nh=diagnostic_state_nh,
prognostic_state_ls=prognostic_state_ls,
prep_adv=prep_adv,
z_fields=z_fields,
nh_constants=nh_constants,
- bdy_divdamp=sp.bdy_divdamp(), # TODO (magdalena) local calculation in solve non-hydro based on nudge_coeff_e and scal_divdamp (also locally calculated)
+ divdamp_fac_o2=initial_divdamp_fac,
dtime=dtime,
idyn_timestep=dyn_timestep,
l_recompute=recompute,
@@ -845,7 +761,7 @@ def test_run_solve_nonhydro_single_step(
lclean_mflx=clean_mflx,
lprep_adv=lprep_adv,
)
-
+ prognostic_state_nnew = prognostic_state_ls[1]
assert dallclose(
sp_step_exit.theta_v_new().asnumpy(),
prognostic_state_nnew.theta_v.asnumpy(),
@@ -914,66 +830,19 @@ def test_run_solve_nonhydro_multi_step(
vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic()
)
- enh_smag_fac = zero_field(icon_grid, KDim)
- a_vec = random_field(icon_grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
- fac = (0.67, 0.5, 1.3, 0.8)
- z = (0.1, 0.2, 0.3, 0.4)
nnow = 0
nnew = 1
recompute = sp_v.get_metadata("recompute").get("recompute")
linit = sp_v.get_metadata("linit").get("linit")
dyn_timestep = sp_v.get_metadata("dyn_timestep").get("dyn_timestep")
- diagnostic_state_nh = DiagnosticStateNonHydro(
- theta_v_ic=sp.theta_v_ic(),
- exner_pr=sp.exner_pr(),
- rho_ic=sp.rho_ic(),
- ddt_exner_phy=sp.ddt_exner_phy(),
- grf_tend_rho=sp.grf_tend_rho(),
- grf_tend_thv=sp.grf_tend_thv(),
- grf_tend_w=sp.grf_tend_w(),
- mass_fl_e=sp.mass_fl_e(),
- ddt_vn_phy=sp.ddt_vn_phy(),
- grf_tend_vn=sp.grf_tend_vn(),
- ddt_vn_apc_ntl1=sp_v.ddt_vn_apc_pc(1),
- ddt_vn_apc_ntl2=sp_v.ddt_vn_apc_pc(2),
- ddt_w_adv_ntl1=sp_v.ddt_w_adv_pc(1),
- ddt_w_adv_ntl2=sp_v.ddt_w_adv_pc(2),
- vt=sp_v.vt(),
- vn_ie=sp_v.vn_ie(),
- w_concorr_c=sp_v.w_concorr_c(),
- rho_incr=None, # sp.rho_incr(),
- vn_incr=None, # sp.vn_incr(),
- exner_incr=None, # sp.exner_incr(),
- )
+ diagnostic_state_nh = construct_diagnostics(sp, sp_v)
prognostic_state_ls = create_prognostic_states(sp)
- z_fields = ZFields(
- z_gradh_exner=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_alpha=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_beta=_allocate(CellDim, KDim, grid=icon_grid),
- z_w_expl=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_exner_expl=_allocate(CellDim, KDim, grid=icon_grid),
- z_q=_allocate(CellDim, KDim, grid=icon_grid),
- z_contr_w_fl_l=_allocate(CellDim, KDim, is_halfdim=True, grid=icon_grid),
- z_rho_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_theta_v_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_graddiv_vn=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_rho_expl=_allocate(CellDim, KDim, grid=icon_grid),
- z_dwdz_dd=_allocate(CellDim, KDim, grid=icon_grid),
- z_kin_hor_e=_allocate(EdgeDim, KDim, grid=icon_grid),
- z_vt_ie=_allocate(EdgeDim, KDim, grid=icon_grid),
- )
+ z_fields = allocate_z_fields(icon_grid)
- nh_constants = NHConstants(
- wgt_nnow_rth=sp.wgt_nnow_rth(),
- wgt_nnew_rth=sp.wgt_nnew_rth(),
- wgt_nnow_vel=sp.wgt_nnow_vel(),
- wgt_nnew_vel=sp.wgt_nnew_vel(),
- scal_divdamp=sp.scal_divdamp(),
- scal_divdamp_o2=sp.scal_divdamp_o2(),
- )
+ nh_constants = create_nh_constants(sp)
interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro()
metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels)
@@ -990,12 +859,8 @@ def test_run_solve_nonhydro_multi_step(
interpolation_state=interpolation_state,
vertical_params=vertical_params,
edge_geometry=edge_geometry,
- cell_areas=cell_geometry.area,
+ cell_geometry=cell_geometry,
owner_mask=grid_savepoint.c_owner_mask(),
- a_vec=a_vec,
- enh_smag_fac=enh_smag_fac,
- fac=fac,
- z=z,
)
for i_substep in range(r_nsubsteps):
@@ -1005,7 +870,7 @@ def test_run_solve_nonhydro_multi_step(
prep_adv=prep_adv,
z_fields=z_fields,
nh_constants=nh_constants,
- bdy_divdamp=sp.bdy_divdamp(),
+ divdamp_fac_o2=sp.divdamp_fac_o2(),
dtime=dtime,
idyn_timestep=dyn_timestep,
l_recompute=recompute,
@@ -1092,6 +957,15 @@ def test_run_solve_nonhydro_multi_step(
)
+def create_nh_constants(sp):
+ return NHConstants(
+ wgt_nnow_rth=sp.wgt_nnow_rth(),
+ wgt_nnew_rth=sp.wgt_nnew_rth(),
+ wgt_nnow_vel=sp.wgt_nnow_vel(),
+ wgt_nnew_vel=sp.wgt_nnew_vel(),
+ )
+
+
def create_prognostic_states(sp):
prognostic_state_nnow = PrognosticState(
w=sp.w_now(),
diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_utils.py b/model/atmosphere/dycore/tests/dycore_tests/test_utils.py
new file mode 100644
index 0000000000..2b8fae5a3b
--- /dev/null
+++ b/model/atmosphere/dycore/tests/dycore_tests/test_utils.py
@@ -0,0 +1,120 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check .
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+from gt4py.next.ffront.fbuiltins import int32
+from gt4py.next.program_processors.runners.gtfn import run_gtfn
+
+from icon4py.model.atmosphere.dycore.state_utils.utils import (
+ _calculate_bdy_divdamp,
+ _calculate_divdamp_fields,
+ _calculate_scal_divdamp,
+)
+from icon4py.model.common import constants
+from icon4py.model.common.dimension import KDim
+from icon4py.model.common.grid.simple import SimpleGrid
+from icon4py.model.common.test_utils.helpers import dallclose, random_field, zero_field
+
+
+backend = run_gtfn
+
+
+def scal_divdamp_for_order_24_numpy(a: np.array, factor: float, mean_cell_area: float):
+ a = np.maximum(0.0, a - 0.25 * factor)
+ return -a * mean_cell_area**2
+
+
+def bdy_divdamp_numpy(coeff: float, field: np.array):
+ return 0.75 / (coeff + constants.dbl_eps) * np.abs(field)
+
+
+def test_caclulate_scal_divdamp_order_24():
+ divdamp_fac_o2 = 3.0
+ divdamp_order = 24
+ mean_cell_area = 1000.0
+ grid = SimpleGrid()
+ enh_divdamp_fac = random_field(grid, KDim)
+ out = random_field(grid, KDim)
+
+ _calculate_scal_divdamp.with_backend(backend)(
+ enh_divdamp_fac=enh_divdamp_fac,
+ divdamp_fac_o2=divdamp_fac_o2,
+ divdamp_order=divdamp_order,
+ mean_cell_area=mean_cell_area,
+ out=out,
+ offset_provider={},
+ )
+
+ ref = scal_divdamp_for_order_24_numpy(enh_divdamp_fac.asnumpy(), divdamp_fac_o2, mean_cell_area)
+ assert dallclose(ref, out.asnumpy())
+
+
+def test_calculate_scal_divdamp_any_order():
+ divdamp_fac_o2 = 4.2
+ divdamp_order = 3
+ mean_cell_area = 1000.0
+ grid = SimpleGrid()
+ enh_divdamp_fac = random_field(grid, KDim)
+ out = random_field(grid, KDim)
+
+ _calculate_scal_divdamp.with_backend(backend)(
+ enh_divdamp_fac=enh_divdamp_fac,
+ divdamp_fac_o2=divdamp_fac_o2,
+ divdamp_order=divdamp_order,
+ mean_cell_area=mean_cell_area,
+ out=out,
+ offset_provider={},
+ )
+ enhanced_factor = -enh_divdamp_fac.asnumpy() * mean_cell_area**2
+ assert dallclose(enhanced_factor, out.asnumpy())
+
+
+def test_calculate_bdy_divdamp():
+ grid = SimpleGrid()
+ scal_divdamp = random_field(grid, KDim)
+ out = zero_field(grid, KDim)
+ coeff = 0.3
+ _calculate_bdy_divdamp.with_backend(backend)(
+ scal_divdamp, coeff, constants.dbl_eps, out=out, offset_provider={}
+ )
+ assert dallclose(out.asnumpy(), bdy_divdamp_numpy(coeff, scal_divdamp.asnumpy()))
+
+
+def test_calculate_divdamp_fields():
+ grid = SimpleGrid()
+ divdamp_field = random_field(grid, KDim)
+ scal_divdamp = zero_field(grid, KDim)
+ boundary_divdamp = zero_field(grid, KDim)
+ divdamp_order = int32(24)
+ mean_cell_area = 1000.0
+ divdamp_fac_o2 = 0.7
+ nudge_max_coeff = 0.3
+
+ scaled_ref = scal_divdamp_for_order_24_numpy(
+ np.asarray(divdamp_field), divdamp_fac_o2, mean_cell_area
+ )
+
+ boundary_ref = bdy_divdamp_numpy(nudge_max_coeff, scaled_ref)
+
+ _calculate_divdamp_fields.with_backend(backend)(
+ divdamp_field,
+ divdamp_order,
+ mean_cell_area,
+ divdamp_fac_o2,
+ nudge_max_coeff,
+ constants.dbl_eps,
+ out=(scal_divdamp, boundary_divdamp),
+ offset_provider={},
+ )
+ dallclose(scal_divdamp.asnumpy(), scaled_ref)
+ dallclose(boundary_divdamp.asnumpy(), boundary_ref)
diff --git a/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py b/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py
index c25201fd45..f1b56b5c5d 100644
--- a/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py
+++ b/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py
@@ -252,7 +252,7 @@ def test_velocity_predictor_step(
atol=5.0e-16,
rtol=1.0e-10,
)
- # stencil 19:
+ # stencil 19
assert dallclose(
icon_result_ddt_vn_apc_pc,
diagnostic_state.ddt_vn_apc_pc[ntnd - 1].asnumpy(),
diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py
index 8bcad50cae..91efb0a86c 100644
--- a/model/common/src/icon4py/model/common/constants.py
+++ b/model/common/src/icon4py/model/common/constants.py
@@ -10,7 +10,7 @@
# distribution for a copy of the license or check .
#
# SPDX-License-Identifier: GPL-3.0-or-later
-
+import sys
from typing import Final
@@ -39,7 +39,7 @@
# Math constants
-dbl_eps = 0.01 # EPSILON(1._wp)
+dbl_eps = sys.float_info.epsilon # EPSILON(1._wp)
# Implementation constants
#: default physics to dynamics time step ratio
diff --git a/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py b/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py
index 6557b18e76..9176019a33 100644
--- a/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py
+++ b/model/common/src/icon4py/model/common/decomposition/mpi_decomposition.py
@@ -16,10 +16,9 @@
import functools
import logging
from dataclasses import dataclass
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, Sequence, Union
-import numpy as np
-from gt4py.next import Dimension
+from gt4py.next import Dimension, Field
from icon4py.model.common.decomposition.definitions import SingleNodeExchange
@@ -191,14 +190,14 @@ def _create_pattern(self, horizontal_dim: Dimension):
)
return pattern
- def exchange(self, dim: definitions.Dimension, *fields: tuple):
+ def exchange(self, dim: definitions.Dimension, *fields: Sequence[Field]):
assert dim in [CellDim, EdgeDim, VertexDim]
pattern = self._patterns[dim]
assert pattern is not None, f"pattern for {dim.value} not found"
domain_descriptor = self._domain_descriptors[dim]
assert domain_descriptor is not None, f"domain descriptor for {dim.value} not found"
applied_patterns = [
- pattern(unstructured.field_descriptor(domain_descriptor, np.asarray(f))) for f in fields
+ pattern(unstructured.field_descriptor(domain_descriptor, f.asnumpy())) for f in fields
]
handle = self._comm.exchange(applied_patterns)
log.info(f"exchange for {len(fields)} fields of dimension ='{dim.value}' initiated.")
diff --git a/model/common/src/icon4py/model/common/grid/horizontal.py b/model/common/src/icon4py/model/common/grid/horizontal.py
index 716b022c7f..6eb62ea959 100644
--- a/model/common/src/icon4py/model/common/grid/horizontal.py
+++ b/model/common/src/icon4py/model/common/grid/horizontal.py
@@ -281,6 +281,8 @@ def __init__(
@dataclass(frozen=True)
class CellParams:
area: Field[[CellDim], float]
+ mean_cell_area: float
+
"""
Area of a cell.
diff --git a/model/common/src/icon4py/model/common/math/__init__.py b/model/common/src/icon4py/model/common/math/__init__.py
new file mode 100644
index 0000000000..15dfdb0098
--- /dev/null
+++ b/model/common/src/icon4py/model/common/math/__init__.py
@@ -0,0 +1,12 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check .
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
diff --git a/model/common/src/icon4py/model/common/math/smagorinsky.py b/model/common/src/icon4py/model/common/math/smagorinsky.py
new file mode 100644
index 0000000000..6fca59d129
--- /dev/null
+++ b/model/common/src/icon4py/model/common/math/smagorinsky.py
@@ -0,0 +1,47 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check .
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from gt4py.next import Field, field_operator
+from gt4py.next.ffront.fbuiltins import broadcast, maximum, minimum
+
+from icon4py.model.common.dimension import KDim, Koff
+
+
+@field_operator
+def en_smag_fac_for_zero_nshift(
+ vect_a: Field[[KDim], float],
+ hdiff_smag_fac: float,
+ hdiff_smag_fac2: float,
+ hdiff_smag_fac3: float,
+ hdiff_smag_fac4: float,
+ hdiff_smag_z: float,
+ hdiff_smag_z2: float,
+ hdiff_smag_z3: float,
+ hdiff_smag_z4: float,
+) -> Field[[KDim], float]:
+ dz21 = hdiff_smag_z2 - hdiff_smag_z
+ alin = (hdiff_smag_fac2 - hdiff_smag_fac) / dz21
+ df32 = hdiff_smag_fac3 - hdiff_smag_fac2
+ df42 = hdiff_smag_fac4 - hdiff_smag_fac2
+ dz32 = hdiff_smag_z3 - hdiff_smag_z2
+ dz42 = hdiff_smag_z4 - hdiff_smag_z2
+
+ bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32))
+ aqdr = df32 / dz32 - bqdr * dz32
+ zf = 0.5 * (vect_a + vect_a(Koff[1]))
+ zero = broadcast(0.0, (KDim,))
+
+ dzlin = minimum(broadcast(dz21, (KDim,)), maximum(zero, zf - hdiff_smag_z))
+ dzqdr = minimum(broadcast(dz42, (KDim,)), maximum(zero, zf - hdiff_smag_z2))
+ enh_smag_fac = hdiff_smag_fac + (dzlin * alin) + dzqdr * (aqdr + dzqdr * bqdr)
+ return enh_smag_fac
diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py
index c2e2ea61b5..46709e51b5 100644
--- a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py
+++ b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py
@@ -24,7 +24,7 @@
# TODO: a run that contains all the fields needed for dycore, diffusion, interpolation fields needs to be consolidated
DATA_URIS = {
- 1: "https://polybox.ethz.ch/index.php/s/psBNNhng0h9KrB4/download",
+ 1: "https://polybox.ethz.ch/index.php/s/re46l1xnJZ4uCMx/download",
2: "https://polybox.ethz.ch/index.php/s/YyC5qDJWyC39y7u/download",
4: "https://polybox.ethz.ch/index.php/s/UIHOVJs6FVPpz9V/download",
}
diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py
index 99f2c1aee6..fc10086e38 100644
--- a/model/common/src/icon4py/model/common/test_utils/helpers.py
+++ b/model/common/src/icon4py/model/common/test_utils/helpers.py
@@ -20,6 +20,8 @@
from gt4py.next import common as gt_common
from gt4py.next.ffront.decorator import Program
+from ..grid.base import BaseGrid
+
try:
import pytest_benchmark
@@ -80,7 +82,7 @@ def random_field(
def zero_field(
- grid: SimpleGrid,
+ grid: BaseGrid,
*dims: gt_common.Dimension,
dtype=float,
extend: Optional[dict[gt_common.Dimension, int]] = None,
diff --git a/model/common/src/icon4py/model/common/test_utils/reference_funcs.py b/model/common/src/icon4py/model/common/test_utils/reference_funcs.py
new file mode 100644
index 0000000000..d2383203e9
--- /dev/null
+++ b/model/common/src/icon4py/model/common/test_utils/reference_funcs.py
@@ -0,0 +1,30 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check .
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+
+
+def enhanced_smagorinski_factor_numpy(factor_in, heigths_in, a_vec):
+ alin = (factor_in[1] - factor_in[0]) / (heigths_in[1] - heigths_in[0])
+ df32 = factor_in[2] - factor_in[1]
+ df42 = factor_in[3] - factor_in[1]
+ dz32 = heigths_in[2] - heigths_in[1]
+ dz42 = heigths_in[3] - heigths_in[1]
+ bqdr = (df42 * dz32 - df32 * dz42) / (dz32 * dz42 * (dz42 - dz32))
+ aqdr = df32 / dz32 - bqdr * dz32
+ zf = 0.5 * (a_vec[:-1] + a_vec[1:])
+ max0 = np.maximum(0.0, zf - heigths_in[0])
+ dzlin = np.minimum(heigths_in[1] - heigths_in[0], max0)
+ max1 = np.maximum(0.0, zf - heigths_in[1])
+ dzqdr = np.minimum(heigths_in[3] - heigths_in[1], max1)
+ return factor_in[0] + dzlin * alin + dzqdr * (aqdr + dzqdr * bqdr)
diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
index 19b4b2af1e..113aed4589 100644
--- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
+++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
@@ -14,7 +14,7 @@
import numpy as np
import serialbox as ser
-from gt4py.next import as_field
+from gt4py.next import Field, as_field
from gt4py.next.common import Dimension, DimensionKind
from gt4py.next.ffront.fbuiltins import int32
@@ -73,6 +73,14 @@ def _get_field(self, name, *dimensions, dtype=float):
self.log.debug(f"{name} {buffer.shape}")
return as_field(dimensions, buffer)
+ def _get_field_component(self, name: str, ntnd: int, dims: tuple[Dimension, Dimension]):
+ buffer = np.squeeze(self.serializer.read(name, self.savepoint).astype(float))[
+ :, :, ntnd - 1
+ ]
+ buffer = self._reduce_to_dim_size(buffer, dims)
+ self.log.debug(f"{name} {buffer.shape}")
+ return as_field(dims, buffer)
+
def _reduce_to_dim_size(self, buffer, dimensions):
buffer_size = (
self.sizes[d] if d.kind is DimensionKind.HORIZONTAL else s
@@ -155,6 +163,9 @@ def dual_normal_cell_y(self):
def cell_areas(self):
return self._get_field("cell_areas", CellDim)
+ def mean_cell_area(self):
+ return self.serializer.read("mean_cell_area", self.savepoint).astype(float)[0]
+
def edge_areas(self):
return self._get_field("edge_areas", EdgeDim)
@@ -379,7 +390,7 @@ def construct_edge_geometry(self) -> EdgeParams:
)
def construct_cell_geometry(self) -> CellParams:
- return CellParams(area=self.cell_areas())
+ return CellParams(area=self.cell_areas(), mean_cell_area=self.mean_cell_area())
class InterpolationSavepoint(IconSavepoint):
@@ -761,6 +772,9 @@ class IconNonHydroInitSavepoint(IconSavepoint):
def bdy_divdamp(self):
return self._get_field("bdy_divdamp", KDim)
+ def divdamp_fac_o2(self):
+ return self.serializer.read("divdamp_fac_o2", self.savepoint).astype(float)[0]
+
def ddt_exner_phy(self):
return self._get_field("ddt_exner_phy", CellDim, KDim)
@@ -798,12 +812,20 @@ def grf_tend_vn(self):
return self._get_field("grf_tend_vn", EdgeDim, KDim)
def ddt_vn_adv_ntl(self, ntl):
- buffer = np.squeeze(self.serializer.read("ddt_vn_adv_ntl", self.savepoint).astype(float))
- return as_field((EdgeDim, KDim), buffer[:, :, ntl - 1])
+ buffer = np.squeeze(self.serializer.read("ddt_vn_adv_ntl", self.savepoint).astype(float))[
+ :, :, ntl - 1
+ ]
+ dims = (EdgeDim, KDim)
+ buffer = self._reduce_to_dim_size(buffer, dims)
+ return as_field(dims, buffer)
def ddt_w_adv_ntl(self, ntl):
- buffer = np.squeeze(self.serializer.read("ddt_w_adv_ntl", self.savepoint).astype(float))
- return as_field((CellDim, KDim), buffer[:, :, ntl - 1])
+ buffer = np.squeeze(self.serializer.read("ddt_w_adv_ntl", self.savepoint).astype(float))[
+ :, :, ntl - 1
+ ]
+ dims = (CellDim, KDim)
+ buffer = self._reduce_to_dim_size(buffer, dims)
+ return as_field(dims, buffer)
def grf_tend_w(self):
return self._get_field("grf_tend_w", CellDim, KDim)
@@ -829,12 +851,12 @@ def exner_incr(self):
def vn_incr(self):
return self._get_field("vn_incr", EdgeDim, KDim)
- def scal_divdamp(self) -> float:
- return self.serializer.read("scal_divdamp", self.savepoint)[0]
-
def scal_divdamp_o2(self) -> float:
return self.serializer.read("scal_divdamp_o2", self.savepoint)[0]
+ def scal_divdamp(self) -> Field[[KDim], float]:
+ return self._get_field("scal_divdamp", KDim)
+
def theta_v_ic(self):
return self._get_field("theta_v_ic", CellDim, KDim)
@@ -907,20 +929,10 @@ def cfl_w_limit(self) -> float:
return self.serializer.read("cfl_w_limit", self.savepoint)[0]
def ddt_vn_apc_pc(self, ntnd):
- buffer = np.squeeze(self.serializer.read("ddt_vn_apc_pc", self.savepoint).astype(float))[
- :, :, ntnd - 1
- ]
- dims = (EdgeDim, KDim)
- buffer = self._reduce_to_dim_size(buffer, dims)
- return as_field(dims, buffer)
+ return self._get_field_component("ddt_vn_apc_pc", ntnd, (EdgeDim, KDim))
def ddt_w_adv_pc(self, ntnd):
- buffer = np.squeeze(self.serializer.read("ddt_w_adv_pc", self.savepoint).astype(float))[
- :, :, ntnd - 1
- ]
- dims = (CellDim, KDim)
- buffer = self._reduce_to_dim_size(buffer, dims)
- return as_field(dims, buffer)
+ return self._get_field_component("ddt_w_adv_pc", ntnd, (CellDim, KDim))
def scalfac_exdiff(self) -> float:
return self.serializer.read("scalfac_exdiff", self.savepoint)[0]
@@ -993,34 +1005,10 @@ def theta_v_now(self):
return self._get_field("x_theta_v_now", CellDim, KDim)
def ddt_vn_apc_pc(self, ntnd):
- buffer = np.squeeze(self.serializer.read("x_ddt_vn_apc_pc", self.savepoint).astype(float))[
- :, :, ntnd - 1
- ]
- dims = (EdgeDim, KDim)
- buffer = self._reduce_to_dim_size(buffer, dims)
- return as_field((EdgeDim, KDim), buffer)
-
- def ddt_vn_apc_pc_19(self, ntnd):
- buffer = np.squeeze(
- self.serializer.read("x_ddt_vn_apc_pc_19", self.savepoint).astype(float)
- )
- return as_field((EdgeDim, KDim), buffer[:, :, ntnd - 1])
+ return self._get_field_component("x_ddt_vn_apc_pc", ntnd, (EdgeDim, KDim))
def ddt_w_adv_pc(self, ntnd):
- buffer = np.squeeze(self.serializer.read("x_ddt_w_adv_pc", self.savepoint).astype(float))[
- :, :, ntnd - 1
- ]
- dims = (CellDim, KDim)
- buffer = self._reduce_to_dim_size(buffer, dims)
- return as_field(dims, buffer)
-
- def ddt_w_adv_pc_16(self, ntnd):
- buffer = np.squeeze(self.serializer.read("x_ddt_w_adv_pc_16", self.savepoint).astype(float))
- return as_field((CellDim, KDim), buffer[:, :, ntnd - 1])
-
- def ddt_w_adv_pc_17(self, ntnd):
- buffer = np.squeeze(self.serializer.read("x_ddt_w_adv_pc_17", self.savepoint).astype(float))
- return as_field((CellDim, KDim, buffer[:, :, ntnd - 1]))
+ return self._get_field_component("x_ddt_w_adv_pc", ntnd, (CellDim, KDim))
def scalfac_exdiff(self) -> float:
return self.serializer.read("scalfac_exdiff", self.savepoint)[0]
@@ -1107,18 +1095,10 @@ def w_new(self):
return self._get_field("x_w_new", CellDim, KDim)
def z_dexner_dz_c(self, ntnd):
- buffer = np.squeeze(self.serializer.read("x_z_dexner_dz_c", self.savepoint).astype(float))[
- :, :, ntnd - 1
- ]
- buffer = self._reduce_to_dim_size(buffer, (CellDim, KDim))
- return as_field((CellDim, KDim), buffer)
+ return self._get_field_component("x_z_dexner_dz_c", ntnd, (CellDim, KDim))
def z_rth_pr(self, ind):
- buffer = np.squeeze(self.serializer.read("x_z_rth_pr", self.savepoint).astype(float))[
- :, :, ind - 1
- ]
- buffer = self._reduce_to_dim_size(buffer, (CellDim, KDim))
- return as_field((CellDim, KDim), buffer)
+ return self._get_field_component("x_z_rth_pr", ind, (CellDim, KDim))
def z_th_ddz_exner_c(self):
return self._get_field("x_z_th_ddz_exner_c", CellDim, KDim)
@@ -1172,12 +1152,7 @@ def z_graddiv_vn(self):
return self._get_field("x_z_graddiv_vn", EdgeDim, KDim)
def z_grad_rth(self, ind):
- buffer = np.squeeze(self.serializer.read("x_z_grad_rth", self.savepoint).astype(float))[
- :, :, ind - 1
- ]
-
- dims = (CellDim, KDim)
- return as_field(dims, self._reduce_to_dim_size(buffer, dims))
+ return self._get_field_component("x_z_grad_rth", ind, (CellDim, KDim))
def z_gradh_exner_18(self):
return self._get_field("x_z_gradh_exner_18", EdgeDim, KDim)
diff --git a/model/common/tests/math_tests/test_smagorinsky.py b/model/common/tests/math_tests/test_smagorinsky.py
new file mode 100644
index 0000000000..7413392932
--- /dev/null
+++ b/model/common/tests/math_tests/test_smagorinsky.py
@@ -0,0 +1,38 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check .
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+
+from icon4py.model.common.dimension import KDim
+from icon4py.model.common.grid.simple import SimpleGrid
+from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift
+from icon4py.model.common.test_utils.helpers import random_field, zero_field
+from icon4py.model.common.test_utils.reference_funcs import enhanced_smagorinski_factor_numpy
+
+
+def test_init_enh_smag_fac():
+ grid = SimpleGrid()
+ enh_smag_fac = zero_field(grid, KDim)
+ a_vec = random_field(grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
+ fac = (0.67, 0.5, 1.3, 0.8)
+ z = (0.1, 0.2, 0.3, 0.4)
+
+ enhanced_smag_fac_np = enhanced_smagorinski_factor_numpy(fac, z, a_vec.asnumpy())
+ en_smag_fac_for_zero_nshift(
+ a_vec,
+ *fac,
+ *z,
+ out=enh_smag_fac,
+ offset_provider={"Koff": KDim},
+ )
+ assert np.allclose(enhanced_smag_fac_np, enh_smag_fac.asnumpy())
diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py
index 5eadabd445..9d30a2dc31 100644
--- a/model/driver/src/icon4py/model/driver/dycore_driver.py
+++ b/model/driver/src/icon4py/model/driver/dycore_driver.py
@@ -17,7 +17,6 @@
import click
from devtools import Timer
-from gt4py.next import Field
from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams
from icon4py.model.atmosphere.diffusion.diffusion_states import DiffusionDiagnosticState
@@ -35,9 +34,7 @@
get_processor_properties,
get_runtype,
)
-from icon4py.model.common.dimension import KDim
from icon4py.model.common.states.prognostic_state import PrognosticState
-from icon4py.model.common.test_utils.helpers import random_field, zero_field
from icon4py.model.driver.icon_configuration import IconRunConfig, read_config
from icon4py.model.driver.io_utils import (
configure_logging,
@@ -151,7 +148,7 @@ def time_integration(
prep_adv: PrepAdvection,
z_fields: ZFields, # local constants in solve_nh
nh_constants: NHConstants,
- bdy_divdamp: Field[[KDim], float],
+ inital_divdamp_fac_o2: float,
do_prep_adv: bool,
):
@@ -196,7 +193,7 @@ def time_integration(
prep_adv,
z_fields,
nh_constants,
- bdy_divdamp,
+ inital_divdamp_fac_o2,
do_prep_adv,
)
timer.capture()
@@ -217,7 +214,7 @@ def _integrate_one_time_step(
prep_adv: PrepAdvection,
z_fields: ZFields,
nh_constants: NHConstants,
- bdy_divdamp: Field[[KDim], float],
+ inital_divdamp_fac_o2: float,
do_prep_adv: bool,
):
@@ -227,7 +224,7 @@ def _integrate_one_time_step(
prep_adv,
z_fields,
nh_constants,
- bdy_divdamp,
+ inital_divdamp_fac_o2,
do_prep_adv,
)
@@ -247,7 +244,7 @@ def _do_dyn_substepping(
prep_adv: PrepAdvection,
z_fields: ZFields,
nh_constants: NHConstants,
- bdy_divdamp: Field[[KDim], float],
+ inital_divdamp_fac_o2: float,
do_prep_adv: bool,
):
@@ -265,7 +262,7 @@ def _do_dyn_substepping(
prep_adv=prep_adv,
z_fields=z_fields,
nh_constants=nh_constants,
- bdy_divdamp=bdy_divdamp,
+ divdamp_fac_o2=inital_divdamp_fac_o2,
dtime=self._substep_timestep,
idyn_timestep=dyn_substep,
l_recompute=do_recompute,
@@ -346,11 +343,6 @@ def initialize(file_path: Path, props: ProcessProperties):
nonhydro_params = NonHydrostaticParams(config.solve_nonhydro_config)
- enh_smag_fac = zero_field(icon_grid, KDim)
- a_vec = random_field(icon_grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
- fac = (0.67, 0.5, 1.3, 0.8)
- z = (0.1, 0.2, 0.3, 0.4)
-
solve_nonhydro = SolveNonhydro()
solve_nonhydro.init(
grid=icon_grid,
@@ -360,12 +352,8 @@ def initialize(file_path: Path, props: ProcessProperties):
interpolation_state=solve_nonhydro_interpolation_state,
vertical_params=vertical_geometry,
edge_geometry=edge_geometry,
- cell_areas=cell_geometry.area,
+ cell_geometry=cell_geometry,
owner_mask=c_owner_mask,
- a_vec=a_vec,
- enh_smag_fac=enh_smag_fac,
- fac=fac,
- z=z,
)
(
@@ -374,7 +362,7 @@ def initialize(file_path: Path, props: ProcessProperties):
z_fields,
nh_constants,
prep_adv,
- bdy_divdamp,
+ inital_divdamp_fac_o2,
prognostic_state_now,
prognostic_state_next,
) = read_initial_state(file_path, rank=props.rank)
@@ -393,7 +381,7 @@ def initialize(file_path: Path, props: ProcessProperties):
z_fields,
nh_constants,
prep_adv,
- bdy_divdamp,
+ inital_divdamp_fac_o2,
)
@@ -431,7 +419,7 @@ def main(input_path, run_path, mpi):
z_fields,
nh_constants,
prep_adv,
- bdy_divdamp,
+ inital_divdamp_fac_o2,
) = initialize(Path(input_path), parallel_props)
configure_logging(run_path, timeloop.simulation_date, parallel_props)
log.info(f"Starting ICON dycore run: {timeloop.simulation_date.isoformat()}")
@@ -451,7 +439,7 @@ def main(input_path, run_path, mpi):
prep_adv,
z_fields,
nh_constants,
- bdy_divdamp,
+ inital_divdamp_fac_o2,
do_prep_adv=False,
)
diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py
index 960a887049..04cc8f0a2c 100644
--- a/model/driver/src/icon4py/model/driver/io_utils.py
+++ b/model/driver/src/icon4py/model/driver/io_utils.py
@@ -184,8 +184,6 @@ def read_initial_state(
wgt_nnew_rth=solve_nonhydro_init_savepoint.wgt_nnew_rth(),
wgt_nnow_vel=solve_nonhydro_init_savepoint.wgt_nnow_vel(),
wgt_nnew_vel=solve_nonhydro_init_savepoint.wgt_nnew_vel(),
- scal_divdamp=solve_nonhydro_init_savepoint.scal_divdamp(),
- scal_divdamp_o2=solve_nonhydro_init_savepoint.scal_divdamp_o2(),
)
prognostic_state_next = PrognosticState(
@@ -208,7 +206,7 @@ def read_initial_state(
z_fields,
nh_constants,
prep_adv,
- solve_nonhydro_init_savepoint.bdy_divdamp(),
+ solve_nonhydro_init_savepoint.divdamp_fac_o2(),
prognostic_state_now,
prognostic_state_next,
)
diff --git a/model/driver/tests/test_timeloop.py b/model/driver/tests/test_timeloop.py
index d2adaac87c..5a94728dd3 100644
--- a/model/driver/tests/test_timeloop.py
+++ b/model/driver/tests/test_timeloop.py
@@ -31,7 +31,7 @@
from icon4py.model.common.grid.horizontal import CellParams, EdgeParams
from icon4py.model.common.grid.vertical import VerticalModelParams
from icon4py.model.common.states.prognostic_state import PrognosticState
-from icon4py.model.common.test_utils.helpers import dallclose, random_field, zero_field
+from icon4py.model.common.test_utils.helpers import dallclose
from icon4py.model.driver.dycore_driver import TimeLoop
@@ -130,11 +130,6 @@ def test_run_timeloop_single_step(
vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic()
)
- enh_smag_fac = zero_field(icon_grid, KDim)
- a_vec = random_field(icon_grid, KDim, low=1.0, high=10.0, extend={KDim: 1})
- fac = (0.67, 0.5, 1.3, 0.8)
- z = (0.1, 0.2, 0.3, 0.4)
-
assert timeloop_diffusion_savepoint_init.fac_bdydiff_v() == diffusion.fac_bdydiff_v
assert r04b09_iconrun_config.dtime == diffusion_dtime
@@ -160,8 +155,6 @@ def test_run_timeloop_single_step(
wgt_nnew_rth=sp.wgt_nnew_rth(),
wgt_nnow_vel=sp.wgt_nnow_vel(),
wgt_nnew_vel=sp.wgt_nnew_vel(),
- scal_divdamp=sp.scal_divdamp(),
- scal_divdamp_o2=sp.scal_divdamp_o2(),
)
nonhydro_interpolation_state = (
@@ -181,12 +174,8 @@ def test_run_timeloop_single_step(
interpolation_state=nonhydro_interpolation_state,
vertical_params=vertical_params,
edge_geometry=edge_geometry,
- cell_areas=cell_geometry.area,
+ cell_geometry=cell_geometry,
owner_mask=grid_savepoint.c_owner_mask(),
- a_vec=a_vec,
- enh_smag_fac=enh_smag_fac,
- fac=fac,
- z=z,
)
nonhydro_diagnostic_state = DiagnosticStateNonHydro(
@@ -243,7 +232,7 @@ def test_run_timeloop_single_step(
prep_adv,
z_fields,
nh_constants,
- sp.bdy_divdamp(),
+ sp.divdamp_fac_o2(),
do_prep_adv,
)
diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py
index 685ce0c2e3..33cf744920 100644
--- a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py
+++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py
@@ -121,6 +121,7 @@ def diffusion_init(
cells_edge_idx: Field[[CellDim, C2EDim], int32],
cells_area: Field[[CellDim], float],
# dsl specific additional args
+ mean_cell_area: float,
mask_hdiff: Field[[CellDim, KDim], bool],
zd_vertoffset: Field[
[CECDim], int32
@@ -182,7 +183,7 @@ def diffusion_init(
primal_normal_vert_y=edges_primal_normal_vert_2,
edge_areas=edges_area_edge,
)
- cell_params = CellParams(cells_area)
+ cell_params = CellParams(area=cells_area, mean_cell_area=mean_cell_area)
config: DiffusionConfig = DiffusionConfig(
diffusion_type=DiffusionType(hdiff_order),
hdiff_w=lhdiff_w,