From b1bfb8a298d73de84013dd0838e69298d73177d9 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Tue, 28 Nov 2023 19:49:57 +0100 Subject: [PATCH] Fix solve nonhydro divdamp fac o2 (#313) (Fix) - pass `divdamp_fac_o2` as argument to `SolveNonHydro.predictor_step` as it changes dynamically in the timeloop - use `mean_cell_area` from `CellParams` - cleanup tests. --- .../model/atmosphere/diffusion/diffusion.py | 2 + .../atmosphere/diffusion/diffusion_utils.py | 36 +- .../tests/diffusion_tests/test_diffusion.py | 11 +- .../diffusion_tests/test_diffusion_utils.py | 20 +- .../diffusion/tests/diffusion_tests/utils.py | 16 - .../dycore/nh_solve/solve_nonhydro.py | 144 +++--- .../dycore/state_utils/interpolation_state.py | 17 - .../dycore/state_utils/nh_constants.py | 2 - .../atmosphere/dycore/state_utils/utils.py | 98 ++-- .../dycore/velocity/velocity_advection.py | 1 - .../mpi_tests/test_parallel_solve_nonhydro.py | 18 +- .../tests/dycore_tests/test_solve_nonhydro.py | 441 +++++++----------- .../dycore/tests/dycore_tests/test_utils.py | 120 +++++ .../dycore_tests/test_velocity_advection.py | 2 +- .../src/icon4py/model/common/constants.py | 4 +- .../common/decomposition/mpi_decomposition.py | 9 +- .../icon4py/model/common/grid/horizontal.py | 2 + .../src/icon4py/model/common/math/__init__.py | 12 + .../icon4py/model/common/math/smagorinsky.py | 47 ++ .../model/common/test_utils/datatest_utils.py | 2 +- .../model/common/test_utils/helpers.py | 4 +- .../common/test_utils/reference_funcs.py | 30 ++ .../common/test_utils/serialbox_utils.py | 101 ++-- .../tests/math_tests/test_smagorinsky.py | 38 ++ .../py2f/wrappers/diffusion_wrapper.py | 3 +- 25 files changed, 571 insertions(+), 609 deletions(-) create mode 100644 model/atmosphere/dycore/tests/dycore_tests/test_utils.py create mode 100644 model/common/src/icon4py/model/common/math/__init__.py create mode 100644 model/common/src/icon4py/model/common/math/smagorinsky.py create mode 100644 model/common/src/icon4py/model/common/test_utils/reference_funcs.py create mode 100644 model/common/tests/math_tests/test_smagorinsky.py 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 7f5b520d24..aca2665509 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, 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,7 +539,9 @@ def run_predictor_step( nnow: int, nnew: int, ): - # TODO (magdalena) fix this! when fixing call of velocity advection in predictor, condition is broken, + 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 @@ -575,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 @@ -635,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 @@ -692,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") @@ -1026,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)( @@ -1342,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( @@ -1353,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 ) @@ -1408,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 @@ -1447,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") @@ -1504,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, @@ -1524,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, @@ -1558,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 +1865,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 fbffa2ec05..b18a66a033 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 @@ -394,7 +394,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 c4b3c8fa7b..acead33ca9 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,10 +186,8 @@ 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) + 1 + CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 2 ) cell_start_nudging = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.nudging(CellDim)) edge_start_lb_plus4 = icon_grid.get_start_index( @@ -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:, :], @@ -235,15 +215,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, ) @@ -277,9 +258,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, ) @@ -327,6 +309,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:, :], @@ -391,11 +374,10 @@ 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(), - solve_nonhydro.z_w_concorr_me.asnumpy(), + sp_exit.z_w_concorr_me().asnumpy()[:, nflatlev:], + solve_nonhydro.z_w_concorr_me.asnumpy()[:, nflatlev:], atol=2e-15, ) # stencils 39,40 @@ -453,7 +435,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:, :], @@ -466,10 +448,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", @@ -488,7 +522,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) @@ -506,51 +542,10 @@ def test_nonhydro_corrector_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 # TODO: @abishekg7 read from serialized data? nnew = 1 - 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(), - ) - - 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(), @@ -569,14 +564,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) @@ -593,27 +582,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, ) - 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, ) @@ -698,17 +684,14 @@ 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_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") @@ -717,80 +700,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(), - ) - - 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=_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) @@ -807,23 +727,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, @@ -833,7 +750,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(), @@ -895,66 +812,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, prognostic_state_nnew = 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) @@ -971,12 +841,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 _ in range(r_nsubsteps): @@ -984,11 +850,9 @@ def test_run_solve_nonhydro_multi_step( diagnostic_state_nh=diagnostic_state_nh, prognostic_state_ls=prognostic_state_ls, prep_adv=prep_adv, - config=config, - params=nonhydro_params, 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, @@ -1064,6 +928,15 @@ def test_run_solve_nonhydro_multi_step( assert dallclose(sp_step_exit.exner_new().asnumpy(), prognostic_state_nnew.exner.asnumpy()) +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 3dd30cc60f..8f4a6ca4c7 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py @@ -250,7 +250,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/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,