From 9181a9ee9940b8aa0b85354f2d53e463be2ec5eb Mon Sep 17 00:00:00 2001 From: Magdalena Date: Fri, 8 Dec 2023 13:14:35 +0100 Subject: [PATCH] Fix greenline dependencies (#327) * move states classes of dycore into one file * fix circular dependency from model/commen to model/atmosphere/dycore through serialbox_utils --- .../dycore/nh_solve/solve_nonhydro.py | 23 ++- .../dycore/state_utils/diagnostic_state.py | 88 ---------- .../dycore/state_utils/interpolation_state.py | 63 ------- .../dycore/state_utils/metric_state.py | 87 ---------- .../dycore/state_utils/prep_adv_state.py | 25 --- .../atmosphere/dycore/state_utils/states.py | 157 ++++++++++++++++++ .../dycore/velocity/velocity_advection.py | 10 +- .../mpi_tests/test_parallel_solve_nonhydro.py | 6 +- .../tests/dycore_tests/test_solve_nonhydro.py | 24 +-- .../dycore_tests/test_velocity_advection.py | 22 +-- .../dycore/tests/dycore_tests/utils.py | 82 +++++++++ .../common/test_utils/serialbox_utils.py | 78 --------- .../src/icon4py/model/driver/dycore_driver.py | 9 +- .../src/icon4py/model/driver/io_utils.py | 11 +- model/driver/tests/test_timeloop.py | 69 +++++++- 15 files changed, 360 insertions(+), 394 deletions(-) delete mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/diagnostic_state.py delete mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/interpolation_state.py delete mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/metric_state.py delete mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/prep_adv_state.py create mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/states.py create mode 100644 model/atmosphere/dycore/tests/dycore_tests/utils.py 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 1c92a21217..137bb20db9 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 @@ -13,7 +13,6 @@ import logging from typing import Final, Optional -import numpy as np from gt4py.next import as_field from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 @@ -138,11 +137,13 @@ from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_68 import ( mo_solve_nonhydro_stencil_68, ) -from icon4py.model.atmosphere.dycore.state_utils.diagnostic_state import DiagnosticStateNonHydro -from icon4py.model.atmosphere.dycore.state_utils.interpolation_state import InterpolationState -from icon4py.model.atmosphere.dycore.state_utils.metric_state import MetricStateNonHydro 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.states import ( + DiagnosticStateNonHydro, + InterpolationState, + MetricStateNonHydro, + PrepAdvection, +) from icon4py.model.atmosphere.dycore.state_utils.utils import ( _allocate, _allocate_indices, @@ -497,6 +498,10 @@ def time_step( vertical_start=0, vertical_end=self.grid.num_levels, offset_provider={}, + horizontal_start=0, + horizontal_end=end_cell_end, + vertical_start=0, + vertical_end=self.grid.num_levels, ) mo_solve_nonhydro_stencil_67.with_backend(backend)( @@ -663,7 +668,7 @@ def run_predictor_step( if self.config.igradp_method == 3: nhsolve_prog.predictor_stencils_4_5_6.with_backend(backend)( - wgtfacq_c_dsl=self.metric_state_nonhydro.wgtfacq_c_dsl, + wgtfacq_c_dsl=self.metric_state_nonhydro.wgtfacq_c, z_exner_ex_pr=self.z_exner_ex_pr, z_exner_ic=self.z_exner_ic, wgtfac_c=self.metric_state_nonhydro.wgtfac_c, @@ -708,7 +713,7 @@ def run_predictor_step( # Perturbation theta at top and surface levels nhsolve_prog.predictor_stencils_11_lower_upper.with_backend(backend)( - wgtfacq_c_dsl=self.metric_state_nonhydro.wgtfacq_c_dsl, + wgtfacq_c_dsl=self.metric_state_nonhydro.wgtfacq_c, z_rth_pr=self.z_rth_pr_2, theta_ref_ic=self.metric_state_nonhydro.theta_ref_ic, z_theta_v_pr_ic=self.z_theta_v_pr_ic, @@ -1073,7 +1078,7 @@ def run_predictor_step( vn_ie=diagnostic_state_nh.vn_ie, z_vt_ie=z_fields.z_vt_ie, z_kin_hor_e=z_fields.z_kin_hor_e, - wgtfacq_e_dsl=self.metric_state_nonhydro.wgtfacq_e_dsl, + wgtfacq_e_dsl=self.metric_state_nonhydro.wgtfacq_e, k_field=self.k_field, nlev=self.grid.num_levels, horizontal_start=start_edge_lb_plus4, @@ -1087,7 +1092,7 @@ def run_predictor_step( e_bln_c_s=self.interpolation_state.e_bln_c_s, z_w_concorr_me=self.z_w_concorr_me, wgtfac_c=self.metric_state_nonhydro.wgtfac_c, - wgtfacq_c_dsl=self.metric_state_nonhydro.wgtfacq_c_dsl, + wgtfacq_c_dsl=self.metric_state_nonhydro.wgtfacq_c, w_concorr_c=diagnostic_state_nh.w_concorr_c, k_field=self.k_field, nflatlev_startindex_plus1=int32(self.vertical_params.nflatlev + 1), diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/diagnostic_state.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/diagnostic_state.py deleted file mode 100644 index f12aa39b55..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/diagnostic_state.py +++ /dev/null @@ -1,88 +0,0 @@ -# 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 dataclasses import dataclass - -from gt4py.next.common import Field - -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim - - -@dataclass -class DiagnosticState: - # fields for 3D elements in turbdiff - hdef_ic: Field[ - [CellDim, KDim], float - ] # ! divergence at half levels(nproma,nlevp1,nblks_c) [1/s] - div_ic: Field[ - [CellDim, KDim], float - ] # ! horizontal wind field deformation (nproma,nlevp1,nblks_c) [1/s^2] - dwdx: Field[ - [CellDim, KDim], float - ] # zonal gradient of vertical wind speed (nproma,nlevp1,nblks_c) [1/s] - - dwdy: Field[ - [CellDim, KDim], float - ] # meridional gradient of vertical wind speed (nproma,nlevp1,nblks_c) - - vt: Field[[EdgeDim, KDim], float] - vn_ie: Field[ - [EdgeDim, KDim], float - ] # normal wind at half levels (nproma,nlevp1,nblks_e) [m/s] # TODO: change this back to KHalfDim, but how do we treat it wrt to field_operators and domain? - w_concorr_c: Field[ - [CellDim, KDim], float - ] # contravariant vert correction (nproma,nlevp1,nblks_c)[m/s] # TODO: change this back to KHalfDim, but how do we treat it wrt to field_operators and domain? - ddt_w_adv_pc: Field[[CellDim, KDim], float] - ddt_vn_apc_pc: Field[[EdgeDim, KDim], float] - ntnd: float - - -@dataclass -class DiagnosticStateNonHydro: - vt: Field[[EdgeDim, KDim], float] - vn_ie: Field[ - [EdgeDim, KDim], float - ] # normal wind at half levels (nproma,nlevp1,nblks_e) [m/s] # TODO: change this back to KHalfDim, but how do we treat it wrt to field_operators and domain? - w_concorr_c: Field[ - [CellDim, KDim], float - ] # contravariant vert correction (nproma,nlevp1,nblks_c)[m/s] # TODO: change this back to KHalfDim, but how do we treat it wrt to field_operators and domain? - theta_v_ic: Field[[CellDim, KDim], float] - exner_pr: Field[[CellDim, KDim], float] - rho_ic: Field[[CellDim, KDim], float] - ddt_exner_phy: Field[[CellDim, KDim], float] - grf_tend_rho: Field[[CellDim, KDim], float] - grf_tend_thv: Field[[CellDim, KDim], float] - grf_tend_w: Field[[CellDim, KDim], float] - mass_fl_e: Field[[EdgeDim, KDim], float] - ddt_vn_phy: Field[[EdgeDim, KDim], float] - grf_tend_vn: Field[[EdgeDim, KDim], float] - ddt_vn_apc_ntl1: Field[[EdgeDim, KDim], float] - ddt_vn_apc_ntl2: Field[[EdgeDim, KDim], float] - ddt_w_adv_ntl1: Field[[CellDim, KDim], float] - ddt_w_adv_ntl2: Field[[CellDim, KDim], float] - - # Analysis increments - rho_incr: Field[[EdgeDim, KDim], float] # moist density increment [kg/m^3] - vn_incr: Field[[EdgeDim, KDim], float] # normal velocity increment [m/s] - exner_incr: Field[[EdgeDim, KDim], float] # exner increment [- ] - - @property - def ddt_vn_apc_pc( - self, - ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: - return (self.ddt_vn_apc_ntl1, self.ddt_vn_apc_ntl2) - - @property - def ddt_w_adv_pc( - self, - ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - return (self.ddt_w_adv_ntl1, self.ddt_w_adv_ntl2) 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 deleted file mode 100644 index 440c4a4b9c..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/interpolation_state.py +++ /dev/null @@ -1,63 +0,0 @@ -# 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 dataclasses import dataclass - -from gt4py.next.common import Field - -from icon4py.model.common.dimension import ( - C2E2CODim, - CEDim, - CellDim, - E2C2EDim, - E2C2EODim, - E2CDim, - ECDim, - EdgeDim, - V2CDim, - V2EDim, - VertexDim, -) - - -@dataclass -class InterpolationState: - """Represents the ICON interpolation state used int SolveNonHydro.""" - - e_bln_c_s: Field[[CEDim], float] # coefficent for bilinear interpolation from edge to cell () - rbf_coeff_1: Field[ - [VertexDim, V2EDim], float - ] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) - rbf_coeff_2: Field[ - [VertexDim, V2EDim], float - ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) - - geofac_div: Field[[CEDim], float] # factor for divergence (nproma,cell_type,nblks_c) - - geofac_n2s: Field[ - [CellDim, C2E2CODim], float - ] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) - geofac_grg_x: Field[[CellDim, C2E2CODim], float] - geofac_grg_y: Field[ - [CellDim, C2E2CODim], float - ] # factors for green gauss gradient (nproma,4,nblks_c,2) - nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges - - c_lin_e: Field[[EdgeDim, E2CDim], float] - geofac_grdiv: Field[[EdgeDim, E2C2EODim], float] - rbf_vec_coeff_e: Field[[EdgeDim, E2C2EDim], float] - c_intp: Field[[VertexDim, V2CDim], float] - geofac_rot: Field[[VertexDim, V2EDim], float] - pos_on_tplane_e_1: Field[[ECDim], float] - pos_on_tplane_e_2: Field[[ECDim], float] - e_flx_avg: Field[[EdgeDim, E2C2EODim], float] diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/metric_state.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/metric_state.py deleted file mode 100644 index 3d32d340b9..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/metric_state.py +++ /dev/null @@ -1,87 +0,0 @@ -# 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 dataclasses import dataclass - -from gt4py.next.common import Field -from numpy import int32 - -from icon4py.model.common.dimension import CECDim, CellDim, ECDim, EdgeDim, KDim - - -@dataclass -class MetricState: - mask_hdiff: Field[[CellDim, KDim], bool] - theta_ref_mc: Field[[CellDim, KDim], float] - wgtfac_c: Field[ - [CellDim, KDim], float - ] # weighting factor for interpolation from full to half levels (nproma,nlevp1,nblks_c) - zd_vertidx: Field[[CECDim, KDim], int32] - zd_diffcoef: Field[[CellDim, KDim], float] - zd_intcoef: Field[[CECDim, KDim], float] - - coeff_gradekin: Field[[ECDim], float] - ddqz_z_full_e: Field[[EdgeDim, KDim], float] - wgtfac_e: Field[[EdgeDim, KDim], float] - wgtfacq_e_dsl: Field[[EdgeDim, KDim], float] - ddxn_z_full: Field[[EdgeDim, KDim], float] - ddxt_z_full: Field[[EdgeDim, KDim], float] - ddqz_z_half: Field[[CellDim, KDim], float] # half KDim ? - coeff1_dwdz: Field[[CellDim, KDim], float] - coeff2_dwdz: Field[[CellDim, KDim], float] - zd_vertoffset: Field[[CECDim, KDim], int32] - - -@dataclass -class MetricStateNonHydro: - bdy_halo_c: Field[[CellDim], bool] - # Finally, a mask field that excludes boundary halo points - mask_prog_halo_c: Field[[CellDim, KDim], bool] - rayleigh_w: Field[[KDim], float] - - wgtfac_c: Field[[CellDim, KDim], float] - wgtfacq_c_dsl: Field[[CellDim, KDim], float] - wgtfac_e: Field[[EdgeDim, KDim], float] - wgtfacq_e_dsl: Field[[EdgeDim, KDim], float] - - exner_exfac: Field[[CellDim, KDim], float] - exner_ref_mc: Field[[CellDim, KDim], float] - rho_ref_mc: Field[[CellDim, KDim], float] - theta_ref_mc: Field[[CellDim, KDim], float] - rho_ref_me: Field[[EdgeDim, KDim], float] - theta_ref_me: Field[[EdgeDim, KDim], float] - theta_ref_ic: Field[[CellDim, KDim], float] - - d_exner_dz_ref_ic: Field[[CellDim, KDim], float] - ddqz_z_half: Field[[CellDim, KDim], float] # half KDim ? - d2dexdz2_fac1_mc: Field[[CellDim, KDim], float] - d2dexdz2_fac2_mc: Field[[CellDim, KDim], float] - ddxn_z_full: Field[[EdgeDim, KDim], float] - ddqz_z_full_e: Field[[EdgeDim, KDim], float] - ddxt_z_full: Field[[EdgeDim, KDim], float] - inv_ddqz_z_full: Field[[CellDim, KDim], float] - - vertoffset_gradp: Field[[ECDim, KDim], float] - zdiff_gradp: Field[[ECDim, KDim], float] - ipeidx_dsl: Field[[EdgeDim, KDim], bool] - pg_exdist: Field[[EdgeDim, KDim], float] - - vwind_expl_wgt: Field[[CellDim], float] - vwind_impl_wgt: Field[[CellDim], float] - - hmask_dd3d: Field[[EdgeDim], float] - scalfac_dd3d: Field[[KDim], float] - - coeff1_dwdz: Field[[CellDim, KDim], float] - coeff2_dwdz: Field[[CellDim, KDim], float] - coeff_gradekin: Field[[ECDim], float] diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/prep_adv_state.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/prep_adv_state.py deleted file mode 100644 index 8a94ccafb4..0000000000 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/prep_adv_state.py +++ /dev/null @@ -1,25 +0,0 @@ -# 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 dataclasses import dataclass - -from gt4py.next.common import Field - -from icon4py.model.common.dimension import CellDim, EdgeDim, KDim - - -@dataclass -class PrepAdvection: - vn_traj: Field[[EdgeDim, KDim], float] - mass_flx_me: Field[[EdgeDim, KDim], float] - mass_flx_ic: Field[[CellDim, KDim], float] diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/states.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/states.py new file mode 100644 index 0000000000..fb413282bf --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/states.py @@ -0,0 +1,157 @@ +# 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 dataclasses import dataclass + +from gt4py.next.common import Field + +from icon4py.model.common.dimension import ( + C2E2CODim, + CEDim, + CellDim, + E2C2EDim, + E2C2EODim, + E2CDim, + ECDim, + EdgeDim, + KDim, + V2CDim, + V2EDim, + VertexDim, +) + + +@dataclass +class DiagnosticStateNonHydro: + vt: Field[[EdgeDim, KDim], float] + vn_ie: Field[ + [EdgeDim, KDim], float + ] # normal wind at half levels (nproma,nlevp1,nblks_e) [m/s] # TODO: change this back to KHalfDim, but how do we treat it wrt to field_operators and domain? + w_concorr_c: Field[ + [CellDim, KDim], float + ] # contravariant vert correction (nproma,nlevp1,nblks_c)[m/s] # TODO: change this back to KHalfDim, but how do we treat it wrt to field_operators and domain? + theta_v_ic: Field[[CellDim, KDim], float] + exner_pr: Field[[CellDim, KDim], float] + rho_ic: Field[[CellDim, KDim], float] + ddt_exner_phy: Field[[CellDim, KDim], float] + grf_tend_rho: Field[[CellDim, KDim], float] + grf_tend_thv: Field[[CellDim, KDim], float] + grf_tend_w: Field[[CellDim, KDim], float] + mass_fl_e: Field[[EdgeDim, KDim], float] + ddt_vn_phy: Field[[EdgeDim, KDim], float] + grf_tend_vn: Field[[EdgeDim, KDim], float] + ddt_vn_apc_ntl1: Field[[EdgeDim, KDim], float] + ddt_vn_apc_ntl2: Field[[EdgeDim, KDim], float] + ddt_w_adv_ntl1: Field[[CellDim, KDim], float] + ddt_w_adv_ntl2: Field[[CellDim, KDim], float] + + # Analysis increments + rho_incr: Field[[EdgeDim, KDim], float] # moist density increment [kg/m^3] + vn_incr: Field[[EdgeDim, KDim], float] # normal velocity increment [m/s] + exner_incr: Field[[EdgeDim, KDim], float] # exner increment [- ] + + @property + def ddt_vn_apc_pc( + self, + ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: + return (self.ddt_vn_apc_ntl1, self.ddt_vn_apc_ntl2) + + @property + def ddt_w_adv_pc( + self, + ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + return (self.ddt_w_adv_ntl1, self.ddt_w_adv_ntl2) + + +@dataclass +class InterpolationState: + """Represents the ICON interpolation state used int SolveNonHydro.""" + + e_bln_c_s: Field[[CEDim], float] # coefficent for bilinear interpolation from edge to cell () + rbf_coeff_1: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_1(nproma, rbf_vec_dim_v, nblks_v) + rbf_coeff_2: Field[ + [VertexDim, V2EDim], float + ] # rbf_vec_coeff_v_2(nproma, rbf_vec_dim_v, nblks_v) + + geofac_div: Field[[CEDim], float] # factor for divergence (nproma,cell_type,nblks_c) + + geofac_n2s: Field[ + [CellDim, C2E2CODim], float + ] # factor for nabla2-scalar (nproma,cell_type+1,nblks_c) + geofac_grg_x: Field[[CellDim, C2E2CODim], float] + geofac_grg_y: Field[ + [CellDim, C2E2CODim], float + ] # factors for green gauss gradient (nproma,4,nblks_c,2) + nudgecoeff_e: Field[[EdgeDim], float] # Nudgeing coeffients for edges + + c_lin_e: Field[[EdgeDim, E2CDim], float] + geofac_grdiv: Field[[EdgeDim, E2C2EODim], float] + rbf_vec_coeff_e: Field[[EdgeDim, E2C2EDim], float] + c_intp: Field[[VertexDim, V2CDim], float] + geofac_rot: Field[[VertexDim, V2EDim], float] + pos_on_tplane_e_1: Field[[ECDim], float] + pos_on_tplane_e_2: Field[[ECDim], float] + e_flx_avg: Field[[EdgeDim, E2C2EODim], float] + + +@dataclass +class MetricStateNonHydro: + bdy_halo_c: Field[[CellDim], bool] + # Finally, a mask field that excludes boundary halo points + mask_prog_halo_c: Field[[CellDim, KDim], bool] + rayleigh_w: Field[[KDim], float] + + wgtfac_c: Field[[CellDim, KDim], float] + wgtfacq_c: Field[[CellDim, KDim], float] + wgtfac_e: Field[[EdgeDim, KDim], float] + wgtfacq_e: Field[[EdgeDim, KDim], float] + + exner_exfac: Field[[CellDim, KDim], float] + exner_ref_mc: Field[[CellDim, KDim], float] + rho_ref_mc: Field[[CellDim, KDim], float] + theta_ref_mc: Field[[CellDim, KDim], float] + rho_ref_me: Field[[EdgeDim, KDim], float] + theta_ref_me: Field[[EdgeDim, KDim], float] + theta_ref_ic: Field[[CellDim, KDim], float] + + d_exner_dz_ref_ic: Field[[CellDim, KDim], float] + ddqz_z_half: Field[[CellDim, KDim], float] # half KDim ? + d2dexdz2_fac1_mc: Field[[CellDim, KDim], float] + d2dexdz2_fac2_mc: Field[[CellDim, KDim], float] + ddxn_z_full: Field[[EdgeDim, KDim], float] + ddqz_z_full_e: Field[[EdgeDim, KDim], float] + ddxt_z_full: Field[[EdgeDim, KDim], float] + inv_ddqz_z_full: Field[[CellDim, KDim], float] + + vertoffset_gradp: Field[[ECDim, KDim], float] + zdiff_gradp: Field[[ECDim, KDim], float] + ipeidx_dsl: Field[[EdgeDim, KDim], bool] + pg_exdist: Field[[EdgeDim, KDim], float] + + vwind_expl_wgt: Field[[CellDim], float] + vwind_impl_wgt: Field[[CellDim], float] + + hmask_dd3d: Field[[EdgeDim], float] + scalfac_dd3d: Field[[KDim], float] + + coeff1_dwdz: Field[[CellDim, KDim], float] + coeff2_dwdz: Field[[CellDim, KDim], float] + coeff_gradekin: Field[[ECDim], float] + + +@dataclass +class PrepAdvection: + vn_traj: Field[[EdgeDim, KDim], float] + mass_flx_me: Field[[EdgeDim, KDim], float] + mass_flx_ic: Field[[CellDim, KDim], float] 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 e6340459d5..7888e09356 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 @@ -55,9 +55,11 @@ from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_20 import ( mo_velocity_advection_stencil_20, ) -from icon4py.model.atmosphere.dycore.state_utils.diagnostic_state import DiagnosticStateNonHydro -from icon4py.model.atmosphere.dycore.state_utils.interpolation_state import InterpolationState -from icon4py.model.atmosphere.dycore.state_utils.metric_state import MetricStateNonHydro +from icon4py.model.atmosphere.dycore.state_utils.states import ( + DiagnosticStateNonHydro, + InterpolationState, + MetricStateNonHydro, +) from icon4py.model.atmosphere.dycore.state_utils.utils import _allocate, _allocate_indices from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim from icon4py.model.common.grid.horizontal import EdgeParams, HorizontalMarkerIndex @@ -249,7 +251,7 @@ def run_predictor_step( offset_provider={}, ) velocity_prog.mo_velocity_advection_stencil_06.with_backend(backend)( - wgtfacq_e=self.metric_state.wgtfacq_e_dsl, + wgtfacq_e=self.metric_state.wgtfacq_e, vn=prognostic_state.vn, vn_ie=diagnostic_state.vn_ie, horizontal_start=start_edge_lb_plus4, 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 2b2d6786db..8b27f60238 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 @@ -20,9 +20,11 @@ NonHydrostaticParams, SolveNonhydro, ) -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.states import ( + DiagnosticStateNonHydro, + PrepAdvection, +) 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 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 7d56058655..6a7fad0b3c 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -20,9 +20,11 @@ NonHydrostaticParams, SolveNonhydro, ) -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.states import ( + DiagnosticStateNonHydro, + PrepAdvection, +) from icon4py.model.atmosphere.dycore.state_utils.utils import ( _allocate, _calculate_bdy_divdamp, @@ -37,6 +39,8 @@ from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.helpers import dallclose +from .utils import construct_interpolation_state_for_nonhydro, construct_nh_metric_state + backend = run_gtfn @@ -143,8 +147,8 @@ def test_nonhydro_predictor_step( 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) + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() @@ -574,8 +578,8 @@ def test_nonhydro_corrector_step( 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) + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() @@ -723,8 +727,8 @@ def test_run_solve_nonhydro_single_step( 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) + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() @@ -844,8 +848,8 @@ def test_run_solve_nonhydro_multi_step( 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) + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() 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 f1b56b5c5d..0d6e99dc90 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_velocity_advection.py @@ -14,13 +14,15 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.state_utils.diagnostic_state import DiagnosticStateNonHydro +from icon4py.model.atmosphere.dycore.state_utils.states import DiagnosticStateNonHydro from icon4py.model.atmosphere.dycore.velocity.velocity_advection import VelocityAdvection 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 +from .utils import construct_interpolation_state_for_nonhydro, construct_nh_metric_state + @pytest.mark.datatest def test_scalfactors(savepoint_velocity_init, icon_grid): @@ -48,8 +50,8 @@ def test_velocity_init( step_date_init, damping_height, ): - interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro() - metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels) + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) vertical_params = VerticalModelParams( vct_a=grid_savepoint.vct_a(), @@ -87,8 +89,8 @@ def test_verify_velocity_init_against_regular_savepoint( savepoint = savepoint_velocity_init dtime = savepoint.get_metadata("dtime").get("dtime") - interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro() - metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels) + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) vertical_params = VerticalModelParams( vct_a=grid_savepoint.vct_a(), rayleigh_damping_height=damping_height, @@ -163,9 +165,9 @@ def test_velocity_predictor_step( rho=None, exner=None, ) - interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro() + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) - metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() @@ -213,7 +215,7 @@ def test_velocity_predictor_step( # stencil 07 if not vn_only: assert dallclose( - icon_result_z_v_grad_w.asnumpy()[3777:31558, :], + icon_result_z_v_grad_w[3777:31558, :], velocity_advection.z_v_grad_w.asnumpy()[3777:31558, :], atol=1.0e-16, ) @@ -313,9 +315,9 @@ def test_velocity_corrector_step( exner=None, ) - interpolation_state = interpolation_savepoint.construct_interpolation_state_for_nonhydro() + interpolation_state = construct_interpolation_state_for_nonhydro(interpolation_savepoint) - metric_state_nonhydro = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels) + metric_state_nonhydro = construct_nh_metric_state(metrics_savepoint, icon_grid.num_levels) cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() diff --git a/model/atmosphere/dycore/tests/dycore_tests/utils.py b/model/atmosphere/dycore/tests/dycore_tests/utils.py new file mode 100644 index 0000000000..a6eee5b0f2 --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore_tests/utils.py @@ -0,0 +1,82 @@ +# 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 icon4py.model.atmosphere.dycore.state_utils.states import ( + InterpolationState, + MetricStateNonHydro, +) +from icon4py.model.common.dimension import CEDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field +from icon4py.model.common.test_utils.serialbox_utils import InterpolationSavepoint, MetricSavepoint + + +def construct_interpolation_state_for_nonhydro( + savepoint: InterpolationSavepoint, +) -> InterpolationState: + grg = savepoint.geofac_grg() + return InterpolationState( + c_lin_e=savepoint.c_lin_e(), + c_intp=savepoint.c_intp(), + e_flx_avg=savepoint.e_flx_avg(), + geofac_grdiv=savepoint.geofac_grdiv(), + geofac_rot=savepoint.geofac_rot(), + pos_on_tplane_e_1=savepoint.pos_on_tplane_e_x(), + pos_on_tplane_e_2=savepoint.pos_on_tplane_e_y(), + rbf_vec_coeff_e=savepoint.rbf_vec_coeff_e(), + e_bln_c_s=as_1D_sparse_field(savepoint.e_bln_c_s(), CEDim), + rbf_coeff_1=savepoint.rbf_vec_coeff_v1(), + rbf_coeff_2=savepoint.rbf_vec_coeff_v2(), + geofac_div=as_1D_sparse_field(savepoint.geofac_div(), CEDim), + geofac_n2s=savepoint.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=savepoint.nudgecoeff_e(), + ) + + +def construct_nh_metric_state(savepoint: MetricSavepoint, num_k_lev) -> MetricStateNonHydro: + return MetricStateNonHydro( + bdy_halo_c=savepoint.bdy_halo_c(), + mask_prog_halo_c=savepoint.mask_prog_halo_c(), + rayleigh_w=savepoint.rayleigh_w(), + exner_exfac=savepoint.exner_exfac(), + exner_ref_mc=savepoint.exner_ref_mc(), + wgtfac_c=savepoint.wgtfac_c(), + wgtfacq_c=savepoint.wgtfacq_c_dsl(), + inv_ddqz_z_full=savepoint.inv_ddqz_z_full(), + rho_ref_mc=savepoint.rho_ref_mc(), + theta_ref_mc=savepoint.theta_ref_mc(), + vwind_expl_wgt=savepoint.vwind_expl_wgt(), + d_exner_dz_ref_ic=savepoint.d_exner_dz_ref_ic(), + ddqz_z_half=savepoint.ddqz_z_half(), + theta_ref_ic=savepoint.theta_ref_ic(), + d2dexdz2_fac1_mc=savepoint.d2dexdz2_fac1_mc(), + d2dexdz2_fac2_mc=savepoint.d2dexdz2_fac2_mc(), + rho_ref_me=savepoint.rho_ref_me(), + theta_ref_me=savepoint.theta_ref_me(), + ddxn_z_full=savepoint.ddxn_z_full(), + zdiff_gradp=savepoint.zdiff_gradp(), + vertoffset_gradp=savepoint.vertoffset_gradp(), + ipeidx_dsl=savepoint.ipeidx_dsl(), + pg_exdist=savepoint.pg_exdist(), + ddqz_z_full_e=savepoint.ddqz_z_full_e(), + ddxt_z_full=savepoint.ddxt_z_full(), + wgtfac_e=savepoint.wgtfac_e(), + wgtfacq_e=savepoint.wgtfacq_e_dsl(num_k_lev), + vwind_impl_wgt=savepoint.vwind_impl_wgt(), + hmask_dd3d=savepoint.hmask_dd3d(), + scalfac_dd3d=savepoint.scalfac_dd3d(), + coeff1_dwdz=savepoint.coeff1_dwdz(), + coeff2_dwdz=savepoint.coeff2_dwdz(), + coeff_gradekin=savepoint.coeff_gradekin(), + ) 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 113aed4589..78e02152cf 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 @@ -24,9 +24,6 @@ DiffusionInterpolationState, DiffusionMetricState, ) -from icon4py.model.atmosphere.dycore.state_utils.diagnostic_state import DiagnosticState -from icon4py.model.atmosphere.dycore.state_utils.interpolation_state import InterpolationState -from icon4py.model.atmosphere.dycore.state_utils.metric_state import MetricStateNonHydro from icon4py.model.common import dimension from icon4py.model.common.decomposition.definitions import DecompositionInfo from icon4py.model.common.dimension import ( @@ -466,27 +463,6 @@ def construct_interpolation_state_for_diffusion( nudgecoeff_e=self.nudgecoeff_e(), ) - def construct_interpolation_state_for_nonhydro(self) -> InterpolationState: - grg = self.geofac_grg() - return InterpolationState( - c_lin_e=self.c_lin_e(), - c_intp=self.c_intp(), - e_flx_avg=self.e_flx_avg(), - geofac_grdiv=self.geofac_grdiv(), - geofac_rot=self.geofac_rot(), - pos_on_tplane_e_1=self.pos_on_tplane_e_x(), - pos_on_tplane_e_2=self.pos_on_tplane_e_y(), - rbf_vec_coeff_e=self.rbf_vec_coeff_e(), - e_bln_c_s=as_1D_sparse_field(self.e_bln_c_s(), CEDim), - rbf_coeff_1=self.rbf_vec_coeff_v1(), - rbf_coeff_2=self.rbf_vec_coeff_v2(), - geofac_div=as_1D_sparse_field(self.geofac_div(), CEDim), - geofac_n2s=self.geofac_n2s(), - geofac_grg_x=grg[0], - geofac_grg_y=grg[1], - nudgecoeff_e=self.nudgecoeff_e(), - ) - class MetricSavepoint(IconSavepoint): def bdy_halo_c(self): @@ -549,9 +525,6 @@ def vwind_impl_wgt(self): def wgtfacq_c_dsl(self): return self._get_field("wgtfacq_c_dsl", CellDim, KDim) - def wgtfacq_c(self): - return self._get_field("wgtfacq_c", CellDim, KDim) - def zdiff_gradp(self): field = self._get_field("zdiff_gradp_dsl", EdgeDim, E2CDim, KDim) return flatten_first_two_dims(ECDim, KDim, field=field) @@ -633,43 +606,6 @@ def zd_vertidx(self): def zd_indlist(self): return np.squeeze(self.serializer.read("zd_indlist", self.savepoint)) - def construct_nh_metric_state(self, num_k_lev) -> MetricStateNonHydro: - return MetricStateNonHydro( - bdy_halo_c=self.bdy_halo_c(), - mask_prog_halo_c=self.mask_prog_halo_c(), - rayleigh_w=self.rayleigh_w(), - exner_exfac=self.exner_exfac(), - exner_ref_mc=self.exner_ref_mc(), - wgtfac_c=self.wgtfac_c(), - wgtfacq_c_dsl=self.wgtfacq_c_dsl(), - inv_ddqz_z_full=self.inv_ddqz_z_full(), - rho_ref_mc=self.rho_ref_mc(), - theta_ref_mc=self.theta_ref_mc(), - vwind_expl_wgt=self.vwind_expl_wgt(), - d_exner_dz_ref_ic=self.d_exner_dz_ref_ic(), - ddqz_z_half=self.ddqz_z_half(), - theta_ref_ic=self.theta_ref_ic(), - d2dexdz2_fac1_mc=self.d2dexdz2_fac1_mc(), - d2dexdz2_fac2_mc=self.d2dexdz2_fac2_mc(), - rho_ref_me=self.rho_ref_me(), - theta_ref_me=self.theta_ref_me(), - ddxn_z_full=self.ddxn_z_full(), - zdiff_gradp=self.zdiff_gradp(), - vertoffset_gradp=self.vertoffset_gradp(), - ipeidx_dsl=self.ipeidx_dsl(), - pg_exdist=self.pg_exdist(), - ddqz_z_full_e=self.ddqz_z_full_e(), - ddxt_z_full=self.ddxt_z_full(), - wgtfac_e=self.wgtfac_e(), - wgtfacq_e_dsl=self.wgtfacq_e_dsl(num_k_lev), - vwind_impl_wgt=self.vwind_impl_wgt(), - hmask_dd3d=self.hmask_dd3d(), - scalfac_dd3d=self.scalfac_dd3d(), - coeff1_dwdz=self.coeff1_dwdz(), - coeff2_dwdz=self.coeff2_dwdz(), - coeff_gradekin=self.coeff_gradekin(), - ) - def construct_metric_state_for_diffusion(self) -> DiffusionMetricState: return DiffusionMetricState( mask_hdiff=self.mask_hdiff(), @@ -753,20 +689,6 @@ def construct_diagnostics_for_diffusion(self) -> DiffusionDiagnosticState: dwdy=self.dwdy(), ) - def construct_diagnostics(self) -> DiagnosticState: - return DiagnosticState( - hdef_ic=self.hdef_ic(), - div_ic=self.div_ic(), - dwdx=self.dwdx(), - dwdy=self.dwdy(), - vt=None, - vn_ie=None, - w_concorr_c=None, - ddt_w_adv_pc_before=None, - ddt_vn_apc_pc_before=None, - ntnd=None, - ) - class IconNonHydroInitSavepoint(IconSavepoint): def bdy_divdamp(self): diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py index 9d30a2dc31..33ef96cdbc 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -24,9 +24,11 @@ NonHydrostaticParams, SolveNonhydro, ) -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.states import ( + DiagnosticStateNonHydro, + PrepAdvection, +) from icon4py.model.atmosphere.dycore.state_utils.z_fields import ZFields from icon4py.model.common.decomposition.definitions import ( ProcessProperties, @@ -151,7 +153,6 @@ def time_integration( inital_divdamp_fac_o2: float, do_prep_adv: bool, ): - log.info( f"starting time loop for dtime={self.run_config.dtime} n_timesteps={self._n_time_steps}" ) @@ -217,7 +218,6 @@ def _integrate_one_time_step( inital_divdamp_fac_o2: float, do_prep_adv: bool, ): - self._do_dyn_substepping( solve_nonhydro_diagnostic_state, prognostic_state_list, @@ -247,7 +247,6 @@ def _do_dyn_substepping( inital_divdamp_fac_o2: float, do_prep_adv: bool, ): - # TODO (Chia Rui): compute airmass for prognostic_state here do_recompute = True diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py index 04cc8f0a2c..b21c206434 100644 --- a/model/driver/src/icon4py/model/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -23,11 +23,13 @@ DiffusionInterpolationState, DiffusionMetricState, ) -from icon4py.model.atmosphere.dycore.state_utils.diagnostic_state import DiagnosticStateNonHydro -from icon4py.model.atmosphere.dycore.state_utils.interpolation_state import InterpolationState -from icon4py.model.atmosphere.dycore.state_utils.metric_state import MetricStateNonHydro 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.states import ( + DiagnosticStateNonHydro, + InterpolationState, + MetricStateNonHydro, + PrepAdvection, +) 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.definitions import DecompositionInfo, ProcessProperties @@ -75,7 +77,6 @@ def read_icon_grid( # TODO (Chia Rui): initialization of prognostic variables and topography of Jablonowski Williamson test def model_initialization(): - # create two prognostic states, nnow and nnew? # at least two prognostic states are global because they are needed in the dycore, AND possibly nesting and restart processes in the future # one is enough for the JW test diff --git a/model/driver/tests/test_timeloop.py b/model/driver/tests/test_timeloop.py index 5a94728dd3..fa72699114 100644 --- a/model/driver/tests/test_timeloop.py +++ b/model/driver/tests/test_timeloop.py @@ -22,16 +22,20 @@ NonHydrostaticParams, SolveNonhydro, ) -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.states import ( + DiagnosticStateNonHydro, + InterpolationState, + MetricStateNonHydro, + PrepAdvection, +) 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.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim 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 +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, dallclose from icon4py.model.driver.dycore_driver import TimeLoop @@ -157,10 +161,60 @@ def test_run_timeloop_single_step( wgt_nnew_vel=sp.wgt_nnew_vel(), ) - nonhydro_interpolation_state = ( - interpolation_savepoint.construct_interpolation_state_for_nonhydro() + grg = interpolation_savepoint.geofac_grg() + nonhydro_interpolation_state = InterpolationState( + c_lin_e=interpolation_savepoint.c_lin_e(), + c_intp=interpolation_savepoint.c_intp(), + e_flx_avg=interpolation_savepoint.e_flx_avg(), + geofac_grdiv=interpolation_savepoint.geofac_grdiv(), + geofac_rot=interpolation_savepoint.geofac_rot(), + pos_on_tplane_e_1=interpolation_savepoint.pos_on_tplane_e_x(), + pos_on_tplane_e_2=interpolation_savepoint.pos_on_tplane_e_y(), + rbf_vec_coeff_e=interpolation_savepoint.rbf_vec_coeff_e(), + e_bln_c_s=as_1D_sparse_field(interpolation_savepoint.e_bln_c_s(), CEDim), + rbf_coeff_1=interpolation_savepoint.rbf_vec_coeff_v1(), + rbf_coeff_2=interpolation_savepoint.rbf_vec_coeff_v2(), + geofac_div=as_1D_sparse_field(interpolation_savepoint.geofac_div(), CEDim), + geofac_n2s=interpolation_savepoint.geofac_n2s(), + geofac_grg_x=grg[0], + geofac_grg_y=grg[1], + nudgecoeff_e=interpolation_savepoint.nudgecoeff_e(), + ) + nonhydro_metric_state = MetricStateNonHydro( + bdy_halo_c=metrics_savepoint.bdy_halo_c(), + mask_prog_halo_c=metrics_savepoint.mask_prog_halo_c(), + rayleigh_w=metrics_savepoint.rayleigh_w(), + exner_exfac=metrics_savepoint.exner_exfac(), + exner_ref_mc=metrics_savepoint.exner_ref_mc(), + wgtfac_c=metrics_savepoint.wgtfac_c(), + wgtfacq_c=metrics_savepoint.wgtfacq_c_dsl(), + inv_ddqz_z_full=metrics_savepoint.inv_ddqz_z_full(), + rho_ref_mc=metrics_savepoint.rho_ref_mc(), + theta_ref_mc=metrics_savepoint.theta_ref_mc(), + vwind_expl_wgt=metrics_savepoint.vwind_expl_wgt(), + d_exner_dz_ref_ic=metrics_savepoint.d_exner_dz_ref_ic(), + ddqz_z_half=metrics_savepoint.ddqz_z_half(), + theta_ref_ic=metrics_savepoint.theta_ref_ic(), + d2dexdz2_fac1_mc=metrics_savepoint.d2dexdz2_fac1_mc(), + d2dexdz2_fac2_mc=metrics_savepoint.d2dexdz2_fac2_mc(), + rho_ref_me=metrics_savepoint.rho_ref_me(), + theta_ref_me=metrics_savepoint.theta_ref_me(), + ddxn_z_full=metrics_savepoint.ddxn_z_full(), + zdiff_gradp=metrics_savepoint.zdiff_gradp(), + vertoffset_gradp=metrics_savepoint.vertoffset_gradp(), + ipeidx_dsl=metrics_savepoint.ipeidx_dsl(), + pg_exdist=metrics_savepoint.pg_exdist(), + ddqz_z_full_e=metrics_savepoint.ddqz_z_full_e(), + ddxt_z_full=metrics_savepoint.ddxt_z_full(), + wgtfac_e=metrics_savepoint.wgtfac_e(), + wgtfacq_e=metrics_savepoint.wgtfacq_e_dsl(icon_grid.num_levels), + vwind_impl_wgt=metrics_savepoint.vwind_impl_wgt(), + hmask_dd3d=metrics_savepoint.hmask_dd3d(), + scalfac_dd3d=metrics_savepoint.scalfac_dd3d(), + coeff1_dwdz=metrics_savepoint.coeff1_dwdz(), + coeff2_dwdz=metrics_savepoint.coeff2_dwdz(), + coeff_gradekin=metrics_savepoint.coeff_gradekin(), ) - nonhydro_metric_state = metrics_savepoint.construct_nh_metric_state(icon_grid.num_levels) cell_geometry: CellParams = grid_savepoint.construct_cell_geometry() edge_geometry: EdgeParams = grid_savepoint.construct_edge_geometry() @@ -270,7 +324,6 @@ def test_run_timeloop_single_step( ) if debug_mode: - script_dir = os.path.dirname(os.path.abspath(__file__)) base_dir = script_dir + "/"