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.

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

Verify global model #315

Merged
merged 53 commits into from
Dec 15, 2023
Merged
Changes from 1 commit
Commits
Show all changes
53 commits
Select commit Hold shift + click to select a range
5e7af70
remove unusde fields from metrics savepoint
halungge Nov 9, 2023
c3acfd8
remove dependency diffusion in common/test_utils
halungge Nov 10, 2023
853f5c9
use system epsilon
halungge Nov 10, 2023
8afbb4e
remove unused fields from serialbox_utils.py
halungge Nov 10, 2023
0237b20
fix imports in io_utils.py
halungge Nov 10, 2023
1b8271b
remove unused typed definition
halungge Nov 10, 2023
e0e8398
fix wrong import in test_diffusion_states.py
halungge Nov 10, 2023
acfc1cc
Merge branch 'main' into verify_global_model
halungge Nov 15, 2023
050cbc9
use dallclose in diffusion
halungge Nov 15, 2023
8516c51
diffusion: make choice of configuration flexible in tests
halungge Nov 15, 2023
85583a9
- pass experiment name a argument to datapath
halungge Nov 16, 2023
877f199
WIP diffusion single_step test for exclaim APE
halungge Nov 16, 2023
d745e5d
WIP (1) make velocity advection and solve nonhdyro tests parametrizab…
halungge Nov 17, 2023
e2d3583
WIP (2) make velocity advection and solve nonhdyro tests parametrizab…
halungge Nov 17, 2023
fcbf9ca
Merge branch 'main' into verify_global_model
halungge Nov 17, 2023
d7bd316
WIP (1) make solve_nonhydro.py choose configuration dynamically
halungge Nov 17, 2023
cf5a1a0
run test_solve_nonhydro.py::test_run_nonhydro_single_step for exclaim…
halungge Nov 20, 2023
c23d7e5
add experiment fixture to common and driver
halungge Nov 20, 2023
0b14265
run on specified backend for test_diffusion_utils.py
halungge Nov 21, 2023
fdbfe63
enable both experiments in test_solve_nonhydro.py
halungge Nov 21, 2023
085ad42
WIP: fix assert in test_diffusion.py
halungge Nov 22, 2023
f0dfaa1
merge main
halungge Nov 28, 2023
a4206ca
test_velocity_advection.py: for both experiments
halungge Nov 29, 2023
ac21490
test_diffusion.py: for both experiments
halungge Nov 29, 2023
e0f4298
fix tests: test_diffusion_utils.py, test_vertical.py
halungge Nov 29, 2023
c3df5aa
test_solve_nonhydro.py: invert args in dallclose, reference should be…
halungge Nov 29, 2023
d1efb33
add missing config parameters in Diffusion config
halungge Nov 30, 2023
eed4bc3
fix imports
halungge Nov 30, 2023
68c8b77
- make data loading fixture dependent
halungge Nov 30, 2023
aed9489
- add logging to solve_nonhydro.py
halungge Nov 30, 2023
da7b2d2
remove wrong and unused grid.nflat_gradp
halungge Nov 30, 2023
8eacbb9
set grid.limited_area according to experiment
halungge Nov 30, 2023
d1eb281
assert bounds in test
halungge Nov 30, 2023
7e49286
pre-commit fixes
halungge Dec 1, 2023
ac8fea5
fix typo in stencil name
halungge Dec 1, 2023
b108aa5
fix typo in stencil name
halungge Dec 1, 2023
767d64d
diffusion_tests:
halungge Dec 1, 2023
e36f8ae
add slow_tests marker to pytest_config.py
halungge Dec 1, 2023
64c001d
get limited_area from grid savepoint metadata
halungge Dec 1, 2023
3a43679
add rtol to diffusion asserts
halungge Dec 1, 2023
d09dfc3
merge from main
halungge Dec 1, 2023
1a9243b
test_velocity_advection.py: fix order of reference and values in dall…
halungge Dec 1, 2023
6ac6d0e
merge main
halungge Dec 5, 2023
450761e
add TODOs for released tolerances
halungge Dec 5, 2023
93c6a64
pre-commit fixes
halungge Dec 5, 2023
ef8be6d
cleanups
halungge Dec 5, 2023
b54f476
fix import of experiment fixture in model/common
halungge Dec 5, 2023
b96b7b4
fix nsubsteps in multistep test
halungge Dec 5, 2023
3149f36
pre-commit
halungge Dec 5, 2023
c720607
merge main
halungge Dec 7, 2023
23e78e4
merge main, resolve conflicts
halungge Dec 15, 2023
7425855
fix merge errors:
halungge Dec 15, 2023
f5fd1e7
fix merge errors: test_velocity_advection.py
halungge Dec 15, 2023
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
Prev Previous commit
Next Next commit
WIP diffusion single_step test for exclaim APE
halungge committed Nov 16, 2023
commit 877f199bde433fa8800627c1259f616c70634a73
Original file line number Diff line number Diff line change
@@ -97,6 +97,7 @@
class DiffusionType(int, Enum):
"""
Order of nabla operator for diffusion.

Note: Called `hdiff_order` in `mo_diffusion_nml.f90`.
Note: We currently only support type 5.
"""
@@ -108,6 +109,23 @@ class DiffusionType(int, Enum):
SMAGORINSKY_4TH_ORDER = 5 #: Smagorinsky diffusion with fourth-order background diffusion


class TurbulenceShearForcingType(int, Enum):
"""
Type of shear forcing used in turbulance.

Note: called `itype_sher` in `mo_turbdiff_nml.f90`
"""

VERTICAL_OF_HORIZONTAL_WIND = 0 #: only vertical shear of horizontal wind
VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND = (
1 #: as `VERTICAL_ONLY` plus horizontal shar correction
)
VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND = (
2 #: as `VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND` plus shear form vertical velocity
)
VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND_LTHESH = 3 #: same as `VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND` but scaling of coarse-grid horizontal shear production term with 1/sqrt(Ri) (if LTKESH = TRUE)


class DiffusionConfig:
"""
Contains necessary parameter to configure a diffusion run.
@@ -139,6 +157,7 @@ def __init__(
temperature_boundary_diffusion_denom: float = 135.0,
max_nudging_coeff: float = 0.02,
nudging_decay_rate: float = 2.0,
shear_type: TurbulenceShearForcingType = TurbulenceShearForcingType.VERTICAL_OF_HORIZONTAL_WIND,
):
"""Set the diffusion configuration parameters with the ICON default values."""
# parameters from namelist diffusion_nml
@@ -221,6 +240,10 @@ def __init__(
#: Called `nudge_efold_width` in mo_interpol_nml.f90
self.nudge_efold_width: float = nudging_decay_rate

#: Type of shear forcing used in turbulence
#: Called itype_shear in `mo_turbdiff_nml.f90
self.shear_type = shear_type

self._validate()

def _validate(self):
@@ -238,9 +261,8 @@ def _validate(self):
else:
self.apply_to_temperature = True
self.apply_to_horizontal_wind = True
# TODO(magdalena) remove: is not set for APE experiment
if not self.apply_zdiffusion_t:
raise NotImplementedError("zdiffu_t = False is not implemented (leaves out stencil_15)")

# TODO validate shear_type

@functools.cached_property
def substep_as_float(self):
@@ -620,7 +642,7 @@ def _do_diffusion_step(

# 2. HALO EXCHANGE -- CALL sync_patch_array_mult u_vert and v_vert
log.debug("communication rbf extrapolation of vn - start")
h = self._exchange.exchange_and_wait(VertexDim, self.u_vert, self.v_vert)
self._exchange.exchange_and_wait(VertexDim, self.u_vert, self.v_vert)
log.debug("communication rbf extrapolation of vn - end")

log.debug("running stencil 01(calculate_nabla2_and_smag_coefficients_for_vn): start")
@@ -651,30 +673,38 @@ def _do_diffusion_step(
},
)
log.debug("running stencil 01 (calculate_nabla2_and_smag_coefficients_for_vn): end")
log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start")
calculate_diagnostic_quantities_for_turbulence.with_backend(backend)(
kh_smag_ec=self.kh_smag_ec,
vn=prognostic_state.vn,
e_bln_c_s=self.interpolation_state.e_bln_c_s,
geofac_div=self.interpolation_state.geofac_div,
diff_multfac_smag=self.diff_multfac_smag,
wgtfac_c=self.metric_state.wgtfac_c,
div_ic=diagnostic_state.div_ic,
hdef_ic=diagnostic_state.hdef_ic,
horizontal_start=cell_start_nudging,
horizontal_end=cell_end_local,
vertical_start=1,
vertical_end=klevels,
offset_provider={
"C2E": self.grid.get_offset_provider("C2E"),
"C2CE": self.grid.get_offset_provider("C2CE"),
"Koff": KDim,
},
)
log.debug("running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end")

# HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array -> false for MCH
if (
self.config.shear_type
>= TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_WIND
): # TODO or ltkesh
log.debug(
"running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): start"
)
calculate_diagnostic_quantities_for_turbulence.with_backend(backend)(
kh_smag_ec=self.kh_smag_ec,
vn=prognostic_state.vn,
e_bln_c_s=self.interpolation_state.e_bln_c_s,
geofac_div=self.interpolation_state.geofac_div,
diff_multfac_smag=self.diff_multfac_smag,
wgtfac_c=self.metric_state.wgtfac_c,
div_ic=diagnostic_state.div_ic,
hdef_ic=diagnostic_state.hdef_ic,
horizontal_start=cell_start_nudging,
horizontal_end=cell_end_local,
vertical_start=1,
vertical_end=klevels,
offset_provider={
"C2E": self.grid.get_offset_provider("C2E"),
"C2CE": self.grid.get_offset_provider("C2CE"),
"Koff": KDim,
},
)
log.debug(
"running stencils 02 03 (calculate_diagnostic_quantities_for_turbulence): end"
)

# HALO EXCHANGE IF (discr_vn > 1) THEN CALL sync_patch_array
# TODO (magdalena) move this up and do asynchronousely
if self.config.type_vn_diffu > 1:
log.debug("communication rbf extrapolation of z_nable2_e - start")
self._exchange.exchange_and_wait(EdgeDim, self.z_nabla2_e)
@@ -744,6 +774,7 @@ def _do_diffusion_step(
geofac_grg_y=self.interpolation_state.geofac_grg_y,
w_old=self.w_tmp,
w=prognostic_state.w,
type_shear=int32(self.config.shear_type.value),
dwdx=diagnostic_state.dwdx,
dwdy=diagnostic_state.dwdy,
diff_multfac_w=self.diff_multfac_w,
Original file line number Diff line number Diff line change
@@ -28,12 +28,13 @@


@field_operator
def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
area: Field[[CellDim], float],
geofac_n2s: Field[[CellDim, C2E2CODim], float],
geofac_grg_x: Field[[CellDim, C2E2CODim], float],
geofac_grg_y: Field[[CellDim, C2E2CODim], float],
w_old: Field[[CellDim, KDim], float],
type_shear: int32,
dwdx: Field[[CellDim, KDim], float],
dwdy: Field[[CellDim, KDim], float],
diff_multfac_w: float,
@@ -49,11 +50,14 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
Field[[CellDim, KDim], float],
]:
vert_idx = broadcast(vert_idx, (CellDim, KDim))

dwdx, dwdy = where(
int32(0) < vert_idx,
_calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y),
(dwdx, dwdy),
dwdx, dwdy = (
where(
int32(0) < vert_idx,
_calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y),
(dwdx, dwdy),
)
if type_shear == int32(2)
else (dwdx, dwdy)
)

z_nabla2_c = _calculate_nabla2_for_w(w_old, geofac_n2s)
@@ -84,6 +88,7 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
geofac_grg_y: Field[[CellDim, C2E2CODim], float],
w_old: Field[[CellDim, KDim], float],
w: Field[[CellDim, KDim], float],
type_shear: int32,
dwdx: Field[[CellDim, KDim], float],
dwdy: Field[[CellDim, KDim], float],
diff_multfac_w: float,
@@ -98,12 +103,13 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
vertical_start: int32,
vertical_end: int32,
):
_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance(
_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence(
area,
geofac_n2s,
geofac_grg_x,
geofac_grg_y,
w_old,
type_shear,
dwdx,
dwdy,
diff_multfac_w,
Original file line number Diff line number Diff line change
@@ -91,7 +91,7 @@ def test_parallel_diffusion(
cell_params=cell_geometry,
)
print(f"rank={processor_props.rank}/{processor_props.comm_size}: diffusion initialized ")
diagnostic_state = construct_diagnostics(diffusion_savepoint_init)
diagnostic_state = construct_diagnostics(diffusion_savepoint_init, grid_savepoint)
prognostic_state = diffusion_savepoint_init.construct_prognostics()
if linit:
diffusion.initial_run(
Original file line number Diff line number Diff line change
@@ -240,10 +240,10 @@ def test_verify_diffusion_init_against_savepoint(

@pytest.mark.datatest
@pytest.mark.parametrize(
"experiment,step_date_init, step_date_exit",
"experiment,step_date_init, step_date_exit, damping_height",
[
("mch_ch_r04b09_dsl", "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000"),
# ("exclaim_ape_R02B04", "2000-01-01T00:00:00.000", "2000-01-01T00:00:00.000"),
# ("mch_ch_r04b09_dsl", "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000", 12500.0),
("exclaim_ape_R02B04", "2000-01-01T00:00:02.000", "2000-01-01T00:00:02.000", 50000.0),
],
)
@pytest.mark.parametrize("ndyn_substeps", (2,))
@@ -263,7 +263,7 @@ def test_run_diffusion_single_step(
cell_geometry: CellParams = grid_savepoint.construct_cell_geometry()
interpolation_state = construct_interpolation_state(interpolation_savepoint)
metric_state = construct_metric_state_for_diffusion(metrics_savepoint)
diagnostic_state = construct_diagnostics(diffusion_savepoint_init)
diagnostic_state = construct_diagnostics(diffusion_savepoint_init, grid_savepoint)
prognostic_state = diffusion_savepoint_init.construct_prognostics()
vertical_params = VerticalModelParams(
vct_a=grid_savepoint.vct_a(),
@@ -312,7 +312,7 @@ def test_run_diffusion_initial_step(
cell_geometry: CellParams = grid_savepoint.construct_cell_geometry()
interpolation_state = construct_interpolation_state(interpolation_savepoint)
metric_state = construct_metric_state_for_diffusion(metrics_savepoint)
diagnostic_state = construct_diagnostics(diffusion_savepoint_init)
diagnostic_state = construct_diagnostics(diffusion_savepoint_init, grid_savepoint)
prognostic_state = diffusion_savepoint_init.construct_prognostics()
vct_a = grid_savepoint.vct_a()
vertical_params = VerticalModelParams(vct_a=vct_a, rayleigh_damping_height=damping_height)
37 changes: 24 additions & 13 deletions model/atmosphere/diffusion/tests/diffusion_tests/utils.py
Original file line number Diff line number Diff line change
@@ -13,20 +13,25 @@

import numpy as np

from icon4py.model.atmosphere.diffusion.diffusion import DiffusionConfig, DiffusionType
from icon4py.model.atmosphere.diffusion.diffusion import (
DiffusionConfig,
DiffusionType,
TurbulenceShearForcingType,
)
from icon4py.model.atmosphere.diffusion.diffusion_states import (
DiffusionDiagnosticState,
DiffusionInterpolationState,
DiffusionMetricState,
)
from icon4py.model.common.dimension import CEDim
from icon4py.model.common.dimension import CEDim, CellDim, KDim
from icon4py.model.common.states.prognostic_state import PrognosticState
from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, dallclose
from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, dallclose, zero_field
from icon4py.model.common.test_utils.serialbox_utils import (
IconDiffusionExitSavepoint,
InterpolationSavepoint,
MetricSavepoint,
IconDiffusionInitSavepoint,
IconGridSavepoint,
)


@@ -40,6 +45,7 @@ def exclaim_ape_diffusion_config(ndyn_substeps):
diffusion_type=DiffusionType.SMAGORINSKY_4TH_ORDER,
hdiff_w=True,
hdiff_vn=True,
zdiffu_t=False,
type_t_diffu=2,
type_vn_diffu=1,
hdiff_efdt_ratio=24.0,
@@ -71,6 +77,7 @@ def r04b09_diffusion_config(
velocity_boundary_diffusion_denom=150.0,
max_nudging_coeff=0.075,
n_substeps=ndyn_substeps,
shear_type=TurbulenceShearForcingType.VERTICAL_HORIZONTAL_OF_HORIZONTAL_VERTICAL_WIND,
)


@@ -90,16 +97,16 @@ def verify_diffusion_fields(
val_div_ic = np.asarray(diagnostic_state.div_ic)
ref_hdef_ic = np.asarray(diffusion_savepoint.hdef_ic())
val_hdef_ic = np.asarray(diagnostic_state.hdef_ic)
assert dallclose(ref_div_ic, val_div_ic, atol=5.0e-18)
assert dallclose(ref_hdef_ic, val_hdef_ic)
# assert dallclose(ref_div_ic, val_div_ic, atol=5.0e-18)
# assert dallclose(ref_hdef_ic, val_hdef_ic)
ref_w = np.asarray(diffusion_savepoint.w())
val_w = np.asarray(prognostic_state.w)
ref_dwdx = np.asarray(diffusion_savepoint.dwdx())
val_dwdx = np.asarray(diagnostic_state.dwdx)
ref_dwdy = np.asarray(diffusion_savepoint.dwdy())
val_dwdy = np.asarray(diagnostic_state.dwdy)
assert dallclose(ref_dwdx, val_dwdx, atol=1e-19)
assert dallclose(ref_dwdy, val_dwdy, atol=1e-19)
# ref_dwdx = np.asarray(diffusion_savepoint.dwdx())
# val_dwdx = np.asarray(diagnostic_state.dwdx)
# ref_dwdy = np.asarray(diffusion_savepoint.dwdy())
# val_dwdy = np.asarray(diagnostic_state.dwdy)
# assert dallclose(ref_dwdx, val_dwdx, atol=1e-19)
# assert dallclose(ref_dwdy, val_dwdy, atol=1e-19)

ref_vn = np.asarray(diffusion_savepoint.vn())
val_vn = np.asarray(prognostic_state.vn)
@@ -167,10 +174,14 @@ def construct_metric_state_for_diffusion(savepoint: MetricSavepoint) -> Diffusio

def construct_diagnostics(
savepoint: IconDiffusionInitSavepoint,
grid_savepoint: IconGridSavepoint,
) -> DiffusionDiagnosticState:
grid = grid_savepoint.construct_icon_grid()
dwdx = savepoint.dwdx() if savepoint.dwdx() else zero_field(grid, CellDim, KDim)
dwdy = savepoint.dwdy() if savepoint.dwdy() else zero_field(grid, CellDim, KDim)
return DiffusionDiagnosticState(
hdef_ic=savepoint.hdef_ic(),
div_ic=savepoint.div_ic(),
dwdx=savepoint.dwdx(),
dwdy=savepoint.dwdy(),
dwdx=dwdx,
dwdy=dwdy,
)
Original file line number Diff line number Diff line change
@@ -63,8 +63,7 @@ def wrapper(self, *args, **kwargs):
return func(self, *args, **kwargs)
except serialbox.SerialboxError:
log.warning(f"{name}: field not registered in savepoint {self.savepoint.metainfo}")

return np
return None

return wrapper

@@ -124,7 +123,7 @@ def _read(self, name: str, offset=0, dtype=int):
return (self.serializer.read(name, self.savepoint) - offset).astype(dtype)


class IconGridSavePoint(IconSavepoint):
class IconGridSavepoint(IconSavepoint):
def vct_a(self):
return self._get_field("vct_a", KDim)

@@ -662,9 +661,11 @@ def hdef_ic(self):
def div_ic(self):
return self._get_field("div_ic", CellDim, KDim)

@optionally_registered
def dwdx(self):
return self._get_field("dwdx", CellDim, KDim)

@optionally_registered
def dwdy(self):
return self._get_field("dwdy", CellDim, KDim)

@@ -1188,9 +1189,9 @@ def _grid_size(self):
}
return grid_sizes

def from_savepoint_grid(self) -> IconGridSavePoint:
def from_savepoint_grid(self) -> IconGridSavepoint:
savepoint = self._get_icon_grid_savepoint()
return IconGridSavePoint(savepoint, self.serializer, size=self.grid_size)
return IconGridSavepoint(savepoint, self.serializer, size=self.grid_size)

def _get_icon_grid_savepoint(self):
savepoint = self.serializer.savepoint["icon-grid"].id[1].as_savepoint()