From 381d378f201ebc3ba48cabe79794599a26916f88 Mon Sep 17 00:00:00 2001 From: David Strassmann Date: Fri, 22 Nov 2024 13:51:13 +0100 Subject: [PATCH] Add vertical advection granule with PPM --- .../model/atmosphere/advection/advection.py | 29 +- .../atmosphere/advection/advection_states.py | 3 + .../advection/advection_vertical.py | 897 +++++++++++++++++- .../compute_ppm4gpu_courant_number.py | 151 +++ .../compute_ppm4gpu_fractional_flux.py | 121 +++ .../stencils/compute_ppm4gpu_integer_flux.py | 78 ++ .../compute_vertical_tracer_flux_upwind.py | 43 + .../copy_cell_kdim_field_koff_plus1.py | 42 + ...limit_vertical_slope_semi_monotonically.py | 51 + ...est_compute_vertical_tracer_flux_upwind.py | 58 ++ ...limit_vertical_slope_semi_monotonically.py | 52 + .../tests/advection_tests/test_advection.py | 44 +- .../advection/tests/advection_tests/utils.py | 67 +- .../common/test_utils/datatest_fixtures.py | 2 +- 14 files changed, 1566 insertions(+), 72 deletions(-) create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py create mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py create mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index ae2aae72b..410a40fd0 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -75,6 +75,10 @@ class VerticalAdvectionType(Enum): #: no vertical advection NO_ADVECTION = auto() + #: 1st order upwind + UPWIND_1ST_ORDER = auto() + #: 3rd order PPM + PPM_3RD_ORDER = auto() class VerticalAdvectionLimiter(Enum): @@ -84,6 +88,8 @@ class VerticalAdvectionLimiter(Enum): #: no vertical limiter NO_LIMITER = auto() + #: semi-monotonic vertical limiter + SEMI_MONOTONIC = auto() @dataclasses.dataclass(frozen=True) @@ -393,7 +399,7 @@ def convert_config_to_horizontal_vertical_advection( match config.horizontal_advection_type: case HorizontalAdvectionType.NO_ADVECTION: - horizontal_advection = advection_horizontal.NoAdvection(grid=grid) + horizontal_advection = advection_horizontal.NoAdvection(grid=grid, backend=backend) case HorizontalAdvectionType.LINEAR_2ND_ORDER: tracer_flux = advection_horizontal.SecondOrderMiura( grid=grid, @@ -417,13 +423,32 @@ def convert_config_to_horizontal_vertical_advection( match config.vertical_advection_limiter: case VerticalAdvectionLimiter.NO_LIMITER: - ... + vertical_limiter = advection_vertical.NoLimiter(grid=grid, backend=backend) + case VerticalAdvectionLimiter.SEMI_MONOTONIC: + vertical_limiter = advection_vertical.SemiMonotonicLimiter(grid=grid, backend=backend) case _: raise NotImplementedError(f"Unknown vertical advection limiter.") match config.vertical_advection_type: case VerticalAdvectionType.NO_ADVECTION: vertical_advection = advection_vertical.NoAdvection(grid=grid, backend=backend) + case VerticalAdvectionType.UPWIND_1ST_ORDER: + boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) + vertical_advection = advection_vertical.FirstOrderUpwind( + boundary_conditions=boundary_conditions, + grid=grid, + metric_state=metric_state, + backend=backend, + ) + case VerticalAdvectionType.PPM_3RD_ORDER: + boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) + vertical_advection = advection_vertical.PiecewiseParabolicMethod( + boundary_conditions=boundary_conditions, + vertical_limiter=vertical_limiter, + grid=grid, + metric_state=metric_state, + backend=backend, + ) case _: raise NotImplementedError(f"Unknown vertical advection type.") diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py index 9a038ecdb..a09ed7daf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py @@ -85,3 +85,6 @@ class AdvectionMetricState: #: metrical modification factor for vertical part of divergence at full levels (KDim) deepatmo_divzu: fa.KField[ta.wpfloat] + + #: vertical grid spacing at full levels + ddqz_z_full: fa.CellKField[ta.wpfloat] diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index 0ddd9daf4..72f4753c1 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -10,13 +10,59 @@ import logging import icon4py.model.common.grid.states as grid_states +import gt4py.next as gtx from gt4py.next import backend from icon4py.model.atmosphere.advection import advection_states +from icon4py.model.atmosphere.advection.stencils.compute_ppm_quadratic_face_values import ( + compute_ppm_quadratic_face_values, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm_quartic_face_values import ( + compute_ppm_quartic_face_values, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm_slope import compute_ppm_slope +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_courant_number import ( + compute_ppm4gpu_courant_number, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_fractional_flux import ( + compute_ppm4gpu_fractional_flux, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_integer_flux import ( + compute_ppm4gpu_integer_flux, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_parabola_coefficients import ( + compute_ppm4gpu_parabola_coefficients, +) +from icon4py.model.atmosphere.advection.stencils.compute_vertical_parabola_limiter_condition import ( + compute_vertical_parabola_limiter_condition, +) +from icon4py.model.atmosphere.advection.stencils.compute_vertical_tracer_flux_upwind import ( + compute_vertical_tracer_flux_upwind, +) from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field import copy_cell_kdim_field +from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field_koff_minus1 import ( + copy_cell_kdim_field_koff_minus1, +) +from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field_koff_plus1 import ( + copy_cell_kdim_field_koff_plus1, +) +from icon4py.model.atmosphere.advection.stencils.init_constant_cell_kdim_field import ( + init_constant_cell_kdim_field, +) +from icon4py.model.atmosphere.advection.stencils.integrate_tracer_vertically import ( + integrate_tracer_vertically, +) +from icon4py.model.atmosphere.advection.stencils.limit_vertical_parabola_semi_monotonically import ( + limit_vertical_parabola_semi_monotonically, +) +from icon4py.model.atmosphere.advection.stencils.limit_vertical_slope_semi_monotonically import ( + limit_vertical_slope_semi_monotonically, +) + from icon4py.model.common import ( + constants, dimension as dims, field_type_aliases as fa, type_alias as ta, @@ -34,6 +80,274 @@ log = logging.getLogger(__name__) +class BoundaryConditions(ABC): + """Class that sets the upper and lower boundary conditions.""" + + @abstractmethod + def run( + self, + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + """ + Set the vertical boundary conditions. + + Args: + p_mflx_tracer_v: output argument, field that contains new vertical tracer mass flux + horizontal_start: input argument, horizontal start index + horizontal_end: input argument, horizontal end index + + """ + ... + + +class NoFluxCondition(BoundaryConditions): + """Class that sets the upper and lower boundary fluxes to zero.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # stencils + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + + def run( + self, + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + log.debug("vertical boundary conditions computation - start") + + # set upper boundary conditions + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=p_mflx_tracer_v, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + # set lower boundary conditions + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=p_mflx_tracer_v, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + log.debug("vertical boundary conditions computation - end") + + +class VerticalLimiter(ABC): + """Class that limits the vertical reconstructed fields and the fluxes.""" + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + +class NoLimiter(VerticalLimiter): + """Class that implements no vertical parabola limiter.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # fields + self._l_limit = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + + # stencils + self._copy_cell_kdim_field = copy_cell_kdim_field.with_backend(self._backend) + self._copy_cell_kdim_field_koff_plus1 = copy_cell_kdim_field_koff_plus1.with_backend( + self._backend + ) + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + # simply copy to up/low face values + log.debug("running stencil copy_cell_kdim_field - start") + self._copy_cell_kdim_field( + field_in=p_face, + field_out=p_face_up, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field - end") + + log.debug("running stencil copy_cell_kdim_field_koff_plus1 - start") + self._copy_cell_kdim_field_koff_plus1( + field_in=p_face, + field_out=p_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field_koff_plus1 - end") + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + +class SemiMonotonicLimiter(VerticalLimiter): + """Class that implements a semi-monotonic vertical parabola limiter.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._l_limit = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, dtype=gtx.int32, backend=self._backend + ) + + # stencils + self._limit_vertical_slope_semi_monotonically = ( + limit_vertical_slope_semi_monotonically.with_backend(self._backend) + ) + self._compute_vertical_parabola_limiter_condition = ( + compute_vertical_parabola_limiter_condition.with_backend(self._backend) + ) + self._limit_vertical_parabola_semi_monotonically = ( + limit_vertical_parabola_semi_monotonically.with_backend(self._backend) + ) + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + log.debug("running stencil limit_vertical_slope_semi_monotonically - start") + self._limit_vertical_slope_semi_monotonically( + p_cc=p_tracer_now, + z_slope=z_slope, + k=self._k_field, + elev=self._grid.num_levels - 1, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil limit_vertical_slope_semi_monotonically - end") + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + # compute semi-monotonic limiter condition + log.debug("running stencil compute_vertical_parabola_limiter_condition - start") + self._compute_vertical_parabola_limiter_condition( + p_face=p_face, + p_cc=p_tracer_now, + l_limit=self._l_limit, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_vertical_parabola_limiter_condition - end") + + # apply semi-monotonic limiter condition and store to up/low face values + log.debug("running stencil limit_vertical_parabola_semi_monotonically - start") + self._limit_vertical_parabola_semi_monotonically( + l_limit=self._l_limit, + p_face=p_face, + p_cc=p_tracer_now, + p_face_up=p_face_up, + p_face_low=p_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil limit_vertical_parabola_semi_monotonically - end") + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + class VerticalAdvection(ABC): """Class that does one vertical advection step.""" @@ -62,6 +376,12 @@ def run( dtime: input argument, the time step even_timestep: input argument, determines whether halo points are included + Note: + Originally, if even step, vertical transport includes all halo points in order to avoid an additional synchronization step, i.e. + if lstep_even: i_rlstart = 2, i_rlend = min_rlcell + else: i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int + Horizontal advection is always called with the same indices though, i.e. + i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int """ ... @@ -90,6 +410,16 @@ def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): log.debug("vertical advection class init - end") + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + def run( self, prep_adv: advection_states.AdvectionPrepAdvState, @@ -103,15 +433,16 @@ def run( ): log.debug("vertical advection run - start") - horizontal_start = ( - self._start_cell_lateral_boundary_level_2 if even_timestep else self._start_cell_nudging + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep ) + log.debug("running stencil copy_cell_kdim_field - start") self._copy_cell_kdim_field( field_in=p_tracer_now, field_out=p_tracer_new, horizontal_start=horizontal_start, - horizontal_end=(self._end_cell_end if even_timestep else self._end_cell_local), + horizontal_end=horizontal_end, vertical_start=0, vertical_end=self._grid.num_levels, offset_provider=self._grid.offset_providers, @@ -135,13 +466,7 @@ def run( dtime: ta.wpfloat, even_timestep: bool = False, ): - log.debug("horizontal advection run - start") - - # TODO (dastrm): maybe change how the indices are handled here? originally: - # if even step, vertical transport includes all halo points in order to avoid an additional synchronization step, i.e. - # if lstep_even: i_rlstart = 2, i_rlend = min_rlcell - # else: i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int - # note: horizontal advection is always called with the same indices, i.e. i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int + log.debug("vertical advection run - start") self._compute_numerical_flux( prep_adv=prep_adv, @@ -149,6 +474,7 @@ def run( rhodz_now=rhodz_now, p_mflx_tracer_v=p_mflx_tracer_v, dtime=dtime, + even_timestep=even_timestep, ) self._update_unknowns( @@ -158,9 +484,10 @@ def run( rhodz_new=rhodz_new, p_mflx_tracer_v=p_mflx_tracer_v, dtime=dtime, + even_timestep=even_timestep, ) - log.debug("horizontal advection run - end") + log.debug("vertical advection run - end") @abstractmethod def _compute_numerical_flux( @@ -170,6 +497,7 @@ def _compute_numerical_flux( rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): ... @@ -182,35 +510,28 @@ def _update_unknowns( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): ... -class SemiLagrangian(FiniteVolume): - """Class that does one vertical semi-Lagrangian finite volume advection step.""" +class FirstOrderUpwind(FiniteVolume): + """Class that does one vertical first-order accurate upwind finite volume advection step.""" def __init__( self, + boundary_conditions: BoundaryConditions, grid: icon_grid.IconGrid, - interpolation_state: advection_states.AdvectionInterpolationState, - least_squares_state: advection_states.AdvectionLeastSquaresState, metric_state: advection_states.AdvectionMetricState, - edge_params: grid_states.EdgeParams, - cell_params: grid_states.CellParams, - backend: backend.Backend, - exchange: decomposition.ExchangeRuntime = decomposition.SingleNodeExchange(), + backend=backend, ): log.debug("vertical advection class init - start") # input arguments + self._boundary_conditions = boundary_conditions self._grid = grid - self._interpolation_state = interpolation_state - self._least_squares_state = least_squares_state self._metric_state = metric_state - self._edge_params = edge_params - self._cell_params = cell_params self._backend = backend - self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -221,8 +542,32 @@ def __init__( self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) self._end_cell_end = self._grid.end_index(cell_domain(h_grid.Zone.END)) + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + + # stencils + self._compute_vertical_tracer_flux_upwind = ( + compute_vertical_tracer_flux_upwind.with_backend(self._backend) + ) + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + self._integrate_tracer_vertically = integrate_tracer_vertically.with_backend(self._backend) + log.debug("vertical advection class init - end") + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + def _compute_numerical_flux( self, prep_adv: advection_states.AdvectionPrepAdvState, @@ -230,9 +575,35 @@ def _compute_numerical_flux( rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): - # TODO (dastrm): implement this - ... + log.debug("vertical numerical flux computation - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + log.debug("running stencil compute_vertical_tracer_flux_upwind - start") + self._compute_vertical_tracer_flux_upwind( + p_cc=p_tracer_now, + p_mflx_contra_v=prep_adv.mass_flx_ic, + p_upflux=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_vertical_tracer_flux_upwind - end") + + # set boundary conditions + self._boundary_conditions.run( + p_mflx_tracer_v=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + log.debug("vertical numerical flux computation - end") def _update_unknowns( self, @@ -242,6 +613,476 @@ def _update_unknowns( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): - # TODO (dastrm): implement this - ... + log.debug("vertical unknowns update - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + # update tracer mass fraction + log.debug("running stencil integrate_tracer_vertically - start") + self._integrate_tracer_vertically( + tracer_now=p_tracer_now, + rhodz_now=rhodz_now, + p_mflx_tracer_v=p_mflx_tracer_v, + deepatmo_divzl=self._metric_state.deepatmo_divzl, + deepatmo_divzu=self._metric_state.deepatmo_divzu, + rhodz_new=rhodz_new, + tracer_new=p_tracer_new, + k=self._k_field, + p_dtime=dtime, + ivadv_tracer=1, + iadv_slev_jt=0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil integrate_tracer_vertically - end") + + log.debug("vertical unknowns update - end") + + +class PiecewiseParabolicMethod(FiniteVolume): + """Class that does one vertical PPM finite volume advection step.""" + + def __init__( + self, + boundary_conditions: BoundaryConditions, + vertical_limiter: VerticalLimiter, + grid: icon_grid.IconGrid, + metric_state: advection_states.AdvectionMetricState, + backend=backend, + ): + log.debug("vertical advection class init - start") + + # input arguments + self._boundary_conditions = boundary_conditions + self._vertical_limiter = vertical_limiter + self._grid = grid + self._metric_state = metric_state + self._backend = backend + + # cell indices + cell_domain = h_grid.domain(dims.CellDim) + self._start_cell_lateral_boundary_level_2 = self._grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + self._start_cell_nudging = self._grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + self._end_cell_end = self._grid.end_index(cell_domain(h_grid.Zone.END)) + + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_cfl = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, is_halfdim=True, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_slope = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_face = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, is_halfdim=True, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_face_up = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_face_low = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_delta_q = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_a1 = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + + # stencils + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + self._compute_ppm4gpu_courant_number = compute_ppm4gpu_courant_number.with_backend( + self._backend + ) + self._compute_ppm_slope = compute_ppm_slope.with_backend(self._backend) + self._compute_ppm_quadratic_face_values = compute_ppm_quadratic_face_values.with_backend( + self._backend + ) + self._compute_ppm_quartic_face_values = compute_ppm_quartic_face_values.with_backend( + self._backend + ) + self._copy_cell_kdim_field = copy_cell_kdim_field.with_backend(self._backend) + self._copy_cell_kdim_field_koff_minus1 = copy_cell_kdim_field_koff_minus1.with_backend( + self._backend + ) + self._compute_ppm4gpu_parabola_coefficients = ( + compute_ppm4gpu_parabola_coefficients.with_backend(self._backend) + ) + self._compute_ppm4gpu_fractional_flux = compute_ppm4gpu_fractional_flux.with_backend( + self._backend + ) + self._compute_ppm4gpu_integer_flux = compute_ppm4gpu_integer_flux.with_backend( + self._backend + ) + self._integrate_tracer_vertically = integrate_tracer_vertically.with_backend(self._backend) + + log.debug("vertical advection class init - end") + + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + + def _compute_numerical_flux( + self, + prep_adv: advection_states.AdvectionPrepAdvState, + p_tracer_now: fa.CellKField[ta.wpfloat], + rhodz_now: fa.CellKField[ta.wpfloat], + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + dtime: ta.wpfloat, + even_timestep: bool, + ): + log.debug("vertical numerical flux computation - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + ## compute density-weighted Courant number + + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=self._z_cfl, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + # TODO (dastrm): write missing stencil here + if 0: + slevp1_ti = 1 + nlev = self._grid.num_levels - 1 + for jc in range(horizontal_start, horizontal_end): + for jk in range(1, self._grid.num_levels): + z_mass = dtime * prep_adv.mass_flx_ic.ndarray[jc, jk] + if z_mass > 0.0: + jks = jk + while (z_mass >= rhodz_now.ndarray[jc, jks]) and (jks <= nlev - 1): + z_mass -= rhodz_now.ndarray[jc, jks] + jks += 1 + self._z_cfl.ndarray[jc, jk] += 1.0 + z_cflfrac = z_mass / rhodz_now.ndarray[jc, jks] + if z_cflfrac < 1.0: + self._z_cfl.ndarray[jc, jk] += z_cflfrac + else: + self._z_cfl.ndarray[jc, jk] += 1.0 - constants.DBL_EPS + else: + jks = jk - 1 + while (abs(z_mass) >= rhodz_now.ndarray[jc, jks]) and (jks >= slevp1_ti): + z_mass += rhodz_now.ndarray[jc, jks] + jks -= 1 + self._z_cfl.ndarray[jc, jk] -= 1.0 + z_cflfrac = z_mass / rhodz_now.ndarray[jc, jks] + if abs(z_cflfrac) < 1.0: + self._z_cfl.ndarray[jc, jk] += z_cflfrac + else: + self._z_cfl.ndarray[jc, jk] += constants.DBL_EPS - 1.0 + else: + log.debug("running stencil compute_ppm4gpu_courant_number - start") + self._compute_ppm4gpu_courant_number( + p_mflx_contra_v=prep_adv.mass_flx_ic, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + k=self._k_field, + slevp1_ti=1, + nlev=self._grid.num_levels - 1, + dbl_eps=constants.DBL_EPS, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_courant_number - end") + + ## reconstruct face values + + # compute slope + log.debug("running stencil compute_ppm_slope - start") + self._compute_ppm_slope( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + k=self._k_field, + z_slope=self._z_slope, + elev=self._grid.num_levels - 1, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_slope - end") + + # limit slope + self._vertical_limiter.limit_slope( + p_tracer_now=p_tracer_now, + z_slope=self._z_slope, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + # compute second highest face value + log.debug("running stencil compute_ppm_quadratic_face_values - start") + self._compute_ppm_quadratic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=2, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quadratic_face_values - end") + + # compute second lowest face value + log.debug("running stencil compute_ppm_quadratic_face_values - start") + self._compute_ppm_quadratic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels - 1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quadratic_face_values - end") + + # compute highest face value + log.debug("running stencil copy_cell_kdim_field - start") + self._copy_cell_kdim_field( + field_in=p_tracer_now, + field_out=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field - end") + + # compute lowest face value + log.debug("running stencil copy_cell_kdim_field_koff_minus1 - start") + self._copy_cell_kdim_field_koff_minus1( + field_in=p_tracer_now, + field_out=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field_koff_minus1 - end") + + # compute all other face values + log.debug("running stencil compute_ppm_quartic_face_values - start") + self._compute_ppm_quartic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + z_slope=self._z_slope, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=2, + vertical_end=self._grid.num_levels - 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quartic_face_values - end") + + ## limit reconstruction + + self._vertical_limiter.limit_parabola( + p_tracer_now=p_tracer_now, + p_face=self._z_face, + p_face_up=self._z_face_up, + p_face_low=self._z_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + ## compute fractional numerical flux + + log.debug("running stencil compute_ppm4gpu_parabola_coefficients - start") + self._compute_ppm4gpu_parabola_coefficients( + z_face_up=self._z_face_up, + z_face_low=self._z_face_low, + p_cc=p_tracer_now, + z_delta_q=self._z_delta_q, + z_a1=self._z_a1, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_parabola_coefficients - end") + + # TODO (dastrm): write missing stencil here + if 0: + slev = 0 + for jc in range(horizontal_start, horizontal_end): + for jk in range(1, self._grid.num_levels): + js = int(abs(self._z_cfl.ndarray[jc, jk])) + z_cflfrac = abs(self._z_cfl.ndarray[jc, jk]) - js + if self._z_cfl.ndarray[jc, jk] > 0.0: + jks = min(jk, self._grid.num_levels - 1) + js + wsign = 1.0 + else: + jks = (jk - 1) - js + wsign = -1.0 + if jks < slev: + p_mflx_tracer_v.ndarray[jc, jk] = 0.0 + continue + z_q_int = ( + p_tracer_now.ndarray[jc, jks] + + wsign * (self._z_delta_q.ndarray[jc, jks] * (1.0 - z_cflfrac)) + - self._z_a1.ndarray[jc, jks] + * (1.0 - 3.0 * z_cflfrac + 2.0 * z_cflfrac * z_cflfrac) + ) + p_mflx_tracer_v.ndarray[jc, jk] = ( + wsign * rhodz_now.ndarray[jc, jks] * z_cflfrac * z_q_int / dtime + ) + else: + log.debug("running stencil compute_ppm4gpu_fractional_flux - start") + self._compute_ppm4gpu_fractional_flux( + p_cc=p_tracer_now, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + z_delta_q=self._z_delta_q, + z_a1=self._z_a1, + p_upflux=p_mflx_tracer_v, + k=self._k_field, + slev=0, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_fractional_flux - end") + + ## compute integer numerical flux + + # TODO (dastrm): write missing stencil here + if 0: + slev = 0 + for jc in range(horizontal_start, horizontal_end): + for jk in range(1, self._grid.num_levels): + js = int(abs(self._z_cfl.ndarray[jc, jk])) + if js == 0: + continue + z_iflx = 0.0 + if self._z_cfl.ndarray[jc, jk] > 0.0: + for n in range(1, js + 1): + jk_shift = jk - 1 + n + z_iflx += ( + p_tracer_now.ndarray[jc, jk_shift] * rhodz_now.ndarray[jc, jk_shift] + ) + else: + for n in range(1, js + 1): + jk_shift = jk - n + if jk_shift < slev: + continue + z_iflx -= ( + p_tracer_now.ndarray[jc, jk_shift] * rhodz_now.ndarray[jc, jk_shift] + ) + p_mflx_tracer_v.ndarray[jc, jk] += z_iflx / dtime + else: + log.debug("running stencil compute_ppm4gpu_integer_flux - start") + self._compute_ppm4gpu_integer_flux( + p_cc=p_tracer_now, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + p_upflux=p_mflx_tracer_v, + k=self._k_field, + slev=0, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_integer_flux - end") + + ## set boundary conditions + + self._boundary_conditions.run( + p_mflx_tracer_v=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + ## apply flux limiter + + self._vertical_limiter.limit_fluxes( + horizontal_start=horizontal_start, horizontal_end=horizontal_end + ) + + log.debug("vertical numerical flux computation - end") + + def _update_unknowns( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_tracer_new: fa.CellKField[ta.wpfloat], + rhodz_now: fa.CellKField[ta.wpfloat], + rhodz_new: fa.CellKField[ta.wpfloat], + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + dtime: ta.wpfloat, + even_timestep: bool, + ): + log.debug("vertical unknowns update - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + # update tracer mass fraction + log.debug("running stencil integrate_tracer_vertically - start") + self._integrate_tracer_vertically( + tracer_now=p_tracer_now, + rhodz_now=rhodz_now, + p_mflx_tracer_v=p_mflx_tracer_v, + deepatmo_divzl=self._metric_state.deepatmo_divzl, + deepatmo_divzu=self._metric_state.deepatmo_divzu, + rhodz_new=rhodz_new, + tracer_new=p_tracer_new, + k=self._k_field, + p_dtime=dtime, + ivadv_tracer=1, + iadv_slev_jt=0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil integrate_tracer_vertically - end") + + log.debug("vertical unknowns update - end") diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py new file mode 100644 index 000000000..59d77fc99 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py @@ -0,0 +1,151 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import abs, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _compute_ppm4gpu_courant_number( + p_mflx_contra_v: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slevp1_ti: gtx.int32, + nlev: gtx.int32, + dbl_eps: ta.wpfloat, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + z_mass_p = p_dtime * p_mflx_contra_v + z_mass_pos = z_mass_p > 0.0 + + # no while loop iterations + p_cellmass_now_jks = p_cellmass_now + + # one while loop iteration + in_bounds = k <= nlev - 1 + mass_gt_cellmass = where(z_mass_pos & in_bounds, z_mass_p >= p_cellmass_now, False) + z_mass_p = z_mass_p - where(mass_gt_cellmass, p_cellmass_now, 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass, 1.0, 0.0) + p_cellmass_now_jks = where(mass_gt_cellmass, p_cellmass_now(Koff[1]), p_cellmass_now_jks) + + # two while loop iterations + in_bounds = k <= nlev - 2 + mass_gt_cellmass = mass_gt_cellmass & where( + z_mass_pos & in_bounds, z_mass_p >= p_cellmass_now(Koff[1]), False + ) + z_mass_p = z_mass_p - where(mass_gt_cellmass, p_cellmass_now(Koff[1]), 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass, 1.0, 0.0) + p_cellmass_now_jks = where(mass_gt_cellmass, p_cellmass_now(Koff[2]), p_cellmass_now_jks) + + # three while loop iterations + in_bounds = k <= nlev - 3 + mass_gt_cellmass = mass_gt_cellmass & where( + z_mass_pos & in_bounds, z_mass_p >= p_cellmass_now(Koff[2]), False + ) + z_mass_p = z_mass_p - where(mass_gt_cellmass, p_cellmass_now(Koff[2]), 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass, 1.0, 0.0) + p_cellmass_now_jks = where(mass_gt_cellmass, p_cellmass_now(Koff[3]), p_cellmass_now_jks) + + # four while loop iterations + in_bounds = k <= nlev - 4 + mass_gt_cellmass = mass_gt_cellmass & where( + z_mass_pos & in_bounds, z_mass_p >= p_cellmass_now(Koff[3]), False + ) + z_mass_p = z_mass_p - where(mass_gt_cellmass, p_cellmass_now(Koff[3]), 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass, 1.0, 0.0) + p_cellmass_now_jks = where(mass_gt_cellmass, p_cellmass_now(Koff[4]), p_cellmass_now_jks) + + z_cflfrac = where(z_mass_pos, z_mass_p / p_cellmass_now_jks, 0.0) + z_cfl = where(z_cflfrac < 1.0, z_cfl + z_cflfrac, z_cfl + 1.0 - dbl_eps) + + z_mass_n = p_dtime * p_mflx_contra_v + z_mass_neg = z_mass_n <= 0.0 + + # no while loop iterations + p_cellmass_now_jks = p_cellmass_now(Koff[-1]) + + # one while loop iteration + in_bounds = k >= slevp1_ti + 1 + mass_gt_cellmass = where( + z_mass_neg & in_bounds, abs(z_mass_n) >= p_cellmass_now(Koff[-1]), False + ) + z_mass_n = z_mass_n + where(mass_gt_cellmass, p_cellmass_now(Koff[-1]), 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass, 1.0, 0.0) + p_cellmass_now_jks = where(mass_gt_cellmass, p_cellmass_now(Koff[-2]), p_cellmass_now_jks) + + # two while loop iterations + in_bounds = k >= slevp1_ti + 2 + mass_gt_cellmass = mass_gt_cellmass & where( + z_mass_neg & in_bounds, abs(z_mass_n) >= p_cellmass_now(Koff[-2]), False + ) + z_mass_n = z_mass_n + where(mass_gt_cellmass, p_cellmass_now(Koff[-2]), 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass, 1.0, 0.0) + p_cellmass_now_jks = where(mass_gt_cellmass, p_cellmass_now(Koff[-3]), p_cellmass_now_jks) + + # three while loop iterations + in_bounds = k >= slevp1_ti + 3 + mass_gt_cellmass = mass_gt_cellmass & where( + z_mass_neg & in_bounds, abs(z_mass_n) >= p_cellmass_now(Koff[-3]), False + ) + z_mass_n = z_mass_n + where(mass_gt_cellmass, p_cellmass_now(Koff[-3]), 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass, 1.0, 0.0) + p_cellmass_now_jks = where(mass_gt_cellmass, p_cellmass_now(Koff[-4]), p_cellmass_now_jks) + + # four while loop iterations + in_bounds = k >= slevp1_ti + 4 + mass_gt_cellmass = mass_gt_cellmass & where( + z_mass_neg & in_bounds, abs(z_mass_n) >= p_cellmass_now(Koff[-4]), False + ) + z_mass_n = z_mass_n + where(mass_gt_cellmass, p_cellmass_now(Koff[-4]), 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass, 1.0, 0.0) + p_cellmass_now_jks = where(mass_gt_cellmass, p_cellmass_now(Koff[-5]), p_cellmass_now_jks) + + z_cflfrac = where(z_mass_neg, z_mass_n / p_cellmass_now_jks, 0.0) + z_cfl = where(abs(z_cflfrac) < 1.0, z_cfl + z_cflfrac, z_cfl + dbl_eps - 1.0) + + return z_cfl + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_courant_number( + p_mflx_contra_v: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slevp1_ti: gtx.int32, + nlev: gtx.int32, + dbl_eps: ta.wpfloat, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_courant_number( + p_mflx_contra_v, + p_cellmass_now, + z_cfl, + k, + slevp1_ti, + nlev, + dbl_eps, + p_dtime, + out=z_cfl, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py new file mode 100644 index 000000000..2ac11aa03 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py @@ -0,0 +1,121 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import abs, astype, floor, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff +from icon4py.model.common.type_alias import wpfloat + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _sum_neighbor_contributions( + mask1: fa.CellKField[bool], + mask2: fa.CellKField[bool], + js: fa.CellKField[ta.wpfloat], + p_cc: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + p_cc_p0 = where(mask1 & (js == 0.0), p_cc, 0.0) + p_cc_p1 = where(mask1 & (js == 1.0), p_cc(Koff[1]), 0.0) + p_cc_p2 = where(mask1 & (js == 2.0), p_cc(Koff[2]), 0.0) + p_cc_p3 = where(mask1 & (js == 3.0), p_cc(Koff[3]), 0.0) + p_cc_p4 = where(mask1 & (js == 4.0), p_cc(Koff[4]), 0.0) + p_cc_m0 = where(mask2 & (js == 0.0), p_cc(Koff[-1]), 0.0) + p_cc_m1 = where(mask2 & (js == 1.0), p_cc(Koff[-2]), 0.0) + p_cc_m2 = where(mask2 & (js == 2.0), p_cc(Koff[-3]), 0.0) + p_cc_m3 = where(mask2 & (js == 3.0), p_cc(Koff[-4]), 0.0) + p_cc_m4 = where(mask2 & (js == 4.0), p_cc(Koff[-5]), 0.0) + p_cc_jks = ( + p_cc_p0 + + p_cc_p1 + + p_cc_p2 + + p_cc_p3 + + p_cc_p4 + + p_cc_m0 + + p_cc_m1 + + p_cc_m2 + + p_cc_m3 + + p_cc_m4 + ) + return p_cc_jks + + +@gtx.field_operator +def _compute_ppm4gpu_fractional_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + z_delta_q: fa.CellKField[ta.wpfloat], + z_a1: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + js = floor(abs(z_cfl)) + z_cflfrac = abs(z_cfl) - js + + z_cfl_pos = z_cfl > 0.0 + z_cfl_neg = not z_cfl_pos + wsign = where(z_cfl_pos, 1.0, -1.0) + + in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat) + + p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc) + p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now) + z_delta_q_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_delta_q) + z_a1_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_a1) + + z_q_int = ( + p_cc_jks + + wsign * (z_delta_q_jks * (1.0 - z_cflfrac)) + - z_a1_jks * (1.0 - 3.0 * z_cflfrac + 2.0 * z_cflfrac * z_cflfrac) + ) + + p_upflux = where( + in_slev_bounds, wsign * p_cellmass_now_jks * z_cflfrac * z_q_int / p_dtime, 0.0 + ) + + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_fractional_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + z_delta_q: fa.CellKField[ta.wpfloat], + z_a1: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_fractional_flux( + p_cc, + p_cellmass_now, + z_cfl, + z_delta_q, + z_a1, + k, + slev, + p_dtime, + out=p_upflux, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py new file mode 100644 index 000000000..b6c5b8c7d --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py @@ -0,0 +1,78 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import abs, astype, floor, where + +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_fractional_flux import ( + _sum_neighbor_contributions, +) +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.type_alias import wpfloat + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _compute_ppm4gpu_integer_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + js = floor(abs(z_cfl)) - 1.0 + + z_cfl_pos = z_cfl > 0.0 + z_cfl_neg = not z_cfl_pos + wsign = where(z_cfl_pos, 1.0, -1.0) + + in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat) + + p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc) + p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now) + + z_iflx = wsign * p_cc_jks * p_cellmass_now_jks + + p_upflux = p_upflux + where(in_slev_bounds, z_iflx / p_dtime, 0.0) + + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_integer_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_integer_flux( + p_cc, + p_cellmass_now, + z_cfl, + p_upflux, + k, + slev, + p_dtime, + out=p_upflux, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py new file mode 100644 index 000000000..556c15c73 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py @@ -0,0 +1,43 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +@gtx.field_operator +def _compute_vertical_tracer_flux_upwind( + p_cc: fa.CellKField[ta.wpfloat], + p_mflx_contra_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim +) -> fa.CellKField[ta.wpfloat]: + p_upflux = where(p_mflx_contra_v >= 0.0, p_cc, p_cc(Koff[-1])) * p_mflx_contra_v + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_vertical_tracer_flux_upwind( + p_cc: fa.CellKField[ta.wpfloat], + p_mflx_contra_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_upflux: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_vertical_tracer_flux_upwind( + p_cc, + p_mflx_contra_v, + out=(p_upflux), + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py new file mode 100644 index 000000000..b3b296dee --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py @@ -0,0 +1,42 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +# TODO (dastrm): move this highly generic stencil to common +# TODO (dastrm): this stencil has no test + + +@gtx.field_operator +def _copy_cell_kdim_field_koff_plus1( + field_in: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + return field_in(Koff[1]) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def copy_cell_kdim_field_koff_plus1( + field_in: fa.CellKField[ta.wpfloat], + field_out: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _copy_cell_kdim_field_koff_plus1( + field_in, + out=field_out, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py new file mode 100644 index 000000000..4737343fa --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py @@ -0,0 +1,51 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import abs, minimum, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +@gtx.field_operator +def _limit_vertical_slope_semi_monotonically( + p_cc: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + elev: gtx.int32, +) -> fa.CellKField[ta.wpfloat]: + p_cc_min_last = minimum(p_cc(Koff[-1]), p_cc) + p_cc_min = where(k == elev, p_cc_min_last, minimum(p_cc_min_last, p_cc(Koff[1]))) + slope_l = minimum(abs(z_slope), 2.0 * (p_cc - p_cc_min)) + slope = where(z_slope >= 0.0, slope_l, -slope_l) + return slope + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def limit_vertical_slope_semi_monotonically( + p_cc: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + elev: gtx.int32, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _limit_vertical_slope_semi_monotonically( + p_cc, + z_slope, + k, + elev, + out=z_slope, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py new file mode 100644 index 000000000..f4d7db02f --- /dev/null +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py @@ -0,0 +1,58 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +import pytest + +import icon4py.model.common.test_utils.helpers as helpers +from icon4py.model.atmosphere.advection.stencils.compute_vertical_tracer_flux_upwind import ( + compute_vertical_tracer_flux_upwind, +) +from icon4py.model.common import dimension as dims +from icon4py.model.common.settings import xp + + +outslice = (slice(None), slice(1, None)) + + +class TestComputeVerticalTracerFluxUpwind(helpers.StencilTest): + PROGRAM = compute_vertical_tracer_flux_upwind + OUTPUTS = (helpers.Output("p_upflux", refslice=outslice, gtslice=outslice),) + + @staticmethod + def reference( + grid, + p_cc: xp.array, + p_mflx_contra_v: xp.array, + **kwargs, + ) -> dict: + p_upflux = p_cc.copy() + p_upflux[:, 1:] = ( + xp.where(p_mflx_contra_v[:, 1:] >= 0.0, p_cc[:, 1:], p_cc[:, :-1]) + * p_mflx_contra_v[:, 1:] + ) + return dict(p_upflux=p_upflux) + + @pytest.fixture + def input_data(self, grid) -> dict: + p_cc = helpers.random_field(grid, dims.CellDim, dims.KDim) + p_mflx_contra_v = helpers.random_field( + grid, dims.CellDim, dims.KDim + ) # TODO (dastrm): should be KHalfDim + p_upflux = helpers.zero_field( + grid, dims.CellDim, dims.KDim + ) # TODO (dastrm): should be KHalfDim + return dict( + p_cc=p_cc, + p_mflx_contra_v=p_mflx_contra_v, + p_upflux=p_upflux, + horizontal_start=0, + horizontal_end=gtx.int32(grid.num_cells), + vertical_start=1, + vertical_end=gtx.int32(grid.num_levels), + ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py new file mode 100644 index 000000000..2e6279653 --- /dev/null +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py @@ -0,0 +1,52 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +import pytest +from gt4py.next import as_field + +import icon4py.model.common.test_utils.helpers as helpers +from icon4py.model.atmosphere.advection.stencils.limit_vertical_slope_semi_monotonically import ( + limit_vertical_slope_semi_monotonically, +) +from icon4py.model.common import dimension as dims +from icon4py.model.common.settings import xp + + +class TestLimitVerticalSlopeSemiMonotonically(helpers.StencilTest): + PROGRAM = limit_vertical_slope_semi_monotonically + OUTPUTS = (helpers.Output("z_slope", gtslice=(slice(None), slice(1, -1))),) + + @staticmethod + def reference( + grid, p_cc: xp.array, z_slope: xp.array, k: xp.array, elev: gtx.int32, **kwargs + ) -> dict: + p_cc_min_last = xp.minimum(p_cc[:, :-2], p_cc[:, 1:-1]) + p_cc_min = xp.where(k[1:-1] == elev, p_cc_min_last, xp.minimum(p_cc_min_last, p_cc[:, 2:])) + slope_l = xp.minimum(xp.abs(z_slope[:, 1:-1]), 2.0 * (p_cc[:, 1:-1] - p_cc_min)) + slope = xp.where(z_slope[:, 1:-1] >= 0.0, slope_l, -slope_l) + return dict(z_slope=slope) + + @pytest.fixture + def input_data(self, grid) -> dict: + p_cc = helpers.random_field(grid, dims.CellDim, dims.KDim) + z_slope = helpers.random_field(grid, dims.CellDim, dims.KDim) + k = as_field( + (dims.KDim,), xp.arange(0, helpers._shape(grid, dims.KDim)[0], dtype=gtx.int32) + ) + elev = k[-2].as_scalar() + return dict( + p_cc=p_cc, + z_slope=z_slope, + k=k, + elev=elev, + horizontal_start=0, + horizontal_end=gtx.int32(grid.num_cells), + vertical_start=1, + vertical_end=gtx.int32(grid.num_levels - 1), + ) diff --git a/model/atmosphere/advection/tests/advection_tests/test_advection.py b/model/atmosphere/advection/tests/advection_tests/test_advection.py index 75ccbbbc5..cb5c6fc68 100644 --- a/model/atmosphere/advection/tests/advection_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection_tests/test_advection.py @@ -25,17 +25,13 @@ ) -# note about ntracer: The first tracer is always dry air which is not advected. Thus, originally -# ntracer=2 is the first tracer in transport_nml, ntracer=3 the second, and so on. -# Here though, ntracer=1 corresponds to the first tracer in transport_nml. - # ntracer legend for the serialization data used here in test_advection: # ------------------------------------ # ntracer | 1, 2, 3, 4, 5 | # ------------------------------------ # ivadv_tracer | 3, 0, 0, 2, 3 | # itype_hlimit | 3, 3, 4, 0, 0 | -# itype_vlimit | 1, 0, 0, 1, 2 | +# itype_vlimit | 1, 0, 0, 2, 1 | # ihadv_tracer | 52, 2, 2, 0, 0 | # ------------------------------------ @@ -47,7 +43,7 @@ ( "2021-06-20T12:00:10.000", False, - 2, + 3, advection.HorizontalAdvectionType.LINEAR_2ND_ORDER, advection.HorizontalAdvectionLimiter.POSITIVE_DEFINITE, advection.VerticalAdvectionType.NO_ADVECTION, @@ -56,12 +52,42 @@ ( "2021-06-20T12:00:20.000", True, - 2, + 3, advection.HorizontalAdvectionType.LINEAR_2ND_ORDER, advection.HorizontalAdvectionLimiter.POSITIVE_DEFINITE, advection.VerticalAdvectionType.NO_ADVECTION, advection.VerticalAdvectionLimiter.NO_LIMITER, ), + ( + "2021-06-20T12:00:10.000", + False, + 5, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), + ( + "2021-06-20T12:00:20.000", + True, + 5, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), + ] + if 0 + else [ + ( + "2021-06-20T12:00:10.000", + False, + 5, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), ], ) def test_advection_run_single_step( @@ -69,6 +95,7 @@ def test_advection_run_single_step( icon_grid, interpolation_savepoint, least_squares_savepoint, + metrics_savepoint, advection_init_savepoint, advection_exit_savepoint, data_provider, @@ -89,7 +116,7 @@ def test_advection_run_single_step( ) interpolation_state = construct_interpolation_state(interpolation_savepoint) least_squares_state = construct_least_squares_state(least_squares_savepoint) - metric_state = construct_metric_state(icon_grid) + metric_state = construct_metric_state(icon_grid, metrics_savepoint) edge_geometry = grid_savepoint.construct_edge_geometry() cell_geometry = grid_savepoint.construct_cell_geometry() @@ -132,4 +159,5 @@ def test_advection_run_single_step( diagnostic_state_ref=diagnostic_state_ref, p_tracer_new=p_tracer_new, p_tracer_new_ref=p_tracer_new_ref, + even_timestep=even_timestep, ) diff --git a/model/atmosphere/advection/tests/advection_tests/utils.py b/model/atmosphere/advection/tests/advection_tests/utils.py index 3e44294c7..6aecba2e6 100644 --- a/model/atmosphere/advection/tests/advection_tests/utils.py +++ b/model/atmosphere/advection/tests/advection_tests/utils.py @@ -8,10 +8,12 @@ import logging +import gt4py.next as gtx from icon4py.model.atmosphere.advection import advection, advection_states from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid +from icon4py.model.common.settings import xp from icon4py.model.common.test_utils import helpers, serialbox_utils as sb from icon4py.model.common.utils import gt4py_field_allocation as field_alloc @@ -54,12 +56,16 @@ def construct_least_squares_state( ) -def construct_metric_state(icon_grid) -> advection_states.AdvectionMetricState: +def construct_metric_state( + icon_grid, savepoint: sb.MetricSavepoint +) -> advection_states.AdvectionMetricState: constant_f = helpers.constant_field(icon_grid, 1.0, dims.KDim) + ddqz_z_full_xp = xp.reciprocal(savepoint.inv_ddqz_z_full().ndarray) return advection_states.AdvectionMetricState( deepatmo_divh=constant_f, deepatmo_divzl=constant_f, deepatmo_divzu=constant_f, + ddqz_z_full=gtx.as_field((dims.CellDim, dims.KDim), ddqz_z_full_xp), ) @@ -128,11 +134,17 @@ def verify_advection_fields( diagnostic_state_ref: advection_states.AdvectionDiagnosticState, p_tracer_new: fa.CellKField[ta.wpfloat], p_tracer_new_ref: fa.CellKField[ta.wpfloat], + even_timestep: bool, ): # cell indices cell_domain = h_grid.domain(dims.CellDim) start_cell_lateral_boundary = grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY)) + start_cell_lateral_boundary_level_2 = grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + end_cell_end = grid.end_index(cell_domain(h_grid.Zone.END)) # edge indices edge_domain = h_grid.domain(dims.EdgeDim) @@ -141,46 +153,35 @@ def verify_advection_fields( ) end_edge_halo = grid.end_index(edge_domain(h_grid.Zone.HALO)) - # log advection output fields - log_dbg( - diagnostic_state.hfl_tracer.asnumpy()[start_edge_lateral_boundary_level_5:end_edge_halo, :], - "hfl_tracer", - ) - log_dbg( - diagnostic_state_ref.hfl_tracer.asnumpy()[ - start_edge_lateral_boundary_level_5:end_edge_halo, : - ], - "hfl_tracer_ref", - ) - log_dbg( - diagnostic_state.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "vfl_tracer", - ) - log_dbg( - diagnostic_state_ref.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "vfl_tracer_ref", - ) - log_dbg(p_tracer_new.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], "p_tracer_new") - log_dbg( - p_tracer_new_ref.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "p_tracer_new_ref", + hfl_tracer_range = xp.arange(start_edge_lateral_boundary_level_5, end_edge_halo) + vfl_tracer_range = ( + xp.arange(start_cell_lateral_boundary_level_2, end_cell_end) + if even_timestep + else xp.arange(start_cell_nudging, end_cell_local) ) + p_tracer_new_range = xp.arange(start_cell_lateral_boundary, end_cell_local) + + # log advection output fields + log_dbg(diagnostic_state.hfl_tracer.asnumpy()[hfl_tracer_range, :], "hfl_tracer") + log_dbg(diagnostic_state_ref.hfl_tracer.asnumpy()[hfl_tracer_range, :], "hfl_tracer_ref") + log_dbg(diagnostic_state.vfl_tracer.asnumpy()[vfl_tracer_range, :], "vfl_tracer") + log_dbg(diagnostic_state_ref.vfl_tracer.asnumpy()[vfl_tracer_range, :], "vfl_tracer_ref") + log_dbg(p_tracer_new.asnumpy()[p_tracer_new_range, :], "p_tracer_new") + log_dbg(p_tracer_new_ref.asnumpy()[p_tracer_new_range, :], "p_tracer_new_ref") # verify advection output fields assert helpers.dallclose( - diagnostic_state.hfl_tracer.asnumpy()[start_edge_lateral_boundary_level_5:end_edge_halo, :], - diagnostic_state_ref.hfl_tracer.asnumpy()[ - start_edge_lateral_boundary_level_5:end_edge_halo, : - ], + diagnostic_state.hfl_tracer.asnumpy()[hfl_tracer_range, :], + diagnostic_state_ref.hfl_tracer.asnumpy()[hfl_tracer_range, :], rtol=1e-10, ) - assert helpers.dallclose( # TODO (dastrm): adjust indices once there is vertical transport - diagnostic_state.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - diagnostic_state_ref.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], + assert helpers.dallclose( + diagnostic_state.vfl_tracer.asnumpy()[vfl_tracer_range, :], + diagnostic_state_ref.vfl_tracer.asnumpy()[vfl_tracer_range, :], rtol=1e-10, ) assert helpers.dallclose( - p_tracer_new.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - p_tracer_new_ref.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], + p_tracer_new.asnumpy()[p_tracer_new_range, :], + p_tracer_new_ref.asnumpy()[p_tracer_new_range, :], atol=1e-16, ) diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py index b84805ad5..c80ccd100 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py @@ -167,7 +167,7 @@ def interpolation_savepoint(data_provider): # F811 @pytest.fixture def metrics_savepoint(data_provider): # F811 - """Load data from ICON mestric state savepoint.""" + """Load data from ICON metric state savepoint.""" return data_provider.from_metrics_savepoint()