Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Small fixes #518

Merged
merged 1 commit into from
Aug 8, 2024
Merged

Small fixes #518

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 9 additions & 11 deletions model/common/src/icon4py/model/common/states/prognostic_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,8 @@
from dataclasses import dataclass

from gt4py.next import as_field
from gt4py.next.common import Field

from icon4py.model.common import field_type_aliases as fa
from icon4py.model.common.dimension import CellDim, EdgeDim, KDim
from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta


@dataclass
Expand All @@ -27,14 +25,14 @@ class PrognosticState:
Corresponds to ICON t_nh_prog
"""

rho: fa.CellKField[float] # density, rho(nproma, nlev, nblks_c) [kg/m^3]
w: fa.CellKField[float] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s]
vn: Field[
[EdgeDim, KDim], float
rho: fa.CellKField[ta.wpfloat] # density, rho(nproma, nlev, nblks_c) [kg/m^3]
w: fa.CellKField[ta.wpfloat] # vertical_wind field, w(nproma, nlevp1, nblks_c) [m/s]
vn: fa.EdgeKField[
ta.wpfloat
] # horizontal wind normal to edges, vn(nproma, nlev, nblks_e) [m/s]
exner: fa.CellKField[float] # exner function, exner(nrpoma, nlev, nblks_c)
theta_v: fa.CellKField[float] # virtual temperature, (nproma, nlev, nlbks_c) [K]
exner: fa.CellKField[ta.wpfloat] # exner function, exner(nrpoma, nlev, nblks_c)
theta_v: fa.CellKField[ta.wpfloat] # virtual temperature, (nproma, nlev, nlbks_c) [K]

@property
def w_1(self) -> fa.CellField[float]:
return as_field((CellDim,), self.w.ndarray[:, 0])
def w_1(self) -> fa.CellField[ta.wpfloat]:
return as_field((dims.CellDim,), self.w.ndarray[:, 0])
4 changes: 2 additions & 2 deletions model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1344,10 +1344,10 @@ def pressure_ifc(self):
def pressure_sfc(self):
return self._get_field("output_diag_pressure_sfc", dims.CellDim)

def zonal_Wind(self):
def zonal_wind(self):
return self._get_field("output_diag_u", dims.CellDim, dims.KDim)

def meridional_Wind(self):
def meridional_wind(self):
return self._get_field("output_diag_v", dims.CellDim, dims.KDim)


Expand Down
137 changes: 72 additions & 65 deletions model/common/tests/diagnostic_calculations_tests/test_diagnostic_calculations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,32 +14,25 @@

import pytest

from icon4py.model.common.constants import CPD_O_RD, GRAV_O_RD, P0REF
from icon4py.model.common.diagnostic_calculations.stencils.diagnose_pressure import (
diagnose_pressure,
import icon4py.model.common.grid.horizontal as h_grid
from icon4py.model.common import constants as constants, dimension as dims
from icon4py.model.common.diagnostic_calculations.stencils import (
diagnose_pressure as pressure,
diagnose_surface_pressure as surface_pressure,
diagnose_temperature as temperature,
)
from icon4py.model.common.diagnostic_calculations.stencils.diagnose_surface_pressure import (
diagnose_surface_pressure,
from icon4py.model.common.interpolation.stencils import (
edge_2_cell_vector_rbf_interpolation as rbf,
)
from icon4py.model.common.diagnostic_calculations.stencils.diagnose_temperature import (
diagnose_temperature,
)
from icon4py.model.common.dimension import CellDim, KDim
from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex
from icon4py.model.common.interpolation.stencils.edge_2_cell_vector_rbf_interpolation import (
edge_2_cell_vector_rbf_interpolation,
)
from icon4py.model.common.states.diagnostic_state import DiagnosticState
from icon4py.model.common.states.prognostic_state import PrognosticState
from icon4py.model.common.test_utils.datatest_utils import JABW_EXPERIMENT
from icon4py.model.common.test_utils.helpers import dallclose, zero_field
from icon4py.model.common.states import diagnostic_state as diagnostics, prognostic_state
from icon4py.model.common.test_utils import datatest_utils as dt_utils, helpers as helpers


@pytest.mark.datatest
@pytest.mark.parametrize(
"experiment",
[
JABW_EXPERIMENT,
dt_utils.JABW_EXPERIMENT,
],
)
def test_diagnose_temperature(
Expand All @@ -49,33 +42,37 @@ def test_diagnose_temperature(
):
sp = data_provider.from_savepoint_jabw_final()
icon_diagnostics_output_sp = data_provider.from_savepoint_jabw_diagnostic()
prognostic_state_now = PrognosticState(
prognostic_state_now = prognostic_state.PrognosticState(
rho=sp.rho(),
w=None,
vn=sp.vn(),
exner=sp.exner(),
theta_v=sp.theta_v(),
)
diagnostic_state = DiagnosticState(
temperature=zero_field(icon_grid, CellDim, KDim, dtype=float),
pressure=zero_field(icon_grid, CellDim, KDim, dtype=float),
pressure_ifc=zero_field(icon_grid, CellDim, KDim, dtype=float, extend={KDim: 1}),
u=zero_field(icon_grid, CellDim, KDim, dtype=float),
v=zero_field(icon_grid, CellDim, KDim, dtype=float),
diagnostic_state = diagnostics.DiagnosticState(
temperature=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
pressure=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
pressure_ifc=helpers.zero_field(
icon_grid, dims.CellDim, dims.KDim, dtype=float, extend={dims.KDim: 1}
),
u=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
v=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
)

diagnose_temperature(
temperature.diagnose_temperature(
prognostic_state_now.theta_v,
prognostic_state_now.exner,
diagnostic_state.temperature,
icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.interior(CellDim)),
icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.end(CellDim)),
icon_grid.get_start_index(
dims.CellDim, h_grid.HorizontalMarkerIndex.interior(dims.CellDim)
),
icon_grid.get_end_index(dims.CellDim, h_grid.HorizontalMarkerIndex.end(dims.CellDim)),
0,
icon_grid.num_levels,
offset_provider={},
)

assert dallclose(
assert helpers.dallclose(
diagnostic_state.temperature.asnumpy(),
icon_diagnostics_output_sp.temperature().asnumpy(),
)
Expand All @@ -85,7 +82,7 @@ def test_diagnose_temperature(
@pytest.mark.parametrize(
"experiment",
[
JABW_EXPERIMENT,
dt_utils.JABW_EXPERIMENT,
],
)
def test_diagnose_meridional_and_zonal_winds(
Expand All @@ -95,28 +92,32 @@ def test_diagnose_meridional_and_zonal_winds(
):
sp = data_provider.from_savepoint_jabw_final()
icon_diagnostics_output_sp = data_provider.from_savepoint_jabw_diagnostic()
prognostic_state_now = PrognosticState(
prognostic_state_now = prognostic_state.PrognosticState(
rho=sp.rho(),
w=None,
vn=sp.vn(),
exner=sp.exner(),
theta_v=sp.theta_v(),
)
diagnostic_state = DiagnosticState(
temperature=zero_field(icon_grid, CellDim, KDim, dtype=float),
pressure=zero_field(icon_grid, CellDim, KDim, dtype=float),
pressure_ifc=zero_field(icon_grid, CellDim, KDim, dtype=float, extend={KDim: 1}),
u=zero_field(icon_grid, CellDim, KDim, dtype=float),
v=zero_field(icon_grid, CellDim, KDim, dtype=float),
diagnostic_state = diagnostics.DiagnosticState(
temperature=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
pressure=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
pressure_ifc=helpers.zero_field(
icon_grid, dims.CellDim, dims.KDim, dtype=float, extend={dims.KDim: 1}
),
u=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
v=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
)

rbv_vec_coeff_c1 = data_provider.from_interpolation_savepoint().rbf_vec_coeff_c1()
rbv_vec_coeff_c2 = data_provider.from_interpolation_savepoint().rbf_vec_coeff_c2()
grid_idx_cell_start_plus1 = icon_grid.get_end_index(
CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) + 1
dims.CellDim, h_grid.HorizontalMarkerIndex.lateral_boundary(dims.CellDim) + 1
)
grid_idx_cell_end = icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.end(CellDim))
edge_2_cell_vector_rbf_interpolation(
grid_idx_cell_end = icon_grid.get_end_index(
dims.CellDim, h_grid.HorizontalMarkerIndex.end(dims.CellDim)
)
rbf.edge_2_cell_vector_rbf_interpolation(
prognostic_state_now.vn,
rbv_vec_coeff_c1,
rbv_vec_coeff_c2,
Expand All @@ -131,14 +132,14 @@ def test_diagnose_meridional_and_zonal_winds(
},
)

assert dallclose(
assert helpers.dallclose(
diagnostic_state.u.asnumpy(),
icon_diagnostics_output_sp.zonal_Wind().asnumpy(),
icon_diagnostics_output_sp.zonal_wind().asnumpy(),
)

assert dallclose(
assert helpers.dallclose(
diagnostic_state.v.asnumpy(),
icon_diagnostics_output_sp.meridional_Wind().asnumpy(),
icon_diagnostics_output_sp.meridional_wind().asnumpy(),
atol=1.0e-13,
)

Expand All @@ -147,7 +148,7 @@ def test_diagnose_meridional_and_zonal_winds(
@pytest.mark.parametrize(
"experiment",
[
JABW_EXPERIMENT,
dt_utils.JABW_EXPERIMENT,
],
)
def test_diagnose_pressure(
Expand All @@ -157,63 +158,69 @@ def test_diagnose_pressure(
):
sp = data_provider.from_savepoint_jabw_final()
icon_diagnostics_output_sp = data_provider.from_savepoint_jabw_diagnostic()
prognostic_state_now = PrognosticState(
prognostic_state_now = prognostic_state.PrognosticState(
rho=sp.rho(),
w=None,
vn=sp.vn(),
exner=sp.exner(),
theta_v=sp.theta_v(),
)
diagnostic_state = DiagnosticState(
diagnostic_state = diagnostics.DiagnosticState(
temperature=sp.temperature(),
pressure=zero_field(icon_grid, CellDim, KDim, dtype=float),
pressure_ifc=zero_field(icon_grid, CellDim, KDim, dtype=float, extend={KDim: 1}),
u=zero_field(icon_grid, CellDim, KDim, dtype=float),
v=zero_field(icon_grid, CellDim, KDim, dtype=float),
pressure=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
pressure_ifc=helpers.zero_field(
icon_grid, dims.CellDim, dims.KDim, dtype=float, extend={dims.KDim: 1}
),
u=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
v=helpers.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
)

diagnose_surface_pressure(
surface_pressure.diagnose_surface_pressure(
prognostic_state_now.exner,
diagnostic_state.temperature,
data_provider.from_metrics_savepoint().ddqz_z_full(),
diagnostic_state.pressure_ifc,
CPD_O_RD,
P0REF,
GRAV_O_RD,
constants.CPD_O_RD,
constants.P0REF,
constants.GRAV_O_RD,
horizontal_start=icon_grid.get_start_index(
CellDim, HorizontalMarkerIndex.interior(CellDim)
dims.CellDim, h_grid.HorizontalMarkerIndex.interior(dims.CellDim)
),
horizontal_end=icon_grid.get_end_index(
dims.CellDim, h_grid.HorizontalMarkerIndex.end(dims.CellDim)
),
horizontal_end=icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.end(CellDim)),
vertical_start=icon_grid.num_levels,
vertical_end=icon_grid.num_levels + 1,
offset_provider={"Koff": KDim},
offset_provider={"Koff": dims.KDim},
)

diagnose_pressure(
pressure.diagnose_pressure(
data_provider.from_metrics_savepoint().ddqz_z_full(),
diagnostic_state.temperature,
diagnostic_state.pressure_sfc,
diagnostic_state.pressure,
diagnostic_state.pressure_ifc,
GRAV_O_RD,
icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.interior(CellDim)),
icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.end(CellDim)),
constants.GRAV_O_RD,
icon_grid.get_start_index(
dims.CellDim, h_grid.HorizontalMarkerIndex.interior(dims.CellDim)
),
icon_grid.get_end_index(dims.CellDim, h_grid.HorizontalMarkerIndex.end(dims.CellDim)),
0,
icon_grid.num_levels,
offset_provider={},
)

assert dallclose(
assert helpers.dallclose(
diagnostic_state.pressure_sfc.asnumpy(),
icon_diagnostics_output_sp.pressure_sfc().asnumpy(),
)

assert dallclose(
assert helpers.dallclose(
diagnostic_state.pressure_ifc.asnumpy(),
icon_diagnostics_output_sp.pressure_ifc().asnumpy(),
)

assert dallclose(
assert helpers.dallclose(
diagnostic_state.pressure.asnumpy(),
icon_diagnostics_output_sp.pressure().asnumpy(),
)
Loading