From ed510b17871216cd1bac0ab2fe04f728a9e06015 Mon Sep 17 00:00:00 2001 From: Chia Rui Ong Date: Mon, 9 Oct 2023 15:52:34 +0200 Subject: [PATCH 1/4] Moving prognostics to common --- .../diffusion/diffusion_tests/utils.py | 2 ++ .../model/atmosphere/diffusion/diffusion.py | 4 ++- .../common/prognostics/core_prognostics.py | 33 +++++++++++++++++++ .../common/test_utils/serialbox_utils.py | 2 ++ .../src/icon4py/model/driver/dycore_driver.py | 2 ++ .../src/icon4py/model/driver/io_utils.py | 2 ++ 6 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 model/common/src/icon4py/model/common/prognostics/core_prognostics.py diff --git a/model/atmosphere/diffusion/diffusion_tests/utils.py b/model/atmosphere/diffusion/diffusion_tests/utils.py index b2fbf987d5..73a25e36c7 100644 --- a/model/atmosphere/diffusion/diffusion_tests/utils.py +++ b/model/atmosphere/diffusion/diffusion_tests/utils.py @@ -15,6 +15,8 @@ from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, +) +from icon4py.model.common.prognostics.core_prognostics import ( PrognosticState, ) from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionExitSavepoint 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 0e885aa12c..1123295c6f 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -32,7 +32,9 @@ from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, DiffusionInterpolationState, - DiffusionMetricState, + DiffusionMetricState +) +from icon4py.model.common.prognostics.core_prognostics import ( PrognosticState, ) from icon4py.model.atmosphere.diffusion.diffusion_utils import ( diff --git a/model/common/src/icon4py/model/common/prognostics/core_prognostics.py b/model/common/src/icon4py/model/common/prognostics/core_prognostics.py new file mode 100644 index 0000000000..4d8d200e2e --- /dev/null +++ b/model/common/src/icon4py/model/common/prognostics/core_prognostics.py @@ -0,0 +1,33 @@ +# 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 PrognosticState: + """Class that contains the prognostic state. + + Corresponds to ICON t_nh_prog + """ + + w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] + vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] + exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) + theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] 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 3d1165d0e7..c4e51cdf52 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 @@ -23,6 +23,8 @@ DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, +) +from icon4py.model.common.prognostics.core_prognostics import ( PrognosticState, ) from icon4py.model.common import dimension diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py index a1b5b492e6..d1e0b5b6bb 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -24,6 +24,8 @@ from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, +) +from icon4py.model.common.prognostics.core_prognostics import ( PrognosticState, ) from icon4py.model.atmosphere.diffusion.diffusion_utils import _identity_c_k, _identity_e_k diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py index 6cf0823fb5..57b77c2279 100644 --- a/model/driver/src/icon4py/model/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -20,6 +20,8 @@ DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, +) +from icon4py.model.common.prognostics.core_prognostics import ( PrognosticState, ) from icon4py.model.common.decomposition.definitions import DecompositionInfo, ProcessProperties From 035cfdbe1435ea58147c7d515fea37aea3ed495d Mon Sep 17 00:00:00 2001 From: Chia Rui Ong Date: Tue, 10 Oct 2023 13:51:52 +0200 Subject: [PATCH 2/4] Move prognostic state module to common/states --- .../diffusion/diffusion_tests/utils.py | 10 ++---- .../model/atmosphere/diffusion/diffusion.py | 10 +++--- .../atmosphere/diffusion/diffusion_states.py | 13 -------- .../dycore/nh_solve/solve_nonydro.py | 2 +- .../dycore/velocity/velocity_advection.py | 2 +- .../dycore/tests/test_solve_nonhydro.py | 10 +----- .../dycore/tests/test_velocity_advection.py | 4 +-- .../common/prognostics/core_prognostics.py | 33 ------------------- .../model/common/states}/prognostic_state.py | 17 +++++----- .../common/test_utils/serialbox_utils.py | 10 ++---- .../src/icon4py/model/driver/dycore_driver.py | 12 +++---- .../src/icon4py/model/driver/io_utils.py | 4 +-- .../py2f/wrappers/diffusion_wrapper.py | 4 +-- 13 files changed, 29 insertions(+), 102 deletions(-) delete mode 100644 model/common/src/icon4py/model/common/prognostics/core_prognostics.py rename model/{atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils => common/src/icon4py/model/common/states}/prognostic_state.py (66%) diff --git a/model/atmosphere/diffusion/diffusion_tests/utils.py b/model/atmosphere/diffusion/diffusion_tests/utils.py index 73a25e36c7..dcb9c24b23 100644 --- a/model/atmosphere/diffusion/diffusion_tests/utils.py +++ b/model/atmosphere/diffusion/diffusion_tests/utils.py @@ -13,12 +13,8 @@ import numpy as np -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, -) -from icon4py.model.common.prognostics.core_prognostics import ( - PrognosticState, -) +from icon4py.model.atmosphere.diffusion.diffusion_states import DiffusionDiagnosticState +from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.serialbox_utils import IconDiffusionExitSavepoint @@ -49,7 +45,7 @@ def verify_diffusion_fields( ref_exner = np.asarray(diffusion_savepoint.exner()) ref_theta_v = np.asarray(diffusion_savepoint.theta_v()) val_theta_v = np.asarray(prognostic_state.theta_v) - val_exner = np.asarray(prognostic_state.exner_pressure) + val_exner = np.asarray(prognostic_state.exner) assert np.allclose(ref_theta_v, val_theta_v) assert np.allclose(ref_exner, val_exner) 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 1123295c6f..f08d7a8fd8 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -32,10 +32,7 @@ from icon4py.model.atmosphere.diffusion.diffusion_states import ( DiffusionDiagnosticState, DiffusionInterpolationState, - DiffusionMetricState -) -from icon4py.model.common.prognostics.core_prognostics import ( - PrognosticState, + DiffusionMetricState, ) from icon4py.model.atmosphere.diffusion.diffusion_utils import ( copy_field, @@ -80,6 +77,7 @@ from icon4py.model.common.interpolation.stencils.mo_intp_rbf_rbf_vec_interpol_vertex import ( mo_intp_rbf_rbf_vec_interpol_vertex, ) +from icon4py.model.common.states.prognostic_state import PrognosticState # flake8: noqa @@ -540,7 +538,7 @@ def _sync_cell_fields(self, prognostic_state): CellDim, prognostic_state.w, prognostic_state.theta_v, - prognostic_state.exner_pressure, + prognostic_state.exner, ) handle_cell_comm.wait() log.debug("communication of prognostic cell fields: theta, w, exner - done") @@ -839,7 +837,7 @@ def _do_diffusion_step( z_temp=self.z_temp, area=self.cell_params.area, theta_v=prognostic_state.theta_v, - exner=prognostic_state.exner_pressure, + exner=prognostic_state.exner, rd_o_cvd=self.rd_o_cvd, horizontal_start=cell_start_nudging, horizontal_end=cell_end_local, diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py index ea423047cb..7d1bd5b4cc 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_states.py @@ -100,16 +100,3 @@ def geofac_n2s_nbh(self) -> Field[[CECDim], float]: old_shape[0] * old_shape[1], ) ) - - -@dataclass -class PrognosticState: - """Class that contains the prognostic state. - - Corresponds to ICON t_nh_prog - """ - - w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] - vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] - exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) - theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonydro.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonydro.py index 4ad10e9f26..529a314a8e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonydro.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonydro.py @@ -142,7 +142,6 @@ 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.prognostic_state import PrognosticState from icon4py.model.atmosphere.dycore.state_utils.utils import ( _allocate, _allocate_indices, @@ -159,6 +158,7 @@ from icon4py.model.common.grid.horizontal import EdgeParams, HorizontalMarkerIndex from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.states.prognostic_state import PrognosticState class NonHydrostaticConfig: 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 0d1ff674e6..29bb1d87a0 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 @@ -59,12 +59,12 @@ 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.prognostic_state import PrognosticState 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 HorizontalMarkerIndex from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.states.prognostic_state import PrognosticState cached_backend = run_gtfn_cached diff --git a/model/atmosphere/dycore/tests/test_solve_nonhydro.py b/model/atmosphere/dycore/tests/test_solve_nonhydro.py index a373fa67b3..8f2b9fab96 100644 --- a/model/atmosphere/dycore/tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/test_solve_nonhydro.py @@ -21,12 +21,12 @@ 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.prognostic_state import PrognosticState 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.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.helpers import dallclose, random_field, zero_field from icon4py.model.common.test_utils.simple_mesh import SimpleMesh @@ -131,7 +131,6 @@ def test_nonhydro_predictor_step( prognostic_state_nnow = PrognosticState( w=sp.w_now(), vn=sp.vn_now(), - exner_pressure=None, theta_v=sp.theta_v_now(), rho=sp.rho_now(), exner=sp.exner_now(), @@ -140,7 +139,6 @@ def test_nonhydro_predictor_step( prognostic_state_nnew = PrognosticState( w=sp.w_new(), vn=sp.vn_new(), - exner_pressure=None, theta_v=sp.theta_v_new(), rho=sp.rho_new(), exner=sp.exner_new(), @@ -526,7 +524,6 @@ def test_nonhydro_corrector_step( prognostic_state_nnow = PrognosticState( w=sp.w_now(), vn=sp.vn_now(), - exner_pressure=None, theta_v=sp.theta_v_now(), rho=sp.rho_now(), exner=sp.exner_now(), @@ -535,7 +532,6 @@ def test_nonhydro_corrector_step( prognostic_state_nnew = PrognosticState( w=sp.w_new(), vn=sp.vn_new(), - exner_pressure=None, theta_v=sp.theta_v_new(), rho=sp.rho_new(), exner=sp.exner_new(), @@ -750,7 +746,6 @@ def test_run_solve_nonhydro_single_step( prognostic_state_nnow = PrognosticState( w=sp.w_now(), vn=sp.vn_now(), - exner_pressure=None, theta_v=sp.theta_v_now(), rho=sp.rho_now(), exner=sp.exner_now(), @@ -759,7 +754,6 @@ def test_run_solve_nonhydro_single_step( prognostic_state_nnew = PrognosticState( w=sp.w_new(), vn=sp.vn_new(), - exner_pressure=None, theta_v=sp.theta_v_new(), rho=sp.rho_new(), exner=sp.exner_new(), @@ -927,7 +921,6 @@ def test_run_solve_nonhydro_multi_step( prognostic_state_nnow = PrognosticState( w=sp.w_now(), vn=sp.vn_now(), - exner_pressure=None, theta_v=sp.theta_v_now(), rho=sp.rho_now(), exner=sp.exner_now(), @@ -936,7 +929,6 @@ def test_run_solve_nonhydro_multi_step( prognostic_state_nnew = PrognosticState( w=sp.w_new(), vn=sp.vn_new(), - exner_pressure=None, theta_v=sp.theta_v_new(), rho=sp.rho_new(), exner=sp.exner_new(), diff --git a/model/atmosphere/dycore/tests/test_velocity_advection.py b/model/atmosphere/dycore/tests/test_velocity_advection.py index 78d33b53ef..323557335a 100644 --- a/model/atmosphere/dycore/tests/test_velocity_advection.py +++ b/model/atmosphere/dycore/tests/test_velocity_advection.py @@ -16,10 +16,10 @@ 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.prognostic_state import PrognosticState 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 @@ -148,7 +148,6 @@ def test_velocity_predictor_step( prognostic_state = PrognosticState( w=sp_v.w(), vn=sp_v.vn(), - exner_pressure=None, theta_v=None, rho=None, exner=None, @@ -310,7 +309,6 @@ def test_velocity_corrector_step( prognostic_state = PrognosticState( w=sp_v.w(), vn=sp_v.vn(), - exner_pressure=None, theta_v=None, rho=None, exner=None, diff --git a/model/common/src/icon4py/model/common/prognostics/core_prognostics.py b/model/common/src/icon4py/model/common/prognostics/core_prognostics.py deleted file mode 100644 index 4d8d200e2e..0000000000 --- a/model/common/src/icon4py/model/common/prognostics/core_prognostics.py +++ /dev/null @@ -1,33 +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 PrognosticState: - """Class that contains the prognostic state. - - Corresponds to ICON t_nh_prog - """ - - w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] - vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] - exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) - theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/prognostic_state.py b/model/common/src/icon4py/model/common/states/prognostic_state.py similarity index 66% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/prognostic_state.py rename to model/common/src/icon4py/model/common/states/prognostic_state.py index b64b91031f..d8a5d91852 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/state_utils/prognostic_state.py +++ b/model/common/src/icon4py/model/common/states/prognostic_state.py @@ -11,9 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import numpy as np from dataclasses import dataclass -import numpy as np from gt4py.next.common import Field from gt4py.next.iterator.embedded import np_as_located_field @@ -24,17 +24,16 @@ class PrognosticState: """Class that contains the prognostic state. - corresponds to ICON t_nh_prog + Corresponds to ICON t_nh_prog """ - w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] - vn: Field[[EdgeDim, KDim], float] # vn(nproma, nlev, nblks_e) [m/s] - exner_pressure: Field[[CellDim, KDim], float] # exner(nrpoma, nlev, nblks_c) - theta_v: Field[[CellDim, KDim], float] # (nproma, nlev, nlbks_c) [K] - - rho: Field[[CellDim, KDim], float] - exner: Field[[CellDim, KDim], float] + rho: Field[[CellDim, KDim], float] # density, rho(nproma, nlev, nblks_c) [m/s] + w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] + vn: Field[[EdgeDim, KDim], float] # horizontal wind normal to edges, vn(nproma, nlev, nblks_e) [m/s] + exner: Field[[CellDim, KDim], float] # exner function, exner(nrpoma, nlev, nblks_c) + theta_v: Field[[CellDim, KDim], float] # virtual temperature, (nproma, nlev, nlbks_c) [K] @property def w_1(self) -> Field[[CellDim], float]: return np_as_located_field(CellDim)(np.asarray(self.w)[:, 0]) + 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 6f8862e9e1..f7a5a6cd97 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,13 +24,9 @@ DiffusionInterpolationState, DiffusionMetricState, ) -from icon4py.model.common.prognostics.core_prognostics import ( - PrognosticState, -) 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.atmosphere.dycore.state_utils.prognostic_state import PrognosticState from icon4py.model.common import dimension from icon4py.model.common.decomposition.definitions import DecompositionInfo from icon4py.model.common.dimension import ( @@ -55,6 +51,7 @@ ) from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid, VerticalGridSize +from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, flatten_first_two_dims @@ -719,10 +716,9 @@ def construct_prognostics(self) -> PrognosticState: return PrognosticState( w=self.w(), vn=self.vn(), - exner_pressure=self.exner(), + exner=self.exner(), theta_v=self.theta_v(), - rho=None, - exner=None, + rho=self.rho(), ) def construct_diagnostics_for_diffusion(self) -> DiffusionDiagnosticState: diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py index eef6de8360..24d4f8162c 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -22,12 +22,7 @@ from gt4py.next.program_processors.runners.gtfn_cpu import run_gtfn from icon4py.model.atmosphere.diffusion.diffusion import Diffusion, DiffusionParams -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, -) -from icon4py.model.common.prognostics.core_prognostics import ( - PrognosticState, -) +from icon4py.model.atmosphere.diffusion.diffusion_states import DiffusionDiagnosticState from icon4py.model.atmosphere.diffusion.diffusion_utils import _identity_c_k, _identity_e_k from icon4py.model.common.decomposition.definitions import ( ProcessProperties, @@ -36,6 +31,7 @@ get_runtype, ) from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils import serialbox_utils as sb from icon4py.model.driver.icon_configuration import IconRunConfig, read_config from icon4py.model.driver.io_utils import ( @@ -125,8 +121,8 @@ def do_dynamics_substepping( prognostic_state.vn, new_p.w, prognostic_state.w, - new_p.exner_pressure, - prognostic_state.exner_pressure, + new_p.exner, + prognostic_state.exner, new_p.theta_v, prognostic_state.theta_v, offset_provider={}, diff --git a/model/driver/src/icon4py/model/driver/io_utils.py b/model/driver/src/icon4py/model/driver/io_utils.py index 57b77c2279..881d9bd597 100644 --- a/model/driver/src/icon4py/model/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -21,14 +21,12 @@ DiffusionInterpolationState, DiffusionMetricState, ) -from icon4py.model.common.prognostics.core_prognostics import ( - PrognosticState, -) from icon4py.model.common.decomposition.definitions import DecompositionInfo, ProcessProperties from icon4py.model.common.decomposition.mpi_decomposition import ParallelLogger from icon4py.model.common.grid.horizontal import CellParams, EdgeParams from icon4py.model.common.grid.icon_grid import IconGrid from icon4py.model.common.grid.vertical import VerticalModelParams +from icon4py.model.common.states.prognostic_state import PrognosticState from icon4py.model.common.test_utils import serialbox_utils as sb diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py index d1b4c24243..98a6ae70d3 100644 --- a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -35,7 +35,6 @@ DiffusionMetricState, DiffusionParams, DiffusionType, - PrognosticState, ) from icon4py.model.common.dimension import ( C2E2CDim, @@ -55,6 +54,7 @@ from icon4py.model.common.grid.horizontal import CellParams, EdgeParams, HorizontalGridSize from icon4py.model.common.grid.icon_grid import GridConfig, IconGrid from icon4py.model.common.grid.vertical import VerticalGridSize, VerticalModelParams +from icon4py.model.common.states.prognostic_state import PrognosticState from icon4pytools.py2f.cffi_utils import CffiMethod, to_fields @@ -249,7 +249,7 @@ def diffusion_run( prognostic_state = PrognosticState( w=w, vn=vn, - exner_pressure=exner, + exner=exner, theta_v=theta_v, ) if linit: From 1d0baeb466ddd871af2c5619a573238e9cd6c890 Mon Sep 17 00:00:00 2001 From: Chia Rui Ong Date: Tue, 10 Oct 2023 14:56:50 +0200 Subject: [PATCH 3/4] Modified diffusion wrapper and test_verify_diffusion_init_against_other_regular_savepoint in test_diffusion.py --- model/atmosphere/diffusion/diffusion_tests/test_diffusion.py | 2 +- tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py index 9069507a8a..da8c5110b3 100644 --- a/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py +++ b/model/atmosphere/diffusion/diffusion_tests/test_diffusion.py @@ -219,7 +219,7 @@ def test_verify_diffusion_init_against_first_regular_savepoint( @pytest.mark.datatest -@pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:50.000"]) +@pytest.mark.parametrize("step_date_init", ["2021-06-20T12:00:20.000"]) def test_verify_diffusion_init_against_other_regular_savepoint( r04b09_diffusion_config, grid_savepoint, diff --git a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py index 98a6ae70d3..809982c41e 100644 --- a/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py +++ b/tools/src/icon4pytools/py2f/wrappers/diffusion_wrapper.py @@ -244,9 +244,11 @@ def diffusion_run( hdef_ic: Field[[CellDim, KDim], float], dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], + rho: Field[[CellDim, KDim], float], ): diagnostic_state = DiffusionDiagnosticState(hdef_ic, div_ic, dwdx, dwdy) prognostic_state = PrognosticState( + rho=rho, w=w, vn=vn, exner=exner, From 4c9e3404d3ef9a83f7ff971130182211eb0f307b Mon Sep 17 00:00:00 2001 From: Chia Rui Ong Date: Tue, 10 Oct 2023 15:36:19 +0200 Subject: [PATCH 4/4] Modified the order of imported modules in prognostic_state.py --- .../src/icon4py/model/common/states/prognostic_state.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/model/common/src/icon4py/model/common/states/prognostic_state.py b/model/common/src/icon4py/model/common/states/prognostic_state.py index d8a5d91852..32e3dc2451 100644 --- a/model/common/src/icon4py/model/common/states/prognostic_state.py +++ b/model/common/src/icon4py/model/common/states/prognostic_state.py @@ -11,9 +11,9 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import numpy as np from dataclasses import dataclass +import numpy as np from gt4py.next.common import Field from gt4py.next.iterator.embedded import np_as_located_field @@ -27,13 +27,14 @@ class PrognosticState: Corresponds to ICON t_nh_prog """ - rho: Field[[CellDim, KDim], float] # density, rho(nproma, nlev, nblks_c) [m/s] + rho: Field[[CellDim, KDim], float] # density, rho(nproma, nlev, nblks_c) [m/s] w: Field[[CellDim, KDim], float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s] - vn: Field[[EdgeDim, KDim], float] # horizontal wind normal to edges, vn(nproma, nlev, nblks_e) [m/s] + vn: Field[ + [EdgeDim, KDim], float + ] # horizontal wind normal to edges, vn(nproma, nlev, nblks_e) [m/s] exner: Field[[CellDim, KDim], float] # exner function, exner(nrpoma, nlev, nblks_c) theta_v: Field[[CellDim, KDim], float] # virtual temperature, (nproma, nlev, nlbks_c) [K] @property def w_1(self) -> Field[[CellDim], float]: return np_as_located_field(CellDim)(np.asarray(self.w)[:, 0]) -