Skip to content

Commit

Permalink
merge main
Browse files Browse the repository at this point in the history
  • Loading branch information
@halungge
halungge committed Feb 6, 2024
2 parents edc5e4f + a5c8e18 commit 083713c
Show file tree
Hide file tree
Showing 40 changed files with 264 additions and 246 deletions.
4 changes: 2 additions & 2 deletions model/atmosphere/diffusion/src/icon4py/model/atmosphere/diffusion/diffusion_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from gt4py.next.ffront.fbuiltins import broadcast, int32, minimum

from icon4py.model.common.dimension import CellDim, EdgeDim, KDim, VertexDim
from icon4py.model.common.math.smagorinsky import en_smag_fac_for_zero_nshift
from icon4py.model.common.math.smagorinsky import _en_smag_fac_for_zero_nshift


# TODO(Magdalena): fix duplication: duplicated from test testutils/utils.py
Expand Down Expand Up @@ -115,7 +115,7 @@ def _init_diffusion_local_fields_for_regular_timestemp(
) -> tuple[Field[[KDim], float], Field[[KDim], float], Field[[KDim], float]]:
diff_multfac_vn = _setup_runtime_diff_multfac_vn(k4, dyn_substeps)
smag_limit = _setup_smag_limit(diff_multfac_vn)
enh_smag_fac = en_smag_fac_for_zero_nshift(
enh_smag_fac = _en_smag_fac_for_zero_nshift(
vect_a,
hdiff_smag_fac,
hdiff_smag_fac2,
Expand Down
2 changes: 1 addition & 1 deletion model/atmosphere/diffusion/tests/diffusion_tests/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ def construct_diagnostics(
savepoint: IconDiffusionInitSavepoint,
grid_savepoint: IconGridSavepoint,
) -> DiffusionDiagnosticState:
grid = grid_savepoint.construct_icon_grid()
grid = grid_savepoint.construct_icon_grid(on_gpu=False)
dwdx = savepoint.dwdx() if savepoint.dwdx() else zero_field(grid, CellDim, KDim)
dwdy = savepoint.dwdy() if savepoint.dwdy() else zero_field(grid, CellDim, KDim)
return DiffusionDiagnosticState(
Expand Down
7 changes: 4 additions & 3 deletions ...re/dycore/mo_solve_nonhydro_stencil_50.py → ...ysis_increments_from_data_assimilation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,14 @@


@field_operator
def _mo_solve_nonhydro_stencil_50(
def _add_analysis_increments_from_data_assimilation(
z_rho_expl: Field[[CellDim, KDim], wpfloat],
z_exner_expl: Field[[CellDim, KDim], wpfloat],
rho_incr: Field[[CellDim, KDim], vpfloat],
exner_incr: Field[[CellDim, KDim], vpfloat],
iau_wgt_dyn: wpfloat,
) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]:
"""Formerly known as _mo_solve_nonhydro_stencil_50."""
rho_incr_wp, exner_incr_wp = astype((rho_incr, exner_incr), wpfloat)

z_rho_expl_wp = z_rho_expl + iau_wgt_dyn * rho_incr_wp
Expand All @@ -35,7 +36,7 @@ def _mo_solve_nonhydro_stencil_50(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_stencil_50(
def add_analysis_increments_from_data_assimilation(
z_rho_expl: Field[[CellDim, KDim], wpfloat],
z_exner_expl: Field[[CellDim, KDim], wpfloat],
rho_incr: Field[[CellDim, KDim], vpfloat],
Expand All @@ -46,7 +47,7 @@ def mo_solve_nonhydro_stencil_50(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_stencil_50(
_add_analysis_increments_from_data_assimilation(
z_rho_expl,
z_exner_expl,
rho_incr,
Expand Down
7 changes: 4 additions & 3 deletions ...re/mo_solve_nonhydro_4th_order_divdamp.py → ...ore/apply_4th_order_divergence_damping.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,19 +20,20 @@


@field_operator
def _mo_solve_nonhydro_4th_order_divdamp(
def _apply_4th_order_divergence_damping(
scal_divdamp: Field[[KDim], wpfloat],
z_graddiv2_vn: Field[[EdgeDim, KDim], vpfloat],
vn: Field[[EdgeDim, KDim], wpfloat],
) -> Field[[EdgeDim, KDim], wpfloat]:
"""Formelry known as _mo_solve_nonhydro_4th_order_divdamp."""
z_graddiv2_vn_wp = astype(z_graddiv2_vn, wpfloat)
scal_divdamp = broadcast(scal_divdamp, (EdgeDim, KDim))
vn_wp = vn + (scal_divdamp * z_graddiv2_vn_wp)
return vn_wp


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_4th_order_divdamp(
def apply_4th_order_divergence_damping(
scal_divdamp: Field[[KDim], wpfloat],
z_graddiv2_vn: Field[[EdgeDim, KDim], vpfloat],
vn: Field[[EdgeDim, KDim], wpfloat],
Expand All @@ -41,7 +42,7 @@ def mo_solve_nonhydro_4th_order_divdamp(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_4th_order_divdamp(
_apply_4th_order_divergence_damping(
scal_divdamp,
z_graddiv2_vn,
vn,
Expand Down
7 changes: 4 additions & 3 deletions ...re/dycore/mo_solve_nonhydro_stencil_54.py → ...ycore/apply_rayleigh_damping_mechanism.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,18 +20,19 @@


@field_operator
def _mo_solve_nonhydro_stencil_54(
def _apply_rayleigh_damping_mechanism(
z_raylfac: Field[[KDim], wpfloat],
w_1: Field[[CellDim], wpfloat],
w: Field[[CellDim, KDim], wpfloat],
) -> Field[[CellDim, KDim], wpfloat]:
"""Formerly known as _mo_solve_nonhydro_stencil_54."""
z_raylfac = broadcast(z_raylfac, (CellDim, KDim))
w_wp = z_raylfac * w + (wpfloat("1.0") - z_raylfac) * w_1
return w_wp


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_stencil_54(
def apply_rayleigh_damping_mechanism(
z_raylfac: Field[[KDim], wpfloat],
w_1: Field[[CellDim], wpfloat],
w: Field[[CellDim, KDim], wpfloat],
Expand All @@ -40,7 +41,7 @@ def mo_solve_nonhydro_stencil_54(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_stencil_54(
_apply_rayleigh_damping_mechanism(
z_raylfac,
w_1,
w,
Expand Down
7 changes: 4 additions & 3 deletions ...dycore/mo_solve_nonhydro_stencil_56_63.py → ...re/compute_dwdz_for_divergence_damping.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,12 @@


@field_operator
def _mo_solve_nonhydro_stencil_56_63(
def _compute_dwdz_for_divergence_damping(
inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat],
w: Field[[CellDim, KDim], wpfloat],
w_concorr_c: Field[[CellDim, KDim], vpfloat],
) -> Field[[CellDim, KDim], vpfloat]:
"""Formerly known as _mo_solve_nonhydro_stencil_56_63."""
inv_ddqz_z_full_wp = astype(inv_ddqz_z_full, wpfloat)

z_dwdz_dd_wp = inv_ddqz_z_full_wp * (
Expand All @@ -34,7 +35,7 @@ def _mo_solve_nonhydro_stencil_56_63(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_stencil_56_63(
def compute_dwdz_for_divergence_damping(
inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat],
w: Field[[CellDim, KDim], wpfloat],
w_concorr_c: Field[[CellDim, KDim], vpfloat],
Expand All @@ -44,7 +45,7 @@ def mo_solve_nonhydro_stencil_56_63(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_stencil_56_63(
_compute_dwdz_for_divergence_damping(
inv_ddqz_z_full,
w,
w_concorr_c,
Expand Down
7 changes: 4 additions & 3 deletions ...re/dycore/mo_solve_nonhydro_stencil_67.py → ...ere/dycore/compute_exner_from_rhotheta.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,20 +20,21 @@


@field_operator
def _mo_solve_nonhydro_stencil_67(
def _compute_exner_from_rhotheta(
rho: Field[[CellDim, KDim], wpfloat],
theta_v: Field[[CellDim, KDim], wpfloat],
exner: Field[[CellDim, KDim], wpfloat],
rd_o_cvd: wpfloat,
rd_o_p0ref: wpfloat,
) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]:
"""Formerly known as _mo_solve_nonhydro_stencil_67."""
theta_v_wp = exner
exner_wp = exp(rd_o_cvd * log(rd_o_p0ref * rho * theta_v_wp))
return theta_v_wp, exner_wp


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_stencil_67(
def compute_exner_from_rhotheta(
rho: Field[[CellDim, KDim], wpfloat],
theta_v: Field[[CellDim, KDim], wpfloat],
exner: Field[[CellDim, KDim], wpfloat],
Expand All @@ -44,7 +45,7 @@ def mo_solve_nonhydro_stencil_67(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_stencil_67(
_compute_exner_from_rhotheta(
rho,
theta_v,
exner,
Expand Down
7 changes: 4 additions & 3 deletions ...re/dycore/mo_solve_nonhydro_stencil_55.py → ...te_results_for_thermodynamic_variables.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@


@field_operator
def _mo_solve_nonhydro_stencil_55(
def _compute_results_for_thermodynamic_variables(
z_rho_expl: Field[[CellDim, KDim], wpfloat],
vwind_impl_wgt: Field[[CellDim], wpfloat],
inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat],
Expand All @@ -40,6 +40,7 @@ def _mo_solve_nonhydro_stencil_55(
Field[[CellDim, KDim], wpfloat],
Field[[CellDim, KDim], wpfloat],
]:
"""Formerly known as _mo_solve_nonhydro_stencil_55."""
inv_ddqz_z_full_wp, exner_ref_mc_wp, z_alpha_wp, z_beta_wp = astype(
(inv_ddqz_z_full, exner_ref_mc, z_alpha, z_beta), wpfloat
)
Expand All @@ -62,7 +63,7 @@ def _mo_solve_nonhydro_stencil_55(


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_stencil_55(
def compute_results_for_thermodynamic_variables(
z_rho_expl: Field[[CellDim, KDim], wpfloat],
vwind_impl_wgt: Field[[CellDim], wpfloat],
inv_ddqz_z_full: Field[[CellDim, KDim], vpfloat],
Expand All @@ -85,7 +86,7 @@ def mo_solve_nonhydro_stencil_55(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_stencil_55(
_compute_results_for_thermodynamic_variables(
z_rho_expl,
vwind_impl_wgt,
inv_ddqz_z_full,
Expand Down
7 changes: 4 additions & 3 deletions ...re/dycore/mo_solve_nonhydro_stencil_66.py → ...osphere/dycore/compute_theta_and_exner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,21 +20,22 @@


@field_operator
def _mo_solve_nonhydro_stencil_66(
def _compute_theta_and_exner(
bdy_halo_c: Field[[CellDim], bool],
rho: Field[[CellDim, KDim], wpfloat],
theta_v: Field[[CellDim, KDim], wpfloat],
exner: Field[[CellDim, KDim], wpfloat],
rd_o_cvd: wpfloat,
rd_o_p0ref: wpfloat,
) -> tuple[Field[[CellDim, KDim], wpfloat], Field[[CellDim, KDim], wpfloat]]:
"""Formelry known as _mo_solve_nonhydro_stencil_66."""
theta_v_wp = where(bdy_halo_c, exner, theta_v)
exner_wp = where(bdy_halo_c, exp(rd_o_cvd * log(rd_o_p0ref * rho * exner)), exner)
return theta_v_wp, exner_wp


@program(grid_type=GridType.UNSTRUCTURED)
def mo_solve_nonhydro_stencil_66(
def compute_theta_and_exner(
bdy_halo_c: Field[[CellDim], bool],
rho: Field[[CellDim, KDim], wpfloat],
theta_v: Field[[CellDim, KDim], wpfloat],
Expand All @@ -46,7 +47,7 @@ def mo_solve_nonhydro_stencil_66(
vertical_start: int32,
vertical_end: int32,
):
_mo_solve_nonhydro_stencil_66(
_compute_theta_and_exner(
bdy_halo_c,
rho,
theta_v,
Expand Down
Loading

0 comments on commit 083713c

Please sign in to comment.