From 688df5672ab42c0817cf91d60165be2960931810 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 15 Dec 2023 17:04:00 +0100 Subject: [PATCH 1/8] add mo_solve_nonhydro_stencil_60 in corrector --- .../dycore/nh_solve/solve_nonhydro.py | 25 ++++++++++++++++--- .../tests/dycore_tests/test_solve_nonhydro.py | 17 +++++++++++++ 2 files changed, 38 insertions(+), 4 deletions(-) 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 45002fbbe..cd20428e9 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 @@ -125,6 +125,9 @@ from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_59 import ( mo_solve_nonhydro_stencil_59, ) +from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_60 import ( + mo_solve_nonhydro_stencil_60, +) from icon4py.model.atmosphere.dycore.mo_solve_nonhydro_stencil_65 import ( mo_solve_nonhydro_stencil_65, ) @@ -421,8 +424,6 @@ def _allocate_local_fields(self): self.z_grad_rth_3 = _allocate(CellDim, KDim, grid=self.grid) self.z_grad_rth_4 = _allocate(CellDim, KDim, grid=self.grid) self.z_dexner_dz_c_2 = _allocate(CellDim, KDim, grid=self.grid) - # TODO (magdalena) missing stencil_60 in corrector remove! this is a field from the diagnostics! - self.exner_dyn_incr = _allocate(CellDim, KDim, grid=self.grid) self.z_hydro_corr = _allocate(EdgeDim, KDim, grid=self.grid) self.z_vn_avg = _allocate(EdgeDim, KDim, grid=self.grid) self.z_theta_v_fl_e = _allocate(EdgeDim, KDim, grid=self.grid) @@ -513,6 +514,7 @@ def time_step( prep_adv=prep_adv, divdamp_fac_o2=divdamp_fac_o2, dtime=dtime, + idyn_timestep=idyn_timestep, nnew=nnew, nnow=nnow, lclean_mflx=lclean_mflx, @@ -1327,7 +1329,7 @@ def run_predictor_step( if idyn_timestep == 1: mo_solve_nonhydro_stencil_59.with_backend(backend)( exner=prognostic_state[nnow].exner, - exner_dyn_incr=self.exner_dyn_incr, + exner_dyn_incr=diagnostic_state_nh.exner_dyn_incr, horizontal_start=start_cell_nudging, horizontal_end=end_cell_local, vertical_start=self.vertical_params.kstart_moist, @@ -1382,6 +1384,7 @@ def run_corrector_step( divdamp_fac_o2: float, prep_adv: PrepAdvection, dtime: float, + dyn_timestep: int, nnew: int, nnow: int, lclean_mflx: bool, @@ -1894,7 +1897,21 @@ def run_corrector_step( vertical_end=self.grid.num_levels, offset_provider={}, ) - # TODO (magdalena) stencil_60 is missing here? + + if dyn_timestep == self.config.ndyn_substeps_var: + mo_solve_nonhydro_stencil_60( + exner=prognostic_state[nnew].exner, + ddt_exner_phy=diagnostic_state_nh.ddt_exner_phy, + exner_dyn_incr=diagnostic_state_nh.exner_dyn_incr, + ndyn_substeps_var=float( + self.config.ndyn_substeps_var + ), # TODO (magdalena) should not be a config parameter... where does it get changed? + dtime=dtime, + horizontal_start=start_cell_nudging, + horizontal_end=end_cell_local, + vertical_start=self.vertical_params.kstart_moist, + vertical_end=self.grid.num_levels, + ) if lprep_adv: if lclean_mflx: 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 5aa50c37a..f0f509fcb 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -674,6 +674,12 @@ def test_nonhydro_corrector_step( savepoint_nonhydro_exit.vn_traj().asnumpy(), rtol=5e-7, # TODO (magdalena) was rtol=1e-10 for local experiment only ) + # stencil 60 only relevant for last substep + assert dallclose( + diagnostic_state_nh.exner_dyn_incr.asnumpy(), + savepoint_nonhydro_exit.exner_dyn_incr().asnumpy(), + atol=1e-14, + ) @pytest.mark.datatest @@ -792,6 +798,12 @@ def test_run_solve_nonhydro_single_step( atol=8e-14, ) + assert dallclose( + diagnostic_state_nh.exner_dyn_incr.asnumpy(), + savepoint_nonhydro_exit.exner_dyn_incr().asnumpy(), + atol=1e-14, + ) + @pytest.mark.datatest @pytest.mark.parametrize("experiment", [REGIONAL_EXPERIMENT]) @@ -952,6 +964,11 @@ def test_run_solve_nonhydro_multi_step( savepoint_nonhydro_exit.vn_new().asnumpy(), atol=5e-13, ) + assert dallclose( + diagnostic_state_nh.exner_dyn_incr.asnumpy(), + savepoint_nonhydro_exit.exner_dyn_incr().asnumpy(), + atol=1e-14, + ) @pytest.mark.datatest From 6bd83ceb5e1cc54c6e2d6686a43c292a117fff1e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 15 Dec 2023 17:39:43 +0100 Subject: [PATCH 2/8] fix call to corrector --- .../icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py | 2 +- .../atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) 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 cd20428e9..011112b7f 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 @@ -514,7 +514,7 @@ def time_step( prep_adv=prep_adv, divdamp_fac_o2=divdamp_fac_o2, dtime=dtime, - idyn_timestep=idyn_timestep, + dyn_timestep=idyn_timestep, nnew=nnew, nnow=nnow, lclean_mflx=lclean_mflx, 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 f0f509fcb..7784d6eb6 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -538,6 +538,7 @@ def test_nonhydro_corrector_step( ) sp_v = savepoint_velocity_init dtime = sp_v.get_metadata("dtime").get("dtime") + dyn_timestep = sp.get_metadata("dyn_timestep").get("dyn_timestep") clean_mflx = sp_v.get_metadata("clean_mflx").get("clean_mflx") lprep_adv = sp_v.get_metadata("prep_adv").get("prep_adv") prep_adv = PrepAdvection( @@ -597,6 +598,7 @@ def test_nonhydro_corrector_step( prep_adv=prep_adv, divdamp_fac_o2=divdamp_fac_o2, dtime=dtime, + dyn_timestep=dyn_timestep, nnew=nnew, nnow=nnow, lclean_mflx=clean_mflx, From dd47d4e549060f1af447c2907ca41e5758e4124d Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 11 Jan 2024 11:55:12 +0100 Subject: [PATCH 3/8] fix fix grid_manager tests add C2V offset provider --- model/common/src/icon4py/model/common/grid/icon.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 5672ea988..6b9ae1ac9 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -34,6 +34,7 @@ V2CDim, V2EDim, VertexDim, + C2VDim, ) from icon4py.model.common.grid.base import BaseGrid from icon4py.model.common.utils import builder @@ -55,6 +56,7 @@ def __init__(self): "E2C2V": (self._get_offset_provider, E2C2VDim, EdgeDim, VertexDim), "V2E": (self._get_offset_provider, V2EDim, VertexDim, EdgeDim), "V2C": (self._get_offset_provider, V2CDim, VertexDim, CellDim), + "C2V": (self._get_offset_provider, C2VDim, CellDim, VertexDim), "E2ECV": (self._get_offset_provider_for_sparse_fields, E2C2VDim, EdgeDim, ECVDim), "C2CEC": (self._get_offset_provider_for_sparse_fields, C2E2CDim, CellDim, CECDim), "C2CE": (self._get_offset_provider_for_sparse_fields, C2EDim, CellDim, CEDim), From ca9c5c1dffbff913d246d3080cf8ef8def2c96ba Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Thu, 11 Jan 2024 17:53:33 +0100 Subject: [PATCH 4/8] fix substep parameter in timestep and runtime checks for stencil_59 and stencil_60 --- .../dycore/nh_solve/solve_nonhydro.py | 23 ++++++++++++------- .../tests/dycore_tests/test_solve_nonhydro.py | 6 ++--- 2 files changed, 17 insertions(+), 12 deletions(-) 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 adc4ed385..2a687d8b1 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 @@ -873,7 +873,7 @@ def run_predictor_step( horizontal_start=start_vertex_lb_plus1, horizontal_end=end_vertex_local_minus1, vertical_start=0, - vertical_end=self.grid.num_levels, # UBOUND(p_cell_in,2) + vertical_end=self.grid.num_levels, offset_provider={ "V2C": self.grid.get_offset_provider("V2C"), }, @@ -1377,7 +1377,8 @@ def run_predictor_step( offset_provider={"Koff": KDim}, ) - if idyn_timestep == 1: + if self._at_first_substep(idyn_timestep): + # TODO (magdalena) this is a simple copy stencil... copies exner to exner_dyn_incr mo_solve_nonhydro_stencil_59.with_backend(backend)( exner=prognostic_state[nnow].exner, exner_dyn_incr=diagnostic_state_nh.exner_dyn_incr, @@ -1427,6 +1428,13 @@ def run_predictor_step( log.debug("exchanging prognostic field 'w'") self._exchange.exchange_and_wait(CellDim, prognostic_state[nnew].w) + @staticmethod + def _at_first_substep(dyn_substep): + return dyn_substep == 0 + + def _at_last_substep(self, dyn_substep): + return dyn_substep == self.config.ndyn_substeps_var - 1 + def run_corrector_step( self, diagnostic_state_nh: DiagnosticStateNonHydro, @@ -1949,20 +1957,19 @@ def run_corrector_step( vertical_end=self.grid.num_levels, offset_provider={}, ) - - if dyn_timestep == self.config.ndyn_substeps_var: + # TODO (magdalena) should not be a config parameter... where does it get changed? + if self._at_last_substep(dyn_timestep): mo_solve_nonhydro_stencil_60( exner=prognostic_state[nnew].exner, ddt_exner_phy=diagnostic_state_nh.ddt_exner_phy, exner_dyn_incr=diagnostic_state_nh.exner_dyn_incr, - ndyn_substeps_var=float( - self.config.ndyn_substeps_var - ), # TODO (magdalena) should not be a config parameter... where does it get changed? + ndyn_substeps_var=float(self.config.ndyn_substeps_var), dtime=dtime, horizontal_start=start_cell_nudging, horizontal_end=end_cell_local, vertical_start=self.vertical_params.kstart_moist, - vertical_end=self.grid.num_levels, + vertical_end=int32(self.grid.num_levels), + offset_provider={}, ) if lprep_adv: 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 2a2d46c03..17d27a462 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -823,8 +823,6 @@ def test_run_solve_nonhydro_multi_step( nnew = 1 recompute = sp_v.get_metadata("recompute").get("recompute") linit = sp_v.get_metadata("linit").get("linit") - dyn_timestep = sp_v.get_metadata("dyn_timestep").get("dyn_timestep") - diagnostic_state_nh = construct_diagnostics(sp, sp_v) prognostic_state_ls = create_prognostic_states(sp) @@ -855,7 +853,7 @@ def test_run_solve_nonhydro_multi_step( prep_adv=prep_adv, divdamp_fac_o2=sp.divdamp_fac_o2(), dtime=dtime, - idyn_timestep=dyn_timestep, + idyn_timestep=i_substep, l_recompute=recompute, l_init=linit, nnew=nnew, @@ -941,7 +939,7 @@ def test_run_solve_nonhydro_multi_step( assert dallclose( diagnostic_state_nh.exner_dyn_incr.asnumpy(), savepoint_nonhydro_exit.exner_dyn_incr().asnumpy(), - atol=1e-14, + atol=1e-5, ) From c7627fcb8091da539264f40b85b5f70993fc2dcf Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 12 Jan 2024 17:14:03 +0100 Subject: [PATCH 5/8] remove substep counter out of solveNonhydro: call determines whether last/first substep is reached. --- .../dycore/nh_solve/solve_nonhydro.py | 51 +++++++++---------- .../tests/dycore_tests/test_solve_nonhydro.py | 17 ++++--- .../src/icon4py/model/common/grid/icon.py | 2 +- .../src/icon4py/model/driver/dycore_driver.py | 16 ++++-- 4 files changed, 47 insertions(+), 39 deletions(-) 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 2a687d8b1..00505e304 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 @@ -14,7 +14,7 @@ from dataclasses import dataclass from typing import Final, Optional -from gt4py.next import Field, as_field +from gt4py.next import as_field from gt4py.next.common import Field from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.program_processors.runners.gtfn import run_gtfn @@ -443,7 +443,7 @@ def init( else: self.jk_start = 0 - en_smag_fac_for_zero_nshift.with_backend(run_gtfn)( + en_smag_fac_for_zero_nshift.with_backend(backend)( self.vertical_params.vct_a, self.config.divdamp_fac, self.config.divdamp_fac2, @@ -510,16 +510,17 @@ def time_step( prep_adv: PrepAdvection, divdamp_fac_o2: float, dtime: float, - idyn_timestep: float, l_recompute: bool, l_init: bool, nnow: int, nnew: int, lclean_mflx: bool, lprep_adv: bool, + at_first_substep: bool, + at_last_substep: bool, ): log.info( - f"running timestep: dtime = {dtime}, dyn_timestep = {idyn_timestep}, init = {l_init}, recompute = {l_recompute}, prep_adv = {lprep_adv} " + f"running timestep: dtime = {dtime}, init = {l_init}, recompute = {l_recompute}, prep_adv = {lprep_adv} clean_mflx={lclean_mflx} " ) start_cell_lb = self.grid.get_start_index( CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) @@ -551,15 +552,13 @@ def time_step( prognostic_state=prognostic_state_ls, z_fields=self.intermediate_fields, dtime=dtime, - idyn_timestep=idyn_timestep, l_recompute=l_recompute, l_init=l_init, + is_first_substep=at_first_substep, nnow=nnow, nnew=nnew, ) - log.info( - f"running corrector step: dtime = {dtime}, dyn_timestep = {idyn_timestep}, prep_adv = {lprep_adv}, divdamp_fac_od = {divdamp_fac_o2} " - ) + self.run_corrector_step( diagnostic_state_nh=diagnostic_state_nh, prognostic_state=prognostic_state_ls, @@ -567,15 +566,13 @@ def time_step( prep_adv=prep_adv, divdamp_fac_o2=divdamp_fac_o2, dtime=dtime, - dyn_timestep=idyn_timestep, nnew=nnew, nnow=nnow, lclean_mflx=lclean_mflx, lprep_adv=lprep_adv, + is_last_substep=at_last_substep, ) - log.info( - f"running corrector step: dtime = {dtime}, dyn_timestep = {idyn_timestep}, prep_adv = {lprep_adv}, divdamp_fac_od = {divdamp_fac_o2} " - ) + start_cell_lb = self.grid.get_start_index( CellDim, HorizontalMarkerIndex.lateral_boundary(CellDim) ) @@ -635,14 +632,14 @@ def run_predictor_step( prognostic_state: list[PrognosticState], z_fields: IntermediateFields, dtime: float, - idyn_timestep: float, l_recompute: bool, l_init: bool, + is_first_substep: bool, nnow: int, nnew: int, ): log.info( - f"running predictor step: dtime = {dtime}, dyn_timestep = {idyn_timestep}, init = {l_init}, recompute = {l_recompute} " + f"running predictor step: dtime = {dtime}, init = {l_init}, recompute = {l_recompute} " ) if l_init or l_recompute: if self.config.itime_scheme == 4 and not l_init: @@ -1377,7 +1374,7 @@ def run_predictor_step( offset_provider={"Koff": KDim}, ) - if self._at_first_substep(idyn_timestep): + if is_first_substep: # TODO (magdalena) this is a simple copy stencil... copies exner to exner_dyn_incr mo_solve_nonhydro_stencil_59.with_backend(backend)( exner=prognostic_state[nnow].exner, @@ -1428,13 +1425,6 @@ def run_predictor_step( log.debug("exchanging prognostic field 'w'") self._exchange.exchange_and_wait(CellDim, prognostic_state[nnew].w) - @staticmethod - def _at_first_substep(dyn_substep): - return dyn_substep == 0 - - def _at_last_substep(self, dyn_substep): - return dyn_substep == self.config.ndyn_substeps_var - 1 - def run_corrector_step( self, diagnostic_state_nh: DiagnosticStateNonHydro, @@ -1443,12 +1433,18 @@ def run_corrector_step( divdamp_fac_o2: float, prep_adv: PrepAdvection, dtime: float, - dyn_timestep: int, nnew: int, nnow: int, lclean_mflx: bool, lprep_adv: bool, + is_last_substep: bool, ): + log.info( + f"running corrector step: dtime = {dtime}, prep_adv = {lprep_adv}, divdamp_fac_o2 = {divdamp_fac_o2} clean_mfxl= {lclean_mflx} " + ) + # TODO (magdalena) is it correcto to use a config parameter here? the actual number of substeps can vary dynmically... + # should this config parameter exist at all in SolveNonHydro? + # Inverse value of ndyn_substeps for tracer advection precomputations r_nsubsteps = 1.0 / self.config.ndyn_substeps_var @@ -1457,7 +1453,7 @@ def run_corrector_step( # Coefficient for reduced fourth-order divergence d scal_divdamp_o2 = divdamp_fac_o2 * self.cell_params.mean_cell_area - _calculate_divdamp_fields.with_backend(run_gtfn)( + _calculate_divdamp_fields.with_backend(backend)( self.enh_divdamp_fac, int32(self.config.divdamp_order), self.cell_params.mean_cell_area, @@ -1810,7 +1806,7 @@ def run_corrector_step( offset_provider={}, ) if not self.l_vert_nested: - mo_solve_nonhydro_stencil_46.with_backend(run_gtfn)( + mo_solve_nonhydro_stencil_46.with_backend(backend)( w_nnew=prognostic_state[nnew].w, z_contr_w_fl_l=z_fields.z_contr_w_fl_l, horizontal_start=start_cell_nudging, @@ -1957,9 +1953,8 @@ def run_corrector_step( vertical_end=self.grid.num_levels, offset_provider={}, ) - # TODO (magdalena) should not be a config parameter... where does it get changed? - if self._at_last_substep(dyn_timestep): - mo_solve_nonhydro_stencil_60( + if is_last_substep: + mo_solve_nonhydro_stencil_60.with_backend(backend)( exner=prognostic_state[nnew].exner, ddt_exner_phy=diagnostic_state_nh.ddt_exner_phy, exner_dyn_incr=diagnostic_state_nh.exner_dyn_incr, 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 17d27a462..5517303f5 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -102,7 +102,6 @@ def test_validate_divdamp_fields_against_savepoint_values( def test_nonhydro_predictor_step( istep_init, istep_exit, - jstep_init, step_date_init, step_date_exit, icon_grid, @@ -161,9 +160,9 @@ def test_nonhydro_predictor_step( prognostic_state=prognostic_state_ls, z_fields=solve_nonhydro.intermediate_fields, dtime=dtime, - idyn_timestep=dyn_timestep, l_recompute=recompute, l_init=linit, + is_first_substep=(dyn_timestep == 1), nnow=nnow, nnew=nnew, ) @@ -576,11 +575,11 @@ def test_nonhydro_corrector_step( prep_adv=prep_adv, divdamp_fac_o2=divdamp_fac_o2, dtime=dtime, - dyn_timestep=dyn_timestep, nnew=nnew, nnow=nnow, lclean_mflx=clean_mflx, lprep_adv=lprep_adv, + is_last_substep=(dyn_timestep == ndyn_substeps), ) if icon_grid.limited_area: assert dallclose(solve_nonhydro._bdy_divdamp.asnumpy(), sp.bdy_divdamp().asnumpy()) @@ -742,13 +741,14 @@ def test_run_solve_nonhydro_single_step( prep_adv=prep_adv, divdamp_fac_o2=initial_divdamp_fac, dtime=dtime, - idyn_timestep=dyn_timestep, l_recompute=recompute, l_init=linit, nnew=nnew, nnow=nnow, lclean_mflx=clean_mflx, lprep_adv=lprep_adv, + at_first_substep=jstep_init == 0, + at_last_substep=jstep_init == ndyn_substeps - 1, ) prognostic_state_nnew = prognostic_state_ls[1] assert dallclose( @@ -847,24 +847,27 @@ def test_run_solve_nonhydro_multi_step( ) for i_substep in range(nsubsteps): + is_first_substep = i_substep == 0 + is_last_substep = i_substep == (nsubsteps - 1) solve_nonhydro.time_step( diagnostic_state_nh=diagnostic_state_nh, prognostic_state_ls=prognostic_state_ls, prep_adv=prep_adv, divdamp_fac_o2=sp.divdamp_fac_o2(), dtime=dtime, - idyn_timestep=i_substep, l_recompute=recompute, l_init=linit, nnew=nnew, nnow=nnow, lclean_mflx=clean_mflx, lprep_adv=lprep_adv, + at_first_substep=is_first_substep, + at_last_substep=is_last_substep, ) linit = False recompute = False clean_mflx = False - if i_substep != nsubsteps - 1: + if not is_last_substep: ntemp = nnow nnow = nnew nnew = ntemp @@ -939,7 +942,7 @@ def test_run_solve_nonhydro_multi_step( assert dallclose( diagnostic_state_nh.exner_dyn_incr.asnumpy(), savepoint_nonhydro_exit.exner_dyn_incr().asnumpy(), - atol=1e-5, + atol=1e-14, ) diff --git a/model/common/src/icon4py/model/common/grid/icon.py b/model/common/src/icon4py/model/common/grid/icon.py index 6b9ae1ac9..c99ad6c93 100644 --- a/model/common/src/icon4py/model/common/grid/icon.py +++ b/model/common/src/icon4py/model/common/grid/icon.py @@ -19,6 +19,7 @@ C2E2CDim, C2E2CODim, C2EDim, + C2VDim, CECDim, CEDim, CellDim, @@ -34,7 +35,6 @@ V2CDim, V2EDim, VertexDim, - C2VDim, ) from icon4py.model.common.grid.base import BaseGrid from icon4py.model.common.utils import builder diff --git a/model/driver/src/icon4py/model/driver/dycore_driver.py b/model/driver/src/icon4py/model/driver/dycore_driver.py index b41a6d5bc..f0bf7ebb2 100644 --- a/model/driver/src/icon4py/model/driver/dycore_driver.py +++ b/model/driver/src/icon4py/model/driver/dycore_driver.py @@ -99,6 +99,13 @@ def _validate_config(self): def _not_first_step(self): self._do_initial_stabilization = False + def _is_last_substep(self, step_nr: int): + return step_nr == (self.n_substeps_var - 1) + + @staticmethod + def _is_first_substep(step_nr: int): + return step_nr == 0 + def _next_simulation_date(self): self._simulation_date += timedelta(seconds=self.run_config.dtime) @@ -241,7 +248,9 @@ def _do_dyn_substepping( do_clean_mflx = True for dyn_substep in range(self._n_substeps_var): log.info( - f"simulation date : {self._simulation_date} sub timestep : {dyn_substep}, initial_stabilization : {self._do_initial_stabilization}, nnow: {self._now}, nnew : {self._next}" + f"simulation date : {self._simulation_date} substep / n_substeps : {dyn_substep} / " + f"{self.n_substeps_var} , initial_stabilization : {self._do_initial_stabilization}, " + f"nnow: {self._now}, nnew : {self._next}" ) self.solve_nonhydro.time_step( solve_nonhydro_diagnostic_state, @@ -249,19 +258,20 @@ def _do_dyn_substepping( prep_adv=prep_adv, divdamp_fac_o2=inital_divdamp_fac_o2, dtime=self._substep_timestep, - idyn_timestep=dyn_substep, l_recompute=do_recompute, l_init=self._do_initial_stabilization, nnew=self._next, nnow=self._now, lclean_mflx=do_clean_mflx, lprep_adv=do_prep_adv, + at_first_substep=self._is_first_substep(dyn_substep), + at_last_substep=self._is_last_substep(dyn_substep), ) do_recompute = False do_clean_mflx = False - if dyn_substep != self._n_substeps_var - 1: + if not self._is_last_substep(dyn_substep): self._swap() self._not_first_step() From 3644e8b4244c351067f07134570ad3ac8dc92041 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 12 Jan 2024 17:33:06 +0100 Subject: [PATCH 6/8] unify parameters and naming --- .../dycore/nh_solve/solve_nonhydro.py | 13 ++++++------- .../tests/dycore_tests/test_solve_nonhydro.py | 19 +++++++++---------- 2 files changed, 15 insertions(+), 17 deletions(-) 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 00505e304..d28482fd8 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 @@ -554,7 +554,7 @@ def time_step( dtime=dtime, l_recompute=l_recompute, l_init=l_init, - is_first_substep=at_first_substep, + at_first_substep=at_first_substep, nnow=nnow, nnew=nnew, ) @@ -570,7 +570,7 @@ def time_step( nnow=nnow, lclean_mflx=lclean_mflx, lprep_adv=lprep_adv, - is_last_substep=at_last_substep, + at_last_substep=at_last_substep, ) start_cell_lb = self.grid.get_start_index( @@ -634,7 +634,7 @@ def run_predictor_step( dtime: float, l_recompute: bool, l_init: bool, - is_first_substep: bool, + at_first_substep: bool, nnow: int, nnew: int, ): @@ -1374,8 +1374,7 @@ def run_predictor_step( offset_provider={"Koff": KDim}, ) - if is_first_substep: - # TODO (magdalena) this is a simple copy stencil... copies exner to exner_dyn_incr + if at_first_substep: mo_solve_nonhydro_stencil_59.with_backend(backend)( exner=prognostic_state[nnow].exner, exner_dyn_incr=diagnostic_state_nh.exner_dyn_incr, @@ -1437,7 +1436,7 @@ def run_corrector_step( nnow: int, lclean_mflx: bool, lprep_adv: bool, - is_last_substep: bool, + at_last_substep: bool, ): log.info( f"running corrector step: dtime = {dtime}, prep_adv = {lprep_adv}, divdamp_fac_o2 = {divdamp_fac_o2} clean_mfxl= {lclean_mflx} " @@ -1953,7 +1952,7 @@ def run_corrector_step( vertical_end=self.grid.num_levels, offset_provider={}, ) - if is_last_substep: + if at_last_substep: mo_solve_nonhydro_stencil_60.with_backend(backend)( exner=prognostic_state[nnew].exner, ddt_exner_phy=diagnostic_state_nh.ddt_exner_phy, 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 5517303f5..ef97537cd 100644 --- a/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py +++ b/model/atmosphere/dycore/tests/dycore_tests/test_solve_nonhydro.py @@ -102,6 +102,7 @@ def test_validate_divdamp_fields_against_savepoint_values( def test_nonhydro_predictor_step( istep_init, istep_exit, + jstep_init, step_date_init, step_date_exit, icon_grid, @@ -125,7 +126,6 @@ def test_nonhydro_predictor_step( sp_v = savepoint_velocity_init dtime = sp_v.get_metadata("dtime").get("dtime") recompute = sp_v.get_metadata("recompute").get("recompute") - dyn_timestep = sp.get_metadata("dyn_timestep").get("dyn_timestep") linit = sp_v.get_metadata("linit").get("linit") nnow = 0 @@ -162,7 +162,7 @@ def test_nonhydro_predictor_step( dtime=dtime, l_recompute=recompute, l_init=linit, - is_first_substep=(dyn_timestep == 1), + at_first_substep=(jstep_init == 0), nnow=nnow, nnew=nnew, ) @@ -489,6 +489,7 @@ def create_vertical_params(damping_height, grid_savepoint): def test_nonhydro_corrector_step( istep_init, istep_exit, + jstep_init, step_date_init, step_date_exit, icon_grid, @@ -515,7 +516,6 @@ def test_nonhydro_corrector_step( ) sp_v = savepoint_velocity_init dtime = sp_v.get_metadata("dtime").get("dtime") - dyn_timestep = sp.get_metadata("dyn_timestep").get("dyn_timestep") clean_mflx = sp_v.get_metadata("clean_mflx").get("clean_mflx") lprep_adv = sp_v.get_metadata("prep_adv").get("prep_adv") prep_adv = PrepAdvection( @@ -579,7 +579,7 @@ def test_nonhydro_corrector_step( nnow=nnow, lclean_mflx=clean_mflx, lprep_adv=lprep_adv, - is_last_substep=(dyn_timestep == ndyn_substeps), + at_last_substep=jstep_init == (ndyn_substeps - 1), ) if icon_grid.limited_area: assert dallclose(solve_nonhydro._bdy_divdamp.asnumpy(), sp.bdy_divdamp().asnumpy()) @@ -709,7 +709,6 @@ def test_run_solve_nonhydro_single_step( nnew = 1 recompute = sp_v.get_metadata("recompute").get("recompute") linit = sp_v.get_metadata("linit").get("linit") - dyn_timestep = sp_v.get_metadata("dyn_timestep").get("dyn_timestep") diagnostic_state_nh = construct_diagnostics(sp, sp_v) @@ -748,7 +747,7 @@ def test_run_solve_nonhydro_single_step( lclean_mflx=clean_mflx, lprep_adv=lprep_adv, at_first_substep=jstep_init == 0, - at_last_substep=jstep_init == ndyn_substeps - 1, + at_last_substep=jstep_init == (ndyn_substeps - 1), ) prognostic_state_nnew = prognostic_state_ls[1] assert dallclose( @@ -804,9 +803,9 @@ def test_run_solve_nonhydro_multi_step( savepoint_nonhydro_exit, savepoint_nonhydro_step_exit, experiment, + ndyn_substeps, ): - nsubsteps = grid_savepoint.get_metadata("nsteps").get("nsteps") - config = construct_config(experiment, nsubsteps) + config = construct_config(experiment, ndyn_substeps=ndyn_substeps) sp = savepoint_nonhydro_init sp_step_exit = savepoint_nonhydro_step_exit nonhydro_params = NonHydrostaticParams(config) @@ -846,9 +845,9 @@ def test_run_solve_nonhydro_multi_step( owner_mask=grid_savepoint.c_owner_mask(), ) - for i_substep in range(nsubsteps): + for i_substep in range(ndyn_substeps): is_first_substep = i_substep == 0 - is_last_substep = i_substep == (nsubsteps - 1) + is_last_substep = i_substep == (ndyn_substeps - 1) solve_nonhydro.time_step( diagnostic_state_nh=diagnostic_state_nh, prognostic_state_ls=prognostic_state_ls, From 7a51bb81da69fa44d06f832506cdbd5001f701f1 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 12 Jan 2024 17:38:48 +0100 Subject: [PATCH 7/8] fix typo --- .../model/atmosphere/dycore/nh_solve/solve_nonhydro.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 d28482fd8..260aead0c 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 @@ -1441,9 +1441,9 @@ def run_corrector_step( log.info( f"running corrector step: dtime = {dtime}, prep_adv = {lprep_adv}, divdamp_fac_o2 = {divdamp_fac_o2} clean_mfxl= {lclean_mflx} " ) - # TODO (magdalena) is it correcto to use a config parameter here? the actual number of substeps can vary dynmically... - # should this config parameter exist at all in SolveNonHydro? + # TODO (magdalena) is it correct to to use a config parameter here? the actual number of substeps can vary dynmically... + # should this config parameter exist at all in SolveNonHydro? # Inverse value of ndyn_substeps for tracer advection precomputations r_nsubsteps = 1.0 / self.config.ndyn_substeps_var From 7cacf1c3a0b70ec307c436323839584b14f90f2e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Fri, 19 Jan 2024 08:29:12 +0100 Subject: [PATCH 8/8] add todo to ndyn_substep in diffusion.py --- .../src/icon4py/model/atmosphere/diffusion/diffusion.py | 5 +++++ 1 file changed, 5 insertions(+) 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 4b4e027d4..0a282ed36 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -217,6 +217,11 @@ def __init__( #: Number of dynamics substeps per fast-physics step #: Called 'ndyn_substeps' in mo_nonhydrostatic_nml.f90 + + # TODO (magdalena) ndyn_substeps may dynamically increase during a model run in order to + # reduce instabilities. Need to figure out whether the parameter is the configured + # (constant!) one or the dynamical one. In the latter case it should be removed from + # DiffusionConfig and init() self.ndyn_substeps: int = n_substeps #: If True, compute horizontal diffusion only at the large time step