diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index 7c7f56159e..182a19ad46 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -835,56 +835,35 @@ def _do_diffusion_step( "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): end" ) - log.debug( - "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" - ) - - self.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_connectivities( - self.compile_time_connectivities - )( - theta_v=prognostic_state.theta_v, - theta_ref_mc=self._metric_state.theta_ref_mc, - thresh_tdiff=self.thresh_tdiff, - smallest_vpfloat=constants.DBL_EPS, - kh_smag_e=self.kh_smag_e, - horizontal_start=self._edge_start_nudging, - horizontal_end=self._edge_end_halo, - vertical_start=(self._grid.num_levels - 2), - vertical_end=self._grid.num_levels, - offset_provider=self._grid.offset_providers, - ) - log.debug( - "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" - ) + if self.config.apply_to_temperature: + log.debug( + "running fused stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): start" + ) - log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") - self.calculate_nabla2_for_theta.with_connectivities(self.compile_time_connectivities)( - kh_smag_e=self.kh_smag_e, - inv_dual_edge_length=self._edge_params.inverse_dual_edge_lengths, - theta_v=prognostic_state.theta_v, - geofac_div=self._interpolation_state.geofac_div, - z_temp=self.z_temp, - horizontal_start=self._cell_start_nudging, - horizontal_end=self._cell_end_local, - vertical_start=0, - vertical_end=self._grid.num_levels, - offset_provider=self._grid.offset_providers, - ) - log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") - log.debug( - "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" - ) - if self.config.apply_zdiffusion_t: - self.truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_connectivities( + self.calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.with_connectivities( self.compile_time_connectivities )( - mask=self._metric_state.mask_hdiff, - zd_vertoffset=self._metric_state.zd_vertoffset, - zd_diffcoef=self._metric_state.zd_diffcoef, - geofac_n2s_c=self._interpolation_state.geofac_n2s_c, - geofac_n2s_nbh=self._interpolation_state.geofac_n2s_nbh, - vcoef=self._metric_state.zd_intcoef, theta_v=prognostic_state.theta_v, + theta_ref_mc=self._metric_state.theta_ref_mc, + thresh_tdiff=self.thresh_tdiff, + smallest_vpfloat=constants.DBL_EPS, + kh_smag_e=self.kh_smag_e, + horizontal_start=self._edge_start_nudging, + horizontal_end=self._edge_end_halo, + vertical_start=(self._grid.num_levels - 2), + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug( + "running stencils 11 12 (calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools): end" + ) + + log.debug("running stencils 13 14 (calculate_nabla2_for_theta): start") + self.calculate_nabla2_for_theta.with_connectivities(self.compile_time_connectivities)( + kh_smag_e=self.kh_smag_e, + inv_dual_edge_length=self._edge_params.inverse_dual_edge_lengths, + theta_v=prognostic_state.theta_v, + geofac_div=self._interpolation_state.geofac_div, z_temp=self.z_temp, horizontal_start=self._cell_start_nudging, horizontal_end=self._cell_end_local, @@ -892,24 +871,46 @@ def _do_diffusion_step( vertical_end=self._grid.num_levels, offset_provider=self._grid.offset_providers, ) - + log.debug("running stencils 13_14 (calculate_nabla2_for_theta): end") log.debug( - "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + "running stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): start" ) - log.debug("running stencil 16 (update_theta_and_exner): start") - self.update_theta_and_exner.with_connectivities(self.compile_time_connectivities)( - z_temp=self.z_temp, - area=self._cell_params.area, - theta_v=prognostic_state.theta_v, - exner=prognostic_state.exner, - rd_o_cvd=self.rd_o_cvd, - horizontal_start=self._cell_start_nudging, - horizontal_end=self._cell_end_local, - vertical_start=0, - vertical_end=self._grid.num_levels, - offset_provider={}, - ) - log.debug("running stencil 16 (update_theta_and_exner): end") + if self.config.apply_zdiffusion_t: + self.truly_horizontal_diffusion_nabla_of_theta_over_steep_points.with_connectivities( + self.compile_time_connectivities + )( + mask=self._metric_state.mask_hdiff, + zd_vertoffset=self._metric_state.zd_vertoffset, + zd_diffcoef=self._metric_state.zd_diffcoef, + geofac_n2s_c=self._interpolation_state.geofac_n2s_c, + geofac_n2s_nbh=self._interpolation_state.geofac_n2s_nbh, + vcoef=self._metric_state.zd_intcoef, + theta_v=prognostic_state.theta_v, + z_temp=self.z_temp, + horizontal_start=self._cell_start_nudging, + horizontal_end=self._cell_end_local, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + + log.debug( + "running fused stencil 15 (truly_horizontal_diffusion_nabla_of_theta_over_steep_points): end" + ) + log.debug("running stencil 16 (update_theta_and_exner): start") + self.update_theta_and_exner.with_connectivities(self.compile_time_connectivities)( + z_temp=self.z_temp, + area=self._cell_params.area, + theta_v=prognostic_state.theta_v, + exner=prognostic_state.exner, + rd_o_cvd=self.rd_o_cvd, + horizontal_start=self._cell_start_nudging, + horizontal_end=self._cell_end_local, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider={}, + ) + log.debug("running stencil 16 (update_theta_and_exner): end") self.halo_exchange_wait( handle_edge_comm diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_wp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_wp.py new file mode 100644 index 0000000000..68de92f94e --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_wp.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 +from gt4py.next.common import GridType +from gt4py.next.ffront.decorator import field_operator, program + +from icon4py.model.common import dimension as dims, field_type_aliases as fa +from icon4py.model.common.dimension import Koff +from icon4py.model.common.settings import backend +from icon4py.model.common.type_alias import wpfloat + + +@field_operator +def _interpolate_to_half_levels_wp( + wgtfac_c: fa.CellKField[wpfloat], + interpolant: fa.CellKField[wpfloat], +) -> fa.CellKField[wpfloat]: + """ + Interpolate a CellDim variable of working precision from full levels to half levels. + The return variable also has working precision. + var_half_k-1/2 = wgt_fac_c_k-1 var_half_k-1 + wgt_fac_c_k var_half_k + + Args: + wgtfac_c: weight factor + interpolant: CellDim variables at full levels + Returns: + CellDim variables at half levels + """ + interpolation_to_half_levels_wp = wgtfac_c * interpolant + ( + wpfloat("1.0") - wgtfac_c + ) * interpolant(Koff[-1]) + return interpolation_to_half_levels_wp + + +@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +def interpolate_to_half_levels_wp( + wgtfac_c: fa.CellKField[wpfloat], + interpolant: fa.CellKField[wpfloat], + interpolation_to_half_levels_wp: fa.CellKField[wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _interpolate_to_half_levels_wp( + wgtfac_c, + interpolant, + out=interpolation_to_half_levels_wp, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py index 36d1403d46..e6559e8900 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/compute_virtual_potential_temperatures_and_pressure_gradient.py @@ -13,6 +13,9 @@ from icon4py.model.atmosphere.dycore.stencils.interpolate_to_half_levels_vp import ( _interpolate_to_half_levels_vp, ) +from icon4py.model.atmosphere.dycore.interpolate_to_half_levels_wp import ( + _interpolate_to_half_levels_wp, +) from icon4py.model.common import dimension as dims, field_type_aliases as fa from icon4py.model.common.dimension import Koff from icon4py.model.common.settings import backend @@ -37,7 +40,7 @@ def _compute_virtual_potential_temperatures_and_pressure_gradient( wgtfac_c_wp, ddqz_z_half_wp = astype((wgtfac_c, ddqz_z_half), wpfloat) z_theta_v_pr_ic_vp = _interpolate_to_half_levels_vp(wgtfac_c=wgtfac_c, interpolant=z_rth_pr_2) - theta_v_ic_wp = wgtfac_c_wp * theta_v + (wpfloat("1.0") - wgtfac_c_wp) * theta_v(Koff[-1]) + theta_v_ic_wp = _interpolate_to_half_levels_wp(wgtfac_c=wgtfac_c_wp, interpolant=theta_v) z_th_ddz_exner_c_wp = vwind_expl_wgt * theta_v_ic_wp * ( exner_pr(Koff[-1]) - exner_pr ) / ddqz_z_half_wp + astype(z_theta_v_pr_ic_vp * d_exner_dz_ref_ic, wpfloat) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_half_levels_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_half_levels_vp.py index d3893178d2..dd12053f7f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_half_levels_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/interpolate_to_half_levels_vp.py @@ -20,7 +20,17 @@ def _interpolate_to_half_levels_vp( wgtfac_c: fa.CellKField[vpfloat], interpolant: fa.CellKField[vpfloat], ) -> fa.CellKField[vpfloat]: - """Formerly known mo_velocity_advection_stencil_10 and as _mo_solve_nonhydro_stencil_05.""" + """ + Interpolate a CellDim variable of floating precision from full levels to half levels. + The return variable also has floating precision. + var_half_k-1/2 = wgt_fac_c_k-1 var_half_k-1 + wgt_fac_c_k var_half_k + + Args: + wgtfac_c: weight factor + interpolant: CellDim variables at full levels + Returns: + CellDim variables at half levels + """ interpolation_to_half_levels_vp = wgtfac_c * interpolant + ( vpfloat("1.0") - wgtfac_c ) * interpolant(Koff[-1]) diff --git a/model/common/src/icon4py/model/common/field_type_aliases.py b/model/common/src/icon4py/model/common/field_type_aliases.py index 3ac4fdd6e2..b2746d7a9d 100644 --- a/model/common/src/icon4py/model/common/field_type_aliases.py +++ b/model/common/src/icon4py/model/common/field_type_aliases.py @@ -8,6 +8,7 @@ from typing import TypeAlias, TypeVar import gt4py.next as gtx +import numpy as np from gt4py.next import Dims, Field from icon4py.model.common import dimension as dims @@ -24,3 +25,6 @@ CellKField: TypeAlias = Field[Dims[dims.CellDim, dims.KDim], T] EdgeKField: TypeAlias = Field[Dims[dims.EdgeDim, dims.KDim], T] VertexKField: TypeAlias = Field[Dims[dims.VertexDim, dims.KDim], T] + +# TODO (anyone): add cupy if cupy is installed? +AnyNDArray: TypeAlias = np.ndarray diff --git a/model/common/src/icon4py/model/common/grid/base.py b/model/common/src/icon4py/model/common/grid/base.py index a8051d66ec..66118f35b9 100644 --- a/model/common/src/icon4py/model/common/grid/base.py +++ b/model/common/src/icon4py/model/common/grid/base.py @@ -17,7 +17,7 @@ from icon4py.model.common import dimension as dims, utils from icon4py.model.common.grid import utils as grid_utils -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import array_allocation as array_alloc class MissingConnectivity(ValueError): @@ -136,6 +136,7 @@ def _get_offset_provider(self, dim, from_dim, to_dim): ), 'Neighbor table\'s "{}" data type must be gtx.int32. Instead it\'s "{}"'.format( dim, self.connectivities[dim].dtype ) + xp = array_alloc.array_ns(self.config.on_gpu) return gtx.NeighborTableOffsetProvider( xp.asarray(self.connectivities[dim]), from_dim, @@ -148,6 +149,7 @@ def _get_offset_provider_for_sparse_fields(self, dim, from_dim, to_dim): if dim not in self.connectivities: raise MissingConnectivity() return grid_utils.neighbortable_offset_provider_for_1d_sparse_fields( + self.config.on_gpu, self.connectivities[dim].shape, from_dim, to_dim, diff --git a/model/common/src/icon4py/model/common/grid/utils.py b/model/common/src/icon4py/model/common/grid/utils.py index 424bd51aca..8416c15928 100644 --- a/model/common/src/icon4py/model/common/grid/utils.py +++ b/model/common/src/icon4py/model/common/grid/utils.py @@ -8,15 +8,17 @@ import gt4py.next as gtx from gt4py.next import Dimension, NeighborTableOffsetProvider -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import array_allocation as array_alloc def neighbortable_offset_provider_for_1d_sparse_fields( + on_gpu: bool, old_shape: tuple[int, int], origin_axis: Dimension, neighbor_axis: Dimension, has_skip_values: bool, ): + xp = array_alloc.array_ns(on_gpu) table = xp.arange(old_shape[0] * old_shape[1], dtype=gtx.int32).reshape(old_shape) assert ( table.dtype == gtx.int32 diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index feedd66526..005df3c457 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -14,10 +14,12 @@ from typing import Final import gt4py.next as gtx +import numpy as np +from gt4py.next import backend as gt4py_backend import icon4py.model.common.states.metadata as data from icon4py.model.common import dimension as dims, exceptions, field_type_aliases as fa -from icon4py.model.common.settings import xp +from icon4py.model.common.utils import array_allocation as array_alloc log = logging.getLogger(__name__) @@ -120,7 +122,7 @@ class VerticalGrid: _end_index_of_flat_layer: Final[gtx.int32] = dataclasses.field(init=False) _min_index_flat_horizontal_grad_pressure: Final[gtx.int32] = None - def __post_init__(self, vct_a, vct_b): + def __post_init__(self, vct_a: fa.KField, vct_b: fa.KField): object.__setattr__( self, "_vct_a", @@ -246,15 +248,19 @@ def size(self, dim: gtx.Dimension) -> int: @classmethod def _determine_start_level_of_moist_physics( - cls, vct_a: xp.ndarray, top_moist_threshold: float, nshift_total: int = 0 + cls, vct_a: fa.AnyNDArray, top_moist_threshold: float, nshift_total: int = 0 ) -> gtx.int32: + xp = array_alloc.array_ns_from_obj(vct_a) n_levels = vct_a.shape[0] interface_height = 0.5 * (vct_a[: n_levels - 1 - nshift_total] + vct_a[1 + nshift_total :]) return gtx.int32(xp.min(xp.where(interface_height < top_moist_threshold)[0]).item()) @classmethod - def _determine_damping_height_index(cls, vct_a: xp.ndarray, damping_height: float) -> gtx.int32: + def _determine_damping_height_index( + cls, vct_a: fa.AnyNDArray, damping_height: float + ) -> gtx.int32: assert damping_height >= 0.0, "Damping height must be positive." + xp = array_alloc.array_ns_from_obj(vct_a) return ( 0 if damping_height > vct_a[0] @@ -263,9 +269,10 @@ def _determine_damping_height_index(cls, vct_a: xp.ndarray, damping_height: floa @classmethod def _determine_end_index_of_flat_layers( - cls, vct_a: xp.ndarray, flat_height: float + cls, vct_a: fa.AnyNDArray, flat_height: float ) -> gtx.int32: assert flat_height >= 0.0, "Flat surface height must be positive." + xp = array_alloc.array_ns_from_obj(vct_a) return ( 0 if flat_height > vct_a[0] @@ -274,7 +281,7 @@ def _determine_end_index_of_flat_layers( def _read_vct_a_and_vct_b_from_file( - file_path: pathlib.Path, num_levels: int + file_path: pathlib.Path, num_levels: int, backend: gt4py_backend.Backend ) -> tuple[fa.KField, fa.KField]: """ Read vct_a and vct_b from a file. @@ -293,8 +300,8 @@ def _read_vct_a_and_vct_b_from_file( Returns: one dimensional vct_a and vct_b arrays. """ num_levels_plus_one = num_levels + 1 - vct_a = xp.zeros(num_levels_plus_one, dtype=float) - vct_b = xp.zeros(num_levels_plus_one, dtype=float) + vct_a = np.zeros(num_levels_plus_one, dtype=float) + vct_b = np.zeros(num_levels_plus_one, dtype=float) try: with open(file_path, "r") as vertical_grid_file: # skip the first line that contains titles @@ -313,10 +320,14 @@ def _read_vct_a_and_vct_b_from_file( ) from err except ValueError as err: raise ValueError(f"data is not float at {k}-th line.") from err - return gtx.as_field((dims.KDim,), vct_a), gtx.as_field((dims.KDim,), vct_b) + return gtx.as_field((dims.KDim,), vct_a, allocator=backend), gtx.as_field( + (dims.KDim,), vct_b, allocator=backend + ) -def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KField, fa.KField]: +def _compute_vct_a_and_vct_b( + vertical_config: VerticalGridConfig, backend: gt4py_backend.Backend +) -> tuple[fa.KField, fa.KField]: """ Compute vct_a and vct_b. @@ -354,16 +365,17 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF Args: vertical_config: Vertical grid configuration + backend: GT4Py backend Returns: one dimensional (dims.KDim) vct_a and vct_b gt4py fields. """ num_levels_plus_one = vertical_config.num_levels + 1 if vertical_config.lowest_layer_thickness > 0.01: - vct_a_exponential_factor = xp.log( + vct_a_exponential_factor = np.log( vertical_config.lowest_layer_thickness / vertical_config.model_top_height - ) / xp.log( + ) / np.log( 2.0 / math.pi - * xp.arccos( + * np.arccos( float(vertical_config.num_levels - 1) ** vertical_config.stretch_factor / float(vertical_config.num_levels) ** vertical_config.stretch_factor ) @@ -374,8 +386,8 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF * ( 2.0 / math.pi - * xp.arccos( - xp.arange(vertical_config.num_levels + 1, dtype=float) + * np.arccos( + np.arange(vertical_config.num_levels + 1, dtype=float) ** vertical_config.stretch_factor / float(vertical_config.num_levels) ** vertical_config.stretch_factor ) @@ -389,10 +401,10 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF < 0.5 * vertical_config.top_height_limit_for_maximal_layer_thickness ): layer_thickness = vct_a[: vertical_config.num_levels] - vct_a[1:] - lowest_level_exceeding_limit = xp.max( - xp.where(layer_thickness > vertical_config.maximal_layer_thickness) + lowest_level_exceeding_limit = np.max( + np.where(layer_thickness > vertical_config.maximal_layer_thickness) ) - modified_vct_a = xp.zeros(num_levels_plus_one, dtype=float) + modified_vct_a = np.zeros(num_levels_plus_one, dtype=float) lowest_level_unmodified_thickness = 0 shifted_levels = 0 for k in range(vertical_config.num_levels - 1, -1, -1): @@ -400,7 +412,7 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF modified_vct_a[k + 1] < vertical_config.top_height_limit_for_maximal_layer_thickness ): - modified_vct_a[k] = modified_vct_a[k + 1] + xp.minimum( + modified_vct_a[k] = modified_vct_a[k + 1] + np.minimum( vertical_config.maximal_layer_thickness, layer_thickness[k] ) elif lowest_level_unmodified_thickness == 0: @@ -431,7 +443,7 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF for k in range(vertical_config.num_levels - 1, -1, -1): if vct_a[k + 1] < vertical_config.top_height_limit_for_maximal_layer_thickness: - vct_a[k] = vct_a[k + 1] + xp.minimum( + vct_a[k] = vct_a[k + 1] + np.minimum( vertical_config.maximal_layer_thickness, layer_thickness[k] ) else: @@ -454,7 +466,7 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF ): modified_vct_a[k] = vct_a[k] else: - modified_layer_thickness = xp.minimum( + modified_layer_thickness = np.minimum( 1.025 * (vct_a[k] - vct_a[k + 1]), 1.025 * ( @@ -467,7 +479,7 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF ) * (modified_vct_a[k + 1] - modified_vct_a[k + 2]), ) - modified_vct_a[k] = xp.minimum( + modified_vct_a[k] = np.minimum( vct_a[k], modified_vct_a[k + 1] + modified_layer_thickness ) if modified_vct_a[0] == vct_a[0]: @@ -484,20 +496,24 @@ def _compute_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KF else: vct_a = ( vertical_config.model_top_height - * (float(vertical_config.num_levels) - xp.arange(num_levels_plus_one, dtype=float)) + * (float(vertical_config.num_levels) - np.arange(num_levels_plus_one, dtype=float)) / float(vertical_config.num_levels) ) - vct_b = xp.exp(-vct_a / 5000.0) + vct_b = np.exp(-vct_a / 5000.0) - if not xp.allclose(vct_a[0], vertical_config.model_top_height): + if not np.allclose(vct_a[0], vertical_config.model_top_height): log.warning( f" Warning. vct_a[0], {vct_a[0]}, is not equal to model top height, {vertical_config.model_top_height}, of vertical configuration. Please consider changing the vertical setting." ) - return gtx.as_field((dims.KDim,), vct_a), gtx.as_field((dims.KDim,), vct_b) + return gtx.as_field((dims.KDim,), vct_a, allocator=backend), gtx.as_field( + (dims.KDim,), vct_b, allocator=backend + ) -def get_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KField, fa.KField]: +def get_vct_a_and_vct_b( + vertical_config: VerticalGridConfig, backend: gt4py_backend.Backend +) -> tuple[fa.KField, fa.KField]: """ get vct_a and vct_b. vct_a is an array that contains the height of grid interfaces (or half levels) from model surface to model top, before terrain-following coordinates are applied. @@ -509,11 +525,14 @@ def get_vct_a_and_vct_b(vertical_config: VerticalGridConfig) -> tuple[fa.KField, Args: vertical_config: Vertical grid configuration + backend: GT4Py backend Returns: one dimensional (dims.KDim) vct_a and vct_b gt4py fields. """ return ( - _read_vct_a_and_vct_b_from_file(vertical_config.file_path, vertical_config.num_levels) + _read_vct_a_and_vct_b_from_file( + vertical_config.file_path, vertical_config.num_levels, backend + ) if vertical_config.file_path - else _compute_vct_a_and_vct_b(vertical_config) + else _compute_vct_a_and_vct_b(vertical_config, backend) ) diff --git a/model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py b/model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py index 647f5a5161..0b4141da16 100644 --- a/model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py +++ b/model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py @@ -12,7 +12,6 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.dimension import C2E2C2E, C2E2C2EDim -from icon4py.model.common.settings import backend @field_operator @@ -41,7 +40,7 @@ def _edge_2_cell_vector_rbf_interpolation( return p_u_out, p_v_out -@program(grid_type=GridType.UNSTRUCTURED, backend=backend) +@program(grid_type=GridType.UNSTRUCTURED) def edge_2_cell_vector_rbf_interpolation( p_e_in: fa.EdgeKField[ta.wpfloat], ptr_coeff_1: gtx.Field[gtx.Dims[dims.CellDim, C2E2C2EDim], ta.wpfloat], 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 b84805ad52..a028db606d 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 @@ -9,10 +9,12 @@ import pytest import icon4py.model.common.decomposition.definitions as decomposition +from icon4py.model.common.test_utils import pytest_config from . import data_handling as data, datatest_utils as dt_utils + @pytest.fixture def experiment(): return dt_utils.REGIONAL_EXPERIMENT @@ -100,13 +102,19 @@ def is_regional(experiment_name): @pytest.fixture -def icon_grid(grid_savepoint): +def icon_grid(request, grid_savepoint): """ Load the icon grid from an ICON savepoint. Uses the special grid_savepoint that contains data from p_patch """ - return grid_savepoint.construct_icon_grid(on_gpu=False) + on_gpu = False + if request.config.getoption("--backend"): + backend = request.config.getoption("--backend") + pytest_config.check_backend_validity(backend) + if backend in pytest_config.gpu_backends: + on_gpu = True + return grid_savepoint.construct_icon_grid(on_gpu=on_gpu) @pytest.fixture diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index d91f264a0d..19c36eafb7 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -10,6 +10,7 @@ import uuid import gt4py.next as gtx +import numpy as np import serialbox import icon4py.model.common.decomposition.definitions as decomposition @@ -18,7 +19,7 @@ import icon4py.model.common.test_utils.helpers as helpers from icon4py.model.common import dimension as dims from icon4py.model.common.grid import base, horizontal, icon -from icon4py.model.common.settings import xp +from icon4py.model.common.settings import backend, xp from icon4py.model.common.states import prognostic_state @@ -45,7 +46,7 @@ def wrapper(self, *args, **kwargs): ) if dims: shp = tuple(self.sizes[d] for d in dims) - return gtx.as_field(dims, xp.zeros(shp)) + return gtx.as_field(dims, xp.zeros(shp), allocator=backend) else: return None @@ -61,14 +62,14 @@ def _get_field(self, name, *dimensions, dtype=float): buffer = self._reduce_to_dim_size(buffer, dimensions) self.log.debug(f"{name} {buffer.shape}") - return gtx.as_field(dimensions, buffer) + return gtx.as_field(dimensions, buffer, allocator=backend) def _get_field_component(self, name: str, ntnd: int, dims: tuple[gtx.Dimension, gtx]): buffer = self.serializer.read(name, self.savepoint).astype(float) buffer = xp.squeeze(buffer)[:, :, ntnd - 1] buffer = self._reduce_to_dim_size(buffer, dims) self.log.debug(f"{name} {buffer.shape}") - return gtx.as_field(dims, buffer) + return gtx.as_field(dims, buffer, allocator=backend) def _reduce_to_dim_size(self, buffer, dimensions): buffer_size = ( @@ -80,7 +81,7 @@ def _reduce_to_dim_size(self, buffer, dimensions): def _get_field_from_ndarray(self, ar, *dimensions, dtype=float): ar = self._reduce_to_dim_size(ar, dimensions) - return gtx.as_field(dimensions, ar) + return gtx.as_field(dimensions, ar, allocator=backend) def get_metadata(self, *names): metadata = self.savepoint.metainfo.to_dict() @@ -326,12 +327,11 @@ def c2e(self): def _get_connectivity_array(self, name: str, target_dim: gtx.Dimension, reverse: bool = False): if reverse: - connectivity = xp.transpose(self._read_int32(name, offset=1))[ + connectivity = np.transpose(self._read_int32(name, offset=1))[ : self.sizes[target_dim], : ] else: connectivity = self._read_int32(name, offset=1)[: self.sizes[target_dim], :] - connectivity = xp.asarray(connectivity) self.log.debug(f" connectivity {name} : {connectivity.shape}") return connectivity @@ -343,7 +343,7 @@ def e2c2e(self): def c2e2c2e(self): if self._c2e2c2e() is None: - return xp.zeros((self.sizes[dims.CellDim], 9), dtype=int) + return np.zeros((self.sizes[dims.CellDim], 9), dtype=int) else: return self._c2e2c2e() @@ -383,6 +383,7 @@ def refin_ctrl(self, dim: gtx.Dimension): xp.squeeze( self._read_field_for_dim(field_name, self._read_int32, dim)[: self.num(dim)], 1 ), + allocator=backend, ) def num(self, dim: gtx.Dimension): @@ -453,8 +454,8 @@ def construct_icon_grid(self, on_gpu: bool) -> icon.IconGrid: ) c2e2c = self.c2e2c() e2c2e = self.e2c2e() - c2e2c0 = xp.column_stack(((range(c2e2c.shape[0])), c2e2c)) - e2c2e0 = xp.column_stack(((range(e2c2e.shape[0])), e2c2e)) + c2e2c0 = np.column_stack((range(c2e2c.shape[0]), c2e2c)) + e2c2e0 = np.column_stack((range(e2c2e.shape[0]), e2c2e)) grid = ( icon.IconGrid(self._grid_id) .with_config(config) @@ -577,9 +578,9 @@ def geofac_grdiv(self): def geofac_grg(self): grg = xp.squeeze(self.serializer.read("geofac_grg", self.savepoint)) num_cells = self.sizes[dims.CellDim] - return gtx.as_field((dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 0]), gtx.as_field( - (dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 1] - ) + return gtx.as_field( + (dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 0], allocator=backend + ), gtx.as_field((dims.CellDim, dims.C2E2CODim), grg[:num_cells, :, 1], allocator=backend) @IconSavepoint.optionally_registered() def zd_intcoef(self): @@ -606,21 +607,21 @@ def rbf_vec_coeff_e(self): buffer = xp.squeeze( self.serializer.read("rbf_vec_coeff_e", self.savepoint).astype(float) ).transpose() - return gtx.as_field((dims.EdgeDim, dims.E2C2EDim), buffer) + return gtx.as_field((dims.EdgeDim, dims.E2C2EDim), buffer, allocator=backend) @IconSavepoint.optionally_registered() def rbf_vec_coeff_c1(self): buffer = xp.squeeze( self.serializer.read("rbf_vec_coeff_c1", self.savepoint).astype(float) ).transpose() - return gtx.as_field((dims.CellDim, dims.C2E2C2EDim), buffer) + return gtx.as_field((dims.CellDim, dims.C2E2C2EDim), buffer, allocator=backend) @IconSavepoint.optionally_registered() def rbf_vec_coeff_c2(self): buffer = xp.squeeze( self.serializer.read("rbf_vec_coeff_c2", self.savepoint).astype(float) ).transpose() - return gtx.as_field((dims.CellDim, dims.C2E2C2EDim), buffer) + return gtx.as_field((dims.CellDim, dims.C2E2C2EDim), buffer, allocator=backend) def rbf_vec_coeff_v1(self): return self._get_field("rbf_vec_coeff_v1", dims.VertexDim, dims.V2EDim) @@ -753,7 +754,7 @@ def wgtfacq_e_dsl( ): # TODO: @abishekg7 Simplify this after serialized data is fixed ar = xp.squeeze(self.serializer.read("wgtfacq_e", self.savepoint)) k = k_level - 3 - ar = xp.pad(ar[:, ::-1], ((0, 0), (k, 0)), "constant", constant_values=(0.0,)) + ar = xp.pad(xp.asarray(ar[:, ::-1]), ((0, 0), (k, 0)), "constant", constant_values=(0.0,)) return self._get_field_from_ndarray(ar, dims.EdgeDim, dims.KDim) @IconSavepoint.optionally_registered(dims.CellDim, dims.KDim) @@ -782,7 +783,9 @@ def _linearize_first_2dims( ): old_shape = data.shape assert old_shape[1] == sparse_size - return gtx.as_field(target_dims, data.reshape(old_shape[0] * old_shape[1], old_shape[2])) + return gtx.as_field( + target_dims, data.reshape(old_shape[0] * old_shape[1], old_shape[2]), allocator=backend + ) @IconSavepoint.optionally_registered() def zd_vertoffset(self): diff --git a/model/common/src/icon4py/model/common/utils/array_allocation.py b/model/common/src/icon4py/model/common/utils/array_allocation.py new file mode 100644 index 0000000000..417583a07a --- /dev/null +++ b/model/common/src/icon4py/model/common/utils/array_allocation.py @@ -0,0 +1,46 @@ +# 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 logging + +import numpy as np +from gt4py._core import definitions as core_defs +from gt4py.next import allocators, backend as gt4py_backend + +from icon4py.model.common import field_type_aliases as fa + + +log = logging.getLogger(__name__) + + +def is_cupy_device(backend: gt4py_backend.Backend) -> bool: + return ( + allocators.is_field_allocator_for(backend.allocator, core_defs.DeviceType.CUDA) + or allocators.is_field_allocator_for(backend.allocator, core_defs.DeviceType.CUDA_MANAGED) + or allocators.is_field_allocator_for(backend.allocator, core_defs.DeviceType.ROCM) + ) + + +def array_ns(try_cupy: bool): + if try_cupy: + try: + import cupy as cp + + return cp + except ImportError: + log.warning("No cupy installed falling back to numpy for array_ns") + return np + + +def array_ns_from_obj(array: fa.AnyNDArray): + if isinstance(array, np.ndarray): + return np + else: + import cupy as cp + + return cp diff --git a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py index 1799fe8cbb..4f9627f2f5 100644 --- a/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py +++ b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py @@ -8,24 +8,27 @@ from typing import Optional import gt4py.next as gtx -from gt4py.next import backend +import numpy as np +from gt4py.next import backend as gt4py_backend, common as gt4py_common -from icon4py.model.common import type_alias as ta -from icon4py.model.common.settings import xp +from icon4py.model.common import dimension as dims, type_alias as ta def allocate_zero_field( - *dims: gtx.Dimension, + *dim: gtx.Dimension, grid, is_halfdim=False, dtype=ta.wpfloat, - backend: Optional[backend.Backend] = None, -): - shapex = tuple(map(lambda x: grid.size[x], dims)) - if is_halfdim: - assert len(shapex) == 2 - shapex = (shapex[0], shapex[1] + 1) - return gtx.as_field(dims, xp.zeros(shapex, dtype=dtype), allocator=backend) + backend: Optional[gt4py_backend.Backend] = None, +) -> gt4py_common.Field: + def size(local_dim: gtx.Dimension, is_half_dim: bool) -> int: + if local_dim == dims.KDim and is_half_dim: + return grid.size[local_dim] + 1 + else: + return grid.size[local_dim] + + dimensions = {d: range(size(d, is_halfdim)) for d in dim} + return gtx.zeros(dimensions, dtype=dtype, allocator=backend) def allocate_indices( @@ -33,7 +36,7 @@ def allocate_indices( grid, is_halfdim=False, dtype=gtx.int32, - backend: Optional[backend.Backend] = None, -): + backend: Optional[gt4py_backend.Backend] = None, +) -> gt4py_common.Field: shapex = grid.size[dim] + 1 if is_halfdim else grid.size[dim] - return gtx.as_field((dim,), xp.arange(shapex, dtype=dtype), allocator=backend) + return gtx.as_field((dim,), np.arange(shapex, dtype=dtype), allocator=backend) diff --git a/model/common/tests/grid_tests/test_vertical.py b/model/common/tests/grid_tests/test_vertical.py index fd508f6b71..ecea4b8ac5 100644 --- a/model/common/tests/grid_tests/test_vertical.py +++ b/model/common/tests/grid_tests/test_vertical.py @@ -277,6 +277,7 @@ def test_vct_a_vct_b_calculation_from_icon_input( stretch_factor, damping_height, htop_moist_proc, + backend, ): vertical_config = v_grid.VerticalGridConfig( num_levels=grid_savepoint.num(dims.KDim), @@ -289,7 +290,7 @@ def test_vct_a_vct_b_calculation_from_icon_input( rayleigh_damping_height=damping_height, htop_moist_proc=htop_moist_proc, ) - vct_a, vct_b = v_grid.get_vct_a_and_vct_b(vertical_config) + vct_a, vct_b = v_grid.get_vct_a_and_vct_b(vertical_config, backend) assert helpers.dallclose(vct_a.ndarray, grid_savepoint.vct_a().ndarray) assert helpers.dallclose(vct_b.ndarray, grid_savepoint.vct_b().ndarray) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index a41a75fa04..093fe14c7a 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -8,7 +8,16 @@ import dataclasses import datetime +import enum import logging +from typing import Final + +from gt4py.next.program_processors.runners.gtfn import ( + run_gtfn, + run_gtfn_cached, + run_gtfn_gpu, + run_gtfn_gpu_cached, +) from icon4py.model.atmosphere.diffusion import diffusion from icon4py.model.atmosphere.dycore import solve_nonhydro as solve_nh @@ -21,6 +30,22 @@ n_substeps_reduced = 2 +class DriverBackends(str, enum.Enum): + GTFN_CPU = "gtfn_cpu" + GTFN_CPU_CACHED = "gtfn_cpu_cached" + GTFN_GPU = "gtfn_gpu" + GTFN_GPU_CACHED = "gtfn_gpu_cached" + + +backend_map: Final[dict] = { + DriverBackends.GTFN_CPU: run_gtfn, + DriverBackends.GTFN_CPU_CACHED: run_gtfn_cached, + DriverBackends.GTFN_GPU: run_gtfn_gpu, + DriverBackends.GTFN_GPU_CACHED: run_gtfn_gpu_cached, +} +gpu_backends = [DriverBackends.GTFN_GPU, DriverBackends.GTFN_GPU_CACHED] + + @dataclasses.dataclass(frozen=True) class Icon4pyRunConfig: dtime: datetime.timedelta = datetime.timedelta(seconds=600.0) # length of a time step @@ -40,6 +65,19 @@ class Icon4pyRunConfig: restart_mode: bool = False + backend_name: str = DriverBackends.GTFN_CPU.value + + def __post_init__(self): + if self.backend_name not in [member.value for member in DriverBackends]: + raise ValueError( + f"Invalid driver backend: {self.backend_name}. \n" + f"Available backends are {', '.join([f'{k}' for k in [member.value for member in DriverBackends]])}" + ) + + @property + def backend(self): + return backend_map[self.backend_name] + @dataclasses.dataclass class Icon4pyConfig: @@ -50,6 +88,7 @@ class Icon4pyConfig: def read_config( + icon4py_driver_backend: str, experiment_type: driver_init.ExperimentType = driver_init.ExperimentType.ANY, ) -> Icon4pyConfig: def _mch_ch_r04b09_vertical_config(): @@ -122,6 +161,7 @@ def _mch_ch_r04b09_config(): end_date=datetime.datetime(2021, 6, 20, 12, 0, 10), n_substeps=n_substeps_reduced, apply_initial_stabilization=True, + backend_name=icon4py_driver_backend, ), _mch_ch_r04b09_vertical_config(), _mch_ch_r04b09_diffusion_config(), @@ -134,6 +174,7 @@ def _jablownoski_Williamson_config(): end_date=datetime.datetime(1, 1, 1, 0, 30, 0), apply_initial_stabilization=False, n_substeps=5, + backend_name=icon4py_driver_backend, ) jabw_vertical_config = _jabw_vertical_config() jabw_diffusion_config = _jabw_diffusion_config(icon_run_config.n_substeps) @@ -168,6 +209,7 @@ def _gauss3d_config(): end_date=datetime.datetime(1, 1, 1, 0, 0, 4), apply_initial_stabilization=False, n_substeps=5, + backend_name=icon4py_driver_backend, ) vertical_config = _gauss3d_vertical_config() diffusion_config = _gauss3d_diffusion_config(icon_run_config.n_substeps) diff --git a/model/driver/src/icon4py/model/driver/icon4py_driver.py b/model/driver/src/icon4py/model/driver/icon4py_driver.py index c567146b3a..2069940171 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_driver.py +++ b/model/driver/src/icon4py/model/driver/icon4py_driver.py @@ -14,7 +14,6 @@ import click from devtools import Timer -from gt4py.next import gtfn_cpu from icon4py.model.atmosphere.diffusion import ( diffusion, @@ -162,13 +161,13 @@ def time_integration( timer = Timer(self._full_name(self._integrate_one_time_step)) for time_step in range(self._n_time_steps): log.info(f"simulation date : {self._simulation_date} run timestep : {time_step}") - log.info( + log.debug( f" MAX VN: {prognostic_state_list[self._now].vn.asnumpy().max():.15e} , MAX W: {prognostic_state_list[self._now].w.asnumpy().max():.15e}" ) - log.info( + log.debug( f" MAX RHO: {prognostic_state_list[self._now].rho.asnumpy().max():.15e} , MAX THETA_V: {prognostic_state_list[self._now].theta_v.asnumpy().max():.15e}" ) - # TODO (Chia Rui): check with Anurag about printing of max and min of variables. + # TODO (Chia Rui): check with Anurag about printing of max and min of variables. Currently, these max values are only output at debug level. There should be namelist parameters to control which variable max should be output. self._next_simulation_date() @@ -274,6 +273,7 @@ def initialize( grid_id: uuid.UUID, grid_root, grid_level, + icon4py_driver_backend: driver_config.DriverBackends, ): """ Inititalize the driver run. @@ -301,7 +301,7 @@ def initialize( """ log.info("initialize parallel runtime") log.info(f"reading configuration: experiment {experiment_type}") - config = driver_config.read_config(experiment_type) + config = driver_config.read_config(icon4py_driver_backend, experiment_type) decomp_info = driver_init.read_decomp_info( file_path, props, serialization_type, grid_id, grid_root, grid_level @@ -310,6 +310,7 @@ def initialize( log.info(f"initializing the grid from '{file_path}'") icon_grid = driver_init.read_icon_grid( file_path, + backend=icon4py_driver_backend, rank=props.rank, ser_type=serialization_type, grid_id=grid_id, @@ -325,6 +326,7 @@ def initialize( ) = driver_init.read_geometry_fields( file_path, vertical_grid_config=config.vertical_grid_config, + backend=config.run_config.backend, rank=props.rank, ser_type=serialization_type, grid_id=grid_id, @@ -357,13 +359,13 @@ def initialize( edge_geometry, cell_geometry, exchange=exchange, - backend=gtfn_cpu, + backend=config.run_config.backend, ) nonhydro_params = solve_nh.NonHydrostaticParams(config.solve_nonhydro_config) solve_nonhydro_granule = solve_nh.SolveNonhydro( - backend=gtfn_cpu, + grid=icon_grid, config=config.solve_nonhydro_config, params=nonhydro_params, metric_state_nonhydro=solve_nonhydro_metric_state, @@ -372,6 +374,7 @@ def initialize( edge_geometry=edge_geometry, cell_geometry=cell_geometry, owner_mask=c_owner_mask, + backend=config.run_config.backend, ) ( @@ -387,6 +390,7 @@ def initialize( cell_geometry, edge_geometry, file_path, + backend=config.run_config.backend, rank=props.rank, experiment_type=experiment_type, ) @@ -423,13 +427,13 @@ def initialize( ) @click.option( "--serialization_type", - default="serialbox", + default=driver_init.SerializationType.SB.value, show_default=True, help="Serialization type for grid info and static fields. This is currently the only possible way to load the grid info and static fields.", ) @click.option( "--experiment_type", - default="any", + default=driver_init.ExperimentType.ANY.value, show_default=True, help="Option for configuration and how the initial state is generated. " "Setting it to the default value will instruct the model to use the default configuration of MeteoSwiss regional experiment and read the initial state from serialized data. " @@ -452,8 +456,29 @@ def initialize( default="af122aca-1dd2-11b2-a7f8-c7bf6bc21eba", help="uuid of the horizontal grid ('uuidOfHGrid' from gridfile)", ) +@click.option( + "--enable_output", + is_flag=True, + help="Enable all debugging messages. Otherwise, only critical error messages are printed.", +) +@click.option( + "--icon4py_driver_backend", + "-b", + default=driver_config.DriverBackends.GTFN_CPU.value, + show_default=True, + help="Backend for all components executed in icon4py driver. Choose between GTFN_CPU or GTFN_GPU. Please see abs_path_to_icon4py/model/driver/src/icon4py/model/driver/icon4py_configuration/) ", +) def icon4py_driver( - input_path, run_path, mpi, serialization_type, experiment_type, grid_id, grid_root, grid_level + input_path, + run_path, + mpi, + serialization_type, + experiment_type, + grid_id, + grid_root, + grid_level, + enable_output, + icon4py_driver_backend, ): """ usage: python dycore_driver.py abs_path_to_icon4py/testdata/ser_icondata/mpitask1/mch_ch_r04b09_dsl/ser_data @@ -477,7 +502,7 @@ def icon4py_driver( """ parallel_props = decomposition.get_processor_properties(decomposition.get_runtype(with_mpi=mpi)) grid_id = uuid.UUID(grid_id) - driver_init.configure_logging(run_path, experiment_type, parallel_props) + driver_init.configure_logging(run_path, experiment_type, enable_output, parallel_props) ( timeloop, diffusion_diagnostic_state, @@ -494,6 +519,7 @@ def icon4py_driver( grid_id, grid_root, grid_level, + icon4py_driver_backend, ) log.info(f"Starting ICON dycore run: {timeloop.simulation_date.isoformat()}") log.info( diff --git a/model/driver/src/icon4py/model/driver/initialization_utils.py b/model/driver/src/icon4py/model/driver/initialization_utils.py index f70669d299..a17b71f495 100644 --- a/model/driver/src/icon4py/model/driver/initialization_utils.py +++ b/model/driver/src/icon4py/model/driver/initialization_utils.py @@ -11,6 +11,8 @@ import logging import pathlib +from gt4py.next import backend as gt4py_backend + from icon4py.model.atmosphere.diffusion import diffusion_states from icon4py.model.atmosphere.dycore import dycore_states from icon4py.model.common import dimension as dims, field_type_aliases as fa @@ -30,6 +32,7 @@ ) from icon4py.model.common.utils import gt4py_field_allocation as field_alloc from icon4py.model.driver import ( + icon4py_configuration as driver_config, serialbox_helpers as driver_sb, ) from icon4py.model.driver.test_cases import gauss3d, jablonowski_williamson @@ -62,6 +65,7 @@ class ExperimentType(str, enum.Enum): def read_icon_grid( path: pathlib.Path, + backend: gt4py_backend.Backend, rank=0, ser_type: SerializationType = SerializationType.SB, grid_id=GLOBAL_GRID_ID, @@ -84,14 +88,14 @@ def read_icon_grid( return ( sb.IconSerialDataProvider("icon_pydycore", str(path.absolute()), False, mpi_rank=rank) .from_savepoint_grid(grid_id, grid_root, grid_level) - .construct_icon_grid(on_gpu=False) + .construct_icon_grid(on_gpu=(backend in driver_config.gpu_backends)) ) else: raise NotImplementedError(SB_ONLY_MSG) def model_initialization_serialbox( - grid: icon_grid.IconGrid, path: pathlib.Path, rank=0 + grid: icon_grid.IconGrid, path: pathlib.Path, backend: gt4py_backend.Backend, rank=0 ) -> tuple[ diffusion_states.DiffusionDiagnosticState, dycore_states.DiagnosticStateNonHydro, @@ -153,14 +157,20 @@ def model_initialization_serialbox( ) diagnostic_state = diagnostics.DiagnosticState( - pressure=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), + pressure=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), pressure_ifc=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + temperature=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), + virtual_temperature=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend ), - temperature=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - virtual_temperature=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - u=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - v=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), + u=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, backend=backend), + v=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, backend=backend), ) prognostic_state_next = prognostics.PrognosticState( @@ -175,7 +185,9 @@ def model_initialization_serialbox( vn_traj=solve_nonhydro_init_savepoint.vn_traj(), mass_flx_me=solve_nonhydro_init_savepoint.mass_flx_me(), mass_flx_ic=solve_nonhydro_init_savepoint.mass_flx_ic(), - vol_flx_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), + vol_flx_ic=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), ) return ( @@ -194,6 +206,7 @@ def read_initial_state( cell_param: grid_states.CellParams, edge_param: grid_states.EdgeParams, path: pathlib.Path, + backend: gt4py_backend.Backend, rank=0, experiment_type: ExperimentType = ExperimentType.ANY, ) -> tuple[ @@ -213,6 +226,7 @@ def read_initial_state( cell_param: cell properties edge_param: edge properties path: path to the serialized input data + backend: GT4Py backend rank: mpi rank of the current compute node experiment_type: (optional) defaults to ANY=any, type of initial condition to be read @@ -230,7 +244,12 @@ def read_initial_state( prognostic_state_now, prognostic_state_next, ) = jablonowski_williamson.model_initialization_jabw( - grid, cell_param, edge_param, path, rank + grid=grid, + cell_param=cell_param, + edge_param=edge_param, + path=path, + backend=backend, + rank=rank, ) elif experiment_type == ExperimentType.GAUSS3D: ( @@ -241,7 +260,13 @@ def read_initial_state( diagnostic_state, prognostic_state_now, prognostic_state_next, - ) = gauss3d.model_initialization_gauss3d(grid, edge_param, path, rank) + ) = gauss3d.model_initialization_gauss3d( + grid=grid, + edge_param=edge_param, + path=path, + backend=backend, + rank=rank, + ) elif experiment_type == ExperimentType.ANY: ( diffusion_diagnostic_state, @@ -251,7 +276,12 @@ def read_initial_state( diagnostic_state, prognostic_state_now, prognostic_state_next, - ) = model_initialization_serialbox(grid, path, rank) + ) = model_initialization_serialbox( + grid=grid, + path=path, + backend=backend, + rank=rank, + ) else: raise NotImplementedError(INITIALIZATION_ERROR_MSG) @@ -269,6 +299,7 @@ def read_initial_state( def read_geometry_fields( path: pathlib.Path, vertical_grid_config: v_grid.VerticalGridConfig, + backend: gt4py_backend.Backend, rank=0, ser_type: SerializationType = SerializationType.SB, grid_id=GLOBAL_GRID_ID, @@ -299,7 +330,7 @@ def read_geometry_fields( sp = _grid_savepoint(path, rank, grid_id, grid_root, grid_level) edge_geometry = sp.construct_edge_geometry() cell_geometry = sp.construct_cell_geometry() - vct_a, vct_b = v_grid.get_vct_a_and_vct_b(vertical_grid_config) + vct_a, vct_b = v_grid.get_vct_a_and_vct_b(vertical_grid_config, backend) vertical_geometry = v_grid.VerticalGrid( config=vertical_grid_config, vct_a=vct_a, @@ -448,7 +479,10 @@ def read_static_fields( def configure_logging( - run_path: str, experiment_name: str, processor_procs: decomposition.ProcessProperties = None + run_path: str, + experiment_name: str, + enable_output: bool = True, + processor_procs: decomposition.ProcessProperties = None, ) -> None: """ Configure logging. @@ -466,8 +500,9 @@ def configure_logging( run_dir.mkdir(exist_ok=True) logfile = run_dir.joinpath(f"dummy_dycore_driver_{experiment_name}.log") logfile.touch(exist_ok=True) + logging_level = logging.DEBUG if enable_output else logging.CRITICAL logging.basicConfig( - level=logging.DEBUG, + level=logging_level, format="%(asctime)s %(filename)-20s (%(lineno)-4d) : %(funcName)-20s: %(levelname)-8s %(message)s", filemode="w", filename=logfile, @@ -479,5 +514,5 @@ def configure_logging( log_format = "{rank} {asctime} - {filename}: {funcName:<20}: {levelname:<7} {message}" formatter = logging.Formatter(fmt=log_format, style="{", defaults={"rank": None}) console_handler.setFormatter(formatter) - console_handler.setLevel(logging.DEBUG) + console_handler.setLevel(logging_level) logging.getLogger("").addHandler(console_handler) diff --git a/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py b/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py index b708225543..b29fe1383b 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py +++ b/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py @@ -9,6 +9,7 @@ import pathlib import gt4py.next as gtx +from gt4py.next import backend as gt4py_backend from icon4py.model.atmosphere.diffusion import diffusion_states from icon4py.model.atmosphere.dycore import dycore_states @@ -18,13 +19,15 @@ cell_2_edge_interpolation, edge_2_cell_vector_rbf_interpolation, ) -from icon4py.model.common.settings import xp from icon4py.model.common.states import ( diagnostic_state as diagnostics, prognostic_state as prognostics, ) from icon4py.model.common.test_utils import serialbox_utils as sb -from icon4py.model.common.utils import gt4py_field_allocation as field_alloc +from icon4py.model.common.utils import ( + array_allocation as array_alloc, + gt4py_field_allocation as field_alloc, +) from icon4py.model.driver.test_cases import utils as testcases_utils @@ -35,6 +38,7 @@ def model_initialization_gauss3d( grid: icon_grid.IconGrid, edge_param: grid_states.EdgeParams, path: pathlib.Path, + backend: gt4py_backend.Backend, rank=0, ) -> tuple[ diffusion_states.DiffusionDiagnosticState, @@ -52,6 +56,7 @@ def model_initialization_gauss3d( grid: IconGrid edge_param: edge properties path: path where to find the input data + backend: GT4Py backend rank: mpi rank of the current compute node Returns: A tuple containing Diagnostic variables for diffusion and solve_nonhydro granules, PrepAdvection, second order divdamp factor, diagnostic variables, and two prognostic @@ -61,15 +66,18 @@ def model_initialization_gauss3d( "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) - wgtfac_c = data_provider.from_metrics_savepoint().wgtfac_c().asnumpy() - ddqz_z_half = data_provider.from_metrics_savepoint().ddqz_z_half().asnumpy() - theta_ref_mc = data_provider.from_metrics_savepoint().theta_ref_mc().asnumpy() - theta_ref_ic = data_provider.from_metrics_savepoint().theta_ref_ic().asnumpy() - exner_ref_mc = data_provider.from_metrics_savepoint().exner_ref_mc().asnumpy() - d_exner_dz_ref_ic = data_provider.from_metrics_savepoint().d_exner_dz_ref_ic().asnumpy() - geopot = data_provider.from_metrics_savepoint().geopot().asnumpy() + is_cupy = array_alloc.is_cupy_device(backend) + xp = array_alloc.array_ns(is_cupy) + + wgtfac_c = data_provider.from_metrics_savepoint().wgtfac_c().ndarray + ddqz_z_half = data_provider.from_metrics_savepoint().ddqz_z_half().ndarray + theta_ref_mc = data_provider.from_metrics_savepoint().theta_ref_mc().ndarray + theta_ref_ic = data_provider.from_metrics_savepoint().theta_ref_ic().ndarray + exner_ref_mc = data_provider.from_metrics_savepoint().exner_ref_mc().ndarray + d_exner_dz_ref_ic = data_provider.from_metrics_savepoint().d_exner_dz_ref_ic().ndarray + geopot = data_provider.from_metrics_savepoint().geopot().ndarray - primal_normal_x = edge_param.primal_normal[0].asnumpy() + primal_normal_x = edge_param.primal_normal[0].ndarray cell_2_edge_coeff = data_provider.from_interpolation_savepoint().c_lin_e() rbf_vec_coeff_c1 = data_provider.from_interpolation_savepoint().rbf_vec_coeff_c1() @@ -90,13 +98,13 @@ def model_initialization_gauss3d( ) end_cell_end = grid.end_index(cell_domain(h_grid.Zone.END)) - w_numpy = xp.zeros((num_cells, num_levels + 1), dtype=float) - exner_numpy = xp.zeros((num_cells, num_levels), dtype=float) - rho_numpy = xp.zeros((num_cells, num_levels), dtype=float) - temperature_numpy = xp.zeros((num_cells, num_levels), dtype=float) - pressure_numpy = xp.zeros((num_cells, num_levels), dtype=float) - theta_v_numpy = xp.zeros((num_cells, num_levels), dtype=float) - eta_v_numpy = xp.zeros((num_cells, num_levels), dtype=float) + w_ndarray = xp.zeros((num_cells, num_levels + 1), dtype=float) + exner_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + rho_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + temperature_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + pressure_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + theta_v_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + eta_v_ndarray = xp.zeros((num_cells, num_levels), dtype=float) mask_array_edge_start_plus1_to_edge_end = xp.ones(num_edges, dtype=bool) mask_array_edge_start_plus1_to_edge_end[0:end_edge_lateral_boundary_level_2] = False @@ -120,44 +128,44 @@ def model_initialization_gauss3d( # Horizontal wind field u = xp.where(mask, nh_u0, 0.0) - vn_numpy = u * primal_normal_x + vn_ndarray = u * primal_normal_x log.info("Wind profile assigned.") # Vertical temperature profile for k_index in range(num_levels - 1, -1, -1): z_help = (nh_brunt_vais / phy_const.GRAV) ** 2 * geopot[:, k_index] # profile of theta is explicitly given - theta_v_numpy[:, k_index] = nh_t0 * xp.exp(z_help) + theta_v_ndarray[:, k_index] = nh_t0 * xp.exp(z_help) # Lower boundary condition for exner pressure if nh_brunt_vais != 0.0: z_help = (nh_brunt_vais / phy_const.GRAV) ** 2 * geopot[:, num_levels - 1] - exner_numpy[:, num_levels - 1] = ( + exner_ndarray[:, num_levels - 1] = ( phy_const.GRAV / nh_brunt_vais ) ** 2 / nh_t0 / phy_const.CPD * (xp.exp(-z_help) - 1.0) + 1.0 else: - exner_numpy[:, num_levels - 1] = 1.0 - geopot[:, num_levels - 1] / phy_const.CPD / nh_t0 + exner_ndarray[:, num_levels - 1] = 1.0 - geopot[:, num_levels - 1] / phy_const.CPD / nh_t0 log.info("Vertical computations completed.") # Compute hydrostatically balanced exner, by integrating the (discretized!) # 3rd equation of motion under the assumption thetav=const. - rho_numpy, exner_numpy = testcases_utils.hydrostatic_adjustment_constant_thetav_numpy( + rho_ndarray, exner_ndarray = testcases_utils.hydrostatic_adjustment_constant_thetav_ndarray( wgtfac_c, ddqz_z_half, exner_ref_mc, d_exner_dz_ref_ic, theta_ref_mc, theta_ref_ic, - rho_numpy, - exner_numpy, - theta_v_numpy, + rho_ndarray, + exner_ndarray, + theta_v_ndarray, num_levels, ) log.info("Hydrostatic adjustment computation completed.") - eta_v = gtx.as_field((dims.CellDim, dims.KDim), eta_v_numpy) - eta_v_e = field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid) - cell_2_edge_interpolation.cell_2_edge_interpolation( + eta_v = gtx.as_field((dims.CellDim, dims.KDim), eta_v_ndarray, allocator=backend) + eta_v_e = field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid, backend=backend) + cell_2_edge_interpolation.cell_2_edge_interpolation.with_backend(backend)( eta_v, cell_2_edge_coeff, eta_v_e, @@ -169,29 +177,38 @@ def model_initialization_gauss3d( ) log.info("Cell-to-edge eta_v computation completed.") - vn = gtx.as_field((dims.EdgeDim, dims.KDim), vn_numpy) - w = gtx.as_field((dims.CellDim, dims.KDim), w_numpy) - exner = gtx.as_field((dims.CellDim, dims.KDim), exner_numpy) - rho = gtx.as_field((dims.CellDim, dims.KDim), rho_numpy) - temperature = gtx.as_field((dims.CellDim, dims.KDim), temperature_numpy) - virtual_temperature = gtx.as_field((dims.CellDim, dims.KDim), temperature_numpy) - pressure = gtx.as_field((dims.CellDim, dims.KDim), pressure_numpy) - theta_v = gtx.as_field((dims.CellDim, dims.KDim), theta_v_numpy) - pressure_ifc_numpy = xp.zeros((num_cells, num_levels + 1), dtype=float) - pressure_ifc_numpy[ - :, -1 - ] = phy_const.P0REF # set surface pressure to the prescribed value (only used for IC in JABW test case, then actually computed in the dycore) - pressure_ifc = gtx.as_field((dims.CellDim, dims.KDim), pressure_ifc_numpy) - - vn_next = gtx.as_field((dims.EdgeDim, dims.KDim), vn_numpy) - w_next = gtx.as_field((dims.CellDim, dims.KDim), w_numpy) - exner_next = gtx.as_field((dims.CellDim, dims.KDim), exner_numpy) - rho_next = gtx.as_field((dims.CellDim, dims.KDim), rho_numpy) - theta_v_next = gtx.as_field((dims.CellDim, dims.KDim), theta_v_numpy) + pressure_ifc_ndarray = xp.zeros((num_cells, num_levels + 1), dtype=float) + ( + vn, + w, + exner, + rho, + theta_v, + vn_next, + w_next, + exner_next, + rho_next, + theta_v_next, + temperature, + virtual_temperature, + pressure, + pressure_ifc, + u, + v, + ) = testcases_utils.create_gt4py_field_for_prognostic_and_diagnostic_variables( + vn_ndarray, + w_ndarray, + exner_ndarray, + rho_ndarray, + theta_v_ndarray, + temperature_ndarray, + pressure_ndarray, + pressure_ifc_ndarray, + grid=grid, + backend=backend, + ) - u = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid) - v = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid) - edge_2_cell_vector_rbf_interpolation.edge_2_cell_vector_rbf_interpolation( + edge_2_cell_vector_rbf_interpolation.edge_2_cell_vector_rbf_interpolation.with_backend(backend)( vn, rbf_vec_coeff_c1, rbf_vec_coeff_c2, @@ -205,8 +222,8 @@ def model_initialization_gauss3d( ) log.info("U, V computation completed.") - exner_pr = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid) - testcases_utils.compute_perturbed_exner( + exner_pr = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, backend=backend) + testcases_utils.compute_perturbed_exner.with_backend(backend)( exner, data_provider.from_metrics_savepoint().exner_ref_mc(), exner_pr, @@ -242,54 +259,14 @@ def model_initialization_gauss3d( exner=exner_next, ) - diffusion_diagnostic_state = diffusion_states.DiffusionDiagnosticState( - hdef_ic=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - div_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, is_halfdim=True), - dwdx=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, is_halfdim=True), - dwdy=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, is_halfdim=True), + diffusion_diagnostic_state = testcases_utils.initialize_diffusion_diagnostic_state( + grid=grid, backend=backend ) - solve_nonhydro_diagnostic_state = dycore_states.DiagnosticStateNonHydro( - theta_v_ic=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - exner_pr=exner_pr, - rho_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, is_halfdim=True), - ddt_exner_phy=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - grf_tend_rho=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - grf_tend_thv=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - grf_tend_w=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - mass_fl_e=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - ddt_vn_phy=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - grf_tend_vn=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - ddt_vn_apc_ntl1=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - ddt_vn_apc_ntl2=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - ddt_w_adv_ntl1=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - ddt_w_adv_ntl2=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - vt=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - vn_ie=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid, is_halfdim=True), - w_concorr_c=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - rho_incr=None, # solve_nonhydro_init_savepoint.rho_incr(), - vn_incr=None, # solve_nonhydro_init_savepoint.vn_incr(), - exner_incr=None, # solve_nonhydro_init_savepoint.exner_incr(), - exner_dyn_incr=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), + solve_nonhydro_diagnostic_state = testcases_utils.initialize_solve_nonhydro_diagnostic_state( + exner_pr=exner_pr, grid=grid, backend=backend ) - prep_adv = dycore_states.PrepAdvection( - vn_traj=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - mass_flx_me=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - mass_flx_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - vol_flx_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - ) + prep_adv = testcases_utils.initialize_prep_advection(grid=grid, backend=backend) log.info("Initialization completed.") return ( diff --git a/model/driver/src/icon4py/model/driver/test_cases/jablonowski_williamson.py b/model/driver/src/icon4py/model/driver/test_cases/jablonowski_williamson.py index 2ffe558d80..6a1ed7cc34 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/jablonowski_williamson.py +++ b/model/driver/src/icon4py/model/driver/test_cases/jablonowski_williamson.py @@ -10,6 +10,7 @@ import pathlib import gt4py.next as gtx +from gt4py.next import backend as gt4py_backend from icon4py.model.atmosphere.diffusion import diffusion_states from icon4py.model.atmosphere.dycore import dycore_states @@ -19,13 +20,15 @@ cell_2_edge_interpolation, edge_2_cell_vector_rbf_interpolation, ) -from icon4py.model.common.settings import xp from icon4py.model.common.states import ( diagnostic_state as diagnostics, prognostic_state as prognostics, ) from icon4py.model.common.test_utils import serialbox_utils as sb -from icon4py.model.common.utils import gt4py_field_allocation as field_alloc +from icon4py.model.common.utils import ( + array_allocation as array_alloc, + gt4py_field_allocation as field_alloc, +) from icon4py.model.driver.test_cases import utils as testcases_utils @@ -37,6 +40,7 @@ def model_initialization_jabw( cell_param: grid_states.CellParams, edge_param: grid_states.EdgeParams, path: pathlib.Path, + backend: gt4py_backend.Backend, rank=0, ) -> tuple[ diffusion_states.DiffusionDiagnosticState, @@ -56,6 +60,7 @@ def model_initialization_jabw( cell_param: cell properties edge_param: edge properties path: path where to find the input data + backend: GT4Py backend rank: mpi rank of the current compute node Returns: A tuple containing Diagnostic variables for diffusion and solve_nonhydro granules, PrepAdvection, second order divdamp factor, diagnostic variables, and two prognostic @@ -65,18 +70,21 @@ def model_initialization_jabw( "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) - wgtfac_c = data_provider.from_metrics_savepoint().wgtfac_c().asnumpy() - ddqz_z_half = data_provider.from_metrics_savepoint().ddqz_z_half().asnumpy() - theta_ref_mc = data_provider.from_metrics_savepoint().theta_ref_mc().asnumpy() - theta_ref_ic = data_provider.from_metrics_savepoint().theta_ref_ic().asnumpy() - exner_ref_mc = data_provider.from_metrics_savepoint().exner_ref_mc().asnumpy() - d_exner_dz_ref_ic = data_provider.from_metrics_savepoint().d_exner_dz_ref_ic().asnumpy() - geopot = data_provider.from_metrics_savepoint().geopot().asnumpy() + is_cupy = array_alloc.is_cupy_device(backend) + xp = array_alloc.array_ns(is_cupy) + + wgtfac_c = data_provider.from_metrics_savepoint().wgtfac_c().ndarray + ddqz_z_half = data_provider.from_metrics_savepoint().ddqz_z_half().ndarray + theta_ref_mc = data_provider.from_metrics_savepoint().theta_ref_mc().ndarray + theta_ref_ic = data_provider.from_metrics_savepoint().theta_ref_ic().ndarray + exner_ref_mc = data_provider.from_metrics_savepoint().exner_ref_mc().ndarray + d_exner_dz_ref_ic = data_provider.from_metrics_savepoint().d_exner_dz_ref_ic().ndarray + geopot = data_provider.from_metrics_savepoint().geopot().ndarray - cell_lat = cell_param.cell_center_lat.asnumpy() - edge_lat = edge_param.edge_center[0].asnumpy() - edge_lon = edge_param.edge_center[1].asnumpy() - primal_normal_x = edge_param.primal_normal[0].asnumpy() + cell_lat = cell_param.cell_center_lat.ndarray + edge_lat = edge_param.edge_center[0].ndarray + edge_lon = edge_param.edge_center[1].ndarray + primal_normal_x = edge_param.primal_normal[0].ndarray cell_2_edge_coeff = data_provider.from_interpolation_savepoint().c_lin_e() rbf_vec_coeff_c1 = data_provider.from_interpolation_savepoint().rbf_vec_coeff_c1() @@ -110,13 +118,13 @@ def model_initialization_jabw( lat_perturbation_center = 2.0 * lon_perturbation_center # latitude of the perturb centre ps_o_p0ref = p_sfc / phy_const.P0REF - w_numpy = xp.zeros((num_cells, num_levels + 1), dtype=float) - exner_numpy = xp.zeros((num_cells, num_levels), dtype=float) - rho_numpy = xp.zeros((num_cells, num_levels), dtype=float) - temperature_numpy = xp.zeros((num_cells, num_levels), dtype=float) - pressure_numpy = xp.zeros((num_cells, num_levels), dtype=float) - theta_v_numpy = xp.zeros((num_cells, num_levels), dtype=float) - eta_v_numpy = xp.zeros((num_cells, num_levels), dtype=float) + w_ndarray = xp.zeros((num_cells, num_levels + 1), dtype=float) + exner_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + rho_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + temperature_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + pressure_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + theta_v_ndarray = xp.zeros((num_cells, num_levels), dtype=float) + eta_v_ndarray = xp.zeros((num_cells, num_levels), dtype=float) sin_lat = xp.sin(cell_lat) cos_lat = xp.cos(cell_lat) @@ -132,9 +140,9 @@ def model_initialization_jabw( log.info(f"In Newton iteration, k = {k_index}") # Newton iteration to determine zeta for _ in range(100): - eta_v_numpy[:, k_index] = (eta_old - eta_0) * math.pi * 0.5 - cos_etav = xp.cos(eta_v_numpy[:, k_index]) - sin_etav = xp.sin(eta_v_numpy[:, k_index]) + eta_v_ndarray[:, k_index] = (eta_old - eta_0) * math.pi * 0.5 + cos_etav = xp.cos(eta_v_ndarray[:, k_index]) + sin_etav = xp.sin(eta_v_ndarray[:, k_index]) temperature_avg = jw_temp0 * (eta_old**lapse_rate) geopot_avg = jw_temp0 * phy_const.GRAV / gamma * (1.0 - eta_old**lapse_rate) @@ -176,23 +184,25 @@ def model_initialization_jabw( eta_old = eta_old - newton_function / newton_function_prime # Final update for zeta_v - eta_v_numpy[:, k_index] = (eta_old - eta_0) * math.pi * 0.5 + eta_v_ndarray[:, k_index] = (eta_old - eta_0) * math.pi * 0.5 # Use analytic expressions at all model level - exner_numpy[:, k_index] = (eta_old * ps_o_p0ref) ** phy_const.RD_O_CPD - theta_v_numpy[:, k_index] = temperature_jw / exner_numpy[:, k_index] - rho_numpy[:, k_index] = ( - exner_numpy[:, k_index] ** phy_const.CVD_O_RD + exner_ndarray[:, k_index] = (eta_old * ps_o_p0ref) ** phy_const.RD_O_CPD + theta_v_ndarray[:, k_index] = temperature_jw / exner_ndarray[:, k_index] + rho_ndarray[:, k_index] = ( + exner_ndarray[:, k_index] ** phy_const.CVD_O_RD * phy_const.P0REF / phy_const.RD - / theta_v_numpy[:, k_index] + / theta_v_ndarray[:, k_index] ) # initialize diagnose pressure and temperature variables - pressure_numpy[:, k_index] = phy_const.P0REF * exner_numpy[:, k_index] ** phy_const.CPD_O_RD - temperature_numpy[:, k_index] = temperature_jw + pressure_ndarray[:, k_index] = ( + phy_const.P0REF * exner_ndarray[:, k_index] ** phy_const.CPD_O_RD + ) + temperature_ndarray[:, k_index] = temperature_jw log.info("Newton iteration completed!") - eta_v = gtx.as_field((dims.CellDim, dims.KDim), eta_v_numpy) - eta_v_e = field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid) + eta_v = gtx.as_field((dims.CellDim, dims.KDim), eta_v_ndarray, allocator=backend) + eta_v_e = field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid, backend=backend) cell_2_edge_interpolation.cell_2_edge_interpolation( eta_v, cell_2_edge_coeff, @@ -205,7 +215,7 @@ def model_initialization_jabw( ) log.info("Cell-to-edge eta_v computation completed.") - vn_numpy = testcases_utils.zonalwind_2_normalwind_numpy( + vn_ndarray = testcases_utils.zonalwind_2_normalwind_ndarray( grid, jw_u0, jw_up, @@ -214,45 +224,59 @@ def model_initialization_jabw( edge_lat, edge_lon, primal_normal_x, - eta_v_e.asnumpy(), + eta_v_e.ndarray, + backend, ) log.info("U2vn computation completed.") - rho_numpy, exner_numpy, theta_v_numpy = testcases_utils.hydrostatic_adjustment_numpy( + rho_ndarray, exner_ndarray, theta_v_ndarray = testcases_utils.hydrostatic_adjustment_ndarray( wgtfac_c, ddqz_z_half, exner_ref_mc, d_exner_dz_ref_ic, theta_ref_mc, theta_ref_ic, - rho_numpy, - exner_numpy, - theta_v_numpy, + rho_ndarray, + exner_ndarray, + theta_v_ndarray, num_levels, + backend, ) log.info("Hydrostatic adjustment computation completed.") - vn = gtx.as_field((dims.EdgeDim, dims.KDim), vn_numpy) - w = gtx.as_field((dims.CellDim, dims.KDim), w_numpy) - exner = gtx.as_field((dims.CellDim, dims.KDim), exner_numpy) - rho = gtx.as_field((dims.CellDim, dims.KDim), rho_numpy) - temperature = gtx.as_field((dims.CellDim, dims.KDim), temperature_numpy) - virutal_temperature = gtx.as_field((dims.CellDim, dims.KDim), temperature_numpy) - pressure = gtx.as_field((dims.CellDim, dims.KDim), pressure_numpy) - theta_v = gtx.as_field((dims.CellDim, dims.KDim), theta_v_numpy) - pressure_ifc_numpy = xp.zeros((num_cells, num_levels + 1), dtype=float) - pressure_ifc_numpy[:, -1] = p_sfc - pressure_ifc = gtx.as_field((dims.CellDim, dims.KDim), pressure_ifc_numpy) - - vn_next = gtx.as_field((dims.EdgeDim, dims.KDim), vn_numpy) - w_next = gtx.as_field((dims.CellDim, dims.KDim), w_numpy) - exner_next = gtx.as_field((dims.CellDim, dims.KDim), exner_numpy) - rho_next = gtx.as_field((dims.CellDim, dims.KDim), rho_numpy) - theta_v_next = gtx.as_field((dims.CellDim, dims.KDim), theta_v_numpy) + pressure_ifc_ndarray = xp.zeros((num_cells, num_levels + 1), dtype=float) + pressure_ifc_ndarray[:, -1] = p_sfc + ( + vn, + w, + exner, + rho, + theta_v, + vn_next, + w_next, + exner_next, + rho_next, + theta_v_next, + temperature, + virtual_temperature, + pressure, + pressure_ifc, + u, + v, + ) = testcases_utils.create_gt4py_field_for_prognostic_and_diagnostic_variables( + vn_ndarray, + w_ndarray, + exner_ndarray, + rho_ndarray, + theta_v_ndarray, + temperature_ndarray, + pressure_ndarray, + pressure_ifc_ndarray, + grid=grid, + backend=backend, + ) - u = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid) - v = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid) - edge_2_cell_vector_rbf_interpolation.edge_2_cell_vector_rbf_interpolation( + edge_2_cell_vector_rbf_interpolation.edge_2_cell_vector_rbf_interpolation.with_backend(backend)( vn, rbf_vec_coeff_c1, rbf_vec_coeff_c2, @@ -267,8 +291,8 @@ def model_initialization_jabw( log.info("U, V computation completed.") - exner_pr = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid) - testcases_utils.compute_perturbed_exner( + exner_pr = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, backend=backend) + testcases_utils.compute_perturbed_exner.with_backend(backend)( exner, data_provider.from_metrics_savepoint().exner_ref_mc(), exner_pr, @@ -284,7 +308,7 @@ def model_initialization_jabw( pressure=pressure, pressure_ifc=pressure_ifc, temperature=temperature, - virtual_temperature=virutal_temperature, + virtual_temperature=virtual_temperature, u=u, v=v, ) @@ -304,54 +328,15 @@ def model_initialization_jabw( exner=exner_next, ) - diffusion_diagnostic_state = diffusion_states.DiffusionDiagnosticState( - hdef_ic=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - div_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, is_halfdim=True), - dwdx=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, is_halfdim=True), - dwdy=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, is_halfdim=True), + diffusion_diagnostic_state = testcases_utils.initialize_diffusion_diagnostic_state( + grid=grid, backend=backend ) - solve_nonhydro_diagnostic_state = dycore_states.DiagnosticStateNonHydro( - theta_v_ic=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), + solve_nonhydro_diagnostic_state = testcases_utils.initialize_solve_nonhydro_diagnostic_state( exner_pr=exner_pr, - rho_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, is_halfdim=True), - ddt_exner_phy=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - grf_tend_rho=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - grf_tend_thv=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - grf_tend_w=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - mass_fl_e=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - ddt_vn_phy=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - grf_tend_vn=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - ddt_vn_apc_ntl1=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - ddt_vn_apc_ntl2=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - ddt_w_adv_ntl1=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - ddt_w_adv_ntl2=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - vt=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - vn_ie=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid, is_halfdim=True), - w_concorr_c=field_alloc.allocate_zero_field( - dims.CellDim, dims.KDim, grid=grid, is_halfdim=True - ), - rho_incr=None, # solve_nonhydro_init_savepoint.rho_incr(), - vn_incr=None, # solve_nonhydro_init_savepoint.vn_incr(), - exner_incr=None, # solve_nonhydro_init_savepoint.exner_incr(), - exner_dyn_incr=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - ) - - prep_adv = dycore_states.PrepAdvection( - vn_traj=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - mass_flx_me=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid), - mass_flx_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), - vol_flx_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid), + grid=grid, + backend=backend, ) + prep_adv = testcases_utils.initialize_prep_advection(grid=grid, backend=backend) log.info("Initialization completed.") return ( diff --git a/model/driver/src/icon4py/model/driver/test_cases/utils.py b/model/driver/src/icon4py/model/driver/test_cases/utils.py index 40b243c119..9c66a859b4 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/utils.py +++ b/model/driver/src/icon4py/model/driver/test_cases/utils.py @@ -6,7 +6,10 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +from gt4py.next import backend as gt4py_backend +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, @@ -14,26 +17,28 @@ type_alias as ta, ) from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid -from icon4py.model.common.settings import backend, xp - - -# TODO: this will have to be removed once domain allows for imports -CellDim = dims.CellDim -KDim = dims.KDim +from icon4py.model.common.utils import ( + array_allocation as array_alloc, + gt4py_field_allocation as field_alloc, +) -def hydrostatic_adjustment_numpy( - wgtfac_c: xp.ndarray, - ddqz_z_half: xp.ndarray, - exner_ref_mc: xp.ndarray, - d_exner_dz_ref_ic: xp.ndarray, - theta_ref_mc: xp.ndarray, - theta_ref_ic: xp.ndarray, - rho: xp.ndarray, - exner: xp.ndarray, - theta_v: xp.ndarray, +def hydrostatic_adjustment_ndarray( + wgtfac_c: fa.AnyNDArray, + ddqz_z_half: fa.AnyNDArray, + exner_ref_mc: fa.AnyNDArray, + d_exner_dz_ref_ic: fa.AnyNDArray, + theta_ref_mc: fa.AnyNDArray, + theta_ref_ic: fa.AnyNDArray, + rho: fa.AnyNDArray, + exner: fa.AnyNDArray, + theta_v: fa.AnyNDArray, num_levels: int, -): + backend: gt4py_backend.Backend, +) -> tuple[fa.AnyNDArray, fa.AnyNDArray, fa.AnyNDArray]: + is_cupy = array_alloc.is_cupy_device(backend) + xp = array_alloc.array_ns(is_cupy) + # virtual temperature temp_v = theta_v * exner @@ -62,21 +67,21 @@ def hydrostatic_adjustment_numpy( return rho, exner, theta_v -def hydrostatic_adjustment_constant_thetav_numpy( - wgtfac_c: xp.ndarray, - ddqz_z_half: xp.ndarray, - exner_ref_mc: xp.ndarray, - d_exner_dz_ref_ic: xp.ndarray, - theta_ref_mc: xp.ndarray, - theta_ref_ic: xp.ndarray, - rho: xp.ndarray, - exner: xp.ndarray, - theta_v: xp.ndarray, +def hydrostatic_adjustment_constant_thetav_ndarray( + wgtfac_c: fa.AnyNDArray, + ddqz_z_half: fa.AnyNDArray, + exner_ref_mc: fa.AnyNDArray, + d_exner_dz_ref_ic: fa.AnyNDArray, + theta_ref_mc: fa.AnyNDArray, + theta_ref_ic: fa.AnyNDArray, + rho: fa.AnyNDArray, + exner: fa.AnyNDArray, + theta_v: fa.AnyNDArray, num_levels: int, -) -> tuple[xp.ndarray, xp.ndarray]: +) -> tuple[fa.AnyNDArray, fa.AnyNDArray]: """ Computes a hydrostatically balanced profile. In constrast to the above - hydrostatic_adjustment_numpy, the virtual temperature is kept (assumed) + hydrostatic_adjustment_ndarray, the virtual temperature is kept (assumed) constant during the adjustment, leading to a simpler formula. """ @@ -102,17 +107,18 @@ def hydrostatic_adjustment_constant_thetav_numpy( return rho, exner -def zonalwind_2_normalwind_numpy( +def zonalwind_2_normalwind_ndarray( grid: icon_grid.IconGrid, jw_u0: float, jw_up: float, lat_perturbation_center: float, lon_perturbation_center: float, - edge_lat: xp.ndarray, - edge_lon: xp.ndarray, - primal_normal_x: xp.ndarray, - eta_v_e: xp.ndarray, -): + edge_lat: fa.AnyNDArray, + edge_lon: fa.AnyNDArray, + primal_normal_x: fa.AnyNDArray, + eta_v_e: fa.AnyNDArray, + backend: gt4py_backend.Backend, +) -> fa.AnyNDArray: """ Compute normal wind at edge center from vertical eta coordinate (eta_v_e). @@ -130,6 +136,9 @@ def zonalwind_2_normalwind_numpy( """ # TODO (Chia Rui) this function needs a test + is_cupy = array_alloc.is_cupy_device(backend) + xp = array_alloc.array_ns(is_cupy) + mask = xp.ones((grid.num_edges, grid.num_levels), dtype=bool) mask[ 0 : grid.end_index(h_grid.domain(dims.EdgeDim)(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)), @@ -188,7 +197,7 @@ def _compute_perturbed_exner( return exner_pr -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED, backend=backend) +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def compute_perturbed_exner( exner: fa.CellKField[ta.wpfloat], exner_ref: fa.CellKField[ta.vpfloat], @@ -203,7 +212,175 @@ def compute_perturbed_exner( exner_ref, out=exner_pr, domain={ - CellDim: (horizontal_start, horizontal_end), - KDim: (vertical_start, vertical_end), + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), }, ) + + +def initialize_diffusion_diagnostic_state( + grid: icon_grid.IconGrid, backend: gt4py_backend.Backend +) -> diffusion_states.DiffusionDiagnosticState: + return diffusion_states.DiffusionDiagnosticState( + hdef_ic=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + div_ic=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + dwdx=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + dwdy=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + ) + + +def initialize_solve_nonhydro_diagnostic_state( + exner_pr: fa.CellKField[ta.wpfloat], grid: icon_grid.IconGrid, backend: gt4py_backend.Backend +) -> dycore_states.DiagnosticStateNonHydro: + return dycore_states.DiagnosticStateNonHydro( + theta_v_ic=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + exner_pr=exner_pr, + rho_ic=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + ddt_exner_phy=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), + grf_tend_rho=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), + grf_tend_thv=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), + grf_tend_w=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + mass_fl_e=field_alloc.allocate_zero_field( + dims.EdgeDim, dims.KDim, grid=grid, backend=backend + ), + ddt_vn_phy=field_alloc.allocate_zero_field( + dims.EdgeDim, dims.KDim, grid=grid, backend=backend + ), + grf_tend_vn=field_alloc.allocate_zero_field( + dims.EdgeDim, dims.KDim, grid=grid, backend=backend + ), + ddt_vn_apc_ntl1=field_alloc.allocate_zero_field( + dims.EdgeDim, dims.KDim, grid=grid, backend=backend + ), + ddt_vn_apc_ntl2=field_alloc.allocate_zero_field( + dims.EdgeDim, dims.KDim, grid=grid, backend=backend + ), + ddt_w_adv_ntl1=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + ddt_w_adv_ntl2=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + vt=field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid, backend=backend), + vn_ie=field_alloc.allocate_zero_field( + dims.EdgeDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + w_concorr_c=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, is_halfdim=True, backend=backend + ), + rho_incr=None, # solve_nonhydro_init_savepoint.rho_incr(), + vn_incr=None, # solve_nonhydro_init_savepoint.vn_incr(), + exner_incr=None, # solve_nonhydro_init_savepoint.exner_incr(), + exner_dyn_incr=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), + ) + + +def initialize_prep_advection( + grid: icon_grid.IconGrid, backend: gt4py_backend.Backend +) -> dycore_states.PrepAdvection: + return dycore_states.PrepAdvection( + vn_traj=field_alloc.allocate_zero_field( + dims.EdgeDim, dims.KDim, grid=grid, backend=backend + ), + mass_flx_me=field_alloc.allocate_zero_field( + dims.EdgeDim, dims.KDim, grid=grid, backend=backend + ), + mass_flx_ic=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), + vol_flx_ic=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=grid, backend=backend + ), + ) + + +def create_gt4py_field_for_prognostic_and_diagnostic_variables( + vn_ndarray: fa.AnyNDArray, + w_ndarray: fa.AnyNDArray, + exner_ndarray: fa.AnyNDArray, + rho_ndarray: fa.AnyNDArray, + theta_v_ndarray: fa.AnyNDArray, + temperature_ndarray: fa.AnyNDArray, + pressure_ndarray: fa.AnyNDArray, + pressure_ifc_ndarray: fa.AnyNDArray, + grid: icon_grid.IconGrid, + backend: gt4py_backend.Backend, +) -> tuple[ + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], +]: + vn = gtx.as_field((dims.EdgeDim, dims.KDim), vn_ndarray, allocator=backend) + w = gtx.as_field((dims.CellDim, dims.KDim), w_ndarray, allocator=backend) + exner = gtx.as_field((dims.CellDim, dims.KDim), exner_ndarray, allocator=backend) + rho = gtx.as_field((dims.CellDim, dims.KDim), rho_ndarray, allocator=backend) + temperature = gtx.as_field((dims.CellDim, dims.KDim), temperature_ndarray, allocator=backend) + virutal_temperature = gtx.as_field( + (dims.CellDim, dims.KDim), temperature_ndarray, allocator=backend + ) + pressure = gtx.as_field((dims.CellDim, dims.KDim), pressure_ndarray, allocator=backend) + theta_v = gtx.as_field((dims.CellDim, dims.KDim), theta_v_ndarray, allocator=backend) + pressure_ifc = gtx.as_field((dims.CellDim, dims.KDim), pressure_ifc_ndarray, allocator=backend) + + vn_next = gtx.as_field((dims.EdgeDim, dims.KDim), vn_ndarray, allocator=backend) + w_next = gtx.as_field((dims.CellDim, dims.KDim), w_ndarray, allocator=backend) + exner_next = gtx.as_field((dims.CellDim, dims.KDim), exner_ndarray, allocator=backend) + rho_next = gtx.as_field((dims.CellDim, dims.KDim), rho_ndarray, allocator=backend) + theta_v_next = gtx.as_field((dims.CellDim, dims.KDim), theta_v_ndarray, allocator=backend) + + u = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, backend=backend) + v = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, backend=backend) + + return ( + vn, + w, + exner, + rho, + theta_v, + vn_next, + w_next, + exner_next, + rho_next, + theta_v_next, + temperature, + virutal_temperature, + pressure, + pressure_ifc, + u, + v, + ) diff --git a/model/driver/tests/driver_tests/test_timeloop.py b/model/driver/tests/driver_tests/test_timeloop.py index 5ed02db104..4709bd6d5b 100644 --- a/model/driver/tests/driver_tests/test_timeloop.py +++ b/model/driver/tests/driver_tests/test_timeloop.py @@ -127,7 +127,11 @@ def test_run_timeloop_single_step( backend, ): if experiment == dt_utils.GAUSS3D_EXPERIMENT: - config = icon4py_configuration.read_config(experiment) + # it does not matter what backend is set here because the granules are set externally in this test + config = icon4py_configuration.read_config( + icon4py_driver_backend=icon4py_configuration.DriverBackends.GTFN_CPU.value, + experiment_type=experiment, + ) diffusion_config = config.diffusion_config nonhydro_config = config.solve_nonhydro_config icon4pyrun_config = config.run_config @@ -259,7 +263,9 @@ def test_run_timeloop_single_step( vn_traj=sp.vn_traj(), mass_flx_me=sp.mass_flx_me(), mass_flx_ic=sp.mass_flx_ic(), - vol_flx_ic=field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=icon_grid), + vol_flx_ic=field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=icon_grid, backend=backend + ), ) nonhydro_diagnostic_state = dycore_states.DiagnosticStateNonHydro( diff --git a/model/driver/tests/initial_condition_tests/test_gauss3d.py b/model/driver/tests/initial_condition_tests/test_gauss3d.py index f3daed8ff5..58b4697742 100644 --- a/model/driver/tests/initial_condition_tests/test_gauss3d.py +++ b/model/driver/tests/initial_condition_tests/test_gauss3d.py @@ -22,6 +22,7 @@ def test_gauss3d_initial_condition( experiment, ranked_data_path, + backend, rank, data_provider, grid_savepoint, @@ -41,6 +42,7 @@ def test_gauss3d_initial_condition( icon_grid, edge_geometry, ranked_data_path.joinpath(f"{experiment}/ser_data"), + backend, rank, ) diff --git a/model/driver/tests/initial_condition_tests/test_jablonowski_williamson.py b/model/driver/tests/initial_condition_tests/test_jablonowski_williamson.py index e2cb58a6af..52e36455ef 100644 --- a/model/driver/tests/initial_condition_tests/test_jablonowski_williamson.py +++ b/model/driver/tests/initial_condition_tests/test_jablonowski_williamson.py @@ -22,6 +22,7 @@ def test_jabw_initial_condition( experiment, ranked_data_path, + backend, rank, data_provider, grid_savepoint, @@ -43,6 +44,7 @@ def test_jabw_initial_condition( cell_geometry, edge_geometry, ranked_data_path.joinpath(f"{experiment}/ser_data"), + backend, rank, ) diff --git a/model/driver/tests/initial_condition_tests/test_utils.py b/model/driver/tests/initial_condition_tests/test_utils.py index f675bfe6b2..95c811ce01 100644 --- a/model/driver/tests/initial_condition_tests/test_utils.py +++ b/model/driver/tests/initial_condition_tests/test_utils.py @@ -11,7 +11,7 @@ from icon4py.model.driver.test_cases import utils -def test_hydrostatic_adjustment_numpy(): +def test_hydrostatic_adjustment_ndarray(backend): # TODO (Jacopo / Chia Rui) these tests could be better num_cells = 10 num_levels = 10 @@ -31,7 +31,7 @@ def test_hydrostatic_adjustment_numpy(): theta_v = theta_v0 * xp.ones((num_cells, num_levels)) # Call the function - r_rho, r_exner, r_theta_v = utils.hydrostatic_adjustment_numpy( + r_rho, r_exner, r_theta_v = utils.hydrostatic_adjustment_ndarray( wgtfac_c, ddqz_z_half, exner_ref_mc, @@ -42,6 +42,7 @@ def test_hydrostatic_adjustment_numpy(): exner, theta_v, num_levels, + backend, ) assert r_rho.shape == (num_cells, num_levels) @@ -66,7 +67,7 @@ def test_hydrostatic_adjustment_numpy(): ) -def test_hydrostatic_adjustment_constant_thetav_numpy(): +def test_hydrostatic_adjustment_constant_thetav_ndarray(): # TODO (Jacopo / Chia Rui) these tests could be better num_cells = 10 num_levels = 10 @@ -86,7 +87,7 @@ def test_hydrostatic_adjustment_constant_thetav_numpy(): theta_v = theta_v0 * xp.ones((num_cells, num_levels)) # Call the function - r_rho, r_exner = utils.hydrostatic_adjustment_constant_thetav_numpy( + r_rho, r_exner = utils.hydrostatic_adjustment_constant_thetav_ndarray( wgtfac_c, ddqz_z_half, exner_ref_mc,