From 0b39c164240526bedfbe85d54e2e5d9e3804cbb0 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 6 Dec 2023 09:25:08 +0100 Subject: [PATCH] Fused velocity advection hud (#286) This PR introduces more fused stencils in the velocity advection. The time limit is increased for the cscs-ci. Unifor naming convection is introduced for k, cell, edge. --- ci/cscs.yml | 2 +- .../advection/btraj_dreg_stencil_01.py | 1 - .../advection/btraj_dreg_stencil_02.py | 1 - .../divide_flux_area_list_stencil_01.py | 4 - .../advection/face_val_ppm_stencil_01.py | 13 +- .../advection/face_val_ppm_stencil_02.py | 18 +- .../advection/face_val_ppm_stencil_02a.py | 1 - .../advection/face_val_ppm_stencil_02b.py | 1 - .../advection/face_val_ppm_stencil_02c.py | 1 - .../advection/face_val_ppm_stencil_05.py | 1 - .../hflux_ffsl_hybrid_stencil_01a.py | 1 - .../advection/hflux_ffsl_hybrid_stencil_02.py | 1 - .../advection/hflx_limiter_mo_stencil_03.py | 3 - .../advection/hflx_limiter_mo_stencil_05.py | 1 - .../prep_gauss_quadrature_c_list_stencil.py | 1 - .../advection/recon_lsq_cell_c_svd_stencil.py | 1 - .../upwind_hflux_miura3_stencil_01.py | 1 - .../upwind_hflux_miura_cycl_stencil_01.py | 1 - .../upwind_hflux_miura_cycl_stencil_02.py | 1 - .../upwind_hflux_miura_stencil_01.py | 1 - .../advection/v_limit_prbl_sm_stencil_02.py | 1 - .../test_divide_flux_area_list_stencil_01.py | 3 - .../test_face_val_ppm_stencil_01.py | 14 +- .../test_face_val_ppm_stencil_02.py | 16 +- ...st_prep_gauss_quadrature_c_list_stencil.py | 1 - .../test_prep_gauss_quadrature_c_stencil.py | 1 - .../model/atmosphere/diffusion/diffusion.py | 10 +- .../stencils/apply_diffusion_to_vn.py | 10 +- ...te_horizontal_gradients_for_turbulence.py} | 29 +- .../tests/diffusion_stencil_tests/__init__.py | 33 +++ ...test_apply_diffusion_to_theta_and_exner.py | 127 +++++++++ .../test_apply_diffusion_to_vn.py | 151 +++++++++++ ...ute_horizontal_gradients_for_turbulence.py | 132 +++++++++ ...st_apply_nabla2_and_nabla4_global_to_vn.py | 15 +- .../test_apply_nabla2_and_nabla4_to_vn.py | 79 +++--- ..._apply_nabla2_to_vn_in_lateral_boundary.py | 13 +- .../test_apply_nabla2_to_w.py | 24 +- ...pply_nabla2_to_w_in_upper_damping_layer.py | 16 +- ...st_calculate_diagnostics_for_turbulence.py | 21 +- ...ate_horizontal_gradients_for_turbulence.py | 23 +- ...ate_nabla2_and_smag_coefficients_for_vn.py | 0 .../test_calculate_nabla2_for_w.py | 17 +- .../test_calculate_nabla2_for_z.py | 22 +- .../test_calculate_nabla2_of_theta.py | 15 +- .../test_calculate_nabla4.py | 81 ++++-- ...n_coefficient_for_grid_point_cold_pools.py | 0 ...d_for_grid_point_cold_pools_enhancement.py | 0 ...orary_fields_for_turbulence_diagnostics.py | 0 ...fusion_nabla_of_theta_over_steep_points.py | 82 ++++-- .../test_update_theta_and_exner.py | 22 +- ...sed_velocity_advection_stencil_15_to_18.py | 213 +++++++++++++++ ...sed_velocity_advection_stencil_19_to_20.py | 160 +++++++++++ ...fused_velocity_advection_stencil_1_to_7.py | 243 +++++++++++++++++ ...used_velocity_advection_stencil_8_to_14.py | 163 +++++++++++ .../__init__.py | 0 .../test_compute_airmass.py | 0 ...test_fused_solve_nonhydro_stencil_39_40.py | 0 ...sed_velocity_advection_stencil_15_to_18.py | 245 +++++++++++++++++ ...sed_velocity_advection_stencil_19_to_20.py | 176 ++++++++++++ ...fused_velocity_advection_stencil_1_to_7.py | 255 ++++++++++++++++++ ...used_velocity_advection_stencil_8_to_14.py | 177 ++++++++++++ ...lation_scalar_cells2verts_scalar_ri_dsl.py | 15 +- .../test_mo_math_divrot_rot_vertex_ri_dsl.py | 13 +- ...ath_gradients_grad_green_gauss_cell_dsl.py | 0 ...est_mo_solve_nonhydro_4th_order_divdamp.py | 0 .../test_mo_solve_nonhydro_stencil_01.py | 0 .../test_mo_solve_nonhydro_stencil_02.py | 0 .../test_mo_solve_nonhydro_stencil_03.py | 0 .../test_mo_solve_nonhydro_stencil_04.py | 0 .../test_mo_solve_nonhydro_stencil_05.py | 0 .../test_mo_solve_nonhydro_stencil_06.py | 0 .../test_mo_solve_nonhydro_stencil_07.py | 0 .../test_mo_solve_nonhydro_stencil_08.py | 0 .../test_mo_solve_nonhydro_stencil_09.py | 0 .../test_mo_solve_nonhydro_stencil_10.py | 0 ...test_mo_solve_nonhydro_stencil_11_lower.py | 0 ...test_mo_solve_nonhydro_stencil_11_upper.py | 0 .../test_mo_solve_nonhydro_stencil_12.py | 0 .../test_mo_solve_nonhydro_stencil_13.py | 0 .../test_mo_solve_nonhydro_stencil_14.py | 0 .../test_mo_solve_nonhydro_stencil_15.py | 0 ...nonhydro_stencil_16_fused_btraj_traj_o1.py | 0 .../test_mo_solve_nonhydro_stencil_17.py | 0 .../test_mo_solve_nonhydro_stencil_18.py | 0 .../test_mo_solve_nonhydro_stencil_19.py | 0 .../test_mo_solve_nonhydro_stencil_20.py | 0 .../test_mo_solve_nonhydro_stencil_21.py | 0 .../test_mo_solve_nonhydro_stencil_22.py | 0 .../test_mo_solve_nonhydro_stencil_23.py | 0 .../test_mo_solve_nonhydro_stencil_24.py | 0 .../test_mo_solve_nonhydro_stencil_25.py | 0 .../test_mo_solve_nonhydro_stencil_26.py | 0 .../test_mo_solve_nonhydro_stencil_27.py | 0 .../test_mo_solve_nonhydro_stencil_28.py | 0 .../test_mo_solve_nonhydro_stencil_29.py | 0 .../test_mo_solve_nonhydro_stencil_30.py | 0 .../test_mo_solve_nonhydro_stencil_31.py | 0 .../test_mo_solve_nonhydro_stencil_32.py | 0 .../test_mo_solve_nonhydro_stencil_33.py | 0 .../test_mo_solve_nonhydro_stencil_34.py | 0 .../test_mo_solve_nonhydro_stencil_35.py | 0 .../test_mo_solve_nonhydro_stencil_36.py | 0 .../test_mo_solve_nonhydro_stencil_37.py | 0 .../test_mo_solve_nonhydro_stencil_38.py | 0 .../test_mo_solve_nonhydro_stencil_39.py | 0 .../test_mo_solve_nonhydro_stencil_40.py | 0 .../test_mo_solve_nonhydro_stencil_41.py | 0 .../test_mo_solve_nonhydro_stencil_42.py | 0 .../test_mo_solve_nonhydro_stencil_43.py | 0 .../test_mo_solve_nonhydro_stencil_44.py | 0 .../test_mo_solve_nonhydro_stencil_45.py | 0 .../test_mo_solve_nonhydro_stencil_46.py | 0 .../test_mo_solve_nonhydro_stencil_47.py | 0 .../test_mo_solve_nonhydro_stencil_48.py | 0 .../test_mo_solve_nonhydro_stencil_49.py | 0 .../test_mo_solve_nonhydro_stencil_50.py | 0 .../test_mo_solve_nonhydro_stencil_51.py | 0 .../test_mo_solve_nonhydro_stencil_52.py | 0 .../test_mo_solve_nonhydro_stencil_53.py | 0 .../test_mo_solve_nonhydro_stencil_54.py | 0 .../test_mo_solve_nonhydro_stencil_55.py | 0 .../test_mo_solve_nonhydro_stencil_56_63.py | 0 .../test_mo_solve_nonhydro_stencil_57.py | 0 .../test_mo_solve_nonhydro_stencil_58.py | 0 .../test_mo_solve_nonhydro_stencil_59.py | 0 .../test_mo_solve_nonhydro_stencil_60.py | 0 .../test_mo_solve_nonhydro_stencil_61.py | 0 .../test_mo_solve_nonhydro_stencil_62.py | 0 .../test_mo_solve_nonhydro_stencil_64.py | 0 .../test_mo_solve_nonhydro_stencil_65.py | 0 .../test_mo_solve_nonhydro_stencil_66.py | 0 .../test_mo_solve_nonhydro_stencil_67.py | 0 .../test_mo_solve_nonhydro_stencil_68.py | 0 .../test_mo_velocity_advection_stencil_01.py | 15 +- .../test_mo_velocity_advection_stencil_02.py | 40 +-- .../test_mo_velocity_advection_stencil_03.py | 13 +- .../test_mo_velocity_advection_stencil_04.py | 9 +- .../test_mo_velocity_advection_stencil_05.py | 11 +- .../test_mo_velocity_advection_stencil_06.py | 30 ++- .../test_mo_velocity_advection_stencil_07.py | 49 +++- .../test_mo_velocity_advection_stencil_08.py | 19 +- .../test_mo_velocity_advection_stencil_09.py | 19 +- .../test_mo_velocity_advection_stencil_10.py | 13 +- .../test_mo_velocity_advection_stencil_11.py | 7 +- .../test_mo_velocity_advection_stencil_12.py | 7 +- .../test_mo_velocity_advection_stencil_13.py | 7 +- .../test_mo_velocity_advection_stencil_14.py | 40 +-- .../test_mo_velocity_advection_stencil_15.py | 9 +- .../test_mo_velocity_advection_stencil_16.py | 23 +- .../test_mo_velocity_advection_stencil_17.py | 18 +- .../test_mo_velocity_advection_stencil_18.py | 90 +++++-- .../test_mo_velocity_advection_stencil_19.py | 57 +++- .../test_mo_velocity_advection_stencil_20.py | 148 ++++++---- .../model/common/test_utils/helpers.py | 7 + tools/tests/liskov/fortran_samples.py | 10 +- 155 files changed, 2854 insertions(+), 460 deletions(-) rename model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/{apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py => apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py} (88%) create mode 100644 model/atmosphere/diffusion/tests/diffusion_stencil_tests/__init__.py create mode 100644 model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_theta_and_exner.py create mode 100644 model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_vn.py create mode 100644 model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_apply_nabla2_and_nabla4_global_to_vn.py (80%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_apply_nabla2_and_nabla4_to_vn.py (55%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_apply_nabla2_to_vn_in_lateral_boundary.py (82%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_apply_nabla2_to_w.py (77%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_apply_nabla2_to_w_in_upper_damping_layer.py (83%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_calculate_diagnostics_for_turbulence.py (70%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_calculate_horizontal_gradients_for_turbulence.py (77%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_calculate_nabla2_and_smag_coefficients_for_vn.py (100%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_calculate_nabla2_for_w.py (80%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_calculate_nabla2_for_z.py (81%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_calculate_nabla2_of_theta.py (80%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_calculate_nabla4.py (62%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py (100%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_temporary_field_for_grid_point_cold_pools_enhancement.py (100%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_temporary_fields_for_turbulence_diagnostics.py (100%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py (66%) rename model/atmosphere/diffusion/tests/{stencil_tests => diffusion_stencil_tests}/test_update_theta_and_exner.py (79%) create mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py create mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_19_to_20.py create mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py create mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/__init__.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_compute_airmass.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_fused_solve_nonhydro_stencil_39_40.py (100%) create mode 100644 model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_15_to_18.py create mode 100644 model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_19_to_20.py create mode 100644 model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py create mode 100644 model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_8_to_14.py rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py (80%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_math_divrot_rot_vertex_ri_dsl.py (82%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_math_gradients_grad_green_gauss_cell_dsl.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_4th_order_divdamp.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_01.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_02.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_03.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_04.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_05.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_06.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_07.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_08.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_09.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_10.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_11_lower.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_11_upper.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_12.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_13.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_14.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_15.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_17.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_18.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_19.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_20.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_21.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_22.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_23.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_24.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_25.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_26.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_27.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_28.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_29.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_30.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_31.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_32.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_33.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_34.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_35.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_36.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_37.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_38.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_39.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_40.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_41.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_42.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_43.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_44.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_45.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_46.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_47.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_48.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_49.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_50.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_51.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_52.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_53.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_54.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_55.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_56_63.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_57.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_58.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_59.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_60.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_61.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_62.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_64.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_65.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_66.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_67.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_solve_nonhydro_stencil_68.py (100%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_01.py (81%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_02.py (73%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_03.py (84%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_04.py (86%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_05.py (87%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_06.py (71%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_07.py (72%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_08.py (81%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_09.py (80%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_10.py (81%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_11.py (90%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_12.py (87%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_13.py (87%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_14.py (71%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_15.py (84%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_16.py (79%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_17.py (80%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_18.py (65%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_19.py (70%) rename model/atmosphere/dycore/tests/{stencil_tests => dycore_stencil_tests}/test_mo_velocity_advection_stencil_20.py (57%) diff --git a/ci/cscs.yml b/ci/cscs.yml index c5e8666fe..2d76634bb 100644 --- a/ci/cscs.yml +++ b/ci/cscs.yml @@ -29,7 +29,7 @@ variables: variables: SLURM_JOB_NUM_NODES: 2 SLURM_NTASKS: 2 - SLURM_TIMELIMIT: '01:00:00' + SLURM_TIMELIMIT: '02:00:00' build_job: extends: .build_template diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py index 17eea2c69..60109d3b6 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py @@ -23,7 +23,6 @@ def _btraj_dreg_stencil_01( p_vn: Field[[EdgeDim, KDim], float], tangent_orientation: Field[[EdgeDim], float], ) -> Field[[EdgeDim, KDim], bool]: - tangent_orientation = broadcast(tangent_orientation, (EdgeDim, KDim)) lvn_sys_pos_true = where(p_vn * tangent_orientation >= 0.0, True, False) mask_lcounterclock = broadcast(lcounterclock, (EdgeDim, KDim)) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py index 5984b77c5..a11392c32 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py @@ -24,7 +24,6 @@ def _btraj_dreg_stencil_02( edge_cell_length: Field[[ECDim], float], p_dt: float, ) -> Field[[EdgeDim, KDim], int32]: - lvn_pos = where(p_vn >= 0.0, True, False) traj_length = sqrt(p_vn * p_vn + p_vt * p_vt) * p_dt e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1])) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py index a753bd238..c18d8eac8 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -33,7 +33,6 @@ def ccw( p2_lon: Field[[EdgeDim, KDim], float], p2_lat: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], int32]: - dx1 = p1_lon - p0_lon dy1 = p1_lat - p0_lat @@ -60,7 +59,6 @@ def lintersect( line2_p2_lon: Field[[EdgeDim, KDim], float], line2_p2_lat: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], bool]: - intersect1 = ccw( line1_p1_lon, line1_p1_lat, @@ -108,7 +106,6 @@ def line_intersect( line2_p2_lon: Field[[EdgeDim, KDim], float], line2_p2_lat: Field[[EdgeDim, KDim], float], ) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: - m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) @@ -159,7 +156,6 @@ def _divide_flux_area_list_stencil_01( Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], ]: - arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py index de605d3dd..f47d3c827 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py @@ -22,7 +22,6 @@ def _face_val_ppm_stencil_01a( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) zfac = (p_cc(Koff[+1]) - p_cc) / (p_cellhgt_mc_now(Koff[+1]) + p_cellhgt_mc_now) z_slope = ( @@ -41,7 +40,6 @@ def _face_val_ppm_stencil_01b( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) zfac = (p_cc - p_cc) / (p_cellhgt_mc_now + p_cellhgt_mc_now) z_slope = ( @@ -58,14 +56,13 @@ def _face_val_ppm_stencil_01b( def _face_val_ppm_stencil_01( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], - vert_idx: Field[[KDim], int32], + k: Field[[KDim], int32], elev: int32, ) -> Field[[CellDim, KDim], float]: - - vert_idx = broadcast(vert_idx, (CellDim, KDim)) + k = broadcast(k, (CellDim, KDim)) z_slope = where( - vert_idx == elev, + k == elev, _face_val_ppm_stencil_01b(p_cc, p_cellhgt_mc_now), _face_val_ppm_stencil_01a(p_cc, p_cellhgt_mc_now), ) @@ -77,14 +74,14 @@ def _face_val_ppm_stencil_01( def face_val_ppm_stencil_01( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], - vert_idx: Field[[KDim], int32], + k: Field[[KDim], int32], elev: int32, z_slope: Field[[CellDim, KDim], float], ): _face_val_ppm_stencil_01( p_cc, p_cellhgt_mc_now, - vert_idx, + k, elev, out=z_slope, ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py index 1fc1cf235..72579f6bf 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py @@ -22,7 +22,6 @@ def _face_val_ppm_stencil_02a( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - p_face = p_cc * (1.0 - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + ( p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) ) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1])) @@ -34,7 +33,6 @@ def _face_val_ppm_stencil_02a( def _face_val_ppm_stencil_02b( p_cc: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - p_face = p_cc return p_face @@ -43,7 +41,6 @@ def _face_val_ppm_stencil_02b( def _face_val_ppm_stencil_02c( p_cc: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - p_face = p_cc(Koff[-1]) return p_face @@ -53,24 +50,23 @@ def _face_val_ppm_stencil_02( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], p_face_in: Field[[CellDim, KDim], float], - vert_idx: Field[[KDim], int32], + k: Field[[KDim], int32], slev: int32, elev: int32, slevp1: int32, elevp1: int32, ) -> Field[[CellDim, KDim], float]: - - vert_idx = broadcast(vert_idx, (CellDim, KDim)) + k = broadcast(k, (CellDim, KDim)) p_face = where( - (vert_idx == slevp1) | (vert_idx == elev), + (k == slevp1) | (k == elev), _face_val_ppm_stencil_02a(p_cc, p_cellhgt_mc_now), p_face_in, ) - p_face = where((vert_idx == slev), _face_val_ppm_stencil_02b(p_cc), p_face) + p_face = where((k == slev), _face_val_ppm_stencil_02b(p_cc), p_face) - p_face = where((vert_idx == elevp1), _face_val_ppm_stencil_02c(p_cc), p_face) + p_face = where((k == elevp1), _face_val_ppm_stencil_02c(p_cc), p_face) return p_face @@ -80,7 +76,7 @@ def face_val_ppm_stencil_02( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], p_face_in: Field[[CellDim, KDim], float], - vert_idx: Field[[KDim], int32], + k: Field[[KDim], int32], slev: int32, elev: int32, slevp1: int32, @@ -91,7 +87,7 @@ def face_val_ppm_stencil_02( p_cc, p_cellhgt_mc_now, p_face_in, - vert_idx, + k, slev, elev, slevp1, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py index 59a728184..3653de1a5 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py @@ -22,7 +22,6 @@ def _face_val_ppm_stencil_02a( p_cc: Field[[CellDim, KDim], float], p_cellhgt_mc_now: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - p_face = p_cc * (1.0 - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + ( p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) ) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1])) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py index bab061cc4..d1a4d1242 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py @@ -21,7 +21,6 @@ def _face_val_ppm_stencil_02b( p_cc: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - p_face = p_cc return p_face diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py index 54821c903..4bb148785 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py @@ -21,7 +21,6 @@ def _face_val_ppm_stencil_02c( p_cc: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - p_face = p_cc(Koff[-1]) return p_face diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py index df6d76839..884937b7c 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py @@ -23,7 +23,6 @@ def _face_val_ppm_stencil_05( p_cellhgt_mc_now: Field[[CellDim, KDim], float], z_slope: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - zgeo1 = p_cellhgt_mc_now(Koff[-1]) / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) zgeo2 = 1.0 / ( p_cellhgt_mc_now(Koff[-2]) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py index 7457dc3eb..6047efeaa 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py @@ -42,7 +42,6 @@ def _hflux_ffsl_hybrid_stencil_01a( z_quad_vector_sum0_10: Field[[EdgeDim, KDim], float], patch0_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - p_out_e_hybrid_1a = ( where( patch0_cell_rel_idx_dsl == int32(1), diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py index fd8ef19b2..6f6477cd8 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py @@ -23,7 +23,6 @@ def _hflux_ffsl_hybrid_stencil_02( p_mass_flx_e: Field[[EdgeDim, KDim], float], z_dreg_area: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area return p_out_e_hybrid_2 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py index a7184d105..0a5cf5fc3 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py @@ -53,7 +53,6 @@ def _hflx_limiter_mo_stencil_03a( z_min: Field[[CellDim, KDim], float], dbl_eps: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) @@ -71,7 +70,6 @@ def _hflx_limiter_mo_stencil_03( z_tracer_new_low: Field[[CellDim, KDim], float], dbl_eps: float, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - z_max, z_min = _hflx_limiter_mo_stencil_03_min_max( z_tracer_max, z_tracer_min, beta_fct, r_beta_fct ) @@ -100,7 +98,6 @@ def hflx_limiter_mo_stencil_03( r_p: Field[[CellDim, KDim], float], r_m: Field[[CellDim, KDim], float], ): - _hflx_limiter_mo_stencil_03( z_tracer_max, z_tracer_min, diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py index d645251f6..af905bc92 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py @@ -25,7 +25,6 @@ def _hflx_limiter_mo_stencil_05( r_m: Field[[CellDim, KDim], float], r_p: Field[[CellDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - z_signum = where((z_anti > 0.0), 1.0, -1.0) r_frac = 0.5 * ( diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py index 957ee74a0..55898089a 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py @@ -72,7 +72,6 @@ def _prep_gauss_quadrature_c_list_stencil( Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float], ]: - z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py index b16a6d96a..6dfb719ab 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py @@ -50,7 +50,6 @@ def _recon_lsq_cell_c_svd_stencil( Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], ]: - p_coeff_10_dsl = ( lsq_pseudoinv_9(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + lsq_pseudoinv_9(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py index 50a1be5a9..1d2e82251 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py @@ -44,7 +44,6 @@ def _upwind_hflux_miura3_stencil_01( p_mass_flx_e: Field[[EdgeDim, KDim], float], cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - p_out_e_miura3 = ( ( where( diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py index eef9ba61e..098617d11 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py @@ -28,7 +28,6 @@ def _upwind_hflux_miura_cycl_stencil_01( p_mass_flx_e: Field[[EdgeDim, KDim], float], cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - z_tracer_mflx_dsl = ( where( cell_rel_idx_dsl == int32(1), diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py index b0026495f..502fc5f57 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py @@ -33,7 +33,6 @@ def _upwind_hflux_miura_cycl_stencil_02( Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], ]: - z_rhofluxdiv_c_out = ( neighbor_sum(p_mass_flx_e(C2E) * geofac_div, axis=C2EDim) if nsub == int32(1) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py index 3637ba94a..5e828de05 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py @@ -28,7 +28,6 @@ def _upwind_hflux_miura_stencil_01( p_mass_flx_e: Field[[EdgeDim, KDim], float], cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], ) -> Field[[EdgeDim, KDim], float]: - p_out_e = ( where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) + distv_bary_1 diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py index cba045918..be7f31fd3 100644 --- a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py @@ -26,7 +26,6 @@ def _v_limit_prbl_sm_stencil_02( p_face: Field[[CellDim, KDim], float], p_cc: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - q_face_up, q_face_low = where( l_limit != int32(0), where( diff --git a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py index 5e7a0f08d..0d060027e 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_divide_flux_area_list_stencil_01.py @@ -39,7 +39,6 @@ def ccw( p2_lon: np.array, p2_lat: np.array, ) -> np.array: - dx1 = p1_lon - p0_lon dy1 = p1_lat - p0_lat @@ -65,7 +64,6 @@ def lintersect( line2_p2_lon: np.array, line2_p2_lat: np.array, ) -> np.array: - intersect1 = ccw( line1_p1_lon, line1_p1_lat, @@ -112,7 +110,6 @@ def line_intersect( line2_p2_lon: np.array, line2_p2_lat: np.array, ) -> tuple[np.array]: - m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py index 7184221b2..8149e7fd2 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_01.py @@ -24,10 +24,10 @@ def face_val_ppm_stencil_01_numpy( p_cc: np.array, p_cellhgt_mc_now: np.array, - vert_idx: np.array, + k: np.array, elev: int32, ): - # this is a comment: vert_idx = np.broadcast_to(vert_idx, p_cc.shape) + # this is a comment: k = np.broadcast_to(k, p_cc.shape) # 01a zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( @@ -55,7 +55,7 @@ def face_val_ppm_stencil_01_numpy( + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 ) - z_slope = np.where(vert_idx[1:-1] < elev, z_slope_a, z_slope_b) + z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b) return z_slope @@ -65,22 +65,22 @@ def test_face_val_ppm_stencil_01(backend): p_cc = random_field(grid, CellDim, KDim, extend={KDim: 1}) p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1}) - vert_idx = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32)) - elev = vert_idx[-2] + k = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32)) + elev = k[-2] z_slope = random_field(grid, CellDim, KDim) ref = face_val_ppm_stencil_01_numpy( p_cc.asnumpy(), p_cellhgt_mc_now.asnumpy(), - vert_idx.asnumpy(), + k.asnumpy(), elev, ) face_val_ppm_stencil_01.with_backend(backend)( p_cc, p_cellhgt_mc_now, - vert_idx, + k, elev, z_slope, offset_provider={"Koff": KDim}, diff --git a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py index c1370dcb1..160d986de 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_face_val_ppm_stencil_02.py @@ -25,7 +25,7 @@ def face_val_ppm_stencil_02_numpy( p_cc: np.array, p_cellhgt_mc_now: np.array, p_face_in: np.array, - vert_idx: np.array, + k: np.array, slev: int32, elev: int32, slevp1: int32, @@ -37,9 +37,9 @@ def face_val_ppm_stencil_02_numpy( p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:]) ) * ((p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1]) - p_face = np.where((vert_idx == slevp1) | (vert_idx == elev), p_face_a, p_face_in) - p_face = np.where((vert_idx == slev), p_cc, p_face) - p_face[:, 1:] = np.where((vert_idx[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 p_face @@ -51,18 +51,18 @@ def test_face_val_ppm_stencil_02(backend): p_face_in = random_field(grid, CellDim, KDim) p_face = random_field(grid, CellDim, KDim) - vert_idx = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) + k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32)) slev = int32(1) slevp1 = slev + int32(1) - elev = vert_idx[-3] + elev = k[-3] elevp1 = elev + int32(1) ref = face_val_ppm_stencil_02_numpy( p_cc.asnumpy(), p_cellhgt_mc_now.asnumpy(), p_face_in.asnumpy(), - vert_idx.asnumpy(), + k.asnumpy(), slev, elev, slevp1, @@ -73,7 +73,7 @@ def test_face_val_ppm_stencil_02(backend): p_cc, p_cellhgt_mc_now, p_face_in, - vert_idx, + k, slev, elev, slevp1, diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py index f7b190a95..8d6e24149 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -65,7 +65,6 @@ def prep_gauss_quadrature_c_list_stencil_numpy( eps: float, p_dreg_area_in: np.array, ) -> tuple[np.ndarray]: - z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 diff --git a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py index 99ab8100a..9c521c56c 100644 --- a/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py +++ b/model/atmosphere/advection/tests/stencil_tests/test_prep_gauss_quadrature_c_stencil.py @@ -62,7 +62,6 @@ def prep_gauss_quadrature_c_stencil_numpy( dbl_eps: float, eps: float, ) -> tuple[np.ndarray]: - z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py index 144e000ed..23d9adba8 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion.py @@ -43,8 +43,8 @@ zero_field, ) from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn -from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance import ( - apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance, +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence import ( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence, ) from icon4py.model.atmosphere.diffusion.stencils.calculate_diagnostic_quantities_for_turbulence import ( calculate_diagnostic_quantities_for_turbulence, @@ -731,11 +731,11 @@ def _do_diffusion_step( handle_edge_comm = self._exchange.exchange(EdgeDim, prognostic_state.vn) log.debug( - "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): start" + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): start" ) # TODO (magdalena) get rid of this copying. So far passing an empty buffer instead did not verify? copy_field.with_backend(backend)(prognostic_state.w, self.w_tmp, offset_provider={}) - apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.with_backend(backend)( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.with_backend(backend)( area=self.cell_params.area, geofac_n2s=self.interpolation_state.geofac_n2s, geofac_grg_x=self.interpolation_state.geofac_grg_x, @@ -762,7 +762,7 @@ def _do_diffusion_step( }, ) log.debug( - "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance): end" + "running stencils 07 08 09 10 (apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence): end" ) log.debug( diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py index 62619743a..0f4a001f1 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_vn.py @@ -42,7 +42,7 @@ def _apply_diffusion_to_vn( diff_multfac_vn: Field[[KDim], wpfloat], nudgecoeff_e: Field[[EdgeDim], wpfloat], vn: Field[[EdgeDim, KDim], wpfloat], - horz_idx: Field[[EdgeDim], int32], + edge: Field[[EdgeDim], int32], nudgezone_diff: vpfloat, fac_bdydiff_v: wpfloat, start_2nd_nudge_line_idx_e: int32, @@ -61,7 +61,7 @@ def _apply_diffusion_to_vn( # TODO: Use if-else statement instead vn = ( where( - start_2nd_nudge_line_idx_e <= horz_idx, + start_2nd_nudge_line_idx_e <= edge, _apply_nabla2_and_nabla4_to_vn( area_edge, kh_smag_e, @@ -76,7 +76,7 @@ def _apply_diffusion_to_vn( ) if limited_area else where( - start_2nd_nudge_line_idx_e <= horz_idx, + start_2nd_nudge_line_idx_e <= edge, _apply_nabla2_and_nabla4_global_to_vn( area_edge, kh_smag_e, @@ -106,7 +106,7 @@ def apply_diffusion_to_vn( diff_multfac_vn: Field[[KDim], wpfloat], nudgecoeff_e: Field[[EdgeDim], wpfloat], vn: Field[[EdgeDim, KDim], wpfloat], - horz_idx: Field[[EdgeDim], int32], + edge: Field[[EdgeDim], int32], nudgezone_diff: vpfloat, fac_bdydiff_v: wpfloat, start_2nd_nudge_line_idx_e: int32, @@ -129,7 +129,7 @@ def apply_diffusion_to_vn( diff_multfac_vn, nudgecoeff_e, vn, - horz_idx, + edge, nudgezone_diff, fac_bdydiff_v, start_2nd_nudge_line_idx_e, diff --git a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py similarity index 88% rename from model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py rename to model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py index e7f90c955..734b61cfd 100644 --- a/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/stencils/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -29,7 +29,7 @@ @field_operator -def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( +def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence( area: Field[[CellDim], wpfloat], geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat], geofac_grg_x: Field[[CellDim, C2E2CODim], wpfloat], @@ -39,8 +39,8 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( dwdy: Field[[CellDim, KDim], vpfloat], diff_multfac_w: wpfloat, diff_multfac_n2w: Field[[KDim], wpfloat], - vert_idx: Field[[KDim], int32], - horz_idx: Field[[CellDim], int32], + k: Field[[KDim], int32], + cell: Field[[CellDim], int32], nrdmax: int32, interior_idx: int32, halo_idx: int32, @@ -49,10 +49,10 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat], ]: - vert_idx = broadcast(vert_idx, (CellDim, KDim)) + k = broadcast(k, (CellDim, KDim)) dwdx, dwdy = where( - int32(0) < vert_idx, + int32(0) < k, _calculate_horizontal_gradients_for_turbulence(w_old, geofac_grg_x, geofac_grg_y), (dwdx, dwdy), ) @@ -60,16 +60,13 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( z_nabla2_c = _calculate_nabla2_for_w(w_old, geofac_n2s) w = where( - (interior_idx <= horz_idx) & (horz_idx < halo_idx), + (interior_idx <= cell) & (cell < halo_idx), _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), w_old, ) w = where( - (int32(0) < vert_idx) - & (vert_idx < nrdmax) - & (interior_idx <= horz_idx) - & (horz_idx < halo_idx), + (int32(0) < k) & (k < nrdmax) & (interior_idx <= cell) & (cell < halo_idx), _apply_nabla2_to_w_in_upper_damping_layer(w, diff_multfac_n2w, area, z_nabla2_c), w, ) @@ -78,7 +75,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( @program -def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( +def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence( area: Field[[CellDim], wpfloat], geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat], geofac_grg_x: Field[[CellDim, C2E2CODim], wpfloat], @@ -89,8 +86,8 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( dwdy: Field[[CellDim, KDim], vpfloat], diff_multfac_w: wpfloat, diff_multfac_n2w: Field[[KDim], wpfloat], - vert_idx: Field[[KDim], int32], - horz_idx: Field[[CellDim], int32], + k: Field[[KDim], int32], + cell: Field[[CellDim], int32], nrdmax: int32, interior_idx: int32, halo_idx: int32, @@ -99,7 +96,7 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( vertical_start: int32, vertical_end: int32, ): - _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( + _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence( area, geofac_n2s, geofac_grg_x, @@ -109,8 +106,8 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( dwdy, diff_multfac_w, diff_multfac_n2w, - vert_idx, - horz_idx, + k, + cell, nrdmax, interior_idx, halo_idx, diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/__init__.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/__init__.py new file mode 100644 index 000000000..dab708955 --- /dev/null +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/__init__.py @@ -0,0 +1,33 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from typing import Final + +from packaging import version as pkg_version + + +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich and individual contributors" +__copyright__: Final = "Copyright (c) 2014-2022 ETH Zurich" +__license__: Final = "GPL-3.0-or-later" + + +__version__: Final = "0.0.6" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_theta_and_exner.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_theta_and_exner.py new file mode 100644 index 000000000..fa1e392aa --- /dev/null +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_theta_and_exner.py @@ -0,0 +1,127 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_theta_and_exner import ( + apply_diffusion_to_theta_and_exner, +) +from icon4py.model.common.dimension import C2E2CDim, CECDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + flatten_first_two_dims, + random_field, + random_mask, + unflatten_first_two_dims, + zero_field, +) + +from .test_calculate_nabla2_for_z import calculate_nabla2_for_z_numpy +from .test_calculate_nabla2_of_theta import calculate_nabla2_of_theta_numpy +from .test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy, +) +from .test_update_theta_and_exner import update_theta_and_exner_numpy + + +class TestApplyDiffusionToThetaAndExner(StencilTest): + PROGRAM = apply_diffusion_to_theta_and_exner + OUTPUTS = ("theta_v", "exner") + + @staticmethod + def reference( + grid, + kh_smag_e, + inv_dual_edge_length, + theta_v_in, + geofac_div, + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + area, + exner, + rd_o_cvd, + **kwargs, + ): + z_nabla2_e = calculate_nabla2_for_z_numpy(grid, kh_smag_e, inv_dual_edge_length, theta_v_in) + z_temp = calculate_nabla2_of_theta_numpy(grid, z_nabla2_e, geofac_div) + + geofac_n2s_nbh = unflatten_first_two_dims(geofac_n2s_nbh) + zd_vertoffset = unflatten_first_two_dims(zd_vertoffset) + + z_temp = truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( + grid, + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + theta_v_in, + z_temp, + ) + theta_v, exner = update_theta_and_exner_numpy( + grid, z_temp, area, theta_v_in, exner, rd_o_cvd + ) + return dict(theta_v=theta_v, exner=exner) + + @pytest.fixture + def input_data(self, grid): + kh_smag_e = random_field(grid, EdgeDim, KDim) + inv_dual_edge_length = random_field(grid, EdgeDim) + theta_v_in = random_field(grid, CellDim, KDim) + geofac_div = random_field(grid, CEDim) + mask = random_mask(grid, CellDim, KDim) + zd_vertoffset = zero_field(grid, CellDim, C2E2CDim, KDim, dtype=int32) + rng = np.random.default_rng() + for k in range(grid.num_levels): + # construct offsets that reach all k-levels except the last (because we are using the entries of this field with `+1`) + zd_vertoffset[:, :, k] = rng.integers( + low=0 - k, + high=grid.num_levels - k - 1, + size=(zd_vertoffset.shape[0], zd_vertoffset.shape[1]), + ) + zd_diffcoef = random_field(grid, CellDim, KDim) + geofac_n2s_c = random_field(grid, CellDim) + geofac_n2s_nbh = random_field(grid, CellDim, C2E2CDim) + vcoef = random_field(grid, CellDim, C2E2CDim, KDim) + area = random_field(grid, CellDim) + theta_v = random_field(grid, CellDim, KDim) + exner = random_field(grid, CellDim, KDim) + rd_o_cvd = 5.0 + + vcoef_new = flatten_first_two_dims(CECDim, KDim, field=vcoef) + zd_vertoffset_new = flatten_first_two_dims(CECDim, KDim, field=zd_vertoffset) + geofac_n2s_nbh_new = flatten_first_two_dims(CECDim, field=geofac_n2s_nbh) + + return dict( + kh_smag_e=kh_smag_e, + inv_dual_edge_length=inv_dual_edge_length, + theta_v_in=theta_v_in, + geofac_div=geofac_div, + mask=mask, + zd_vertoffset=zd_vertoffset_new, + zd_diffcoef=zd_diffcoef, + geofac_n2s_c=geofac_n2s_c, + geofac_n2s_nbh=geofac_n2s_nbh_new, + vcoef=vcoef_new, + area=area, + theta_v=theta_v, + exner=exner, + rd_o_cvd=rd_o_cvd, + ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_vn.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_vn.py new file mode 100644 index 000000000..e121b428c --- /dev/null +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_vn.py @@ -0,0 +1,151 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_vn import apply_diffusion_to_vn +from icon4py.model.atmosphere.dycore.state_utils.utils import indices_field +from icon4py.model.common.dimension import E2C2VDim, ECVDim, EdgeDim, KDim, VertexDim +from icon4py.model.common.test_utils.helpers import StencilTest, as_1D_sparse_field, random_field + +from .test_apply_nabla2_and_nabla4_global_to_vn import apply_nabla2_and_nabla4_global_to_vn_numpy +from .test_apply_nabla2_and_nabla4_to_vn import apply_nabla2_and_nabla4_to_vn_numpy +from .test_apply_nabla2_to_vn_in_lateral_boundary import ( + apply_nabla2_to_vn_in_lateral_boundary_numpy, +) +from .test_calculate_nabla4 import calculate_nabla4_numpy + + +class TestApplyDiffusionToVn(StencilTest): + PROGRAM = apply_diffusion_to_vn + OUTPUTS = ("vn",) + + @staticmethod + def reference( + grid, + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, + area_edge, + kh_smag_e, + diff_multfac_vn, + nudgecoeff_e, + vn, + edge, + nudgezone_diff, + fac_bdydiff_v, + start_2nd_nudge_line_idx_e, + limited_area, + **kwargs, + ): + + z_nabla4_e2 = calculate_nabla4_numpy( + grid, + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, + ) + + condition = start_2nd_nudge_line_idx_e <= edge[:, np.newaxis] + + if limited_area: + vn = np.where( + condition, + apply_nabla2_and_nabla4_to_vn_numpy( + grid, + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + ), + apply_nabla2_to_vn_in_lateral_boundary_numpy( + grid, z_nabla2_e, area_edge, vn, fac_bdydiff_v + ), + ) + else: + vn = np.where( + condition, + apply_nabla2_and_nabla4_global_to_vn_numpy( + grid, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn + ), + vn, + ) + + return dict(vn=vn) + + @pytest.fixture + def input_data(self, grid): + edge = indices_field(EdgeDim, grid, is_halfdim=False, dtype=int32) + + u_vert = random_field(grid, VertexDim, KDim) + v_vert = random_field(grid, VertexDim, KDim) + + primal_normal_vert_v1 = random_field(grid, EdgeDim, E2C2VDim) + primal_normal_vert_v2 = random_field(grid, EdgeDim, E2C2VDim) + + primal_normal_vert_v1_new = as_1D_sparse_field(primal_normal_vert_v1, ECVDim) + primal_normal_vert_v2_new = as_1D_sparse_field(primal_normal_vert_v2, ECVDim) + + inv_vert_vert_length = random_field(grid, EdgeDim) + inv_primal_edge_length = random_field(grid, EdgeDim) + + area_edge = random_field(grid, EdgeDim) + kh_smag_e = random_field(grid, EdgeDim, KDim) + z_nabla2_e = random_field(grid, EdgeDim, KDim) + diff_multfac_vn = random_field(grid, KDim) + vn = random_field(grid, EdgeDim, KDim) + nudgecoeff_e = random_field(grid, EdgeDim) + + limited_area = True + fac_bdydiff_v = 5.0 + nudgezone_diff = 9.0 + + start_2nd_nudge_line_idx_e = 6 + + return dict( + u_vert=u_vert, + v_vert=v_vert, + primal_normal_vert_v1=primal_normal_vert_v1_new, + primal_normal_vert_v2=primal_normal_vert_v2_new, + z_nabla2_e=z_nabla2_e, + inv_vert_vert_length=inv_vert_vert_length, + inv_primal_edge_length=inv_primal_edge_length, + area_edge=area_edge, + kh_smag_e=kh_smag_e, + diff_multfac_vn=diff_multfac_vn, + nudgecoeff_e=nudgecoeff_e, + vn=vn, + edge=edge, + nudgezone_diff=nudgezone_diff, + fac_bdydiff_v=fac_bdydiff_v, + start_2nd_nudge_line_idx_e=start_2nd_nudge_line_idx_e, + limited_area=limited_area, + horizontal_start=0, + horizontal_end=grid.num_edges, + vertical_start=0, + vertical_end=grid.num_levels, + ) diff --git a/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py new file mode 100644 index 000000000..a117d6a5e --- /dev/null +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence.py @@ -0,0 +1,132 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.diffusion.stencils.apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence import ( + apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence, +) +from icon4py.model.common.dimension import C2E2CODim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + +from .test_apply_nabla2_to_w import apply_nabla2_to_w_numpy +from .test_apply_nabla2_to_w_in_upper_damping_layer import ( + apply_nabla2_to_w_in_upper_damping_layer_numpy, +) +from .test_calculate_horizontal_gradients_for_turbulence import ( + calculate_horizontal_gradients_for_turbulence_numpy, +) +from .test_calculate_nabla2_for_w import calculate_nabla2_for_w_numpy + + +class TestApplyDiffusionToWAndComputeHorizontalGradientsForTurbulence(StencilTest): + PROGRAM = apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence + OUTPUTS = ("w", "dwdx", "dwdy") + + @staticmethod + def reference( + grid, + area, + geofac_n2s, + geofac_grg_x, + geofac_grg_y, + w_old, + dwdx, + dwdy, + diff_multfac_w, + diff_multfac_n2w, + k, + cell, + nrdmax, + interior_idx, + halo_idx, + **kwargs, + ): + reshaped_k = k[np.newaxis, :] + reshaped_cell = cell[:, np.newaxis] + + dwdx, dwdy = np.where( + 0 < reshaped_k, + calculate_horizontal_gradients_for_turbulence_numpy( + grid, w_old, geofac_grg_x, geofac_grg_y + ), + (dwdx, dwdy), + ) + + z_nabla2_c = calculate_nabla2_for_w_numpy(grid, w_old, geofac_n2s) + + w = np.where( + (interior_idx <= reshaped_cell) & (reshaped_cell < halo_idx), + apply_nabla2_to_w_numpy(grid, area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), + w_old, + ) + + w = np.where( + (0 < reshaped_k) + & (reshaped_k < nrdmax) + & (interior_idx <= reshaped_cell) + & (reshaped_cell < halo_idx), + apply_nabla2_to_w_in_upper_damping_layer_numpy(w, diff_multfac_n2w, area, z_nabla2_c), + w, + ) + return dict(w=w, dwdx=dwdx, dwdy=dwdy) + + @pytest.fixture + def input_data(self, grid): + k = zero_field(grid, KDim, dtype=int32) + for lev in range(grid.num_levels): + k[lev] = lev + + cell = zero_field(grid, CellDim, dtype=int32) + for c in range(grid.num_cells): + cell[c] = c + + nrdmax = 13 + interior_idx = 1 + halo_idx = 5 + + geofac_grg_x = random_field(grid, CellDim, C2E2CODim) + geofac_grg_y = random_field(grid, CellDim, C2E2CODim) + diff_multfac_n2w = random_field(grid, KDim) + area = random_field(grid, CellDim) + geofac_n2s = random_field(grid, CellDim, C2E2CODim) + w_old = random_field(grid, CellDim, KDim) + diff_multfac_w = 5.0 + + w = zero_field(grid, CellDim, KDim) + dwdx = zero_field(grid, CellDim, KDim) + dwdy = zero_field(grid, CellDim, KDim) + + return dict( + area=area, + geofac_n2s=geofac_n2s, + geofac_grg_x=geofac_grg_x, + geofac_grg_y=geofac_grg_y, + w_old=w_old, + diff_multfac_w=diff_multfac_w, + diff_multfac_n2w=diff_multfac_n2w, + k=k, + cell=cell, + nrdmax=nrdmax, + interior_idx=interior_idx, + halo_idx=halo_idx, + w=w, + dwdx=dwdx, + dwdy=dwdy, + horizontal_start=0, + horizontal_end=grid.num_cells, + vertical_start=0, + vertical_end=grid.num_levels, + ) diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py similarity index 80% rename from model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py index aaa34f22f..e2b8eb9f4 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_global_to_vn.py @@ -23,6 +23,15 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def apply_nabla2_and_nabla4_global_to_vn_numpy( + mesh, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn +): + area_edge = np.expand_dims(area_edge, axis=-1) + diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) + vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + return vn + + class TestApplyNabla2AndNabla4GlobalToVn(StencilTest): PROGRAM = apply_nabla2_and_nabla4_global_to_vn OUTPUTS = ("vn",) @@ -53,9 +62,9 @@ def input_data(self, grid): def reference( grid, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn, **kwargs ): - area_edge = np.expand_dims(area_edge, axis=-1) - diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) + vn = apply_nabla2_and_nabla4_global_to_vn_numpy( + grid, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn + ) return dict( vn=vn, ) diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py similarity index 55% rename from model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py index 33c4ef6a8..2111f2901 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_and_nabla4_to_vn.py @@ -15,9 +15,6 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_global_to_vn import ( - apply_nabla2_and_nabla4_global_to_vn, -) from icon4py.model.atmosphere.diffusion.stencils.apply_nabla2_and_nabla4_to_vn import ( apply_nabla2_and_nabla4_to_vn, ) @@ -26,6 +23,27 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def apply_nabla2_and_nabla4_to_vn_numpy( + grid, + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, +): + area_edge = np.expand_dims(area_edge, axis=-1) + diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) + nudgecoeff_e = np.expand_dims(nudgecoeff_e, axis=-1) + vn = vn + area_edge * ( + np.maximum(nudgezone_diff * nudgecoeff_e, kh_smag_e) * z_nabla2_e + - diff_multfac_vn * z_nabla4_e2 * area_edge + ) + return vn + + class TestApplyNabla2AndNabla4ToVn(StencilTest): PROGRAM = apply_nabla2_and_nabla4_to_vn OUTPUTS = ("vn",) @@ -69,50 +87,15 @@ def reference( nudgezone_diff, **kwargs, ): - area_edge = np.expand_dims(area_edge, axis=-1) - diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - nudgecoeff_e = np.expand_dims(nudgecoeff_e, axis=-1) - vn = vn + area_edge * ( - np.maximum(nudgezone_diff * nudgecoeff_e, kh_smag_e) * z_nabla2_e - - diff_multfac_vn * z_nabla4_e2 * area_edge + vn = apply_nabla2_and_nabla4_to_vn_numpy( + grid, + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, ) return dict(vn=vn) - - -class TestApplyNabla2AndNabla4ToVnGlobalMode(StencilTest): - PROGRAM = apply_nabla2_and_nabla4_global_to_vn - - OUTPUTS = ("vn",) - - @pytest.fixture - def input_data(self, grid): - area_edge = random_field(grid, EdgeDim, dtype=wpfloat) - kh_smag_e = random_field(grid, EdgeDim, KDim, dtype=vpfloat) - z_nabla2_e = random_field(grid, EdgeDim, KDim, dtype=wpfloat) - z_nabla4_e2 = random_field(grid, EdgeDim, KDim, dtype=vpfloat) - diff_multfac_vn = random_field(grid, KDim, dtype=wpfloat) - vn = random_field(grid, EdgeDim, KDim, dtype=wpfloat) - - return dict( - area_edge=area_edge, - kh_smag_e=kh_smag_e, - z_nabla2_e=z_nabla2_e, - z_nabla4_e2=z_nabla4_e2, - diff_multfac_vn=diff_multfac_vn, - vn=vn, - horizontal_start=int32(0), - horizontal_end=int32(grid.num_edges), - vertical_start=int32(0), - vertical_end=int32(grid.num_levels), - ) - - @staticmethod - def reference( - grid, area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn, **kwargs - ): - area_edge = np.expand_dims(area_edge, axis=-1) - diff_multfac_vn = np.expand_dims(diff_multfac_vn, axis=0) - vn = vn + area_edge * (kh_smag_e * z_nabla2_e - diff_multfac_vn * z_nabla4_e2 * area_edge) - return dict( - vn=vn, - ) diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py similarity index 82% rename from model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py index fd3b73f9c..8936a9154 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_vn_in_lateral_boundary.py @@ -23,6 +23,14 @@ from icon4py.model.common.type_alias import wpfloat +def apply_nabla2_to_vn_in_lateral_boundary_numpy( + grid, z_nabla2_e: np.array, area_edge: np.array, vn: np.array, fac_bdydiff_v +) -> np.array: + area_edge = np.expand_dims(area_edge, axis=-1) + vn = vn + (z_nabla2_e * area_edge * fac_bdydiff_v) + return vn + + class TestApplyNabla2ToVnInLateralBoundary(StencilTest): PROGRAM = apply_nabla2_to_vn_in_lateral_boundary OUTPUTS = ("vn",) @@ -48,6 +56,7 @@ def input_data(self, grid): def reference( grid, z_nabla2_e: np.array, area_edge: np.array, vn: np.array, fac_bdydiff_v, **kwargs ) -> dict: - area_edge = np.expand_dims(area_edge, axis=-1) - vn = vn + (z_nabla2_e * area_edge * fac_bdydiff_v) + vn = apply_nabla2_to_vn_in_lateral_boundary_numpy( + grid, z_nabla2_e, area_edge, vn, fac_bdydiff_v + ) return dict(vn=vn) diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_w.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w.py similarity index 77% rename from model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_w.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w.py index 3756eeb9c..03f333d94 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_w.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w.py @@ -21,6 +21,23 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def apply_nabla2_to_w_numpy( + grid, + area: np.array, + z_nabla2_c: np.array, + geofac_n2s: np.array, + w: np.array, + diff_multfac_w: float, +) -> np.array: + c2e2cO = grid.connectivities[C2E2CODim] + geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) + area = np.expand_dims(area, axis=-1) + w = w - diff_multfac_w * area * area * np.sum( + np.where((c2e2cO != -1)[:, :, np.newaxis], z_nabla2_c[c2e2cO] * geofac_n2s, 0.0), axis=1 + ) + return w + + class TestMoApplyNabla2ToW(StencilTest): PROGRAM = apply_nabla2_to_w OUTPUTS = ("w",) @@ -35,12 +52,7 @@ def reference( diff_multfac_w: float, **kwargs, ) -> dict: - c2e2cO = grid.connectivities[C2E2CODim] - geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) - area = np.expand_dims(area, axis=-1) - w = w - diff_multfac_w * area * area * np.sum( - np.where((c2e2cO != -1)[:, :, np.newaxis], z_nabla2_c[c2e2cO] * geofac_n2s, 0.0), axis=1 - ) + w = apply_nabla2_to_w_numpy(grid, area, z_nabla2_c, geofac_n2s, w, diff_multfac_w) return dict(w=w) @pytest.fixture diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py similarity index 83% rename from model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py index a2f4115bb..9b6852f43 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_apply_nabla2_to_w_in_upper_damping_layer.py @@ -23,6 +23,17 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def apply_nabla2_to_w_in_upper_damping_layer_numpy( + w: np.array, + diff_multfac_n2w: np.array, + cell_area: np.array, + z_nabla2_c: np.array, +): + cell_area = np.expand_dims(cell_area, axis=-1) + w = w + diff_multfac_n2w * cell_area * z_nabla2_c + return w + + class TestApplyNabla2ToWInUpperDampingLayer(StencilTest): PROGRAM = apply_nabla2_to_w_in_upper_damping_layer OUTPUTS = ("w",) @@ -54,6 +65,7 @@ def reference( z_nabla2_c: np.array, **kwargs, ) -> dict: - cell_area = np.expand_dims(cell_area, axis=-1) - w = w + diff_multfac_n2w * cell_area * z_nabla2_c + w = apply_nabla2_to_w_in_upper_damping_layer_numpy( + w, diff_multfac_n2w, cell_area, z_nabla2_c + ) return dict(w=w) diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_diagnostics_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_diagnostics_for_turbulence.py similarity index 70% rename from model/atmosphere/diffusion/tests/stencil_tests/test_calculate_diagnostics_for_turbulence.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_diagnostics_for_turbulence.py index 88bf7c39c..8a7d9d056 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_diagnostics_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_diagnostics_for_turbulence.py @@ -22,18 +22,25 @@ from icon4py.model.common.type_alias import vpfloat +def calculate_diagnostics_for_turbulence_numpy( + wgtfac_c: np.array, div: np.array, kh_c: np.array, div_ic, hdef_ic +) -> tuple[np.array, np.array]: + kc_offset_1 = np.roll(kh_c, shift=1, axis=1) + div_offset_1 = np.roll(div, shift=1, axis=1) + div_ic[:, 1:] = (wgtfac_c * div + (1.0 - wgtfac_c) * div_offset_1)[:, 1:] + hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[:, 1:] + return div_ic, hdef_ic + + class TestCalculateDiagnosticsForTurbulence(StencilTest): PROGRAM = calculate_diagnostics_for_turbulence OUTPUTS = ("div_ic", "hdef_ic") @staticmethod - def reference( - grid, wgtfac_c: np.array, div: np.array, kh_c: np.array, div_ic, hdef_ic - ) -> tuple[np.array, np.array]: - kc_offset_1 = np.roll(kh_c, shift=1, axis=1) - div_offset_1 = np.roll(div, shift=1, axis=1) - div_ic[:, 1:] = (wgtfac_c * div + (1.0 - wgtfac_c) * div_offset_1)[:, 1:] - hdef_ic[:, 1:] = ((wgtfac_c * kh_c + (1.0 - wgtfac_c) * kc_offset_1) ** 2)[:, 1:] + def reference(grid, wgtfac_c: np.array, div: np.array, kh_c: np.array, div_ic, hdef_ic) -> dict: + div_ic, hdef_ic = calculate_diagnostics_for_turbulence_numpy( + wgtfac_c, div, kh_c, div_ic, hdef_ic + ) return dict(div_ic=div_ic, hdef_ic=hdef_ic) @pytest.fixture diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py similarity index 77% rename from model/atmosphere/diffusion/tests/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py index 717f582ca..9449e0890 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_horizontal_gradients_for_turbulence.py @@ -23,6 +23,16 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def calculate_horizontal_gradients_for_turbulence_numpy(grid, w, geofac_grg_x, geofac_grg_y): + c2e2cO = grid.connectivities[C2E2CODim] + geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) + dwdx = np.sum(np.where((c2e2cO != -1)[:, :, np.newaxis], geofac_grg_x * w[c2e2cO], 0.0), axis=1) + + geofac_grg_y = np.expand_dims(geofac_grg_y, axis=-1) + dwdy = np.sum(np.where((c2e2cO != -1)[:, :, np.newaxis], geofac_grg_y * w[c2e2cO], 0.0), axis=1) + return dwdx, dwdy + + class TestCalculateHorizontalGradientsForTurbulence(StencilTest): PROGRAM = calculate_horizontal_gradients_for_turbulence OUTPUTS = ("dwdx", "dwdy") @@ -30,16 +40,9 @@ class TestCalculateHorizontalGradientsForTurbulence(StencilTest): @staticmethod def reference( grid, w: np.array, geofac_grg_x: np.array, geofac_grg_y: np.array, **kwargs - ) -> tuple[np.array]: - c2e2cO = grid.connectivities[C2E2CODim] - geofac_grg_x = np.expand_dims(geofac_grg_x, axis=-1) - dwdx = np.sum( - np.where((c2e2cO != -1)[:, :, np.newaxis], geofac_grg_x * w[c2e2cO], 0.0), axis=1 - ) - - geofac_grg_y = np.expand_dims(geofac_grg_y, axis=-1) - dwdy = np.sum( - np.where((c2e2cO != -1)[:, :, np.newaxis], geofac_grg_y * w[c2e2cO], 0.0), axis=1 + ) -> dict: + dwdx, dwdy = calculate_horizontal_gradients_for_turbulence_numpy( + grid, w, geofac_grg_x, geofac_grg_y ) return dict(dwdx=dwdx, dwdy=dwdy) diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py similarity index 100% rename from model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_and_smag_coefficients_for_vn.py diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_for_w.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_w.py similarity index 80% rename from model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_for_w.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_w.py index 3e15eba48..defcdc80d 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_for_w.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_w.py @@ -22,17 +22,22 @@ from icon4py.model.common.test_utils.helpers import StencilTest, constant_field, zero_field +def calculate_nabla2_for_w_numpy(grid, w: np.array, geofac_n2s: np.array): + c2e2cO = grid.connectivities[C2E2CODim] + geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) + z_nabla2_c = np.sum( + np.where((c2e2cO != -1)[:, :, np.newaxis], w[c2e2cO] * geofac_n2s, 0), axis=1 + ) + return z_nabla2_c + + class TestCalculateNabla2ForW(StencilTest): PROGRAM = calculate_nabla2_for_w OUTPUTS = ("z_nabla2_c",) @staticmethod - def reference(grid, w: np.array, geofac_n2s: np.array, **kwargs) -> np.array: - c2e2cO = grid.connectivities[C2E2CODim] - geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) - z_nabla2_c = np.sum( - np.where((c2e2cO != -1)[:, :, np.newaxis], w[c2e2cO] * geofac_n2s, 0), axis=1 - ) + def reference(grid, w: np.array, geofac_n2s: np.array, **kwargs) -> dict: + z_nabla2_c = calculate_nabla2_for_w_numpy(grid, w, geofac_n2s) return dict(z_nabla2_c=z_nabla2_c) @pytest.fixture diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_for_z.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py similarity index 81% rename from model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_for_z.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py index 4a64550d6..b44cd8a26 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_for_z.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_for_z.py @@ -24,6 +24,21 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def calculate_nabla2_for_z_numpy( + grid, + kh_smag_e: np.array, + inv_dual_edge_length: np.array, + theta_v: np.array, +) -> np.array: + inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) + + theta_v_e2c = theta_v[grid.connectivities[E2CDim]] + theta_v_weighted = theta_v_e2c[:, 1] - theta_v_e2c[:, 0] + + z_nabla2_e = kh_smag_e * inv_dual_edge_length * theta_v_weighted + return z_nabla2_e + + class TestCalculateNabla2ForZ(StencilTest): PROGRAM = calculate_nabla2_for_z OUTPUTS = ("z_nabla2_e",) @@ -36,12 +51,7 @@ def reference( theta_v: np.array, **kwargs, ) -> dict: - inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - - theta_v_e2c = theta_v[grid.connectivities[E2CDim]] - theta_v_weighted = theta_v_e2c[:, 1] - theta_v_e2c[:, 0] - - z_nabla2_e = kh_smag_e * inv_dual_edge_length * theta_v_weighted + z_nabla2_e = calculate_nabla2_for_z_numpy(grid, kh_smag_e, inv_dual_edge_length, theta_v) return dict(z_nabla2_e=z_nabla2_e) @pytest.fixture diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_of_theta.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_of_theta.py similarity index 80% rename from model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_of_theta.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_of_theta.py index 0aecbfe00..0e37ca9b7 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla2_of_theta.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla2_of_theta.py @@ -28,16 +28,21 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def calculate_nabla2_of_theta_numpy(grid, z_nabla2_e: np.array, geofac_div: np.array) -> np.array: + c2e = grid.connectivities[C2EDim] + geofac_div = geofac_div.reshape(c2e.shape) + geofac_div = np.expand_dims(geofac_div, axis=-1) + z_temp = np.sum(z_nabla2_e[c2e] * geofac_div, axis=1) # sum along edge dimension + return z_temp + + class TestCalculateNabla2OfTheta(StencilTest): PROGRAM = calculate_nabla2_of_theta OUTPUTS = ("z_temp",) @staticmethod - def reference(grid, z_nabla2_e: np.array, geofac_div: np.array, **kwargs) -> np.array: - c2e = grid.connectivities[C2EDim] - geofac_div = geofac_div.reshape(c2e.shape) - geofac_div = np.expand_dims(geofac_div, axis=-1) - z_temp = np.sum(z_nabla2_e[c2e] * geofac_div, axis=1) # sum along edge dimension + def reference(grid, z_nabla2_e: np.array, geofac_div: np.array, **kwargs) -> dict: + z_temp = calculate_nabla2_of_theta_numpy(grid, z_nabla2_e, geofac_div) return dict(z_temp=z_temp) @pytest.fixture diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla4.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla4.py similarity index 62% rename from model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla4.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla4.py index 172d6f1b2..97ba06642 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_calculate_nabla4.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_calculate_nabla4.py @@ -26,6 +26,49 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def calculate_nabla4_numpy( + grid, + u_vert: np.array, + v_vert: np.array, + primal_normal_vert_v1: np.array, + primal_normal_vert_v2: np.array, + z_nabla2_e: np.array, + inv_vert_vert_length: np.array, + inv_primal_edge_length: np.array, +) -> np.array: + e2c2v = grid.connectivities[E2C2VDim] + u_vert_e2c2v = u_vert[e2c2v] + v_vert_e2c2v = v_vert[e2c2v] + + primal_normal_vert_v1 = primal_normal_vert_v1.reshape(e2c2v.shape) + primal_normal_vert_v2 = primal_normal_vert_v2.reshape(e2c2v.shape) + + primal_normal_vert_v1 = np.expand_dims(primal_normal_vert_v1, axis=-1) + primal_normal_vert_v2 = np.expand_dims(primal_normal_vert_v2, axis=-1) + inv_vert_vert_length = np.expand_dims(inv_vert_vert_length, axis=-1) + inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) + + nabv_tang = ( + u_vert_e2c2v[:, 0] * primal_normal_vert_v1[:, 0] + + v_vert_e2c2v[:, 0] * primal_normal_vert_v2[:, 0] + ) + ( + u_vert_e2c2v[:, 1] * primal_normal_vert_v1[:, 1] + + v_vert_e2c2v[:, 1] * primal_normal_vert_v2[:, 1] + ) + nabv_norm = ( + u_vert_e2c2v[:, 2] * primal_normal_vert_v1[:, 2] + + v_vert_e2c2v[:, 2] * primal_normal_vert_v2[:, 2] + ) + ( + u_vert_e2c2v[:, 3] * primal_normal_vert_v1[:, 3] + + v_vert_e2c2v[:, 3] * primal_normal_vert_v2[:, 3] + ) + z_nabla4_e2 = 4.0 * ( + (nabv_norm - 2.0 * z_nabla2_e) * inv_vert_vert_length**2 + + (nabv_tang - 2.0 * z_nabla2_e) * inv_primal_edge_length**2 + ) + return z_nabla4_e2 + + class TestCalculateNabla4(StencilTest): PROGRAM = calculate_nabla4 OUTPUTS = ("z_nabla4_e2",) @@ -42,35 +85,15 @@ def reference( inv_primal_edge_length: np.array, **kwargs, ) -> dict: - e2c2v = grid.connectivities[E2C2VDim] - u_vert_e2c2v = u_vert[e2c2v] - v_vert_e2c2v = v_vert[e2c2v] - - primal_normal_vert_v1 = primal_normal_vert_v1.reshape(e2c2v.shape) - primal_normal_vert_v2 = primal_normal_vert_v2.reshape(e2c2v.shape) - - primal_normal_vert_v1 = np.expand_dims(primal_normal_vert_v1, axis=-1) - primal_normal_vert_v2 = np.expand_dims(primal_normal_vert_v2, axis=-1) - inv_vert_vert_length = np.expand_dims(inv_vert_vert_length, axis=-1) - inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) - - nabv_tang = ( - u_vert_e2c2v[:, 0] * primal_normal_vert_v1[:, 0] - + v_vert_e2c2v[:, 0] * primal_normal_vert_v2[:, 0] - ) + ( - u_vert_e2c2v[:, 1] * primal_normal_vert_v1[:, 1] - + v_vert_e2c2v[:, 1] * primal_normal_vert_v2[:, 1] - ) - nabv_norm = ( - u_vert_e2c2v[:, 2] * primal_normal_vert_v1[:, 2] - + v_vert_e2c2v[:, 2] * primal_normal_vert_v2[:, 2] - ) + ( - u_vert_e2c2v[:, 3] * primal_normal_vert_v1[:, 3] - + v_vert_e2c2v[:, 3] * primal_normal_vert_v2[:, 3] - ) - z_nabla4_e2 = 4.0 * ( - (nabv_norm - 2.0 * z_nabla2_e) * inv_vert_vert_length**2 - + (nabv_tang - 2.0 * z_nabla2_e) * inv_primal_edge_length**2 + z_nabla4_e2 = calculate_nabla4_numpy( + grid, + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, ) return dict(z_nabla4_e2=z_nabla4_e2) diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py similarity index 100% rename from model/atmosphere/diffusion/tests/stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_enhance_diffusion_coefficient_for_grid_point_cold_pools.py diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py similarity index 100% rename from model/atmosphere/diffusion/tests/stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_field_for_grid_point_cold_pools_enhancement.py diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py similarity index 100% rename from model/atmosphere/diffusion/tests/stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_temporary_fields_for_turbulence_diagnostics.py diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 66% rename from model/atmosphere/diffusion/tests/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 51b3bac65..0df5cdb68 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -29,6 +29,49 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( + grid, + mask: np.array, + zd_vertoffset: np.array, + zd_diffcoef: np.array, + geofac_n2s_c: np.array, + geofac_n2s_nbh: np.array, + vcoef: np.array, + theta_v: np.array, + z_temp: np.array, + **kwargs, +) -> np.array: + c2e2c = grid.connectivities[C2E2CDim] + shape = c2e2c.shape + vcoef.shape[1:] + vcoef = vcoef.reshape(shape) + zd_vertoffset = zd_vertoffset.reshape(shape) + geofac_n2s_nbh = geofac_n2s_nbh.reshape(c2e2c.shape) + full_shape = vcoef.shape + + geofac_n2s_nbh = np.expand_dims(geofac_n2s_nbh, axis=2) + + theta_v_at_zd_vertidx = np.zeros_like(vcoef) + theta_v_at_zd_vertidx_p1 = np.zeros_like(vcoef) + for ic in range(full_shape[0]): + for isparse in range(full_shape[1]): + for ik in range(full_shape[2]): + theta_v_at_zd_vertidx[ic, isparse, ik] = theta_v[ + c2e2c[ic, isparse], ik + zd_vertoffset[ic, isparse, ik] + ] + theta_v_at_zd_vertidx_p1[ic, isparse, ik] = theta_v[ + c2e2c[ic, isparse], ik + zd_vertoffset[ic, isparse, ik] + 1 + ] + + sum_over = np.sum( + geofac_n2s_nbh * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), + axis=1, + ) + + geofac_n2s_c = np.expand_dims(geofac_n2s_c, axis=1) # add KDim + z_temp = np.where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp) + return z_temp + + class TestTrulyHorizontalDiffusionNablaOfThetaOverSteepPoints(StencilTest): PROGRAM = truly_horizontal_diffusion_nabla_of_theta_over_steep_points OUTPUTS = ("z_temp",) @@ -46,35 +89,18 @@ def reference( z_temp: np.array, **kwargs, ) -> dict: - c2e2c = grid.connectivities[C2E2CDim] - shape = c2e2c.shape + vcoef.shape[1:] - vcoef = vcoef.reshape(shape) - zd_vertoffset = zd_vertoffset.reshape(shape) - geofac_n2s_nbh = geofac_n2s_nbh.reshape(c2e2c.shape) - full_shape = vcoef.shape - - geofac_n2s_nbh = np.expand_dims(geofac_n2s_nbh, axis=2) - - theta_v_at_zd_vertidx = np.zeros_like(vcoef) - theta_v_at_zd_vertidx_p1 = np.zeros_like(vcoef) - for ic in range(full_shape[0]): - for isparse in range(full_shape[1]): - for ik in range(full_shape[2]): - theta_v_at_zd_vertidx[ic, isparse, ik] = theta_v[ - c2e2c[ic, isparse], ik + zd_vertoffset[ic, isparse, ik] - ] - theta_v_at_zd_vertidx_p1[ic, isparse, ik] = theta_v[ - c2e2c[ic, isparse], ik + zd_vertoffset[ic, isparse, ik] + 1 - ] - - sum_over = np.sum( - geofac_n2s_nbh - * (vcoef * theta_v_at_zd_vertidx + (1.0 - vcoef) * theta_v_at_zd_vertidx_p1), - axis=1, - ) - geofac_n2s_c = np.expand_dims(geofac_n2s_c, axis=1) # add KDim - z_temp = np.where(mask, z_temp + zd_diffcoef * (theta_v * geofac_n2s_c + sum_over), z_temp) + z_temp = truly_horizontal_diffusion_nabla_of_theta_over_steep_points_numpy( + grid, + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + theta_v, + z_temp, + ) return dict(z_temp=z_temp) @pytest.fixture diff --git a/model/atmosphere/diffusion/tests/stencil_tests/test_update_theta_and_exner.py b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_update_theta_and_exner.py similarity index 79% rename from model/atmosphere/diffusion/tests/stencil_tests/test_update_theta_and_exner.py rename to model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_update_theta_and_exner.py index 2f999ee56..7ed4bd178 100644 --- a/model/atmosphere/diffusion/tests/stencil_tests/test_update_theta_and_exner.py +++ b/model/atmosphere/diffusion/tests/diffusion_stencil_tests/test_update_theta_and_exner.py @@ -23,6 +23,21 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def update_theta_and_exner_numpy( + mesh, + z_temp: np.array, + area: np.array, + theta_v: np.array, + exner: np.array, + rd_o_cvd: float, +) -> tuple[np.array]: + area = np.expand_dims(area, axis=0) + z_theta = theta_v + theta_v = theta_v + (np.expand_dims(area, axis=-1) * z_temp) + exner = exner * (1.0 + rd_o_cvd * (theta_v / z_theta - 1.0)) + return theta_v, exner + + class TestUpdateThetaAndExner(StencilTest): PROGRAM = update_theta_and_exner OUTPUTS = ("theta_v", "exner") @@ -34,13 +49,10 @@ def reference( area: np.array, theta_v: np.array, exner: np.array, - rd_o_cvd, + rd_o_cvd: float, **kwargs, ) -> tuple[np.array]: - area = np.expand_dims(area, axis=0) - z_theta = theta_v - theta_v = theta_v + (np.expand_dims(area, axis=-1) * z_temp) - exner = exner * (1.0 + rd_o_cvd * (theta_v / z_theta - 1.0)) + theta_v, exner = update_theta_and_exner_numpy(grid, z_temp, area, theta_v, exner, rd_o_cvd) return dict(theta_v=theta_v, exner=exner) @pytest.fixture diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py new file mode 100644 index 000000000..b0568dbf2 --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_15_to_18.py @@ -0,0 +1,213 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import Field, GridType +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import broadcast, int32, maximum, where + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_15 import ( + _mo_velocity_advection_stencil_15, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_16 import ( + _mo_velocity_advection_stencil_16, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_17 import ( + _mo_velocity_advection_stencil_17, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_18 import ( + _mo_velocity_advection_stencil_18, +) +from icon4py.model.common.dimension import C2E2CODim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.type_alias import vpfloat, wpfloat + + +@field_operator +def _fused_velocity_advection_stencil_16_to_18( + z_w_con_c: Field[[CellDim, KDim], vpfloat], + w: Field[[CellDim, KDim], wpfloat], + coeff1_dwdz: Field[[CellDim, KDim], vpfloat], + coeff2_dwdz: Field[[CellDim, KDim], vpfloat], + ddt_w_adv: Field[[CellDim, KDim], vpfloat], + e_bln_c_s: Field[[CEDim], wpfloat], + z_v_grad_w: Field[[EdgeDim, KDim], vpfloat], + levelmask: Field[[KDim], bool], + cfl_clipping: Field[[CellDim, KDim], bool], + owner_mask: Field[[CellDim], bool], + ddqz_z_half: Field[[CellDim, KDim], vpfloat], + area: Field[[CellDim], wpfloat], + geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat], + cell: Field[[CellDim], int32], + k: Field[[KDim], int32], + scalfac_exdiff: wpfloat, + cfl_w_limit: vpfloat, + dtime: wpfloat, + cell_lower_bound: int32, + cell_upper_bound: int32, + nlev: int32, + nrdmax: int32, + extra_diffu: bool, +) -> Field[[CellDim, KDim], vpfloat]: + k = broadcast(k, (CellDim, KDim)) + + ddt_w_adv = where( + (cell_lower_bound <= cell < cell_upper_bound) & (int32(1) <= k), + _mo_velocity_advection_stencil_16(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz), + ddt_w_adv, + ) + ddt_w_adv = where( + (cell_lower_bound <= cell < cell_upper_bound) & (int32(1) <= k), + _mo_velocity_advection_stencil_17(e_bln_c_s, z_v_grad_w, ddt_w_adv), + ddt_w_adv, + ) + ddt_w_adv = ( + where( + (cell_lower_bound <= cell < cell_upper_bound) + & (maximum(2, nrdmax - 2) <= k < nlev - 3), + _mo_velocity_advection_stencil_18( + levelmask, + cfl_clipping, + owner_mask, + z_w_con_c, + ddqz_z_half, + area, + geofac_n2s, + w, + ddt_w_adv, + scalfac_exdiff, + cfl_w_limit, + dtime, + ), + ddt_w_adv, + ) + if extra_diffu + else ddt_w_adv + ) + + return ddt_w_adv + + +@field_operator +def _fused_velocity_advection_stencil_15_to_18( + z_w_con_c: Field[[CellDim, KDim], vpfloat], + w: Field[[CellDim, KDim], wpfloat], + coeff1_dwdz: Field[[CellDim, KDim], vpfloat], + coeff2_dwdz: Field[[CellDim, KDim], vpfloat], + ddt_w_adv: Field[[CellDim, KDim], vpfloat], + e_bln_c_s: Field[[CEDim], wpfloat], + z_v_grad_w: Field[[EdgeDim, KDim], vpfloat], + levelmask: Field[[KDim], bool], + cfl_clipping: Field[[CellDim, KDim], bool], + owner_mask: Field[[CellDim], bool], + ddqz_z_half: Field[[CellDim, KDim], vpfloat], + area: Field[[CellDim], wpfloat], + geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat], + cell: Field[[CellDim], int32], + k: Field[[KDim], int32], + scalfac_exdiff: wpfloat, + cfl_w_limit: vpfloat, + dtime: wpfloat, + cell_lower_bound: int32, + cell_upper_bound: int32, + nlev: int32, + nrdmax: int32, + lvn_only: bool, + extra_diffu: bool, +) -> tuple[Field[[CellDim, KDim], vpfloat], Field[[CellDim, KDim], vpfloat]]: + z_w_con_c_full = _mo_velocity_advection_stencil_15(z_w_con_c) + ddt_w_adv = ( + _fused_velocity_advection_stencil_16_to_18( + z_w_con_c, + w, + coeff1_dwdz, + coeff2_dwdz, + ddt_w_adv, + e_bln_c_s, + z_v_grad_w, + levelmask, + cfl_clipping, + owner_mask, + ddqz_z_half, + area, + geofac_n2s, + cell, + k, + scalfac_exdiff, + cfl_w_limit, + dtime, + cell_lower_bound, + cell_upper_bound, + nlev, + nrdmax, + extra_diffu, + ) + if not lvn_only + else ddt_w_adv + ) + + return (z_w_con_c_full, ddt_w_adv) + + +@program(grid_type=GridType.UNSTRUCTURED) +def fused_velocity_advection_stencil_15_to_18( + z_w_con_c: Field[[CellDim, KDim], vpfloat], + w: Field[[CellDim, KDim], wpfloat], + coeff1_dwdz: Field[[CellDim, KDim], vpfloat], + coeff2_dwdz: Field[[CellDim, KDim], vpfloat], + ddt_w_adv: Field[[CellDim, KDim], vpfloat], + e_bln_c_s: Field[[CEDim], wpfloat], + z_v_grad_w: Field[[EdgeDim, KDim], vpfloat], + levelmask: Field[[KDim], bool], + cfl_clipping: Field[[CellDim, KDim], bool], + owner_mask: Field[[CellDim], bool], + ddqz_z_half: Field[[CellDim, KDim], vpfloat], + area: Field[[CellDim], wpfloat], + geofac_n2s: Field[[CellDim, C2E2CODim], wpfloat], + z_w_con_c_full: Field[[CellDim, KDim], vpfloat], + cell: Field[[CellDim], int32], + k: Field[[KDim], int32], + scalfac_exdiff: wpfloat, + cfl_w_limit: vpfloat, + dtime: wpfloat, + cell_lower_bound: int32, + cell_upper_bound: int32, + nlev: int32, + nrdmax: int32, + lvn_only: bool, + extra_diffu: bool, +): + _fused_velocity_advection_stencil_15_to_18( + z_w_con_c, + w, + coeff1_dwdz, + coeff2_dwdz, + ddt_w_adv, + e_bln_c_s, + z_v_grad_w, + levelmask, + cfl_clipping, + owner_mask, + ddqz_z_half, + area, + geofac_n2s, + cell, + k, + scalfac_exdiff, + cfl_w_limit, + dtime, + cell_lower_bound, + cell_upper_bound, + nlev, + nrdmax, + lvn_only, + extra_diffu, + out=(z_w_con_c_full, ddt_w_adv), + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_19_to_20.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_19_to_20.py new file mode 100644 index 000000000..f8dd74c15 --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_19_to_20.py @@ -0,0 +1,160 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import Field, GridType +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, maximum, where + +from icon4py.model.atmosphere.dycore.mo_math_divrot_rot_vertex_ri_dsl import ( + _mo_math_divrot_rot_vertex_ri_dsl, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_19 import ( + _mo_velocity_advection_stencil_19, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_20 import ( + _mo_velocity_advection_stencil_20, +) +from icon4py.model.common.dimension import ( + CellDim, + E2C2EODim, + E2CDim, + ECDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.type_alias import vpfloat, wpfloat + + +@field_operator +def _fused_velocity_advection_stencil_19_to_20( + vn: Field[[EdgeDim, KDim], wpfloat], + geofac_rot: Field[[VertexDim, V2EDim], wpfloat], + z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], + coeff_gradekin: Field[[ECDim], vpfloat], + z_ekinh: Field[[CellDim, KDim], vpfloat], + vt: Field[[EdgeDim, KDim], vpfloat], + f_e: Field[[EdgeDim], wpfloat], + c_lin_e: Field[[EdgeDim, E2CDim], wpfloat], + z_w_con_c_full: Field[[CellDim, KDim], vpfloat], + vn_ie: Field[[EdgeDim, KDim], vpfloat], + ddqz_z_full_e: Field[[EdgeDim, KDim], vpfloat], + levelmask: Field[[KDim], bool], + area_edge: Field[[EdgeDim], wpfloat], + tangent_orientation: Field[[EdgeDim], wpfloat], + inv_primal_edge_length: Field[[EdgeDim], wpfloat], + geofac_grdiv: Field[[EdgeDim, E2C2EODim], wpfloat], + k: Field[[KDim], int32], + cfl_w_limit: vpfloat, + scalfac_exdiff: wpfloat, + d_time: wpfloat, + extra_diffu: bool, + nlev: int32, + nrdmax: int32, +) -> Field[[EdgeDim, KDim], vpfloat]: + zeta = _mo_math_divrot_rot_vertex_ri_dsl(vn, geofac_rot) + + ddt_vn_apc = _mo_velocity_advection_stencil_19( + z_kin_hor_e, + coeff_gradekin, + z_ekinh, + zeta, + vt, + f_e, + c_lin_e, + z_w_con_c_full, + vn_ie, + ddqz_z_full_e, + ) + + ddt_vn_apc = ( + where( + maximum(2, nrdmax - 2) <= k < nlev - 3, + _mo_velocity_advection_stencil_20( + levelmask, + c_lin_e, + z_w_con_c_full, + ddqz_z_full_e, + area_edge, + tangent_orientation, + inv_primal_edge_length, + zeta, + geofac_grdiv, + vn, + ddt_vn_apc, + cfl_w_limit, + scalfac_exdiff, + d_time, + ), + ddt_vn_apc, + ) + if extra_diffu + else ddt_vn_apc + ) + + return ddt_vn_apc + + +@program(grid_type=GridType.UNSTRUCTURED) +def fused_velocity_advection_stencil_19_to_20( + vn: Field[[EdgeDim, KDim], wpfloat], + geofac_rot: Field[[VertexDim, V2EDim], wpfloat], + z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], + coeff_gradekin: Field[[ECDim], vpfloat], + z_ekinh: Field[[CellDim, KDim], vpfloat], + vt: Field[[EdgeDim, KDim], vpfloat], + f_e: Field[[EdgeDim], wpfloat], + c_lin_e: Field[[EdgeDim, E2CDim], wpfloat], + z_w_con_c_full: Field[[CellDim, KDim], vpfloat], + vn_ie: Field[[EdgeDim, KDim], vpfloat], + ddqz_z_full_e: Field[[EdgeDim, KDim], vpfloat], + levelmask: Field[[KDim], bool], + area_edge: Field[[EdgeDim], wpfloat], + tangent_orientation: Field[[EdgeDim], wpfloat], + inv_primal_edge_length: Field[[EdgeDim], wpfloat], + geofac_grdiv: Field[[EdgeDim, E2C2EODim], wpfloat], + ddt_vn_apc: Field[[EdgeDim, KDim], vpfloat], + k: Field[[KDim], int32], + cfl_w_limit: vpfloat, + scalfac_exdiff: wpfloat, + d_time: wpfloat, + extra_diffu: bool, + nlev: int32, + nrdmax: int32, +): + _fused_velocity_advection_stencil_19_to_20( + vn, + geofac_rot, + z_kin_hor_e, + coeff_gradekin, + z_ekinh, + vt, + f_e, + c_lin_e, + z_w_con_c_full, + vn_ie, + ddqz_z_full_e, + levelmask, + area_edge, + tangent_orientation, + inv_primal_edge_length, + geofac_grdiv, + k, + cfl_w_limit, + scalfac_exdiff, + d_time, + extra_diffu, + nlev, + nrdmax, + out=ddt_vn_apc, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py new file mode 100644 index 000000000..f20d94b01 --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_1_to_7.py @@ -0,0 +1,243 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import Field, GridType +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import broadcast, int32, where + +from icon4py.model.atmosphere.dycore.mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( + _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_01 import ( + _mo_velocity_advection_stencil_01, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_02 import ( + _mo_velocity_advection_stencil_02, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_03 import ( + _mo_velocity_advection_stencil_03, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_04 import ( + _mo_velocity_advection_stencil_04, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_05 import ( + _mo_velocity_advection_stencil_05, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_06 import ( + _mo_velocity_advection_stencil_06, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_07 import ( + _mo_velocity_advection_stencil_07, +) +from icon4py.model.common.dimension import CellDim, E2C2EDim, EdgeDim, KDim, V2CDim, VertexDim +from icon4py.model.common.type_alias import vpfloat, wpfloat + + +@field_operator +def _fused_velocity_advection_stencil_1_to_6( + vn: Field[[EdgeDim, KDim], wpfloat], + rbf_vec_coeff_e: Field[[EdgeDim, E2C2EDim], wpfloat], + wgtfac_e: Field[[EdgeDim, KDim], vpfloat], + ddxn_z_full: Field[[EdgeDim, KDim], vpfloat], + ddxt_z_full: Field[[EdgeDim, KDim], vpfloat], + z_w_concorr_me: Field[[EdgeDim, KDim], vpfloat], + wgtfacq_e_dsl: Field[[EdgeDim, KDim], vpfloat], + nflatlev: int32, + z_vt_ie: Field[[EdgeDim, KDim], vpfloat], + vt: Field[[EdgeDim, KDim], vpfloat], + vn_ie: Field[[EdgeDim, KDim], vpfloat], + z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], + k: Field[[KDim], int32], + nlevp1: int32, + lvn_only: bool, +) -> tuple[ + Field[[EdgeDim, KDim], vpfloat], + Field[[EdgeDim, KDim], vpfloat], + Field[[EdgeDim, KDim], vpfloat], + Field[[EdgeDim, KDim], vpfloat], +]: + vt = where( + k < nlevp1, + _mo_velocity_advection_stencil_01(vn, rbf_vec_coeff_e), + vt, + ) + + vn_ie, z_kin_hor_e = where( + 1 < k < nlevp1, + _mo_velocity_advection_stencil_02(wgtfac_e, vn, vt), + (vn_ie, z_kin_hor_e), + ) + + z_vt_ie = ( + where( + 1 < k < nlevp1, + _mo_velocity_advection_stencil_03(wgtfac_e, vt), + z_vt_ie, + ) + if not lvn_only + else z_vt_ie + ) + + (vn_ie, z_vt_ie, z_kin_hor_e) = where( + k == int32(0), + _mo_velocity_advection_stencil_05(vn, vt), + (vn_ie, z_vt_ie, z_kin_hor_e), + ) + + vn_ie = where(k == nlevp1, _mo_velocity_advection_stencil_06(wgtfacq_e_dsl, vn), vn_ie) + + z_w_concorr_me = where( + nflatlev < k < nlevp1, + _mo_velocity_advection_stencil_04(vn, ddxn_z_full, ddxt_z_full, vt), + z_w_concorr_me, + ) + + return vt, vn_ie, z_kin_hor_e, z_w_concorr_me + + +@field_operator +def _fused_velocity_advection_stencil_1_to_7( + vn: Field[[EdgeDim, KDim], wpfloat], + rbf_vec_coeff_e: Field[[EdgeDim, E2C2EDim], wpfloat], + wgtfac_e: Field[[EdgeDim, KDim], vpfloat], + ddxn_z_full: Field[[EdgeDim, KDim], vpfloat], + ddxt_z_full: Field[[EdgeDim, KDim], vpfloat], + z_w_concorr_me: Field[[EdgeDim, KDim], vpfloat], + wgtfacq_e_dsl: Field[[EdgeDim, KDim], vpfloat], + nflatlev: int32, + c_intp: Field[[VertexDim, V2CDim], wpfloat], + w: Field[[CellDim, KDim], wpfloat], + inv_dual_edge_length: Field[[EdgeDim], wpfloat], + inv_primal_edge_length: Field[[EdgeDim], wpfloat], + tangent_orientation: Field[[EdgeDim], wpfloat], + z_vt_ie: Field[[EdgeDim, KDim], wpfloat], + vt: Field[[EdgeDim, KDim], vpfloat], + vn_ie: Field[[EdgeDim, KDim], vpfloat], + z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], + z_v_grad_w: Field[[EdgeDim, KDim], vpfloat], + k: Field[[KDim], int32], + istep: int32, + nlevp1: int32, + lvn_only: bool, + edge: Field[[EdgeDim], int32], + lateral_boundary_7: int32, + halo_1: int32, +) -> tuple[ + Field[[EdgeDim, KDim], vpfloat], + Field[[EdgeDim, KDim], vpfloat], + Field[[EdgeDim, KDim], vpfloat], + Field[[EdgeDim, KDim], vpfloat], + Field[[EdgeDim, KDim], vpfloat], +]: + vt, vn_ie, z_kin_hor_e, z_w_concorr_me = ( + _fused_velocity_advection_stencil_1_to_6( + vn, + rbf_vec_coeff_e, + wgtfac_e, + ddxn_z_full, + ddxt_z_full, + z_w_concorr_me, + wgtfacq_e_dsl, + nflatlev, + z_vt_ie, + vt, + vn_ie, + z_kin_hor_e, + k, + nlevp1, + lvn_only, + ) + if istep == 1 + else (vt, vn_ie, z_kin_hor_e, z_w_concorr_me) + ) + + k = broadcast(k, (EdgeDim, KDim)) + + z_w_v = _mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl(w, c_intp) + + z_v_grad_w = ( + where( + (lateral_boundary_7 < edge < halo_1) & (k < nlevp1), + _mo_velocity_advection_stencil_07( + vn_ie, + inv_dual_edge_length, + w, + z_vt_ie, + inv_primal_edge_length, + tangent_orientation, + z_w_v, + ), + z_v_grad_w, + ) + if not lvn_only + else z_v_grad_w + ) + + return vt, vn_ie, z_kin_hor_e, z_w_concorr_me, z_v_grad_w + + +@program(grid_type=GridType.UNSTRUCTURED) +def fused_velocity_advection_stencil_1_to_7( + vn: Field[[EdgeDim, KDim], wpfloat], + rbf_vec_coeff_e: Field[[EdgeDim, E2C2EDim], wpfloat], + wgtfac_e: Field[[EdgeDim, KDim], vpfloat], + ddxn_z_full: Field[[EdgeDim, KDim], vpfloat], + ddxt_z_full: Field[[EdgeDim, KDim], vpfloat], + z_w_concorr_me: Field[[EdgeDim, KDim], vpfloat], + wgtfacq_e_dsl: Field[[EdgeDim, KDim], vpfloat], + nflatlev: int32, + c_intp: Field[[VertexDim, V2CDim], wpfloat], + w: Field[[CellDim, KDim], wpfloat], + inv_dual_edge_length: Field[[EdgeDim], wpfloat], + inv_primal_edge_length: Field[[EdgeDim], wpfloat], + tangent_orientation: Field[[EdgeDim], wpfloat], + z_vt_ie: Field[[EdgeDim, KDim], wpfloat], + vt: Field[[EdgeDim, KDim], vpfloat], + vn_ie: Field[[EdgeDim, KDim], vpfloat], + z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], + z_v_grad_w: Field[[EdgeDim, KDim], vpfloat], + k: Field[[KDim], int32], + istep: int32, + nlevp1: int32, + lvn_only: bool, + edge: Field[[EdgeDim], int32], + lateral_boundary_7: int32, + halo_1: int32, +): + _fused_velocity_advection_stencil_1_to_7( + vn, + rbf_vec_coeff_e, + wgtfac_e, + ddxn_z_full, + ddxt_z_full, + z_w_concorr_me, + wgtfacq_e_dsl, + nflatlev, + c_intp, + w, + inv_dual_edge_length, + inv_primal_edge_length, + tangent_orientation, + z_vt_ie, + vt, + vn_ie, + z_kin_hor_e, + z_v_grad_w, + k, + istep, + nlevp1, + lvn_only, + edge, + lateral_boundary_7, + halo_1, + out=(vt, vn_ie, z_kin_hor_e, z_w_concorr_me, z_v_grad_w), + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py new file mode 100644 index 000000000..7226b5542 --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/fused_velocity_advection_stencil_8_to_14.py @@ -0,0 +1,163 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +from gt4py.next.common import Field, GridType +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, maximum, where + +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_08 import ( + _mo_velocity_advection_stencil_08, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_09 import ( + _mo_velocity_advection_stencil_09, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_10 import ( + _mo_velocity_advection_stencil_10, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_11 import ( + _mo_velocity_advection_stencil_11, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_12 import ( + _mo_velocity_advection_stencil_12, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_13 import ( + _mo_velocity_advection_stencil_13, +) +from icon4py.model.atmosphere.dycore.mo_velocity_advection_stencil_14 import ( + _mo_velocity_advection_stencil_14, +) +from icon4py.model.common.dimension import CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.type_alias import vpfloat, wpfloat + + +@field_operator +def _fused_velocity_advection_stencil_8_to_14( + z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], + e_bln_c_s: Field[[CEDim], wpfloat], + z_w_concorr_me: Field[[EdgeDim, KDim], vpfloat], + wgtfac_c: Field[[CellDim, KDim], vpfloat], + w: Field[[CellDim, KDim], wpfloat], + ddqz_z_half: Field[[CellDim, KDim], vpfloat], + cfl_clipping: Field[[CellDim, KDim], bool], + pre_levelmask: Field[[CellDim, KDim], bool], + vcfl: Field[[CellDim, KDim], vpfloat], + z_w_concorr_mc: Field[[CellDim, KDim], vpfloat], + w_concorr_c: Field[[CellDim, KDim], vpfloat], + z_ekinh: Field[[CellDim, KDim], vpfloat], + k: Field[[KDim], int32], + istep: int32, + cfl_w_limit: vpfloat, + dtime: wpfloat, + nlevp1: int32, + nlev: int32, + nflatlev: int32, + nrdmax: int32, +) -> tuple[ + Field[[CellDim, KDim], vpfloat], + Field[[CellDim, KDim], bool], + Field[[CellDim, KDim], bool], + Field[[CellDim, KDim], vpfloat], + Field[[CellDim, KDim], vpfloat], +]: + z_ekinh = where( + k < nlev, + _mo_velocity_advection_stencil_08(z_kin_hor_e, e_bln_c_s), + z_ekinh, + ) + + z_w_concorr_mc = ( + where( + nflatlev < k < nlev, + _mo_velocity_advection_stencil_09(z_w_concorr_me, e_bln_c_s), + z_w_concorr_mc, + ) + if istep == 1 + else z_w_concorr_mc + ) + + w_concorr_c = ( + where( + nflatlev + 1 < k < nlev, + _mo_velocity_advection_stencil_10(z_w_concorr_mc, wgtfac_c), + w_concorr_c, + ) + if istep == 1 + else w_concorr_c + ) + + z_w_con_c = where( + k < nlevp1, + _mo_velocity_advection_stencil_11(w), + _mo_velocity_advection_stencil_12(), + ) + + z_w_con_c = where( + nflatlev + 1 < k < nlev, + _mo_velocity_advection_stencil_13(z_w_con_c, w_concorr_c), + z_w_con_c, + ) + cfl_clipping, vcfl, z_w_con_c = where( + maximum(3, nrdmax - 2) < k < nlev - 3, + _mo_velocity_advection_stencil_14(ddqz_z_half, z_w_con_c, cfl_w_limit, dtime), + (cfl_clipping, vcfl, z_w_con_c), + ) + + return z_ekinh, cfl_clipping, pre_levelmask, vcfl, z_w_con_c + + +@program(grid_type=GridType.UNSTRUCTURED) +def fused_velocity_advection_stencil_8_to_14( + z_kin_hor_e: Field[[EdgeDim, KDim], vpfloat], + e_bln_c_s: Field[[CEDim], wpfloat], + z_w_concorr_me: Field[[EdgeDim, KDim], vpfloat], + wgtfac_c: Field[[CellDim, KDim], vpfloat], + w: Field[[CellDim, KDim], wpfloat], + ddqz_z_half: Field[[CellDim, KDim], vpfloat], + cfl_clipping: Field[[CellDim, KDim], bool], + pre_levelmask: Field[[CellDim, KDim], bool], + vcfl: Field[[CellDim, KDim], vpfloat], + z_w_concorr_mc: Field[[CellDim, KDim], vpfloat], + w_concorr_c: Field[[CellDim, KDim], vpfloat], + z_ekinh: Field[[CellDim, KDim], vpfloat], + z_w_con_c: Field[[CellDim, KDim], vpfloat], + k: Field[[KDim], int32], + istep: int32, + cfl_w_limit: wpfloat, + dtime: wpfloat, + nlevp1: int32, + nlev: int32, + nflatlev: int32, + nrdmax: int32, +): + _fused_velocity_advection_stencil_8_to_14( + z_kin_hor_e, + e_bln_c_s, + z_w_concorr_me, + wgtfac_c, + w, + ddqz_z_half, + cfl_clipping, + pre_levelmask, + vcfl, + z_w_concorr_mc, + w_concorr_c, + z_ekinh, + k, + istep, + cfl_w_limit, + dtime, + nlevp1, + nlev, + nflatlev, + nrdmax, + out=(z_ekinh, cfl_clipping, pre_levelmask, vcfl, z_w_con_c), + ) diff --git a/model/atmosphere/dycore/tests/stencil_tests/__init__.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/__init__.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/__init__.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/__init__.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_airmass.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_compute_airmass.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_compute_airmass.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_fused_solve_nonhydro_stencil_39_40.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_solve_nonhydro_stencil_39_40.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_fused_solve_nonhydro_stencil_39_40.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_solve_nonhydro_stencil_39_40.py diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_15_to_18.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_15_to_18.py new file mode 100644 index 000000000..a8785512f --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_15_to_18.py @@ -0,0 +1,245 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.dycore.fused_velocity_advection_stencil_15_to_18 import ( + fused_velocity_advection_stencil_15_to_18, +) +from icon4py.model.common.dimension import C2E2CODim, C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + random_mask, + zero_field, +) + +from .test_mo_velocity_advection_stencil_15 import mo_velocity_advection_stencil_15_numpy +from .test_mo_velocity_advection_stencil_16 import mo_velocity_advection_stencil_16_numpy +from .test_mo_velocity_advection_stencil_17 import mo_velocity_advection_stencil_17_numpy +from .test_mo_velocity_advection_stencil_18 import mo_velocity_advection_stencil_18_numpy + + +class TestFusedVelocityAdvectionStencil15To18(StencilTest): + PROGRAM = fused_velocity_advection_stencil_15_to_18 + OUTPUTS = ( + "z_w_con_c_full", + "ddt_w_adv", + ) + + @staticmethod + def _fused_velocity_advection_stencil_16_to_18( + grid, + z_w_con_c, + w, + coeff1_dwdz, + coeff2_dwdz, + ddt_w_adv, + e_bln_c_s, + z_v_grad_w, + levelmask, + cfl_clipping, + owner_mask, + ddqz_z_half, + area, + geofac_n2s, + cell, + k, + scalfac_exdiff, + cfl_w_limit, + dtime, + cell_lower_bound, + cell_upper_bound, + nlev, + nrdmax, + extra_diffu, + ): + cell = cell[:, np.newaxis] + + condition1 = (cell_lower_bound <= cell) & (cell < cell_upper_bound) & (k >= 1) + + ddt_w_adv = np.where( + condition1, + mo_velocity_advection_stencil_16_numpy(z_w_con_c[:, :-1], w, coeff1_dwdz, coeff2_dwdz), + ddt_w_adv, + ) + + ddt_w_adv = np.where( + condition1, + mo_velocity_advection_stencil_17_numpy(grid, e_bln_c_s, z_v_grad_w, ddt_w_adv), + ddt_w_adv, + ) + + condition2 = ( + (cell_lower_bound <= cell) + & (cell < cell_upper_bound) + & (np.maximum(2, nrdmax - 2) <= k) + & (k < nlev - 3) + ) + + if extra_diffu: + ddt_w_adv = np.where( + condition2, + mo_velocity_advection_stencil_18_numpy( + grid, + levelmask, + cfl_clipping, + owner_mask, + z_w_con_c[:, :-1], + ddqz_z_half, + area, + geofac_n2s, + w[:, :-1], + ddt_w_adv, + scalfac_exdiff, + cfl_w_limit, + dtime, + ), + ddt_w_adv, + ) + + return ddt_w_adv + + @classmethod + def reference( + cls, + grid, + z_w_con_c, + w, + coeff1_dwdz, + coeff2_dwdz, + ddt_w_adv, + e_bln_c_s, + z_v_grad_w, + levelmask, + cfl_clipping, + owner_mask, + ddqz_z_half, + area, + geofac_n2s, + cell, + k, + scalfac_exdiff, + cfl_w_limit, + dtime, + cell_lower_bound, + cell_upper_bound, + nlev, + nrdmax, + lvn_only, + extra_diffu, + **kwargs, + ): + z_w_con_c_full = mo_velocity_advection_stencil_15_numpy(grid, z_w_con_c) + + if not lvn_only: + ddt_w_adv = cls._fused_velocity_advection_stencil_16_to_18( + grid, + z_w_con_c, + w, + coeff1_dwdz, + coeff2_dwdz, + ddt_w_adv, + e_bln_c_s, + z_v_grad_w, + levelmask, + cfl_clipping, + owner_mask, + ddqz_z_half, + area, + geofac_n2s, + cell, + k, + scalfac_exdiff, + cfl_w_limit, + dtime, + cell_lower_bound, + cell_upper_bound, + nlev, + nrdmax, + extra_diffu, + ) + + return dict(z_w_con_c_full=z_w_con_c_full, ddt_w_adv=ddt_w_adv) + + @pytest.fixture + def input_data(self, grid): + z_w_con_c = random_field(grid, CellDim, KDim, extend={KDim: 1}) + w = random_field(grid, CellDim, KDim, extend={KDim: 1}) + coeff1_dwdz = random_field(grid, CellDim, KDim) + coeff2_dwdz = random_field(grid, CellDim, KDim) + + z_v_grad_w = random_field(grid, EdgeDim, KDim) + e_bln_c_s = as_1D_sparse_field(random_field(grid, CellDim, C2EDim), CEDim) + + levelmask = random_mask(grid, KDim) + cfl_clipping = random_mask(grid, CellDim, KDim) + owner_mask = random_mask(grid, CellDim) + ddqz_z_half = random_field(grid, CellDim, KDim) + area = random_field(grid, CellDim) + geofac_n2s = random_field(grid, CellDim, C2E2CODim) + + z_w_con_c_full = zero_field(grid, CellDim, KDim) + ddt_w_adv = zero_field(grid, CellDim, KDim) + + scalfac_exdiff = 10.0 + cfl_w_limit = 3.0 + dtime = 2.0 + + k = zero_field(grid, KDim, dtype=int32) + for level in range(grid.num_levels): + k[level] = level + + cell = zero_field(grid, CellDim, dtype=int32) + for c in range(grid.num_cells): + cell[c] = c + + nlev = grid.num_levels + nrdmax = 5 + extra_diffu = True + + cell_lower_bound = 2 + cell_upper_bound = 4 + + lvn_only = False + + return dict( + z_w_con_c=z_w_con_c, + w=w, + coeff1_dwdz=coeff1_dwdz, + coeff2_dwdz=coeff2_dwdz, + ddt_w_adv=ddt_w_adv, + e_bln_c_s=e_bln_c_s, + z_v_grad_w=z_v_grad_w, + levelmask=levelmask, + cfl_clipping=cfl_clipping, + owner_mask=owner_mask, + ddqz_z_half=ddqz_z_half, + area=area, + geofac_n2s=geofac_n2s, + cell=cell, + k=k, + scalfac_exdiff=scalfac_exdiff, + cfl_w_limit=cfl_w_limit, + dtime=dtime, + cell_lower_bound=cell_lower_bound, + cell_upper_bound=cell_upper_bound, + nlev=nlev, + nrdmax=nrdmax, + lvn_only=lvn_only, + extra_diffu=extra_diffu, + z_w_con_c_full=z_w_con_c_full, + ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_19_to_20.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_19_to_20.py new file mode 100644 index 000000000..eaf3571d2 --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_19_to_20.py @@ -0,0 +1,176 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.dycore.fused_velocity_advection_stencil_19_to_20 import ( + fused_velocity_advection_stencil_19_to_20, +) +from icon4py.model.common.dimension import ( + CellDim, + E2C2EODim, + E2CDim, + ECDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + random_mask, + zero_field, +) + +from .test_mo_math_divrot_rot_vertex_ri_dsl import mo_math_divrot_rot_vertex_ri_dsl_numpy +from .test_mo_velocity_advection_stencil_19 import mo_velocity_advection_stencil_19_numpy +from .test_mo_velocity_advection_stencil_20 import mo_velocity_advection_stencil_20_numpy + + +class TestFusedVelocityAdvectionStencil19To20(StencilTest): + PROGRAM = fused_velocity_advection_stencil_19_to_20 + OUTPUTS = ("ddt_vn_apc",) + + @staticmethod + def reference( + grid, + vn, + geofac_rot, + z_kin_hor_e, + coeff_gradekin, + z_ekinh, + vt, + f_e, + c_lin_e, + z_w_con_c_full, + vn_ie, + ddqz_z_full_e, + levelmask, + area_edge, + tangent_orientation, + inv_primal_edge_length, + geofac_grdiv, + k, + cfl_w_limit, + scalfac_exdiff, + d_time, + extra_diffu, + nlev, + nrdmax, + **kwargs, + ): + zeta = mo_math_divrot_rot_vertex_ri_dsl_numpy(grid, vn, geofac_rot) + + coeff_gradekin = np.reshape(coeff_gradekin, (grid.num_edges, 2)) + + ddt_vn_apc = mo_velocity_advection_stencil_19_numpy( + grid, + z_kin_hor_e, + coeff_gradekin, + z_ekinh, + zeta, + vt, + f_e, + c_lin_e, + z_w_con_c_full, + vn_ie, + ddqz_z_full_e, + ) + + condition = (np.maximum(2, nrdmax - 2) <= k) & (k < nlev - 3) + + ddt_vn_apc_extra_diffu = mo_velocity_advection_stencil_20_numpy( + grid, + levelmask, + c_lin_e, + z_w_con_c_full, + ddqz_z_full_e, + area_edge, + tangent_orientation, + inv_primal_edge_length, + zeta, + geofac_grdiv, + vn, + ddt_vn_apc, + cfl_w_limit, + scalfac_exdiff, + d_time, + ) + + ddt_vn_apc = np.where(condition & extra_diffu, ddt_vn_apc_extra_diffu, ddt_vn_apc) + + return dict(ddt_vn_apc=ddt_vn_apc) + + @pytest.fixture + def input_data(self, grid): + z_kin_hor_e = random_field(grid, EdgeDim, KDim) + coeff_gradekin = random_field(grid, EdgeDim, E2CDim) + coeff_gradekin_new = as_1D_sparse_field(coeff_gradekin, ECDim) + z_ekinh = random_field(grid, CellDim, KDim) + vt = random_field(grid, EdgeDim, KDim) + f_e = random_field(grid, EdgeDim) + c_lin_e = random_field(grid, EdgeDim, E2CDim) + z_w_con_c_full = random_field(grid, CellDim, KDim) + vn_ie = random_field(grid, EdgeDim, KDim, extend={KDim: 1}) + ddqz_z_full_e = random_field(grid, EdgeDim, KDim) + ddt_vn_apc = zero_field(grid, EdgeDim, KDim) + levelmask = random_mask(grid, KDim, extend={KDim: 1}) + area_edge = random_field(grid, EdgeDim) + tangent_orientation = random_field(grid, EdgeDim) + inv_primal_edge_length = random_field(grid, EdgeDim) + geofac_grdiv = random_field(grid, EdgeDim, E2C2EODim) + vn = random_field(grid, EdgeDim, KDim) + geofac_rot = random_field(grid, VertexDim, V2EDim) + cfl_w_limit = 4.0 + scalfac_exdiff = 6.0 + d_time = 2.0 + + k = zero_field(grid, KDim, dtype=int32) + nlev = grid.num_levels + + for level in range(nlev): + k[level] = level + + nrdmax = 5 + extra_diffu = True + + return dict( + vn=vn, + geofac_rot=geofac_rot, + z_kin_hor_e=z_kin_hor_e, + coeff_gradekin=coeff_gradekin_new, + z_ekinh=z_ekinh, + vt=vt, + f_e=f_e, + c_lin_e=c_lin_e, + z_w_con_c_full=z_w_con_c_full, + vn_ie=vn_ie, + ddqz_z_full_e=ddqz_z_full_e, + levelmask=levelmask, + area_edge=area_edge, + tangent_orientation=tangent_orientation, + inv_primal_edge_length=inv_primal_edge_length, + geofac_grdiv=geofac_grdiv, + k=k, + cfl_w_limit=cfl_w_limit, + scalfac_exdiff=scalfac_exdiff, + d_time=d_time, + extra_diffu=extra_diffu, + nlev=nlev, + nrdmax=nrdmax, + ddt_vn_apc=ddt_vn_apc, + ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py new file mode 100644 index 000000000..0b4f40aab --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_1_to_7.py @@ -0,0 +1,255 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.dycore.fused_velocity_advection_stencil_1_to_7 import ( + fused_velocity_advection_stencil_1_to_7, +) +from icon4py.model.atmosphere.dycore.state_utils.utils import indices_field +from icon4py.model.common.dimension import CellDim, E2C2EDim, EdgeDim, KDim, V2CDim, VertexDim +from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field + +from .test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl import ( + mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy, +) +from .test_mo_velocity_advection_stencil_01 import mo_velocity_advection_stencil_01_numpy +from .test_mo_velocity_advection_stencil_02 import mo_velocity_advection_stencil_02_numpy +from .test_mo_velocity_advection_stencil_03 import mo_velocity_advection_stencil_03_numpy +from .test_mo_velocity_advection_stencil_04 import mo_velocity_advection_stencil_04_numpy +from .test_mo_velocity_advection_stencil_05 import mo_velocity_advection_stencil_05_numpy +from .test_mo_velocity_advection_stencil_06 import mo_velocity_advection_stencil_06_numpy +from .test_mo_velocity_advection_stencil_07 import mo_velocity_advection_stencil_07_numpy + + +class TestFusedVelocityAdvectionStencil1To7(StencilTest): + PROGRAM = fused_velocity_advection_stencil_1_to_7 + OUTPUTS = ( + "vt", + "vn_ie", + "z_kin_hor_e", + "z_w_concorr_me", + "z_v_grad_w", + ) + + @staticmethod + def _fused_velocity_advection_stencil_1_to_6_numpy( + grid, + vn, + rbf_vec_coeff_e, + wgtfac_e, + ddxn_z_full, + ddxt_z_full, + z_w_concorr_me, + wgtfacq_e_dsl, + nflatlev, + z_vt_ie, + vt, + vn_ie, + z_kin_hor_e, + k, + nlevp1, + lvn_only, + ): + + k = k[np.newaxis, :] + + condition1 = k < nlevp1 + vt = np.where( + condition1, mo_velocity_advection_stencil_01_numpy(grid, vn, rbf_vec_coeff_e), vt + ) + + condition2 = (1 < k) & (k < nlevp1) + vn_ie, z_kin_hor_e = np.where( + condition2, + mo_velocity_advection_stencil_02_numpy(grid, wgtfac_e, vn, vt), + (vn_ie, z_kin_hor_e), + ) + + if not lvn_only: + z_vt_ie = np.where( + condition2, mo_velocity_advection_stencil_03_numpy(grid, wgtfac_e, vt), z_vt_ie + ) + + condition3 = k == 0 + vn_ie, z_vt_ie, z_kin_hor_e = np.where( + condition3, + mo_velocity_advection_stencil_05_numpy(vn, vt), + (vn_ie, z_vt_ie, z_kin_hor_e), + ) + + condition4 = k == nlevp1 + vn_ie = np.where( + condition4, mo_velocity_advection_stencil_06_numpy(wgtfacq_e_dsl, vn), vn_ie + ) + + condition5 = (nflatlev < k) & (k < nlevp1) + z_w_concorr_me = np.where( + condition5, + mo_velocity_advection_stencil_04_numpy(vn, ddxn_z_full, ddxt_z_full, vt), + z_w_concorr_me, + ) + + return vt, vn_ie, z_kin_hor_e, z_w_concorr_me + + @classmethod + def reference( + cls, + grid, + vn, + rbf_vec_coeff_e, + wgtfac_e, + ddxn_z_full, + ddxt_z_full, + z_w_concorr_me, + wgtfacq_e_dsl, + nflatlev, + c_intp, + w, + inv_dual_edge_length, + inv_primal_edge_length, + tangent_orientation, + z_vt_ie, + vt, + vn_ie, + z_kin_hor_e, + z_v_grad_w, + k, + istep, + nlevp1, + lvn_only, + edge, + lateral_boundary_7, + halo_1, + **kwargs, + ): + + if istep == 1: + ( + vt, + vn_ie, + z_kin_hor_e, + z_w_concorr_me, + ) = cls._fused_velocity_advection_stencil_1_to_6_numpy( + grid, + vn, + rbf_vec_coeff_e, + wgtfac_e, + ddxn_z_full, + ddxt_z_full, + z_w_concorr_me, + wgtfacq_e_dsl, + nflatlev, + z_vt_ie, + vt, + vn_ie, + z_kin_hor_e, + k, + nlevp1, + lvn_only, + ) + + edge = edge[:, np.newaxis] + + z_w_v = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy(grid, w, c_intp) + + condition_mask = (lateral_boundary_7 < edge) & (edge < halo_1) & (k < nlevp1) + + if not lvn_only: + z_v_grad_w = np.where( + condition_mask, + mo_velocity_advection_stencil_07_numpy( + grid, + vn_ie, + inv_dual_edge_length, + w, + z_vt_ie, + inv_primal_edge_length, + tangent_orientation, + z_w_v, + ), + z_v_grad_w, + ) + + return dict( + vt=vt, + vn_ie=vn_ie, + z_kin_hor_e=z_kin_hor_e, + z_w_concorr_me=z_w_concorr_me, + z_v_grad_w=z_v_grad_w, + ) + + @pytest.fixture + def input_data(self, grid): + c_intp = random_field(grid, VertexDim, V2CDim) + vn = random_field(grid, EdgeDim, KDim) + rbf_vec_coeff_e = random_field(grid, EdgeDim, E2C2EDim) + vt = zero_field(grid, EdgeDim, KDim) + wgtfac_e = random_field(grid, EdgeDim, KDim) + vn_ie = zero_field(grid, EdgeDim, KDim) + z_kin_hor_e = zero_field(grid, EdgeDim, KDim) + z_vt_ie = zero_field(grid, EdgeDim, KDim) + ddxn_z_full = random_field(grid, EdgeDim, KDim) + ddxt_z_full = random_field(grid, EdgeDim, KDim) + z_w_concorr_me = zero_field(grid, EdgeDim, KDim) + inv_dual_edge_length = random_field(grid, EdgeDim) + w = random_field(grid, CellDim, KDim) + inv_primal_edge_length = random_field(grid, EdgeDim) + tangent_orientation = random_field(grid, EdgeDim) + z_v_grad_w = zero_field(grid, EdgeDim, KDim) + wgtfacq_e = random_field(grid, EdgeDim, KDim) + + k = indices_field(KDim, grid, is_halfdim=False, dtype=int32) + + edge = zero_field(grid, EdgeDim, dtype=int32) + for e in range(grid.num_edges): + edge[e] = e + + nlevp1 = grid.num_levels + 1 + nflatlev = 13 + + istep = 1 + lvn_only = False + + lateral_boundary_7 = 2 + halo_1 = 6 + + return dict( + vn=vn, + rbf_vec_coeff_e=rbf_vec_coeff_e, + wgtfac_e=wgtfac_e, + ddxn_z_full=ddxn_z_full, + ddxt_z_full=ddxt_z_full, + z_w_concorr_me=z_w_concorr_me, + wgtfacq_e_dsl=wgtfacq_e, + nflatlev=nflatlev, + c_intp=c_intp, + w=w, + inv_dual_edge_length=inv_dual_edge_length, + inv_primal_edge_length=inv_primal_edge_length, + tangent_orientation=tangent_orientation, + z_vt_ie=z_vt_ie, + vt=vt, + vn_ie=vn_ie, + z_kin_hor_e=z_kin_hor_e, + z_v_grad_w=z_v_grad_w, + k=k, + istep=istep, + nlevp1=nlevp1, + lvn_only=lvn_only, + edge=edge, + lateral_boundary_7=lateral_boundary_7, + halo_1=halo_1, + ) diff --git a/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_8_to_14.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_8_to_14.py new file mode 100644 index 000000000..8b81cf1f5 --- /dev/null +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_fused_velocity_advection_stencil_8_to_14.py @@ -0,0 +1,177 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.dycore.fused_velocity_advection_stencil_8_to_14 import ( + fused_velocity_advection_stencil_8_to_14, +) +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + StencilTest, + as_1D_sparse_field, + random_field, + random_mask, + zero_field, +) + +from .test_mo_velocity_advection_stencil_08 import mo_velocity_advection_stencil_08_numpy +from .test_mo_velocity_advection_stencil_09 import mo_velocity_advection_stencil_09_numpy +from .test_mo_velocity_advection_stencil_10 import mo_velocity_advection_stencil_10_numpy +from .test_mo_velocity_advection_stencil_11 import mo_velocity_advection_stencil_11_numpy +from .test_mo_velocity_advection_stencil_12 import mo_velocity_advection_stencil_12_numpy +from .test_mo_velocity_advection_stencil_13 import mo_velocity_advection_stencil_13_numpy +from .test_mo_velocity_advection_stencil_14 import mo_velocity_advection_stencil_14_numpy + + +class TestFusedVelocityAdvectionStencil8To14(StencilTest): + PROGRAM = fused_velocity_advection_stencil_8_to_14 + OUTPUTS = ( + "z_ekinh", + "cfl_clipping", + "pre_levelmask", + "vcfl", + "z_w_con_c", + ) + + @staticmethod + def reference( + grid, + z_kin_hor_e, + e_bln_c_s, + z_w_concorr_me, + wgtfac_c, + w, + ddqz_z_half, + cfl_clipping, + pre_levelmask, + vcfl, + z_w_concorr_mc, + w_concorr_c, + z_ekinh, + k, + istep, + cfl_w_limit, + dtime, + nlevp1, + nlev, + nflatlev, + nrdmax, + z_w_con_c, + **kwargs, + ): + + z_ekinh = np.where( + k < nlev, + mo_velocity_advection_stencil_08_numpy(grid, z_kin_hor_e, e_bln_c_s), + z_ekinh, + ) + + if istep == 1: + z_w_concorr_mc = np.where( + (nflatlev < k) & (k < nlev), + mo_velocity_advection_stencil_09_numpy(grid, z_w_concorr_me, e_bln_c_s), + z_w_concorr_mc, + ) + + w_concorr_c = np.where( + (nflatlev + 1 < k) & (k < nlev), + mo_velocity_advection_stencil_10_numpy(grid, z_w_concorr_mc, wgtfac_c), + w_concorr_c, + ) + + z_w_con_c = np.where( + k < nlevp1, + mo_velocity_advection_stencil_11_numpy(w), + mo_velocity_advection_stencil_12_numpy(z_w_con_c), + ) + + z_w_con_c = np.where( + (nflatlev + 1 < k) & (k < nlev), + mo_velocity_advection_stencil_13_numpy(z_w_con_c, w_concorr_c), + z_w_con_c, + ) + + condition = (np.maximum(3, nrdmax - 2) < k) & (k < nlev - 3) + cfl_clipping_new, vcfl_new, z_w_con_c_new = mo_velocity_advection_stencil_14_numpy( + grid, ddqz_z_half, z_w_con_c, cfl_w_limit, dtime + ) + + cfl_clipping = np.where(condition, cfl_clipping_new, cfl_clipping) + vcfl = np.where(condition, vcfl_new, vcfl) + z_w_con_c = np.where(condition, z_w_con_c_new, z_w_con_c) + + return dict( + z_ekinh=z_ekinh, + cfl_clipping=cfl_clipping, + pre_levelmask=pre_levelmask, + vcfl=vcfl, + z_w_con_c=z_w_con_c, + ) + + @pytest.fixture + def input_data(self, grid): + z_kin_hor_e = random_field(grid, EdgeDim, KDim) + e_bln_c_s = random_field(grid, CellDim, C2EDim) + z_ekinh = zero_field(grid, CellDim, KDim) + z_w_concorr_me = random_field(grid, EdgeDim, KDim) + z_w_concorr_mc = zero_field(grid, CellDim, KDim) + wgtfac_c = random_field(grid, CellDim, KDim) + w_concorr_c = zero_field(grid, CellDim, KDim) + w = random_field(grid, CellDim, KDim) + z_w_con_c = zero_field(grid, CellDim, KDim) + ddqz_z_half = random_field(grid, CellDim, KDim) + cfl_clipping = random_mask(grid, CellDim, KDim, dtype=bool) + pre_levelmask = random_mask( + grid, CellDim, KDim, dtype=bool + ) # TODO should be just a K field + + vcfl = zero_field(grid, CellDim, KDim) + cfl_w_limit = 5.0 + dtime = 9.0 + + k = zero_field(grid, KDim, dtype=int32) + for level in range(grid.num_levels): + k[level] = level + + nlevp1 = grid.num_levels + 1 + nlev = grid.num_levels + nflatlev = 13 + nrdmax = 10 + + istep = 1 + return dict( + z_kin_hor_e=z_kin_hor_e, + e_bln_c_s=as_1D_sparse_field(e_bln_c_s, CEDim), + z_w_concorr_me=z_w_concorr_me, + wgtfac_c=wgtfac_c, + w=w, + ddqz_z_half=ddqz_z_half, + cfl_clipping=cfl_clipping, + pre_levelmask=pre_levelmask, + vcfl=vcfl, + z_w_concorr_mc=z_w_concorr_mc, + w_concorr_c=w_concorr_c, + z_ekinh=z_ekinh, + k=k, + istep=istep, + cfl_w_limit=cfl_w_limit, + dtime=dtime, + nlevp1=nlevp1, + nlev=nlev, + nflatlev=nflatlev, + nrdmax=nrdmax, + z_w_con_c=z_w_con_c, + ) diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py similarity index 80% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py index 2d3e77a8c..fb1dcfe6b 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl.py @@ -23,16 +23,23 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( + grid, p_cell_in: np.array, c_intp: np.array +) -> np.array: + v2c = grid.connectivities[V2CDim] + c_intp = np.expand_dims(c_intp, axis=-1) + p_vert_out = np.sum(np.where((v2c != -1)[:, :, np.newaxis], p_cell_in[v2c] * c_intp, 0), axis=1) + return p_vert_out + + class TestMoIconInterpolationScalarCells2vertsScalarRiDsl(StencilTest): PROGRAM = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl OUTPUTS = ("p_vert_out",) @staticmethod def reference(grid, p_cell_in: np.array, c_intp: np.array, **kwargs) -> dict: - v2c = grid.connectivities[V2CDim] - c_intp = np.expand_dims(c_intp, axis=-1) - p_vert_out = np.sum( - np.where((v2c != -1)[:, :, np.newaxis], p_cell_in[v2c] * c_intp, 0), axis=1 + p_vert_out = mo_icon_interpolation_scalar_cells2verts_scalar_ri_dsl_numpy( + grid, p_cell_in, c_intp ) return dict( p_vert_out=p_vert_out, diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py similarity index 82% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py index 895c6a463..d44025b13 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_divrot_rot_vertex_ri_dsl.py @@ -23,17 +23,20 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_math_divrot_rot_vertex_ri_dsl_numpy(grid, vec_e: np.array, geofac_rot: np.array) -> np.array: + v2e = grid.connectivities[V2EDim] + geofac_rot = np.expand_dims(geofac_rot, axis=-1) + rot_vec = np.sum(np.where((v2e != -1)[:, :, np.newaxis], vec_e[v2e] * geofac_rot, 0), axis=1) + return rot_vec + + class TestMoMathDivrotRotVertexRiDsl(StencilTest): PROGRAM = mo_math_divrot_rot_vertex_ri_dsl OUTPUTS = ("rot_vec",) @staticmethod def reference(grid, vec_e: np.array, geofac_rot: np.array, **kwargs) -> dict: - v2e = grid.connectivities[V2EDim] - geofac_rot = np.expand_dims(geofac_rot, axis=-1) - rot_vec = np.sum( - np.where((v2e != -1)[:, :, np.newaxis], vec_e[v2e] * geofac_rot, 0), axis=1 - ) + rot_vec = mo_math_divrot_rot_vertex_ri_dsl_numpy(grid, vec_e, geofac_rot) return dict(rot_vec=rot_vec) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_math_gradients_grad_green_gauss_cell_dsl.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_4th_order_divdamp.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_4th_order_divdamp.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_4th_order_divdamp.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_4th_order_divdamp.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_01.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_01.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_01.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_01.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_02.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_02.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_02.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_02.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_03.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_03.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_03.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_03.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_04.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_04.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_04.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_04.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_05.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_05.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_05.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_05.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_06.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_06.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_06.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_06.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_07.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_07.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_07.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_07.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_08.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_08.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_08.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_08.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_09.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_09.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_09.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_09.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_10.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_10.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_10.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_10.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_11_lower.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_11_lower.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_11_lower.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_11_lower.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_11_upper.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_11_upper.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_11_upper.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_11_upper.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_12.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_12.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_12.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_12.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_13.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_13.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_13.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_13.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_14.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_14.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_14.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_14.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_15.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_15.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_15.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_15.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_16_fused_btraj_traj_o1.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_17.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_17.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_17.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_17.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_18.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_18.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_18.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_18.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_19.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_19.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_19.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_19.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_20.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_20.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_20.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_20.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_21.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_21.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_21.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_21.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_22.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_22.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_22.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_22.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_23.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_23.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_23.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_23.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_24.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_24.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_24.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_24.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_25.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_25.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_25.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_25.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_26.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_26.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_26.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_26.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_27.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_27.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_27.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_27.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_28.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_28.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_28.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_28.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_29.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_29.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_29.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_29.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_30.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_30.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_30.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_30.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_31.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_31.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_31.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_31.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_32.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_32.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_32.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_32.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_33.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_33.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_33.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_33.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_34.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_34.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_34.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_34.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_35.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_35.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_35.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_35.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_36.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_36.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_36.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_36.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_37.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_37.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_37.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_37.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_38.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_38.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_38.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_38.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_39.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_39.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_39.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_39.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_40.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_40.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_40.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_40.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_41.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_41.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_41.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_41.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_42.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_42.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_42.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_42.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_43.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_43.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_43.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_43.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_44.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_44.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_44.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_44.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_45.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_45.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_45.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_45.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_46.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_46.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_46.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_46.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_47.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_47.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_47.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_47.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_48.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_48.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_48.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_48.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_49.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_49.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_49.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_49.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_50.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_50.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_50.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_50.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_51.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_51.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_51.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_51.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_52.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_52.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_52.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_52.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_53.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_53.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_53.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_53.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_54.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_54.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_54.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_54.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_55.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_55.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_55.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_55.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_56_63.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_56_63.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_56_63.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_56_63.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_57.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_57.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_57.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_57.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_58.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_58.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_58.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_58.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_59.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_59.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_59.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_59.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_60.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_60.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_60.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_60.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_61.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_61.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_61.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_61.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_62.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_62.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_62.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_62.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_64.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_64.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_64.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_64.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_65.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_65.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_65.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_65.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_66.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_66.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_66.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_66.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_67.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_67.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_67.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_67.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_68.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_68.py similarity index 100% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_solve_nonhydro_stencil_68.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_solve_nonhydro_stencil_68.py diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_01.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_01.py similarity index 81% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_01.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_01.py index b6cefb8a7..33035a58f 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_01.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_01.py @@ -23,17 +23,22 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_01_numpy( + grid, vn: np.array, rbf_vec_coeff_e: np.array +) -> np.array: + rbf_vec_coeff_e = np.expand_dims(rbf_vec_coeff_e, axis=-1) + e2c2e = grid.connectivities[E2C2EDim] + vt = np.sum(np.where((e2c2e != -1)[:, :, np.newaxis], vn[e2c2e] * rbf_vec_coeff_e, 0), axis=1) + return vt + + class TestMoVelocityAdvectionStencil01(StencilTest): PROGRAM = mo_velocity_advection_stencil_01 OUTPUTS = ("vt",) @staticmethod def reference(grid, vn: np.array, rbf_vec_coeff_e: np.array, **kwargs) -> dict: - rbf_vec_coeff_e = np.expand_dims(rbf_vec_coeff_e, axis=-1) - e2c2e = grid.connectivities[E2C2EDim] - vt = np.sum( - np.where((e2c2e != -1)[:, :, np.newaxis], vn[e2c2e] * rbf_vec_coeff_e, 0), axis=1 - ) + vt = mo_velocity_advection_stencil_01_numpy(grid, vn, rbf_vec_coeff_e) return dict(vt=vt) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_02.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_02.py similarity index 73% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_02.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_02.py index 817f4fe64..5ba7a5460 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_02.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_02.py @@ -23,27 +23,37 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: + vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) + vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 + vn_ie[:, 0] = 0 + return vn_ie + + +def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: + z_kin_hor_e = 0.5 * (vn * vn + vt * vt) + z_kin_hor_e[:, 0] = 0 + return z_kin_hor_e + + +def mo_velocity_advection_stencil_02_numpy( + grid, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs +) -> tuple: + vn_ie = mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e, vn) + z_kin_hor_e = mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn, vt) + return ( + vn_ie, + z_kin_hor_e, + ) + + class TestMoVelocityAdvectionStencil02VnIe(StencilTest): PROGRAM = mo_velocity_advection_stencil_02 OUTPUTS = ("vn_ie", "z_kin_hor_e") - @staticmethod - def mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e: np.array, vn: np.array) -> np.array: - vn_ie_k_minus_1 = np.roll(vn, shift=1, axis=1) - vn_ie = wgtfac_e * vn + (1.0 - wgtfac_e) * vn_ie_k_minus_1 - vn_ie[:, 0] = 0 - return vn_ie - - @staticmethod - def mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn: np.array, vt: np.array) -> np.array: - z_kin_hor_e = 0.5 * (vn * vn + vt * vt) - z_kin_hor_e[:, 0] = 0 - return z_kin_hor_e - @classmethod def reference(cls, grid, wgtfac_e: np.array, vn: np.array, vt: np.array, **kwargs) -> dict: - vn_ie = cls.mo_velocity_advection_stencil_02_vn_ie_numpy(wgtfac_e, vn) - z_kin_hor_e = cls.mo_velocity_advection_stencil_02_z_kin_hor_e_numpy(vn, vt) + vn_ie, z_kin_hor_e = mo_velocity_advection_stencil_02_numpy(grid, wgtfac_e, vn, vt) return dict( vn_ie=vn_ie[int32(1) : int32(grid.num_cells), int32(1) : int32(grid.num_levels)], z_kin_hor_e=z_kin_hor_e[ diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_03.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_03.py similarity index 84% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_03.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_03.py index 41f64af9d..f4e32000d 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_03.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_03.py @@ -23,15 +23,22 @@ from icon4py.model.common.type_alias import vpfloat +def mo_velocity_advection_stencil_03_numpy( + grid, wgtfac_e: np.array, vt: np.array, **kwargs +) -> np.array: + vt_k_minus_1 = np.roll(vt, shift=1, axis=1) + z_vt_ie = wgtfac_e * vt + (1.0 - wgtfac_e) * vt_k_minus_1 + z_vt_ie[:, 0] = 0 + return z_vt_ie + + class TestMoVelocityAdvectionStencil03(StencilTest): PROGRAM = mo_velocity_advection_stencil_03 OUTPUTS = ("z_vt_ie",) @staticmethod def reference(grid, wgtfac_e: np.array, vt: np.array, **kwargs) -> dict: - vt_k_minus_1 = np.roll(vt, shift=1, axis=1) - z_vt_ie = wgtfac_e * vt + (1.0 - wgtfac_e) * vt_k_minus_1 - z_vt_ie[:, 0] = 0 + z_vt_ie = mo_velocity_advection_stencil_03_numpy(grid, wgtfac_e, vt) return dict( z_vt_ie=z_vt_ie[int32(1) : int32(grid.num_cells), int32(1) : int32(grid.num_levels)] ) diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_04.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_04.py similarity index 86% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_04.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_04.py index 44ac442ff..2dc2ccb76 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_04.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_04.py @@ -23,6 +23,13 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_04_numpy( + vn: np.array, ddxn_z_full: np.array, ddxt_z_full: np.array, vt: np.array +) -> np.array: + z_w_concorr_me = vn * ddxn_z_full + vt * ddxt_z_full + return z_w_concorr_me + + class TestMoVelocityAdvectionStencil04(StencilTest): PROGRAM = mo_velocity_advection_stencil_04 OUTPUTS = ("z_w_concorr_me",) @@ -36,7 +43,7 @@ def reference( vt: np.array, **kwargs, ) -> dict: - z_w_concorr_me = vn * ddxn_z_full + vt * ddxt_z_full + z_w_concorr_me = mo_velocity_advection_stencil_04_numpy(vn, ddxn_z_full, ddxt_z_full, vt) return dict(z_w_concorr_me=z_w_concorr_me) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_05.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_05.py similarity index 87% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_05.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_05.py index 39285ed3e..0c8b34279 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_05.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_05.py @@ -23,15 +23,20 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_05_numpy(vn: np.array, vt: np.array) -> tuple: + vn_ie = vn + z_vt_ie = vt + z_kin_hor_e = 0.5 * ((vn * vn) + (vt * vt)) + return vn_ie, z_vt_ie, z_kin_hor_e + + class TestMoVelocityAdvectionStencil05(StencilTest): PROGRAM = mo_velocity_advection_stencil_05 OUTPUTS = ("vn_ie", "z_vt_ie", "z_kin_hor_e") @staticmethod def reference(grid, vn: np.array, vt: np.array, **kwargs) -> dict: - vn_ie = vn - z_vt_ie = vt - z_kin_hor_e = 0.5 * ((vn * vn) + (vt * vt)) + vn_ie, z_vt_ie, z_kin_hor_e = mo_velocity_advection_stencil_05_numpy(vn, vt) return dict(vn_ie=vn_ie, z_vt_ie=z_vt_ie, z_kin_hor_e=z_kin_hor_e) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_06.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_06.py similarity index 71% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_06.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_06.py index 8dbc1ba58..ba8890ce1 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_06.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_06.py @@ -23,25 +23,29 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_06_numpy(wgtfacq_e: np.array, vn: np.array) -> np.array: + vn_k_minus_1 = np.roll(vn, shift=1, axis=1) + vn_k_minus_2 = np.roll(vn, shift=2, axis=1) + vn_k_minus_3 = np.roll(vn, shift=3, axis=1) + wgtfacq_e_k_minus_1 = np.roll(wgtfacq_e, shift=1, axis=1) + wgtfacq_e_k_minus_2 = np.roll(wgtfacq_e, shift=2, axis=1) + wgtfacq_e_k_minus_3 = np.roll(wgtfacq_e, shift=3, axis=1) + vn_ie = np.zeros_like(vn) + vn_ie[:, -1] = ( + wgtfacq_e_k_minus_1 * vn_k_minus_1 + + wgtfacq_e_k_minus_2 * vn_k_minus_2 + + wgtfacq_e_k_minus_3 * vn_k_minus_3 + )[:, -1] + return vn_ie + + class TestMoVelocityAdvectionStencil06(StencilTest): PROGRAM = mo_velocity_advection_stencil_06 OUTPUTS = ("vn_ie",) @staticmethod def reference(grid, wgtfacq_e: np.array, vn: np.array, **kwargs) -> dict: - vn_k_minus_1 = np.roll(vn, shift=1, axis=1) - vn_k_minus_2 = np.roll(vn, shift=2, axis=1) - vn_k_minus_3 = np.roll(vn, shift=3, axis=1) - wgtfacq_e_k_minus_1 = np.roll(wgtfacq_e, shift=1, axis=1) - wgtfacq_e_k_minus_2 = np.roll(wgtfacq_e, shift=2, axis=1) - wgtfacq_e_k_minus_3 = np.roll(wgtfacq_e, shift=3, axis=1) - vn_ie = np.zeros_like(vn) - vn_ie[:, -1] = ( - wgtfacq_e_k_minus_1 * vn_k_minus_1 - + wgtfacq_e_k_minus_2 * vn_k_minus_2 - + wgtfacq_e_k_minus_3 * vn_k_minus_3 - )[:, -1] - + vn_ie = mo_velocity_advection_stencil_06_numpy(wgtfacq_e, vn) return dict(vn_ie=vn_ie) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_07.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_07.py similarity index 72% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_07.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_07.py index 3952b6105..b1a7897af 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_07.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_07.py @@ -23,6 +23,33 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_07_numpy( + grid, + vn_ie: np.array, + inv_dual_edge_length: np.array, + w: np.array, + z_vt_ie: np.array, + inv_primal_edge_length: np.array, + tangent_orientation: np.array, + z_w_v: np.array, +) -> np.array: + inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) + inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) + + w_e2c = w[grid.connectivities[E2CDim]] + z_w_v_e2v = z_w_v[grid.connectivities[E2VDim]] + + red_w = w_e2c[:, 0] - w_e2c[:, 1] + red_z_w_v = z_w_v_e2v[:, 0] - z_w_v_e2v[:, 1] + + z_v_grad_w = ( + vn_ie * inv_dual_edge_length * red_w + + z_vt_ie * inv_primal_edge_length * tangent_orientation * red_z_w_v + ) + return z_v_grad_w + + class TestMoVelocityAdvectionStencil07(StencilTest): PROGRAM = mo_velocity_advection_stencil_07 OUTPUTS = ("z_v_grad_w",) @@ -39,19 +66,15 @@ def reference( z_w_v: np.array, **kwargs, ) -> dict: - inv_dual_edge_length = np.expand_dims(inv_dual_edge_length, axis=-1) - inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) - tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - - w_e2c = w[grid.connectivities[E2CDim]] - z_w_v_e2v = z_w_v[grid.connectivities[E2VDim]] - - red_w = w_e2c[:, 0] - w_e2c[:, 1] - red_z_w_v = z_w_v_e2v[:, 0] - z_w_v_e2v[:, 1] - - z_v_grad_w = ( - vn_ie * inv_dual_edge_length * red_w - + z_vt_ie * inv_primal_edge_length * tangent_orientation * red_z_w_v + z_v_grad_w = mo_velocity_advection_stencil_07_numpy( + grid, + vn_ie, + inv_dual_edge_length, + w, + z_vt_ie, + inv_primal_edge_length, + tangent_orientation, + z_w_v, ) return dict(z_v_grad_w=z_v_grad_w) diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_08.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_08.py similarity index 81% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_08.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_08.py index b824d3744..22729ebb8 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_08.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_08.py @@ -28,18 +28,25 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_08_numpy( + grid, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs +) -> np.array: + e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) + z_ekinh = np.sum( + z_kin_hor_e[grid.connectivities[C2EDim]] + * e_bln_c_s[grid.get_offset_provider("C2CE").table], + axis=1, + ) + return z_ekinh + + class TestMoVelocityAdvectionStencil08(StencilTest): PROGRAM = mo_velocity_advection_stencil_08 OUTPUTS = ("z_ekinh",) @staticmethod def reference(grid, z_kin_hor_e: np.array, e_bln_c_s: np.array, **kwargs) -> dict: - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - z_ekinh = np.sum( - z_kin_hor_e[grid.connectivities[C2EDim]] - * e_bln_c_s[grid.get_offset_provider("C2CE").table], - axis=1, - ) + z_ekinh = mo_velocity_advection_stencil_08_numpy(grid, z_kin_hor_e, e_bln_c_s) return dict(z_ekinh=z_ekinh) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_09.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_09.py similarity index 80% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_09.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_09.py index afe3ba4fa..de6d09fc7 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_09.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_09.py @@ -28,18 +28,25 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_09_numpy( + grid, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs +) -> np.array: + e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) + z_w_concorr_mc = np.sum( + z_w_concorr_me[grid.connectivities[C2EDim]] + * e_bln_c_s[grid.get_offset_provider("C2CE").table], + axis=1, + ) + return z_w_concorr_mc + + class TestMoVelocityAdvectionStencil09(StencilTest): PROGRAM = mo_velocity_advection_stencil_09 OUTPUTS = ("z_w_concorr_mc",) @staticmethod def reference(grid, z_w_concorr_me: np.array, e_bln_c_s: np.array, **kwargs) -> dict: - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - z_w_concorr_mc = np.sum( - z_w_concorr_me[grid.connectivities[C2EDim]] - * e_bln_c_s[grid.get_offset_provider("C2CE").table], - axis=1, - ) + z_w_concorr_mc = mo_velocity_advection_stencil_09_numpy(grid, z_w_concorr_me, e_bln_c_s) return dict(z_w_concorr_mc=z_w_concorr_mc) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_10.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_10.py similarity index 81% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_10.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_10.py index c8abc1850..e452928d5 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_10.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_10.py @@ -23,15 +23,22 @@ from icon4py.model.common.type_alias import vpfloat +def mo_velocity_advection_stencil_10_numpy( + grid, wgtfac_c: np.array, z_w_concorr_mc: np.array +) -> np.array: + z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) + w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 + w_concorr_c[:, 0] = 0 + return w_concorr_c + + class TestMoVelocityAdvectionStencil10(StencilTest): PROGRAM = mo_velocity_advection_stencil_10 OUTPUTS = ("w_concorr_c",) @staticmethod def reference(grid, wgtfac_c: np.array, z_w_concorr_mc: np.array, **kwargs) -> dict: - z_w_concorr_mc_k_minus_1 = np.roll(z_w_concorr_mc, shift=1, axis=1) - w_concorr_c = wgtfac_c * z_w_concorr_mc + (1.0 - wgtfac_c) * z_w_concorr_mc_k_minus_1 - w_concorr_c[:, 0] = 0 + w_concorr_c = mo_velocity_advection_stencil_10_numpy(grid, wgtfac_c, z_w_concorr_mc) return dict(w_concorr_c=w_concorr_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_11.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_11.py similarity index 90% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_11.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_11.py index 4e02b857b..106e4fc66 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_11.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_11.py @@ -23,13 +23,18 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_11_numpy(w: np.array) -> np.array: + z_w_con_c = w + return z_w_con_c + + class TestMoVelocityAdvectionStencil11(StencilTest): PROGRAM = mo_velocity_advection_stencil_11 OUTPUTS = ("z_w_con_c",) @staticmethod def reference(grid, w: np.array, **kwargs) -> dict: - z_w_con_c = w + z_w_con_c = mo_velocity_advection_stencil_11_numpy(w) return dict(z_w_con_c=z_w_con_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_12.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_12.py similarity index 87% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_12.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_12.py index fd46d6d82..db7969eee 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_12.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_12.py @@ -23,13 +23,18 @@ from icon4py.model.common.type_alias import vpfloat +def mo_velocity_advection_stencil_12_numpy(z_w_con_c: np.array) -> np.array: + z_w_con_c = np.zeros_like(z_w_con_c) + return z_w_con_c + + class TestMoVelocityAdvectionStencil12(StencilTest): PROGRAM = mo_velocity_advection_stencil_12 OUTPUTS = ("z_w_con_c",) @staticmethod def reference(grid, z_w_con_c: np.array, **kwargs) -> dict: - z_w_con_c = np.zeros_like(z_w_con_c) + z_w_con_c = mo_velocity_advection_stencil_12_numpy(z_w_con_c) return dict(z_w_con_c=z_w_con_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_13.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_13.py similarity index 87% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_13.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_13.py index b265b9082..05a11791b 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_13.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_13.py @@ -23,13 +23,18 @@ from icon4py.model.common.type_alias import vpfloat +def mo_velocity_advection_stencil_13_numpy(w_concorr_c: np.array, z_w_con_c: np.array) -> np.array: + z_w_con_c = z_w_con_c - w_concorr_c + return z_w_con_c + + class TestMoVelocityAdvectionStencil13(StencilTest): PROGRAM = mo_velocity_advection_stencil_13 OUTPUTS = ("z_w_con_c",) @staticmethod def reference(grid, w_concorr_c: np.array, z_w_con_c: np.array, **kwargs) -> dict: - z_w_con_c = z_w_con_c - w_concorr_c + z_w_con_c = mo_velocity_advection_stencil_13_numpy(w_concorr_c, z_w_con_c) return dict(z_w_con_c=z_w_con_c) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_14.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_14.py similarity index 71% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_14.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_14.py index 623d70b50..7cf6a1d5e 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_14.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_14.py @@ -28,6 +28,29 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_14_numpy( + mesh, ddqz_z_half: np.array, z_w_con_c: np.array, cfl_w_limit, dtime +) -> tuple: + num_rows, num_cols = z_w_con_c.shape + cfl_clipping = np.where( + np.abs(z_w_con_c) > cfl_w_limit * ddqz_z_half, + np.ones([num_rows, num_cols]), + np.zeros_like(z_w_con_c), + ) + num_rows, num_cols = cfl_clipping.shape + vcfl = np.where(cfl_clipping == 1.0, z_w_con_c * dtime / ddqz_z_half, 0.0) + z_w_con_c = np.where( + (cfl_clipping == 1.0) & (vcfl < -0.85), + -0.85 * ddqz_z_half / dtime, + z_w_con_c, + ) + z_w_con_c = np.where( + (cfl_clipping == 1.0) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c + ) + + return cfl_clipping, vcfl, z_w_con_c + + class TestMoVelocityAdvectionStencil14(StencilTest): PROGRAM = mo_velocity_advection_stencil_14 OUTPUTS = ("cfl_clipping", "vcfl", "z_w_con_c") @@ -36,21 +59,8 @@ class TestMoVelocityAdvectionStencil14(StencilTest): def reference( grid, ddqz_z_half: np.array, z_w_con_c: np.array, cfl_w_limit, dtime, **kwargs ) -> dict: - num_rows, num_cols = z_w_con_c.shape - cfl_clipping = np.where( - np.abs(z_w_con_c) > cfl_w_limit * ddqz_z_half, - np.ones([num_rows, num_cols]), - np.zeros_like(z_w_con_c), - ) - num_rows, num_cols = cfl_clipping.shape - vcfl = np.where(cfl_clipping == 1.0, z_w_con_c * dtime / ddqz_z_half, 0.0) - z_w_con_c = np.where( - (cfl_clipping == 1.0) & (vcfl < -0.85), - -0.85 * ddqz_z_half / dtime, - z_w_con_c, - ) - z_w_con_c = np.where( - (cfl_clipping == 1.0) & (vcfl > 0.85), 0.85 * ddqz_z_half / dtime, z_w_con_c + cfl_clipping, vcfl, z_w_con_c = mo_velocity_advection_stencil_14_numpy( + grid, ddqz_z_half, z_w_con_c, cfl_w_limit, dtime ) return dict( diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_15.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_15.py similarity index 84% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_15.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_15.py index 3e4f62d25..2d15ea1dd 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_15.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_15.py @@ -23,13 +23,18 @@ from icon4py.model.common.type_alias import vpfloat +def mo_velocity_advection_stencil_15_numpy(grid, z_w_con_c: np.array): + z_w_con_c_full = 0.5 * (z_w_con_c[:, :-1] + z_w_con_c[:, 1:]) + return z_w_con_c_full + + class TestMoVelocityAdvectionStencil15(StencilTest): PROGRAM = mo_velocity_advection_stencil_15 OUTPUTS = ("z_w_con_c_full",) @staticmethod - def reference(grid, z_w_con_c: np.array, **kwargs): - z_w_con_c_full = 0.5 * (z_w_con_c[:, :-1] + z_w_con_c[:, 1:]) + def reference(grid, z_w_con_c: np.array, **kwargs) -> dict: + z_w_con_c_full = mo_velocity_advection_stencil_15_numpy(grid, z_w_con_c) return dict(z_w_con_c_full=z_w_con_c_full) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_16.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_16.py similarity index 79% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_16.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_16.py index 9a19f007c..025d9cd24 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_16.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_16.py @@ -23,6 +23,22 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_16_numpy( + z_w_con_c: np.array, + w: np.array, + coeff1_dwdz: np.array, + coeff2_dwdz: np.array, + **kwargs, +) -> np.array: + ddt_w_adv = np.zeros_like(coeff1_dwdz) + ddt_w_adv[:, 1:] = -z_w_con_c[:, 1:] * ( + w[:, :-2] * coeff1_dwdz[:, 1:] + - w[:, 2:] * coeff2_dwdz[:, 1:] + + w[:, 1:-1] * (coeff2_dwdz[:, 1:] - coeff1_dwdz[:, 1:]) + ) + return ddt_w_adv + + class TestMoVelocityAdvectionStencil16(StencilTest): PROGRAM = mo_velocity_advection_stencil_16 OUTPUTS = ("ddt_w_adv",) @@ -36,12 +52,7 @@ def reference( coeff2_dwdz: np.array, **kwargs, ) -> dict: - ddt_w_adv = np.zeros_like(coeff1_dwdz) - ddt_w_adv[:, 1:] = -z_w_con_c[:, 1:] * ( - w[:, :-2] * coeff1_dwdz[:, 1:] - - w[:, 2:] * coeff2_dwdz[:, 1:] - + w[:, 1:-1] * (coeff2_dwdz[:, 1:] - coeff1_dwdz[:, 1:]) - ) + ddt_w_adv = mo_velocity_advection_stencil_16_numpy(z_w_con_c, w, coeff1_dwdz, coeff2_dwdz) return dict(ddt_w_adv=ddt_w_adv) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_17.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_17.py similarity index 80% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_17.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_17.py index e3f437b73..e888f25c6 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_17.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_17.py @@ -23,6 +23,17 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_17_numpy( + grid, e_bln_c_s: np.array, z_v_grad_w: np.array, ddt_w_adv: np.array, **kwargs +) -> np.array: + e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) + ddt_w_adv = ddt_w_adv + np.sum( + z_v_grad_w[grid.connectivities[C2EDim]] * e_bln_c_s[grid.get_offset_provider("C2CE").table], + axis=1, + ) + return ddt_w_adv + + class TestMoVelocityAdvectionStencil17(StencilTest): PROGRAM = mo_velocity_advection_stencil_17 OUTPUTS = ("ddt_w_adv",) @@ -31,12 +42,7 @@ class TestMoVelocityAdvectionStencil17(StencilTest): def reference( grid, e_bln_c_s: np.array, z_v_grad_w: np.array, ddt_w_adv: np.array, **kwargs ) -> dict: - e_bln_c_s = np.expand_dims(e_bln_c_s, axis=-1) - ddt_w_adv = ddt_w_adv + np.sum( - z_v_grad_w[grid.connectivities[C2EDim]] - * e_bln_c_s[grid.get_offset_provider("C2CE").table], - axis=1, - ) + ddt_w_adv = mo_velocity_advection_stencil_17_numpy(grid, e_bln_c_s, z_v_grad_w, ddt_w_adv) return dict(ddt_w_adv=ddt_w_adv) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_18.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_18.py similarity index 65% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_18.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_18.py index 7c54d9cee..1ff9221ba 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_18.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_18.py @@ -23,6 +23,54 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_18_numpy( + grid, + levmask: np.array, + cfl_clipping: np.array, + owner_mask: np.array, + z_w_con_c: np.array, + ddqz_z_half: np.array, + area: np.array, + geofac_n2s: np.array, + w: np.array, + ddt_w_adv: np.array, + scalfac_exdiff: float, + cfl_w_limit: float, + dtime: float, +) -> np.array: + levmask = np.expand_dims(levmask, axis=0) + owner_mask = np.expand_dims(owner_mask, axis=-1) + area = np.expand_dims(area, axis=-1) + geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) + + difcoef = np.where( + (levmask == 1) & (cfl_clipping == 1) & (owner_mask == 1), + scalfac_exdiff + * np.minimum( + 0.85 - cfl_w_limit * dtime, + np.abs(z_w_con_c) * dtime / ddqz_z_half - cfl_w_limit * dtime, + ), + 0, + ) + + ddt_w_adv = np.where( + (levmask == 1) & (cfl_clipping == 1) & (owner_mask == 1), + ddt_w_adv + + difcoef + * area + * np.sum( + np.where( + (grid.connectivities[C2E2CODim] != -1)[:, :, np.newaxis], + w[grid.connectivities[C2E2CODim]] * geofac_n2s, + 0, + ), + axis=1, + ), + ddt_w_adv, + ) + return ddt_w_adv + + class TestMoVelocityAdvectionStencil18(StencilTest): PROGRAM = mo_velocity_advection_stencil_18 OUTPUTS = ("ddt_w_adv",) @@ -44,37 +92,21 @@ def reference( dtime: wpfloat, **kwargs, ): - levmask = np.expand_dims(levmask, axis=0) - owner_mask = np.expand_dims(owner_mask, axis=-1) - area = np.expand_dims(area, axis=-1) - geofac_n2s = np.expand_dims(geofac_n2s, axis=-1) - - difcoef = np.where( - (levmask == 1) & (cfl_clipping == 1) & (owner_mask == 1), - scalfac_exdiff - * np.minimum( - 0.85 - cfl_w_limit * dtime, - np.abs(z_w_con_c) * dtime / ddqz_z_half - cfl_w_limit * dtime, - ), - 0, - ) - - ddt_w_adv = np.where( - (levmask == 1) & (cfl_clipping == 1) & (owner_mask == 1), - ddt_w_adv - + difcoef - * area - * np.sum( - np.where( - (grid.connectivities[C2E2CODim] != -1)[:, :, np.newaxis], - w[grid.connectivities[C2E2CODim]] * geofac_n2s, - 0, - ), - axis=1, - ), + ddt_w_adv = mo_velocity_advection_stencil_18_numpy( + grid, + levmask, + cfl_clipping, + owner_mask, + z_w_con_c, + ddqz_z_half, + area, + geofac_n2s, + w, ddt_w_adv, + scalfac_exdiff, + cfl_w_limit, + dtime, ) - return dict(ddt_w_adv=ddt_w_adv) @pytest.fixture diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_19.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_19.py similarity index 70% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_19.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_19.py index c9610b1d9..3216cab75 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_19.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_19.py @@ -28,6 +28,37 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_19_numpy( + grid, + z_kin_hor_e: np.array, + coeff_gradekin: np.array, + z_ekinh: np.array, + zeta: np.array, + vt: np.array, + f_e: np.array, + c_lin_e: np.array, + z_w_con_c_full: np.array, + vn_ie: np.array, + ddqz_z_full_e: np.array, +) -> np.array: + e2c = grid.connectivities[E2CDim] + z_ekinh_e2c = z_ekinh[e2c] + coeff_gradekin = coeff_gradekin.reshape(e2c.shape) + coeff_gradekin = np.expand_dims(coeff_gradekin, axis=-1) + f_e = np.expand_dims(f_e, axis=-1) + c_lin_e = np.expand_dims(c_lin_e, axis=-1) + + ddt_vn_apc = -( + (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) * z_kin_hor_e + + (-coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1]) + + vt * (f_e + 0.5 * np.sum(zeta[grid.connectivities[E2VDim]], axis=1)) + + np.sum(z_w_con_c_full[e2c] * c_lin_e, axis=1) + * (vn_ie[:, :-1] - vn_ie[:, 1:]) + / ddqz_z_full_e + ) + return ddt_vn_apc + + class TestMoVelocityAdvectionStencil19(StencilTest): PROGRAM = mo_velocity_advection_stencil_19 OUTPUTS = ("ddt_vn_apc",) @@ -47,20 +78,18 @@ def reference( ddqz_z_full_e: np.array, **kwargs, ) -> dict: - e2c = grid.connectivities[E2CDim] - z_ekinh_e2c = z_ekinh[e2c] - coeff_gradekin = coeff_gradekin.reshape(e2c.shape) - coeff_gradekin = np.expand_dims(coeff_gradekin, axis=-1) - f_e = np.expand_dims(f_e, axis=-1) - c_lin_e = np.expand_dims(c_lin_e, axis=-1) - - ddt_vn_apc = -( - (coeff_gradekin[:, 0] - coeff_gradekin[:, 1]) * z_kin_hor_e - + (-coeff_gradekin[:, 0] * z_ekinh_e2c[:, 0] + coeff_gradekin[:, 1] * z_ekinh_e2c[:, 1]) - + vt * (f_e + 0.5 * np.sum(zeta[grid.connectivities[E2VDim]], axis=1)) - + np.sum(z_w_con_c_full[e2c] * c_lin_e, axis=1) - * (vn_ie[:, :-1] - vn_ie[:, 1:]) - / ddqz_z_full_e + ddt_vn_apc = mo_velocity_advection_stencil_19_numpy( + grid, + z_kin_hor_e, + coeff_gradekin, + z_ekinh, + zeta, + vt, + f_e, + c_lin_e, + z_w_con_c_full, + vn_ie, + ddqz_z_full_e, ) return dict(ddt_vn_apc=ddt_vn_apc) diff --git a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_20.py b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_20.py similarity index 57% rename from model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_20.py rename to model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_20.py index 4e42c7879..3a497700c 100644 --- a/model/atmosphere/dycore/tests/stencil_tests/test_mo_velocity_advection_stencil_20.py +++ b/model/atmosphere/dycore/tests/dycore_stencil_tests/test_mo_velocity_advection_stencil_20.py @@ -31,6 +31,81 @@ from icon4py.model.common.type_alias import vpfloat, wpfloat +def mo_velocity_advection_stencil_20_numpy( + grid, + levelmask: np.array, + c_lin_e: np.array, + z_w_con_c_full: np.array, + ddqz_z_full_e: np.array, + area_edge: np.array, + tangent_orientation: np.array, + inv_primal_edge_length: np.array, + zeta: np.array, + geofac_grdiv: np.array, + vn: np.array, + ddt_vn_apc: np.array, + cfl_w_limit, + scalfac_exdiff, + dtime, +) -> np.array: + w_con_e = np.zeros_like(vn) + difcoef = np.zeros_like(vn) + + levelmask_offset_0 = levelmask[:-1] + levelmask_offset_1 = levelmask[1:] + + c_lin_e = np.expand_dims(c_lin_e, axis=-1) + geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) + area_edge = np.expand_dims(area_edge, axis=-1) + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) + inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) + + w_con_e = np.where( + (levelmask_offset_0) | (levelmask_offset_1), + np.sum( + np.where( + (grid.connectivities[E2CDim] != -1)[:, :, np.newaxis], + c_lin_e * z_w_con_c_full[grid.connectivities[E2CDim]], + 0, + ), + axis=1, + ), + w_con_e, + ) + difcoef = np.where( + ((levelmask_offset_0) | (levelmask_offset_1)) + & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), + scalfac_exdiff + * np.minimum( + 0.85 - cfl_w_limit * dtime, + np.abs(w_con_e) * dtime / ddqz_z_full_e - cfl_w_limit * dtime, + ), + difcoef, + ) + ddt_vn_apc = np.where( + ((levelmask_offset_0) | (levelmask_offset_1)) + & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), + ddt_vn_apc + + difcoef + * area_edge + * ( + np.sum( + np.where( + (grid.connectivities[E2C2EODim] != -1)[:, :, np.newaxis], + geofac_grdiv * vn[grid.connectivities[E2C2EODim]], + 0, + ), + axis=1, + ) + + tangent_orientation + * inv_primal_edge_length + * (zeta[grid.connectivities[E2VDim]][:, 1] - zeta[grid.connectivities[E2VDim]][:, 0]) + ), + ddt_vn_apc, + ) + return ddt_vn_apc + + class TestMoVelocityAdvectionStencil20(StencilTest): PROGRAM = mo_velocity_advection_stencil_20 OUTPUTS = ("ddt_vn_apc",) @@ -90,63 +165,22 @@ def reference( scalfac_exdiff, dtime, **kwargs, - ): - w_con_e = np.zeros_like(vn) - difcoef = np.zeros_like(vn) - - levelmask_offset_0 = levelmask[:-1] - levelmask_offset_1 = levelmask[1:] - - c_lin_e = np.expand_dims(c_lin_e, axis=-1) - geofac_grdiv = np.expand_dims(geofac_grdiv, axis=-1) - area_edge = np.expand_dims(area_edge, axis=-1) - tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) - inv_primal_edge_length = np.expand_dims(inv_primal_edge_length, axis=-1) - - w_con_e = np.where( - (levelmask_offset_0) | (levelmask_offset_1), - np.sum( - np.where( - (grid.connectivities[E2CDim] != -1)[:, :, np.newaxis], - c_lin_e * z_w_con_c_full[grid.connectivities[E2CDim]], - 0, - ), - axis=1, - ), - w_con_e, - ) - difcoef = np.where( - ((levelmask_offset_0) | (levelmask_offset_1)) - & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), - scalfac_exdiff - * np.minimum( - 0.85 - cfl_w_limit * dtime, - np.abs(w_con_e) * dtime / ddqz_z_full_e - cfl_w_limit * dtime, - ), - difcoef, - ) - ddt_vn_apc = np.where( - ((levelmask_offset_0) | (levelmask_offset_1)) - & (np.abs(w_con_e) > cfl_w_limit * ddqz_z_full_e), - ddt_vn_apc - + difcoef - * area_edge - * ( - np.sum( - np.where( - (grid.connectivities[E2C2EODim] != -1)[:, :, np.newaxis], - geofac_grdiv * vn[grid.connectivities[E2C2EODim]], - 0, - ), - axis=1, - ) - + tangent_orientation - * inv_primal_edge_length - * ( - zeta[grid.connectivities[E2VDim]][:, 1] - - zeta[grid.connectivities[E2VDim]][:, 0] - ) - ), + ) -> dict: + ddt_vn_apc = mo_velocity_advection_stencil_20_numpy( + grid, + levelmask, + c_lin_e, + z_w_con_c_full, + ddqz_z_full_e, + area_edge, + tangent_orientation, + inv_primal_edge_length, + zeta, + geofac_grdiv, + vn, ddt_vn_apc, + cfl_w_limit, + scalfac_exdiff, + dtime, ) return dict(ddt_vn_apc=ddt_vn_apc) diff --git a/model/common/src/icon4py/model/common/test_utils/helpers.py b/model/common/src/icon4py/model/common/test_utils/helpers.py index fc10086e3..5bc0d31c3 100644 --- a/model/common/src/icon4py/model/common/test_utils/helpers.py +++ b/model/common/src/icon4py/model/common/test_utils/helpers.py @@ -124,6 +124,13 @@ def flatten_first_two_dims(*dims: gt_common.Dimension, field: gt_common.Field) - return as_field(dims, newarray) +def unflatten_first_two_dims(field: gt_common.Field) -> np.array: + """Convert a (n-1)-D flattened (Felix-style) sparse field back to a n-D sparse field.""" + old_shape = np.asarray(field).shape + new_shape = (old_shape[0] // 3, 3) + old_shape[1:] + return np.asarray(field).reshape(new_shape) + + def dallclose(a, b, rtol=1.0e-12, atol=0.0, equal_nan=False): return np.allclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) diff --git a/tools/tests/liskov/fortran_samples.py b/tools/tests/liskov/fortran_samples.py index 7ed24e595..87529b0b5 100644 --- a/tools/tests/liskov/fortran_samples.py +++ b/tools/tests/liskov/fortran_samples.py @@ -351,7 +351,7 @@ !$DSL diff_multfac_vn=diff_multfac_vn(:); & !$DSL nudgecoeff_e=p_int%nudgecoeff_e(:,1); & !$DSL vn=p_nh_prog%vn(:,:,1); & - !$DSL horz_idx=horz_idx(:); & + !$DSL edge=horizontal_idx(:); & !$DSL nudgezone_diff=nudgezone_diff; & !$DSL fac_bdydiff_v=fac_bdydiff_v; & !$DSL start_2nd_nudge_line_idx_e=start_2nd_nudge_line_idx_e-1; & @@ -403,12 +403,12 @@ !$DSL INSERT(w_old(:,:,:) = p_nh_prog%w(:,:,:)) !$DSL INSERT(!$ACC END PARALLEL) - !$DSL START FUSED STENCIL(name=apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance; & + !$DSL START FUSED STENCIL(name=apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence; & !$DSL area=p_patch%cells%area(:,1); geofac_grg_x=p_int%geofac_grg(:,:,1,1); & !$DSL geofac_grg_y=p_int%geofac_grg(:,:,1,2); geofac_n2s=p_int%geofac_n2s(:,:,1); & !$DSL w_old=w_old(:,:,1); w=p_nh_prog%w(:,:,1); diff_multfac_w=diff_multfac_w; & - !$DSL diff_multfac_n2w=diff_multfac_n2w(:); vert_idx=vert_idx(:); & - !$DSL horz_idx=horz_idx(:); nrdmax=nrdmax(jg); interior_idx=start_interior_idx_c-1; & + !$DSL diff_multfac_n2w=diff_multfac_n2w(:); k=vertical_idx(:); & + !$DSL cell=horizontal_idx(:); nrdmax=nrdmax(jg); interior_idx=start_interior_idx_c-1; & !$DSL halo_idx=end_interior_idx_c; dwdx=p_nh_diag%dwdx(:,:,1); & !$DSL dwdy=p_nh_diag%dwdy(:,:,1); & !$DSL w_rel_tol=1e-09_wp; dwdx_rel_tol=1e-09_wp; dwdy_abs_tol=1e-09_wp; & @@ -443,7 +443,7 @@ !$DSL END STENCIL(name=apply_nabla2_to_w_in_upper_damping_layer) - !$DSL END FUSED STENCIL(name=apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance) + !$DSL END FUSED STENCIL(name=apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulence) """