diff --git a/ci/distributed.yml b/ci/distributed.yml index 916afe1734..c0a38d8558 100644 --- a/ci/distributed.yml +++ b/ci/distributed.yml @@ -79,7 +79,7 @@ build_distributed_cpu: - ci/scripts/ci-mpi-wrapper.sh pytest -sv -k mpi_tests --with-mpi --backend=$BACKEND model/$COMPONENT parallel: matrix: - - COMPONENT: [atmosphere/diffusion, atmosphere/dycore, common] + - COMPONENT: [atmosphere/diffusion, atmosphere/dycore, atmosphere/advection, common] BACKEND: [embedded, gtfn_cpu, dace_cpu] rules: - if: $COMPONENT == 'atmosphere/diffusion' @@ -91,6 +91,9 @@ build_distributed_cpu: - if: $COMPONENT == 'atmosphere/dycore' variables: SLURM_TIMELIMIT: '00:15:00' + - if: $COMPONENT == 'atmosphere/advection' + variables: + SLURM_TIMELIMIT: '00:15:00' - when: on_success variables: SLURM_TIMELIMIT: '00:30:00' diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index 23318563de..2ee13c0da5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -11,6 +11,7 @@ from abc import ABC, abstractmethod from enum import Enum +import gt4py.next as gtx import gt4py.next.typing as gtx_typing import icon4py.model.common.grid.states as grid_states @@ -29,6 +30,7 @@ from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid +from icon4py.model.common.model_options import setup_program from icon4py.model.common.utils import data_allocation as data_alloc @@ -160,7 +162,7 @@ def __init__( self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None, - exchange: decomposition.ExchangeRuntime | None = None, + exchange: decomposition.ExchangeRuntime | None = decomposition.single_node_default, ): log.debug("advection class init - start") @@ -175,7 +177,19 @@ def __init__( self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) # stencils - self._copy_cell_kdim_field = copy_cell_kdim_field.with_backend(self._backend) + self._copy_cell_kdim_field = setup_program( + backend=self._backend, + program=copy_cell_kdim_field, + horizontal_sizes={ + "horizontal_start": self._start_cell_nudging, + "horizontal_end": self._end_cell_local, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": self._grid.num_levels, + }, + offset_provider=self._grid.connectivities, + ) def run( self, @@ -187,19 +201,10 @@ def run( ): log.debug("advection run - start") - log.debug("communication of prep_adv cell field: mass_flx_ic - start") - self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) - log.debug("communication of prep_adv cell field: mass_flx_ic - end") - log.debug("running stencil copy_cell_kdim_field - start") self._copy_cell_kdim_field( field_in=p_tracer_now, field_out=p_tracer_new, - horizontal_start=self._start_cell_nudging, - horizontal_end=self._end_cell_local, - vertical_start=0, - vertical_end=self._grid.num_levels, - offset_provider=self._grid.connectivities, ) log.debug("running stencil copy_cell_kdim_field - end") @@ -216,7 +221,7 @@ def __init__( grid: icon_grid.IconGrid, metric_state: advection_states.AdvectionMetricState, backend: gtx_typing.Backend | None, - exchange: decomposition.ExchangeRuntime | None = None, + exchange: decomposition.ExchangeRuntime | None = decomposition.single_node_default, even_timestep: bool = False, ): log.debug("advection class init - start") @@ -237,9 +242,33 @@ def __init__( ) self._determine_local_domains() # stencils - self._apply_density_increment = apply_density_increment.with_backend(self._backend) - self._apply_interpolated_tracer_time_tendency = ( - apply_interpolated_tracer_time_tendency.with_backend(self._backend) + self._apply_density_increment = setup_program( + backend=self._backend, + program=apply_density_increment, + constant_args={ + "deepatmo_divzl": self._metric_state.deepatmo_divzl, + "deepatmo_divzu": self._metric_state.deepatmo_divzu, + }, + horizontal_sizes={ + "horizontal_end": self._end_cell_end, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(self._grid.num_levels), + }, + offset_provider=self._grid.connectivities, + ) + self._apply_interpolated_tracer_time_tendency = setup_program( + backend=self._backend, + program=apply_interpolated_tracer_time_tendency, + horizontal_sizes={ + "horizontal_start": self._start_cell_lateral_boundary, + "horizontal_end": self._end_cell_lateral_boundary_level_4, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(self._grid.num_levels), + }, ) log.debug("advection class init - end") @@ -272,7 +301,6 @@ def run( log.debug("advection run - start") log.debug("communication of prep_adv cell field: mass_flx_ic - start") - self._exchange.exchange_and_wait(dims.CellDim, prep_adv.mass_flx_ic) log.debug("communication of prep_adv cell field: mass_flx_ic - end") # reintegrate density for conservation of mass @@ -281,20 +309,15 @@ def run( if self._even_timestep else (diagnostic_state.airmass_new, self._start_cell_lateral_boundary_level_3) ) + log.debug("running stencil apply_density_increment - start") self._apply_density_increment( rhodz_in=rhodz_in, p_mflx_contra_v=prep_adv.mass_flx_ic, - deepatmo_divzl=self._metric_state.deepatmo_divzl, - deepatmo_divzu=self._metric_state.deepatmo_divzu, rhodz_out=self._rhodz_ast2, p_dtime=dtime, even_timestep=self._even_timestep, horizontal_start=horizontal_start, - horizontal_end=self._end_cell_end, - vertical_start=0, - vertical_end=self._grid.num_levels, - offset_provider=self._grid.connectivities, ) log.debug("running stencil apply_density_increment - end") @@ -355,19 +378,9 @@ def run( p_grf_tend_tracer=diagnostic_state.grf_tend_tracer, p_tracer_new=p_tracer_new, p_dtime=dtime, - horizontal_start=self._start_cell_lateral_boundary, - horizontal_end=self._end_cell_lateral_boundary_level_4, - vertical_start=0, - vertical_end=self._grid.num_levels, - offset_provider=self._grid.connectivities, ) log.debug("running stencil apply_interpolated_tracer_time_tendency - end") - # exchange updated tracer values, originally happens only if iforcing /= inwp - log.debug("communication of advection cell field: p_tracer_new - start") - self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) - log.debug("communication of advection cell field: p_tracer_new - end") - # finalize step self._even_timestep = not self._even_timestep @@ -402,13 +415,16 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- match config.horizontal_advection_type: case HorizontalAdvectionType.NO_ADVECTION: - horizontal_advection = advection_horizontal.NoAdvection(grid=grid, backend=backend) + horizontal_advection = advection_horizontal.NoAdvection( + grid=grid, backend=backend, exchange=exchange + ) case HorizontalAdvectionType.LINEAR_2ND_ORDER: tracer_flux = advection_horizontal.SecondOrderMiura( grid=grid, least_squares_state=least_squares_state, horizontal_limiter=horizontal_limiter, backend=backend, + exchange=exchange, ) horizontal_advection = advection_horizontal.SemiLagrangian( tracer_flux=tracer_flux, @@ -433,7 +449,9 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- match config.vertical_advection_type: case VerticalAdvectionType.NO_ADVECTION: - vertical_advection = advection_vertical.NoAdvection(grid=grid, backend=backend) + vertical_advection = advection_vertical.NoAdvection( + grid=grid, backend=backend, exchange=exchange + ) case VerticalAdvectionType.UPWIND_1ST_ORDER: boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) vertical_advection = advection_vertical.FirstOrderUpwind( @@ -441,6 +459,7 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- grid=grid, metric_state=metric_state, backend=backend, + exchange=exchange, ) case VerticalAdvectionType.PPM_3RD_ORDER: boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) @@ -450,6 +469,7 @@ def convert_config_to_horizontal_vertical_advection( # noqa: PLR0912 [too-many- grid=grid, metric_state=metric_state, backend=backend, + exchange=exchange, ) case _: raise NotImplementedError("Unknown vertical advection type.") @@ -466,7 +486,7 @@ def convert_config_to_advection( edge_params: grid_states.EdgeParams, cell_params: grid_states.CellParams, backend: gtx_typing.Backend | None, - exchange: decomposition.ExchangeRuntime | None = None, + exchange: decomposition.ExchangeRuntime = decomposition.single_node_default, even_timestep: bool = False, ) -> Advection: exchange = exchange or decomposition.SingleNodeExchange() diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py index 200bed37db..8be6071ce3 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_horizontal.py @@ -196,11 +196,13 @@ def __init__( least_squares_state: advection_states.AdvectionLeastSquaresState, backend: model_backends.BackendLike, horizontal_limiter: HorizontalFluxLimiter | None = None, + exchange: decomposition.ExchangeRuntime | None = None, ): self._grid = grid self._least_squares_state = least_squares_state self._backend = backend self._horizontal_limiter = horizontal_limiter or HorizontalFluxLimiter() + self._exchange = exchange or decomposition.SingleNodeExchange() # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -340,7 +342,7 @@ def run( class NoAdvection(HorizontalAdvection): """Class that implements disabled horizontal advection.""" - def __init__(self, grid: icon_grid.IconGrid, backend: model_backends.BackendLike): + def __init__(self, grid: icon_grid.IconGrid, backend: model_backends.BackendLike, exchange): log.debug("horizontal advection class init - start") # input arguments @@ -350,6 +352,7 @@ def __init__(self, grid: icon_grid.IconGrid, backend: model_backends.BackendLike cell_domain = h_grid.domain(dims.CellDim) self._start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) self._end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + self._exchange = exchange # stencils self._copy_cell_kdim_field = setup_program( @@ -386,7 +389,7 @@ def run( field_out=p_tracer_new, ) log.debug("running stencil copy_cell_kdim_field - end") - + self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("horizontal advection run - end") @@ -421,7 +424,6 @@ def run( p_mflx_tracer_h=p_mflx_tracer_h, dtime=dtime, ) - log.debug("horizontal advection run - end") @abstractmethod diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py deleted file mode 100644 index 01c0cc80bf..0000000000 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_lsq_coeffs.py +++ /dev/null @@ -1,142 +0,0 @@ -# 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 numpy as np -import scipy - -from icon4py.model.common.grid import base as base_grid -from icon4py.model.common.math.projection import ( - gnomonic_proj_single_val, - plane_torus_closest_coordinates, -) -from icon4py.model.common.utils import data_allocation as data_alloc - - -def compute_lsq_pseudoinv( - cell_owner_mask: data_alloc.NDArray, - lsq_pseudoinv: data_alloc.NDArray, - z_lsq_mat_c: data_alloc.NDArray, - lsq_weights_c: data_alloc.NDArray, - start_idx: int, - min_rlcell_int: int, - lsq_dim_unk: int, - lsq_dim_c: int, -) -> data_alloc.NDArray: - for jjb in range(lsq_dim_c): - for jjk in range(lsq_dim_unk): - for jc in range(start_idx, min_rlcell_int): - u, s, v_t, _ = scipy.linalg.lapack.dgesdd(z_lsq_mat_c[jc, :, :]) - if cell_owner_mask[jc]: - lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = ( - lsq_pseudoinv[jc, :lsq_dim_unk, jjb] - + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb] - ) - return lsq_pseudoinv - - -def compute_lsq_weights_c( - z_dist_g: data_alloc.NDArray, - lsq_weights_c_jc: data_alloc.NDArray, - lsq_dim_stencil: int, - lsq_wgt_exp: int, -) -> data_alloc.NDArray: - for js in range(lsq_dim_stencil): - z_norm = np.sqrt(np.dot(z_dist_g[js, :], z_dist_g[js, :])) - lsq_weights_c_jc[js] = 1.0 / (z_norm**lsq_wgt_exp) - return lsq_weights_c_jc / np.max(lsq_weights_c_jc) - - -def compute_z_lsq_mat_c( - cell_owner_mask: data_alloc.NDArray, - z_lsq_mat_c: data_alloc.NDArray, - lsq_weights_c: data_alloc.NDArray, - z_dist_g: data_alloc.NDArray, - jc: int, - lsq_dim_unk: int, - lsq_dim_c: int, -) -> data_alloc.NDArray: - min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) - if cell_owner_mask[jc]: - z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 - - for js in range(lsq_dim_c): - z_lsq_mat_c[jc, js, :lsq_dim_unk] = lsq_weights_c[jc, js] * z_dist_g[js, :] - - return z_lsq_mat_c[jc, js, :lsq_dim_unk] - - -def lsq_compute_coeffs( - cell_center_x: data_alloc.NDArray, - cell_center_y: data_alloc.NDArray, - cell_lat: data_alloc.NDArray, - cell_lon: data_alloc.NDArray, - c2e2c: data_alloc.NDArray, - cell_owner_mask: data_alloc.NDArray, - domain_length: float, - domain_height: float, - grid_sphere_radius: float, - lsq_dim_unk: int, - lsq_dim_c: int, - lsq_wgt_exp: int, - lsq_dim_stencil: int, - start_idx: int, - min_rlcell_int: int, - geometry_type: base_grid.GeometryType, -) -> data_alloc.NDArray: - lsq_weights_c = np.zeros((min_rlcell_int, lsq_dim_stencil)) - lsq_pseudoinv = np.zeros((min_rlcell_int, lsq_dim_unk, lsq_dim_c)) - z_lsq_mat_c = np.zeros((min_rlcell_int, lsq_dim_c, lsq_dim_c)) - - for jc in range(start_idx, min_rlcell_int): - match geometry_type: - case base_grid.GeometryType.ICOSAHEDRON: - z_dist_g = np.zeros((lsq_dim_c, 2)) - for js in range(lsq_dim_stencil): - z_dist_g[js, :] = gnomonic_proj_single_val( - cell_lon[jc], cell_lat[jc], cell_lon[c2e2c[jc, js]], cell_lat[c2e2c[jc, js]] - ) - z_dist_g *= grid_sphere_radius - - min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) - if cell_owner_mask[jc]: - z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 - case base_grid.GeometryType.TORUS: - ilc_s = c2e2c[jc, :lsq_dim_stencil] - cc_cell = np.zeros((lsq_dim_stencil, 2)) - - cc_cv = (cell_center_x[jc], cell_center_y[jc]) - for js in range(lsq_dim_stencil): - cc_cell[js, :] = plane_torus_closest_coordinates( - cell_center_x[jc], - cell_center_y[jc], - cell_center_x[ilc_s][js], - cell_center_y[ilc_s][js], - domain_length, - domain_height, - ) - z_dist_g = cc_cell - cc_cv - - lsq_weights_c[jc, :] = compute_lsq_weights_c( - z_dist_g, lsq_weights_c[jc, :], lsq_dim_stencil, lsq_wgt_exp - ) - z_lsq_mat_c[jc, js, :lsq_dim_unk] = compute_z_lsq_mat_c( - cell_owner_mask, z_lsq_mat_c, lsq_weights_c, z_dist_g, jc, lsq_dim_unk, lsq_dim_c - ) - - lsq_pseudoinv = compute_lsq_pseudoinv( - cell_owner_mask, - lsq_pseudoinv, - z_lsq_mat_c, - lsq_weights_c, - start_idx, - min_rlcell_int, - lsq_dim_unk, - lsq_dim_c, - ) - - return lsq_pseudoinv diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index af595436fa..c4b70028d6 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -63,6 +63,7 @@ field_type_aliases as fa, type_alias as ta, ) +from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid from icon4py.model.common.model_options import setup_program from icon4py.model.common.utils import data_allocation as data_alloc @@ -400,12 +401,13 @@ def run( class NoAdvection(VerticalAdvection): """Class that implements disabled vertical advection.""" - def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None): + def __init__(self, grid: icon_grid.IconGrid, backend: gtx_typing.Backend | None, exchange): log.debug("vertical advection class init - start") # input arguments self._grid = grid self._backend = backend + self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -464,7 +466,6 @@ def run( horizontal_end=horizontal_end, ) log.debug("running stencil copy_cell_kdim_field - end") - log.debug("vertical advection run - end") @@ -538,6 +539,7 @@ def __init__( grid: icon_grid.IconGrid, metric_state: advection_states.AdvectionMetricState, backend: gtx_typing.Backend | None, + exchange: decomposition.ExchangeRuntime | None = None, ): log.debug("vertical advection class init - start") @@ -546,6 +548,7 @@ def __init__( self._grid = grid self._metric_state = metric_state self._backend = backend + self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -625,7 +628,6 @@ def _compute_numerical_flux( horizontal_start, horizontal_end = self._get_horizontal_start_end( even_timestep=even_timestep ) - log.debug("running stencil compute_vertical_tracer_flux_upwind - start") self._compute_vertical_tracer_flux_upwind( p_cc=p_tracer_now, @@ -674,7 +676,7 @@ def _update_unknowns( horizontal_end=horizontal_end, ) log.debug("running stencil integrate_tracer_vertically - end") - + self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("vertical unknowns update - end") @@ -688,6 +690,7 @@ def __init__( grid: icon_grid.IconGrid, metric_state: advection_states.AdvectionMetricState, backend: gtx_typing.Backend | None, + exchange: decomposition.ExchangeRuntime | None = None, ): log.debug("vertical advection class init - start") @@ -697,6 +700,7 @@ def __init__( self._grid = grid self._metric_state = metric_state self._backend = backend + self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -1052,7 +1056,6 @@ def _compute_numerical_flux( log.debug("running stencil compute_ppm4gpu_integer_flux - end") ## set boundary conditions - self._boundary_conditions.run( p_mflx_tracer_v=p_mflx_tracer_v, horizontal_start=horizontal_start, @@ -1060,7 +1063,6 @@ def _compute_numerical_flux( ) ## apply flux limiter - self._vertical_limiter.limit_fluxes( horizontal_start=horizontal_start, horizontal_end=horizontal_end ) @@ -1095,6 +1097,6 @@ def _update_unknowns( horizontal_start=horizontal_start, horizontal_end=horizontal_end, ) + self._exchange.exchange_and_wait(dims.CellDim, p_tracer_new) log.debug("running stencil integrate_tracer_vertically - end") - log.debug("vertical unknowns update - end") diff --git a/model/atmosphere/advection/tests/advection/fixtures.py b/model/atmosphere/advection/tests/advection/fixtures.py index aabae893bb..f039d29ed4 100644 --- a/model/atmosphere/advection/tests/advection/fixtures.py +++ b/model/atmosphere/advection/tests/advection/fixtures.py @@ -10,7 +10,9 @@ import pytest -from icon4py.model.testing.fixtures.datatest import data_provider +from icon4py.model.atmosphere.advection import advection_states +from icon4py.model.testing import serialbox +from icon4py.model.testing.fixtures.datatest import data_provider, decomposition_info @pytest.fixture @@ -44,3 +46,13 @@ def advection_exit_savepoint(data_provider, date): fixture, passing 'date='. """ return data_provider.from_advection_exit_savepoint(size=data_provider.grid_size, date=date) + + +@pytest.fixture +def advection_lsq_state( + interpolation_savepoint: serialbox.InterpolationSavepoint, +) -> advection_states.AdvectionLeastSquaresState: + return advection_states.AdvectionLeastSquaresState( + lsq_pseudoinv_1=interpolation_savepoint.lsq_pseudoinv_1(), + lsq_pseudoinv_2=interpolation_savepoint.lsq_pseudoinv_2(), + ) diff --git a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py index ee43d3ae3d..16931b41ad 100644 --- a/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection/integration_tests/test_advection.py @@ -10,14 +10,14 @@ import pytest import icon4py.model.testing.test_utils as test_helpers -from icon4py.model.atmosphere.advection import advection, advection_states -from icon4py.model.atmosphere.advection.advection_lsq_coeffs import lsq_compute_coeffs +from icon4py.model.atmosphere.advection import advection from icon4py.model.common import constants, dimension as dims from icon4py.model.common.grid import ( base as base_grid, geometry_attributes as geometry_attrs, horizontal as h_grid, ) +from icon4py.model.common.interpolation.interpolation_fields import compute_lsq_coeffs from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing import ( definitions, @@ -142,7 +142,7 @@ def test_advection_run_single_step( ) interpolation_state = construct_interpolation_state(interpolation_savepoint, backend=backend) geometry = gridtest_utils.get_grid_geometry(backend, experiment) - least_squares_coeffs = lsq_compute_coeffs( + least_squares_coeffs = compute_lsq_coeffs( cell_center_x=geometry.get(geometry_attrs.CELL_CENTER_X).asnumpy(), cell_center_y=geometry.get(geometry_attrs.CELL_CENTER_Y).asnumpy(), cell_lat=geometry.get(geometry_attrs.CELL_LAT).asnumpy(), @@ -216,7 +216,7 @@ def test_advection_run_single_step( @pytest.mark.level("unit") @pytest.mark.datatest -def test_lsq_compute_coeffs( +def test_compute_lsq_coeffs( icon_grid: base_grid.Grid, grid_savepoint: sb.IconGridSavepoint, backend: gtx_typing.Backend, @@ -251,7 +251,7 @@ def test_lsq_compute_coeffs( coordinates = gm.coordinates cell_lat = coordinates[dims.CellDim]["lat"].asnumpy() cell_lon = coordinates[dims.CellDim]["lon"].asnumpy() - lsq_pseudoinv = lsq_compute_coeffs( + lsq_pseudoinv = compute_lsq_coeffs( cell_center_x, cell_center_y, cell_lat, diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/__init__.py b/model/atmosphere/advection/tests/advection/mpi_tests/__init__.py new file mode 100644 index 0000000000..de9850de36 --- /dev/null +++ b/model/atmosphere/advection/tests/advection/mpi_tests/__init__.py @@ -0,0 +1,7 @@ +# 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 diff --git a/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py new file mode 100644 index 0000000000..3b36a57891 --- /dev/null +++ b/model/atmosphere/advection/tests/advection/mpi_tests/test_parallel_advection.py @@ -0,0 +1,275 @@ +# 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 pytest +from gt4py.next import typing as gtx_typing + +import icon4py.model.testing.test_utils as test_helpers +from icon4py.model.atmosphere.advection import advection +from icon4py.model.common import constants, dimension as dims +from icon4py.model.common.decomposition import definitions, mpi_decomposition +from icon4py.model.common.grid import ( + base as base_grid, + geometry_attributes as geometry_attrs, + horizontal as h_grid, +) +from icon4py.model.common.interpolation.interpolation_fields import compute_lsq_coeffs +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing import ( + definitions as test_defs, + grid_utils, + parallel_helpers, + serialbox as sb, + test_utils, +) +from icon4py.model.testing.fixtures.datatest import ( + backend, + backend_like, + data_provider, + download_ser_data, + experiment, + grid_savepoint, + icon_grid, + interpolation_savepoint, + metrics_savepoint, + processor_props, +) + +from ..fixtures import * # noqa: F403 +from ..utils import ( + construct_config, + construct_diagnostic_exit_state, + construct_diagnostic_init_state, + construct_interpolation_state, + construct_least_squares_state, + construct_metric_state, + construct_prep_adv, + log_serialized, + verify_advection_fields, +) + + +try: + import mpi4py + + mpi_decomposition.init_mpi() +except ImportError: + pytest.skip("Skipping parallel on single node installation", allow_module_level=True) + + +@pytest.mark.parametrize("processor_props", [True], indirect=True) +@pytest.mark.datatest +@pytest.mark.parametrize("experiment", [test_defs.Experiments.MCH_CH_R04B09]) +@pytest.mark.parametrize( + "date, even_timestep, ntracer, horizontal_advection_type, horizontal_advection_limiter, vertical_advection_type, vertical_advection_limiter", + [ + ( + "2021-06-20T12:00:10.000", + False, + 1, + advection.HorizontalAdvectionType.LINEAR_2ND_ORDER, + advection.HorizontalAdvectionLimiter.POSITIVE_DEFINITE, + advection.VerticalAdvectionType.NO_ADVECTION, + advection.VerticalAdvectionLimiter.NO_LIMITER, + ), + ( + "2021-06-20T12:00:20.000", + True, + 1, + advection.HorizontalAdvectionType.LINEAR_2ND_ORDER, + advection.HorizontalAdvectionLimiter.POSITIVE_DEFINITE, + advection.VerticalAdvectionType.NO_ADVECTION, + advection.VerticalAdvectionLimiter.NO_LIMITER, + ), + ( + "2021-06-20T12:00:10.000", + False, + 4, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), + ( + "2021-06-20T12:00:20.000", + True, + 4, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), + ], +) +@pytest.mark.mpi +def test_advection_run_single_step( + date, + even_timestep, + ntracer, + horizontal_advection_type, + horizontal_advection_limiter, + vertical_advection_type, + vertical_advection_limiter, + *, + grid_savepoint, + icon_grid, + interpolation_savepoint, + metrics_savepoint, + # data_provider, + backend, + advection_init_savepoint, + advection_exit_savepoint, + experiment: test_defs.Experiment, + processor_props: definitions.ProcessProperties, + decomposition_info: definitions.DecompositionInfo, # : F811 fixture + advection_lsq_state, +): + if test_utils.is_embedded(backend): + # https://github.com/GridTools/gt4py/issues/1583 + pytest.xfail("ValueError: axes don't match array") + # TODO(OngChia): the last datatest fails on GPU (or even CPU) backend when there is no advection because the horizontal flux is not zero. Further check required. + if ( + even_timestep + and horizontal_advection_type == advection.HorizontalAdvectionType.NO_ADVECTION + ): + pytest.xfail( + "This test is skipped until the cause of nonzero horizontal advection if revealed." + ) + + parallel_helpers.check_comm_size(processor_props) + parallel_helpers.log_process_properties(processor_props) + parallel_helpers.log_local_field_size(decomposition_info) + config = construct_config( + horizontal_advection_type=horizontal_advection_type, + horizontal_advection_limiter=horizontal_advection_limiter, + vertical_advection_type=vertical_advection_type, + vertical_advection_limiter=vertical_advection_limiter, + ) + + interpolation_state = construct_interpolation_state(interpolation_savepoint, backend=backend) + + metric_state = construct_metric_state(icon_grid, metrics_savepoint, backend=backend) + edge_geometry = grid_savepoint.construct_edge_geometry() + cell_geometry = grid_savepoint.construct_cell_geometry() + exchange = definitions.create_exchange(processor_props, decomposition_info) + + advection_granule = advection.convert_config_to_advection( + config=config, + grid=icon_grid, + interpolation_state=interpolation_state, + least_squares_state=advection_lsq_state, + metric_state=metric_state, + edge_params=edge_geometry, + cell_params=cell_geometry, + even_timestep=even_timestep, + backend=backend, + exchange=exchange, + ) + + diagnostic_state = construct_diagnostic_init_state( + icon_grid, advection_init_savepoint, ntracer, backend=backend + ) + prep_adv = construct_prep_adv(advection_init_savepoint) + p_tracer_now = advection_init_savepoint.tracer(ntracer) + + p_tracer_new = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, allocator=backend) + dtime = advection_init_savepoint.get_metadata("dtime").get("dtime") + + log_serialized(diagnostic_state, prep_adv, p_tracer_now, dtime) + + advection_granule.run( + diagnostic_state=diagnostic_state, + prep_adv=prep_adv, + p_tracer_now=p_tracer_now, + p_tracer_new=p_tracer_new, + dtime=dtime, + ) + + diagnostic_state_ref = construct_diagnostic_exit_state( + icon_grid, advection_exit_savepoint, ntracer, backend=backend + ) + p_tracer_new_ref = advection_exit_savepoint.tracer(ntracer) + + assert test_helpers.dallclose( + diagnostic_state.hfl_tracer.asnumpy(), diagnostic_state_ref.hfl_tracer.asnumpy(), atol=1e-8 + ) + assert test_utils.dallclose( + diagnostic_state.vfl_tracer.asnumpy(), + diagnostic_state_ref.vfl_tracer.asnumpy(), + rtol=1e-10, + ) + assert test_helpers.dallclose(p_tracer_new_ref.asnumpy(), p_tracer_new.asnumpy(), atol=1e-10) + + +@pytest.mark.level("unit") +@pytest.mark.datatest +@pytest.mark.mpi +def test_compute_lsq_coeffs( + icon_grid: base_grid.Grid, + grid_savepoint: sb.IconGridSavepoint, + backend: gtx_typing.Backend, + interpolation_savepoint: sb.InterpolationSavepoint, + experiment: test_defs.Experiment, + processor_props: definitions.ProcessProperties, + decomposition_info: definitions.DecompositionInfo, # : F811 fixture +) -> None: + gm = grid_utils.get_grid_manager_from_identifier( + experiment.grid, + num_levels=1, + keep_skip_values=True, + allocator=backend, + ) + + c2e2c = gm.grid.connectivities["C2E2C"].asnumpy() + cell_owner_mask = grid_savepoint.c_owner_mask().asnumpy() + grid_sphere_radius = constants.EARTH_RADIUS + lsq_dim_unk = 2 + lsq_dim_c = 3 + lsq_wgt_exp = 2 + cell_domain = h_grid.domain(dims.CellDim) + + min_rlcell_int = gm.grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + start_idx = gm.grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2)) + + grid_geometry = grid_utils.get_grid_geometry(backend, experiment) + cell_center_x = grid_geometry.get(geometry_attrs.CELL_CENTER_X).asnumpy() + cell_center_y = grid_geometry.get(geometry_attrs.CELL_CENTER_Y).asnumpy() + domain_length = gm.grid.global_properties.domain_length + domain_height = gm.grid.global_properties.domain_height + lsq_dim_stencil = 3 + exchange = definitions.create_exchange(processor_props, decomposition_info) + + coordinates = gm.coordinates + cell_lat = coordinates[dims.CellDim]["lat"].asnumpy() + cell_lon = coordinates[dims.CellDim]["lon"].asnumpy() + lsq_pseudoinv = compute_lsq_coeffs( + cell_center_x, + cell_center_y, + cell_lat, + cell_lon, + c2e2c, + cell_owner_mask, + domain_length, + domain_height, + grid_sphere_radius, + lsq_dim_unk, + lsq_dim_c, + lsq_wgt_exp, + lsq_dim_stencil, + start_idx, + min_rlcell_int, + icon_grid.geometry_type.value, + exchange, + ) + + assert test_helpers.dallclose( + interpolation_savepoint.lsq_pseudoinv_1().asnumpy(), lsq_pseudoinv[:, 0, :], atol=1e-15 + ) + assert test_helpers.dallclose( + interpolation_savepoint.lsq_pseudoinv_2().asnumpy(), lsq_pseudoinv[:, 1, :], atol=1e-15 + ) diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py index d50765ea3f..5ec0f04224 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_attributes.py @@ -34,6 +34,7 @@ RBF_SCALE_CELL: Final[str] = "rbf_scale_cell" RBF_SCALE_EDGE: Final[str] = "rbf_scale_edge" RBF_SCALE_VERTEX: Final[str] = "rbf_scale_vertex" +LSQ_PSEUDOINV: Final[str] = "lsq_interpolation_coefficient" attrs: dict[str, model.FieldMetaData] = { C_LIN_E: dict( @@ -212,4 +213,12 @@ icon_var_name="rbf_vec_scale_v", dtype=ta.wpfloat, ), + LSQ_PSEUDOINV: dict( + standard_name=LSQ_PSEUDOINV, + long_name="lsq_pseudoinv", + units="", + dims=(dims.CellDim, dims.C2E2CDim), + icon_var_name="ptr_int_lsq%lsq_pseudoinv", + dtype=ta.wpfloat, + ), } diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py index a22b909fd8..343ce8a8a1 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py @@ -57,6 +57,8 @@ def __init__( self._providers: dict[str, factory.FieldProvider] = {} self._geometry = geometry_source self._exchange = exchange + domain_length = self.grid.global_properties.domain_length + domain_height = self.grid.global_properties.domain_height # TODO @halungge: Dummy config dict - to be replaced by real configuration self._config = { "divergence_averaging_central_cell_weight": 0.5, # divavg_cntrwgt in ICON @@ -67,6 +69,12 @@ def __init__( "rbf_kernel_cell": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.CELL], "rbf_kernel_edge": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.EDGE], "rbf_kernel_vertex": rbf.DEFAULT_RBF_KERNEL[rbf.RBFDimension.VERTEX], + "lsq_dim_unk": 2, + "lsq_dim_c": 3, + "lsq_wgt_exp": 2, + "lsq_dim_stencil": 3, + "domain_length": 0.0 if domain_length is None else domain_length, + "domain_height": 0.0 if domain_height is None else domain_height, } log.info( f"initialized interpolation factory for backend = '{self._backend_name()}' and grid = '{self._grid}'" @@ -228,6 +236,39 @@ def _register_computed_fields(self) -> None: ) self.register_provider(rbf_scale_vertex_np) + lsq_pseudoinv = factory.NumpyDataProvider( + func=functools.partial( + interpolation_fields.compute_lsq_coeffs, + exchange=functools.partial(self._exchange.exchange_and_wait, dims.CellDim), + array_ns=self._xp, + ), + fields=(attrs.LSQ_PSEUDOINV,), + domain=(), + deps={ + "cell_center_x": geometry_attrs.CELL_CENTER_X, + "cell_center_y": geometry_attrs.CELL_CENTER_Y, + "cell_lat": geometry_attrs.CELL_LAT, + "cell_lon": geometry_attrs.CELL_LON, + "cell_owner_mask": "cell_owner_mask", + }, + connectivities={"c2e2c": dims.C2E2CDim}, + params={ + "domain_length": self._config["domain_length"], + "domain_height": self._config["domain_height"], + "grid_sphere_radius": constants.EARTH_RADIUS, + "lsq_dim_unk": self._config["lsq_dim_unk"], + "lsq_dim_c": self._config["lsq_dim_c"], + "lsq_wgt_exp": self._config["lsq_wgt_exp"], + "lsq_dim_stencil": self._config["lsq_dim_stencil"], + "start_idx": self.grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ), + "min_rlcell_int": self.grid.end_index(cell_domain(h_grid.Zone.HALO_LEVEL_2)), + "geometry_type": self.grid.global_properties.geometry_type.value, + }, + ) + self.register_provider(lsq_pseudoinv) + match self.grid.global_properties.geometry_type: case base.GeometryType.ICOSAHEDRON: cell_average_weight = factory.NumpyDataProvider( diff --git a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py index a4e30d4659..ee4f661657 100644 --- a/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py +++ b/model/common/src/icon4py/model/common/interpolation/interpolation_fields.py @@ -13,6 +13,7 @@ from typing import Final import numpy as np +import scipy from gt4py import next as gtx from gt4py.next import where @@ -20,9 +21,11 @@ import icon4py.model.common.math.projection as proj import icon4py.model.common.type_alias as ta from icon4py.model.common import dimension as dims +from icon4py.model.common.decomposition import definitions as decomposition from icon4py.model.common.dimension import C2E, V2E -from icon4py.model.common.grid import gridfile +from icon4py.model.common.grid import base as base_grid, gridfile from icon4py.model.common.grid.geometry_stencils import compute_primal_cart_normal +from icon4py.model.common.math.projection import diff_on_edges_torus_numpy, gnomonic_proj from icon4py.model.common.utils import data_allocation as data_alloc @@ -36,7 +39,7 @@ def compute_c_lin_e( inv_dual_edge_length: data_alloc.NDArray, edge_owner_mask: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -110,7 +113,7 @@ def compute_geofac_n2s( e2c: data_alloc.NDArray, c2e2c: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -171,7 +174,7 @@ def compute_geofac_grg( e2c: data_alloc.NDArray, c2e2c: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]: owned = array_ns.stack((owner_mask, owner_mask, owner_mask)).T @@ -218,7 +221,7 @@ def compute_geofac_grdiv( e2c: data_alloc.NDArray, e2c2e: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -439,7 +442,7 @@ def _force_mass_conservation_to_c_bln_avg( cell_owner_mask: data_alloc.NDArray, divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, niter: int = 1000, ) -> data_alloc.NDArray: @@ -617,7 +620,7 @@ def compute_mass_conserving_bilinear_cell_average_weight( divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, horizontal_start_level_3: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: c_bln_avg = _compute_c_bln_avg( @@ -650,7 +653,7 @@ def compute_mass_conserving_bilinear_cell_average_weight_torus( divergence_averaging_central_cell_weight: ta.wpfloat, horizontal_start: gtx.int32, horizontal_start_level_3: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: c_bln_avg = _compute_uniform_c_bln_avg( @@ -721,7 +724,7 @@ def compute_e_flx_avg( e2c2e: data_alloc.NDArray, horizontal_start_p3: gtx.int32, horizontal_start_p4: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -886,7 +889,7 @@ def compute_cells_aw_verts( v2c: data_alloc.NDArray, e2c: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -1019,7 +1022,7 @@ def compute_pos_on_tplane_e_x_y( owner_mask: data_alloc.NDArray, e2c: data_alloc.NDArray, horizontal_start: gtx.int32, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]: """ @@ -1106,7 +1109,7 @@ def compute_pos_on_tplane_e_x_y( def compute_pos_on_tplane_e_x_y_torus( dual_edge_length: data_alloc.NDArray, e2c: data_alloc.NDArray, - exchange: Callable[[data_alloc.NDArray], None], + exchange: Callable[[data_alloc.NDArray], None] = decomposition.single_node_default, array_ns: ModuleType = np, ) -> data_alloc.NDArray: """ @@ -1149,3 +1152,145 @@ def compute_pos_on_tplane_e_x_y_torus( exchange(pos_on_tplane_e_x, pos_on_tplane_e_y) return pos_on_tplane_e_x, pos_on_tplane_e_y + + +def compute_lsq_pseudoinv( + cell_owner_mask: data_alloc.NDArray, + lsq_pseudoinv: data_alloc.NDArray, + z_lsq_mat_c: data_alloc.NDArray, + lsq_weights_c: data_alloc.NDArray, + start_idx: int, + min_rlcell_int: int, + lsq_dim_unk: int, + lsq_dim_c: int, +) -> data_alloc.NDArray: + for jjb in range(lsq_dim_c): + for jjk in range(lsq_dim_unk): + for jc in range(start_idx, min_rlcell_int): + u, s, v_t, _ = scipy.linalg.lapack.dgesdd(z_lsq_mat_c[jc, :, :]) + if cell_owner_mask[jc]: + lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = ( + lsq_pseudoinv[jc, :lsq_dim_unk, jjb] + + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb] + ) + return lsq_pseudoinv + + +def compute_lsq_weights_c( + z_dist_g: data_alloc.NDArray, + lsq_weights_c_jc: data_alloc.NDArray, + lsq_dim_stencil: int, + lsq_wgt_exp: int, +) -> data_alloc.NDArray: + for js in range(lsq_dim_stencil): + z_norm = np.sqrt(np.dot(z_dist_g[js, :], z_dist_g[js, :])) + lsq_weights_c_jc[js] = 1.0 / (z_norm**lsq_wgt_exp) + return lsq_weights_c_jc / np.max(lsq_weights_c_jc) + + +def compute_z_lsq_mat_c( + cell_owner_mask: data_alloc.NDArray, + z_lsq_mat_c: data_alloc.NDArray, + lsq_weights_c: data_alloc.NDArray, + z_dist_g: data_alloc.NDArray, + jc: int, + lsq_dim_unk: int, + lsq_dim_c: int, +) -> data_alloc.NDArray: + min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) + if cell_owner_mask[jc]: + z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 + + for js in range(lsq_dim_c): + z_lsq_mat_c[jc, js, :lsq_dim_unk] = lsq_weights_c[jc, js] * z_dist_g[js, :] + + return z_lsq_mat_c[jc, js, :lsq_dim_unk] + + +def compute_lsq_coeffs( + cell_center_x: data_alloc.NDArray, + cell_center_y: data_alloc.NDArray, + cell_lat: data_alloc.NDArray, + cell_lon: data_alloc.NDArray, + c2e2c: data_alloc.NDArray, + cell_owner_mask: data_alloc.NDArray, + domain_length: float, + domain_height: float, + grid_sphere_radius: float, + lsq_dim_unk: int, + lsq_dim_c: int, + lsq_wgt_exp: int, + lsq_dim_stencil: int, + start_idx: int, + min_rlcell_int: int, + geometry_type: int, + exchange: decomposition.ExchangeRuntime = decomposition.single_node_default, + array_ns: ModuleType = np, +) -> data_alloc.NDArray: + lsq_weights_c = array_ns.zeros((min_rlcell_int, lsq_dim_stencil)) + lsq_pseudoinv = array_ns.zeros((min_rlcell_int, lsq_dim_unk, lsq_dim_c)) + z_lsq_mat_c = array_ns.zeros((min_rlcell_int, lsq_dim_c, lsq_dim_c)) + z_dist_g = array_ns.zeros((min_rlcell_int, lsq_dim_c, 2)) + match base_grid.GeometryType(geometry_type): + case base_grid.GeometryType.ICOSAHEDRON: + for js in range(lsq_dim_stencil): + z_dist_g[:, js, :] = np.asarray( + gnomonic_proj( + cell_lon[:], cell_lat[:], cell_lon[c2e2c[:, js]], cell_lat[c2e2c[:, js]] + ) + ).T + + z_dist_g *= grid_sphere_radius + min_lsq_bound = min(lsq_dim_unk, lsq_dim_c) + + for jc in range(start_idx, min_rlcell_int): + if cell_owner_mask[jc]: + z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0 + case base_grid.GeometryType.TORUS: + for jc in range(start_idx, min_rlcell_int): + ilc_s = c2e2c[jc, :lsq_dim_stencil] + cc_cell = array_ns.zeros((lsq_dim_stencil, 2)) + + cc_cv = (cell_center_x[jc], cell_center_y[jc]) + for js in range(lsq_dim_stencil): + cc_cell[js, :] = diff_on_edges_torus_numpy( + cell_center_x[jc], + cell_center_y[jc], + cell_center_x[ilc_s][js], + cell_center_y[ilc_s][js], + domain_length, + domain_height, + ) + z_dist_g[jc, :, :] = cc_cell - cc_cv + + for jc in range(start_idx, min_rlcell_int): + lsq_weights_c[jc, :] = compute_lsq_weights_c( + z_dist_g[jc, :, :], lsq_weights_c[jc, :], lsq_dim_stencil, lsq_wgt_exp + ) + z_lsq_mat_c[jc, js, :lsq_dim_unk] = compute_z_lsq_mat_c( + cell_owner_mask, + z_lsq_mat_c, + lsq_weights_c, + z_dist_g[jc, :, :], + jc, + lsq_dim_unk, + lsq_dim_c, + ) + + exchange(lsq_weights_c, dim=dims.CellDim) + + lsq_pseudoinv = compute_lsq_pseudoinv( + cell_owner_mask, + lsq_pseudoinv, + z_lsq_mat_c, + lsq_weights_c, + start_idx, + min_rlcell_int, + lsq_dim_unk, + lsq_dim_c, + ) + + exchange(lsq_pseudoinv[:, 0, :], dim=dims.CellDim) + exchange(lsq_pseudoinv[:, 1, :], dim=dims.CellDim) + + return lsq_pseudoinv diff --git a/model/common/src/icon4py/model/common/math/projection.py b/model/common/src/icon4py/model/common/math/projection.py index 9df9650ee1..fcec8fbc5f 100644 --- a/model/common/src/icon4py/model/common/math/projection.py +++ b/model/common/src/icon4py/model/common/math/projection.py @@ -6,7 +6,6 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -import math import numpy as np @@ -48,20 +47,7 @@ def gnomonic_proj( return x, y -def gnomonic_proj_single_val( - lon_c: float, lat_c: float, lon: float, lat: float -) -> tuple[float, float]: - cosc = math.sin(lat_c) * math.sin(lat) + math.cos(lat_c) * math.cos(lat) * math.cos(lon - lon_c) - zk = 1.0 / cosc - - x = zk * math.cos(lat) * math.sin(lon - lon_c) - y = zk * ( - math.cos(lat_c) * math.sin(lat) - math.sin(lat_c) * math.cos(lat) * math.cos(lon - lon_c) - ) - return x, y - - -def plane_torus_closest_coordinates( +def diff_on_edges_torus_numpy( cc_cv_x: float, cc_cv_y: float, cc_cell_x: float, diff --git a/model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py b/model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py index b3a6e699bc..5c2be786b1 100644 --- a/model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py +++ b/model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py @@ -222,3 +222,27 @@ def test_distributed_interpolation_rbf( field_ref = interpolation_savepoint.__getattribute__(intrp_name)().asnumpy() field = factory.get(attrs_name).asnumpy() test_utils.dallclose(field, field_ref, atol=RBF_TOLERANCES[dim][experiment.name]) + + +@pytest.mark.datatest +@pytest.mark.mpi +@pytest.mark.parametrize("processor_props", [True], indirect=True) +def test_distributed_interpolation_lsq_pseudoinv( + backend: gtx_typing.Backend, + interpolation_savepoint: sb.InterpolationSavepoint, + grid_savepoint: sb.IconGridSavepoint, + experiment: test_defs.Experiment, + processor_props: decomposition.ProcessProperties, + decomposition_info: decomposition.DecompositionInfo, + interpolation_factory_from_savepoint: interpolation_factory.InterpolationFieldsFactory, +) -> None: + parallel_helpers.check_comm_size(processor_props) + parallel_helpers.log_process_properties(processor_props) + parallel_helpers.log_local_field_size(decomposition_info) + factory = interpolation_factory_from_savepoint + field_ref_1 = interpolation_savepoint.__getattribute__("lsq_pseudoinv_1")().asnumpy() + field_ref_2 = interpolation_savepoint.__getattribute__("lsq_pseudoinv_2")().asnumpy() + field_1 = factory.get(attrs.LSQ_PSEUDOINV)[:, 0, :] + field_2 = factory.get(attrs.LSQ_PSEUDOINV)[:, 1, :] + test_utils.dallclose(field_1, field_ref_1, atol=1e-15) # type: ignore[arg-type] # mypy does not recognize sliced array as still an array + test_utils.dallclose(field_2, field_ref_2, atol=1e-15) # type: ignore[arg-type] # mypy does not recognize sliced array as still an array diff --git a/uv.lock b/uv.lock index cf5cb697b5..bde0dd5bea 100644 --- a/uv.lock +++ b/uv.lock @@ -1741,18 +1741,14 @@ source = { editable = "model/atmosphere/advection" } dependencies = [ { name = "gt4py" }, { name = "icon4py-common" }, - { name = "numpy" }, { name = "packaging" }, - { name = "scipy" }, ] [package.metadata] requires-dist = [ { name = "gt4py", specifier = "==1.1.4" }, { name = "icon4py-common", editable = "model/common" }, - { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, - { name = "scipy", specifier = ">=1.14.1" }, ] [[package]]