From 4db692345efe04f3baa416f4a6002a28311e31a4 Mon Sep 17 00:00:00 2001 From: David Strassmann <107603957+dastrm@users.noreply.github.com> Date: Mon, 2 Dec 2024 06:45:34 +0100 Subject: [PATCH] Add vertical advection granule with PPM (#605) * Add vertical advection granule with PPM --------- Co-authored-by: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com> --- .../model/atmosphere/advection/advection.py | 29 +- .../atmosphere/advection/advection_states.py | 3 + .../advection/advection_vertical.py | 828 +++++++++++++++++- .../compute_ppm4gpu_courant_number.py | 174 ++++ .../compute_ppm4gpu_fractional_flux.py | 121 +++ .../stencils/compute_ppm4gpu_integer_flux.py | 78 ++ .../compute_vertical_tracer_flux_upwind.py | 43 + .../copy_cell_kdim_field_koff_plus1.py | 42 + ...limit_vertical_slope_semi_monotonically.py | 51 ++ ...test_apply_horizontal_density_increment.py | 12 +- ...apply_interpolated_tracer_time_tendency.py | 10 +- .../test_apply_vertical_density_increment.py | 10 +- ...st_average_horizontal_flux_subcycling_2.py | 6 +- ...st_average_horizontal_flux_subcycling_3.py | 8 +- ...e_antidiffusive_cell_fluxes_and_min_max.py | 36 +- ...test_compute_barycentric_backtrajectory.py | 54 +- ..._compute_barycentric_backtrajectory_alt.py | 38 +- .../test_compute_ffsl_backtrajectory.py | 100 +-- ...cktrajectory_counterclockwise_indicator.py | 14 +- ...te_ffsl_backtrajectory_length_indicator.py | 18 +- ...tal_tracer_flux_from_cubic_coefficients.py | 8 +- ...al_tracer_flux_from_linear_coefficients.py | 24 +- ...racer_flux_from_linear_coefficients_alt.py | 24 +- ...e_horizontal_multiplicative_flux_factor.py | 24 +- ...t_compute_ppm4gpu_parabola_coefficients.py | 4 +- .../test_compute_ppm_all_face_values.py | 18 +- .../test_compute_ppm_quadratic_face_values.py | 4 +- .../test_compute_ppm_quartic_face_values.py | 4 +- .../test_compute_ppm_slope.py | 8 +- .../test_compute_tendency.py | 8 +- ...ute_vertical_parabola_limiter_condition.py | 6 +- ...est_compute_vertical_tracer_flux_upwind.py | 58 ++ ...t_integrate_tracer_density_horizontally.py | 20 +- .../test_integrate_tracer_horizontally.py | 20 +- .../test_integrate_tracer_vertically.py | 22 +- ...it_vertical_parabola_semi_monotonically.py | 12 +- ...limit_vertical_slope_semi_monotonically.py | 52 ++ ...s_antidiffusive_cell_fluxes_and_min_max.py | 28 +- ...est_prepare_ffsl_flux_area_patches_list.py | 422 ++++----- ...cal_quadrature_for_cubic_reconstruction.py | 26 +- ...uadrature_list_for_cubic_reconstruction.py | 50 +- ...test_reconstruct_cubic_coefficients_svd.py | 112 +-- .../tests/advection_tests/test_advection.py | 30 +- .../advection/tests/advection_tests/utils.py | 67 +- .../common/test_utils/datatest_fixtures.py | 2 +- .../model/common/test_utils/datatest_utils.py | 2 +- 46 files changed, 2083 insertions(+), 647 deletions(-) create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py create mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py create mode 100644 model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py index ae2aae72b2..410a40fd03 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection.py @@ -75,6 +75,10 @@ class VerticalAdvectionType(Enum): #: no vertical advection NO_ADVECTION = auto() + #: 1st order upwind + UPWIND_1ST_ORDER = auto() + #: 3rd order PPM + PPM_3RD_ORDER = auto() class VerticalAdvectionLimiter(Enum): @@ -84,6 +88,8 @@ class VerticalAdvectionLimiter(Enum): #: no vertical limiter NO_LIMITER = auto() + #: semi-monotonic vertical limiter + SEMI_MONOTONIC = auto() @dataclasses.dataclass(frozen=True) @@ -393,7 +399,7 @@ def convert_config_to_horizontal_vertical_advection( match config.horizontal_advection_type: case HorizontalAdvectionType.NO_ADVECTION: - horizontal_advection = advection_horizontal.NoAdvection(grid=grid) + horizontal_advection = advection_horizontal.NoAdvection(grid=grid, backend=backend) case HorizontalAdvectionType.LINEAR_2ND_ORDER: tracer_flux = advection_horizontal.SecondOrderMiura( grid=grid, @@ -417,13 +423,32 @@ def convert_config_to_horizontal_vertical_advection( match config.vertical_advection_limiter: case VerticalAdvectionLimiter.NO_LIMITER: - ... + vertical_limiter = advection_vertical.NoLimiter(grid=grid, backend=backend) + case VerticalAdvectionLimiter.SEMI_MONOTONIC: + vertical_limiter = advection_vertical.SemiMonotonicLimiter(grid=grid, backend=backend) case _: raise NotImplementedError(f"Unknown vertical advection limiter.") match config.vertical_advection_type: case VerticalAdvectionType.NO_ADVECTION: vertical_advection = advection_vertical.NoAdvection(grid=grid, backend=backend) + case VerticalAdvectionType.UPWIND_1ST_ORDER: + boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) + vertical_advection = advection_vertical.FirstOrderUpwind( + boundary_conditions=boundary_conditions, + grid=grid, + metric_state=metric_state, + backend=backend, + ) + case VerticalAdvectionType.PPM_3RD_ORDER: + boundary_conditions = advection_vertical.NoFluxCondition(grid=grid, backend=backend) + vertical_advection = advection_vertical.PiecewiseParabolicMethod( + boundary_conditions=boundary_conditions, + vertical_limiter=vertical_limiter, + grid=grid, + metric_state=metric_state, + backend=backend, + ) case _: raise NotImplementedError(f"Unknown vertical advection type.") diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py index 9a038ecdb1..a09ed7dafe 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py @@ -85,3 +85,6 @@ class AdvectionMetricState: #: metrical modification factor for vertical part of divergence at full levels (KDim) deepatmo_divzu: fa.KField[ta.wpfloat] + + #: vertical grid spacing at full levels + ddqz_z_full: fa.CellKField[ta.wpfloat] diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py index 0ddd9daf41..374db87949 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_vertical.py @@ -10,13 +10,59 @@ import logging import icon4py.model.common.grid.states as grid_states +import gt4py.next as gtx from gt4py.next import backend from icon4py.model.atmosphere.advection import advection_states +from icon4py.model.atmosphere.advection.stencils.compute_ppm_quadratic_face_values import ( + compute_ppm_quadratic_face_values, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm_quartic_face_values import ( + compute_ppm_quartic_face_values, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm_slope import compute_ppm_slope +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_courant_number import ( + compute_ppm4gpu_courant_number, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_fractional_flux import ( + compute_ppm4gpu_fractional_flux, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_integer_flux import ( + compute_ppm4gpu_integer_flux, +) +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_parabola_coefficients import ( + compute_ppm4gpu_parabola_coefficients, +) +from icon4py.model.atmosphere.advection.stencils.compute_vertical_parabola_limiter_condition import ( + compute_vertical_parabola_limiter_condition, +) +from icon4py.model.atmosphere.advection.stencils.compute_vertical_tracer_flux_upwind import ( + compute_vertical_tracer_flux_upwind, +) from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field import copy_cell_kdim_field +from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field_koff_minus1 import ( + copy_cell_kdim_field_koff_minus1, +) +from icon4py.model.atmosphere.advection.stencils.copy_cell_kdim_field_koff_plus1 import ( + copy_cell_kdim_field_koff_plus1, +) +from icon4py.model.atmosphere.advection.stencils.init_constant_cell_kdim_field import ( + init_constant_cell_kdim_field, +) +from icon4py.model.atmosphere.advection.stencils.integrate_tracer_vertically import ( + integrate_tracer_vertically, +) +from icon4py.model.atmosphere.advection.stencils.limit_vertical_parabola_semi_monotonically import ( + limit_vertical_parabola_semi_monotonically, +) +from icon4py.model.atmosphere.advection.stencils.limit_vertical_slope_semi_monotonically import ( + limit_vertical_slope_semi_monotonically, +) + from icon4py.model.common import ( + constants, dimension as dims, field_type_aliases as fa, type_alias as ta, @@ -34,6 +80,274 @@ log = logging.getLogger(__name__) +class BoundaryConditions(ABC): + """Class that sets the upper and lower boundary conditions.""" + + @abstractmethod + def run( + self, + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + """ + Set the vertical boundary conditions. + + Args: + p_mflx_tracer_v: output argument, field that contains new vertical tracer mass flux + horizontal_start: input argument, horizontal start index + horizontal_end: input argument, horizontal end index + + """ + ... + + +class NoFluxCondition(BoundaryConditions): + """Class that sets the upper and lower boundary fluxes to zero.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # stencils + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + + def run( + self, + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + log.debug("vertical boundary conditions computation - start") + + # set upper boundary conditions + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=p_mflx_tracer_v, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + # set lower boundary conditions + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=p_mflx_tracer_v, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + log.debug("vertical boundary conditions computation - end") + + +class VerticalLimiter(ABC): + """Class that limits the vertical reconstructed fields and the fluxes.""" + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + +class NoLimiter(VerticalLimiter): + """Class that implements no vertical parabola limiter.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # fields + self._l_limit = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + + # stencils + self._copy_cell_kdim_field = copy_cell_kdim_field.with_backend(self._backend) + self._copy_cell_kdim_field_koff_plus1 = copy_cell_kdim_field_koff_plus1.with_backend( + self._backend + ) + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + # simply copy to up/low face values + log.debug("running stencil copy_cell_kdim_field - start") + self._copy_cell_kdim_field( + field_in=p_face, + field_out=p_face_up, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field - end") + + log.debug("running stencil copy_cell_kdim_field_koff_plus1 - start") + self._copy_cell_kdim_field_koff_plus1( + field_in=p_face, + field_out=p_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field_koff_plus1 - end") + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + +class SemiMonotonicLimiter(VerticalLimiter): + """Class that implements a semi-monotonic vertical parabola limiter.""" + + def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): + # input arguments + self._grid = grid + self._backend = backend + + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._l_limit = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, dtype=gtx.int32, backend=self._backend + ) + + # stencils + self._limit_vertical_slope_semi_monotonically = ( + limit_vertical_slope_semi_monotonically.with_backend(self._backend) + ) + self._compute_vertical_parabola_limiter_condition = ( + compute_vertical_parabola_limiter_condition.with_backend(self._backend) + ) + self._limit_vertical_parabola_semi_monotonically = ( + limit_vertical_parabola_semi_monotonically.with_backend(self._backend) + ) + + def limit_slope( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + log.debug("running stencil limit_vertical_slope_semi_monotonically - start") + self._limit_vertical_slope_semi_monotonically( + p_cc=p_tracer_now, + z_slope=z_slope, + k=self._k_field, + elev=self._grid.num_levels - 1, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil limit_vertical_slope_semi_monotonically - end") + + def limit_parabola( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_face: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_face_up: fa.CellKField[ta.wpfloat], + p_face_low: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + # compute semi-monotonic limiter condition + log.debug("running stencil compute_vertical_parabola_limiter_condition - start") + self._compute_vertical_parabola_limiter_condition( + p_face=p_face, + p_cc=p_tracer_now, + l_limit=self._l_limit, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_vertical_parabola_limiter_condition - end") + + # apply semi-monotonic limiter condition and store to up/low face values + log.debug("running stencil limit_vertical_parabola_semi_monotonically - start") + self._limit_vertical_parabola_semi_monotonically( + l_limit=self._l_limit, + p_face=p_face, + p_cc=p_tracer_now, + p_face_up=p_face_up, + p_face_low=p_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil limit_vertical_parabola_semi_monotonically - end") + + def limit_fluxes( + self, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + ): + ... + + class VerticalAdvection(ABC): """Class that does one vertical advection step.""" @@ -62,6 +376,12 @@ def run( dtime: input argument, the time step even_timestep: input argument, determines whether halo points are included + Note: + Originally, if even step, vertical transport includes all halo points in order to avoid an additional synchronization step, i.e. + if lstep_even: i_rlstart = 2, i_rlend = min_rlcell + else: i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int + Horizontal advection is always called with the same indices though, i.e. + i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int """ ... @@ -90,6 +410,16 @@ def __init__(self, grid: icon_grid.IconGrid, backend: backend.Backend): log.debug("vertical advection class init - end") + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + def run( self, prep_adv: advection_states.AdvectionPrepAdvState, @@ -103,15 +433,16 @@ def run( ): log.debug("vertical advection run - start") - horizontal_start = ( - self._start_cell_lateral_boundary_level_2 if even_timestep else self._start_cell_nudging + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep ) + log.debug("running stencil copy_cell_kdim_field - start") self._copy_cell_kdim_field( field_in=p_tracer_now, field_out=p_tracer_new, horizontal_start=horizontal_start, - horizontal_end=(self._end_cell_end if even_timestep else self._end_cell_local), + horizontal_end=horizontal_end, vertical_start=0, vertical_end=self._grid.num_levels, offset_provider=self._grid.offset_providers, @@ -135,13 +466,7 @@ def run( dtime: ta.wpfloat, even_timestep: bool = False, ): - log.debug("horizontal advection run - start") - - # TODO (dastrm): maybe change how the indices are handled here? originally: - # if even step, vertical transport includes all halo points in order to avoid an additional synchronization step, i.e. - # if lstep_even: i_rlstart = 2, i_rlend = min_rlcell - # else: i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int - # note: horizontal advection is always called with the same indices, i.e. i_rlstart = grf_bdywidth_c+1, i_rlend = min_rlcell_int + log.debug("vertical advection run - start") self._compute_numerical_flux( prep_adv=prep_adv, @@ -149,6 +474,7 @@ def run( rhodz_now=rhodz_now, p_mflx_tracer_v=p_mflx_tracer_v, dtime=dtime, + even_timestep=even_timestep, ) self._update_unknowns( @@ -158,9 +484,10 @@ def run( rhodz_new=rhodz_new, p_mflx_tracer_v=p_mflx_tracer_v, dtime=dtime, + even_timestep=even_timestep, ) - log.debug("horizontal advection run - end") + log.debug("vertical advection run - end") @abstractmethod def _compute_numerical_flux( @@ -170,6 +497,7 @@ def _compute_numerical_flux( rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): ... @@ -182,35 +510,28 @@ def _update_unknowns( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): ... -class SemiLagrangian(FiniteVolume): - """Class that does one vertical semi-Lagrangian finite volume advection step.""" +class FirstOrderUpwind(FiniteVolume): + """Class that does one vertical first-order accurate upwind finite volume advection step.""" def __init__( self, + boundary_conditions: BoundaryConditions, grid: icon_grid.IconGrid, - interpolation_state: advection_states.AdvectionInterpolationState, - least_squares_state: advection_states.AdvectionLeastSquaresState, metric_state: advection_states.AdvectionMetricState, - edge_params: grid_states.EdgeParams, - cell_params: grid_states.CellParams, - backend: backend.Backend, - exchange: decomposition.ExchangeRuntime = decomposition.SingleNodeExchange(), + backend=backend, ): log.debug("vertical advection class init - start") # input arguments + self._boundary_conditions = boundary_conditions self._grid = grid - self._interpolation_state = interpolation_state - self._least_squares_state = least_squares_state self._metric_state = metric_state - self._edge_params = edge_params - self._cell_params = cell_params self._backend = backend - self._exchange = exchange # cell indices cell_domain = h_grid.domain(dims.CellDim) @@ -221,8 +542,36 @@ def __init__( self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) self._end_cell_end = self._grid.end_index(cell_domain(h_grid.Zone.END)) + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + + # stencils + self._compute_vertical_tracer_flux_upwind = ( + compute_vertical_tracer_flux_upwind.with_backend(self._backend) + ) + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + self._integrate_tracer_vertically = integrate_tracer_vertically.with_backend(self._backend) + + # misc + self._ivadv_tracer = 1 + self._iadv_slev_jt = 0 + log.debug("vertical advection class init - end") + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + def _compute_numerical_flux( self, prep_adv: advection_states.AdvectionPrepAdvState, @@ -230,9 +579,35 @@ def _compute_numerical_flux( rhodz_now: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): - # TODO (dastrm): implement this - ... + log.debug("vertical numerical flux computation - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + log.debug("running stencil compute_vertical_tracer_flux_upwind - start") + self._compute_vertical_tracer_flux_upwind( + p_cc=p_tracer_now, + p_mflx_contra_v=prep_adv.mass_flx_ic, + p_upflux=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_vertical_tracer_flux_upwind - end") + + # set boundary conditions + self._boundary_conditions.run( + p_mflx_tracer_v=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + log.debug("vertical numerical flux computation - end") def _update_unknowns( self, @@ -242,6 +617,403 @@ def _update_unknowns( rhodz_new: fa.CellKField[ta.wpfloat], p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim dtime: ta.wpfloat, + even_timestep: bool, ): - # TODO (dastrm): implement this - ... + log.debug("vertical unknowns update - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + # update tracer mass fraction + log.debug("running stencil integrate_tracer_vertically - start") + self._integrate_tracer_vertically( + tracer_now=p_tracer_now, + rhodz_now=rhodz_now, + p_mflx_tracer_v=p_mflx_tracer_v, + deepatmo_divzl=self._metric_state.deepatmo_divzl, + deepatmo_divzu=self._metric_state.deepatmo_divzu, + rhodz_new=rhodz_new, + tracer_new=p_tracer_new, + k=self._k_field, + p_dtime=dtime, + ivadv_tracer=self._ivadv_tracer, + iadv_slev_jt=self._iadv_slev_jt, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil integrate_tracer_vertically - end") + + log.debug("vertical unknowns update - end") + + +class PiecewiseParabolicMethod(FiniteVolume): + """Class that does one vertical PPM finite volume advection step.""" + + def __init__( + self, + boundary_conditions: BoundaryConditions, + vertical_limiter: VerticalLimiter, + grid: icon_grid.IconGrid, + metric_state: advection_states.AdvectionMetricState, + backend=backend, + ): + log.debug("vertical advection class init - start") + + # input arguments + self._boundary_conditions = boundary_conditions + self._vertical_limiter = vertical_limiter + self._grid = grid + self._metric_state = metric_state + self._backend = backend + + # cell indices + cell_domain = h_grid.domain(dims.CellDim) + self._start_cell_lateral_boundary_level_2 = self._grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + self._start_cell_nudging = self._grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + self._end_cell_end = self._grid.end_index(cell_domain(h_grid.Zone.END)) + + # fields + self._k_field = field_alloc.allocate_indices( + dims.KDim, grid=self._grid, is_halfdim=True, dtype=gtx.int32, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_cfl = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, is_halfdim=True, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_slope = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_face = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, is_halfdim=True, backend=self._backend + ) # TODO (dastrm): should be KHalfDim + self._z_face_up = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_face_low = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_delta_q = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + self._z_a1 = field_alloc.allocate_zero_field( + dims.CellDim, dims.KDim, grid=self._grid, backend=self._backend + ) + + # stencils + self._init_constant_cell_kdim_field = init_constant_cell_kdim_field.with_backend( + self._backend + ) + self._compute_ppm4gpu_courant_number = compute_ppm4gpu_courant_number.with_backend( + self._backend + ) + self._compute_ppm_slope = compute_ppm_slope.with_backend(self._backend) + self._compute_ppm_quadratic_face_values = compute_ppm_quadratic_face_values.with_backend( + self._backend + ) + self._compute_ppm_quartic_face_values = compute_ppm_quartic_face_values.with_backend( + self._backend + ) + self._copy_cell_kdim_field = copy_cell_kdim_field.with_backend(self._backend) + self._copy_cell_kdim_field_koff_minus1 = copy_cell_kdim_field_koff_minus1.with_backend( + self._backend + ) + self._compute_ppm4gpu_parabola_coefficients = ( + compute_ppm4gpu_parabola_coefficients.with_backend(self._backend) + ) + self._compute_ppm4gpu_fractional_flux = compute_ppm4gpu_fractional_flux.with_backend( + self._backend + ) + self._compute_ppm4gpu_integer_flux = compute_ppm4gpu_integer_flux.with_backend( + self._backend + ) + self._integrate_tracer_vertically = integrate_tracer_vertically.with_backend(self._backend) + + # misc + self._slev = 0 + self._slevp1_ti = 1 + self._elev = self._grid.num_levels - 1 + self._nlev = self._grid.num_levels - 1 + self._ivadv_tracer = 1 + self._iadv_slev_jt = 0 + + log.debug("vertical advection class init - end") + + def _get_horizontal_start_end(self, even_timestep: bool): + if even_timestep: + horizontal_start = self._start_cell_lateral_boundary_level_2 + horizontal_end = self._end_cell_end + else: + horizontal_start = self._start_cell_nudging + horizontal_end = self._end_cell_local + + return horizontal_start, horizontal_end + + def _compute_numerical_flux( + self, + prep_adv: advection_states.AdvectionPrepAdvState, + p_tracer_now: fa.CellKField[ta.wpfloat], + rhodz_now: fa.CellKField[ta.wpfloat], + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + dtime: ta.wpfloat, + even_timestep: bool, + ): + log.debug("vertical numerical flux computation - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + ## compute density-weighted Courant number + + log.debug("running stencil init_constant_cell_kdim_field - start") + self._init_constant_cell_kdim_field( + field=self._z_cfl, + value=0.0, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil init_constant_cell_kdim_field - end") + + log.debug("running stencil compute_ppm4gpu_courant_number - start") + self._compute_ppm4gpu_courant_number( + p_mflx_contra_v=prep_adv.mass_flx_ic, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + k=self._k_field, + slevp1_ti=self._slevp1_ti, + nlev=self._nlev, + dbl_eps=constants.DBL_EPS, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_courant_number - end") + + ## reconstruct face values + + # compute slope + log.debug("running stencil compute_ppm_slope - start") + self._compute_ppm_slope( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + k=self._k_field, + z_slope=self._z_slope, + elev=self._elev, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_slope - end") + + # limit slope + self._vertical_limiter.limit_slope( + p_tracer_now=p_tracer_now, + z_slope=self._z_slope, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + # compute second highest face value + log.debug("running stencil compute_ppm_quadratic_face_values - start") + self._compute_ppm_quadratic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=2, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quadratic_face_values - end") + + # compute second lowest face value + log.debug("running stencil compute_ppm_quadratic_face_values - start") + self._compute_ppm_quadratic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels - 1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quadratic_face_values - end") + + # compute highest face value + log.debug("running stencil copy_cell_kdim_field - start") + self._copy_cell_kdim_field( + field_in=p_tracer_now, + field_out=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field - end") + + # compute lowest face value + log.debug("running stencil copy_cell_kdim_field_koff_minus1 - start") + self._copy_cell_kdim_field_koff_minus1( + field_in=p_tracer_now, + field_out=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=self._grid.num_levels, + vertical_end=self._grid.num_levels + 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil copy_cell_kdim_field_koff_minus1 - end") + + # compute all other face values + log.debug("running stencil compute_ppm_quartic_face_values - start") + self._compute_ppm_quartic_face_values( + p_cc=p_tracer_now, + p_cellhgt_mc_now=self._metric_state.ddqz_z_full, + z_slope=self._z_slope, + p_face=self._z_face, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=2, + vertical_end=self._grid.num_levels - 1, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm_quartic_face_values - end") + + ## limit reconstruction + + self._vertical_limiter.limit_parabola( + p_tracer_now=p_tracer_now, + p_face=self._z_face, + p_face_up=self._z_face_up, + p_face_low=self._z_face_low, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + ## compute fractional numerical flux + + log.debug("running stencil compute_ppm4gpu_parabola_coefficients - start") + self._compute_ppm4gpu_parabola_coefficients( + z_face_up=self._z_face_up, + z_face_low=self._z_face_low, + p_cc=p_tracer_now, + z_delta_q=self._z_delta_q, + z_a1=self._z_a1, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_parabola_coefficients - end") + + log.debug("running stencil compute_ppm4gpu_fractional_flux - start") + self._compute_ppm4gpu_fractional_flux( + p_cc=p_tracer_now, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + z_delta_q=self._z_delta_q, + z_a1=self._z_a1, + p_upflux=p_mflx_tracer_v, + k=self._k_field, + slev=self._slev, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_fractional_flux - end") + + ## compute integer numerical flux + + log.debug("running stencil compute_ppm4gpu_integer_flux - start") + self._compute_ppm4gpu_integer_flux( + p_cc=p_tracer_now, + p_cellmass_now=rhodz_now, + z_cfl=self._z_cfl, + p_upflux=p_mflx_tracer_v, + k=self._k_field, + slev=self._slev, + p_dtime=dtime, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=1, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil compute_ppm4gpu_integer_flux - end") + + ## set boundary conditions + + self._boundary_conditions.run( + p_mflx_tracer_v=p_mflx_tracer_v, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + ) + + ## apply flux limiter + + self._vertical_limiter.limit_fluxes( + horizontal_start=horizontal_start, horizontal_end=horizontal_end + ) + + log.debug("vertical numerical flux computation - end") + + def _update_unknowns( + self, + p_tracer_now: fa.CellKField[ta.wpfloat], + p_tracer_new: fa.CellKField[ta.wpfloat], + rhodz_now: fa.CellKField[ta.wpfloat], + rhodz_new: fa.CellKField[ta.wpfloat], + p_mflx_tracer_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + dtime: ta.wpfloat, + even_timestep: bool, + ): + log.debug("vertical unknowns update - start") + + horizontal_start, horizontal_end = self._get_horizontal_start_end( + even_timestep=even_timestep + ) + + # update tracer mass fraction + log.debug("running stencil integrate_tracer_vertically - start") + self._integrate_tracer_vertically( + tracer_now=p_tracer_now, + rhodz_now=rhodz_now, + p_mflx_tracer_v=p_mflx_tracer_v, + deepatmo_divzl=self._metric_state.deepatmo_divzl, + deepatmo_divzu=self._metric_state.deepatmo_divzu, + rhodz_new=rhodz_new, + tracer_new=p_tracer_new, + k=self._k_field, + p_dtime=dtime, + ivadv_tracer=self._ivadv_tracer, + iadv_slev_jt=self._iadv_slev_jt, + horizontal_start=horizontal_start, + horizontal_end=horizontal_end, + vertical_start=0, + vertical_end=self._grid.num_levels, + offset_provider=self._grid.offset_providers, + ) + log.debug("running stencil integrate_tracer_vertically - end") + + log.debug("vertical unknowns update - end") diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py new file mode 100644 index 0000000000..2c45621b29 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_courant_number.py @@ -0,0 +1,174 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import abs, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _compute_courant_number_below( + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_mass: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + nlev: gtx.int32, + dbl_eps: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + z_mass_pos = z_mass > 0.0 + + in_bounds_p0 = k <= nlev - 1 + in_bounds_p1 = k <= nlev - 2 + in_bounds_p2 = k <= nlev - 3 + in_bounds_p3 = k <= nlev - 4 + + mass_gt_cellmass_p0 = where(z_mass_pos & in_bounds_p0, z_mass >= p_cellmass_now, False) + z_mass = z_mass - where(mass_gt_cellmass_p0, p_cellmass_now, 0.0) + + mass_gt_cellmass_p1 = mass_gt_cellmass_p0 & where( + z_mass_pos & in_bounds_p1, z_mass >= p_cellmass_now(Koff[1]), False + ) + z_mass = z_mass - where(mass_gt_cellmass_p1, p_cellmass_now(Koff[1]), 0.0) + + mass_gt_cellmass_p2 = mass_gt_cellmass_p1 & where( + z_mass_pos & in_bounds_p2, z_mass >= p_cellmass_now(Koff[2]), False + ) + z_mass = z_mass - where(mass_gt_cellmass_p2, p_cellmass_now(Koff[2]), 0.0) + + mass_gt_cellmass_p3 = mass_gt_cellmass_p2 & where( + z_mass_pos & in_bounds_p3, z_mass >= p_cellmass_now(Koff[3]), False + ) + z_mass = z_mass - where(mass_gt_cellmass_p3, p_cellmass_now(Koff[3]), 0.0) + + z_cfl = z_cfl + where(mass_gt_cellmass_p0, 1.0, 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass_p1, 1.0, 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass_p2, 1.0, 0.0) + z_cfl = z_cfl + where(mass_gt_cellmass_p3, 1.0, 0.0) + + p_cellmass_now_jks = p_cellmass_now + p_cellmass_now_jks = where(mass_gt_cellmass_p0, p_cellmass_now(Koff[1]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_p1, p_cellmass_now(Koff[2]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_p2, p_cellmass_now(Koff[3]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_p3, p_cellmass_now(Koff[4]), p_cellmass_now_jks) + + z_cflfrac = where(z_mass_pos, z_mass / p_cellmass_now_jks, 0.0) + z_cfl = z_cfl + where(z_cflfrac < 1.0, z_cflfrac, 1.0 - dbl_eps) + + return z_cfl + + +@gtx.field_operator +def _compute_courant_number_above( + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_mass: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slevp1_ti: gtx.int32, + dbl_eps: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + z_mass_neg = z_mass <= 0.0 + + in_bounds_m0 = k >= slevp1_ti + 1 + in_bounds_m1 = k >= slevp1_ti + 2 + in_bounds_m2 = k >= slevp1_ti + 3 + in_bounds_m3 = k >= slevp1_ti + 4 + + mass_gt_cellmass_m0 = where( + z_mass_neg & in_bounds_m0, abs(z_mass) >= p_cellmass_now(Koff[-1]), False + ) + z_mass = z_mass + where(mass_gt_cellmass_m0, p_cellmass_now(Koff[-1]), 0.0) + + mass_gt_cellmass_m1 = mass_gt_cellmass_m0 & where( + z_mass_neg & in_bounds_m1, abs(z_mass) >= p_cellmass_now(Koff[-2]), False + ) + z_mass = z_mass + where(mass_gt_cellmass_m1, p_cellmass_now(Koff[-2]), 0.0) + + mass_gt_cellmass_m2 = mass_gt_cellmass_m1 & where( + z_mass_neg & in_bounds_m2, abs(z_mass) >= p_cellmass_now(Koff[-3]), False + ) + z_mass = z_mass + where(mass_gt_cellmass_m2, p_cellmass_now(Koff[-3]), 0.0) + + mass_gt_cellmass_m3 = mass_gt_cellmass_m2 & where( + z_mass_neg & in_bounds_m3, abs(z_mass) >= p_cellmass_now(Koff[-4]), False + ) + z_mass = z_mass + where(mass_gt_cellmass_m3, p_cellmass_now(Koff[-4]), 0.0) + + p_cellmass_now_jks = p_cellmass_now(Koff[-1]) + p_cellmass_now_jks = where(mass_gt_cellmass_m0, p_cellmass_now(Koff[-2]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_m1, p_cellmass_now(Koff[-3]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_m2, p_cellmass_now(Koff[-4]), p_cellmass_now_jks) + p_cellmass_now_jks = where(mass_gt_cellmass_m3, p_cellmass_now(Koff[-5]), p_cellmass_now_jks) + + z_cfl = z_cfl - where(mass_gt_cellmass_m0, 1.0, 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass_m1, 1.0, 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass_m2, 1.0, 0.0) + z_cfl = z_cfl - where(mass_gt_cellmass_m3, 1.0, 0.0) + + z_cflfrac = where(z_mass_neg, z_mass / p_cellmass_now_jks, 0.0) + z_cfl = z_cfl + where(abs(z_cflfrac) < 1.0, z_cflfrac, dbl_eps - 1.0) + + return z_cfl + + +@gtx.field_operator +def _compute_ppm4gpu_courant_number( + p_mflx_contra_v: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slevp1_ti: gtx.int32, + nlev: gtx.int32, + dbl_eps: ta.wpfloat, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + z_mass = p_dtime * p_mflx_contra_v + + cfl_below = _compute_courant_number_below(p_cellmass_now, z_mass, z_cfl, k, nlev, dbl_eps) + cfl_above = _compute_courant_number_above(p_cellmass_now, z_mass, z_cfl, k, slevp1_ti, dbl_eps) + + z_cfl = cfl_below + cfl_above + + return z_cfl + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_courant_number( + p_mflx_contra_v: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slevp1_ti: gtx.int32, + nlev: gtx.int32, + dbl_eps: ta.wpfloat, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_courant_number( + p_mflx_contra_v, + p_cellmass_now, + z_cfl, + k, + slevp1_ti, + nlev, + dbl_eps, + p_dtime, + out=z_cfl, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py new file mode 100644 index 0000000000..2ac11aa036 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_fractional_flux.py @@ -0,0 +1,121 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import abs, astype, floor, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff +from icon4py.model.common.type_alias import wpfloat + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _sum_neighbor_contributions( + mask1: fa.CellKField[bool], + mask2: fa.CellKField[bool], + js: fa.CellKField[ta.wpfloat], + p_cc: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + p_cc_p0 = where(mask1 & (js == 0.0), p_cc, 0.0) + p_cc_p1 = where(mask1 & (js == 1.0), p_cc(Koff[1]), 0.0) + p_cc_p2 = where(mask1 & (js == 2.0), p_cc(Koff[2]), 0.0) + p_cc_p3 = where(mask1 & (js == 3.0), p_cc(Koff[3]), 0.0) + p_cc_p4 = where(mask1 & (js == 4.0), p_cc(Koff[4]), 0.0) + p_cc_m0 = where(mask2 & (js == 0.0), p_cc(Koff[-1]), 0.0) + p_cc_m1 = where(mask2 & (js == 1.0), p_cc(Koff[-2]), 0.0) + p_cc_m2 = where(mask2 & (js == 2.0), p_cc(Koff[-3]), 0.0) + p_cc_m3 = where(mask2 & (js == 3.0), p_cc(Koff[-4]), 0.0) + p_cc_m4 = where(mask2 & (js == 4.0), p_cc(Koff[-5]), 0.0) + p_cc_jks = ( + p_cc_p0 + + p_cc_p1 + + p_cc_p2 + + p_cc_p3 + + p_cc_p4 + + p_cc_m0 + + p_cc_m1 + + p_cc_m2 + + p_cc_m3 + + p_cc_m4 + ) + return p_cc_jks + + +@gtx.field_operator +def _compute_ppm4gpu_fractional_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + z_delta_q: fa.CellKField[ta.wpfloat], + z_a1: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + js = floor(abs(z_cfl)) + z_cflfrac = abs(z_cfl) - js + + z_cfl_pos = z_cfl > 0.0 + z_cfl_neg = not z_cfl_pos + wsign = where(z_cfl_pos, 1.0, -1.0) + + in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat) + + p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc) + p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now) + z_delta_q_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_delta_q) + z_a1_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, z_a1) + + z_q_int = ( + p_cc_jks + + wsign * (z_delta_q_jks * (1.0 - z_cflfrac)) + - z_a1_jks * (1.0 - 3.0 * z_cflfrac + 2.0 * z_cflfrac * z_cflfrac) + ) + + p_upflux = where( + in_slev_bounds, wsign * p_cellmass_now_jks * z_cflfrac * z_q_int / p_dtime, 0.0 + ) + + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_fractional_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + z_delta_q: fa.CellKField[ta.wpfloat], + z_a1: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_fractional_flux( + p_cc, + p_cellmass_now, + z_cfl, + z_delta_q, + z_a1, + k, + slev, + p_dtime, + out=p_upflux, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py new file mode 100644 index 0000000000..b6c5b8c7d0 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_ppm4gpu_integer_flux.py @@ -0,0 +1,78 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import abs, astype, floor, where + +from icon4py.model.atmosphere.advection.stencils.compute_ppm4gpu_fractional_flux import ( + _sum_neighbor_contributions, +) +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.type_alias import wpfloat + + +# TODO (dastrm): this stencil has no test +# TODO (dastrm): this stencil does not strictly match the fortran code + + +@gtx.field_operator +def _compute_ppm4gpu_integer_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + js = floor(abs(z_cfl)) - 1.0 + + z_cfl_pos = z_cfl > 0.0 + z_cfl_neg = not z_cfl_pos + wsign = where(z_cfl_pos, 1.0, -1.0) + + in_slev_bounds = astype(k, wpfloat) - js >= astype(slev, wpfloat) + + p_cc_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cc) + p_cellmass_now_jks = _sum_neighbor_contributions(z_cfl_pos, z_cfl_neg, js, p_cellmass_now) + + z_iflx = wsign * p_cc_jks * p_cellmass_now_jks + + p_upflux = p_upflux + where(in_slev_bounds, z_iflx / p_dtime, 0.0) + + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_ppm4gpu_integer_flux( + p_cc: fa.CellKField[ta.wpfloat], + p_cellmass_now: fa.CellKField[ta.wpfloat], + z_cfl: fa.CellKField[ta.wpfloat], + p_upflux: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + slev: gtx.int32, + p_dtime: ta.wpfloat, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_ppm4gpu_integer_flux( + p_cc, + p_cellmass_now, + z_cfl, + p_upflux, + k, + slev, + p_dtime, + out=p_upflux, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py new file mode 100644 index 0000000000..7e4a55d7ab --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/compute_vertical_tracer_flux_upwind.py @@ -0,0 +1,43 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +@gtx.field_operator +def _compute_vertical_tracer_flux_upwind( + p_cc: fa.CellKField[ta.wpfloat], + p_mflx_contra_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim +) -> fa.CellKField[ta.wpfloat]: + p_upflux = where(p_mflx_contra_v >= 0.0, p_cc, p_cc(Koff[-1])) * p_mflx_contra_v + return p_upflux + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def compute_vertical_tracer_flux_upwind( + p_cc: fa.CellKField[ta.wpfloat], + p_mflx_contra_v: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + p_upflux: fa.CellKField[ta.wpfloat], # TODO (dastrm): should be KHalfDim + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _compute_vertical_tracer_flux_upwind( + p_cc, + p_mflx_contra_v, + out=p_upflux, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py new file mode 100644 index 0000000000..b3b296deed --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/copy_cell_kdim_field_koff_plus1.py @@ -0,0 +1,42 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +# TODO (dastrm): move this highly generic stencil to common +# TODO (dastrm): this stencil has no test + + +@gtx.field_operator +def _copy_cell_kdim_field_koff_plus1( + field_in: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + return field_in(Koff[1]) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def copy_cell_kdim_field_koff_plus1( + field_in: fa.CellKField[ta.wpfloat], + field_out: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _copy_cell_kdim_field_koff_plus1( + field_in, + out=field_out, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py new file mode 100644 index 0000000000..4737343fa0 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/limit_vertical_slope_semi_monotonically.py @@ -0,0 +1,51 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next.ffront.fbuiltins import abs, minimum, where + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.dimension import Koff + + +@gtx.field_operator +def _limit_vertical_slope_semi_monotonically( + p_cc: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + elev: gtx.int32, +) -> fa.CellKField[ta.wpfloat]: + p_cc_min_last = minimum(p_cc(Koff[-1]), p_cc) + p_cc_min = where(k == elev, p_cc_min_last, minimum(p_cc_min_last, p_cc(Koff[1]))) + slope_l = minimum(abs(z_slope), 2.0 * (p_cc - p_cc_min)) + slope = where(z_slope >= 0.0, slope_l, -slope_l) + return slope + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def limit_vertical_slope_semi_monotonically( + p_cc: fa.CellKField[ta.wpfloat], + z_slope: fa.CellKField[ta.wpfloat], + k: fa.KField[gtx.int32], + elev: gtx.int32, + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _limit_vertical_slope_semi_monotonically( + p_cc, + z_slope, + k, + elev, + out=z_slope, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_horizontal_density_increment.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_horizontal_density_increment.py index 33f793cb94..2492b7f8f6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_horizontal_density_increment.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_horizontal_density_increment.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ apply_horizontal_density_increment, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestApplyHorizontalDensityIncrement(helpers.StencilTest): @@ -30,15 +30,15 @@ class TestApplyHorizontalDensityIncrement(helpers.StencilTest): @staticmethod def reference( grid, - p_rhodz_new: xp.array, - p_mflx_contra_v: xp.array, - deepatmo_divzl: xp.array, - deepatmo_divzu: xp.array, + p_rhodz_new: np.array, + p_mflx_contra_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, p_dtime: float, **kwargs, ) -> dict: tmp = p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu - rhodz_ast2 = xp.maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * tmp + rhodz_ast2 = np.maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * tmp return dict(rhodz_ast2=rhodz_ast2) @pytest.fixture diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_interpolated_tracer_time_tendency.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_interpolated_tracer_time_tendency.py index 778a5db7f0..0526018051 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_interpolated_tracer_time_tendency.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_interpolated_tracer_time_tendency.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ apply_interpolated_tracer_time_tendency, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestApplyInterpolatedTracerTimeTendency(helpers.StencilTest): @@ -24,13 +24,13 @@ class TestApplyInterpolatedTracerTimeTendency(helpers.StencilTest): @staticmethod def reference( grid, - p_tracer_now: xp.array, - p_grf_tend_tracer: xp.array, + p_tracer_now: np.array, + p_grf_tend_tracer: np.array, p_dtime, **kwargs, ) -> dict: p_tracer_new = p_tracer_now + p_dtime * p_grf_tend_tracer - p_tracer_new = xp.where(p_tracer_new < 0.0, 0.0, p_tracer_new) + p_tracer_new = np.where(p_tracer_new < 0.0, 0.0, p_tracer_new) return dict(p_tracer_new=p_tracer_new) @@ -39,7 +39,7 @@ def input_data(self, grid) -> dict: p_tracer_now = helpers.random_field(grid, dims.CellDim, dims.KDim) p_grf_tend_tracer = helpers.random_field(grid, dims.CellDim, dims.KDim) p_tracer_new = helpers.random_field(grid, dims.CellDim, dims.KDim) - p_dtime = xp.float64(5.0) + p_dtime = np.float64(5.0) return dict( p_tracer_now=p_tracer_now, p_grf_tend_tracer=p_grf_tend_tracer, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_vertical_density_increment.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_vertical_density_increment.py index 434ff83c37..dc512a74bd 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_vertical_density_increment.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_apply_vertical_density_increment.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ apply_vertical_density_increment, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestApplyVerticalDensityIncrement(helpers.StencilTest): @@ -30,10 +30,10 @@ class TestApplyVerticalDensityIncrement(helpers.StencilTest): @staticmethod def reference( grid, - rhodz_ast: xp.array, - p_mflx_contra_v: xp.array, - deepatmo_divzl: xp.array, - deepatmo_divzu: xp.array, + rhodz_ast: np.array, + p_mflx_contra_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, p_dtime, **kwargs, ) -> dict: diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_2.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_2.py index bd1b5ba5ef..c4b0da2ba6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_2.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_2.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ average_horizontal_flux_subcycling_2, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestAverageHorizontalFluxSubcycling2(helpers.StencilTest): @@ -24,8 +24,8 @@ class TestAverageHorizontalFluxSubcycling2(helpers.StencilTest): @staticmethod def reference( grid, - z_tracer_mflx_1_dsl: xp.array, - z_tracer_mflx_2_dsl: xp.array, + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, **kwargs, ) -> dict: p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float(2) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_3.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_3.py index 2efcd635d7..44873bd564 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_3.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_average_horizontal_flux_subcycling_3.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ average_horizontal_flux_subcycling_3, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestAverageHorizontalFluxSubcycling3(helpers.StencilTest): @@ -24,9 +24,9 @@ class TestAverageHorizontalFluxSubcycling3(helpers.StencilTest): @staticmethod def reference( grid, - z_tracer_mflx_1_dsl: xp.array, - z_tracer_mflx_2_dsl: xp.array, - z_tracer_mflx_3_dsl: xp.array, + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, + z_tracer_mflx_3_dsl: np.array, **kwargs, ) -> dict: p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py index 7b0c393770..490b4bb0e9 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_antidiffusive_cell_fluxes_and_min_max.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_antidiffusive_cell_fluxes_and_min_max, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeAntidiffusiveCellFluxesAndMinMax(helpers.StencilTest): @@ -30,12 +30,12 @@ class TestComputeAntidiffusiveCellFluxesAndMinMax(helpers.StencilTest): @staticmethod def reference( grid, - geofac_div: xp.ndarray, - p_rhodz_now: xp.ndarray, - p_rhodz_new: xp.ndarray, - z_mflx_low: xp.ndarray, - z_anti: xp.ndarray, - p_cc: xp.ndarray, + geofac_div: np.ndarray, + p_rhodz_now: np.ndarray, + p_rhodz_new: np.ndarray, + z_mflx_low: np.ndarray, + z_anti: np.ndarray, + p_cc: np.ndarray, p_dtime: float, **kwargs, ) -> dict: @@ -43,31 +43,31 @@ def reference( z_anti_c2e = z_anti[c2e] geofac_div = helpers.reshape(geofac_div, c2e.shape) - geofac_div = xp.expand_dims(geofac_div, axis=-1) + geofac_div = np.expand_dims(geofac_div, axis=-1) - zero_array = xp.zeros(p_rhodz_now.shape) + zero_array = np.zeros(p_rhodz_now.shape) z_mflx_anti_1 = p_dtime * geofac_div[:, 0] / p_rhodz_new * z_anti_c2e[:, 0] z_mflx_anti_2 = p_dtime * geofac_div[:, 1] / p_rhodz_new * z_anti_c2e[:, 1] z_mflx_anti_3 = p_dtime * geofac_div[:, 2] / p_rhodz_new * z_anti_c2e[:, 2] z_mflx_anti_in = -1.0 * ( - xp.minimum(zero_array, z_mflx_anti_1) - + xp.minimum(zero_array, z_mflx_anti_2) - + xp.minimum(zero_array, z_mflx_anti_3) + np.minimum(zero_array, z_mflx_anti_1) + + np.minimum(zero_array, z_mflx_anti_2) + + np.minimum(zero_array, z_mflx_anti_3) ) z_mflx_anti_out = ( - xp.maximum(zero_array, z_mflx_anti_1) - + xp.maximum(zero_array, z_mflx_anti_2) - + xp.maximum(zero_array, z_mflx_anti_3) + np.maximum(zero_array, z_mflx_anti_1) + + np.maximum(zero_array, z_mflx_anti_2) + + np.maximum(zero_array, z_mflx_anti_3) ) - z_fluxdiv_c = xp.sum(z_mflx_low[c2e] * geofac_div, axis=1) + z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new - z_tracer_max = xp.maximum(p_cc, z_tracer_new_low) - z_tracer_min = xp.minimum(p_cc, z_tracer_new_low) + z_tracer_max = np.maximum(p_cc, z_tracer_new_low) + z_tracer_min = np.minimum(p_cc, z_tracer_new_low) return dict( z_mflx_anti_in=z_mflx_anti_in, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py index bda355f300..ad8d7757a6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_barycentric_backtrajectory, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeBarycentricBacktrajectory(helpers.StencilTest): @@ -24,16 +24,16 @@ class TestComputeBarycentricBacktrajectory(helpers.StencilTest): @staticmethod def reference( grid, - p_vn: xp.array, - p_vt: xp.array, - cell_idx: xp.array, - cell_blk: xp.array, - pos_on_tplane_e_1: xp.array, - pos_on_tplane_e_2: xp.array, - primal_normal_cell_1: xp.array, - dual_normal_cell_1: xp.array, - primal_normal_cell_2: xp.array, - dual_normal_cell_2: xp.array, + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + pos_on_tplane_e_1: np.array, + pos_on_tplane_e_2: np.array, + primal_normal_cell_1: np.array, + dual_normal_cell_1: np.array, + primal_normal_cell_2: np.array, + dual_normal_cell_2: np.array, p_dthalf: float, **kwargs, ) -> dict: @@ -48,27 +48,27 @@ def reference( dual_normal_cell_2 = dual_normal_cell_2.reshape(e2c.shape) lvn_pos = p_vn >= 0.0 - cell_idx = xp.expand_dims(cell_idx, axis=-1) - cell_blk = xp.expand_dims(cell_blk, axis=-1) - pos_on_tplane_e_1 = xp.expand_dims(pos_on_tplane_e_1, axis=-1) - pos_on_tplane_e_2 = xp.expand_dims(pos_on_tplane_e_2, axis=-1) - primal_normal_cell_1 = xp.expand_dims(primal_normal_cell_1, axis=-1) - dual_normal_cell_1 = xp.expand_dims(dual_normal_cell_1, axis=-1) - primal_normal_cell_2 = xp.expand_dims(primal_normal_cell_2, axis=-1) - dual_normal_cell_2 = xp.expand_dims(dual_normal_cell_2, axis=-1) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + pos_on_tplane_e_1 = np.expand_dims(pos_on_tplane_e_1, axis=-1) + pos_on_tplane_e_2 = np.expand_dims(pos_on_tplane_e_2, axis=-1) + primal_normal_cell_1 = np.expand_dims(primal_normal_cell_1, axis=-1) + dual_normal_cell_1 = np.expand_dims(dual_normal_cell_1, axis=-1) + primal_normal_cell_2 = np.expand_dims(primal_normal_cell_2, axis=-1) + dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) - p_cell_idx = xp.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) - p_cell_rel_idx_dsl = xp.where(lvn_pos, 0, 1) - p_cell_blk = xp.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, 0, 1) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + xp.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + xp.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) - p_distv_bary_1 = xp.where( + p_distv_bary_1 = np.where( lvn_pos, z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], @@ -76,7 +76,7 @@ def reference( + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) - p_distv_bary_2 = xp.where( + p_distv_bary_2 = np.where( lvn_pos, z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], @@ -96,7 +96,7 @@ def reference( def input_data(self, grid) -> dict: p_vn = helpers.random_field(grid, dims.EdgeDim, dims.KDim) p_vt = helpers.random_field(grid, dims.EdgeDim, dims.KDim) - cell_idx = xp.asarray(grid.connectivities[dims.E2CDim], dtype=gtx.int32) + cell_idx = np.asarray(grid.connectivities[dims.E2CDim], dtype=gtx.int32) cell_idx_new = helpers.numpy_to_1D_sparse_field(cell_idx, dims.ECDim) cell_blk = helpers.constant_field(grid, 1, dims.EdgeDim, dims.E2CDim, dtype=gtx.int32) cell_blk_new = helpers.as_1D_sparse_field(cell_blk, dims.ECDim) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory_alt.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory_alt.py index 6a923c8c6b..ac65277bb2 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory_alt.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_barycentric_backtrajectory_alt.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_barycentric_backtrajectory_alt, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeBarycentricBacktrajectoryAlt(helpers.StencilTest): @@ -24,14 +24,14 @@ class TestComputeBarycentricBacktrajectoryAlt(helpers.StencilTest): @staticmethod def reference( grid, - p_vn: xp.array, - p_vt: xp.array, - pos_on_tplane_e_1: xp.array, - pos_on_tplane_e_2: xp.array, - primal_normal_cell_1: xp.array, - dual_normal_cell_1: xp.array, - primal_normal_cell_2: xp.array, - dual_normal_cell_2: xp.array, + p_vn: np.array, + p_vt: np.array, + pos_on_tplane_e_1: np.array, + pos_on_tplane_e_2: np.array, + primal_normal_cell_1: np.array, + dual_normal_cell_1: np.array, + primal_normal_cell_2: np.array, + dual_normal_cell_2: np.array, p_dthalf: float, **kwargs, ) -> dict: @@ -44,21 +44,21 @@ def reference( dual_normal_cell_2 = dual_normal_cell_2.reshape(e2c.shape) lvn_pos = p_vn >= 0.0 - pos_on_tplane_e_1 = xp.expand_dims(pos_on_tplane_e_1, axis=-1) - pos_on_tplane_e_2 = xp.expand_dims(pos_on_tplane_e_2, axis=-1) - primal_normal_cell_1 = xp.expand_dims(primal_normal_cell_1, axis=-1) - dual_normal_cell_1 = xp.expand_dims(dual_normal_cell_1, axis=-1) - primal_normal_cell_2 = xp.expand_dims(primal_normal_cell_2, axis=-1) - dual_normal_cell_2 = xp.expand_dims(dual_normal_cell_2, axis=-1) + pos_on_tplane_e_1 = np.expand_dims(pos_on_tplane_e_1, axis=-1) + pos_on_tplane_e_2 = np.expand_dims(pos_on_tplane_e_2, axis=-1) + primal_normal_cell_1 = np.expand_dims(primal_normal_cell_1, axis=-1) + dual_normal_cell_1 = np.expand_dims(dual_normal_cell_1, axis=-1) + primal_normal_cell_2 = np.expand_dims(primal_normal_cell_2, axis=-1) + dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) z_ntdistv_bary_1 = -( - p_vn * p_dthalf + xp.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) + p_vn * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_1[:, 0], pos_on_tplane_e_1[:, 1]) ) z_ntdistv_bary_2 = -( - p_vt * p_dthalf + xp.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) + p_vt * p_dthalf + np.where(lvn_pos, pos_on_tplane_e_2[:, 0], pos_on_tplane_e_2[:, 1]) ) - p_distv_bary_1 = xp.where( + p_distv_bary_1 = np.where( lvn_pos, z_ntdistv_bary_1 * primal_normal_cell_1[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_1[:, 0], @@ -66,7 +66,7 @@ def reference( + z_ntdistv_bary_2 * dual_normal_cell_1[:, 1], ) - p_distv_bary_2 = xp.where( + p_distv_bary_2 = np.where( lvn_pos, z_ntdistv_bary_1 * primal_normal_cell_2[:, 0] + z_ntdistv_bary_2 * dual_normal_cell_2[:, 0], diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory.py index b3c81bf201..cb1e3aecc8 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ffsl_backtrajectory, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeFfslBacktrajectory(helpers.StencilTest): @@ -36,23 +36,23 @@ class TestComputeFfslBacktrajectory(helpers.StencilTest): @staticmethod def reference( grid, - p_vn: xp.array, - p_vt: xp.array, - cell_idx: xp.array, - cell_blk: xp.array, - edge_verts_1_x: xp.array, - edge_verts_2_x: xp.array, - edge_verts_1_y: xp.array, - edge_verts_2_y: xp.array, - pos_on_tplane_e_1_x: xp.array, - pos_on_tplane_e_2_x: xp.array, - pos_on_tplane_e_1_y: xp.array, - pos_on_tplane_e_2_y: xp.array, - primal_normal_cell_x: xp.array, - primal_normal_cell_y: xp.array, - dual_normal_cell_x: xp.array, - dual_normal_cell_y: xp.array, - lvn_sys_pos: xp.array, + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + edge_verts_1_x: np.array, + edge_verts_2_x: np.array, + edge_verts_1_y: np.array, + edge_verts_2_y: np.array, + pos_on_tplane_e_1_x: np.array, + pos_on_tplane_e_2_x: np.array, + pos_on_tplane_e_1_y: np.array, + pos_on_tplane_e_2_y: np.array, + primal_normal_cell_x: np.array, + primal_normal_cell_y: np.array, + dual_normal_cell_x: np.array, + dual_normal_cell_y: np.array, + lvn_sys_pos: np.array, p_dt: float, **kwargs, ) -> dict: @@ -65,54 +65,54 @@ def reference( dual_normal_cell_y = helpers.reshape(dual_normal_cell_y, e2c_shape) lvn_pos = p_vn >= 0.0 - cell_idx = xp.expand_dims(cell_idx, axis=-1) - cell_blk = xp.expand_dims(cell_blk, axis=-1) - primal_normal_cell_x = xp.expand_dims(primal_normal_cell_x, axis=-1) - dual_normal_cell_x = xp.expand_dims(dual_normal_cell_x, axis=-1) - primal_normal_cell_y = xp.expand_dims(primal_normal_cell_y, axis=-1) - dual_normal_cell_y = xp.expand_dims(dual_normal_cell_y, axis=-1) - edge_verts_1_x = xp.expand_dims(edge_verts_1_x, axis=-1) - edge_verts_1_y = xp.expand_dims(edge_verts_1_y, axis=-1) - edge_verts_2_x = xp.expand_dims(edge_verts_2_x, axis=-1) - edge_verts_2_y = xp.expand_dims(edge_verts_2_y, axis=-1) - pos_on_tplane_e_1_x = xp.expand_dims(pos_on_tplane_e_1_x, axis=-1) - pos_on_tplane_e_1_y = xp.expand_dims(pos_on_tplane_e_1_y, axis=-1) - pos_on_tplane_e_2_x = xp.expand_dims(pos_on_tplane_e_2_x, axis=-1) - pos_on_tplane_e_2_y = xp.expand_dims(pos_on_tplane_e_2_y, axis=-1) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + primal_normal_cell_x = np.expand_dims(primal_normal_cell_x, axis=-1) + dual_normal_cell_x = np.expand_dims(dual_normal_cell_x, axis=-1) + primal_normal_cell_y = np.expand_dims(primal_normal_cell_y, axis=-1) + dual_normal_cell_y = np.expand_dims(dual_normal_cell_y, axis=-1) + edge_verts_1_x = np.expand_dims(edge_verts_1_x, axis=-1) + edge_verts_1_y = np.expand_dims(edge_verts_1_y, axis=-1) + edge_verts_2_x = np.expand_dims(edge_verts_2_x, axis=-1) + edge_verts_2_y = np.expand_dims(edge_verts_2_y, axis=-1) + pos_on_tplane_e_1_x = np.expand_dims(pos_on_tplane_e_1_x, axis=-1) + pos_on_tplane_e_1_y = np.expand_dims(pos_on_tplane_e_1_y, axis=-1) + pos_on_tplane_e_2_x = np.expand_dims(pos_on_tplane_e_2_x, axis=-1) + pos_on_tplane_e_2_y = np.expand_dims(pos_on_tplane_e_2_y, axis=-1) - p_cell_idx = xp.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) - p_cell_blk = xp.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) - p_cell_rel_idx_dsl = xp.where(lvn_pos, 0, 1) + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, 0, 1) - depart_pts_1_x = xp.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt - depart_pts_1_y = xp.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt - depart_pts_2_x = xp.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt - depart_pts_2_y = xp.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt + depart_pts_1_x = np.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt + depart_pts_1_y = np.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt + depart_pts_2_x = np.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt + depart_pts_2_y = np.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt - pos_on_tplane_e_x = xp.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) - pos_on_tplane_e_y = xp.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + pos_on_tplane_e_x = np.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) + pos_on_tplane_e_y = np.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y pos_dreg_vert_c_2_x = ( - xp.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x ) pos_dreg_vert_c_2_y = ( - xp.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y ) pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y pos_dreg_vert_c_4_x = ( - xp.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x ) pos_dreg_vert_c_4_y = ( - xp.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y ) - pn_cell_1 = xp.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) - pn_cell_2 = xp.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) - dn_cell_1 = xp.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) - dn_cell_2 = xp.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) + pn_cell_1 = np.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) + pn_cell_2 = np.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) + dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) + dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) p_coords_dreg_v_1_lon_dsl = ( pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 @@ -156,7 +156,7 @@ def reference( def input_data(self, grid) -> dict: p_vn = helpers.random_field(grid, dims.EdgeDim, dims.KDim) p_vt = helpers.random_field(grid, dims.EdgeDim, dims.KDim) - cell_idx = xp.asarray(grid.connectivities[dims.E2CDim], dtype=gtx.int32) + cell_idx = np.asarray(grid.connectivities[dims.E2CDim], dtype=gtx.int32) cell_idx_new = helpers.numpy_to_1D_sparse_field(cell_idx, dims.ECDim) cell_blk = helpers.constant_field(grid, 1, dims.EdgeDim, dims.E2CDim, dtype=gtx.int32) cell_blk_new = helpers.as_1D_sparse_field(cell_blk, dims.ECDim) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py index 13e112b2b0..42829df90a 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_counterclockwise_indicator.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ffsl_backtrajectory_counterclockwise_indicator, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeFfslBacktrajectoryCounterclockwiseIndicator(helpers.StencilTest): @@ -23,17 +23,17 @@ class TestComputeFfslBacktrajectoryCounterclockwiseIndicator(helpers.StencilTest @staticmethod def reference( - grid, p_vn: xp.array, tangent_orientation: xp.array, lcounterclock: bool, **kwargs + grid, p_vn: np.array, tangent_orientation: np.array, lcounterclock: bool, **kwargs ) -> dict: - tangent_orientation = xp.expand_dims(tangent_orientation, axis=-1) + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - tangent_orientation = xp.broadcast_to(tangent_orientation, p_vn.shape) + tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) - lvn_sys_pos_true = xp.where(tangent_orientation * p_vn >= 0.0, True, False) + lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False) - mask_lcounterclock = xp.broadcast_to(lcounterclock, p_vn.shape) + mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) - lvn_sys_pos = xp.where(mask_lcounterclock, lvn_sys_pos_true, False) + lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) return dict(lvn_sys_pos=lvn_sys_pos) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py index 302cbb0f1c..94efbdf63d 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ffsl_backtrajectory_length_indicator.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ffsl_backtrajectory_length_indicator, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeFfslBacktrajectoryLengthIndicator(helpers.StencilTest): @@ -22,18 +22,18 @@ class TestComputeFfslBacktrajectoryLengthIndicator(helpers.StencilTest): OUTPUTS = ("opt_famask_dsl",) @staticmethod - def reference(grid, p_vn: xp.array, p_vt: xp.array, p_dt: float, **kwargs) -> dict: + def reference(grid, p_vn: np.array, p_vt: np.array, p_dt: float, **kwargs) -> dict: lvn_pos = p_vn >= 0.0 - traj_length = xp.sqrt(p_vn**2 + p_vt**2) * p_dt + traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt - edge_cell_length = xp.expand_dims( - xp.asarray(grid.connectivities[dims.E2CDim], dtype=float), axis=-1 + edge_cell_length = np.expand_dims( + np.asarray(grid.connectivities[dims.E2CDim], dtype=float), axis=-1 ) - e2c_length = xp.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) + e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) - opt_famask_dsl = xp.where( - traj_length > (1.25 * xp.broadcast_to(e2c_length, p_vn.shape)), + opt_famask_dsl = np.where( + traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), 1, 0, ) @@ -44,7 +44,7 @@ def reference(grid, p_vn: xp.array, p_vt: xp.array, p_dt: float, **kwargs) -> di def input_data(self, grid) -> dict: p_vn = helpers.random_field(grid, dims.EdgeDim, dims.KDim) p_vt = helpers.random_field(grid, dims.EdgeDim, dims.KDim) - edge_cell_length = xp.asarray(grid.connectivities[dims.E2CDim], dtype=float) + edge_cell_length = np.asarray(grid.connectivities[dims.E2CDim], dtype=float) edge_cell_length_new = helpers.numpy_to_1D_sparse_field(edge_cell_length, dims.ECDim) opt_famask_dsl = helpers.zero_field(grid, dims.EdgeDim, dims.KDim, dtype=gtx.int32) p_dt = 1.0 diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py index 5b30b04957..9984ebe88c 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_cubic_coefficients.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_horizontal_tracer_flux_from_cubic_coefficients, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeHorizontalTracerFluxFromCubicCoefficients(helpers.StencilTest): @@ -24,9 +24,9 @@ class TestComputeHorizontalTracerFluxFromCubicCoefficients(helpers.StencilTest): @staticmethod def reference( grid, - p_out_e_hybrid_2: xp.array, - p_mass_flx_e: xp.array, - z_dreg_area: xp.array, + p_out_e_hybrid_2: np.array, + p_mass_flx_e: np.array, + z_dreg_area: np.array, **kwargs, ) -> dict: p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py index f5aa12f44b..c26ba6212e 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -15,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import horizontal as h_grid -from icon4py.model.common.settings import xp class TestComputeHorizontalTracerFluxFromLinearCoefficients(helpers.StencilTest): @@ -25,14 +25,14 @@ class TestComputeHorizontalTracerFluxFromLinearCoefficients(helpers.StencilTest) @staticmethod def reference( grid, - z_lsq_coeff_1: xp.array, - z_lsq_coeff_2: xp.array, - z_lsq_coeff_3: xp.array, - distv_bary_1: xp.array, - distv_bary_2: xp.array, - p_mass_flx_e: xp.array, - cell_rel_idx_dsl: xp.array, - p_out_e: xp.array, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, + p_out_e: np.array, **kwargs, ) -> dict: p_out_e_cp = p_out_e.copy() @@ -42,11 +42,11 @@ def reference( z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] p_out_e = ( - xp.where(cell_rel_idx_dsl == 1, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) + np.where(cell_rel_idx_dsl == 1, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) + distv_bary_1 - * xp.where(cell_rel_idx_dsl == 1, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) + * np.where(cell_rel_idx_dsl == 1, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) + distv_bary_2 - * xp.where(cell_rel_idx_dsl == 1, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) + * np.where(cell_rel_idx_dsl == 1, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) ) * p_mass_flx_e # restriction of execution domain diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py index 6bca883f67..21725ce7e6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_horizontal_tracer_flux_from_linear_coefficients_alt.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -15,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import horizontal as h_grid -from icon4py.model.common.settings import xp class TestComputeHorizontalTracerFluxFromLinearCoefficientsAlt(helpers.StencilTest): @@ -25,14 +25,14 @@ class TestComputeHorizontalTracerFluxFromLinearCoefficientsAlt(helpers.StencilTe @staticmethod def reference( grid, - z_lsq_coeff_1: xp.array, - z_lsq_coeff_2: xp.array, - z_lsq_coeff_3: xp.array, - distv_bary_1: xp.array, - distv_bary_2: xp.array, - p_mass_flx_e: xp.array, - p_vn: xp.array, - p_out_e: xp.array, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + p_vn: np.array, + p_out_e: np.array, **kwargs, ) -> dict: p_out_e_cp = p_out_e.copy() @@ -44,9 +44,9 @@ def reference( lvn_pos_inv = p_vn < 0.0 p_out_e = ( - xp.where(lvn_pos_inv, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) - + distv_bary_1 * xp.where(lvn_pos_inv, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) - + distv_bary_2 * xp.where(lvn_pos_inv, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) + np.where(lvn_pos_inv, z_lsq_coeff_1_e2c[:, 1], z_lsq_coeff_1_e2c[:, 0]) + + distv_bary_1 * np.where(lvn_pos_inv, z_lsq_coeff_2_e2c[:, 1], z_lsq_coeff_2_e2c[:, 0]) + + distv_bary_2 * np.where(lvn_pos_inv, z_lsq_coeff_3_e2c[:, 1], z_lsq_coeff_3_e2c[:, 0]) ) * p_mass_flx_e # restriction of execution domain diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py index 93e9cb70cb..b3d15bfc20 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_positive_definite_horizontal_multiplicative_flux_factor.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_positive_definite_horizontal_multiplicative_flux_factor, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePositiveDefiniteHorizontalMultiplicativeFluxFactor(helpers.StencilTest): @@ -24,31 +24,31 @@ class TestComputePositiveDefiniteHorizontalMultiplicativeFluxFactor(helpers.Sten @staticmethod def reference( grid, - geofac_div: xp.ndarray, - p_cc: xp.ndarray, - p_rhodz_now: xp.ndarray, - p_mflx_tracer_h: xp.ndarray, + geofac_div: np.ndarray, + p_cc: np.ndarray, + p_rhodz_now: np.ndarray, + p_mflx_tracer_h: np.ndarray, p_dtime, dbl_eps, **kwargs, ) -> dict: geofac_div = helpers.reshape(geofac_div, grid.connectivities[dims.C2EDim].shape) - geofac_div = xp.expand_dims(geofac_div, axis=-1) - p_m_0 = xp.maximum( + geofac_div = np.expand_dims(geofac_div, axis=-1) + p_m_0 = np.maximum( 0.0, p_mflx_tracer_h[grid.connectivities[dims.C2EDim][:, 0]] * geofac_div[:, 0] * p_dtime, ) - p_m_1 = xp.maximum( + p_m_1 = np.maximum( 0.0, p_mflx_tracer_h[grid.connectivities[dims.C2EDim][:, 1]] * geofac_div[:, 1] * p_dtime, ) - p_m_2 = xp.maximum( + p_m_2 = np.maximum( 0.0, p_mflx_tracer_h[grid.connectivities[dims.C2EDim][:, 2]] * geofac_div[:, 2] * p_dtime, ) p_m = p_m_0 + p_m_1 + p_m_2 - r_m = xp.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) + r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) return dict(r_m=r_m) @@ -60,8 +60,8 @@ def input_data(self, grid) -> dict: p_rhodz_now = helpers.random_field(grid, dims.CellDim, dims.KDim) p_mflx_tracer_h = helpers.random_field(grid, dims.EdgeDim, dims.KDim) r_m = helpers.zero_field(grid, dims.CellDim, dims.KDim) - p_dtime = xp.float64(5) - dbl_eps = xp.float64(1e-9) + p_dtime = np.float64(5) + dbl_eps = np.float64(1e-9) return dict( geofac_div=geofac_div_new, p_cc=p_cc, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py index bcf8c0569f..b1455b6994 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm4gpu_parabola_coefficients.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ppm4gpu_parabola_coefficients, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePpm4gpuParabolaCoefficients(helpers.StencilTest): @@ -23,7 +23,7 @@ class TestComputePpm4gpuParabolaCoefficients(helpers.StencilTest): @staticmethod def reference( - grid, z_face_up: xp.ndarray, z_face_low: xp.ndarray, p_cc: xp.ndarray, **kwargs + grid, z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray, **kwargs ) -> dict: z_delta_q = 0.5 * (z_face_up - z_face_low) z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_all_face_values.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_all_face_values.py index 6483e06466..e6f25381dd 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_all_face_values.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_all_face_values.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest from gt4py.next import as_field @@ -15,7 +16,6 @@ compute_ppm_all_face_values, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePpmAllFaceValues(helpers.StencilTest): @@ -25,10 +25,10 @@ class TestComputePpmAllFaceValues(helpers.StencilTest): @staticmethod def reference( grid, - p_cc: xp.array, - p_cellhgt_mc_now: xp.array, - p_face_in: xp.array, - k: xp.array, + p_cc: np.array, + p_cellhgt_mc_now: np.array, + p_face_in: np.array, + k: np.array, slev: gtx.int32, elev: gtx.int32, slevp1: gtx.int32, @@ -42,9 +42,9 @@ def reference( (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1] ) - p_face = xp.where((k == slevp1) | (k == elev), p_face_a, p_face_in) - p_face = xp.where((k == slev), p_cc, p_face) - p_face[:, 1:] = xp.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) + p_face = np.where((k == slevp1) | (k == elev), p_face_a, p_face_in) + p_face = np.where((k == slev), p_cc, p_face) + p_face[:, 1:] = np.where((k[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) return dict(p_face=p_face) @pytest.fixture @@ -55,7 +55,7 @@ def input_data(self, grid) -> dict: p_face = helpers.zero_field(grid, dims.CellDim, dims.KDim) k = as_field( - (dims.KDim,), xp.arange(0, helpers._shape(grid, dims.KDim)[0], dtype=gtx.int32) + (dims.KDim,), np.arange(0, helpers._shape(grid, dims.KDim)[0], dtype=gtx.int32) ) slev = gtx.int32(1) slevp1 = gtx.int32(2) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quadratic_face_values.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quadratic_face_values.py index f7aef94aaf..7ed3c06254 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quadratic_face_values.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quadratic_face_values.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ppm_quadratic_face_values, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp outslice = (slice(None), slice(1, None)) @@ -25,7 +25,7 @@ class TestComputePpmQuadraticFaceValues(helpers.StencilTest): OUTPUTS = (helpers.Output("p_face", refslice=outslice, gtslice=outslice),) @staticmethod - def reference(grid, p_cc: xp.array, p_cellhgt_mc_now: xp.array, **kwargs) -> dict: + def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, **kwargs) -> dict: p_face = p_cc.copy() p_face[:, 1:] = p_cc[:, 1:] * ( 1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quartic_face_values.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quartic_face_values.py index 8be8b6ec73..dc65a506ce 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quartic_face_values.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_quartic_face_values.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_ppm_quartic_face_values, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePpmQuarticFaceValues(helpers.StencilTest): @@ -23,7 +23,7 @@ class TestComputePpmQuarticFaceValues(helpers.StencilTest): @staticmethod def reference( - grid, p_cc: xp.array, p_cellhgt_mc_now: xp.array, z_slope: xp.array, **kwargs + grid, p_cc: np.array, p_cellhgt_mc_now: np.array, z_slope: np.array, **kwargs ) -> dict: p_cellhgt_mc_now_k_minus_1 = p_cellhgt_mc_now[:, 1:-2] p_cellhgt_mc_now_k_minus_2 = p_cellhgt_mc_now[:, 0:-3] diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_slope.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_slope.py index b9139805ef..f1b920944f 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_slope.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_ppm_slope.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest from gt4py.next import as_field @@ -15,7 +16,6 @@ compute_ppm_slope, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputePpmSlope(helpers.StencilTest): @@ -28,7 +28,7 @@ class TestComputePpmSlope(helpers.StencilTest): @staticmethod def reference( - grid, p_cc: xp.array, p_cellhgt_mc_now: xp.array, k: xp.array, elev: gtx.int32, **kwargs + grid, p_cc: np.array, p_cellhgt_mc_now: np.array, k: np.array, elev: gtx.int32, **kwargs ) -> dict: zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] @@ -56,7 +56,7 @@ def reference( + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 ) - z_slope = xp.where(k[1:-1] < elev, z_slope_a, z_slope_b) + z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b) return dict(z_slope=z_slope) @pytest.fixture @@ -68,7 +68,7 @@ def input_data(self, grid) -> dict: ) k = as_field( (dims.KDim,), - xp.arange( + np.arange( 0, helpers._shape(grid, dims.KDim, extend={dims.KDim: 1})[0], dtype=gtx.int32 ), ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_tendency.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_tendency.py index 30fd8f84f5..96179b99e6 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_tendency.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_tendency.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_tendency, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeTendency(helpers.StencilTest): @@ -24,8 +24,8 @@ class TestComputeTendency(helpers.StencilTest): @staticmethod def reference( grid, - p_tracer_now: xp.array, - p_tracer_new: xp.array, + p_tracer_now: np.array, + p_tracer_new: np.array, p_dtime, **kwargs, ) -> dict: @@ -38,7 +38,7 @@ def input_data(self, grid) -> dict: p_tracer_now = helpers.random_field(grid, dims.CellDim, dims.KDim) p_tracer_new = helpers.random_field(grid, dims.CellDim, dims.KDim) opt_ddt_tracer_adv = helpers.zero_field(grid, dims.CellDim, dims.KDim) - p_dtime = xp.float64(5.0) + p_dtime = np.float64(5.0) return dict( p_tracer_now=p_tracer_now, p_tracer_new=p_tracer_new, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_parabola_limiter_condition.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_parabola_limiter_condition.py index 56d27c46fb..54e87d6f33 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_parabola_limiter_condition.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_parabola_limiter_condition.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ compute_vertical_parabola_limiter_condition, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestComputeVerticalParabolaLimiterCondition(helpers.StencilTest): @@ -22,10 +22,10 @@ class TestComputeVerticalParabolaLimiterCondition(helpers.StencilTest): OUTPUTS = ("l_limit",) @staticmethod - def reference(grid, p_face: xp.array, p_cc: xp.array, **kwargs) -> dict: + def reference(grid, p_face: np.array, p_cc: np.array, **kwargs) -> dict: z_delta = p_face[:, :-1] - p_face[:, 1:] z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) - l_limit = xp.where(xp.abs(z_delta) < -1 * z_a6i, 1, 0) + l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, 1, 0) return dict(l_limit=l_limit) @pytest.fixture diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py new file mode 100644 index 0000000000..f0e11e9a63 --- /dev/null +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_compute_vertical_tracer_flux_upwind.py @@ -0,0 +1,58 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +import numpy as np +import pytest + +import icon4py.model.common.test_utils.helpers as helpers +from icon4py.model.atmosphere.advection.stencils.compute_vertical_tracer_flux_upwind import ( + compute_vertical_tracer_flux_upwind, +) +from icon4py.model.common import dimension as dims + + +outslice = (slice(None), slice(1, None)) + + +class TestComputeVerticalTracerFluxUpwind(helpers.StencilTest): + PROGRAM = compute_vertical_tracer_flux_upwind + OUTPUTS = (helpers.Output("p_upflux", refslice=outslice, gtslice=outslice),) + + @staticmethod + def reference( + grid, + p_cc: np.array, + p_mflx_contra_v: np.array, + **kwargs, + ) -> dict: + p_upflux = p_cc.copy() + p_upflux[:, 1:] = ( + np.where(p_mflx_contra_v[:, 1:] >= 0.0, p_cc[:, 1:], p_cc[:, :-1]) + * p_mflx_contra_v[:, 1:] + ) + return dict(p_upflux=p_upflux) + + @pytest.fixture + def input_data(self, grid) -> dict: + p_cc = helpers.random_field(grid, dims.CellDim, dims.KDim) + p_mflx_contra_v = helpers.random_field( + grid, dims.CellDim, dims.KDim + ) # TODO (dastrm): should be KHalfDim + p_upflux = helpers.zero_field( + grid, dims.CellDim, dims.KDim + ) # TODO (dastrm): should be KHalfDim + return dict( + p_cc=p_cc, + p_mflx_contra_v=p_mflx_contra_v, + p_upflux=p_upflux, + horizontal_start=0, + horizontal_end=gtx.int32(grid.num_cells), + vertical_start=1, + vertical_end=gtx.int32(grid.num_levels), + ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_density_horizontally.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_density_horizontally.py index feb60a15e7..9ad0a6526b 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_density_horizontally.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_density_horizontally.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ integrate_tracer_density_horizontally, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestIntegrateTracerDensityHorizontally(helpers.StencilTest): @@ -29,25 +29,25 @@ class TestIntegrateTracerDensityHorizontally(helpers.StencilTest): @staticmethod def reference( grid, - p_mass_flx_e: xp.array, - geofac_div: xp.array, - z_rhofluxdiv_c: xp.array, - z_tracer_mflx: xp.array, - z_rho_now: xp.array, - z_tracer_now: xp.array, + p_mass_flx_e: np.array, + geofac_div: np.array, + z_rhofluxdiv_c: np.array, + z_tracer_mflx: np.array, + z_rho_now: np.array, + z_tracer_now: np.array, z_dtsub: float, nsub: gtx.int32, **kwargs, ) -> dict: c2e = grid.connectivities[dims.C2EDim] p_mass_flx_e_c2e = p_mass_flx_e[c2e] - geofac_div = xp.expand_dims(geofac_div, axis=-1) + geofac_div = np.expand_dims(geofac_div, axis=-1) z_tracer_mflx_c2e = z_tracer_mflx[c2e] z_rhofluxdiv_c_out = ( - xp.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == 1 else z_rhofluxdiv_c + np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == 1 else z_rhofluxdiv_c ) - z_fluxdiv_c_dsl = xp.sum(z_tracer_mflx_c2e * geofac_div, axis=1) + z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_horizontally.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_horizontally.py index a71ede7b21..f270f55c83 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_horizontally.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_horizontally.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ integrate_tracer_horizontally, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestIntegrateTracerHorizontally(helpers.StencilTest): @@ -24,22 +24,22 @@ class TestIntegrateTracerHorizontally(helpers.StencilTest): @staticmethod def reference( grid, - p_mflx_tracer_h: xp.array, - deepatmo_divh: xp.array, - tracer_now: xp.array, - rhodz_now: xp.array, - rhodz_new: xp.array, - geofac_div: xp.array, + p_mflx_tracer_h: np.array, + deepatmo_divh: np.array, + tracer_now: np.array, + rhodz_now: np.array, + rhodz_new: np.array, + geofac_div: np.array, p_dtime, **kwargs, ) -> dict: geofac_div = helpers.reshape(geofac_div, grid.connectivities[dims.C2EDim].shape) - geofac_div = xp.expand_dims(geofac_div, axis=-1) + geofac_div = np.expand_dims(geofac_div, axis=-1) tracer_new_hor = ( tracer_now * rhodz_now - p_dtime * deepatmo_divh - * xp.sum(p_mflx_tracer_h[grid.connectivities[dims.C2EDim]] * geofac_div, axis=1) + * np.sum(p_mflx_tracer_h[grid.connectivities[dims.C2EDim]] * geofac_div, axis=1) ) / rhodz_new return dict(tracer_new_hor=tracer_new_hor) @@ -52,7 +52,7 @@ def input_data(self, grid) -> dict: rhodz_new = helpers.random_field(grid, dims.CellDim, dims.KDim) geofac_div = helpers.random_field(grid, dims.CellDim, dims.C2EDim) geofac_div_new = helpers.as_1D_sparse_field(geofac_div, dims.CEDim) - p_dtime = xp.float64(5.0) + p_dtime = np.float64(5.0) tracer_new_hor = helpers.zero_field(grid, dims.CellDim, dims.KDim) return dict( p_mflx_tracer_h=p_mflx_tracer_h, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_vertically.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_vertically.py index 82e4f2ff5c..07ca62df76 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_vertically.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_integrate_tracer_vertically.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest from gt4py.next import as_field @@ -15,7 +16,6 @@ integrate_tracer_vertically, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestIntegrateTracerVertically(helpers.StencilTest): @@ -25,20 +25,20 @@ class TestIntegrateTracerVertically(helpers.StencilTest): @staticmethod def reference( grid, - tracer_now: xp.array, - rhodz_now: xp.array, - p_mflx_tracer_v: xp.array, - deepatmo_divzl: xp.array, - deepatmo_divzu: xp.array, - rhodz_new: xp.array, - k: xp.array, + tracer_now: np.array, + rhodz_now: np.array, + p_mflx_tracer_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + rhodz_new: np.array, + k: np.array, ivadv_tracer: gtx.int32, iadv_slev_jt: gtx.int32, p_dtime: float, **kwargs, ) -> dict: if ivadv_tracer != 0: - tracer_new = xp.where( + tracer_new = np.where( (iadv_slev_jt <= k), ( tracer_now * rhodz_now @@ -65,8 +65,8 @@ def input_data(self, grid) -> dict: deepatmo_divzu = helpers.random_field(grid, dims.KDim) rhodz_new = helpers.random_field(grid, dims.CellDim, dims.KDim) tracer_new = helpers.zero_field(grid, dims.CellDim, dims.KDim) - k = as_field((dims.KDim,), xp.arange(grid.num_levels, dtype=gtx.int32)) - p_dtime = xp.float64(5.0) + k = as_field((dims.KDim,), np.arange(grid.num_levels, dtype=gtx.int32)) + p_dtime = np.float64(5.0) ivadv_tracer = 1 iadv_slev_jt = 4 return dict( diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_parabola_semi_monotonically.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_parabola_semi_monotonically.py index c587c51c9b..e1fc0066f9 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_parabola_semi_monotonically.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_parabola_semi_monotonically.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ limit_vertical_parabola_semi_monotonically, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestLimitVerticalParabolaSemiMonotonically(helpers.StencilTest): @@ -22,13 +22,13 @@ class TestLimitVerticalParabolaSemiMonotonically(helpers.StencilTest): OUTPUTS = ("p_face_up", "p_face_low") @staticmethod - def reference(grid, l_limit: xp.array, p_face: xp.array, p_cc: xp.array, **kwargs) -> dict: - q_face_up, q_face_low = xp.where( + def reference(grid, l_limit: np.array, p_face: np.array, p_cc: np.array, **kwargs) -> dict: + q_face_up, q_face_low = np.where( l_limit != 0, - xp.where( - (p_cc < xp.minimum(p_face[:, :-1], p_face[:, 1:])), + np.where( + (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), (p_cc, p_cc), - xp.where( + np.where( p_face[:, :-1] > p_face[:, 1:], (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py new file mode 100644 index 0000000000..5db041d352 --- /dev/null +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_limit_vertical_slope_semi_monotonically.py @@ -0,0 +1,52 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +import numpy as np +import pytest +from gt4py.next import as_field + +import icon4py.model.common.test_utils.helpers as helpers +from icon4py.model.atmosphere.advection.stencils.limit_vertical_slope_semi_monotonically import ( + limit_vertical_slope_semi_monotonically, +) +from icon4py.model.common import dimension as dims + + +class TestLimitVerticalSlopeSemiMonotonically(helpers.StencilTest): + PROGRAM = limit_vertical_slope_semi_monotonically + OUTPUTS = (helpers.Output("z_slope", gtslice=(slice(None), slice(1, -1))),) + + @staticmethod + def reference( + grid, p_cc: np.array, z_slope: np.array, k: np.array, elev: gtx.int32, **kwargs + ) -> dict: + p_cc_min_last = np.minimum(p_cc[:, :-2], p_cc[:, 1:-1]) + p_cc_min = np.where(k[1:-1] == elev, p_cc_min_last, np.minimum(p_cc_min_last, p_cc[:, 2:])) + slope_l = np.minimum(np.abs(z_slope[:, 1:-1]), 2.0 * (p_cc[:, 1:-1] - p_cc_min)) + slope = np.where(z_slope[:, 1:-1] >= 0.0, slope_l, -slope_l) + return dict(z_slope=slope) + + @pytest.fixture + def input_data(self, grid) -> dict: + p_cc = helpers.random_field(grid, dims.CellDim, dims.KDim) + z_slope = helpers.random_field(grid, dims.CellDim, dims.KDim) + k = as_field( + (dims.KDim,), np.arange(0, helpers._shape(grid, dims.KDim)[0], dtype=gtx.int32) + ) + elev = k[-2].as_scalar() + return dict( + p_cc=p_cc, + z_slope=z_slope, + k=k, + elev=elev, + horizontal_start=0, + horizontal_end=gtx.int32(grid.num_cells), + vertical_start=1, + vertical_end=gtx.int32(grid.num_levels - 1), + ) diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py index 3637ffa74e..fd7991126e 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_postprocess_antidiffusive_cell_fluxes_and_min_max.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ postprocess_antidiffusive_cell_fluxes_and_min_max, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp class TestPostprocessAntidiffusiveCellFluxesAndMinMax(helpers.StencilTest): @@ -24,27 +24,27 @@ class TestPostprocessAntidiffusiveCellFluxesAndMinMax(helpers.StencilTest): @staticmethod def reference( grid, - refin_ctrl: xp.ndarray, - p_cc: xp.ndarray, - z_tracer_new_low: xp.ndarray, - z_tracer_max: xp.ndarray, - z_tracer_min: xp.ndarray, + refin_ctrl: np.ndarray, + p_cc: np.ndarray, + z_tracer_new_low: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, lo_bound: float, hi_bound: float, **kwargs, ) -> dict: - refin_ctrl = xp.expand_dims(refin_ctrl, axis=1) - condition = xp.logical_or( - xp.equal(refin_ctrl, lo_bound * xp.ones(refin_ctrl.shape, dtype=gtx.int32)), - xp.equal(refin_ctrl, hi_bound * xp.ones(refin_ctrl.shape, dtype=gtx.int32)), + refin_ctrl = np.expand_dims(refin_ctrl, axis=1) + condition = np.logical_or( + np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=gtx.int32)), + np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=gtx.int32)), ) - z_tracer_new_out = xp.where( + z_tracer_new_out = np.where( condition, - xp.minimum(1.1 * p_cc, xp.maximum(0.9 * p_cc, z_tracer_new_low)), + np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), z_tracer_new_low, ) - z_tracer_max_out = xp.where(condition, xp.maximum(p_cc, z_tracer_new_out), z_tracer_max) - z_tracer_min_out = xp.where(condition, xp.minimum(p_cc, z_tracer_new_out), z_tracer_min) + z_tracer_max_out = np.where(condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max) + z_tracer_min_out = np.where(condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min) return dict( z_tracer_new_low=z_tracer_new_out, z_tracer_max=z_tracer_max_out, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_ffsl_flux_area_patches_list.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_ffsl_flux_area_patches_list.py index 5ce482bd96..143aaaba60 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_ffsl_flux_area_patches_list.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_ffsl_flux_area_patches_list.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ prepare_ffsl_flux_area_patches_list, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp # Check whether lines inters. @@ -35,8 +35,8 @@ def _ccw_numpy( dx1dy2 = dx1 * dy2 dy1dx2 = dy1 * dx2 - lccw = xp.where(dx1dy2 > dy1dx2, True, False) - ccw_out = xp.where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise + lccw = np.where(dx1dy2 > dy1dx2, True, False) + ccw_out = np.where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise return ccw_out @@ -81,7 +81,7 @@ def _lintersect_numpy( line1_p2_lon, line1_p2_lat, ) - lintersect_out = xp.where((intersect1 + intersect2) == -2, True, False) + lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) return lintersect_out @@ -98,10 +98,10 @@ def _line_intersect_numpy( line2_p2_lat, ): d1 = line1_p2_lon - line1_p1_lon - d1 = xp.where(d1 != 0.0, d1, line1_p2_lon) + d1 = np.where(d1 != 0.0, d1, line1_p2_lon) d2 = line2_p2_lon - line2_p1_lon - d2 = xp.where(d2 != 0.0, d2, line2_p2_lon) + d2 = np.where(d2 != 0.0, d2, line2_p2_lon) m1 = (line1_p2_lat - line1_p1_lat) / d1 m2 = (line2_p2_lat - line2_p1_lat) / d2 @@ -174,28 +174,28 @@ def _generate_flux_area_geometry( tri_line1_p1_lon = arrival_pts_1_lon_dsl tri_line1_p1_lat = arrival_pts_1_lat_dsl - tri_line1_p2_lon = xp.where( + tri_line1_p2_lon = np.where( lvn_pos, - xp.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - xp.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), ) - tri_line1_p2_lat = xp.where( + tri_line1_p2_lat = np.where( lvn_pos, - xp.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - xp.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), ) tri_line2_p1_lon = arrival_pts_2_lon_dsl tri_line2_p1_lat = arrival_pts_2_lat_dsl - tri_line2_p2_lon = xp.where( + tri_line2_p2_lon = np.where( lvn_pos, - xp.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), - xp.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), ) - tri_line2_p2_lat = xp.where( + tri_line2_p2_lat = np.where( lvn_pos, - xp.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), - xp.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), ) return ( @@ -240,26 +240,26 @@ def _apply_case1_patch0( ): dreg_patch0_1_lon_dsl = arrival_pts_1_lon_dsl dreg_patch0_1_lat_dsl = arrival_pts_1_lat_dsl - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_2_lon_dsl = np.where( mask_case1, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), arrival_pts_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case1, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), arrival_pts_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) - dreg_patch0_3_lat_dsl = xp.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, depart_pts_2_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, depart_pts_2_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( mask_case1, - xp.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), depart_pts_1_lon_dsl, ) - dreg_patch0_4_lat_dsl = xp.where( + dreg_patch0_4_lat_dsl = np.where( mask_case1, - xp.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), depart_pts_1_lat_dsl, ) @@ -285,21 +285,21 @@ def _apply_case1_patch1( ps1_x, ps1_y, ): - dreg_patch1_1_lon_vmask = xp.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_1_lat_vmask = xp.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_4_lon_vmask = xp.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) - dreg_patch1_4_lat_vmask = xp.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) - dreg_patch1_2_lon_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 ) - dreg_patch1_2_lat_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + dreg_patch1_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 ) - dreg_patch1_3_lon_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + dreg_patch1_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 ) - dreg_patch1_3_lat_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + dreg_patch1_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 ) return ( @@ -325,21 +325,21 @@ def _apply_case1_patch2( ps2_y, ): # Case 1 - patch 2 - dreg_patch2_1_lon_vmask = xp.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) - dreg_patch2_1_lat_vmask = xp.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_4_lon_vmask = xp.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) - dreg_patch2_4_lat_vmask = xp.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) - dreg_patch2_2_lon_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 + dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 ) - dreg_patch2_2_lat_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 + dreg_patch2_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 ) - dreg_patch2_3_lon_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 + dreg_patch2_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 ) - dreg_patch2_3_lat_vmask = xp.where( - mask_case1, xp.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 + dreg_patch2_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 ) return ( @@ -374,28 +374,28 @@ def _apply_case2a_patch0( dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, ): - dreg_patch0_1_lon_dsl = xp.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = xp.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case2a, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case2a, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = xp.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( mask_case2a, - xp.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl, ) - dreg_patch0_4_lat_dsl = xp.where( + dreg_patch0_4_lat_dsl = np.where( mask_case2a, - xp.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, ) @@ -429,36 +429,36 @@ def _apply_case2a_patch1( dreg_patch1_3_lon_vmask, dreg_patch1_3_lat_vmask, ): - dreg_patch1_1_lon_vmask = xp.where( + dreg_patch1_1_lon_vmask = np.where( mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask ) - dreg_patch1_1_lat_vmask = xp.where( + dreg_patch1_1_lat_vmask = np.where( mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask ) - dreg_patch1_4_lon_vmask = xp.where( + dreg_patch1_4_lon_vmask = np.where( mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask ) - dreg_patch1_4_lat_vmask = xp.where( + dreg_patch1_4_lat_vmask = np.where( mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask ) - dreg_patch1_2_lon_vmask = xp.where( + dreg_patch1_2_lon_vmask = np.where( mask_case2a, - xp.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), dreg_patch1_2_lon_vmask, ) - dreg_patch1_2_lat_vmask = xp.where( + dreg_patch1_2_lat_vmask = np.where( mask_case2a, - xp.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), + np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), dreg_patch1_2_lat_vmask, ) - dreg_patch1_3_lon_vmask = xp.where( + dreg_patch1_3_lon_vmask = np.where( mask_case2a, - xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), dreg_patch1_3_lon_vmask, ) - dreg_patch1_3_lat_vmask = xp.where( + dreg_patch1_3_lat_vmask = np.where( mask_case2a, - xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), dreg_patch1_3_lat_vmask, ) @@ -494,28 +494,28 @@ def _apply_case2b_patch0( dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, ): - dreg_patch0_1_lon_dsl = xp.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = xp.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case2b, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case2b, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = xp.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( mask_case2b, - xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl, ) - dreg_patch0_4_lat_dsl = xp.where( + dreg_patch0_4_lat_dsl = np.where( mask_case2b, - xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, ) @@ -542,16 +542,16 @@ def _apply_case2b_patch1( dreg_patch1_4_lon_vmask, dreg_patch1_4_lat_vmask, ): - zeros_array = xp.zeros_like(mask_case2b) + zeros_array = np.zeros_like(mask_case2b) - dreg_patch1_1_lon_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = xp.where(mask_case2b, zeros_array, dreg_patch1_4_lat_vmask) + dreg_patch1_1_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2b, zeros_array, dreg_patch1_4_lat_vmask) return ( dreg_patch1_1_lon_vmask, @@ -583,36 +583,36 @@ def _apply_case2b_patch2( dreg_patch2_3_lon_vmask, dreg_patch2_3_lat_vmask, ): - dreg_patch2_1_lon_vmask = xp.where( + dreg_patch2_1_lon_vmask = np.where( mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask ) - dreg_patch2_1_lat_vmask = xp.where( + dreg_patch2_1_lat_vmask = np.where( mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask ) - dreg_patch2_4_lon_vmask = xp.where( + dreg_patch2_4_lon_vmask = np.where( mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask ) - dreg_patch2_4_lat_vmask = xp.where( + dreg_patch2_4_lat_vmask = np.where( mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask ) - dreg_patch2_2_lon_vmask = xp.where( + dreg_patch2_2_lon_vmask = np.where( mask_case2b, - xp.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), dreg_patch2_2_lon_vmask, ) - dreg_patch2_2_lat_vmask = xp.where( + dreg_patch2_2_lat_vmask = np.where( mask_case2b, - xp.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), dreg_patch2_2_lat_vmask, ) - dreg_patch2_3_lon_vmask = xp.where( + dreg_patch2_3_lon_vmask = np.where( mask_case2b, - xp.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), + np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), dreg_patch2_3_lon_vmask, ) - dreg_patch2_3_lat_vmask = xp.where( + dreg_patch2_3_lat_vmask = np.where( mask_case2b, - xp.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), + np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), dreg_patch2_3_lat_vmask, ) @@ -648,28 +648,28 @@ def _apply_case3a_patch0( dreg_patch0_4_lon_dsl, dreg_patch0_4_lat_dsl, ): - dreg_patch0_1_lon_dsl = xp.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = xp.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case3a, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case3a, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) - dreg_patch0_3_lat_dsl = xp.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( mask_case3a, - xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), dreg_patch0_4_lon_dsl, ) - dreg_patch0_4_lat_dsl = xp.where( + dreg_patch0_4_lat_dsl = np.where( mask_case3a, - xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), dreg_patch0_4_lat_dsl, ) @@ -705,36 +705,36 @@ def _apply_case3a_patch1( dreg_patch1_3_lon_vmask, dreg_patch1_3_lat_vmask, ): - dreg_patch1_1_lon_vmask = xp.where( + dreg_patch1_1_lon_vmask = np.where( mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask ) - dreg_patch1_1_lat_vmask = xp.where( + dreg_patch1_1_lat_vmask = np.where( mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask ) - dreg_patch1_2_lon_vmask = xp.where( + dreg_patch1_2_lon_vmask = np.where( mask_case3a, - xp.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), dreg_patch1_2_lon_vmask, ) - dreg_patch1_2_lat_vmask = xp.where( + dreg_patch1_2_lat_vmask = np.where( mask_case3a, - xp.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), dreg_patch1_2_lat_vmask, ) - dreg_patch1_3_lon_vmask = xp.where( + dreg_patch1_3_lon_vmask = np.where( mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask ) - dreg_patch1_3_lat_vmask = xp.where( + dreg_patch1_3_lat_vmask = np.where( mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask ) - dreg_patch1_4_lon_vmask = xp.where( + dreg_patch1_4_lon_vmask = np.where( mask_case3a, - xp.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), dreg_patch1_4_lon_vmask, ) - dreg_patch1_4_lat_vmask = xp.where( + dreg_patch1_4_lat_vmask = np.where( mask_case3a, - xp.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), dreg_patch1_4_lat_vmask, ) @@ -768,28 +768,28 @@ def _apply_case3b_patch0( dreg_patch0_3_lon_dsl, dreg_patch0_3_lat_dsl, ): - dreg_patch0_1_lon_dsl = xp.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) - dreg_patch0_1_lat_dsl = xp.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) - dreg_patch0_4_lon_dsl = xp.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) - dreg_patch0_4_lat_dsl = xp.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) - dreg_patch0_2_lon_dsl = xp.where( + dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( mask_case3b, - xp.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), dreg_patch0_2_lon_dsl, ) - dreg_patch0_2_lat_dsl = xp.where( + dreg_patch0_2_lat_dsl = np.where( mask_case3b, - xp.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), dreg_patch0_2_lat_dsl, ) - dreg_patch0_3_lon_dsl = xp.where( + dreg_patch0_3_lon_dsl = np.where( mask_case3b, - xp.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), dreg_patch0_3_lon_dsl, ) - dreg_patch0_3_lat_dsl = xp.where( + dreg_patch0_3_lat_dsl = np.where( mask_case3b, - xp.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), dreg_patch0_3_lat_dsl, ) @@ -825,36 +825,36 @@ def _apply_case3b_patch2( dreg_patch2_4_lon_vmask, dreg_patch2_4_lat_vmask, ): - dreg_patch2_1_lon_vmask = xp.where( + dreg_patch2_1_lon_vmask = np.where( mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask ) - dreg_patch2_1_lat_vmask = xp.where( + dreg_patch2_1_lat_vmask = np.where( mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask ) - dreg_patch2_2_lon_vmask = xp.where( + dreg_patch2_2_lon_vmask = np.where( mask_case3b, - xp.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), dreg_patch2_2_lon_vmask, ) - dreg_patch2_2_lat_vmask = xp.where( + dreg_patch2_2_lat_vmask = np.where( mask_case3b, - xp.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), dreg_patch2_2_lat_vmask, ) - dreg_patch2_3_lon_vmask = xp.where( + dreg_patch2_3_lon_vmask = np.where( mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask ) - dreg_patch2_3_lat_vmask = xp.where( + dreg_patch2_3_lat_vmask = np.where( mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask ) - dreg_patch2_4_lon_vmask = xp.where( + dreg_patch2_4_lon_vmask = np.where( mask_case3b, - xp.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), dreg_patch2_4_lon_vmask, ) - dreg_patch2_4_lat_vmask = xp.where( + dreg_patch2_4_lat_vmask = np.where( mask_case3b, - xp.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), dreg_patch2_4_lat_vmask, ) @@ -890,10 +890,10 @@ def reference( ) -> dict: e2c = grid.connectivities[dims.E2CDim] ptr_v3_lon = helpers.reshape(ptr_v3_lon, e2c.shape) - ptr_v3_lon_e = xp.expand_dims(ptr_v3_lon, axis=-1) + ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) ptr_v3_lat = helpers.reshape(ptr_v3_lat, e2c.shape) - ptr_v3_lat_e = xp.expand_dims(ptr_v3_lat, axis=-1) - tangent_orientation_dsl = xp.expand_dims(tangent_orientation_dsl, axis=-1) + ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) + tangent_orientation_dsl = np.expand_dims(tangent_orientation_dsl, axis=-1) result_tuple = cls._generate_flux_area_geometry( dreg_patch0_1_lon_dsl, @@ -955,13 +955,13 @@ def reference( tri_line2_p2_lat, ) - lvn_sys_pos = xp.where( - (p_vn * xp.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, + lvn_sys_pos = np.where( + (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, True, False, ) - famask_bool = xp.where(famask_int == 1, True, False) - mask_case1 = xp.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) + famask_bool = np.where(famask_int == 1, True, False) + mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) ps1_x, ps1_y = _line_intersect_numpy( fl_line_p1_lon, fl_line_p1_lat, @@ -1052,8 +1052,8 @@ def reference( ) # ------------------------------------------------- Case 2a - mask_case2a = xp.logical_and.reduce( - [lintersect_line1, xp.logical_not(lintersect_line2), famask_bool] + mask_case2a = np.logical_and.reduce( + [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] ) # Case 2a - patch 0 result_tuple_patch0 = cls._apply_case2a_patch0( @@ -1118,18 +1118,18 @@ def reference( dreg_patch1_3_lat_vmask, ) = result_tuple_patch1 # Case 2a - patch 2 - dreg_patch2_1_lon_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = xp.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) + dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) # -------------------------------------------------- Case 2b - mask_case2b = xp.logical_and.reduce( - [lintersect_line2, xp.logical_not(lintersect_line1), famask_bool] + mask_case2b = np.logical_and.reduce( + [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] ) # Case 2b - patch 0 result_tuple_patch0_case2b = cls._apply_case2b_patch0( @@ -1241,7 +1241,7 @@ def reference( tri_line1_p2_lon, tri_line1_p2_lat, ) - mask_case3a = xp.logical_and(lintersect_e2_line1, famask_bool) + mask_case3a = np.logical_and(lintersect_e2_line1, famask_bool) pi1_x, pi1_y = _line_intersect_numpy( fl_e2_p1_lon, fl_e2_p1_lat, @@ -1316,14 +1316,14 @@ def reference( dreg_patch1_3_lat_vmask, ) # Case 3a - patch 2 - dreg_patch2_1_lon_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) - dreg_patch2_4_lat_vmask = xp.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) + dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) # ------------------------------------------------ Case 3b # Check whether flux area edge 1 intersects with triangle edge 2 @@ -1377,14 +1377,14 @@ def reference( dreg_patch0_3_lat_dsl, ) # Case 3b - patch 1 - dreg_patch1_1_lon_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = xp.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) + dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) # Case 3b - patch 2 ( @@ -1419,31 +1419,31 @@ def reference( # --------------------------------------------- Case 4 # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions # can't be overwritten by this new condition. - indices_previously_matched = xp.logical_or.reduce( + indices_previously_matched = np.logical_or.reduce( [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] ) # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices - mask_case4 = xp.logical_and.reduce( - [famask_bool, xp.logical_not(indices_previously_matched)] + mask_case4 = np.logical_and.reduce( + [famask_bool, np.logical_not(indices_previously_matched)] ) # Case 4 - patch 0 - no change # Case 4 - patch 1 - dreg_patch1_1_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) - dreg_patch1_1_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) - dreg_patch1_2_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) - dreg_patch1_2_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) - dreg_patch1_3_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) - dreg_patch1_3_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) - dreg_patch1_4_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) - dreg_patch1_4_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) + dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) # Case 4 - patch 2 - dreg_patch2_1_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) - dreg_patch2_1_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) - dreg_patch2_2_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) - dreg_patch2_2_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) - dreg_patch2_3_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) - dreg_patch2_3_lat_vmask = xp.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) - dreg_patch2_4_lon_vmask = xp.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) return dict( dreg_patch0_1_lon_dsl=dreg_patch0_1_lon_dsl, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py index 2e19beac72..08b651a24b 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_for_cubic_reconstruction.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ prepare_numerical_quadrature_for_cubic_reconstruction, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp @pytest.mark.slow_tests @@ -337,14 +337,14 @@ def _compute_vector_sums( def reference( cls, grid, - p_coords_dreg_v_1_x: xp.array, - p_coords_dreg_v_2_x: xp.array, - p_coords_dreg_v_3_x: xp.array, - p_coords_dreg_v_4_x: xp.array, - p_coords_dreg_v_1_y: xp.array, - p_coords_dreg_v_2_y: xp.array, - p_coords_dreg_v_3_y: xp.array, - p_coords_dreg_v_4_y: xp.array, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, shape_func_1_1: float, shape_func_2_1: float, shape_func_3_1: float, @@ -464,10 +464,10 @@ def reference( ) z_area = p_quad_vector_sum_1 - p_dreg_area_out = xp.where( + p_dreg_area_out = np.where( z_area >= 0.0, - xp.maximum(eps, xp.absolute(z_area)), - -xp.maximum(eps, xp.absolute(z_area)), + np.maximum(eps, np.absolute(z_area)), + -np.maximum(eps, np.absolute(z_area)), ) return dict( p_quad_vector_sum_1=p_quad_vector_sum_1, @@ -532,7 +532,7 @@ def input_data(self, grid) -> dict: wgt_zeta_2 = 0.003 wgt_eta_1 = 0.002 wgt_eta_2 = 0.007 - dbl_eps = xp.float64(0.1) + dbl_eps = np.float64(0.1) eps = 0.1 return dict( p_coords_dreg_v_1_x=p_coords_dreg_v_1_x, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py index 9191ed3b29..4b78309142 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_prepare_numerical_quadrature_list_for_cubic_reconstruction.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -14,7 +15,6 @@ prepare_numerical_quadrature_list_for_cubic_reconstruction, ) from icon4py.model.common import dimension as dims -from icon4py.model.common.settings import xp @pytest.mark.slow_tests @@ -89,18 +89,18 @@ def _compute_wgt_t_detjac( 1.0 + zeta_4, ) - famask_bool = xp.where(famask_int == 1, True, False) + famask_bool = np.where(famask_int == 1, True, False) - p_coords_dreg_v_1_x = xp.where(famask_bool, p_coords_dreg_v_1_x, 0.0) - p_coords_dreg_v_2_x = xp.where(famask_bool, p_coords_dreg_v_2_x, 0.0) - p_coords_dreg_v_3_x = xp.where(famask_bool, p_coords_dreg_v_3_x, 0.0) - p_coords_dreg_v_4_x = xp.where(famask_bool, p_coords_dreg_v_4_x, 0.0) - p_coords_dreg_v_1_y = xp.where(famask_bool, p_coords_dreg_v_1_y, 0.0) - p_coords_dreg_v_2_y = xp.where(famask_bool, p_coords_dreg_v_2_y, 0.0) - p_coords_dreg_v_3_y = xp.where(famask_bool, p_coords_dreg_v_3_y, 0.0) - p_coords_dreg_v_4_y = xp.where(famask_bool, p_coords_dreg_v_4_y, 0.0) + p_coords_dreg_v_1_x = np.where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = np.where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = np.where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = np.where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = np.where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = np.where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) - wgt_t_detjac_1 = xp.where( + wgt_t_detjac_1 = np.where( famask_bool, dbl_eps + z_wgt_1 @@ -125,7 +125,7 @@ def _compute_wgt_t_detjac( 0.0, ) - wgt_t_detjac_2 = xp.where( + wgt_t_detjac_2 = np.where( famask_bool, dbl_eps + z_wgt_2 @@ -150,7 +150,7 @@ def _compute_wgt_t_detjac( 0.0, ) - wgt_t_detjac_3 = xp.where( + wgt_t_detjac_3 = np.where( famask_bool, dbl_eps + z_wgt_3 @@ -174,7 +174,7 @@ def _compute_wgt_t_detjac( ), 0.0, ) - wgt_t_detjac_4 = xp.where( + wgt_t_detjac_4 = np.where( famask_bool, dbl_eps + z_wgt_4 @@ -374,16 +374,16 @@ def _compute_vector_sums( def reference( cls, grid, - famask_int: xp.array, - p_coords_dreg_v_1_x: xp.array, - p_coords_dreg_v_2_x: xp.array, - p_coords_dreg_v_3_x: xp.array, - p_coords_dreg_v_4_x: xp.array, - p_coords_dreg_v_1_y: xp.array, - p_coords_dreg_v_2_y: xp.array, - p_coords_dreg_v_3_y: xp.array, - p_coords_dreg_v_4_y: xp.array, - p_dreg_area_in: xp.array, + famask_int: np.array, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + p_dreg_area_in: np.array, shape_func_1_1: float, shape_func_2_1: float, shape_func_3_1: float, @@ -569,7 +569,7 @@ def input_data(self, grid) -> dict: wgt_zeta_2 = 0.003 wgt_eta_1 = 0.002 wgt_eta_2 = 0.007 - dbl_eps = xp.float64(0.1) + dbl_eps = np.float64(0.1) eps = 0.1 return dict( famask_int=famask_int, diff --git a/model/atmosphere/advection/tests/advection_stencil_tests/test_reconstruct_cubic_coefficients_svd.py b/model/atmosphere/advection/tests/advection_stencil_tests/test_reconstruct_cubic_coefficients_svd.py index 4bf6ad8f73..b4bc666c47 100644 --- a/model/atmosphere/advection/tests/advection_stencil_tests/test_reconstruct_cubic_coefficients_svd.py +++ b/model/atmosphere/advection/tests/advection_stencil_tests/test_reconstruct_cubic_coefficients_svd.py @@ -7,6 +7,7 @@ # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.test_utils.helpers as helpers @@ -15,7 +16,6 @@ ) from icon4py.model.common import dimension as dims from icon4py.model.common.grid import horizontal as h_grid -from icon4py.model.common.settings import xp @pytest.mark.slow_tests @@ -37,35 +37,35 @@ class TestReconstructCubicCoefficientsSvd(helpers.StencilTest): @staticmethod def reference( grid, - p_cc: xp.array, - lsq_pseudoinv_1: xp.array, - lsq_pseudoinv_2: xp.array, - lsq_pseudoinv_3: xp.array, - lsq_pseudoinv_4: xp.array, - lsq_pseudoinv_5: xp.array, - lsq_pseudoinv_6: xp.array, - lsq_pseudoinv_7: xp.array, - lsq_pseudoinv_8: xp.array, - lsq_pseudoinv_9: xp.array, - lsq_moments_1: xp.array, - lsq_moments_2: xp.array, - lsq_moments_3: xp.array, - lsq_moments_4: xp.array, - lsq_moments_5: xp.array, - lsq_moments_6: xp.array, - lsq_moments_7: xp.array, - lsq_moments_8: xp.array, - lsq_moments_9: xp.array, - p_coeff_1_dsl: xp.array, - p_coeff_2_dsl: xp.array, - p_coeff_3_dsl: xp.array, - p_coeff_4_dsl: xp.array, - p_coeff_5_dsl: xp.array, - p_coeff_6_dsl: xp.array, - p_coeff_7_dsl: xp.array, - p_coeff_8_dsl: xp.array, - p_coeff_9_dsl: xp.array, - p_coeff_10_dsl: xp.array, + p_cc: np.array, + lsq_pseudoinv_1: np.array, + lsq_pseudoinv_2: np.array, + lsq_pseudoinv_3: np.array, + lsq_pseudoinv_4: np.array, + lsq_pseudoinv_5: np.array, + lsq_pseudoinv_6: np.array, + lsq_pseudoinv_7: np.array, + lsq_pseudoinv_8: np.array, + lsq_pseudoinv_9: np.array, + lsq_moments_1: np.array, + lsq_moments_2: np.array, + lsq_moments_3: np.array, + lsq_moments_4: np.array, + lsq_moments_5: np.array, + lsq_moments_6: np.array, + lsq_moments_7: np.array, + lsq_moments_8: np.array, + lsq_moments_9: np.array, + p_coeff_1_dsl: np.array, + p_coeff_2_dsl: np.array, + p_coeff_3_dsl: np.array, + p_coeff_4_dsl: np.array, + p_coeff_5_dsl: np.array, + p_coeff_6_dsl: np.array, + p_coeff_7_dsl: np.array, + p_coeff_8_dsl: np.array, + p_coeff_9_dsl: np.array, + p_coeff_10_dsl: np.array, **kwargs, ) -> dict: p_coeff_1_dsl_cp = p_coeff_1_dsl.copy() @@ -80,41 +80,41 @@ def reference( p_coeff_10_dsl_cp = p_coeff_10_dsl.copy() c2e2c2e2c = grid.connectivities[dims.C2E2C2E2CDim] - lsq_moments_1 = xp.expand_dims(lsq_moments_1, axis=-1) - lsq_moments_2 = xp.expand_dims(lsq_moments_2, axis=-1) - lsq_moments_3 = xp.expand_dims(lsq_moments_3, axis=-1) - lsq_moments_4 = xp.expand_dims(lsq_moments_4, axis=-1) - lsq_moments_5 = xp.expand_dims(lsq_moments_5, axis=-1) - lsq_moments_6 = xp.expand_dims(lsq_moments_6, axis=-1) - lsq_moments_7 = xp.expand_dims(lsq_moments_7, axis=-1) - lsq_moments_8 = xp.expand_dims(lsq_moments_8, axis=-1) - lsq_moments_9 = xp.expand_dims(lsq_moments_9, axis=-1) - lsq_moments_1 = xp.broadcast_to(lsq_moments_1, p_cc.shape) - lsq_moments_2 = xp.broadcast_to(lsq_moments_2, p_cc.shape) - lsq_moments_3 = xp.broadcast_to(lsq_moments_3, p_cc.shape) - lsq_moments_4 = xp.broadcast_to(lsq_moments_4, p_cc.shape) - lsq_moments_5 = xp.broadcast_to(lsq_moments_5, p_cc.shape) - lsq_moments_6 = xp.broadcast_to(lsq_moments_6, p_cc.shape) - lsq_moments_7 = xp.broadcast_to(lsq_moments_7, p_cc.shape) - lsq_moments_8 = xp.broadcast_to(lsq_moments_8, p_cc.shape) + lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) + lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) + lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) + lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) + lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) + lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) + lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) + lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) + lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) + lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) + lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) + lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) + lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) + lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) + lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) + lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) + lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) lsq_pseudoinv_9 = helpers.reshape(lsq_pseudoinv_9, c2e2c2e2c.shape) - lsq_pseudoinv_9 = xp.expand_dims(lsq_pseudoinv_9, axis=-1) + lsq_pseudoinv_9 = np.expand_dims(lsq_pseudoinv_9, axis=-1) lsq_pseudoinv_8 = helpers.reshape(lsq_pseudoinv_8, c2e2c2e2c.shape) - lsq_pseudoinv_8 = xp.expand_dims(lsq_pseudoinv_8, axis=-1) + lsq_pseudoinv_8 = np.expand_dims(lsq_pseudoinv_8, axis=-1) lsq_pseudoinv_7 = helpers.reshape(lsq_pseudoinv_7, c2e2c2e2c.shape) - lsq_pseudoinv_7 = xp.expand_dims(lsq_pseudoinv_7, axis=-1) + lsq_pseudoinv_7 = np.expand_dims(lsq_pseudoinv_7, axis=-1) lsq_pseudoinv_6 = helpers.reshape(lsq_pseudoinv_6, c2e2c2e2c.shape) - lsq_pseudoinv_6 = xp.expand_dims(lsq_pseudoinv_6, axis=-1) + lsq_pseudoinv_6 = np.expand_dims(lsq_pseudoinv_6, axis=-1) lsq_pseudoinv_5 = helpers.reshape(lsq_pseudoinv_5, c2e2c2e2c.shape) - lsq_pseudoinv_5 = xp.expand_dims(lsq_pseudoinv_5, axis=-1) + lsq_pseudoinv_5 = np.expand_dims(lsq_pseudoinv_5, axis=-1) lsq_pseudoinv_4 = helpers.reshape(lsq_pseudoinv_4, c2e2c2e2c.shape) - lsq_pseudoinv_4 = xp.expand_dims(lsq_pseudoinv_4, axis=-1) + lsq_pseudoinv_4 = np.expand_dims(lsq_pseudoinv_4, axis=-1) lsq_pseudoinv_3 = helpers.reshape(lsq_pseudoinv_3, c2e2c2e2c.shape) - lsq_pseudoinv_3 = xp.expand_dims(lsq_pseudoinv_3, axis=-1) + lsq_pseudoinv_3 = np.expand_dims(lsq_pseudoinv_3, axis=-1) lsq_pseudoinv_2 = helpers.reshape(lsq_pseudoinv_2, c2e2c2e2c.shape) - lsq_pseudoinv_2 = xp.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) lsq_pseudoinv_1 = helpers.reshape(lsq_pseudoinv_1, c2e2c2e2c.shape) - lsq_pseudoinv_1 = xp.expand_dims(lsq_pseudoinv_1, axis=-1) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) p_coeff_10_dsl = ( lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) diff --git a/model/atmosphere/advection/tests/advection_tests/test_advection.py b/model/atmosphere/advection/tests/advection_tests/test_advection.py index 75ccbbbc5b..0afdba32f1 100644 --- a/model/atmosphere/advection/tests/advection_tests/test_advection.py +++ b/model/atmosphere/advection/tests/advection_tests/test_advection.py @@ -25,17 +25,13 @@ ) -# note about ntracer: The first tracer is always dry air which is not advected. Thus, originally -# ntracer=2 is the first tracer in transport_nml, ntracer=3 the second, and so on. -# Here though, ntracer=1 corresponds to the first tracer in transport_nml. - # ntracer legend for the serialization data used here in test_advection: # ------------------------------------ # ntracer | 1, 2, 3, 4, 5 | # ------------------------------------ # ivadv_tracer | 3, 0, 0, 2, 3 | -# itype_hlimit | 3, 3, 4, 0, 0 | -# itype_vlimit | 1, 0, 0, 1, 2 | +# itype_hlimit | 3, 4, 3, 0, 0 | +# itype_vlimit | 1, 0, 0, 2, 1 | # ihadv_tracer | 52, 2, 2, 0, 0 | # ------------------------------------ @@ -62,6 +58,24 @@ advection.VerticalAdvectionType.NO_ADVECTION, advection.VerticalAdvectionLimiter.NO_LIMITER, ), + ( + "2021-06-20T12:00:10.000", + False, + 5, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), + ( + "2021-06-20T12:00:20.000", + True, + 5, + advection.HorizontalAdvectionType.NO_ADVECTION, + advection.HorizontalAdvectionLimiter.NO_LIMITER, + advection.VerticalAdvectionType.PPM_3RD_ORDER, + advection.VerticalAdvectionLimiter.SEMI_MONOTONIC, + ), ], ) def test_advection_run_single_step( @@ -69,6 +83,7 @@ def test_advection_run_single_step( icon_grid, interpolation_savepoint, least_squares_savepoint, + metrics_savepoint, advection_init_savepoint, advection_exit_savepoint, data_provider, @@ -89,7 +104,7 @@ def test_advection_run_single_step( ) interpolation_state = construct_interpolation_state(interpolation_savepoint) least_squares_state = construct_least_squares_state(least_squares_savepoint) - metric_state = construct_metric_state(icon_grid) + metric_state = construct_metric_state(icon_grid, metrics_savepoint) edge_geometry = grid_savepoint.construct_edge_geometry() cell_geometry = grid_savepoint.construct_cell_geometry() @@ -132,4 +147,5 @@ def test_advection_run_single_step( diagnostic_state_ref=diagnostic_state_ref, p_tracer_new=p_tracer_new, p_tracer_new_ref=p_tracer_new_ref, + even_timestep=even_timestep, ) diff --git a/model/atmosphere/advection/tests/advection_tests/utils.py b/model/atmosphere/advection/tests/advection_tests/utils.py index 3e44294c74..dbd884c25f 100644 --- a/model/atmosphere/advection/tests/advection_tests/utils.py +++ b/model/atmosphere/advection/tests/advection_tests/utils.py @@ -8,6 +8,8 @@ import logging +import gt4py.next as gtx +import numpy as np from icon4py.model.atmosphere.advection import advection, advection_states from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta @@ -54,12 +56,16 @@ def construct_least_squares_state( ) -def construct_metric_state(icon_grid) -> advection_states.AdvectionMetricState: +def construct_metric_state( + icon_grid, savepoint: sb.MetricSavepoint +) -> advection_states.AdvectionMetricState: constant_f = helpers.constant_field(icon_grid, 1.0, dims.KDim) + ddqz_z_full_np = np.reciprocal(savepoint.inv_ddqz_z_full().asnumpy()) return advection_states.AdvectionMetricState( deepatmo_divh=constant_f, deepatmo_divzl=constant_f, deepatmo_divzu=constant_f, + ddqz_z_full=gtx.as_field((dims.CellDim, dims.KDim), ddqz_z_full_np), ) @@ -128,11 +134,17 @@ def verify_advection_fields( diagnostic_state_ref: advection_states.AdvectionDiagnosticState, p_tracer_new: fa.CellKField[ta.wpfloat], p_tracer_new_ref: fa.CellKField[ta.wpfloat], + even_timestep: bool, ): # cell indices cell_domain = h_grid.domain(dims.CellDim) start_cell_lateral_boundary = grid.start_index(cell_domain(h_grid.Zone.LATERAL_BOUNDARY)) + start_cell_lateral_boundary_level_2 = grid.start_index( + cell_domain(h_grid.Zone.LATERAL_BOUNDARY_LEVEL_2) + ) + start_cell_nudging = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) end_cell_local = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + end_cell_end = grid.end_index(cell_domain(h_grid.Zone.END)) # edge indices edge_domain = h_grid.domain(dims.EdgeDim) @@ -141,46 +153,35 @@ def verify_advection_fields( ) end_edge_halo = grid.end_index(edge_domain(h_grid.Zone.HALO)) - # log advection output fields - log_dbg( - diagnostic_state.hfl_tracer.asnumpy()[start_edge_lateral_boundary_level_5:end_edge_halo, :], - "hfl_tracer", - ) - log_dbg( - diagnostic_state_ref.hfl_tracer.asnumpy()[ - start_edge_lateral_boundary_level_5:end_edge_halo, : - ], - "hfl_tracer_ref", - ) - log_dbg( - diagnostic_state.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "vfl_tracer", - ) - log_dbg( - diagnostic_state_ref.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "vfl_tracer_ref", - ) - log_dbg(p_tracer_new.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], "p_tracer_new") - log_dbg( - p_tracer_new_ref.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - "p_tracer_new_ref", + hfl_tracer_range = np.arange(start_edge_lateral_boundary_level_5, end_edge_halo) + vfl_tracer_range = ( + np.arange(start_cell_lateral_boundary_level_2, end_cell_end) + if even_timestep + else np.arange(start_cell_nudging, end_cell_local) ) + p_tracer_new_range = np.arange(start_cell_lateral_boundary, end_cell_local) + + # log advection output fields + log_dbg(diagnostic_state.hfl_tracer.asnumpy()[hfl_tracer_range, :], "hfl_tracer") + log_dbg(diagnostic_state_ref.hfl_tracer.asnumpy()[hfl_tracer_range, :], "hfl_tracer_ref") + log_dbg(diagnostic_state.vfl_tracer.asnumpy()[vfl_tracer_range, :], "vfl_tracer") + log_dbg(diagnostic_state_ref.vfl_tracer.asnumpy()[vfl_tracer_range, :], "vfl_tracer_ref") + log_dbg(p_tracer_new.asnumpy()[p_tracer_new_range, :], "p_tracer_new") + log_dbg(p_tracer_new_ref.asnumpy()[p_tracer_new_range, :], "p_tracer_new_ref") # verify advection output fields assert helpers.dallclose( - diagnostic_state.hfl_tracer.asnumpy()[start_edge_lateral_boundary_level_5:end_edge_halo, :], - diagnostic_state_ref.hfl_tracer.asnumpy()[ - start_edge_lateral_boundary_level_5:end_edge_halo, : - ], + diagnostic_state.hfl_tracer.asnumpy()[hfl_tracer_range, :], + diagnostic_state_ref.hfl_tracer.asnumpy()[hfl_tracer_range, :], rtol=1e-10, ) - assert helpers.dallclose( # TODO (dastrm): adjust indices once there is vertical transport - diagnostic_state.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - diagnostic_state_ref.vfl_tracer.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], + assert helpers.dallclose( + diagnostic_state.vfl_tracer.asnumpy()[vfl_tracer_range, :], + diagnostic_state_ref.vfl_tracer.asnumpy()[vfl_tracer_range, :], rtol=1e-10, ) assert helpers.dallclose( - p_tracer_new.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], - p_tracer_new_ref.asnumpy()[start_cell_lateral_boundary:end_cell_local, :], + p_tracer_new.asnumpy()[p_tracer_new_range, :], + p_tracer_new_ref.asnumpy()[p_tracer_new_range, :], atol=1e-16, ) diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py index b84805ad52..c80ccd100f 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_fixtures.py @@ -167,7 +167,7 @@ def interpolation_savepoint(data_provider): # F811 @pytest.fixture def metrics_savepoint(data_provider): # F811 - """Load data from ICON mestric state savepoint.""" + """Load data from ICON metric state savepoint.""" return data_provider.from_metrics_savepoint() diff --git a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py index c1f348eaca..967bb91229 100644 --- a/model/common/src/icon4py/model/common/test_utils/datatest_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/datatest_utils.py @@ -66,7 +66,7 @@ def get_test_data_root_path() -> pathlib.Path: DATA_URIS_JABW = {1: "https://polybox.ethz.ch/index.php/s/kp9Rab00guECrEd/download"} DATA_URIS_GAUSS3D = {1: "https://polybox.ethz.ch/index.php/s/IiRimdJH2ZBZ1od/download"} DATA_URIS_WK = {1: "https://polybox.ethz.ch/index.php/s/91DEUGmAkBgrXO6/download"} -DATA_URIS_ADVECTION = {1: "https://polybox.ethz.ch/index.php/s/KV6FYstcGysNDOj/download"} +DATA_URIS_ADVECTION = {1: "https://polybox.ethz.ch/index.php/s/3rTia1A41YW7KX9/download"} def get_global_grid_params(experiment: str) -> tuple[int, int]: