From 9f8e5826bfbc80a4a759af6ea93d3d514abd346a Mon Sep 17 00:00:00 2001 From: Chia Rui Ong Date: Wed, 30 Oct 2024 13:52:02 +0100 Subject: [PATCH 01/12] add backend to as_field in driver initialization and serialbox_utils.py --- ...tial_temperatures_and_pressure_gradient.py | 5 +- .../dycore/interpolate_to_half_levels_wp.py | 48 +++++ .../edge_2_cell_vector_rbf_interpolation.py | 3 +- .../common/test_utils/serialbox_utils.py | 27 +-- .../common/utils/gt4py_field_allocation.py | 15 ++ .../model/driver/icon4py_configuration.py | 27 ++- .../icon4py/model/driver/icon4py_driver.py | 34 ++- .../model/driver/initialization_utils.py | 38 ++-- .../model/driver/test_cases/gauss3d.py | 155 ++++++-------- .../test_cases/jablonowski_williamson.py | 184 ++++++++--------- .../icon4py/model/driver/test_cases/utils.py | 193 ++++++++++++++++-- .../tests/driver_tests/test_timeloop.py | 20 +- .../initial_condition_tests/test_gauss3d.py | 12 +- .../initial_condition_tests/test_utils.py | 8 +- 14 files changed, 509 insertions(+), 260 deletions(-) create mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_wp.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_virtual_potential_temperatures_and_pressure_gradient.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_virtual_potential_temperatures_and_pressure_gradient.py index eaf32aef0a..662f02b82e 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_virtual_potential_temperatures_and_pressure_gradient.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_virtual_potential_temperatures_and_pressure_gradient.py @@ -13,6 +13,9 @@ from icon4py.model.atmosphere.dycore.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/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..25d0c9ba36 --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_wp.py @@ -0,0 +1,48 @@ +# 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]: + """Formerly known mo_velocity_advection_stencil_10 and as _mo_solve_nonhydro_stencil_05.""" + 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/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/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index e83a460aca..cb20330716 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 @@ -18,7 +18,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 +45,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 +61,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 +80,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() @@ -358,6 +358,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): @@ -547,9 +548,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): @@ -576,21 +577,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) @@ -752,7 +753,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/gt4py_field_allocation.py b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py index 1799fe8cbb..d38f2e514f 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 @@ -37,3 +37,18 @@ def allocate_indices( ): shapex = grid.size[dim] + 1 if is_halfdim else grid.size[dim] return gtx.as_field((dim,), xp.arange(shapex, dtype=dtype), allocator=backend) + + +def allocate_field_from_array( + *dims: gtx.Dimension, + grid, + input_array: xp.ndarray, + 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) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index 66e1073559..f5c4ed6fd1 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -8,8 +8,14 @@ import dataclasses import datetime +import enum import logging +from gt4py.next.program_processors.runners.gtfn import ( + run_gtfn_cached, + run_gtfn_gpu_cached, +) + from icon4py.model.atmosphere.diffusion import diffusion from icon4py.model.atmosphere.dycore.nh_solve import solve_nonhydro as solve_nh from icon4py.model.common.grid import vertical as v_grid @@ -21,7 +27,12 @@ n_substeps_reduced = 2 -@dataclasses.dataclass(frozen=True) +class DriverBackends(enum.Enum): + GTFN_CPU = "gtfn_cpu" + GTFN_GPU = "gtfn_gpu" + + +@dataclasses.dataclass() class Icon4pyRunConfig: dtime: datetime.timedelta = datetime.timedelta(seconds=600.0) # length of a time step start_date: datetime.datetime = datetime.datetime(1, 1, 1, 0, 0, 0) @@ -40,6 +51,16 @@ class Icon4pyRunConfig: restart_mode: bool = False + backend_name: DriverBackends = DriverBackends.GTFN_CPU + + @property + def backend(self): + backend_map = { + DriverBackends.GTFN_CPU.name: run_gtfn_cached, + DriverBackends.GTFN_GPU.name: run_gtfn_gpu_cached, + } + return backend_map[self.backend_name] + @dataclasses.dataclass class Icon4pyConfig: @@ -50,6 +71,7 @@ class Icon4pyConfig: def read_config( + icon4py_driver_backend: DriverBackends, experiment_type: driver_init.ExperimentType = driver_init.ExperimentType.ANY, ) -> Icon4pyConfig: def _mch_ch_r04b09_vertical_config(): @@ -122,6 +144,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 +157,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 +192,7 @@ def _gauss3d_config(): end_date=datetime.datetime(1, 1, 1, 0, 0, 4), apply_initial_stabilization=False, n_substeps=5, + backend=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 23089ca240..3d0ca75b68 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, @@ -279,6 +278,7 @@ def initialize( grid_id: uuid.UUID, grid_root, grid_level, + icon4py_driver_backend: driver_config.DriverBackends, ): """ Inititalize the driver run. @@ -306,7 +306,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 @@ -352,7 +352,7 @@ def initialize( log.info("initializing diffusion") diffusion_params = diffusion.DiffusionParams(config.diffusion_config) exchange = decomposition.create_exchange(props, decomp_info) - diffusion_granule = diffusion.Diffusion(exchange=exchange, backend=gtfn_cpu) + diffusion_granule = diffusion.Diffusion(exchange=exchange, backend=config.run_config.backend) diffusion_granule.init( icon_grid, config.diffusion_config, @@ -366,7 +366,7 @@ def initialize( nonhydro_params = solve_nh.NonHydrostaticParams(config.solve_nonhydro_config) - solve_nonhydro_granule = solve_nh.SolveNonhydro(backend=gtfn_cpu) + solve_nonhydro_granule = solve_nh.SolveNonhydro(backend=config.run_config.backend) solve_nonhydro_granule.init( grid=icon_grid, config=config.solve_nonhydro_config, @@ -392,6 +392,7 @@ def initialize( cell_geometry, edge_geometry, file_path, + backend=config.run_config.backend, rank=props.rank, experiment_type=experiment_type, ) @@ -457,8 +458,28 @@ 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", + default=driver_config.DriverBackends.GTFN_CPU, + 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 @@ -482,7 +503,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, @@ -499,6 +520,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 6964b2b28b..6cb8c8741d 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 as diffus_states from icon4py.model.atmosphere.dycore.state_utils import states as solve_nh_states from icon4py.model.common import dimension as dims, field_type_aliases as fa @@ -91,7 +93,7 @@ def read_icon_grid( 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[ diffus_states.DiffusionDiagnosticState, solve_nh_states.DiagnosticStateNonHydro, @@ -153,14 +155,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( @@ -194,6 +202,7 @@ def read_initial_state( cell_param: geometry.CellParams, edge_param: geometry.EdgeParams, path: pathlib.Path, + backend: gt4py_backend.Backend, rank=0, experiment_type: ExperimentType = ExperimentType.ANY, ) -> tuple[ @@ -213,6 +222,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 +240,7 @@ def read_initial_state( prognostic_state_now, prognostic_state_next, ) = jablonowski_williamson.model_initialization_jabw( - grid, cell_param, edge_param, path, rank + grid, cell_param, edge_param, path, backend, rank ) elif experiment_type == ExperimentType.GAUSS3D: ( @@ -241,7 +251,7 @@ 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, edge_param, path, backend, rank) elif experiment_type == ExperimentType.ANY: ( diffusion_diagnostic_state, @@ -448,7 +458,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 +479,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 +493,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 eb7aebb8ff..e5ed6d2b02 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 as diffus_states from icon4py.model.atmosphere.dycore.state_utils import states as solve_nh_states @@ -35,6 +36,7 @@ def model_initialization_gauss3d( grid: icon_grid.IconGrid, edge_param: geometry.EdgeParams, path: pathlib.Path, + backend: gt4py_backend.Backend, rank=0, ) -> tuple[ diffus_states.DiffusionDiagnosticState, @@ -52,6 +54,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 +64,15 @@ 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() + 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 +93,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 +123,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 = gtx.as_field((dims.CellDim, dims.KDim), eta_v_ndarray) eta_v_e = field_alloc.allocate_zero_field(dims.EdgeDim, dims.KDim, grid=grid) - cell_2_edge_interpolation.cell_2_edge_interpolation( + cell_2_edge_interpolation.cell_2_edge_interpolation.with_backend(backend)( eta_v, cell_2_edge_coeff, eta_v_e, @@ -169,28 +172,40 @@ 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[ + pressure_ifc_ndarray = xp.zeros((num_cells, num_levels + 1), dtype=float) + pressure_ifc_ndarray[ :, -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) + ( + 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( vn, rbf_vec_coeff_c1, @@ -206,7 +221,7 @@ 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( + testcases_utils.compute_perturbed_exner.with_backend(backend)( exner, data_provider.from_metrics_savepoint().exner_ref_mc(), exner_pr, @@ -242,54 +257,14 @@ def model_initialization_gauss3d( exner=exner_next, ) - diffusion_diagnostic_state = diffus_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 = solve_nh_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 = solve_nh_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 469e637728..fb8f5efdda 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 as diffus_states from icon4py.model.atmosphere.dycore.state_utils import states as solve_nh_states @@ -37,6 +38,7 @@ def model_initialization_jabw( cell_param: geometry.CellParams, edge_param: geometry.EdgeParams, path: pathlib.Path, + backend: gt4py_backend.Backend, rank=0, ) -> tuple[ diffus_states.DiffusionDiagnosticState, @@ -56,6 +58,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 +68,18 @@ 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() + 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 +113,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 +135,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,22 +179,24 @@ 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 = 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) cell_2_edge_interpolation.cell_2_edge_interpolation( eta_v, @@ -205,7 +210,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 +219,57 @@ def model_initialization_jabw( edge_lat, edge_lon, primal_normal_x, - eta_v_e.asnumpy(), + eta_v_e.ndarray, ) 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, ) 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,7 +284,7 @@ def model_initialization_jabw( log.info("U, V computation completed.") - exner_pr = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid) + exner_pr = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, backend=backend) testcases_utils.compute_perturbed_exner( exner, data_provider.from_metrics_savepoint().exner_ref_mc(), @@ -284,7 +301,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 +321,15 @@ def model_initialization_jabw( exner=exner_next, ) - diffusion_diagnostic_state = diffus_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 = solve_nh_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 = solve_nh_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..4a584eee2b 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 as diffus_states +from icon4py.model.atmosphere.dycore.state_utils import states as solve_nh_states from icon4py.model.common import ( constants as phy_const, dimension as dims, @@ -14,15 +17,11 @@ 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 +from icon4py.model.common.settings import xp +from icon4py.model.common.utils import gt4py_field_allocation as field_alloc -# TODO: this will have to be removed once domain allows for imports -CellDim = dims.CellDim -KDim = dims.KDim - - -def hydrostatic_adjustment_numpy( +def hydrostatic_adjustment_ndarray( wgtfac_c: xp.ndarray, ddqz_z_half: xp.ndarray, exner_ref_mc: xp.ndarray, @@ -62,7 +61,7 @@ def hydrostatic_adjustment_numpy( return rho, exner, theta_v -def hydrostatic_adjustment_constant_thetav_numpy( +def hydrostatic_adjustment_constant_thetav_ndarray( wgtfac_c: xp.ndarray, ddqz_z_half: xp.ndarray, exner_ref_mc: xp.ndarray, @@ -76,7 +75,7 @@ def hydrostatic_adjustment_constant_thetav_numpy( ) -> tuple[xp.ndarray, xp.ndarray]: """ 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,7 +101,7 @@ 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, @@ -188,7 +187,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 +202,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 +) -> diffus_states.DiffusionDiagnosticState: + return diffus_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 +) -> solve_nh_states.DiagnosticStateNonHydro: + return solve_nh_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 +) -> solve_nh_states.PrepAdvection: + return solve_nh_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: xp.ndarray, + w_ndarray: xp.ndarray, + exner_ndarray: xp.ndarray, + rho_ndarray: xp.ndarray, + theta_v_ndarray: xp.ndarray, + temperature_ndarray: xp.ndarray, + pressure_ndarray: xp.ndarray, + pressure_ifc_ndarray: xp.ndarray, + 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 15d3baef0e..f3a53a5656 100644 --- a/model/driver/tests/driver_tests/test_timeloop.py +++ b/model/driver/tests/driver_tests/test_timeloop.py @@ -325,29 +325,29 @@ def test_run_timeloop_single_step( w_sp = timeloop_diffusion_savepoint_exit.w() assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].vn.asnumpy(), - vn_sp.asnumpy(), + prognostic_state_list[timeloop.prognostic_now].vn.ndarray, + vn_sp.ndarray, atol=6e-12, ) assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].w.asnumpy(), - w_sp.asnumpy(), + prognostic_state_list[timeloop.prognostic_now].w.ndarray, + w_sp.ndarray, atol=8e-14, ) assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].exner.asnumpy(), - exner_sp.asnumpy(), + prognostic_state_list[timeloop.prognostic_now].exner.ndarray, + exner_sp.ndarray, ) assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].theta_v.asnumpy(), - theta_sp.asnumpy(), + prognostic_state_list[timeloop.prognostic_now].theta_v.ndarray, + theta_sp.ndarray, atol=4e-12, ) assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].rho.asnumpy(), - rho_sp.asnumpy(), + prognostic_state_list[timeloop.prognostic_now].rho.ndarray, + rho_sp.ndarray, ) diff --git a/model/driver/tests/initial_condition_tests/test_gauss3d.py b/model/driver/tests/initial_condition_tests/test_gauss3d.py index f3daed8ff5..bbbf5ea30c 100644 --- a/model/driver/tests/initial_condition_tests/test_gauss3d.py +++ b/model/driver/tests/initial_condition_tests/test_gauss3d.py @@ -46,22 +46,22 @@ def test_gauss3d_initial_condition( # only verifying those assigned in the IC rather than all (at least for now) assert helpers.dallclose( - prognostic_state_now.rho.asnumpy(), + prognostic_state_now.rho.ndarray, data_provider.from_savepoint_nonhydro_init(1, "2001-01-01T00:00:04.000", 0) .rho_now() - .asnumpy(), + .ndarray, ) assert helpers.dallclose( - prognostic_state_now.exner.asnumpy(), + prognostic_state_now.exner.ndarray, data_provider.from_savepoint_nonhydro_init(1, "2001-01-01T00:00:04.000", 0) .exner_now() - .asnumpy(), + .ndarray, ) assert helpers.dallclose( - prognostic_state_now.theta_v.asnumpy(), + prognostic_state_now.theta_v.ndarray, data_provider.from_savepoint_nonhydro_init(1, "2001-01-01T00:00:04.000", 0) .theta_v_now() - .asnumpy(), + .ndarray, ) diff --git a/model/driver/tests/initial_condition_tests/test_utils.py b/model/driver/tests/initial_condition_tests/test_utils.py index f675bfe6b2..00d011c90b 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(): # 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, @@ -66,7 +66,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 +86,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, From 9c01cceb2c19efbaa9b72372612d56d6c158f9c0 Mon Sep 17 00:00:00 2001 From: Chia Rui Ong Date: Wed, 30 Oct 2024 15:57:59 +0100 Subject: [PATCH 02/12] diffusion apply_to_temperature fix --- .../model/atmosphere/diffusion/diffusion.py | 125 +++++++++--------- 1 file changed, 63 insertions(+), 62 deletions(-) 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 1fb9456728..54564cf373 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 From f55530f36e83c90535c00f73eaa8a22dc7e10213 Mon Sep 17 00:00:00 2001 From: OngChia Date: Mon, 4 Nov 2024 20:38:57 +0100 Subject: [PATCH 03/12] debug --- .../common/test_utils/serialbox_utils.py | 6 ++-- .../model/driver/icon4py_configuration.py | 24 ++++++++++----- .../icon4py/model/driver/icon4py_driver.py | 15 +++++----- .../model/driver/initialization_utils.py | 26 +++++++++++++--- .../model/driver/test_cases/gauss3d.py | 11 +++---- .../test_cases/jablonowski_williamson.py | 4 +-- .../tests/driver_tests/test_timeloop.py | 30 +++++++++++-------- .../initial_condition_tests/test_gauss3d.py | 14 +++++---- .../test_jablonowski_williamson.py | 2 ++ 9 files changed, 83 insertions(+), 49 deletions(-) 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 5402a18b23..fde489aef4 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 @@ -429,8 +429,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 = xp.column_stack((xp.asarray(range(c2e2c.shape[0])), c2e2c)) + e2c2e0 = xp.column_stack((xp.asarray(range(e2c2e.shape[0])), e2c2e)) grid = ( icon.IconGrid(self._grid_id) .with_config(config) @@ -729,7 +729,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) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index f5c4ed6fd1..698e7ef70f 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -10,6 +10,7 @@ import datetime import enum import logging +from functools import cached_property from gt4py.next.program_processors.runners.gtfn import ( run_gtfn_cached, @@ -27,12 +28,12 @@ n_substeps_reduced = 2 -class DriverBackends(enum.Enum): +class DriverBackends(str, enum.Enum): GTFN_CPU = "gtfn_cpu" GTFN_GPU = "gtfn_gpu" -@dataclasses.dataclass() +@dataclasses.dataclass class Icon4pyRunConfig: dtime: datetime.timedelta = datetime.timedelta(seconds=600.0) # length of a time step start_date: datetime.datetime = datetime.datetime(1, 1, 1, 0, 0, 0) @@ -51,13 +52,20 @@ class Icon4pyRunConfig: restart_mode: bool = False - backend_name: DriverBackends = DriverBackends.GTFN_CPU + 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 + @cached_property def backend(self): backend_map = { - DriverBackends.GTFN_CPU.name: run_gtfn_cached, - DriverBackends.GTFN_GPU.name: run_gtfn_gpu_cached, + DriverBackends.GTFN_CPU.value: run_gtfn_cached, + DriverBackends.GTFN_GPU.value: run_gtfn_gpu_cached, } return backend_map[self.backend_name] @@ -71,7 +79,7 @@ class Icon4pyConfig: def read_config( - icon4py_driver_backend: DriverBackends, + icon4py_driver_backend: str, experiment_type: driver_init.ExperimentType = driver_init.ExperimentType.ANY, ) -> Icon4pyConfig: def _mch_ch_r04b09_vertical_config(): @@ -192,7 +200,7 @@ def _gauss3d_config(): end_date=datetime.datetime(1, 1, 1, 0, 0, 4), apply_initial_stabilization=False, n_substeps=5, - backend=icon4py_driver_backend, + 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 052bb6af62..3e4ee8159c 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_driver.py +++ b/model/driver/src/icon4py/model/driver/icon4py_driver.py @@ -162,13 +162,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() @@ -364,7 +364,7 @@ def initialize( nonhydro_params = solve_nh.NonHydrostaticParams(config.solve_nonhydro_config) solve_nonhydro_granule = solve_nh.SolveNonhydro( - backend=config.run_config.backend, + grid=icon_grid, config=config.solve_nonhydro_config, params=nonhydro_params, metric_state_nonhydro=solve_nonhydro_metric_state, @@ -373,6 +373,7 @@ def initialize( edge_geometry=edge_geometry, cell_geometry=cell_geometry, owner_mask=c_owner_mask, + backend=config.run_config.backend, ) ( @@ -425,13 +426,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. " @@ -461,7 +462,7 @@ def initialize( ) @click.option( "--icon4py_driver_backend", - default=driver_config.DriverBackends.GTFN_CPU, + 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/) ", ) diff --git a/model/driver/src/icon4py/model/driver/initialization_utils.py b/model/driver/src/icon4py/model/driver/initialization_utils.py index 6cb8c8741d..d6f6ff1908 100644 --- a/model/driver/src/icon4py/model/driver/initialization_utils.py +++ b/model/driver/src/icon4py/model/driver/initialization_utils.py @@ -183,7 +183,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 ( @@ -240,7 +242,12 @@ def read_initial_state( prognostic_state_now, prognostic_state_next, ) = jablonowski_williamson.model_initialization_jabw( - grid, cell_param, edge_param, path, backend, rank + grid=grid, + cell_param=cell_param, + edge_param=edge_param, + path=path, + backend=backend, + rank=rank, ) elif experiment_type == ExperimentType.GAUSS3D: ( @@ -251,7 +258,13 @@ def read_initial_state( diagnostic_state, prognostic_state_now, prognostic_state_next, - ) = gauss3d.model_initialization_gauss3d(grid, edge_param, path, backend, 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, @@ -261,7 +274,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) 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 e5ed6d2b02..59777e4297 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py +++ b/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py @@ -158,8 +158,8 @@ def model_initialization_gauss3d( ) log.info("Hydrostatic adjustment computation completed.") - eta_v = gtx.as_field((dims.CellDim, dims.KDim), eta_v_ndarray) - 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.with_backend(backend)( eta_v, cell_2_edge_coeff, @@ -173,9 +173,6 @@ def model_initialization_gauss3d( log.info("Cell-to-edge eta_v computation completed.") pressure_ifc_ndarray = xp.zeros((num_cells, num_levels + 1), dtype=float) - pressure_ifc_ndarray[ - :, -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) ( vn, w, @@ -206,7 +203,7 @@ def model_initialization_gauss3d( backend=backend, ) - 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, @@ -220,7 +217,7 @@ def model_initialization_gauss3d( ) log.info("U, V computation completed.") - exner_pr = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid) + 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(), 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 fb8f5efdda..1aac80a697 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 @@ -197,7 +197,7 @@ def model_initialization_jabw( log.info("Newton iteration completed!") 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) + 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, @@ -285,7 +285,7 @@ def model_initialization_jabw( log.info("U, V computation completed.") exner_pr = field_alloc.allocate_zero_field(dims.CellDim, dims.KDim, grid=grid, backend=backend) - testcases_utils.compute_perturbed_exner( + testcases_utils.compute_perturbed_exner.with_backend(backend)( exner, data_provider.from_metrics_savepoint().exner_ref_mc(), exner_pr, diff --git a/model/driver/tests/driver_tests/test_timeloop.py b/model/driver/tests/driver_tests/test_timeloop.py index 6a867c38aa..aecea9a408 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=None, + 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 = solve_nh_states.DiagnosticStateNonHydro( @@ -325,29 +331,29 @@ def test_run_timeloop_single_step( w_sp = timeloop_diffusion_savepoint_exit.w() assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].vn.ndarray, - vn_sp.ndarray, + prognostic_state_list[timeloop.prognostic_now].vn.asnumpy(), + vn_sp.asnumpy(), atol=6e-12, ) assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].w.ndarray, - w_sp.ndarray, + prognostic_state_list[timeloop.prognostic_now].w.asnumpy(), + w_sp.asnumpy(), atol=8e-14, ) assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].exner.ndarray, - exner_sp.ndarray, + prognostic_state_list[timeloop.prognostic_now].exner.asnumpy(), + exner_sp.asnumpy(), ) assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].theta_v.ndarray, - theta_sp.ndarray, + prognostic_state_list[timeloop.prognostic_now].theta_v.asnumpy(), + theta_sp.asnumpy(), atol=4e-12, ) assert helpers.dallclose( - prognostic_state_list[timeloop.prognostic_now].rho.ndarray, - rho_sp.ndarray, + prognostic_state_list[timeloop.prognostic_now].rho.asnumpy(), + rho_sp.asnumpy(), ) diff --git a/model/driver/tests/initial_condition_tests/test_gauss3d.py b/model/driver/tests/initial_condition_tests/test_gauss3d.py index bbbf5ea30c..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,27 +42,28 @@ def test_gauss3d_initial_condition( icon_grid, edge_geometry, ranked_data_path.joinpath(f"{experiment}/ser_data"), + backend, rank, ) # only verifying those assigned in the IC rather than all (at least for now) assert helpers.dallclose( - prognostic_state_now.rho.ndarray, + prognostic_state_now.rho.asnumpy(), data_provider.from_savepoint_nonhydro_init(1, "2001-01-01T00:00:04.000", 0) .rho_now() - .ndarray, + .asnumpy(), ) assert helpers.dallclose( - prognostic_state_now.exner.ndarray, + prognostic_state_now.exner.asnumpy(), data_provider.from_savepoint_nonhydro_init(1, "2001-01-01T00:00:04.000", 0) .exner_now() - .ndarray, + .asnumpy(), ) assert helpers.dallclose( - prognostic_state_now.theta_v.ndarray, + prognostic_state_now.theta_v.asnumpy(), data_provider.from_savepoint_nonhydro_init(1, "2001-01-01T00:00:04.000", 0) .theta_v_now() - .ndarray, + .asnumpy(), ) 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, ) From fdca0f911261d59f8d1c8d1a309867e31d50ec3d Mon Sep 17 00:00:00 2001 From: OngChia Date: Fri, 8 Nov 2024 13:17:21 +0100 Subject: [PATCH 04/12] add backend in driver functions --- .../src/icon4py/model/common/grid/vertical.py | 86 +++++++++++-------- .../common/utils/gt4py_field_allocation.py | 10 +-- .../model/driver/icon4py_configuration.py | 41 +++++++-- .../icon4py/model/driver/icon4py_driver.py | 1 + .../model/driver/initialization_utils.py | 3 +- .../model/driver/test_cases/gauss3d.py | 4 +- .../test_cases/jablonowski_williamson.py | 6 +- .../icon4py/model/driver/test_cases/utils.py | 75 ++++++++-------- .../initial_condition_tests/test_utils.py | 3 +- 9 files changed, 142 insertions(+), 87 deletions(-) diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 3a11143c83..c2d5d824b2 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -14,9 +14,10 @@ from typing import Final import gt4py.next as gtx +import numpy as np +from gt4py.next import backend as gt4py_backend from icon4py.model.common import dimension as dims, exceptions, field_type_aliases as fa -from icon4py.model.common.settings import xp log = logging.getLogger(__name__) @@ -111,7 +112,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: np.ndarray, vct_b: np.ndarray): object.__setattr__( self, "_vct_a", @@ -122,21 +123,20 @@ def __post_init__(self, vct_a, vct_b): "_vct_b", vct_b, ) - vct_a_array = self._vct_a.ndarray object.__setattr__( self, "_end_index_of_damping_layer", - self._determine_damping_height_index(vct_a_array, self.config.rayleigh_damping_height), + self._determine_damping_height_index(vct_a, self.config.rayleigh_damping_height), ) object.__setattr__( self, "_start_index_for_moist_physics", - self._determine_start_level_of_moist_physics(vct_a_array, self.config.htop_moist_proc), + self._determine_start_level_of_moist_physics(vct_a, self.config.htop_moist_proc), ) object.__setattr__( self, "_end_index_of_flat_layer", - self._determine_end_index_of_flat_layers(vct_a_array, self.config.flat_height), + self._determine_end_index_of_flat_layers(vct_a, self.config.flat_height), ) log.info(f"computation of moist physics start on layer: {self.kstart_moist}") log.info(f"end index of Rayleigh damping layer for w: {self.nrdmax} ") @@ -231,36 +231,36 @@ 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: np.ndarray, top_moist_threshold: float, nshift_total: int = 0 ) -> gtx.int32: 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()) + return gtx.int32(np.min(np.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: np.ndarray, damping_height: float) -> gtx.int32: assert damping_height >= 0.0, "Damping height must be positive." return ( 0 if damping_height > vct_a[0] - else gtx.int32(xp.argmax(xp.where(vct_a >= damping_height)[0]).item()) + else gtx.int32(np.argmax(np.where(vct_a >= damping_height)[0]).item()) ) @classmethod def _determine_end_index_of_flat_layers( - cls, vct_a: xp.ndarray, flat_height: float + cls, vct_a: np.ndarray, flat_height: float ) -> gtx.int32: assert flat_height >= 0.0, "Flat surface height must be positive." return ( 0 if flat_height > vct_a[0] - else gtx.int32(xp.max(xp.where(vct_a >= flat_height)[0]).item()) + else gtx.int32(np.max(np.where(vct_a >= flat_height)[0]).item()) ) def _read_vct_a_and_vct_b_from_file( - file_path: pathlib.Path, num_levels: int -) -> tuple[fa.KField, fa.KField]: + file_path: pathlib.Path, num_levels: int, backend: gt4py_backend.Backend +) -> tuple[np.ndarray, np.ndarray]: """ Read vct_a and vct_b from a file. The file format should be as follows (the same format used for icon): @@ -278,8 +278,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 @@ -298,10 +298,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[np.ndarray, np.ndarray]: """ Compute vct_a and vct_b. @@ -339,16 +343,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 ) @@ -359,8 +364,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 ) @@ -374,10 +379,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): @@ -385,7 +390,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: @@ -416,7 +421,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: @@ -439,7 +444,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 * ( @@ -452,7 +457,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]: @@ -469,20 +474,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[np.ndarray, np.ndarray]: """ 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. @@ -494,11 +503,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/utils/gt4py_field_allocation.py b/model/common/src/icon4py/model/common/utils/gt4py_field_allocation.py index d38f2e514f..6e9d02fd27 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,10 +8,10 @@ from typing import Optional import gt4py.next as gtx +import numpy as np from gt4py.next import backend from icon4py.model.common import type_alias as ta -from icon4py.model.common.settings import xp def allocate_zero_field( @@ -25,7 +25,7 @@ def allocate_zero_field( 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) + return gtx.as_field(dims, np.zeros(shapex, dtype=dtype), allocator=backend) def allocate_indices( @@ -36,13 +36,13 @@ def allocate_indices( backend: Optional[backend.Backend] = None, ): 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) def allocate_field_from_array( *dims: gtx.Dimension, grid, - input_array: xp.ndarray, + input_array: np.ndarray, is_halfdim=False, dtype=ta.wpfloat, backend: Optional[backend.Backend] = None, @@ -51,4 +51,4 @@ def allocate_field_from_array( 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) + return gtx.as_field(dims, np.zeros(shapex, dtype=dtype), allocator=backend) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index 698e7ef70f..da2390c476 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -10,10 +10,14 @@ import datetime import enum import logging -from functools import cached_property +from typing import Final +import numpy as np +from gt4py.next import backend as gt4py_backend from gt4py.next.program_processors.runners.gtfn import ( + run_gtfn, run_gtfn_cached, + run_gtfn_gpu, run_gtfn_gpu_cached, ) @@ -30,10 +34,37 @@ 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.value: run_gtfn, + DriverBackends.GTFN_CPU_CACHED.value: run_gtfn_cached, + DriverBackends.GTFN_GPU.value: run_gtfn_gpu, + DriverBackends.GTFN_GPU_CACHED.value: run_gtfn_gpu_cached, +} + + +def host_or_device_array(backend: gt4py_backend.Backend): + if backend in ( + backend_map[DriverBackends.GTFN_CPU.value], + backend_map[DriverBackends.GTFN_CPU_CACHED.value], + ): + return np + elif backend in ( + backend_map[DriverBackends.GTFN_GPU.value], + backend_map[DriverBackends.GTFN_GPU_CACHED.value], + ): + import cupy as cp # type: ignore[import-untyped] + + return cp + else: + raise ValueError(f"Unknown backend {backend}.") -@dataclasses.dataclass +@dataclasses.dataclass(frozen=True) class Icon4pyRunConfig: dtime: datetime.timedelta = datetime.timedelta(seconds=600.0) # length of a time step start_date: datetime.datetime = datetime.datetime(1, 1, 1, 0, 0, 0) @@ -61,12 +92,8 @@ def __post_init__(self): f"Available backends are {', '.join([f'{k}' for k in [member.value for member in DriverBackends]])}" ) - @cached_property + @property def backend(self): - backend_map = { - DriverBackends.GTFN_CPU.value: run_gtfn_cached, - DriverBackends.GTFN_GPU.value: run_gtfn_gpu_cached, - } return backend_map[self.backend_name] diff --git a/model/driver/src/icon4py/model/driver/icon4py_driver.py b/model/driver/src/icon4py/model/driver/icon4py_driver.py index 3e4ee8159c..0f2beb9228 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_driver.py +++ b/model/driver/src/icon4py/model/driver/icon4py_driver.py @@ -326,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, diff --git a/model/driver/src/icon4py/model/driver/initialization_utils.py b/model/driver/src/icon4py/model/driver/initialization_utils.py index d6f6ff1908..441e984990 100644 --- a/model/driver/src/icon4py/model/driver/initialization_utils.py +++ b/model/driver/src/icon4py/model/driver/initialization_utils.py @@ -297,6 +297,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, @@ -327,7 +328,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, 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 59777e4297..f632e7479f 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py +++ b/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py @@ -19,13 +19,13 @@ 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.driver import icon4py_configuration as driver_config from icon4py.model.driver.test_cases import utils as testcases_utils @@ -64,6 +64,8 @@ def model_initialization_gauss3d( "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) + xp = driver_config.host_or_device_array(backend) + 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 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 1aac80a697..264de2d82b 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 @@ -20,13 +20,13 @@ 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.driver import icon4py_configuration as driver_config from icon4py.model.driver.test_cases import utils as testcases_utils @@ -68,6 +68,8 @@ def model_initialization_jabw( "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) + xp = driver_config.host_or_device_array(backend) + 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 @@ -220,6 +222,7 @@ def model_initialization_jabw( edge_lon, primal_normal_x, eta_v_e.ndarray, + backend, ) log.info("U2vn computation completed.") @@ -234,6 +237,7 @@ def model_initialization_jabw( exner_ndarray, theta_v_ndarray, num_levels, + backend, ) log.info("Hydrostatic adjustment computation completed.") 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 4a584eee2b..7686a84af8 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/utils.py +++ b/model/driver/src/icon4py/model/driver/test_cases/utils.py @@ -6,6 +6,7 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np from gt4py.next import backend as gt4py_backend from icon4py.model.atmosphere.diffusion import diffusion_states as diffus_states @@ -17,22 +18,25 @@ type_alias as ta, ) from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid -from icon4py.model.common.settings import xp from icon4py.model.common.utils import gt4py_field_allocation as field_alloc +from icon4py.model.driver import icon4py_configuration as driver_config def hydrostatic_adjustment_ndarray( - 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, + wgtfac_c: np.ndarray, + ddqz_z_half: np.ndarray, + exner_ref_mc: np.ndarray, + d_exner_dz_ref_ic: np.ndarray, + theta_ref_mc: np.ndarray, + theta_ref_ic: np.ndarray, + rho: np.ndarray, + exner: np.ndarray, + theta_v: np.ndarray, num_levels: int, -): + backend: gt4py_backend.Backend, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + xp = driver_config.host_or_device_array(backend) + # virtual temperature temp_v = theta_v * exner @@ -62,17 +66,17 @@ def hydrostatic_adjustment_ndarray( def hydrostatic_adjustment_constant_thetav_ndarray( - 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, + wgtfac_c: np.ndarray, + ddqz_z_half: np.ndarray, + exner_ref_mc: np.ndarray, + d_exner_dz_ref_ic: np.ndarray, + theta_ref_mc: np.ndarray, + theta_ref_ic: np.ndarray, + rho: np.ndarray, + exner: np.ndarray, + theta_v: np.ndarray, num_levels: int, -) -> tuple[xp.ndarray, xp.ndarray]: +) -> tuple[np.ndarray, np.ndarray]: """ Computes a hydrostatically balanced profile. In constrast to the above hydrostatic_adjustment_ndarray, the virtual temperature is kept (assumed) @@ -107,11 +111,12 @@ def zonalwind_2_normalwind_ndarray( 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: np.ndarray, + edge_lon: np.ndarray, + primal_normal_x: np.ndarray, + eta_v_e: np.ndarray, + backend: gt4py_backend.Backend, +) -> np.ndarray: """ Compute normal wind at edge center from vertical eta coordinate (eta_v_e). @@ -129,6 +134,8 @@ def zonalwind_2_normalwind_ndarray( """ # TODO (Chia Rui) this function needs a test + xp = driver_config.host_or_device_array(backend) + 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)), @@ -307,14 +314,14 @@ def initialize_prep_advection( def create_gt4py_field_for_prognostic_and_diagnostic_variables( - vn_ndarray: xp.ndarray, - w_ndarray: xp.ndarray, - exner_ndarray: xp.ndarray, - rho_ndarray: xp.ndarray, - theta_v_ndarray: xp.ndarray, - temperature_ndarray: xp.ndarray, - pressure_ndarray: xp.ndarray, - pressure_ifc_ndarray: xp.ndarray, + vn_ndarray: np.ndarray, + w_ndarray: np.ndarray, + exner_ndarray: np.ndarray, + rho_ndarray: np.ndarray, + theta_v_ndarray: np.ndarray, + temperature_ndarray: np.ndarray, + pressure_ndarray: np.ndarray, + pressure_ifc_ndarray: np.ndarray, grid: icon_grid.IconGrid, backend: gt4py_backend.Backend, ) -> tuple[ diff --git a/model/driver/tests/initial_condition_tests/test_utils.py b/model/driver/tests/initial_condition_tests/test_utils.py index 00d011c90b..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_ndarray(): +def test_hydrostatic_adjustment_ndarray(backend): # TODO (Jacopo / Chia Rui) these tests could be better num_cells = 10 num_levels = 10 @@ -42,6 +42,7 @@ def test_hydrostatic_adjustment_ndarray(): exner, theta_v, num_levels, + backend, ) assert r_rho.shape == (num_cells, num_levels) From aeeac65e8ee91db5d894a048ac48c828b236a31d Mon Sep 17 00:00:00 2001 From: OngChia Date: Fri, 8 Nov 2024 13:21:19 +0100 Subject: [PATCH 05/12] remove redundant function --- .../model/common/utils/gt4py_field_allocation.py | 15 --------------- 1 file changed, 15 deletions(-) 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 6e9d02fd27..7d6e906383 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 @@ -37,18 +37,3 @@ def allocate_indices( ): shapex = grid.size[dim] + 1 if is_halfdim else grid.size[dim] return gtx.as_field((dim,), np.arange(shapex, dtype=dtype), allocator=backend) - - -def allocate_field_from_array( - *dims: gtx.Dimension, - grid, - input_array: np.ndarray, - 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, np.zeros(shapex, dtype=dtype), allocator=backend) From 27b50620a3f3c44f48023873719f04ddee629c56 Mon Sep 17 00:00:00 2001 From: OngChia Date: Fri, 8 Nov 2024 14:17:29 +0100 Subject: [PATCH 06/12] add back vct array in VerticalGrid --- model/common/src/icon4py/model/common/grid/vertical.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/model/common/src/icon4py/model/common/grid/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index c2d5d824b2..9f92a195f9 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -123,20 +123,21 @@ def __post_init__(self, vct_a: np.ndarray, vct_b: np.ndarray): "_vct_b", vct_b, ) + vct_a_array = self._vct_a.asnumpy() object.__setattr__( self, "_end_index_of_damping_layer", - self._determine_damping_height_index(vct_a, self.config.rayleigh_damping_height), + self._determine_damping_height_index(vct_a_array, self.config.rayleigh_damping_height), ) object.__setattr__( self, "_start_index_for_moist_physics", - self._determine_start_level_of_moist_physics(vct_a, self.config.htop_moist_proc), + self._determine_start_level_of_moist_physics(vct_a_array, self.config.htop_moist_proc), ) object.__setattr__( self, "_end_index_of_flat_layer", - self._determine_end_index_of_flat_layers(vct_a, self.config.flat_height), + self._determine_end_index_of_flat_layers(vct_a_array, self.config.flat_height), ) log.info(f"computation of moist physics start on layer: {self.kstart_moist}") log.info(f"end index of Rayleigh damping layer for w: {self.nrdmax} ") From 77543da34231a668963978711134889992bd8cd7 Mon Sep 17 00:00:00 2001 From: OngChia Date: Fri, 8 Nov 2024 15:08:02 +0100 Subject: [PATCH 07/12] fix backend in gauss test --- .../model/driver/icon4py_configuration.py | 16 ++++++++-------- model/driver/tests/driver_tests/test_timeloop.py | 2 +- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index da2390c476..1be3e5a4e9 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -40,22 +40,22 @@ class DriverBackends(str, enum.Enum): backend_map: Final[dict] = { - DriverBackends.GTFN_CPU.value: run_gtfn, - DriverBackends.GTFN_CPU_CACHED.value: run_gtfn_cached, - DriverBackends.GTFN_GPU.value: run_gtfn_gpu, - DriverBackends.GTFN_GPU_CACHED.value: run_gtfn_gpu_cached, + 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, } def host_or_device_array(backend: gt4py_backend.Backend): if backend in ( - backend_map[DriverBackends.GTFN_CPU.value], - backend_map[DriverBackends.GTFN_CPU_CACHED.value], + backend_map[DriverBackends.GTFN_CPU], + backend_map[DriverBackends.GTFN_CPU_CACHED], ): return np elif backend in ( - backend_map[DriverBackends.GTFN_GPU.value], - backend_map[DriverBackends.GTFN_GPU_CACHED.value], + backend_map[DriverBackends.GTFN_GPU], + backend_map[DriverBackends.GTFN_GPU_CACHED], ): import cupy as cp # type: ignore[import-untyped] diff --git a/model/driver/tests/driver_tests/test_timeloop.py b/model/driver/tests/driver_tests/test_timeloop.py index aecea9a408..13b074cb67 100644 --- a/model/driver/tests/driver_tests/test_timeloop.py +++ b/model/driver/tests/driver_tests/test_timeloop.py @@ -129,7 +129,7 @@ def test_run_timeloop_single_step( if experiment == dt_utils.GAUSS3D_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=None, + icon4py_driver_backend=icon4py_configuration.DriverBackends.GTFN_CPU.value, experiment_type=experiment, ) diffusion_config = config.diffusion_config From ef9a00dbcfc040a46eda3167a2241cb3e8ac522b Mon Sep 17 00:00:00 2001 From: OngChia Date: Fri, 8 Nov 2024 15:57:47 +0100 Subject: [PATCH 08/12] add backend to get_vt_a_and_vct_b in the test --- model/common/tests/grid_tests/test_vertical.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/model/common/tests/grid_tests/test_vertical.py b/model/common/tests/grid_tests/test_vertical.py index 8c22174328..ecd88f2e83 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.asnumpy(), grid_savepoint.vct_a().asnumpy()) assert helpers.dallclose(vct_b.asnumpy(), grid_savepoint.vct_b().asnumpy()) From 2d8328e1fe69b573fd139239e0ae8d4618839087 Mon Sep 17 00:00:00 2001 From: OngChia Date: Fri, 15 Nov 2024 12:37:52 +0100 Subject: [PATCH 09/12] add docstring, xp function in common, fix changes in vertical.py --- .../dycore/interpolate_to_half_levels_vp.py | 12 ++++- .../dycore/interpolate_to_half_levels_wp.py | 12 ++++- .../model/common/field_type_aliases.py | 4 ++ .../src/icon4py/model/common/grid/vertical.py | 28 ++++++----- .../model/common/utils/array_allocation.py | 46 +++++++++++++++++++ .../common/utils/gt4py_field_allocation.py | 27 ++++++----- .../model/driver/icon4py_configuration.py | 21 +-------- .../model/driver/test_cases/gauss3d.py | 9 ++-- .../test_cases/jablonowski_williamson.py | 10 ++-- .../icon4py/model/driver/test_cases/utils.py | 12 +++-- 10 files changed, 126 insertions(+), 55 deletions(-) create mode 100644 model/common/src/icon4py/model/common/utils/array_allocation.py diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_vp.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_vp.py index d3893178d2..dd12053f7f 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/interpolate_to_half_levels_vp.py +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/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/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 index 25d0c9ba36..68de92f94e 100644 --- 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 @@ -20,7 +20,17 @@ def _interpolate_to_half_levels_wp( wgtfac_c: fa.CellKField[wpfloat], interpolant: fa.CellKField[wpfloat], ) -> fa.CellKField[wpfloat]: - """Formerly known mo_velocity_advection_stencil_10 and as _mo_solve_nonhydro_stencil_05.""" + """ + 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]) 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/vertical.py b/model/common/src/icon4py/model/common/grid/vertical.py index 745be96c7c..005df3c457 100644 --- a/model/common/src/icon4py/model/common/grid/vertical.py +++ b/model/common/src/icon4py/model/common/grid/vertical.py @@ -19,6 +19,7 @@ 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.utils import array_allocation as array_alloc log = logging.getLogger(__name__) @@ -121,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: np.ndarray, vct_b: np.ndarray): + def __post_init__(self, vct_a: fa.KField, vct_b: fa.KField): object.__setattr__( self, "_vct_a", @@ -132,7 +133,7 @@ def __post_init__(self, vct_a: np.ndarray, vct_b: np.ndarray): "_vct_b", vct_b, ) - vct_a_array = self._vct_a.asnumpy() + vct_a_array = self._vct_a.ndarray object.__setattr__( self, "_end_index_of_damping_layer", @@ -247,36 +248,41 @@ def size(self, dim: gtx.Dimension) -> int: @classmethod def _determine_start_level_of_moist_physics( - cls, vct_a: np.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(np.min(np.where(interface_height < top_moist_threshold)[0]).item()) + return gtx.int32(xp.min(xp.where(interface_height < top_moist_threshold)[0]).item()) @classmethod - def _determine_damping_height_index(cls, vct_a: np.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] - else gtx.int32(np.argmax(np.where(vct_a >= damping_height)[0]).item()) + else gtx.int32(xp.argmax(xp.where(vct_a >= damping_height)[0]).item()) ) @classmethod def _determine_end_index_of_flat_layers( - cls, vct_a: np.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] - else gtx.int32(np.max(np.where(vct_a >= flat_height)[0]).item()) + else gtx.int32(xp.max(xp.where(vct_a >= flat_height)[0]).item()) ) def _read_vct_a_and_vct_b_from_file( file_path: pathlib.Path, num_levels: int, backend: gt4py_backend.Backend -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[fa.KField, fa.KField]: """ Read vct_a and vct_b from a file. The file format should be as follows (the same format used for icon): @@ -321,7 +327,7 @@ def _read_vct_a_and_vct_b_from_file( def _compute_vct_a_and_vct_b( vertical_config: VerticalGridConfig, backend: gt4py_backend.Backend -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[fa.KField, fa.KField]: """ Compute vct_a and vct_b. @@ -507,7 +513,7 @@ def _compute_vct_a_and_vct_b( def get_vct_a_and_vct_b( vertical_config: VerticalGridConfig, backend: gt4py_backend.Backend -) -> tuple[np.ndarray, np.ndarray]: +) -> 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. 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 7d6e906383..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 @@ -9,23 +9,26 @@ import gt4py.next as gtx import numpy as np -from gt4py.next import backend +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 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, np.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,), np.arange(shapex, dtype=dtype), allocator=backend) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index 1be3e5a4e9..ca61046343 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -12,8 +12,6 @@ import logging from typing import Final -import numpy as np -from gt4py.next import backend as gt4py_backend from gt4py.next.program_processors.runners.gtfn import ( run_gtfn, run_gtfn_cached, @@ -47,23 +45,6 @@ class DriverBackends(str, enum.Enum): } -def host_or_device_array(backend: gt4py_backend.Backend): - if backend in ( - backend_map[DriverBackends.GTFN_CPU], - backend_map[DriverBackends.GTFN_CPU_CACHED], - ): - return np - elif backend in ( - backend_map[DriverBackends.GTFN_GPU], - backend_map[DriverBackends.GTFN_GPU_CACHED], - ): - import cupy as cp # type: ignore[import-untyped] - - return cp - else: - raise ValueError(f"Unknown backend {backend}.") - - @dataclasses.dataclass(frozen=True) class Icon4pyRunConfig: dtime: datetime.timedelta = datetime.timedelta(seconds=600.0) # length of a time step @@ -189,7 +170,7 @@ def _mch_ch_r04b09_config(): def _jablownoski_Williamson_config(): icon_run_config = Icon4pyRunConfig( dtime=datetime.timedelta(seconds=300.0), - end_date=datetime.datetime(1, 1, 1, 0, 30, 0), + end_date=datetime.datetime(1, 1, 1, 0, 5, 0), apply_initial_stabilization=False, n_substeps=5, backend_name=icon4py_driver_backend, 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 2f1cbfa3cc..a738b62c3f 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py +++ b/model/driver/src/icon4py/model/driver/test_cases/gauss3d.py @@ -24,8 +24,10 @@ 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.driver import icon4py_configuration as driver_config +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 @@ -64,7 +66,8 @@ def model_initialization_gauss3d( "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) - xp = driver_config.host_or_device_array(backend) + 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 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 61fbe17b7b..0b2e3c14f3 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 @@ -25,8 +25,10 @@ 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.driver import icon4py_configuration as driver_config +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 @@ -68,7 +70,8 @@ def model_initialization_jabw( "icon_pydycore", str(path.absolute()), False, mpi_rank=rank ) - xp = driver_config.host_or_device_array(backend) + 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 @@ -199,6 +202,7 @@ def model_initialization_jabw( log.info("Newton iteration completed!") eta_v = gtx.as_field((dims.CellDim, dims.KDim), eta_v_ndarray, allocator=backend) + log.debug(f"field device: {eta_v}") 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, 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 7686a84af8..954812c036 100644 --- a/model/driver/src/icon4py/model/driver/test_cases/utils.py +++ b/model/driver/src/icon4py/model/driver/test_cases/utils.py @@ -18,8 +18,10 @@ type_alias as ta, ) from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid -from icon4py.model.common.utils import gt4py_field_allocation as field_alloc -from icon4py.model.driver import icon4py_configuration as driver_config +from icon4py.model.common.utils import ( + array_allocation as array_alloc, + gt4py_field_allocation as field_alloc, +) def hydrostatic_adjustment_ndarray( @@ -35,7 +37,8 @@ def hydrostatic_adjustment_ndarray( num_levels: int, backend: gt4py_backend.Backend, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - xp = driver_config.host_or_device_array(backend) + is_cupy = array_alloc.is_cupy_device(backend) + xp = array_alloc.array_ns(is_cupy) # virtual temperature temp_v = theta_v * exner @@ -134,7 +137,8 @@ def zonalwind_2_normalwind_ndarray( """ # TODO (Chia Rui) this function needs a test - xp = driver_config.host_or_device_array(backend) + 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[ From 32c1cf42543ba5f0bbd7ea3f5ece002b8915bca5 Mon Sep 17 00:00:00 2001 From: OngChia Date: Sun, 17 Nov 2024 20:17:04 +0100 Subject: [PATCH 10/12] add xp allocation in get_offset_provider --- model/driver/src/icon4py/model/driver/icon4py_configuration.py | 2 +- model/driver/src/icon4py/model/driver/icon4py_driver.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index ca61046343..a9a36bf8a9 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -170,7 +170,7 @@ def _mch_ch_r04b09_config(): def _jablownoski_Williamson_config(): icon_run_config = Icon4pyRunConfig( dtime=datetime.timedelta(seconds=300.0), - end_date=datetime.datetime(1, 1, 1, 0, 5, 0), + end_date=datetime.datetime(1, 1, 1, 0, 30, 0), apply_initial_stabilization=False, n_substeps=5, backend_name=icon4py_driver_backend, diff --git a/model/driver/src/icon4py/model/driver/icon4py_driver.py b/model/driver/src/icon4py/model/driver/icon4py_driver.py index 0f2beb9228..3ab2c09c19 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_driver.py +++ b/model/driver/src/icon4py/model/driver/icon4py_driver.py @@ -462,7 +462,7 @@ def initialize( help="Enable all debugging messages. Otherwise, only critical error messages are printed.", ) @click.option( - "--icon4py_driver_backend", + "--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/) ", From c49cabc61a75b638320746249f8958b2b4a13096 Mon Sep 17 00:00:00 2001 From: OngChia Date: Sun, 17 Nov 2024 20:18:50 +0100 Subject: [PATCH 11/12] add xp allocation in get_offset_provider --- model/common/src/icon4py/model/common/grid/base.py | 4 +++- model/common/src/icon4py/model/common/grid/utils.py | 4 +++- .../icon4py/model/common/test_utils/serialbox_utils.py | 10 +++++----- .../src/icon4py/model/driver/icon4py_configuration.py | 1 + .../driver/src/icon4py/model/driver/icon4py_driver.py | 4 +++- .../src/icon4py/model/driver/initialization_utils.py | 4 +++- .../model/driver/test_cases/jablonowski_williamson.py | 1 - 7 files changed, 18 insertions(+), 10 deletions(-) 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/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index 3715e6cce7..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 @@ -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() @@ -454,8 +454,8 @@ def construct_icon_grid(self, on_gpu: bool) -> icon.IconGrid: ) c2e2c = self.c2e2c() e2c2e = self.e2c2e() - c2e2c0 = xp.column_stack((xp.asarray(range(c2e2c.shape[0])), c2e2c)) - e2c2e0 = xp.column_stack((xp.asarray(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) diff --git a/model/driver/src/icon4py/model/driver/icon4py_configuration.py b/model/driver/src/icon4py/model/driver/icon4py_configuration.py index a9a36bf8a9..bf015dbbe2 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_configuration.py +++ b/model/driver/src/icon4py/model/driver/icon4py_configuration.py @@ -43,6 +43,7 @@ class DriverBackends(str, enum.Enum): 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) diff --git a/model/driver/src/icon4py/model/driver/icon4py_driver.py b/model/driver/src/icon4py/model/driver/icon4py_driver.py index 3ab2c09c19..869d65a2b7 100644 --- a/model/driver/src/icon4py/model/driver/icon4py_driver.py +++ b/model/driver/src/icon4py/model/driver/icon4py_driver.py @@ -311,6 +311,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, @@ -462,7 +463,8 @@ def initialize( help="Enable all debugging messages. Otherwise, only critical error messages are printed.", ) @click.option( - "--icon4py_driver_backend", "-b", + "--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/) ", diff --git a/model/driver/src/icon4py/model/driver/initialization_utils.py b/model/driver/src/icon4py/model/driver/initialization_utils.py index 1ec36e74b0..8effc19a95 100644 --- a/model/driver/src/icon4py/model/driver/initialization_utils.py +++ b/model/driver/src/icon4py/model/driver/initialization_utils.py @@ -32,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 @@ -64,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, @@ -86,7 +88,7 @@ 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) 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 0b2e3c14f3..1fbe9d2b2c 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 @@ -202,7 +202,6 @@ def model_initialization_jabw( log.info("Newton iteration completed!") eta_v = gtx.as_field((dims.CellDim, dims.KDim), eta_v_ndarray, allocator=backend) - log.debug(f"field device: {eta_v}") 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, From 9a21467870d9e629d8b7e72384821a41bcca746d Mon Sep 17 00:00:00 2001 From: OngChia Date: Mon, 18 Nov 2024 09:06:59 +0100 Subject: [PATCH 12/12] remove xp.ndarray in driver test_utils --- .../icon4py/model/driver/test_cases/utils.py | 67 +++++++++---------- 1 file changed, 33 insertions(+), 34 deletions(-) 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 954812c036..d04213cde1 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,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx -import numpy as np from gt4py.next import backend as gt4py_backend from icon4py.model.atmosphere.diffusion import diffusion_states as diffus_states @@ -25,18 +24,18 @@ def hydrostatic_adjustment_ndarray( - wgtfac_c: np.ndarray, - ddqz_z_half: np.ndarray, - exner_ref_mc: np.ndarray, - d_exner_dz_ref_ic: np.ndarray, - theta_ref_mc: np.ndarray, - theta_ref_ic: np.ndarray, - rho: np.ndarray, - exner: np.ndarray, - theta_v: np.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[np.ndarray, np.ndarray, np.ndarray]: +) -> tuple[fa.AnyNDArray, fa.AnyNDArray, fa.AnyNDArray]: is_cupy = array_alloc.is_cupy_device(backend) xp = array_alloc.array_ns(is_cupy) @@ -69,17 +68,17 @@ def hydrostatic_adjustment_ndarray( def hydrostatic_adjustment_constant_thetav_ndarray( - wgtfac_c: np.ndarray, - ddqz_z_half: np.ndarray, - exner_ref_mc: np.ndarray, - d_exner_dz_ref_ic: np.ndarray, - theta_ref_mc: np.ndarray, - theta_ref_ic: np.ndarray, - rho: np.ndarray, - exner: np.ndarray, - theta_v: np.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[np.ndarray, np.ndarray]: +) -> tuple[fa.AnyNDArray, fa.AnyNDArray]: """ Computes a hydrostatically balanced profile. In constrast to the above hydrostatic_adjustment_ndarray, the virtual temperature is kept (assumed) @@ -114,12 +113,12 @@ def zonalwind_2_normalwind_ndarray( jw_up: float, lat_perturbation_center: float, lon_perturbation_center: float, - edge_lat: np.ndarray, - edge_lon: np.ndarray, - primal_normal_x: np.ndarray, - eta_v_e: np.ndarray, + edge_lat: fa.AnyNDArray, + edge_lon: fa.AnyNDArray, + primal_normal_x: fa.AnyNDArray, + eta_v_e: fa.AnyNDArray, backend: gt4py_backend.Backend, -) -> np.ndarray: +) -> fa.AnyNDArray: """ Compute normal wind at edge center from vertical eta coordinate (eta_v_e). @@ -318,14 +317,14 @@ def initialize_prep_advection( def create_gt4py_field_for_prognostic_and_diagnostic_variables( - vn_ndarray: np.ndarray, - w_ndarray: np.ndarray, - exner_ndarray: np.ndarray, - rho_ndarray: np.ndarray, - theta_v_ndarray: np.ndarray, - temperature_ndarray: np.ndarray, - pressure_ndarray: np.ndarray, - pressure_ifc_ndarray: np.ndarray, + 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[