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/model/atmosphere/diffusion/diffusion_tests/utils.py b/model/atmosphere/diffusion/diffusion_tests/utils.py index b2fbf987d5..dcb9c24b23 100644 --- a/model/atmosphere/diffusion/diffusion_tests/utils.py +++ b/model/atmosphere/diffusion/diffusion_tests/utils.py @@ -13,10 +13,8 @@ import numpy as np -from icon4py.model.atmosphere.diffusion.diffusion_states import ( - DiffusionDiagnosticState, - 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 @@ -47,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 0e885aa12c..f08d7a8fd8 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -33,7 +33,6 @@ DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, - PrognosticState, ) from icon4py.model.atmosphere.diffusion.diffusion_utils import ( copy_field, @@ -78,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 @@ -538,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") @@ -837,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/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 65% 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..32e3dc2451 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 @@ -24,16 +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]: 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 caf3867923..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 @@ -27,7 +27,6 @@ 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 ( @@ -52,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 @@ -716,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 93ba9460b7..24d4f8162c 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -22,10 +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, - 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, @@ -34,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 ( @@ -123,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 6cf0823fb5..881d9bd597 100644 --- a/model/driver/src/icon4py/model/driver/io_utils.py +++ b/model/driver/src/icon4py/model/driver/io_utils.py @@ -20,13 +20,13 @@ DiffusionDiagnosticState, DiffusionInterpolationState, DiffusionMetricState, - 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..809982c41e 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 @@ -244,12 +244,14 @@ 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_pressure=exner, + exner=exner, theta_v=theta_v, ) if linit: