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 2e6bb58e5d..39082a794c 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -27,7 +27,12 @@ apply_interpolated_tracer_time_tendency, ) from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field import copy_cell_kdim_field -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common import ( + dimension as dims, + field_type_aliases as fa, + model_backends, + type_alias as ta, +) from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid from icon4py.model.common.model_options import setup_program @@ -98,10 +103,10 @@ class AdvectionConfig: Contains necessary parameters to configure an advection run. """ - def __post_init__(self): + def __post_init__(self) -> None: self._validate() - def _validate(self): + def _validate(self) -> None: """Apply consistency checks and validation on configuration parameters.""" if not hasattr(HorizontalAdvectionType, self.horizontal_advection_type.name): @@ -140,7 +145,7 @@ def run( p_tracer_now: fa.CellKField[ta.wpfloat], p_tracer_new: fa.CellKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: """ Run an advection step. @@ -198,7 +203,7 @@ def run( p_tracer_now: fa.CellKField[ta.wpfloat], p_tracer_new: fa.CellKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: log.debug("advection run - start") log.debug("communication of prep_adv cell field: mass_flx_ic - start") @@ -242,7 +247,10 @@ def __init__( # density fields #: intermediate density times cell thickness, includes either the horizontal or vertical advective density increment [kg/m^2] self._rhodz_ast2 = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, + dims.CellDim, + dims.KDim, + allocator=model_backends.get_allocator(self._backend), ) self._determine_local_domains() # stencils @@ -277,7 +285,7 @@ def __init__( log.debug("advection class init - end") - def _determine_local_domains(self): + def _determine_local_domains(self) -> None: # cell indices cell_domain = h_grid.domain(dims.CellDim) self._start_cell_lateral_boundary = self._grid.start_index( @@ -301,7 +309,7 @@ def run( p_tracer_now: fa.CellKField[ta.wpfloat], p_tracer_new: fa.CellKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: log.debug("advection run - start") log.debug("communication of prep_adv cell field: mass_flx_ic - start") @@ -410,9 +418,10 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- ) -> tuple[advection_horizontal.HorizontalAdvection, advection_vertical.VerticalAdvection]: exchange = exchange or decomposition.SingleNodeExchange() assert exchange is not None, "Exchange runtime must not be None." + horizontal_limiter: advection_horizontal.HorizontalFluxLimiter | None match config.horizontal_advection_limiter: case HorizontalAdvectionLimiter.NO_LIMITER: - horizontal_limiter = advection_horizontal.HorizontalFluxLimiter() + horizontal_limiter = advection_horizontal.NoLimiter() case HorizontalAdvectionLimiter.POSITIVE_DEFINITE: horizontal_limiter = advection_horizontal.PositiveDefinite( grid=grid, @@ -423,6 +432,7 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- case _: raise NotImplementedError("Unknown horizontal advection limiter.") + horizontal_advection: advection_horizontal.HorizontalAdvection match config.horizontal_advection_type: case HorizontalAdvectionType.NO_ADVECTION: horizontal_advection = advection_horizontal.NoAdvection(grid=grid, backend=backend) @@ -446,6 +456,7 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- case _: raise NotImplementedError("Unknown horizontal advection type.") + vertical_limiter: advection_vertical.VerticalLimiter match config.vertical_advection_limiter: case VerticalAdvectionLimiter.NO_LIMITER: vertical_limiter = advection_vertical.NoLimiter(grid=grid, backend=backend) @@ -454,6 +465,7 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- case _: raise NotImplementedError("Unknown vertical advection limiter.") + vertical_advection: advection_vertical.VerticalAdvection match config.vertical_advection_type: case VerticalAdvectionType.NO_ADVECTION: vertical_advection = advection_vertical.NoAdvection(grid=grid, backend=backend) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py index 241f71182d..bb6b0888e3 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py @@ -40,11 +40,11 @@ dimension as dims, field_type_aliases as fa, model_backends, + model_options, type_alias as ta, ) from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid -from icon4py.model.common.model_options import setup_program from icon4py.model.common.utils import data_allocation as data_alloc @@ -63,7 +63,19 @@ def apply_flux_limiter( p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], rhodz_now: fa.CellKField[ta.wpfloat], dtime: ta.wpfloat, - ): ... + ) -> None: ... + + +class NoLimiter(HorizontalFluxLimiter): + """Do not apply any limiting.""" + + def apply_flux_limiter( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], + rhodz_now: fa.CellKField[ta.wpfloat], + dtime: ta.wpfloat, + ) -> None: ... class PositiveDefinite(HorizontalFluxLimiter): @@ -73,7 +85,7 @@ def __init__( self, grid: icon_grid.IconGrid, interpolation_state: advection_states.AdvectionInterpolationState, - backend: model_backends.BackendLike, + backend: gtx.typing.Backend | None, exchange: decomposition.ExchangeRuntime | None = decomposition.single_node_default, ): self._grid = grid @@ -97,39 +109,46 @@ def __init__( # limiter fields self._r_m = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, + dims.CellDim, + dims.KDim, + allocator=model_backends.get_allocator(self._backend), ) # stencils - self._compute_positive_definite_horizontal_multiplicative_flux_factor = setup_program( - backend=self._backend, - program=compute_positive_definite_horizontal_multiplicative_flux_factor, - constant_args={ - "geofac_div": self._interpolation_state.geofac_div, - "dbl_eps": constants.DBL_EPS, - }, - horizontal_sizes={ - "horizontal_start": self._start_cell_lateral_boundary_level_2, - "horizontal_end": self._end_cell_local, - }, - vertical_sizes={ - "vertical_start": gtx.int32(0), - "vertical_end": gtx.int32(self._grid.num_levels), - }, - offset_provider=self._grid.connectivities, - ) - self._apply_positive_definite_horizontal_multiplicative_flux_factor = setup_program( - backend=self._backend, - program=apply_positive_definite_horizontal_multiplicative_flux_factor, - horizontal_sizes={ - "horizontal_start": self._start_edge_lateral_boundary_level_5, - "horizontal_end": self._end_edge_halo, - }, - vertical_sizes={ - "vertical_start": gtx.int32(0), - "vertical_end": gtx.int32(self._grid.num_levels), - }, - offset_provider=self._grid.connectivities, + self._compute_positive_definite_horizontal_multiplicative_flux_factor = ( + model_options.setup_program( + backend=self._backend, + program=compute_positive_definite_horizontal_multiplicative_flux_factor, + constant_args={ + "geofac_div": self._interpolation_state.geofac_div, + "dbl_eps": constants.DBL_EPS, + }, + horizontal_sizes={ + "horizontal_start": self._start_cell_lateral_boundary_level_2, + "horizontal_end": self._end_cell_local, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(self._grid.num_levels), + }, + offset_provider=self._grid.connectivities, + ) + ) + self._apply_positive_definite_horizontal_multiplicative_flux_factor = ( + model_options.setup_program( + backend=self._backend, + program=apply_positive_definite_horizontal_multiplicative_flux_factor, + horizontal_sizes={ + "horizontal_start": self._start_edge_lateral_boundary_level_5, + "horizontal_end": self._end_edge_halo, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(self._grid.num_levels), + }, + offset_provider=self._grid.connectivities, + ) ) def apply_flux_limiter( @@ -138,7 +157,7 @@ def apply_flux_limiter( p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], rhodz_now: fa.CellKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: # compute multiplicative flux factor to guarantee no undershoot log.debug( "running stencil compute_positive_definite_horizontal_multiplicative_flux_factor - start" @@ -180,11 +199,11 @@ def compute_tracer_flux( prep_adv: advection_states.AdvectionPrepAdvState, p_tracer_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], - p_distv_bary_1: fa.EdgeKField[ta.vpfloat], - p_distv_bary_2: fa.EdgeKField[ta.vpfloat], + p_distv_bary_1: fa.EdgeKField[ta.anyfloat], + p_distv_bary_2: fa.EdgeKField[ta.anyfloat], rhodz_now: fa.CellKField[ta.wpfloat], dtime: ta.wpfloat, - ): ... + ) -> None: ... class SecondOrderMiura(SemiLagrangianTracerFlux): @@ -194,13 +213,13 @@ def __init__( self, grid: icon_grid.IconGrid, least_squares_state: advection_states.AdvectionLeastSquaresState, - backend: model_backends.BackendLike, + backend: gtx.typing.Backend | None, horizontal_limiter: HorizontalFluxLimiter | None = None, ): self._grid = grid self._least_squares_state = least_squares_state self._backend = backend - self._horizontal_limiter = horizontal_limiter or HorizontalFluxLimiter() + self._horizontal_limiter = horizontal_limiter or NoLimiter() # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -217,18 +236,19 @@ def __init__( self._end_edge_halo = self._grid.end_index(edge_domain(h_grid.Zone.HALO)) # reconstruction fields + allocator = model_backends.get_allocator(self._backend) self._p_coeff_1 = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, allocator=allocator ) self._p_coeff_2 = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, allocator=allocator ) self._p_coeff_3 = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, allocator=allocator ) # stencils - self._reconstruct_linear_coefficients_svd = setup_program( + self._reconstruct_linear_coefficients_svd = model_options.setup_program( backend=self._backend, program=reconstruct_linear_coefficients_svd, horizontal_sizes={ @@ -241,18 +261,20 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._compute_horizontal_tracer_flux_from_linear_coefficients_alt = setup_program( - backend=self._backend, - program=compute_horizontal_tracer_flux_from_linear_coefficients_alt, - horizontal_sizes={ - "horizontal_start": self._start_edge_lateral_boundary_level_5, - "horizontal_end": self._end_edge_halo, - }, - vertical_sizes={ - "vertical_start": gtx.int32(0), - "vertical_end": gtx.int32(self._grid.num_levels), - }, - offset_provider=self._grid.connectivities, + self._compute_horizontal_tracer_flux_from_linear_coefficients_alt = ( + model_options.setup_program( + backend=self._backend, + program=compute_horizontal_tracer_flux_from_linear_coefficients_alt, + horizontal_sizes={ + "horizontal_start": self._start_edge_lateral_boundary_level_5, + "horizontal_end": self._end_edge_halo, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(self._grid.num_levels), + }, + offset_provider=self._grid.connectivities, + ) ) def compute_tracer_flux( @@ -260,11 +282,11 @@ def compute_tracer_flux( prep_adv: advection_states.AdvectionPrepAdvState, p_tracer_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], - p_distv_bary_1: fa.EdgeKField[ta.vpfloat], - p_distv_bary_2: fa.EdgeKField[ta.vpfloat], + p_distv_bary_1: fa.EdgeKField[ta.anyfloat], + p_distv_bary_2: fa.EdgeKField[ta.anyfloat], rhodz_now: fa.CellKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: log.debug("horizontal tracer flux computation - start") # linear reconstruction using singular value decomposition @@ -320,7 +342,7 @@ def run( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: """ Run a horizontal advection step. @@ -340,11 +362,11 @@ def run( class NoAdvection(HorizontalAdvection): """Class that implements disabled horizontal advection.""" - def __init__(self, grid: icon_grid.IconGrid, backend: model_backends.BackendLike): + def __init__(self, grid: icon_grid.IconGrid, backend: gtx.typing.Backend | None): log.debug("horizontal advection class init - start") # input arguments - self._backend = backend + self._backend = model_options.customize_backend(program=None, backend=backend) # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -352,7 +374,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: model_backends.BackendLike self._end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) # stencils - self._copy_cell_kdim_field = setup_program( + self._copy_cell_kdim_field = model_options.setup_program( backend=self._backend, program=copy_cell_kdim_field, horizontal_sizes={ @@ -377,7 +399,7 @@ def run( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: log.debug("horizontal advection run - start") log.debug("running stencil copy_cell_kdim_field - start") @@ -402,7 +424,7 @@ def run( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: log.debug("horizontal advection run - start") self._compute_numerical_flux( @@ -432,7 +454,7 @@ def _compute_numerical_flux( rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], dtime: ta.wpfloat, - ): ... + ) -> None: ... @abstractmethod def _update_unknowns( @@ -443,7 +465,7 @@ def _update_unknowns( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], dtime: ta.wpfloat, - ): ... + ) -> None: ... class SemiLagrangian(FiniteVolume): @@ -457,7 +479,7 @@ def __init__( metric_state: advection_states.AdvectionMetricState, edge_params: grid_states.EdgeParams, cell_params: grid_states.CellParams, - backend: model_backends.BackendLike, + backend: gtx.typing.Backend | None, exchange: decomposition.ExchangeRuntime | None = decomposition.single_node_default, ): log.debug("horizontal advection class init - start") @@ -488,18 +510,19 @@ def __init__( self._end_edge_halo = self._grid.end_index(edge_domain(h_grid.Zone.HALO)) # backtrajectory fields + allocator = model_backends.get_allocator(self._backend) self._z_real_vt = data_alloc.zero_field( - self._grid, dims.EdgeDim, dims.KDim, allocator=self._backend + self._grid, dims.EdgeDim, dims.KDim, allocator=allocator ) self._p_distv_bary_1 = data_alloc.zero_field( - self._grid, dims.EdgeDim, dims.KDim, allocator=self._backend + self._grid, dims.EdgeDim, dims.KDim, allocator=allocator ) self._p_distv_bary_2 = data_alloc.zero_field( - self._grid, dims.EdgeDim, dims.KDim, allocator=self._backend + self._grid, dims.EdgeDim, dims.KDim, allocator=allocator ) # stencils - self._compute_edge_tangential = setup_program( + self._compute_edge_tangential = model_options.setup_program( backend=self._backend, program=compute_edge_tangential, constant_args={ @@ -516,7 +539,7 @@ def __init__( offset_provider=self._grid.connectivities, ) - self._compute_barycentric_backtrajectory_alt = setup_program( + self._compute_barycentric_backtrajectory_alt = model_options.setup_program( backend=self._backend, program=compute_barycentric_backtrajectory_alt, constant_args={ @@ -537,7 +560,7 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._integrate_tracer_horizontally = setup_program( + self._integrate_tracer_horizontally = model_options.setup_program( backend=self._backend, program=integrate_tracer_horizontally, constant_args={ @@ -564,7 +587,7 @@ def _compute_numerical_flux( rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: log.debug("horizontal numerical flux computation - start") ## tracer-independent part @@ -610,7 +633,7 @@ def _update_unknowns( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_h: fa.EdgeKField[ta.wpfloat], dtime: ta.wpfloat, - ): + ) -> None: log.debug("horizontal unknowns update - start") # update tracer mass fraction 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 501f6a6659..ed4129118c 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 @@ -5,12 +5,21 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations import dataclasses +from typing import TYPE_CHECKING 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.utils import data_allocation as data_alloc + + +if TYPE_CHECKING: + import gt4py.next.typing as gtx_typing + + from icon4py.model.common.grid import icon as icon_grid @dataclasses.dataclass(frozen=True) @@ -88,3 +97,26 @@ class AdvectionMetricState: #: vertical grid spacing at full levels ddqz_z_full: fa.CellKField[ta.wpfloat] + + +def initialize_advection_diagnostic_state( + grid: icon_grid.IconGrid, + allocator: gtx_typing.Allocator, +) -> AdvectionDiagnosticState: + return AdvectionDiagnosticState( + airmass_now=data_alloc.zero_field( + grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat + ), + airmass_new=data_alloc.zero_field( + grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat + ), + grf_tend_tracer=data_alloc.zero_field( + grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat + ), + hfl_tracer=data_alloc.zero_field( + grid, dims.EdgeDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat + ), + vfl_tracer=data_alloc.zero_field( + grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=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 af595436fa..4623eae3da 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 @@ -61,10 +61,11 @@ constants, dimension as dims, field_type_aliases as fa, + model_backends, + model_options, type_alias as ta, ) from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid -from icon4py.model.common.model_options import setup_program from icon4py.model.common.utils import data_allocation as data_alloc @@ -84,7 +85,7 @@ def run( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): + ) -> None: """ Set the vertical boundary conditions. @@ -115,7 +116,7 @@ def run( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): + ) -> None: log.debug("vertical boundary conditions computation - start") # set upper boundary conditions @@ -157,7 +158,7 @@ def limit_slope( z_slope: fa.CellKField[ta.wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): ... + ) -> None: ... @abc.abstractmethod def limit_parabola( @@ -168,14 +169,14 @@ def limit_parabola( p_face_low: fa.CellKField[ta.wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): ... + ) -> None: ... @abc.abstractmethod def limit_fluxes( self, horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): ... + ) -> None: ... class NoLimiter(VerticalLimiter): @@ -188,11 +189,14 @@ def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None) # fields self._l_limit = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, + dims.CellDim, + dims.KDim, + allocator=model_backends.get_allocator(self._backend), ) # stencils - self._copy_cell_kdim_field = setup_program( + self._copy_cell_kdim_field = model_options.setup_program( backend=backend, program=copy_cell_kdim_field, vertical_sizes={ @@ -201,7 +205,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None) }, offset_provider=self._grid.connectivities, ) - self._copy_cell_kdim_field_koff_plus1 = setup_program( + self._copy_cell_kdim_field_koff_plus1 = model_options.setup_program( backend=backend, program=copy_cell_kdim_field_koff_plus1, vertical_sizes={ @@ -217,7 +221,7 @@ def limit_slope( z_slope: fa.CellKField[ta.wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): ... + ) -> None: ... def limit_parabola( self, @@ -227,7 +231,7 @@ def limit_parabola( p_face_low: fa.CellKField[ta.wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): + ) -> None: # simply copy to up/low face values log.debug("running stencil copy_cell_kdim_field - start") self._copy_cell_kdim_field( @@ -251,7 +255,7 @@ def limit_fluxes( self, horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): ... + ) -> None: ... class SemiMonotonicLimiter(VerticalLimiter): @@ -263,15 +267,16 @@ def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None) self._backend = backend # fields + allocator = model_backends.get_allocator(self._backend) self._k_field = data_alloc.index_field( - self._grid, dims.KDim, extend={dims.KDim: 1}, dtype=gtx.int32, allocator=self._backend + self._grid, dims.KDim, extend={dims.KDim: 1}, dtype=gtx.int32, allocator=allocator ) # TODO(dastrm): should be KHalfDim self._l_limit = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, dtype=gtx.int32, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, dtype=gtx.int32, allocator=allocator ) # stencils - self._limit_vertical_slope_semi_monotonically = setup_program( + self._limit_vertical_slope_semi_monotonically = model_options.setup_program( backend=backend, program=limit_vertical_slope_semi_monotonically, constant_args={ @@ -284,7 +289,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None) }, offset_provider=self._grid.connectivities, ) - self._compute_vertical_parabola_limiter_condition = setup_program( + self._compute_vertical_parabola_limiter_condition = model_options.setup_program( backend=backend, program=compute_vertical_parabola_limiter_condition, vertical_sizes={ @@ -293,7 +298,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None) }, offset_provider=self._grid.connectivities, ) - self._limit_vertical_parabola_semi_monotonically = setup_program( + self._limit_vertical_parabola_semi_monotonically = model_options.setup_program( backend=backend, program=limit_vertical_parabola_semi_monotonically, vertical_sizes={ @@ -309,7 +314,7 @@ def limit_slope( z_slope: fa.CellKField[ta.wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): + ) -> None: log.debug("running stencil limit_vertical_slope_semi_monotonically - start") self._limit_vertical_slope_semi_monotonically( p_cc=p_tracer_now, @@ -327,7 +332,7 @@ def limit_parabola( p_face_low: fa.CellKField[ta.wpfloat], horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): + ) -> None: # compute semi-monotonic limiter condition log.debug("running stencil compute_vertical_parabola_limiter_condition - start") self._compute_vertical_parabola_limiter_condition( @@ -356,7 +361,7 @@ def limit_fluxes( self, horizontal_start: gtx.int32, horizontal_end: gtx.int32, - ): ... + ) -> None: ... class VerticalAdvection(abc.ABC): @@ -373,7 +378,7 @@ def run( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool = False, - ): + ) -> None: """ Run a vertical advection step. @@ -417,7 +422,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None) self._end_cell_end = self._grid.end_index(cell_domain(h_grid.Zone.END)) # stencils - self._copy_cell_kdim_field = setup_program( + self._copy_cell_kdim_field = model_options.setup_program( backend=backend, program=copy_cell_kdim_field, vertical_sizes={ @@ -429,7 +434,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None) log.debug("vertical advection class init - end") - def _get_horizontal_start_end(self, even_timestep: bool): + def _get_horizontal_start_end(self, even_timestep: bool) -> tuple[gtx.int32, gtx.int32]: if even_timestep: horizontal_start = self._start_cell_lateral_boundary_level_2 horizontal_end = self._end_cell_end @@ -449,7 +454,7 @@ def run( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool = False, - ): + ) -> None: log.debug("vertical advection run - start") horizontal_start, horizontal_end = self._get_horizontal_start_end( @@ -481,7 +486,7 @@ def run( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool = False, - ): + ) -> None: log.debug("vertical advection run - start") self._compute_numerical_flux( @@ -514,7 +519,7 @@ def _compute_numerical_flux( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool, - ): ... + ) -> None: ... @abc.abstractmethod def _update_unknowns( @@ -526,7 +531,7 @@ def _update_unknowns( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool, - ): ... + ) -> None: ... class FirstOrderUpwind(FiniteVolume): @@ -563,14 +568,13 @@ def __init__( self._k_field = data_alloc.index_field( self._grid, dims.KDim, - grid=self._grid, extend={dims.KDim: 1}, dtype=gtx.int32, - allocator=self._backend, + allocator=model_backends.get_allocator(self._backend), ) # TODO(dastrm): should be KHalfDim # stencils - self._compute_vertical_tracer_flux_upwind = setup_program( + self._compute_vertical_tracer_flux_upwind = model_options.setup_program( backend=backend, program=compute_vertical_tracer_flux_upwind, vertical_sizes={ @@ -582,7 +586,7 @@ def __init__( self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( self._backend ) - self._integrate_tracer_vertically = setup_program( + self._integrate_tracer_vertically = model_options.setup_program( backend=backend, program=integrate_tracer_vertically, constant_args={ @@ -601,7 +605,7 @@ def __init__( log.debug("vertical advection class init - end") - def _get_horizontal_start_end(self, even_timestep: bool): + def _get_horizontal_start_end(self, even_timestep: bool) -> tuple[gtx.int32, gtx.int32]: if even_timestep: horizontal_start = self._start_cell_lateral_boundary_level_2 horizontal_end = self._end_cell_end @@ -619,7 +623,7 @@ def _compute_numerical_flux( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool, - ): + ) -> None: log.debug("vertical numerical flux computation - start") horizontal_start, horizontal_end = self._get_horizontal_start_end( @@ -654,7 +658,7 @@ def _update_unknowns( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool, - ): + ) -> None: log.debug("vertical unknowns update - start") horizontal_start, horizontal_end = self._get_horizontal_start_end( @@ -708,30 +712,29 @@ def __init__( self._end_cell_end = self._grid.end_index(cell_domain(h_grid.Zone.END)) # fields + allocator = model_backends.get_allocator(self._backend) self._k_field = data_alloc.index_field( - self._grid, dims.KDim, extend={dims.KDim: 1}, dtype=gtx.int32, allocator=self._backend + self._grid, dims.KDim, extend={dims.KDim: 1}, dtype=gtx.int32, allocator=allocator ) # TODO(dastrm): should be KHalfDim self._z_cfl = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=allocator ) # TODO(dastrm): should be KHalfDim self._z_slope = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, allocator=allocator ) self._z_face = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=allocator ) # TODO(dastrm): should be KHalfDim self._z_face_up = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, allocator=allocator ) self._z_face_low = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, allocator=allocator ) self._z_delta_q = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend - ) - self._z_a1 = data_alloc.zero_field( - self._grid, dims.CellDim, dims.KDim, allocator=self._backend + self._grid, dims.CellDim, dims.KDim, allocator=allocator ) + self._z_a1 = data_alloc.zero_field(self._grid, dims.CellDim, dims.KDim, allocator=allocator) # misc self._slev = 0 @@ -742,7 +745,7 @@ def __init__( self._iadv_slev_jt = 0 # stencils - self._init_constant_cell_kdim_field = setup_program( + self._init_constant_cell_kdim_field = model_options.setup_program( backend=self._backend, program=init_constant_cell_kdim_field, constant_args={ @@ -755,7 +758,7 @@ def __init__( offset_provider=self._grid.connectivities, ) - self._compute_ppm4gpu_courant_number = setup_program( + self._compute_ppm4gpu_courant_number = model_options.setup_program( backend=self._backend, program=compute_ppm4gpu_courant_number, constant_args={ @@ -770,7 +773,7 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._compute_ppm_slope = setup_program( + self._compute_ppm_slope = model_options.setup_program( backend=self._backend, program=compute_ppm_slope, constant_args={ @@ -783,7 +786,7 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._compute_ppm_quadratic_face_values = setup_program( + self._compute_ppm_quadratic_face_values = model_options.setup_program( backend=self._backend, program=compute_ppm_quadratic_face_values, constant_args={ @@ -791,7 +794,7 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._compute_ppm_quartic_face_values = setup_program( + self._compute_ppm_quartic_face_values = model_options.setup_program( backend=self._backend, program=compute_ppm_quartic_face_values, constant_args={ @@ -803,7 +806,7 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._copy_cell_kdim_field = setup_program( + self._copy_cell_kdim_field = model_options.setup_program( backend=self._backend, program=copy_cell_kdim_field, vertical_sizes={ @@ -813,7 +816,7 @@ def __init__( offset_provider=self._grid.connectivities, ) - self._copy_cell_kdim_field_koff_minus1 = setup_program( + self._copy_cell_kdim_field_koff_minus1 = model_options.setup_program( backend=self._backend, program=copy_cell_kdim_field_koff_minus1, vertical_sizes={ @@ -823,7 +826,7 @@ def __init__( offset_provider=self._grid.connectivities, ) - self._compute_ppm4gpu_parabola_coefficients = setup_program( + self._compute_ppm4gpu_parabola_coefficients = model_options.setup_program( backend=self._backend, program=compute_ppm4gpu_parabola_coefficients, vertical_sizes={ @@ -832,7 +835,7 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._compute_ppm4gpu_fractional_flux = setup_program( + self._compute_ppm4gpu_fractional_flux = model_options.setup_program( backend=self._backend, program=compute_ppm4gpu_fractional_flux, constant_args={ @@ -845,7 +848,7 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._compute_ppm4gpu_integer_flux = setup_program( + self._compute_ppm4gpu_integer_flux = model_options.setup_program( backend=self._backend, program=compute_ppm4gpu_integer_flux, constant_args={ @@ -858,7 +861,7 @@ def __init__( }, offset_provider=self._grid.connectivities, ) - self._integrate_tracer_vertically = setup_program( + self._integrate_tracer_vertically = model_options.setup_program( backend=self._backend, program=integrate_tracer_vertically, constant_args={ @@ -877,7 +880,7 @@ def __init__( log.debug("vertical advection class init - end") - def _get_horizontal_start_end(self, even_timestep: bool): + def _get_horizontal_start_end(self, even_timestep: bool) -> tuple[gtx.int32, gtx.int32]: if even_timestep: horizontal_start = self._start_cell_lateral_boundary_level_2 horizontal_end = self._end_cell_end @@ -895,7 +898,7 @@ def _compute_numerical_flux( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool, - ): + ) -> None: log.debug("vertical numerical flux computation - start") horizontal_start, horizontal_end = self._get_horizontal_start_end( @@ -1076,7 +1079,7 @@ def _update_unknowns( p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO(dastrm): should be KHalfDim dtime: ta.wpfloat, even_timestep: bool, - ): + ) -> None: log.debug("vertical unknowns update - start") horizontal_start, horizontal_end = self._get_horizontal_start_end( diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm_all_face_values.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm_all_face_values.py index 196ee12b2c..c80b9614d1 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm_all_face_values.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm_all_face_values.py @@ -40,7 +40,7 @@ def _compute_ppm_all_face_values( p_face = concat_where(dims.KDim == elevp1, p_cc(Koff[-1]), p_face) - return p_face + return p_face # type: ignore[return-value] # concat_where leads to static type-erasure @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm_slope.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm_slope.py index af2b5f8992..e5e7927975 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm_slope.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm_slope.py @@ -58,7 +58,7 @@ def _compute_ppm_slope( _compute_ppm_slope_a(p_cc, p_cellhgt_mc_now), ) - return z_slope + return z_slope # type: ignore[return-value] # concat_where leads to static type erasure @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) diff --git a/model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics/integration_tests/test_single_moment_six_class_gscp_graupel.py b/model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics/integration_tests/test_single_moment_six_class_gscp_graupel.py index 6c571a0e54..bc36c53094 100644 --- a/model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics/integration_tests/test_single_moment_six_class_gscp_graupel.py +++ b/model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics/integration_tests/test_single_moment_six_class_gscp_graupel.py @@ -88,11 +88,7 @@ def test_graupel( qg=entry_savepoint.qg(), ) prognostic_state = prognostics.PrognosticState( - rho=entry_savepoint.rho(), - vn=None, - w=None, - exner=None, - theta_v=None, + rho=entry_savepoint.rho(), vn=None, w=None, exner=None, theta_v=None ) diagnostic_state = diagnostics.DiagnosticState( temperature=entry_savepoint.temperature(), diff --git a/model/common/src/icon4py/model/common/grid/states.py b/model/common/src/icon4py/model/common/grid/states.py index 37836c8dfc..97be2c4d1c 100644 --- a/model/common/src/icon4py/model/common/grid/states.py +++ b/model/common/src/icon4py/model/common/grid/states.py @@ -206,9 +206,9 @@ def __init__( @dataclasses.dataclass(frozen=True) class CellParams: #: Latitude at the cell center. The cell center is defined to be the circumcenter of a triangle. - cell_center_lat: fa.CellField[float] = None + cell_center_lat: fa.CellField[float] | None = None #: Longitude at the cell center. The cell center is defined to be the circumcenter of a triangle. - cell_center_lon: fa.CellField[float] = None + cell_center_lon: fa.CellField[float] | None = None #: Area of a cell, defined in ICON in mo_model_domain.f90:t_grid_cells%area - area: fa.CellField[float] = None - mean_cell_area: float = None + area: fa.CellField[float] | None = None + mean_cell_area: float | None = None diff --git a/model/common/src/icon4py/model/common/metrics/compute_advection_metrics.py b/model/common/src/icon4py/model/common/metrics/compute_advection_metrics.py new file mode 100644 index 0000000000..62ca8838fd --- /dev/null +++ b/model/common/src/icon4py/model/common/metrics/compute_advection_metrics.py @@ -0,0 +1,85 @@ +# 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 +from __future__ import annotations + +from gt4py import next as gtx + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta + + +@gtx.field_operator # type: ignore[call-overload] # see https://github.com/GridTools/gt4py/issues/2496 +def _compute_advection_deepatmo_fields( + height_u: fa.KField[ta.wpfloat], + height_l: fa.KField[ta.wpfloat], + grid_sphere_radius: float, +) -> tuple[fa.KField[ta.wpfloat], fa.KField[ta.wpfloat], fa.KField[ta.wpfloat]]: + """ + Compute 'deepatmo_divh', 'deepatmo_divzL', 'deepatmo_divzU' from 'vct_a' and 'grid_sphere_radius'. + + # Input Fields: + - height_u: expects vct_a[:nlev] + - height_l: expects vct_a[1:nlev+1] + + # Output Fields: + - deepatmo_divh + - deepatmo_divzL + - deepatmo_divzU + """ + height = 0.5 * (height_l + height_u) + radial_distance = height + grid_sphere_radius + radial_distance_l = grid_sphere_radius + height_l + radial_distance_u = grid_sphere_radius + height_u + deepatmo_gradh = grid_sphere_radius / radial_distance + deepatmo_divh = ( + deepatmo_gradh + * 3.0 + / 4.0 + / ( + 1.0 + - radial_distance_l * radial_distance_u / (radial_distance_l + radial_distance_u) ** 2 + ) + ) + deepatmo_divzL = 3.0 / ( + 1.0 + radial_distance_u / radial_distance_l + (radial_distance_u / radial_distance_l) ** 2 + ) + deepatmo_divzU = 3.0 / ( + 1.0 + radial_distance_l / radial_distance_u + (radial_distance_l / radial_distance_u) ** 2 + ) + return deepatmo_divh, deepatmo_divzL, deepatmo_divzU + + +@gtx.program # type: ignore[call-overload] # see https://github.com/GridTools/gt4py/issues/2496 +def compute_advection_deepatmo_fields( + height_u: fa.KField[ta.wpfloat], + height_l: fa.KField[ta.wpfloat], + deepatmo_divh: fa.KField[ta.wpfloat], + deepatmo_divzL: fa.KField[ta.wpfloat], + deepatmo_divzU: fa.KField[ta.wpfloat], + grid_sphere_radius: float, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +) -> None: + """ + Compute 'deepatmo_divh', 'deepatmo_divzL', 'deepatmo_divzU' from 'vct_a' and 'grid_sphere_radius'. + + # Input Fields: + - height_u: expects vct_a[:nlev] + - height_l: expects vct_a[1:nlev+1] + + # Inout Fields: + - deepatmo_divh + - deepatmo_divzL + - deepatmo_divzU + """ + _compute_advection_deepatmo_fields( + height_u=height_u, + height_l=height_l, + grid_sphere_radius=grid_sphere_radius, + out=(deepatmo_divh, deepatmo_divzL, deepatmo_divzU), + domain={dims.KDim: (vertical_start, vertical_end)}, + ) diff --git a/model/common/src/icon4py/model/common/metrics/metrics_attributes.py b/model/common/src/icon4py/model/common/metrics/metrics_attributes.py index 6979e50270..a822c35f52 100644 --- a/model/common/src/icon4py/model/common/metrics/metrics_attributes.py +++ b/model/common/src/icon4py/model/common/metrics/metrics_attributes.py @@ -68,6 +68,9 @@ ZD_INTCOEF_DSL: Final[str] = "zd_intcoef_dsl" ZD_VERTOFFSET_DSL: Final[str] = "zd_vertoffset_dsl" CELL_HEIGHT_ON_HALF_LEVEL: Final[str] = "vertical_coordinates_on_half_levels" +DEEPATMO_DIVH: Final[str] = "deepatmo_divh" +DEEPATMO_DIVZL: Final[str] = "deepatmo_divzL" +DEEPATMO_DIVZU: Final[str] = "deepatmo_divzU" attrs: dict[str, model.FieldMetaData] = { @@ -454,4 +457,28 @@ icon_var_name="z_ifc", dtype=ta.wpfloat, ), + DEEPATMO_DIVH: dict( + standard_name=DEEPATMO_DIVH, + long_name="", + units="", + dims=(dims.KDim), + icon_var_name="deepatmo_divh_mc", + dtype=ta.wpfloat, + ), + DEEPATMO_DIVZL: dict( + standard_name=DEEPATMO_DIVZL, + long_name="", + units="", + dims=(dims.KDim), + icon_var_name="deepatmo_divzL_mc", + dtype=ta.wpfloat, + ), + DEEPATMO_DIVZU: dict( + standard_name=DEEPATMO_DIVZU, + long_name="", + units="", + dims=(dims.KDim), + icon_var_name="deepatmo_divzU_mc", + dtype=ta.wpfloat, + ), } diff --git a/model/common/src/icon4py/model/common/metrics/metrics_factory.py b/model/common/src/icon4py/model/common/metrics/metrics_factory.py index 386ecbd3dc..44e16f6137 100644 --- a/model/common/src/icon4py/model/common/metrics/metrics_factory.py +++ b/model/common/src/icon4py/model/common/metrics/metrics_factory.py @@ -33,6 +33,7 @@ from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory from icon4py.model.common.interpolation.stencils import cell_2_edge_interpolation from icon4py.model.common.metrics import ( + compute_advection_metrics, compute_coeff_gradekin, compute_diffusion_metrics, compute_zdiff_gradp_dsl, @@ -128,6 +129,12 @@ def __init__( { "topography": topography, "vct_a": self._vertical_grid.interface_physical_height, + "height_u": self._vertical_grid.interface_physical_height[ + : self._grid.num_levels + ], + "height_l": self._vertical_grid.interface_physical_height[ + 1 : self._grid.num_levels + 1 + ], "c_refin_ctrl": c_refin_ctrl, "e_refin_ctrl": e_refin_ctrl, "e_owner_mask": e_owner_mask, @@ -951,6 +958,29 @@ def _register_computed_fields(self) -> None: # noqa: PLR0915 [too-many-statemen self.register_provider(compute_diffusion_intcoef_and_vertoffset) + compute_advection_deepatmo_fields = factory.ProgramFieldProvider( + func=compute_advection_metrics.compute_advection_deepatmo_fields, + domain={ + dims.KDim: ( + vertical_domain(v_grid.Zone.TOP), + vertical_domain(v_grid.Zone.BOTTOM), + ), + }, + fields={ + attrs.DEEPATMO_DIVH: attrs.DEEPATMO_DIVH, + attrs.DEEPATMO_DIVZL: attrs.DEEPATMO_DIVZL, + attrs.DEEPATMO_DIVZU: attrs.DEEPATMO_DIVZU, + }, + deps={ + "height_u": "height_u", + "height_l": "height_l", + }, + params={"grid_sphere_radius": constants.EARTH_RADIUS}, + do_exchange=False, + ) + + self.register_provider(compute_advection_deepatmo_fields) + @property def metadata(self) -> dict[str, model.FieldMetaData]: return self._attrs diff --git a/model/common/src/icon4py/model/common/states/factory.py b/model/common/src/icon4py/model/common/states/factory.py index c3d0f7656c..8fdbeaf1c0 100644 --- a/model/common/src/icon4py/model/common/states/factory.py +++ b/model/common/src/icon4py/model/common/states/factory.py @@ -156,6 +156,7 @@ class RetrievalType(enum.Enum): FIELD = 0 DATA_ARRAY = 1 METADATA = 2 + SCALAR = 3 class FieldSource(GridProvider, Protocol): @@ -195,7 +196,7 @@ def get( @overload def get( - self, field_name: str, type_: Literal[RetrievalType.FIELD] = RetrievalType.FIELD + self, field_name: str, type_: Literal[RetrievalType.SCALAR] = RetrievalType.SCALAR ) -> state_utils.ScalarType: ... @overload @@ -228,7 +229,7 @@ def get( match type_: case RetrievalType.METADATA: return self.metadata[field_name] - case RetrievalType.FIELD | RetrievalType.DATA_ARRAY: + case RetrievalType.FIELD | RetrievalType.DATA_ARRAY | RetrievalType.SCALAR: provider = self._providers[field_name] if field_name not in provider.fields: raise ValueError( @@ -238,7 +239,7 @@ def get( buffer = provider(field_name, self._sources, self.backend, self, self._exchange) return ( buffer - if type_ == RetrievalType.FIELD + if type_ in (RetrievalType.FIELD, RetrievalType.SCALAR) else state_utils.to_data_array(buffer, self.metadata[field_name]) ) case _: diff --git a/model/common/src/icon4py/model/common/states/prognostic_state.py b/model/common/src/icon4py/model/common/states/prognostic_state.py index 8506998e26..3e220e0d39 100644 --- a/model/common/src/icon4py/model/common/states/prognostic_state.py +++ b/model/common/src/icon4py/model/common/states/prognostic_state.py @@ -35,6 +35,9 @@ class PrognosticState: ] # horizontal wind normal to edges, vn(nproma, nlev, nblks_e) [m/s] exner: fa.CellKField[ta.wpfloat] # exner function, exner(nrpoma, nlev, nblks_c) theta_v: fa.CellKField[ta.wpfloat] # virtual temperature, (nproma, nlev, nlbks_c) [K] + tracer: list[fa.CellKField[ta.wpfloat]] = dataclasses.field( + default_factory=list + ) # tracer concentration (nproma,nlev,nblks_c,ntracer) [kg/kg] @property def w_1(self) -> fa.CellField[ta.wpfloat]: @@ -44,6 +47,7 @@ def w_1(self) -> fa.CellField[ta.wpfloat]: def initialize_prognostic_state( grid: icon_grid.IconGrid, allocator: gtx_typing.Allocator, + ntracer: int = 0, ) -> PrognosticState: """Initialize the prognostic state with zero fields.""" rho = data_alloc.zero_field( @@ -82,4 +86,14 @@ def initialize_prognostic_state( allocator=allocator, dtype=ta.wpfloat, ) - return PrognosticState(rho=rho, w=w, vn=vn, exner=exner, theta_v=theta_v) + tracer = [ + data_alloc.zero_field( + grid, + dims.CellDim, + dims.KDim, + allocator=allocator, + dtype=ta.wpfloat, + ) + for _ in range(ntracer) + ] + return PrognosticState(rho=rho, w=w, vn=vn, exner=exner, theta_v=theta_v, tracer=tracer) diff --git a/model/common/src/icon4py/model/common/type_alias.py b/model/common/src/icon4py/model/common/type_alias.py index 02b933bfa3..e8192d3605 100644 --- a/model/common/src/icon4py/model/common/type_alias.py +++ b/model/common/src/icon4py/model/common/type_alias.py @@ -16,6 +16,7 @@ wpfloat: TypeAlias = gtx.float64 vpfloat: type[gtx.float32] | type[gtx.float64] = wpfloat +anyfloat: TypeAlias = gtx.float32 | gtx.float64 precision = os.environ.get("FLOAT_PRECISION", DEFAULT_PRECISION).lower() diff --git a/model/common/tests/common/grid/unit_tests/test_topography.py b/model/common/tests/common/grid/unit_tests/test_topography.py index 26fb9b683c..0ed305d658 100644 --- a/model/common/tests/common/grid/unit_tests/test_topography.py +++ b/model/common/tests/common/grid/unit_tests/test_topography.py @@ -36,6 +36,9 @@ def test_topography_smoothing_with_serialized_data( backend: gtx_typing.Backend | None, ) -> None: cell_geometry = grid_savepoint.construct_cell_geometry() + assert ( + cell_geometry.area is not None + ), "Broken assumption: this test assumes it's running from a savepoint containing a 'cell_area' field." geofac_n2s = interpolation_savepoint.geofac_n2s() num_iterations = 25 diff --git a/model/common/tests/common/grid/unit_tests/test_vertical.py b/model/common/tests/common/grid/unit_tests/test_vertical.py index b9df2e773d..a37ba4a436 100644 --- a/model/common/tests/common/grid/unit_tests/test_vertical.py +++ b/model/common/tests/common/grid/unit_tests/test_vertical.py @@ -400,6 +400,7 @@ def test_compute_vertical_coordinate( ) assert vertical_geometry.nflatlev == grid_savepoint.nflatlev() + topography = None if experiment in (definitions.Experiments.MCH_CH_R04B09, definitions.Experiments.GAUSS3D): topography = topography_savepoint.topo_c() elif experiment == definitions.Experiments.EXCLAIM_APE: @@ -409,6 +410,9 @@ def test_compute_vertical_coordinate( geofac_n2s = interpolation_savepoint.geofac_n2s() + assert cell_geometry.area is not None + assert topography is not None + vertical_coordinates_on_half_levels = v_grid.compute_vertical_coordinate( vct_a=vct_a.ndarray, topography=topography.ndarray, diff --git a/model/driver/src/icon4py/model/driver/testcases/gauss3d.py b/model/driver/src/icon4py/model/driver/testcases/gauss3d.py index c0cc32da25..d308283548 100644 --- a/model/driver/src/icon4py/model/driver/testcases/gauss3d.py +++ b/model/driver/src/icon4py/model/driver/testcases/gauss3d.py @@ -250,18 +250,10 @@ def model_initialization_gauss3d( # noqa: PLR0915 [too-many-statements] ) prognostic_state_now = prognostics.PrognosticState( - w=w, - vn=vn, - theta_v=theta_v, - rho=rho, - exner=exner, + w=w, vn=vn, theta_v=theta_v, rho=rho, exner=exner ) prognostic_state_next = prognostics.PrognosticState( - w=w_next, - vn=vn_next, - theta_v=theta_v_next, - rho=rho_next, - exner=exner_next, + w=w_next, vn=vn_next, theta_v=theta_v_next, rho=rho_next, exner=exner_next ) diffusion_diagnostic_state = diffusion_states.initialize_diffusion_diagnostic_state( diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/config.py b/model/standalone_driver/src/icon4py/model/standalone_driver/config.py index b016af5576..7d0cbe7453 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/config.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/config.py @@ -24,6 +24,12 @@ class ProfilingStats: @dataclasses.dataclass(frozen=True) class DriverConfig: + """ + Standalone driver configuration. + + Default values should correspond to default values in ICON. + """ + experiment_name: str output_path: pathlib.Path profiling_stats: ProfilingStats | None @@ -34,3 +40,4 @@ class DriverConfig: vertical_cfl_threshold: ta.wpfloat = 0.85 ndyn_substeps: int = 5 enable_statistics_output: bool = False + ntracer: int = 0 diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py index 648cd46456..473f176d74 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py @@ -17,6 +17,7 @@ import devtools import icon4py.model.common.utils as common_utils +from icon4py.model.atmosphere.advection import advection_states from icon4py.model.atmosphere.diffusion import diffusion_states from icon4py.model.atmosphere.dycore import dycore_states from icon4py.model.common import type_alias as ta @@ -62,6 +63,8 @@ class DriverStates(NamedTuple): prep_advection_prognostic: dycore_states.PrepAdvection solve_nonhydro_diagnostic: dycore_states.DiagnosticStateNonHydro + tracer_advection_diagnostic: advection_states.AdvectionDiagnosticState + prep_tracer_advection_prognostic: advection_states.AdvectionPrepAdvState diffusion_diagnostic: diffusion_states.DiffusionDiagnosticState prognostics: common_utils.TimeStepPair[prognostics.PrognosticState] diagnostic: diagnostics.DiagnosticState diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py index a2b66ccce3..818c0de177 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py @@ -13,11 +13,12 @@ import sys import time from types import ModuleType -from typing import Any +from typing import Any, Literal import gt4py.next as gtx import gt4py.next.typing as gtx_typing +from icon4py.model.atmosphere.advection import advection, advection_states from icon4py.model.atmosphere.diffusion import diffusion, diffusion_states from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as solve_nh from icon4py.model.common import ( @@ -25,7 +26,7 @@ dimension as dims, field_type_aliases as fa, model_backends, - model_options, + type_alias as ta, ) from icon4py.model.common.decomposition import ( definitions as decomposition_defs, @@ -41,6 +42,7 @@ ) from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory from icon4py.model.common.metrics import metrics_attributes, metrics_factory +from icon4py.model.common.states import factory as states_factory from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.standalone_driver import config as driver_config, driver_states @@ -110,14 +112,13 @@ def create_static_field_factories( grid_manager: gm.GridManager, decomposition_info: decomposition_defs.DecompositionInfo, vertical_grid: v_grid.VerticalGrid, - cell_topography: data_alloc.NDArray, - backend: model_backends.BackendLike, + cell_topography: fa.CellField[ta.wpfloat], + backend: gtx_typing.Backend | None, ) -> driver_states.StaticFieldFactories: - concrete_backend = model_options.customize_backend(program=None, backend=backend) geometry_field_source = grid_geometry.GridGeometry( grid=grid_manager.grid, decomposition_info=decomposition_info, - backend=concrete_backend, + backend=backend, coordinates=grid_manager.coordinates, extra_fields=grid_manager.geometry_fields, metadata=geometry_meta.attrs, @@ -127,7 +128,7 @@ def create_static_field_factories( grid=grid_manager.grid, decomposition_info=decomposition_info, geometry_source=geometry_field_source, - backend=concrete_backend, + backend=backend, metadata=interpolation_attributes.attrs, ) @@ -138,7 +139,7 @@ def create_static_field_factories( geometry_source=geometry_field_source, topography=cell_topography, interpolation_source=interpolation_field_source, - backend=concrete_backend, + backend=backend, metadata=metrics_attributes.attrs, rayleigh_type=constants.RayleighType.KLEMP, rayleigh_coeff=0.1, @@ -158,14 +159,12 @@ def initialize_granules( vertical_grid: v_grid.VerticalGrid, diffusion_config: diffusion.DiffusionConfig, solve_nh_config: solve_nh.NonHydrostaticConfig, + advection_config: advection.AdvectionConfig, static_field_factories: driver_states.StaticFieldFactories, exchange: decomposition_defs.ExchangeRuntime, owner_mask: fa.CellField[bool], - backend: model_backends.BackendLike, -) -> tuple[ - diffusion.Diffusion, - solve_nh.SolveNonhydro, -]: + backend: gtx_typing.Backend | None, +) -> tuple[diffusion.Diffusion, solve_nh.SolveNonhydro, advection.Advection]: geometry_field_source = static_field_factories.geometry_field_source interpolation_field_source = static_field_factories.interpolation_field_source metrics_field_source = static_field_factories.metrics_field_source @@ -175,7 +174,9 @@ def initialize_granules( cell_center_lat=geometry_field_source.get(geometry_meta.CELL_LAT), cell_center_lon=geometry_field_source.get(geometry_meta.CELL_LON), area=geometry_field_source.get(geometry_meta.CELL_AREA), - mean_cell_area=geometry_field_source.get(geometry_meta.MEAN_CELL_AREA), + mean_cell_area=geometry_field_source.get( + geometry_meta.MEAN_CELL_AREA, states_factory.RetrievalType.SCALAR + ), ) log.info("creating edge geometry") @@ -343,7 +344,42 @@ def initialize_granules( owner_mask=owner_mask, ) - return diffusion_granule, solve_nonhydro_granule + advection_granule = advection.convert_config_to_advection( + grid=grid, + backend=backend, + config=advection_config, + interpolation_state=advection_states.AdvectionInterpolationState( + geofac_div=interpolation_field_source.get(interpolation_attributes.GEOFAC_DIV), + rbf_vec_coeff_e=interpolation_field_source.get( + interpolation_attributes.RBF_VEC_COEFF_E + ), + pos_on_tplane_e_1=interpolation_field_source.get( + interpolation_attributes.POS_ON_TPLANE_E_X + ), + pos_on_tplane_e_2=interpolation_field_source.get( + interpolation_attributes.POS_ON_TPLANE_E_Y + ), + ), + least_squares_state=advection_states.AdvectionLeastSquaresState( + lsq_pseudoinv_1=interpolation_field_source.get(interpolation_attributes.LSQ_PSEUDOINV)[ + :, 0, : + ], + lsq_pseudoinv_2=interpolation_field_source.get(interpolation_attributes.LSQ_PSEUDOINV)[ + :, 1, : + ], + ), + metric_state=advection_states.AdvectionMetricState( + deepatmo_divh=metrics_field_source.get(metrics_attributes.DEEPATMO_DIVH), + deepatmo_divzl=metrics_field_source.get(metrics_attributes.DEEPATMO_DIVZL), + deepatmo_divzu=metrics_field_source.get(metrics_attributes.DEEPATMO_DIVZU), + ddqz_z_full=metrics_field_source.get(metrics_attributes.DDQZ_Z_FULL), + ), + edge_params=edge_geometry, + cell_params=cell_geometry, + exchange=exchange, + ) + + return diffusion_granule, solve_nonhydro_granule, advection_granule def find_maximum_from_field( @@ -353,7 +389,7 @@ def find_maximum_from_field( array_ns.abs(input_field.ndarray).argmax(), input_field.ndarray.shape, ) - return max_indices, input_field.ndarray[max_indices] + return max_indices, input_field.ndarray[max_indices] # type: ignore[return-value] ## this is congruent with observed numpy behavior def display_icon4py_logo_in_log_file() -> None: @@ -412,13 +448,13 @@ def display_icon4py_logo_in_log_file() -> None: for _ in range(3): icon4py_signature += empty_line icon4py_signature += boundary_line - icon4py_signature = "\n".join(icon4py_signature) - log.info(f"{icon4py_signature}") + icon4py_signature_str = "\n".join(icon4py_signature) + log.info(icon4py_signature_str) def display_driver_setup_in_log_file( n_time_steps: int, - vertical_params, + vertical_params: v_grid.VerticalGrid, config: driver_config.DriverConfig, ) -> None: log.info("===== ICON4Py Driver Configuration =====") @@ -431,6 +467,7 @@ def display_driver_setup_in_log_file( log.info(f"Vertical CFL threshold : {config.vertical_cfl_threshold}") log.info(f"Second-order divdamp : {config.apply_extra_second_order_divdamp}") log.info(f"Statistics enabled : {config.enable_statistics_output}") + log.info(f"Number of tracers : {config.ntracer}") log.info("") log.info("==== Vertical Grid Parameters ====") @@ -445,14 +482,14 @@ def display_driver_setup_in_log_file( @dataclasses.dataclass class _InfoFormatter(logging.Formatter): - style: str + style: Literal["%", "{", "$"] default_fmt: str info_fmt: str defaults: dict[str, Any] | None _info_formatter: logging.Formatter = dataclasses.field(init=False) - def __post_init__(self): + def __post_init__(self) -> None: super().__init__(fmt=self.default_fmt, style=self.style, defaults=self.defaults) self._info_formatter = logging.Formatter( fmt=self.info_fmt, @@ -489,7 +526,7 @@ def make_handler( def configure_logging( logging_level: str, - processor_procs: decomposition_defs.ProcessProperties = None, + processor_procs: decomposition_defs.ProcessProperties | None = None, ) -> None: """ Configure logging. diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/main.py b/model/standalone_driver/src/icon4py/model/standalone_driver/main.py index e0d9efd307..df22513ad6 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/main.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/main.py @@ -6,6 +6,7 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import logging +import pathlib from typing import Annotated import typer @@ -19,8 +20,10 @@ def main( - configuration_file_path: Annotated[str, typer.Argument(help="Configuration file path.")], - grid_file_path: Annotated[str, typer.Option(help="Grid file path.")], + configuration_file_path: Annotated[ + pathlib.Path, typer.Argument(help="Configuration file path.") + ], + grid_file_path: Annotated[pathlib.Path, typer.Option(help="Grid file path.")], # it may be better to split device from backend, # or only asking for cpu or gpu and the best backend for perfornamce is handled inside icon4py, # whether to automatically use gpu if cupy is installed can be discussed further @@ -31,8 +34,8 @@ def main( ), ], output_path: Annotated[ - str, typer.Option(help="Folder path that holds the output and log files.") - ] = "./output", + pathlib.Path, typer.Option(help="Folder path that holds the output and log files.") + ] = pathlib.Path("./output"), log_level: Annotated[ str, typer.Option( @@ -78,10 +81,5 @@ def main( log.info("time loop: DONE") -def click(): - """Entry point for the standalone driver CLI.""" - typer.run(main) - - if __name__ == "__main__": - click() + typer.run(main) diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/py.typed b/model/standalone_driver/src/icon4py/model/standalone_driver/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py b/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py index 9836e8d7f6..da7474f905 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py @@ -10,6 +10,7 @@ import functools import logging import pathlib +import types from collections.abc import Callable import gt4py.next as gtx @@ -17,6 +18,7 @@ from gt4py.next.instrumentation import metrics as gtx_metrics import icon4py.model.common.utils as common_utils +from icon4py.model.atmosphere.advection import advection, advection_states from icon4py.model.atmosphere.diffusion import diffusion, diffusion_states from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as solve_nh from icon4py.model.common import dimension as dims, model_backends, model_options, type_alias as ta @@ -42,11 +44,12 @@ class Icon4pyDriver: def __init__( self, config: driver_config.DriverConfig, - backend: model_backends.BackendLike, + backend: gtx.typing.Backend | None, grid: IconGrid, static_field_factories: driver_states.StaticFieldFactories, diffusion_granule: diffusion.Diffusion, solve_nonhydro_granule: solve_nh.SolveNonhydro, + tracer_advection_granule: advection.Advection, ): self.config = config self.backend = backend @@ -55,6 +58,7 @@ def __init__( self.diffusion = diffusion_granule self.solve_nonhydro = solve_nonhydro_granule self.model_time_variables = driver_states.ModelTimeVariables(config=config) + self.tracer_advection = tracer_advection_granule self.timer_collection = driver_states.TimerCollection( [timer.value for timer in driver_states.DriverTimers] ) @@ -66,17 +70,13 @@ def __init__( ) @functools.cached_property - def _allocator(self): + def _allocator(self) -> gtx.typing.Backend: return model_backends.get_allocator(self.backend) @functools.cached_property - def _xp(self): + def _xp(self) -> types.ModuleType: return data_alloc.import_array_ns(self._allocator) - @functools.cached_property - def _concrete_backend(self): - return model_options.customize_backend(program=None, backend=self.backend) - def _is_last_substep(self, step_nr: int) -> bool: return step_nr == (self.model_time_variables.ndyn_substeps_var - 1) @@ -94,8 +94,10 @@ def time_integration( ) -> None: diffusion_diagnostic_state = ds.diffusion_diagnostic solve_nonhydro_diagnostic_state = ds.solve_nonhydro_diagnostic + tracer_advection_diagnostic_state = ds.tracer_advection_diagnostic prognostic_states = ds.prognostics prep_adv = ds.prep_advection_prognostic + tracer_prep_adv = ds.prep_tracer_advection_prognostic log.debug( f"starting time loop for dtime = {self.model_time_variables.dtime_in_seconds} s, substep_timestep = {self.model_time_variables.substep_timestep} s, n_timesteps = {self.model_time_variables.n_time_steps}" @@ -123,11 +125,13 @@ def time_integration( self._integrate_one_time_step( diffusion_diagnostic_state, solve_nonhydro_diagnostic_state, + tracer_advection_diagnostic_state, prognostic_states, prep_adv, do_prep_adv, + tracer_prep_adv, ) - device_utils.sync(self._concrete_backend) + device_utils.sync(self.backend) self.model_time_variables.is_first_step_in_simulation = False @@ -149,10 +153,12 @@ def _integrate_one_time_step( self, diffusion_diagnostic_state: diffusion_states.DiffusionDiagnosticState, solve_nonhydro_diagnostic_state: dycore_states.DiagnosticStateNonHydro, + tracer_advection_diagnostic_state: advection_states.AdvectionDiagnosticState, prognostic_states: common_utils.TimeStepPair[prognostics.PrognosticState], prep_adv: dycore_states.PrepAdvection, do_prep_adv: bool, - ): + tracer_prep_adv: advection_states.AdvectionPrepAdvState, + ) -> None: log.debug(f"Running {self.solve_nonhydro.__class__}") self._do_dyn_substepping( solve_nonhydro_diagnostic_state, @@ -176,6 +182,17 @@ def _integrate_one_time_step( ) timer_diffusion.capture() + # TODO(ricoh): [c34] optionally move the loop into the granule (for efficiency gains) + # Precondition: passing data test with ntracer > 0 + for tracer_idx in range(self.config.ntracer): + self.tracer_advection.run( + diagnostic_state=tracer_advection_diagnostic_state, + prep_adv=tracer_prep_adv, + p_tracer_now=prognostic_states.current.tracer[tracer_idx], + p_tracer_new=prognostic_states.next.tracer[tracer_idx], + dtime=self.model_time_variables.dtime_in_seconds, + ) + prognostic_states.swap() def _update_time_levels_for_velocity_tendencies( @@ -183,7 +200,7 @@ def _update_time_levels_for_velocity_tendencies( diagnostic_state_nh: dycore_states.DiagnosticStateNonHydro, at_first_substep: bool, at_initial_timestep: bool, - ): + ) -> None: """ Set time levels of advective tendency fields for call to velocity_tendencies. @@ -223,7 +240,7 @@ def _do_dyn_substepping( prognostic_states: common_utils.TimeStepPair[prognostics.PrognosticState], prep_adv: dycore_states.PrepAdvection, do_prep_adv: bool, - ): + ) -> None: # TODO(OngChia): compute airmass for prognostic_state here timer_solve_nh = ( @@ -457,7 +474,7 @@ def _compute_mean_at_final_time_step( log.info( f"{interface_physical_height_ndarray[k]:12.3f}: {self._xp.mean(rho_ndarray[:, k]):.5e} " f"{self._xp.mean(vn_ndarray[:, k]):.5e} " - f"{self._xp.mean(w_ndarray[:, k+1]):.5e} " + f"{self._xp.mean(w_ndarray[:, k + 1]):.5e} " f"{self._xp.mean(theta_v_ndarray[:, k]):.5e} " f"{self._xp.mean(exner_ndarray[:, k]):.5e} " ) @@ -471,6 +488,7 @@ def _read_config( driver_config.DriverConfig, v_grid.VerticalGridConfig, diffusion.DiffusionConfig, + advection.AdvectionConfig, solve_nh.NonHydrostaticConfig, ]: vertical_grid_config = v_grid.VerticalGridConfig( @@ -493,6 +511,15 @@ def _read_config( velocity_boundary_diffusion_denom=200.0, ) + # NOTE(ricoh): adjust when switching experiments! + # These are ICON defaults, irrelevant for Jablonowski_Williamson (no tracers) + advection_config = advection.AdvectionConfig( + horizontal_advection_limiter=advection.HorizontalAdvectionLimiter.POSITIVE_DEFINITE, + horizontal_advection_type=advection.HorizontalAdvectionType.LINEAR_2ND_ORDER, + vertical_advection_limiter=advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + vertical_advection_type=advection.VerticalAdvectionType.PPM_3RD_ORDER, + ) + nonhydro_config = solve_nh.NonHydrostaticConfig( fourth_order_divdamp_factor=0.0025, ) @@ -515,6 +542,7 @@ def _read_config( icon4py_driver_config, vertical_grid_config, diffusion_config, + advection_config, nonhydro_config, ) @@ -552,27 +580,29 @@ def initialize_driver( processor_procs=parallel_props, ) - configuration_file_path = pathlib.Path(configuration_file_path) global_reductions = decomposition_defs.create_reduction(parallel_props) - grid_file_path = pathlib.Path(grid_file_path) - if pathlib.Path(output_path).exists(): + if output_path.exists(): current_time = datetime.datetime.now() log.warning(f"output path {output_path} already exists, a time stamp will be added") - output_path = pathlib.Path( - output_path - + f"_{datetime.date.today()}_{current_time.hour}h_{current_time.minute}m_{current_time.second}s" + output_path = ( + output_path.parent + / f"{output_path.name}_{datetime.date.today()}_{current_time.hour}h_{current_time.minute}m_{current_time.second}s" ) + else: - output_path = pathlib.Path(output_path) - output_path.mkdir(parents=True, exist_ok=False) + output_path.mkdir(parents=True, exist_ok=False) - backend = driver_utils.get_backend_from_name(backend_name) + backend = model_options.customize_backend( + program=None, backend=driver_utils.get_backend_from_name(backend_name) + ) allocator = model_backends.get_allocator(backend) log.info("Initializing the driver") - driver_config, vertical_grid_config, diffusion_config, solve_nh_config = _read_config( - output_path=output_path, - enable_profiling=False, + driver_config, vertical_grid_config, diffusion_config, advection_config, solve_nh_config = ( + _read_config( + output_path=output_path, + enable_profiling=False, + ) ) log.info(f"initializing the grid manager from '{grid_file_path}'") @@ -609,7 +639,7 @@ def initialize_driver( grid_manager=grid_manager, decomposition_info=decomposition_info, vertical_grid=vertical_grid, - cell_topography=gtx.as_field((dims.CellDim,), data=cell_topography, allocator=allocator), + cell_topography=gtx.as_field((dims.CellDim,), data=cell_topography, allocator=allocator), # type: ignore[arg-type] # due to array_ns opacity backend=backend, ) @@ -617,17 +647,19 @@ def initialize_driver( ( diffusion_granule, solve_nonhydro_granule, + tracer_advection_granule, ) = driver_utils.initialize_granules( grid=grid_manager.grid, vertical_grid=vertical_grid, diffusion_config=diffusion_config, solve_nh_config=solve_nh_config, + advection_config=advection_config, static_field_factories=static_field_factories, exchange=exchange, owner_mask=gtx.as_field( (dims.CellDim,), - decomposition_info.owner_mask(dims.CellDim), - allocator=allocator, # type: ignore[arg-type] + decomposition_info.owner_mask(dims.CellDim), # type: ignore[arg-type] # due to array_ns opacity + allocator=allocator, ), backend=backend, ) @@ -638,6 +670,7 @@ def initialize_driver( static_field_factories=static_field_factories, diffusion_granule=diffusion_granule, solve_nonhydro_granule=solve_nonhydro_granule, + tracer_advection_granule=tracer_advection_granule, ) return icon4py_driver diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/testcases/initial_condition.py b/model/standalone_driver/src/icon4py/model/standalone_driver/testcases/initial_condition.py index 769719c688..163618e363 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/testcases/initial_condition.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/testcases/initial_condition.py @@ -9,14 +9,16 @@ import logging import math +from gt4py import next as gtx + import icon4py.model.common.utils as common_utils +from icon4py.model.atmosphere.advection import advection_states from icon4py.model.atmosphere.diffusion import diffusion_states from icon4py.model.atmosphere.dycore import dycore_states from icon4py.model.common import ( constants as phy_const, dimension as dims, model_backends, - model_options, type_alias as ta, ) from icon4py.model.common.grid import ( @@ -49,7 +51,7 @@ def jablonowski_williamson( # noqa: PLR0915 [too-many-statements] geometry_field_source: grid_geometry.GridGeometry, interpolation_field_source: interpolation_factory.InterpolationFieldsFactory, metrics_field_source: metrics_factory.MetricsFieldsFactory, - backend: model_backends.BackendLike, + backend: gtx.typing.Backend | None, ) -> driver_states.DriverStates: """ Initial condition of Jablonowski-Williamson test. Set jw_baroclinic_amplitude to values larger than 0.01 if @@ -62,10 +64,12 @@ def jablonowski_williamson( # noqa: PLR0915 [too-many-statements] metrics_field_source: metric field factory backend: GT4Py backend Returns: driver state + + The reference experiment config for this is in icon-exclaim/run/exp.exclaim_nh35_tri_jws_sb. """ - concrete_backend = model_options.customize_backend(program=None, backend=backend) - xp = data_alloc.import_array_ns(concrete_backend.allocator) + allocator = model_backends.get_allocator(backend) + xp = data_alloc.import_array_ns(allocator) wgtfac_c = metrics_field_source.get(metrics_attributes.WGTFAC_C).ndarray ddqz_z_half = metrics_field_source.get(metrics_attributes.DDQZ_Z_HALF).ndarray @@ -118,21 +122,18 @@ def jablonowski_williamson( # noqa: PLR0915 [too-many-statements] # Initialize prognostic state, diagnostic state and other local fields prognostic_state_now = prognostics.initialize_prognostic_state( - grid=grid, allocator=concrete_backend.allocator - ) - diagnostic_state = diagnostics.initialize_diagnostic_state( - grid=grid, allocator=concrete_backend.allocator + grid=grid, + allocator=allocator, ) + diagnostic_state = diagnostics.initialize_diagnostic_state(grid=grid, allocator=allocator) eta_v = data_alloc.zero_field( grid, dims.CellDim, dims.KDim, - allocator=concrete_backend.allocator, + allocator=allocator, dtype=ta.wpfloat, ) - eta_v_at_edge = data_alloc.zero_field( - grid, dims.EdgeDim, dims.KDim, allocator=concrete_backend.allocator - ) + eta_v_at_edge = data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, allocator=allocator) exner_ndarray = prognostic_state_now.exner.ndarray rho_ndarray = prognostic_state_now.rho.ndarray @@ -224,7 +225,7 @@ def jablonowski_williamson( # noqa: PLR0915 [too-many-statements] temperature_ndarray[:, k_index] = temperature_jw log.info("Newton iteration completed.") - cell_2_edge_interpolation.cell_2_edge_interpolation.with_backend(concrete_backend)( + cell_2_edge_interpolation.cell_2_edge_interpolation.with_backend(backend)( in_field=eta_v, coeff=cell_2_edge_coeff, out_field=eta_v_at_edge, @@ -267,19 +268,15 @@ def jablonowski_williamson( # noqa: PLR0915 [too-many-statements] log.info("Hydrostatic adjustment computation completed.") prognostic_state_next = prognostics.PrognosticState( - vn=data_alloc.as_field(prognostic_state_now.vn, allocator=concrete_backend.allocator), - w=data_alloc.as_field(prognostic_state_now.w, allocator=concrete_backend.allocator), - exner=data_alloc.as_field(prognostic_state_now.exner, allocator=concrete_backend.allocator), - rho=data_alloc.as_field(prognostic_state_now.rho, allocator=concrete_backend.allocator), - theta_v=data_alloc.as_field( - prognostic_state_now.theta_v, allocator=concrete_backend.allocator - ), + vn=data_alloc.as_field(prognostic_state_now.vn, allocator=allocator), + w=data_alloc.as_field(prognostic_state_now.w, allocator=allocator), + exner=data_alloc.as_field(prognostic_state_now.exner, allocator=allocator), + rho=data_alloc.as_field(prognostic_state_now.rho, allocator=allocator), + theta_v=data_alloc.as_field(prognostic_state_now.theta_v, allocator=allocator), ) prognostic_states = common_utils.TimeStepPair(prognostic_state_now, prognostic_state_next) - edge_2_cell_vector_rbf_interpolation.edge_2_cell_vector_rbf_interpolation.with_backend( - concrete_backend - )( + edge_2_cell_vector_rbf_interpolation.edge_2_cell_vector_rbf_interpolation.with_backend(backend)( p_e_in=prognostic_states.current.vn, ptr_coeff_1=rbf_vec_coeff_c1, ptr_coeff_2=rbf_vec_coeff_c2, @@ -294,10 +291,8 @@ def jablonowski_williamson( # noqa: PLR0915 [too-many-statements] log.info("U, V computation completed.") - perturbed_exner = data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, allocator=concrete_backend.allocator - ) - gt4py_math_op.compute_difference_on_cell_k.with_backend(concrete_backend)( + perturbed_exner = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, allocator=allocator) + gt4py_math_op.compute_difference_on_cell_k.with_backend(backend)( field_a=prognostic_states.current.exner, field_b=metrics_field_source.get(metrics_attributes.EXNER_REF_MC), output_field=perturbed_exner, @@ -310,21 +305,29 @@ def jablonowski_williamson( # noqa: PLR0915 [too-many-statements] log.info("perturbed_exner initialization completed.") diffusion_diagnostic_state = diffusion_states.initialize_diffusion_diagnostic_state( - grid=grid, allocator=concrete_backend.allocator + grid=grid, allocator=allocator ) solve_nonhydro_diagnostic_state = dycore_states.initialize_solve_nonhydro_diagnostic_state( perturbed_exner_at_cells_on_model_levels=perturbed_exner, grid=grid, - allocator=concrete_backend.allocator, + allocator=allocator, + ) + prep_adv = dycore_states.initialize_prep_advection(grid=grid, allocator=allocator) + tracer_advection_diagnostic_state = advection_states.initialize_advection_diagnostic_state( + grid=grid, allocator=allocator ) - prep_adv = dycore_states.initialize_prep_advection( - grid=grid, allocator=concrete_backend.allocator + prep_tracer_adv = advection_states.AdvectionPrepAdvState( + vn_traj=data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, allocator=allocator), + mass_flx_me=data_alloc.zero_field(grid, dims.EdgeDim, dims.KDim, allocator=allocator), + mass_flx_ic=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, allocator=allocator), ) log.info("Initialization completed.") ds = driver_states.DriverStates( prep_advection_prognostic=prep_adv, solve_nonhydro_diagnostic=solve_nonhydro_diagnostic_state, + prep_tracer_advection_prognostic=prep_tracer_adv, + tracer_advection_diagnostic=tracer_advection_diagnostic_state, diffusion_diagnostic=diffusion_diagnostic_state, prognostics=prognostic_states, diagnostic=diagnostic_state, diff --git a/model/standalone_driver/tests/standalone_driver/integration_tests/test_standalone_driver.py b/model/standalone_driver/tests/standalone_driver/integration_tests/test_standalone_driver.py index bd38e6dce9..44d3212fa5 100644 --- a/model/standalone_driver/tests/standalone_driver/integration_tests/test_standalone_driver.py +++ b/model/standalone_driver/tests/standalone_driver/integration_tests/test_standalone_driver.py @@ -17,15 +17,15 @@ @pytest.mark.embedded_remap_error def test_standalone_driver( - backend_like, + backend_like: model_backends.BackendLike, tmp_path: pathlib.Path, -): +) -> None: """ Currently, this is a only test to check if the driver runs from a grid file without verifying the end result. TODO(anyone): Modify this test for scientific validation after IO is ready. """ - backend_name = None + backend_name = "embedded" for k, v in model_backends.BACKENDS.items(): if backend_like == v: backend_name = k @@ -34,8 +34,8 @@ def test_standalone_driver( output_path = tmp_path / f"ci_driver_output_for_backend_{backend_name}" main.main( - configuration_file_path="./", + configuration_file_path=pathlib.Path("./"), grid_file_path=grid_file_path, icon4py_backend=backend_name, - output_path=str(output_path), + output_path=output_path, ) diff --git a/pyproject.toml b/pyproject.toml index e349356eb7..d617268d68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,8 @@ test = [ typing = [ "mypy[faster-cache]>=1.13.0", "typing-extensions>=4.11.0", - "types-cffi>=1.16.0" + "types-cffi>=1.16.0", + "scipy-stubs>=1.15.3.0" ] # -- Standard project description options (PEP 621) -- @@ -135,7 +136,7 @@ disallow_untyped_defs = true explicit_package_bases = true files = [ # TODO(egparedes): activate for all sources and tests - # '$MYPY_CONFIG_FILE_DIR/model/atmosphere/advection/src', + '$MYPY_CONFIG_FILE_DIR/model/atmosphere/advection/src', # '$MYPY_CONFIG_FILE_DIR/model/atmosphere/advection/tests', # '$MYPY_CONFIG_FILE_DIR/model/atmosphere/diffusion/src', # '$MYPY_CONFIG_FILE_DIR/model/atmosphere/diffusion/tests', @@ -157,11 +158,14 @@ files = [ '$MYPY_CONFIG_FILE_DIR/model/common/src/icon4py/model/common/initialization', '$MYPY_CONFIG_FILE_DIR/model/common/src/icon4py/model/common/io', '$MYPY_CONFIG_FILE_DIR/model/common/src/icon4py/model/common/components', + '$MYPY_CONFIG_FILE_DIR/model/common/src/icon4py/model/common/metrics/compute_advection_metrics.py', # '$MYPY_CONFIG_FILE_DIR/model/driver/src/', # '$MYPY_CONFIG_FILE_DIR/model/driver/tests/', '$MYPY_CONFIG_FILE_DIR/model/testing/src/', # '$MYPY_CONFIG_FILE_DIR/model/testing/tests/', - '$MYPY_CONFIG_FILE_DIR/tools/src/' + '$MYPY_CONFIG_FILE_DIR/tools/src/', + '$MYPY_CONFIG_FILE_DIR/model/standalone_driver/src', + '$MYPY_CONFIG_FILE_DIR/model/standalone_driver/tests' ] ignore_missing_imports = false implicit_reexport = true @@ -173,6 +177,8 @@ $MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/microphysics/src: $MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/muphys/src: $MYPY_CONFIG_FILE_DIR/model/common/src: $MYPY_CONFIG_FILE_DIR/model/driver/src: +$MYPY_CONFIG_FILE_DIR/model/standalone_driver/src: +$MYPY_CONFIG_FILE_DIR/model/standalone_driver/tests: $MYPY_CONFIG_FILE_DIR/model/testing/src: $MYPY_CONFIG_FILE_DIR/tools/src: ''' @@ -188,7 +194,12 @@ warn_unused_ignores = true # TODO(): fix gtx.int32 export in gt4py [[tool.mypy.overrides]] disable_error_code = ["attr-defined", "name-defined"] -module = ['*.stencil_tests.*', '*.unit_tests.*'] +module = [ + '*.stencil_tests.*', + '*.unit_tests.*', + "icon4py.model.atmosphere.advection.advection_vertical", # https://github.com/GridTools/gt4py/pull/2484 + "icon4py.model.common.metrics.compute_advection_metrics" # https://github.com/GridTools/gt4py/pull/2484 +] [[tool.mypy.overrides]] ignore_errors = true @@ -204,6 +215,36 @@ module = [ disallow_untyped_defs = false module = ['*.dycore.integration_tests.*', '*.dycore.mpi_tests.*'] +[[tool.mypy.overrides]] # complains about assigning to indices of NDArrayObject +disable_error_code = ["index"] +module = ["icon4py.model.standalone_driver.testcases.initial_condition"] + +[[tool.mypy.overrides]] +disable_error_code = ["attr-defined"] # https://github.com/GridTools/gt4py/pull/2484 +module = [ + "icon4py.model.atmosphere.advection.advection_horizontal", + "icon4py.model.atmosphere.advection.advection" +] + +[[tool.mypy.overrides]] +disable_error_code = ["valid-type"] # complains about dims not being valid types +module = [ + "icon4py.model.atmosphere.advection.advection_states" +] + +[[tool.mypy.overrides]] +disable_error_code = [ + "no-untyped-def", # programs are not annotated with -> None in icon4py + "call-overload", # https://github.com/GridTools/gt4py/issues/2496 + "operator", # mypy thinks float < Field[..., float] is not allowed + "attr-defined", # https://github.com/GridTools/gt4py/pull/2484 + "name-defined", # https://github.com/GridTools/gt4py/pull/2484 + "valid-type" # dims are not valid types according to mypy +] +module = [ + "icon4py.model.atmosphere.advection.stencils.*" +] + # -- pytest -- [tool.pytest] diff --git a/uv.lock b/uv.lock index f1dd438fda..6a4a75f68f 100644 --- a/uv.lock +++ b/uv.lock @@ -1594,6 +1594,7 @@ dev = [ { name = "pytest-unused-fixtures" }, { name = "pytest-xdist", extra = ["psutil"] }, { name = "ruff" }, + { name = "scipy-stubs" }, { name = "setuptools" }, { name = "sphinx" }, { name = "sphinx-math-dollar" }, @@ -1641,6 +1642,7 @@ test = [ ] typing = [ { name = "mypy", extra = ["faster-cache"] }, + { name = "scipy-stubs" }, { name = "types-cffi" }, { name = "typing-extensions" }, ] @@ -1693,6 +1695,7 @@ dev = [ { name = "pytest-unused-fixtures", specifier = ">=0.2.0" }, { name = "pytest-xdist", extras = ["psutil"], specifier = ">=3.5.0" }, { name = "ruff", specifier = ">=0.8.0" }, + { name = "scipy-stubs", specifier = ">=1.15.3.0" }, { name = "setuptools", specifier = ">=70.1.1" }, { name = "sphinx", specifier = "==7.3.7" }, { name = "sphinx-math-dollar", specifier = ">=1.2.1" }, @@ -1740,6 +1743,7 @@ test = [ ] typing = [ { name = "mypy", extras = ["faster-cache"], specifier = ">=1.13.0" }, + { name = "scipy-stubs", specifier = ">=1.15.3.0" }, { name = "types-cffi", specifier = ">=1.16.0" }, { name = "typing-extensions", specifier = ">=4.11.0" }, ] @@ -2880,6 +2884,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932, upload-time = "2024-09-26T14:33:23.039Z" }, ] +[[package]] +name = "optype" +version = "0.9.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-7-icon4py-cuda11' and extra == 'extra-7-icon4py-cuda12')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/88/3c/9d59b0167458b839273ad0c4fc5f62f787058d8f5aed7f71294963a99471/optype-0.9.3.tar.gz", hash = "sha256:5f09d74127d316053b26971ce441a4df01f3a01943601d3712dd6f34cdfbaf48", size = 96143, upload-time = "2025-03-31T17:00:08.392Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/d8/ac50e2982bdc2d3595dc2bfe3c7e5a0574b5e407ad82d70b5f3707009671/optype-0.9.3-py3-none-any.whl", hash = "sha256:2935c033265938d66cc4198b0aca865572e635094e60e6e79522852f029d9e8d", size = 84357, upload-time = "2025-03-31T17:00:06.464Z" }, +] + [[package]] name = "orderly-set" version = "5.2.3" @@ -3925,6 +3941,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f5/1b/6ee032251bf4cdb0cc50059374e86a9f076308c1512b61c4e003e241efb7/scipy-1.14.1-cp313-cp313-win_amd64.whl", hash = "sha256:baff393942b550823bfce952bb62270ee17504d02a1801d7fd0719534dfb9c84", size = 44469524, upload-time = "2024-08-21T00:07:15.381Z" }, ] +[[package]] +name = "scipy-stubs" +version = "1.15.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "optype" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0b/5f/35c43bd7d412add4adcd68475702571b2489b50c40b6564f808b2355e452/scipy_stubs-1.15.3.0.tar.gz", hash = "sha256:e8f76c9887461cf9424c1e2ad78ea5dac71dd4cbb383dc85f91adfe8f74d1e17", size = 275699, upload-time = "2025-05-08T16:58:35.139Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/42/cd8dc81f8060de1f14960885ad5b2d2651f41de8b93d09f3f919d6567a5a/scipy_stubs-1.15.3.0-py3-none-any.whl", hash = "sha256:a251254cf4fd6e7fb87c55c1feee92d32ddbc1f542ecdf6a0159cdb81c2fb62d", size = 459062, upload-time = "2025-05-08T16:58:33.356Z" }, +] + [[package]] name = "seaborn" version = "0.13.2"