diff --git a/.github/workflows/icon4py-test-model.yml b/.github/workflows/icon4py-test-model.yml index 6d09d89fc4..05d5a3a05a 100644 --- a/.github/workflows/icon4py-test-model.yml +++ b/.github/workflows/icon4py-test-model.yml @@ -24,7 +24,7 @@ jobs: matrix: python-version: ["3.12", "3.13", "3.14"] backend: ["embedded", "dace_cpu", "gtfn_cpu"] - component: ["tracer_advection", "diffusion", "dycore", "microphysics", "muphys", "common", "standalone_driver"] + component: ["tracer_advection", "diffusion", "dycore", "microphysics", "muphys", "physics_driver", "common", "standalone_driver"] steps: - name: Checkout uses: actions/checkout@v7 diff --git a/ci/default.yml b/ci/default.yml index 4178d715d6..4e2e7dd2a6 100644 --- a/ci/default.yml +++ b/ci/default.yml @@ -10,7 +10,7 @@ variables: # # cscs-ci run default;BACKENDS=gtfn_cpu;LEVELS=unit;MODEL_SUBPACKAGES=common:standalone_driver;MODEL_MPI_SUBPACKAGES=common;SESSIONS=model;MODEL_SUBSETS=datatest SESSIONS: "model:model_mpi:tools" - MODEL_SUBPACKAGES: "tracer_advection:diffusion:dycore:microphysics:muphys:common:standalone_driver" + MODEL_SUBPACKAGES: "tracer_advection:diffusion:dycore:microphysics:muphys:physics_driver:common:standalone_driver" MODEL_SUBSETS: "datatest" MODEL_MPI_SUBPACKAGES: "tracer_advection:diffusion:dycore:common:standalone_driver" MODEL_MPI_SUBSETS: "datatest" diff --git a/ci/merge.yml b/ci/merge.yml index e7541faa36..53bab0e819 100644 --- a/ci/merge.yml +++ b/ci/merge.yml @@ -5,7 +5,7 @@ variables: # See default.yml and generate_ci_pipeline.py for details on how these are # used. This pipeline gates merges and variables should not be overridden. SESSIONS: "model:model_mpi:tools" - MODEL_SUBPACKAGES: "tracer_advection:diffusion:dycore:microphysics:muphys:common:standalone_driver" + MODEL_SUBPACKAGES: "tracer_advection:diffusion:dycore:microphysics:muphys:physics_driver:common:standalone_driver" MODEL_SUBSETS: "stencils:datatest" MODEL_MPI_SUBPACKAGES: "tracer_advection:diffusion:dycore:common:standalone_driver" MODEL_MPI_SUBSETS: "datatest" diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/component.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/component.py new file mode 100644 index 0000000000..0a1201244c --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/component.py @@ -0,0 +1,209 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import types +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, cast + +import gt4py.next as gtx + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys import ( + config as muphys_config, + data as muphys_data, +) +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import SPECIES, Q +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver.run_full_muphys import ( + setup_muphys, +) +from icon4py.model.common import ( + dimension as dims, + field_type_aliases as fa, + model_backends, + model_options, + time, + type_alias as ta, +) +from icon4py.model.common.diagnostic_calculations.stencils import calculate_tendency +from icon4py.model.common.grid import horizontal as h_grid +from icon4py.model.common.math.stencils import generic_math_operations + + +if TYPE_CHECKING: + import gt4py.next.typing as gtx_typing + + from icon4py.model.common.grid import icon as icon_grid + from icon4py.model.common.states import model + + +class MuphysComponent: + """Per-process adapter wrapping the muphys microphysics program.""" + + # TODO (Yilu): inherit the Component protocol once it is formalized (deferred to a separate PR). + inputs_properties = muphys_data.INPUTS_PROPERTIES + outputs_properties = muphys_data.OUTPUTS_PROPERTIES + + def __init__( + self, + grid: icon_grid.IconGrid, + dtime: time.RelativeTime, + qnc: float, + backend: gtx_typing.Backend | None = None, + *, + scheme: muphys_config.MuphysScheme = muphys_config.MuphysScheme.KOKKOS_MUPHYS, + step: Callable[..., Any] | None = None, + ) -> None: + self._ncells = grid.num_cells + self._nlev = grid.num_levels + self._dt_seconds = dtime.total_seconds() + self._qnc = qnc + self._backend = model_options.customize_backend(program=None, backend=backend) + + cell_domain = h_grid.domain(dims.CellDim) + # ICON applies physics on the prognostic cells only: + # grf_bdywidth_c+1 .. min_rlcell_int (NUDGING start .. LOCAL end). + cell_start = grid.start_index(cell_domain(h_grid.Zone.NUDGING)) + cell_end = grid.end_index(cell_domain(h_grid.Zone.LOCAL)) + + # Inputs are copied over the full range: the muphys step reads whole fields. + full_horizontal_sizes = { + "horizontal_start": gtx.int32(0), + "horizontal_end": gtx.int32(self._ncells), + } + # Tendencies only on the prognostic subdomain -- the tendency buffers stay + # zero outside, so scattering is a no-op on boundary/halo rows, as in ICON. + prognostic_horizontal_sizes = { + "horizontal_start": cell_start, + "horizontal_end": cell_end, + } + vertical_sizes = {"vertical_start": gtx.int32(0), "vertical_end": gtx.int32(self._nlev)} + + self._calculate_tendency = model_options.setup_program( + program=calculate_tendency.calculate_cell_kdim_field_tendency, + backend=self._backend, + horizontal_sizes=prognostic_horizontal_sizes, + vertical_sizes=vertical_sizes, + offset_provider={}, + ) + self._copy_field = model_options.setup_program( + program=generic_math_operations.copy_field_on_cell_k, + backend=self._backend, + horizontal_sizes=full_horizontal_sizes, + vertical_sizes=vertical_sizes, + offset_provider={}, + ) + + allocator = model_backends.get_allocator(backend) + + if step is None: + sizes = types.SimpleNamespace(ncells=self._ncells, nlev=self._nlev) + step = setup_muphys( + inp=sizes, # type: ignore[arg-type] # only .ncells/.nlev are read + dt=self._dt_seconds, + qnc=qnc, + backend=backend, + single_program=False, + scheme=scheme, + ) + self._step = step + + cell_k_domain = gtx.domain({dims.CellDim: self._ncells, dims.KDim: self._nlev}) + self._t_out = gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator) + self._q_out = Q( + v=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + c=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + r=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + s=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + i=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + g=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + ) + self._pflx = gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator) + self._pr = gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator) + self._ps = gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator) + self._pi = gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator) + self._pg = gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator) + self._pre = gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator) + + self._tendencies: dict[str, fa.CellKField[ta.wpfloat]] = { + "tend_temperature": gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + "tend_qv": gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + "tend_qc": gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + "tend_qr": gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + "tend_qs": gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + "tend_qi": gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + "tend_qg": gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + } + + self._te_in = gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator) + self._q_in = Q( + v=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + c=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + r=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + s=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + i=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + g=gtx.zeros(cell_k_domain, dtype=ta.wpfloat, allocator=allocator), + ) + + def __call__( + self, state: dict[str, model.DataField], time_step: time.AbsoluteTime + ) -> dict[str, model.DataField]: + """Run muphys, then convert its updated state into tendencies. + + muphys returns updated state (t_out, q_out); this boundary converts it + to tendencies ``(new - old) / dt`` s. Precip outputs are diagnostics, passed straight through. + """ + # cast from generic ``DataField`` to bare gt4py fields + fields = cast("dict[str, fa.CellKField[ta.wpfloat]]", state) + + self._copy_field(field=fields["te"], output_field=self._te_in) + for s in SPECIES: + self._copy_field(field=fields[f"q{s}"], output_field=getattr(self._q_in, s)) + + self._step( + dz=fields["dz"], + te=self._te_in, + p=fields["p"], + rho=fields["rho"], + q_in=self._q_in, + q_out=self._q_out, + t_out=self._t_out, + pflx=self._pflx, + pr=self._pr, + ps=self._ps, + pi=self._pi, + pg=self._pg, + pre=self._pre, + ) + + self._calculate_tendency( + dtime=self._dt_seconds, + old_field=fields["te"], + new_field=self._t_out, + tendency=self._tendencies["tend_temperature"], + ) + for s in SPECIES: + self._calculate_tendency( + dtime=self._dt_seconds, + old_field=fields[f"q{s}"], + new_field=getattr(self._q_out, s), + tendency=self._tendencies[f"tend_q{s}"], + ) + + return cast( + "dict[str, model.DataField]", + { + **self._tendencies, + "pflx": self._pflx, + "pr": self._pr, + "ps": self._ps, + "pi": self._pi, + "pg": self._pg, + "pre": self._pre, + }, + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/config.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/config.py new file mode 100644 index 0000000000..89cdaef381 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/config.py @@ -0,0 +1,34 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import dataclasses +import enum + + +class MuphysScheme(enum.Enum): + """Selects between the two graupel microphysics formulations. + + KOKKOS_MUPHYS follows the muphys C++/Kokkos reference implementation (the + original source of this port, validated by the netCDF-based muphys tests). + AES_GRAUPEL follows the newer MPIM rain-microphysics revisions in + ICON (mo_aes_graupel.f90), validated against the aes-graupel serialbox + savepoints. + """ + + KOKKOS_MUPHYS = "kokkos_muphys" + AES_GRAUPEL = "aes_graupel" + + +@dataclasses.dataclass(frozen=True) +class MuphysConfig: + """Configuration for the muphys microphysics component.""" + + qnc: float = 50.0e6 # cloud droplet number concentration [m^-3] + scheme: MuphysScheme = MuphysScheme.AES_GRAUPEL diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/common/constants.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/common/constants.py index 918a53a991..9a4a82dbc7 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/common/constants.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/common/constants.py @@ -78,8 +78,45 @@ class IndexConsts(ta.wpfloat, enum.Enum): exponent_i = 0.160 offset_i = 1.0e-12 prefactor_s = 57.80 - exponent_s = 0.16666666666666666 + exponent_s = 1 / 6 offset_s = 1.0e-12 prefactor_g = 12.24 exponent_g = 0.217 offset_g = 1.0e-08 + + +class AesGraupelConsts(ta.wpfloat, enum.Enum): + # Constants of the newer MPIM rain-microphysics revisions in ICON + # (mo_aes_graupel.f90), used by the AES_GRAUPEL scheme variants only. + # Hydrometeor-density clamp bounds, shared by all polynomial fits and fall speeds + rhox_mn = 3.26216e-08 + rhox_mx = 6.97604e-03 + # cloud_to_rain: accretion kernel, degree-4 polynomial in log(clamped rho*qr) + a_ac_1 = -2.155543e00 + a_ac_2 = -1.148491e00 + a_ac_3 = -1.882563e-02 + a_ac_4 = 2.941391e-03 + a_ac_5 = 5.575598e-05 + # rain_to_vapor: evaporation, exp of degree-4 polynomial in log(clamped rho*qr) + a_ev_1 = -5.532194e00 + a_ev_2 = 2.432848e-01 + a_ev_3 = -4.145391e-02 + a_ev_4 = -1.798439e-03 + a_ev_5 = -1.405764e-05 + # rain fall speed: degree-4 polynomial in log(clamped rho*qr); the constant term + # vm_a_r_1 is NOT multiplied by sqrt(rho_00/rho) in the Fortran (operator + # precedence in mo_aes_graupel.f90 vm()), reproduced faithfully + vm_a_r_1 = -5.91051e-01 + vm_a_r_2 = -5.37440e00 + vm_a_r_3 = -1.00459e00 + vm_a_r_4 = -6.44895e-02 + vm_a_r_5 = -1.40361e-03 + # ice / snow / graupel fall speeds: prefactor * clamped_rhox**exponent * density correction + vm_prefactor_i = 0.80 + vm_exponent_i = 0.160 + vm_prefactor_s = 115.60 # 2 * 57.80 + vm_exponent_s = 1 / 6 + vm_prefactor_g = 12.24 + vm_exponent_g = 0.217 + # snow_number: lower clamp on the snow mass density rho*qs + rho_s_mn = 2.0e-7 diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py index 5db19381a3..d3ec459763 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py @@ -27,3 +27,11 @@ class Q_scalar(NamedTuple): s: ta.wpfloat # Specific snow water i: ta.wpfloat # Specific ice water content g: ta.wpfloat # Specific graupel water content + + +SPECIES: tuple[str, ...] = Q._fields + +# muphys precip diagnostic field names (pure diagnostics -- never applied to the +# prognostic state): total precip flux, surface rain / snow / ice / graupel rates, +# and surface precip energy flux. +PRECIP_DIAGNOSTICS: tuple[str, ...] = ("pflx", "pr", "ps", "pi", "pg", "pre") diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/properties.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/properties.py index 4f9e4e67cd..2f4596f29d 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/properties.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/properties.py @@ -6,9 +6,10 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx -from gt4py.next import exp, maximum, minimum, power, where +from gt4py.next import exp, log, maximum, minimum, power, sqrt, where from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.constants import ( + AesGraupelConsts, GraupelConsts, ThermodynamicConsts, ) @@ -554,3 +555,247 @@ def vel_scale_factor_snow( scale_factor: fa.CellKField[ta.wpfloat], # output ) -> None: _vel_scale_factor_snow(xrho=xrho, rho=rho, t=t, qs=qs, out=scale_factor) + + +@gtx.field_operator +def _snow_number_aes_graupel( + t: fa.CellKField[ta.wpfloat], + rho_s: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + """ + Compute the snow number, AES_GRAUPEL scheme. + + Same fit as _snow_number, but operates on the snow mass density rho*qs with a + lower clamp MAX(rho_s, rho_s_mn) instead of the additive (qs + 2e-6)*rho offset, + and branches on rho_s > qmin (ICON mo_aes_graupel.f90 snow_number). + + Args: + t: Temperature + rho_s: Snow mass density rho*qs + + Result: Snow number + """ + TMIN = ThermodynamicConsts.tmelt - wpfloat(40.0) + TMAX = ThermodynamicConsts.tmelt + XA1 = wpfloat(-1.65e0) + XA2 = wpfloat(5.45e-2) + XA3 = wpfloat(3.27e-4) + XB1 = wpfloat(1.42e0) + XB2 = wpfloat(1.19e-2) + XB3 = wpfloat(9.60e-5) + N0S0 = wpfloat(8.00e5) + N0S1 = wpfloat(13.5) * wpfloat(5.65e05) + N0S2 = wpfloat(-0.107) + N0S3 = wpfloat(13.5) + N0S4 = wpfloat(0.5) * N0S1 + N0S5 = wpfloat(1.0e6) + N0S6 = wpfloat(1.0e2) * N0S1 + N0S7 = wpfloat(1.0e9) + + tc = maximum(minimum(t, TMAX), TMIN) - ThermodynamicConsts.tmelt + alf = power(wpfloat(10.0), (XA1 + tc * (XA2 + tc * XA3))) + bet = XB1 + tc * (XB2 + tc * XB3) + n0s = ( + N0S3 + * power( + (maximum(rho_s, AesGraupelConsts.rho_s_mn) / GraupelConsts.ams), + (wpfloat(4.0) - wpfloat(3.0) * bet), + ) + / (alf * alf * alf) + ) + y = exp(N0S2 * tc) + n0smn = maximum(N0S4 * y, N0S5) + n0smx = minimum(N0S6 * y, N0S7) + return where(rho_s > GraupelConsts.qmin, minimum(n0smx, maximum(n0smn, n0s)), N0S0) + + +@gtx.field_operator +def _snow_number_aes_graupel_scalar( + t: ta.wpfloat, + rho_s: ta.wpfloat, +) -> ta.wpfloat: + """ + Compute the snow number, AES_GRAUPEL scheme (scalar version for scan operators). + + Args: + t: Temperature + rho_s: Snow mass density rho*qs + + Result: Snow number + """ + TMIN = ThermodynamicConsts.tmelt - wpfloat(40.0) + TMAX = ThermodynamicConsts.tmelt + XA1 = wpfloat(-1.65e0) + XA2 = wpfloat(5.45e-2) + XA3 = wpfloat(3.27e-4) + XB1 = wpfloat(1.42e0) + XB2 = wpfloat(1.19e-2) + XB3 = wpfloat(9.60e-5) + N0S0 = wpfloat(8.00e5) + N0S1 = wpfloat(13.5) * wpfloat(5.65e05) + N0S2 = wpfloat(-0.107) + N0S3 = wpfloat(13.5) + N0S4 = wpfloat(0.5) * N0S1 + N0S5 = wpfloat(1.0e6) + N0S6 = wpfloat(1.0e2) * N0S1 + N0S7 = wpfloat(1.0e9) + + tc = maximum(minimum(t, TMAX), TMIN) - ThermodynamicConsts.tmelt + alf = power(wpfloat(10.0), (XA1 + tc * (XA2 + tc * XA3))) + bet = XB1 + tc * (XB2 + tc * XB3) + n0s = ( + N0S3 + * power( + (maximum(rho_s, AesGraupelConsts.rho_s_mn) / GraupelConsts.ams), + (wpfloat(4.0) - wpfloat(3.0) * bet), + ) + / (alf * alf * alf) + ) + y = exp(N0S2 * tc) + n0smn = maximum(N0S4 * y, N0S5) + n0smx = minimum(N0S6 * y, N0S7) + return minimum(n0smx, maximum(n0smn, n0s)) if rho_s > GraupelConsts.qmin else N0S0 + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def snow_number_aes_graupel( + t: fa.CellKField[ta.wpfloat], # Temperature + rho_s: fa.CellKField[ta.wpfloat], # Snow mass density rho*qs + number: fa.CellKField[ta.wpfloat], # output +): + _snow_number_aes_graupel(t=t, rho_s=rho_s, out=number) + + +@gtx.field_operator +def _snow_lambda_aes_graupel( + rho_s: fa.CellKField[ta.wpfloat], + ns: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + """ + Compute the snow slope parameter, AES_GRAUPEL scheme. + + Same formula as _snow_lambda but operating on the snow mass density rho*qs + and branching on rho_s > qmin (ICON mo_aes_graupel.f90 snow_lambda). + + Args: + rho_s: Snow mass density rho*qs + ns: Snow number + + Result: Snow slope parameter (lambda) + """ + A2 = GraupelConsts.ams * wpfloat(2.0) # (with ams*gam(bms+1.0_wp) where gam(3) = 2) + LMD_0 = wpfloat(1.0e10) # no snow value of lambda + BX = wpfloat(1.0) / (GraupelConsts.bms + wpfloat(1.0)) # Exponent + + return where(rho_s > GraupelConsts.qmin, power((A2 * ns / rho_s), BX), LMD_0) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def snow_lambda_aes_graupel( + rho_s: fa.CellKField[ta.wpfloat], # Snow mass density rho*qs + ns: fa.CellKField[ta.wpfloat], # Snow number + riming_snow_rate: fa.CellKField[ta.wpfloat], # output +): + _snow_lambda_aes_graupel(rho_s=rho_s, ns=ns, out=riming_snow_rate) + + +@gtx.field_operator +def _vm_rain_aes_graupel_scalar( + rho_x: ta.wpfloat, + rho: ta.wpfloat, +) -> ta.wpfloat: + """ + Rain fall speed, AES_GRAUPEL scheme (scalar version for scan operators). + + Degree-4 polynomial in log of the clamped rain mass density. As in the Fortran + (mo_aes_graupel.f90 vm), the density correction sqrt(rho_00/rho) multiplies + only the polynomial part, not the constant term. + + Args: + rho_x: Rain mass density rho*qr + rho: Ambient air density + + Result: Fall speed + """ + x = log(minimum(AesGraupelConsts.rhox_mx, maximum(AesGraupelConsts.rhox_mn, rho_x))) + return AesGraupelConsts.vm_a_r_1 + x * ( + AesGraupelConsts.vm_a_r_2 + + x + * ( + AesGraupelConsts.vm_a_r_3 + + x * (AesGraupelConsts.vm_a_r_4 + x * AesGraupelConsts.vm_a_r_5) + ) + ) * sqrt(GraupelConsts.rho_00 / rho) + + +@gtx.field_operator +def _vm_ice_aes_graupel_scalar( + rho_x: ta.wpfloat, + rho: ta.wpfloat, +) -> ta.wpfloat: + """ + Ice fall speed, AES_GRAUPEL scheme (scalar version for scan operators). + + Args: + rho_x: Ice mass density rho*qi + rho: Ambient air density + + Result: Fall speed + """ + B_I = wpfloat(0.33333333333333333) + x = minimum(AesGraupelConsts.rhox_mx, maximum(AesGraupelConsts.rhox_mn, rho_x)) + return ( + AesGraupelConsts.vm_prefactor_i + * power(x, AesGraupelConsts.vm_exponent_i) + * power(GraupelConsts.rho_00 / rho, B_I) + ) + + +@gtx.field_operator +def _vm_snow_aes_graupel_scalar( + rho_x: ta.wpfloat, + rho: ta.wpfloat, + t: ta.wpfloat, +) -> ta.wpfloat: + """ + Snow fall speed, AES_GRAUPEL scheme (scalar version for scan operators). + + snow_number is evaluated at the clamped falling-mass density, as in the Fortran. + + Args: + rho_x: Snow mass density (rho*qs, or the level-averaged value for vt) + rho: Ambient air density + t: Temperature + + Result: Fall speed + """ + B_S = wpfloat(-0.16666666666666667) + x = minimum(AesGraupelConsts.rhox_mx, maximum(AesGraupelConsts.rhox_mn, rho_x)) + return ( + AesGraupelConsts.vm_prefactor_s + * power(x, AesGraupelConsts.vm_exponent_s) + * sqrt(GraupelConsts.rho_00 / rho) + * power(_snow_number_aes_graupel_scalar(t=t, rho_s=x), B_S) + ) + + +@gtx.field_operator +def _vm_graupel_aes_graupel_scalar( + rho_x: ta.wpfloat, + rho: ta.wpfloat, +) -> ta.wpfloat: + """ + Graupel fall speed, AES_GRAUPEL scheme (scalar version for scan operators). + + Args: + rho_x: Graupel mass density rho*qg + rho: Ambient air density + + Result: Fall speed + """ + x = minimum(AesGraupelConsts.rhox_mx, maximum(AesGraupelConsts.rhox_mn, rho_x)) + return ( + AesGraupelConsts.vm_prefactor_g + * power(x, AesGraupelConsts.vm_exponent_g) + * sqrt(GraupelConsts.rho_00 / rho) + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/transitions.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/transitions.py index d120ca1bc2..2f786265ad 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/transitions.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/transitions.py @@ -6,9 +6,10 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx -from gt4py.next import exp, maximum, minimum, power, sqrt, where +from gt4py.next import exp, log, maximum, minimum, power, sqrt, where from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.constants import ( + AesGraupelConsts, GraupelConsts, ThermodynamicConsts, ) @@ -36,11 +37,10 @@ def _cloud_to_graupel( """ A_RIM = wpfloat(4.43) B_RIM = wpfloat(0.94878) - ZERO = wpfloat(0.0) return where( (minimum(qc, qg) > GraupelConsts.qmin) & (t > GraupelConsts.tfrz_hom), A_RIM * qc * power(qg * rho, B_RIM), - ZERO, + wpfloat(0.0), ) @@ -83,20 +83,27 @@ def _cloud_to_rain( X3 = wpfloat(2.0e0) # gamma exponent for cloud distribution X2 = wpfloat(2.6e-10) # separating mass between cloud and rain X1 = wpfloat(9.44e9) # kernel coeff for SB2001 autoconversion - ZERO = wpfloat(0.0) - ONE = wpfloat(1.0) - TWO = wpfloat(2.0) - THREE = wpfloat(3.0) - FOUR = wpfloat(4.0) - AU_KERNEL = X1 / (wpfloat(20.0) * X2) * (X3 + TWO) * (X3 + FOUR) / ((X3 + ONE) * (X3 + ONE)) + AU_KERNEL = ( + X1 + / (wpfloat(20.0) * X2) + * (X3 + wpfloat(2.0)) + * (X3 + wpfloat(4.0)) + / ((X3 + wpfloat(1.0)) * (X3 + wpfloat(1.0))) + ) # TO-DO: put as much of this into the WHERE statement as possible - tau = maximum(TAU_MIN, minimum(ONE - qc / (qc + qr), TAU_MAX)) # temporary cannot go in where + tau = maximum( + TAU_MIN, minimum(wpfloat(1.0) - qc / (qc + qr), TAU_MAX) + ) # temporary cannot go in where phi = power(tau, B_PHI) - phi = A_PHI * phi * power(ONE - phi, THREE) - xau = AU_KERNEL * power(qc * qc / nc, TWO) * (ONE + phi / power(ONE - tau, TWO)) - xac = AC_KERNEL * qc * qr * power(tau / (tau + C_PHI), FOUR) - return where((qc > QMIN_AC) & (t > GraupelConsts.tfrz_hom), xau + xac, ZERO) + phi = A_PHI * phi * power(wpfloat(1.0) - phi, wpfloat(3.0)) + xau = ( + AU_KERNEL + * power(qc * qc / nc, wpfloat(2.0)) + * (wpfloat(1.0) + phi / power(wpfloat(1.0) - tau, wpfloat(2.0))) + ) + xac = AC_KERNEL * qc * qr * power(tau / (tau + C_PHI), wpfloat(4.0)) + return where((qc > QMIN_AC) & (t > GraupelConsts.tfrz_hom), xau + xac, wpfloat(0.0)) @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) @@ -132,12 +139,12 @@ def _cloud_to_snow( """ ECS = wpfloat(0.9) B_RIM = -(wpfloat(GraupelConsts.v1s) + wpfloat(3.0)) - C_RIM = wpfloat(2.61) * ECS * GraupelConsts.v0s # (with pi*gam(v1s+3)/4 = 2.610) - ZERO = wpfloat(0.0) + # ICON hardcodes the rounded 2.61 (pi*gam(v1s+3)/4 = 2.6102); keep literal for bit-exactness + C_RIM = wpfloat(2.61) * ECS * GraupelConsts.v0s return where( (minimum(qc, qs) > GraupelConsts.qmin) & (t > GraupelConsts.tfrz_hom), C_RIM * ns * qc * power(lam, B_RIM), - ZERO, + wpfloat(0.0), ) @@ -171,8 +178,7 @@ def _cloud_x_ice( Return: Freezing rate """ - ZERO = wpfloat(0.0) - result = where((qc > GraupelConsts.qmin) & (t < GraupelConsts.tfrz_hom), qc / dt, ZERO) + result = where((qc > GraupelConsts.qmin) & (t < GraupelConsts.tfrz_hom), qc / dt, wpfloat(0.0)) result = where((qi > GraupelConsts.qmin) & (t > ThermodynamicConsts.tmelt), -qi / dt, result) return result @@ -212,7 +218,6 @@ def _graupel_to_rain( B_MELT = wpfloat(0.6) # melting exponent C1_MELT = wpfloat(12.31698) # Constants in melting formula C2_MELT = wpfloat(7.39441e-05) # Constants in melting formula - ZERO = wpfloat(0.0) return where( ( t @@ -225,7 +230,7 @@ def _graupel_to_rain( (C1_MELT / p + C2_MELT) * (t - ThermodynamicConsts.tmelt + A_MELT * dvsw0) * power(qg * rho, B_MELT), - ZERO, + wpfloat(0.0), ) @@ -265,11 +270,10 @@ def _ice_to_graupel( B_CT = wpfloat(0.875) # Exponent = 7/8 C_AGG_CT = wpfloat(2.46) B_AGG_CT = wpfloat(0.94878) # Exponent - ZERO = wpfloat(0.0) result = where( (qi > GraupelConsts.qmin) & (qg > GraupelConsts.qmin), sticking_eff * qi * C_AGG_CT * power(rho * qg, B_AGG_CT), - ZERO, + wpfloat(0.0), ) result = where( (qi > GraupelConsts.qmin) & (qr > GraupelConsts.qmin), @@ -311,14 +315,15 @@ def _ice_to_snow( """ QI0 = wpfloat(0.0) # Critical ice required for autoconversion C_IAU = wpfloat(1.0e-3) # Coefficient of auto conversion - C_AGG = wpfloat(2.61) * GraupelConsts.v0s # Coeff of aggregation (2.610 = pi*gam(v1s+3)/4) + # ICON hardcodes the rounded 2.61 (pi*gam(v1s+3)/4 = 2.6102); keep literal for bit-exactness + C_AGG = wpfloat(2.61) * GraupelConsts.v0s # Coeff of aggregation B_AGG = -(GraupelConsts.v1s + wpfloat(3.0)) # Aggregation exponent - ZERO = wpfloat(0.0) return where( (qi > GraupelConsts.qmin), - sticking_eff * (C_IAU * maximum(ZERO, (qi - QI0)) + qi * (C_AGG * ns) * power(lam, B_AGG)), - ZERO, + sticking_eff + * (C_IAU * maximum(wpfloat(0.0), (qi - QI0)) + qi * (C_AGG * ns) * power(lam, B_AGG)), + wpfloat(0.0), ) @@ -371,14 +376,13 @@ def _rain_to_graupel( # noqa: PLR0917 [too-many-positional-arguments] A2 = wpfloat(1.24e-3) # (PI/24)*EIR*V0R*Gamma(6.5)*AR**(-5/8) B2 = wpfloat(1.625) # exponent for rho*qr QS_CRIT = wpfloat(1.0e-7) # critical humidity of snow - ZERO = wpfloat(0.0) - maskinner = (dvsw + qc <= ZERO) | (qr > C4 * qc) + maskinner = (dvsw + qc <= wpfloat(0.0)) | (qr > C4 * qc) mask = (qr > GraupelConsts.qmin) & (t < TFRZ_RAIN) result = where( mask & (t > GraupelConsts.tfrz_hom) & maskinner, (exp(C2 * (TFRZ_RAIN - t)) - C3) * (A1 * power((qr * rho), B1)), - ZERO, + wpfloat(0.0), ) result = where(mask & (t <= GraupelConsts.tfrz_hom), qr / dt, result) @@ -437,18 +441,17 @@ def _rain_to_vapor( # noqa: PLR0917 [too-many-positional-arguments] A1_RV = wpfloat(1.536e-3) # coefficient 1 in qr reconstruction A2_RV = wpfloat(1.0e0) # coefficient 2 in qr reconstruction A3_RV = wpfloat(19.0621e0) # coefficient 3 in qr reconstruction - ZERO = wpfloat(0.0) # TO-DO: move as much as possible into WHERE statement tc = t - ThermodynamicConsts.tmelt evap_max = (C1_RV + tc * (C2_RV + C3_RV * tc)) * (-dvsw) / dt return where( - (qr > GraupelConsts.qmin) & (dvsw + qc <= ZERO), + (qr > GraupelConsts.qmin) & (dvsw + qc <= wpfloat(0.0)), minimum( A1_RV * (A2_RV + A3_RV * power(qr * rho, B1_RV)) * (-dvsw) * power(qr * rho, B2_RV), evap_max, ), - ZERO, + wpfloat(0.0), ) @@ -485,11 +488,10 @@ def _snow_to_graupel( """ A_RIM_CT = wpfloat(0.5) # Constants in riming formula B_RIM_CT = wpfloat(0.75) - ZERO = wpfloat(0.0) return where( (minimum(qc, qs) > GraupelConsts.qmin) & (t > GraupelConsts.tfrz_hom), A_RIM_CT * qc * power(qs * rho, B_RIM_CT), - ZERO, + wpfloat(0.0), ) @@ -528,7 +530,6 @@ def _snow_to_rain( C2_SR = wpfloat(0.612654e-3) # Constants in melting formula A_SR = GraupelConsts.tx - wpfloat(389.5) # Melting prefactor B_SR = wpfloat(0.8) # Melting exponent - ZERO = wpfloat(0.0) return where( ( t @@ -541,7 +542,7 @@ def _snow_to_rain( (C1_SR / p + C2_SR) * (t - ThermodynamicConsts.tmelt + A_SR * dvsw0) * power(qs * rho, B_SR), - ZERO, + wpfloat(0.0), ) @@ -592,17 +593,16 @@ def _vapor_x_graupel( # noqa: PLR0917 [too-many-positional-arguments] A7 = wpfloat(0.0418521) A8 = wpfloat(-4.7524e-8) B_VG = wpfloat(0.6) - ZERO = wpfloat(0.0) result = where( (t < ThermodynamicConsts.tmelt), (A1_VG + A2_VG * t + A3 / p + A4 * p) * dvsi * power(qg * rho, B_VG), where( (t > (ThermodynamicConsts.tmelt - GraupelConsts.tx * dvsw0)), - (A5 + A6 * p) * minimum(ZERO, dvsw0) * power(qg * rho, B_VG), + (A5 + A6 * p) * minimum(wpfloat(0.0), dvsw0) * power(qg * rho, B_VG), (A7 + A8 * p) * dvsw * power(qg * rho, B_VG), ), ) - return where(qg > GraupelConsts.qmin, maximum(result, -qg / dt), ZERO) + return where(qg > GraupelConsts.qmin, maximum(result, -qg / dt), wpfloat(0.0)) @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) @@ -649,15 +649,14 @@ def _vapor_x_ice( # noqa: PLR0917 [too-many-positional-arguments] A_FACT = wpfloat(4.0) * AMI ** ( wpfloat(-1.0) / wpfloat(3.0) ) # Is (1.0/5.065797019100886) * 4.0 - ZERO = wpfloat(0.0) # TO-DO: see if this can be folded into the WHERE statement result = (A_FACT * eta) * rho * qi * power(mi, B_EXP) * dvsi result = where( - result > ZERO, + result > wpfloat(0.0), minimum(result, dvsi / dt), maximum(maximum(result, dvsi / dt), -qi / dt), ) - return where(qi > GraupelConsts.qmin, result, ZERO) + return where(qi > GraupelConsts.qmin, result, wpfloat(0.0)) @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) @@ -719,29 +718,30 @@ def _vapor_x_snow( # noqa: PLR0917 [too-many-positional-arguments] C2_VS = wpfloat(0.241897) C3_VS = wpfloat(0.28003) C4_VS = wpfloat(-0.146293e-6) - ZERO = wpfloat(0.0) # See if this can be incorporated into WHERE statement result = where( (t < ThermodynamicConsts.tmelt), (CNX * ns * eta / rho) * (A0_VS + A1_VS * power(lam, A2_VS)) * dvsi / (lam * lam + EPS), - ZERO, + wpfloat(0.0), ) # GZ: This mask>0 limitation, which was missing in the original graupel scheme, # is crucial for numerical stability in the tropics! # a meaningful distinction between cloud ice and snow result = where( - (t < ThermodynamicConsts.tmelt) & (result > ZERO), + (t < ThermodynamicConsts.tmelt) & (result > wpfloat(0.0)), minimum(result, dvsi / dt - ice_dep), result, ) - result = where((t < ThermodynamicConsts.tmelt) & (qs <= QS_LIM), minimum(result, ZERO), result) + result = where( + (t < ThermodynamicConsts.tmelt) & (qs <= QS_LIM), minimum(result, wpfloat(0.0)), result + ) # ELSE section result = where( (t >= ThermodynamicConsts.tmelt) & (t > (ThermodynamicConsts.tmelt - GraupelConsts.tx * dvsw0)), - (C1_VS / p + C2_VS) * minimum(ZERO, dvsw0) * power(qs * rho, B_VS), + (C1_VS / p + C2_VS) * minimum(wpfloat(0.0), dvsw0) * power(qs * rho, B_VS), result, ) result = where( @@ -750,7 +750,7 @@ def _vapor_x_snow( # noqa: PLR0917 [too-many-positional-arguments] (C3_VS + C4_VS * p) * dvsw * power(qs * rho, B_VS), result, ) - return where((qs > GraupelConsts.qmin), maximum(result, -qs / dt), ZERO) + return where((qs > GraupelConsts.qmin), maximum(result, -qs / dt), wpfloat(0.0)) @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) @@ -784,3 +784,189 @@ def vapor_x_snow( # noqa: PLR0917 [too-many-positional-arguments] dt=dt, out=vapor_deposition_rate, ) + + +@gtx.field_operator +def _cloud_to_rain_aes_graupel( + t: fa.CellKField[ta.wpfloat], + rho: fa.CellKField[ta.wpfloat], + qc: fa.CellKField[ta.wpfloat], + qr: fa.CellKField[ta.wpfloat], + nc: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + """ + Compute the conversion rate from cloud to rain, AES_GRAUPEL scheme. + + Same autoconversion as _cloud_to_rain, but the SB2001 accretion kernel is a + degree-4 polynomial in log(clamped rho*qr) instead of a constant + (ICON mo_aes_graupel.f90 cloud_to_rain). + + Args: + t: Temperature + rho: Ambient density + qc: Cloud specific mass + qr: Rain water specific mass + nc: Cloud water number concentration + + Return: Conversion rate + """ + QMIN_AC = wpfloat(1.0e-6) # threshold for auto conversion + TAU_MAX = wpfloat(0.90e0) # maximum allowed value of tau + TAU_MIN = wpfloat(1.0e-30) # minimum allowed value of tau + A_PHI = wpfloat(6.0e2) # constant in phi-function for autoconversion + B_PHI = wpfloat(0.68e0) # exponent in phi-function for autoconversion + C_PHI = wpfloat(5.0e-5) # exponent in phi-function for accretion + X3 = wpfloat(2.0e0) # gamma exponent for cloud distribution + X2 = wpfloat(2.6e-10) # separating mass between cloud and rain + X1 = wpfloat(9.44e9) # kernel coeff for SB2001 autoconversion + AU_KERNEL = ( + X1 + / (wpfloat(20.0) * X2) + * (X3 + wpfloat(2.0)) + * (X3 + wpfloat(4.0)) + / ((X3 + wpfloat(1.0)) * (X3 + wpfloat(1.0))) + ) + + x = log(minimum(AesGraupelConsts.rhox_mx, maximum(AesGraupelConsts.rhox_mn, rho * qr))) + ac_kernel = AesGraupelConsts.a_ac_1 + x * ( + AesGraupelConsts.a_ac_2 + + x + * (AesGraupelConsts.a_ac_3 + x * (AesGraupelConsts.a_ac_4 + x * AesGraupelConsts.a_ac_5)) + ) + tau = maximum(TAU_MIN, minimum(wpfloat(1.0) - qc / (qc + qr), TAU_MAX)) + phi = power(tau, B_PHI) + phi = A_PHI * phi * power(wpfloat(1.0) - phi, wpfloat(3.0)) + xau = ( + AU_KERNEL + * power(qc * qc / nc, wpfloat(2.0)) + * (wpfloat(1.0) + phi / power(wpfloat(1.0) - tau, wpfloat(2.0))) + ) + xac = ac_kernel * qc * qr * power(tau / (tau + C_PHI), wpfloat(4.0)) + return where((qc > QMIN_AC) & (t > GraupelConsts.tfrz_hom), xau + xac, wpfloat(0.0)) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def cloud_to_rain_aes_graupel( # noqa: PLR0917 [too-many-positional-arguments] + t: fa.CellKField[ta.wpfloat], # Temperature + rho: fa.CellKField[ta.wpfloat], # Ambient density + qc: fa.CellKField[ta.wpfloat], # Cloud specific mass + qr: fa.CellKField[ta.wpfloat], # Rain water specific mass + nc: ta.wpfloat, # Cloud water number concentration + conversion_rate: fa.CellKField[ta.wpfloat], # output +): + _cloud_to_rain_aes_graupel(t=t, rho=rho, qc=qc, qr=qr, nc=nc, out=conversion_rate) + + +@gtx.field_operator +def _cloud_to_snow_aes_graupel( + t: fa.CellKField[ta.wpfloat], + qc: fa.CellKField[ta.wpfloat], + qs: fa.CellKField[ta.wpfloat], + ns: fa.CellKField[ta.wpfloat], + lam: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + """ + Compute the conversion rate from cloud to snow, AES_GRAUPEL scheme. + + Same as _cloud_to_snow with the additional riming tuning factor 3.0 + (ICON mo_aes_graupel.f90 cloud_to_snow). + + Args: + t: Temperature + qc: Cloud specific mass + qs: Snow specific mass + ns: Snow number + lam: Snow slope parameter (lambda) + + Return: Conversion rate + """ + ECS = wpfloat(0.9) + B_RIM = -(wpfloat(GraupelConsts.v1s) + wpfloat(3.0)) + # ICON hardcodes the rounded 2.61 (pi*gam(v1s+3)/4 = 2.6102); keep literal for bit-exactness. + # The extra 3.0 is the AES riming tuning factor (ICON). + C_RIM = wpfloat(2.61) * ECS * GraupelConsts.v0s * wpfloat(3.0) + return where( + (minimum(qc, qs) > GraupelConsts.qmin) & (t > GraupelConsts.tfrz_hom), + C_RIM * ns * qc * power(lam, B_RIM), + wpfloat(0.0), + ) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def cloud_to_snow_aes_graupel( # noqa: PLR0917 [too-many-positional-arguments] + t: fa.CellKField[ta.wpfloat], # Temperature + qc: fa.CellKField[ta.wpfloat], # Cloud specific mass + qs: fa.CellKField[ta.wpfloat], # Snow specific mass + ns: fa.CellKField[ta.wpfloat], # Snow number + lam: fa.CellKField[ta.wpfloat], # Snow slope parameter + riming_snow_rate: fa.CellKField[ta.wpfloat], # output +): + _cloud_to_snow_aes_graupel(t=t, qc=qc, qs=qs, ns=ns, lam=lam, out=riming_snow_rate) + + +@gtx.field_operator +def _rain_to_vapor_aes_graupel( # noqa: PLR0917 [too-many-positional-arguments] + t: fa.CellKField[ta.wpfloat], + rho: fa.CellKField[ta.wpfloat], + qc: fa.CellKField[ta.wpfloat], + qr: fa.CellKField[ta.wpfloat], + dvsw: fa.CellKField[ta.wpfloat], + dt: ta.wpfloat, +) -> fa.CellKField[ta.wpfloat]: + """ + Compute the conversion rate from rain to vapor, AES_GRAUPEL scheme. + + Same trigger and evaporation cap as _rain_to_vapor, but the evaporation rate + is the exponential of a degree-4 polynomial in log(clamped rho*qr) instead of + a power law (ICON mo_aes_graupel.f90 rain_to_vapor). + + Args: + t: Temperature + rho: Ambient density + qc: Cloud specific mass + qr: Rain specific mass + dvsw: qv-qsat_water (T) + dt: Time step + + Return: Conversion rate + """ + C1_RV = wpfloat(0.61) # coefficient for tc^0 in quadratic expansion + C2_RV = wpfloat(-0.0163) # coefficient of tc^1 in quadratic expansion + C3_RV = wpfloat(1.111e-4) # coefficient of tc^2 in quadratic expansion + + tc = t - ThermodynamicConsts.tmelt + evap_max = (C1_RV + tc * (C2_RV + C3_RV * tc)) * (-dvsw) / dt + x = log(minimum(AesGraupelConsts.rhox_mx, maximum(AesGraupelConsts.rhox_mn, qr * rho))) + evap = ( + -exp( + AesGraupelConsts.a_ev_1 + + x + * ( + AesGraupelConsts.a_ev_2 + + x + * ( + AesGraupelConsts.a_ev_3 + + x * (AesGraupelConsts.a_ev_4 + x * AesGraupelConsts.a_ev_5) + ) + ) + ) + * dvsw + ) + return where( + (qr > GraupelConsts.qmin) & (dvsw + qc <= wpfloat(0.0)), + minimum(evap, evap_max), + wpfloat(0.0), + ) + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def rain_to_vapor_aes_graupel( # noqa: PLR0917 [too-many-positional-arguments] + t: fa.CellKField[ta.wpfloat], # Temperature + rho: fa.CellKField[ta.wpfloat], # Ambient density + qc: fa.CellKField[ta.wpfloat], # Cloud specific mass + qr: fa.CellKField[ta.wpfloat], # Rain specific mass + dvsw: fa.CellKField[ta.wpfloat], # qv-qsat_water(T) + dt: ta.wpfloat, # time step + conversion_rate: fa.CellKField[ta.wpfloat], # output +): + _rain_to_vapor_aes_graupel(t=t, rho=rho, qc=qc, qr=qr, dvsw=dvsw, dt=dt, out=conversion_rate) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/data.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/data.py new file mode 100644 index 0000000000..f13f95d41b --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/data.py @@ -0,0 +1,39 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +"""Field metadata for the MuphysComponent input/output contract.""" + +from __future__ import annotations + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import SPECIES +from icon4py.model.common.states import data, model + + +# muphys precip port -> common precip-registry key. +_PRECIP_KEY = dict( + pflx="precipitation_flux", + pr="rainfall_flux", + ps="snowfall_flux", + pi="icefall_flux", + pg="graupelfall_flux", + pre="precipitation_energy_flux", +) + +INPUTS_PROPERTIES: dict[str, model.FieldMetaData] = { + "dz": model.FieldMetaData(standard_name="layer_thickness", units="m"), + "te": data.DIAGNOSTIC_CF_ATTRIBUTES["temperature"], + "p": data.DIAGNOSTIC_CF_ATTRIBUTES["pressure"], + "rho": data.PROGNOSTIC_CF_ATTRIBUTES["air_density"], + **{f"q{s}": data.COMMON_TRACER_CF_ATTRIBUTES[f"q{s}"] for s in SPECIES}, +} + +OUTPUTS_PROPERTIES: dict[str, model.FieldMetaData] = { + "tend_temperature": data.TENDENCY_CF_ATTRIBUTES["temperature"], + **{f"tend_q{s}": data.TENDENCY_CF_ATTRIBUTES[f"q{s}"] for s in SPECIES}, + **{port: data.PRECIPITATION_CF_ATTRIBUTES[key] for port, key in _PRECIP_KEY.items()}, +} diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py index 2acffbe229..d3c6ce974b 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py @@ -17,6 +17,7 @@ from gt4py import next as gtx +from icon4py.model.atmosphere.subgrid_scale_physics.muphys import config from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core import saturation_adjustment from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import ( common, @@ -115,14 +116,28 @@ def setup_muphys( backend: model_backends.BackendLike, *, single_program: bool = False, + scheme: config.MuphysScheme = config.MuphysScheme.KOKKOS_MUPHYS, ): + # the GT4Py operators branch on a plain bool (the DSL has no match statement) + match scheme: + case config.MuphysScheme.AES_GRAUPEL: + use_aes_graupel = True + case config.MuphysScheme.KOKKOS_MUPHYS: + use_aes_graupel = False + case _: + raise ValueError(f"unknown muphys scheme: {scheme}") + if single_program: # TODO(havogt): make an option in gt4py for thread-safety? with utils.recursion_limit(10**5): muphys_program = model_options.setup_program( backend=backend, program=muphys.muphys_run, - constant_args={"dt": ta.wpfloat(dt), "qnc": ta.wpfloat(qnc)}, + constant_args={ + "dt": ta.wpfloat(dt), + "qnc": ta.wpfloat(qnc), + "use_aes_graupel": use_aes_graupel, + }, horizontal_sizes={ "horizontal_start": gtx.int32(0), "horizontal_end": inp.ncells, @@ -145,6 +160,7 @@ def setup_muphys( vertical_start=0, vertical_end=inp.nlev, enable_masking=True, + scheme=scheme, ) with utils.recursion_limit(10**5): # TODO(havogt): make an option in gt4py? saturation_adjustment_program = model_options.setup_program( diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py index 6c2bd00acc..2f9d3aebf5 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py @@ -19,6 +19,7 @@ from gt4py.next.instrumentation import metrics as gtx_metrics from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations +from icon4py.model.atmosphere.subgrid_scale_physics.muphys import config from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import common, utils from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations import ( graupel, @@ -67,9 +68,24 @@ def setup_graupel( vertical_end: int, enable_masking: bool = True, enable_dace_hooks: bool = True, + scheme: config.MuphysScheme = config.MuphysScheme.KOKKOS_MUPHYS, ): - if enable_dace_hooks: - assert model_backends.is_backend_descriptor(backend) + # the GT4Py operators branch on a plain bool (the DSL has no match statement) + match scheme: + case config.MuphysScheme.AES_GRAUPEL: + use_aes_graupel = True + case config.MuphysScheme.KOKKOS_MUPHYS: + use_aes_graupel = False + case _: + raise ValueError(f"unknown muphys scheme: {scheme}") + + if enable_dace_hooks and model_backends.is_backend_descriptor(backend): + # The graupel scan needs two dace auto-opt hooks. They can only be injected into + # the backend *descriptor* (a dict), before it is turned into a concrete backend. + # A concrete/resolved backend (e.g. the gtfn backend the standalone driver threads + # in) is not a descriptor: gtfn does not use these hooks, so we pass it through + # unmodified. The dace path always drives muphys with a descriptor, so it still + # gets the hooks. backend = copy.deepcopy(backend) if "optimization_args" not in backend: backend["optimization_args"] = {} @@ -85,6 +101,7 @@ def setup_graupel( "dt": ta.wpfloat(dt), "qnc": ta.wpfloat(qnc), "enable_masking": enable_masking, + "use_aes_graupel": use_aes_graupel, }, horizontal_sizes={ "horizontal_start": horizontal_start, diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py index ead3e0fb0c..cae7bf62a5 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py @@ -25,10 +25,16 @@ _ice_number, _ice_sticking, _snow_lambda, + _snow_lambda_aes_graupel, _snow_number, + _snow_number_aes_graupel, _vel_scale_factor_default_scalar, _vel_scale_factor_ice_scalar, _vel_scale_factor_snow_scalar, + _vm_graupel_aes_graupel_scalar, + _vm_ice_aes_graupel_scalar, + _vm_rain_aes_graupel_scalar, + _vm_snow_aes_graupel_scalar, ) from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import ( _internal_energy_scalar, @@ -39,13 +45,16 @@ from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import ( _cloud_to_graupel, _cloud_to_rain, + _cloud_to_rain_aes_graupel, _cloud_to_snow, + _cloud_to_snow_aes_graupel, _cloud_x_ice, _graupel_to_rain, _ice_to_graupel, _ice_to_snow, _rain_to_graupel, _rain_to_vapor, + _rain_to_vapor_aes_graupel, _snow_to_graupel, _snow_to_rain, _vapor_x_graupel, @@ -78,6 +87,9 @@ class IntegrationState(NamedTuple): t_state: TempState rho: ta.wpfloat pflx_tot: ta.wpfloat + # input temperature of the level, carried for the AES_GRAUPEL snow fall speed of + # the next level (snow_number depends on the previous level's temperature) + t_in: ta.wpfloat @gtx.field_operator @@ -121,6 +133,40 @@ def precip_qx_level_update( # noqa: PLR0917 [too-many-positional-arguments] ) +@gtx.field_operator +def precip_qx_level_update_aes_graupel( # noqa: PLR0917 [too-many-positional-arguments] + previous_level_q: PrecipStateQx, + zeta: ta.wpfloat, # dt/(2dz) + q: ta.wpfloat, # specific mass of hydrometeor + rho: ta.wpfloat, # density + fall_speed: ta.wpfloat, # vm at this level's hydrometeor density + vt_prev: ta.wpfloat, # vm at the level-averaged density of the level above + mask: bool, +) -> PrecipStateQx: + # AES_GRAUPEL variant: the fall speeds are evaluated by the caller with the full + # species-specific vm formulas (clamped densities, previous-level rho and t), + # so no vc factor is carried in the state + current_level_activated = previous_level_q.activated | mask + rho_x = q * rho + flx_eff = (rho_x / zeta) + wpfloat(2.0) * previous_level_q.p + flx_partial = minimum(rho_x * fall_speed, flx_eff) + + if current_level_activated: + vt = vt_prev if previous_level_q.activated else wpfloat(0.0) + x = (zeta * (flx_eff - flx_partial)) / ((wpfloat(1.0) + zeta * vt) * rho) # q update + p = (x * rho * vt + flx_partial) * wpfloat(0.5) # flux + else: + x = q + p = wpfloat(0.0) + + return PrecipStateQx( + x=x, + p=p, + vc=wpfloat(0.0), + activated=current_level_activated, + ) + + @gtx.field_operator def _temperature_update( # noqa: PLR0917 [too-many-positional-arguments] previous_level: TempState, @@ -183,6 +229,7 @@ def _temperature_update( # noqa: PLR0917 [too-many-positional-arguments] t_state=TempState(t=0.0, eflx=0.0, activated=False), rho=0.0, pflx_tot=0.0, + t_in=0.0, ), ) def _precip_and_t( # noqa: PLR0917 [too-many-positional-arguments] @@ -197,6 +244,7 @@ def _precip_and_t( # noqa: PLR0917 [too-many-positional-arguments] mask_g: bool, dt: ta.wpfloat, dz: ta.wpfloat, + use_aes_graupel: bool, ) -> IntegrationState: zeta = dt / (wpfloat(2.0) * dz) xrho = sqrt(GraupelConsts.rho_00 / rho) @@ -216,54 +264,107 @@ def _precip_and_t( # noqa: PLR0917 [too-many-positional-arguments] current_level_activated = any_mask | previous_level_activated # TODO(): Use of combined if-statement to reduce checks in case any of the masks or previous levels are not activated. Can be made unnecessary with future transformations. if current_level_activated: - r_update = precip_qx_level_update( - previous_level.r, - previous_level.rho, - IndexConsts.prefactor_r, - IndexConsts.exponent_r, - IndexConsts.offset_r, - zeta, - vc_r, - q.r, - rho, - mask_r, - ) - s_update = precip_qx_level_update( - previous_level.s, - previous_level.rho, - IndexConsts.prefactor_s, - IndexConsts.exponent_s, - IndexConsts.offset_s, - zeta, - vc_s, - q.s, - rho, - mask_s, - ) - i_update = precip_qx_level_update( - previous_level.i, - previous_level.rho, - IndexConsts.prefactor_i, - IndexConsts.exponent_i, - IndexConsts.offset_i, - zeta, - vc_i, - q.i, - rho, - mask_i, - ) - g_update = precip_qx_level_update( - previous_level.g, - previous_level.rho, - IndexConsts.prefactor_g, - IndexConsts.exponent_g, - IndexConsts.offset_g, - zeta, - vc_g, - q.g, - rho, - mask_g, - ) + if use_aes_graupel: + # full vm formulas: fall_speed at this level's hydrometeor density, + # vt at the level-averaged density using the previous level's rho and t + r_update = precip_qx_level_update_aes_graupel( + previous_level.r, + zeta, + q.r, + rho, + _vm_rain_aes_graupel_scalar(q.r * rho, rho), + _vm_rain_aes_graupel_scalar( + (previous_level.r.x + q.r) * wpfloat(0.5) * previous_level.rho, + previous_level.rho, + ), + mask_r, + ) + s_update = precip_qx_level_update_aes_graupel( + previous_level.s, + zeta, + q.s, + rho, + _vm_snow_aes_graupel_scalar(q.s * rho, rho, t), + _vm_snow_aes_graupel_scalar( + (previous_level.s.x + q.s) * wpfloat(0.5) * previous_level.rho, + previous_level.rho, + previous_level.t_in, + ), + mask_s, + ) + i_update = precip_qx_level_update_aes_graupel( + previous_level.i, + zeta, + q.i, + rho, + _vm_ice_aes_graupel_scalar(q.i * rho, rho), + _vm_ice_aes_graupel_scalar( + (previous_level.i.x + q.i) * wpfloat(0.5) * previous_level.rho, + previous_level.rho, + ), + mask_i, + ) + g_update = precip_qx_level_update_aes_graupel( + previous_level.g, + zeta, + q.g, + rho, + _vm_graupel_aes_graupel_scalar(q.g * rho, rho), + _vm_graupel_aes_graupel_scalar( + (previous_level.g.x + q.g) * wpfloat(0.5) * previous_level.rho, + previous_level.rho, + ), + mask_g, + ) + else: # KOKKOS_MUPHYS (gt4py cannot prove an if/elif on a bool exhaustive) + r_update = precip_qx_level_update( + previous_level.r, + previous_level.rho, + IndexConsts.prefactor_r, + IndexConsts.exponent_r, + IndexConsts.offset_r, + zeta, + vc_r, + q.r, + rho, + mask_r, + ) + s_update = precip_qx_level_update( + previous_level.s, + previous_level.rho, + IndexConsts.prefactor_s, + IndexConsts.exponent_s, + IndexConsts.offset_s, + zeta, + vc_s, + q.s, + rho, + mask_s, + ) + i_update = precip_qx_level_update( + previous_level.i, + previous_level.rho, + IndexConsts.prefactor_i, + IndexConsts.exponent_i, + IndexConsts.offset_i, + zeta, + vc_i, + q.i, + rho, + mask_i, + ) + g_update = precip_qx_level_update( + previous_level.g, + previous_level.rho, + IndexConsts.prefactor_g, + IndexConsts.exponent_g, + IndexConsts.offset_g, + zeta, + vc_g, + q.g, + rho, + mask_g, + ) qliq = q.c + r_update.x qice = s_update.x + i_update.x + g_update.x @@ -296,6 +397,7 @@ def _precip_and_t( # noqa: PLR0917 [too-many-positional-arguments] t_state=t_update, rho=rho, pflx_tot=s_update.p + i_update.p + g_update.p + r_update.p, + t_in=t, ) @@ -339,7 +441,7 @@ def sink_saturation( @gtx.field_operator -def _q_t_update( # noqa: PLR0917 [too-many-positional-arguments] +def _q_t_update( # noqa: PLR0917, PLR0915 [too-many-positional-arguments, too-many-statements] t: fa.CellKField[ta.wpfloat], p: fa.CellKField[ta.wpfloat], rho: fa.CellKField[ta.wpfloat], @@ -347,6 +449,7 @@ def _q_t_update( # noqa: PLR0917 [too-many-positional-arguments] dt: ta.wpfloat, qnc: ta.wpfloat, enable_masking: bool, + use_aes_graupel: bool, ) -> tuple[ Q, fa.CellKField[ta.wpfloat], @@ -359,19 +462,29 @@ def _q_t_update( # noqa: PLR0917 [too-many-positional-arguments] dvsw = q.v - _qsat_rho(t, rho) qvsi = _qsat_ice_rho(t, rho) dvsi = q.v - qvsi - n_snow = _snow_number(t, rho, q.s) - - l_snow = _snow_lambda(rho, q.s, n_snow) + if use_aes_graupel: + n_snow = _snow_number_aes_graupel(t, rho * q.s) + l_snow = _snow_lambda_aes_graupel(rho * q.s, n_snow) + else: # KOKKOS_MUPHYS (gt4py cannot prove an if/elif on a bool exhaustive) + n_snow = _snow_number(t, rho, q.s) + l_snow = _snow_lambda(rho, q.s, n_snow) t_below_tmelt = t < ThermodynamicConsts.tmelt t_at_least_tmelt = ~t_below_tmelt # Define conversion 'matrix' - c2r = _cloud_to_rain(t, q.c, q.r, qnc) - r2v = _rain_to_vapor(t, rho, q.c, q.r, dvsw, dt) + if use_aes_graupel: + c2r = _cloud_to_rain_aes_graupel(t, rho, q.c, q.r, qnc) + r2v = _rain_to_vapor_aes_graupel(t, rho, q.c, q.r, dvsw, dt) + else: # KOKKOS_MUPHYS (gt4py cannot prove an if/elif on a bool exhaustive) + c2r = _cloud_to_rain(t, q.c, q.r, qnc) + r2v = _rain_to_vapor(t, rho, q.c, q.r, dvsw, dt) c2i, i2c = symmetric(_cloud_x_ice(t, q.c, q.i, dt)) - c2s = _cloud_to_snow(t, q.c, q.s, n_snow, l_snow) + if use_aes_graupel: + c2s = _cloud_to_snow_aes_graupel(t, q.c, q.s, n_snow, l_snow) + else: # KOKKOS_MUPHYS (gt4py cannot prove an if/elif on a bool exhaustive) + c2s = _cloud_to_snow(t, q.c, q.s, n_snow, l_snow) c2g = _cloud_to_graupel(t, rho, q.c, q.g) c2r = where(t_at_least_tmelt, c2r + c2s + c2g, c2r) @@ -386,11 +499,22 @@ def _q_t_update( # noqa: PLR0917 [too-many-positional-arguments] v2i, i2v = cond_symmetric( t_below_tmelt & is_sig_present, _vapor_x_ice(q.i, m_ice, eta, dvsi, rho, dt) ) - v2i = where( - t_below_tmelt, v2i + _ice_deposition_nucleation(t, q.c, q.i, n_ice, dvsi, dt), wpfloat(0.0) - ) # 0.0 or v2i both OK - - ice_dep = where(t_below_tmelt, minimum(v2i, dvsi / dt), wpfloat(0.0)) + if use_aes_graupel: + # ICON computes ice_dep from the vapor_x_ice deposition alone, BEFORE + # the nucleation contribution is added (and only where is_sig_present) + ice_dep = where(t_below_tmelt & is_sig_present, minimum(v2i, dvsi / dt), wpfloat(0.0)) + v2i = where( + t_below_tmelt, + v2i + _ice_deposition_nucleation(t, q.c, q.i, n_ice, dvsi, dt), + wpfloat(0.0), + ) + else: # KOKKOS_MUPHYS (gt4py cannot prove an if/elif on a bool exhaustive) + v2i = where( + t_below_tmelt, + v2i + _ice_deposition_nucleation(t, q.c, q.i, n_ice, dvsi, dt), + wpfloat(0.0), + ) # 0.0 or v2i both OK + ice_dep = where(t_below_tmelt, minimum(v2i, dvsi / dt), wpfloat(0.0)) # TODO(): _deposition_auto_conversion yields roundoff differences in i2s i2s = where( t_below_tmelt & is_sig_present, @@ -495,6 +619,7 @@ def _precipitation_effects( # noqa: PLR0917 [too-many-positional-arguments] rho: fa.CellKField[ta.wpfloat], # density dz: fa.CellKField[ta.wpfloat], dt: ta.wpfloat, + use_aes_graupel: bool, ) -> tuple[ fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], @@ -521,6 +646,7 @@ def _precipitation_effects( # noqa: PLR0917 [too-many-positional-arguments] kmin_g, dt, dz, + use_aes_graupel, ) qr = precip_state.r.x pr = precip_state.r.p @@ -550,6 +676,7 @@ def graupel( # noqa: PLR0917 [too-many-positional-arguments] dt: ta.wpfloat, qnc: ta.wpfloat, enable_masking: bool, + use_aes_graupel: bool, ) -> tuple[ fa.CellKField[ta.wpfloat], Q, @@ -569,9 +696,15 @@ def graupel( # noqa: PLR0917 [too-many-positional-arguments] | ((te < GraupelConsts.tfrz_het2) & (q.v > _qsat_ice_rho(te, rho))) | ~enable_masking ) - q, t = where(mask, _q_t_update(te, p, rho, q, dt, qnc, enable_masking=enable_masking), (q, te)) + q, t = where( + mask, + _q_t_update( + te, p, rho, q, dt, qnc, enable_masking=enable_masking, use_aes_graupel=use_aes_graupel + ), + (q, te), + ) qr, qs, qi, qg, t, pflx, pr, ps, pi, pg, pre = _precipitation_effects( - last_level, kmin_r, kmin_i, kmin_s, kmin_g, q, t, rho, dz, dt + last_level, kmin_r, kmin_i, kmin_s, kmin_g, q, t, rho, dz, dt, use_aes_graupel ) return t, Q(v=q.v, c=q.c, r=qr, s=qs, i=qi, g=qg), pflx, pr, ps, pi, pg, pre @@ -599,6 +732,7 @@ def graupel_run( # noqa: PLR0917 [too-many-positional-arguments] vertical_start: gtx.int32, vertical_end: gtx.int32, enable_masking: bool, + use_aes_graupel: bool, ) -> None: graupel( last_level=vertical_end - 1, @@ -610,6 +744,7 @@ def graupel_run( # noqa: PLR0917 [too-many-positional-arguments] dt=dt, qnc=qnc, enable_masking=enable_masking, + use_aes_graupel=use_aes_graupel, out=(t_out, q_out, pflx, pr, ps, pi, pg, pre), domain=( # t_out diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py index 550d1e8d57..b812996053 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py @@ -26,6 +26,7 @@ def _muphys( # noqa: PLR0917 [too-many-positional-arguments] q_in: Q, dt: ta.wpfloat, qnc: ta.wpfloat, + use_aes_graupel: bool, ) -> tuple[ fa.CellKField[ta.wpfloat], Q, @@ -52,6 +53,7 @@ def _muphys( # noqa: PLR0917 [too-many-positional-arguments] dt=dt, qnc=qnc, enable_masking=True, # TODO(havogt): expose this option when optimizing full muphys + use_aes_graupel=use_aes_graupel, ) te, qve, qce = _saturation_adjustment( @@ -60,7 +62,9 @@ def _muphys( # noqa: PLR0917 [too-many-positional-arguments] rho=rho, ) - return t, Q(v=qve, c=qce, r=q.r, s=q.s, i=q.i, g=q.g), pflx, pr, ps, pi, pg, pre + # return the temperature updated by the final saturation adjustment (as the + # Fortran does), consistent with the adjusted qv/qc + return te, Q(v=qve, c=qce, r=q.r, s=q.s, i=q.i, g=q.g), pflx, pr, ps, pi, pg, pre @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) @@ -84,6 +88,7 @@ def muphys_run( # noqa: PLR0917 [too-many-positional-arguments] horizontal_end: gtx.int32, vertical_start: gtx.int32, vertical_end: gtx.int32, + use_aes_graupel: bool, ) -> None: _muphys( last_level=vertical_end - 1, @@ -94,6 +99,7 @@ def muphys_run( # noqa: PLR0917 [too-many-positional-arguments] q_in=q_in, dt=dt, qnc=qnc, + use_aes_graupel=use_aes_graupel, out=(t_out, q_out, pflx, pr, ps, pi, pg, pre), domain=( # t_out diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/state.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/state.py new file mode 100644 index 0000000000..e3127d7c12 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/state.py @@ -0,0 +1,275 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import datetime +from typing import TYPE_CHECKING + +import gt4py.next as gtx + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import ( + PRECIP_DIAGNOSTICS, + SPECIES, +) +from icon4py.model.common import ( + dimension as dims, + field_type_aliases as fa, + model_options, + type_alias as ta, +) +from icon4py.model.common.components.physics_state import PhysicsState +from icon4py.model.common.diagnostic_calculations.stencils import ( + calculate_tendency, + diagnose_pressure, + diagnose_surface_pressure, + diagnose_temperature, + update_exner_and_theta_v, +) +from icon4py.model.common.math.stencils import generic_math_operations +from icon4py.model.common.metrics import metrics_attributes +from icon4py.model.common.utils import data_allocation as data_alloc + + +if TYPE_CHECKING: + import gt4py.next.typing as gtx_typing + + from icon4py.model.common.grid import base as base_grid + from icon4py.model.common.states import factory, prognostic_state as prognostics, tracer_states + + +class State(PhysicsState): + """The muphys physics State adapter. + + Bridges the dycore's prognostic state and the muphys Component contract. + Two independent axes describe each field: + + muphys role + - input : fed to muphys via ``as_component_input`` -- dz, rho, q, te, p + - returned : muphys updates it; te/q changes come back as tendencies. tend_q + updates the tracers; tend_T drives the exner + theta_v update + (via the exact EOS, mirroring ICON's phy2dyn coupling). + rho and p are input-only. + - internal : not a muphys fields -- tv, pressure_on_cells_half_levels -- used only + to diagnose p and to recompute exner/theta_v from tend_T. + - diagnostic : a muphys output stored for reporting -- pflx, pr, ps, pi, pg, pre. + + memory ownership + - reference : a pointer into the dycore state, no copy -- dz, rho, q. + - owned : a buffer allocated once here and overwritten in place each + step -- te, p, tv, pressure_on_cells_half_levels, and the scratch buffers. + """ + + def __init__( + self, + grid: base_grid.Grid, + metrics: factory.FieldSource, + backend: gtx_typing.Backend | None = None, + ) -> None: + self._num_cells = grid.num_cells + self._num_levels = grid.num_levels + self._backend = backend + + full_horizontal = { + "horizontal_start": gtx.int32(0), + "horizontal_end": gtx.int32(self._num_cells), + } + full_vertical = { + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(self._num_levels), + } + + self._diagnose_temperature = model_options.setup_program( + program=diagnose_temperature.diagnose_virtual_temperature_and_temperature, + backend=self._backend, + horizontal_sizes=full_horizontal, + vertical_sizes=full_vertical, + offset_provider={}, + ) + self._diagnose_surface_pressure = model_options.setup_program( + program=diagnose_surface_pressure.diagnose_surface_pressure, + backend=self._backend, + horizontal_sizes=full_horizontal, + vertical_sizes={ + "vertical_start": gtx.int32(self._num_levels), + "vertical_end": gtx.int32(self._num_levels + 1), + }, + offset_provider={}, + ) + self._diagnose_pressure = model_options.setup_program( + program=diagnose_pressure.diagnose_pressure, + backend=self._backend, + horizontal_sizes=full_horizontal, + vertical_sizes=full_vertical, + offset_provider={}, + ) + self._apply_tendency = model_options.setup_program( + program=generic_math_operations.compute_field_a_plus_coeff_times_field_b_on_cell_k, + backend=self._backend, + horizontal_sizes=full_horizontal, + vertical_sizes=full_vertical, + offset_provider={}, + ) + self._calculate_virtual_temperature_tendency = model_options.setup_program( + program=calculate_tendency.calculate_virtual_temperature_tendency, + backend=self._backend, + horizontal_sizes=full_horizontal, + vertical_sizes=full_vertical, + offset_provider={}, + ) + self._update_exner_and_theta_v = model_options.setup_program( + program=update_exner_and_theta_v.update_exner_and_theta_v, + backend=self._backend, + horizontal_sizes=full_horizontal, + vertical_sizes=full_vertical, + offset_provider={}, + ) + + self.dz = metrics.get(metrics_attributes.DDQZ_Z_FULL) + self.rho: fa.CellKField[ta.wpfloat] | None = None + self._tracers: tracer_states.TracerState | None = None + self.te = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, allocator=backend) + self.p = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, allocator=backend) + self.tv = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, allocator=backend) + self.pressure_on_cells_half_levels = data_alloc.zero_field( + grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, allocator=backend + ) + + # INTERNAL + self._new_te = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, allocator=backend) + self._tv_tendency = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, allocator=backend) + self._precip_diagnostics: dict[str, fa.CellKField[ta.wpfloat]] | None = None + + def gather_from_prognostic( + self, prognostic: prognostics.PrognosticState, tracers: tracer_states.TracerState + ) -> None: + """ + prepare the input fields for muphys from the prognostic state. This includes: + - binding the references for rho and q (the muphys input fields that are stored prognostically in the dycore state) + - diagnosing the muphys input fields that aren't stored prognostically (te, p) from the prognostic state. + """ + self.rho = prognostic.rho + # muphys needs all six moisture species; TracerState fields are optional (a + # tracer may be inactive per TracerConfig), so fail loudly once here rather + # than feed None into the microphysics. + missing = [f"q{s}" for s in SPECIES if getattr(tracers, f"q{s}") is None] + if missing: + raise ValueError( + f"muphys requires all moisture species active in the TracerState; missing: {missing}" + ) + self._tracers = tracers + + # Diagnose virtual temperature and temperature (te is not stored prognostically). + self._diagnose_temperature( + qv=tracers.qv, + qc=tracers.qc, + qi=tracers.qi, + qr=tracers.qr, + qs=tracers.qs, + qg=tracers.qg, + theta_v=prognostic.theta_v, + exner=prognostic.exner, + virtual_temperature=self.tv, + temperature=self.te, + ) + + self._diagnose_surface_pressure( + exner=prognostic.exner, + virtual_temperature=self.tv, + ddqz_z_full=self.dz, + surface_pressure=self.pressure_on_cells_half_levels, + ) + surface_pressure = gtx.as_field( + (dims.CellDim,), + self.pressure_on_cells_half_levels.ndarray[:, -1], + allocator=self._backend, + ) + self._diagnose_pressure( + ddqz_z_full=self.dz, + virtual_temperature=self.tv, + surface_pressure=surface_pressure, + pressure=self.p, + pressure_ifc=self.pressure_on_cells_half_levels, + ) + + def scatter_to_prognostic( + self, + prognostic: prognostics.PrognosticState, + outputs: dict[str, fa.CellKField[ta.wpfloat]], + dtime: datetime.timedelta, + ) -> None: + """Outbound translation: apply muphys output (tendencies) back to the prognostic state. + + This will be called after calling the muphys. + output is got from muphys, and the tendencies in output will be applied to the prognostic state. + """ + assert self._tracers is not None, "gather_from_prognostic must be called first" + # convert to seconds only at the gt4py boundary (stencils take a scalar dt) + dt_seconds = dtime.total_seconds() + # 1. Apply moisture tendencies to the tracers (in place; tracers were bound in gather). + for s in SPECIES: + tracer = getattr(self._tracers, f"q{s}") + self._apply_tendency( + field_a=tracer, + coeff=dt_seconds, + field_b=outputs[f"tend_q{s}"], + output_field=tracer, + ) + + # 2. tend_T -> new temperature: new_te = te + tend_T*dt + self._apply_tendency( + field_a=self.te, + coeff=dt_seconds, + field_b=outputs["tend_temperature"], + output_field=self._new_te, + ) + + # dTv/dt from the new temperature and the species just updated in step 1 + self._calculate_virtual_temperature_tendency( + dtime=dt_seconds, + qv=self._tracers.qv, + qc=self._tracers.qc, + qi=self._tracers.qi, + qr=self._tracers.qr, + qs=self._tracers.qs, + qg=self._tracers.qg, + temperature=self._new_te, + virtual_temperature=self.tv, + virtual_temperature_tendency=self._tv_tendency, + ) + # Recompute exner via the exact EOS from the updated virtual temperature and + # diagnose theta_v = Tv/exner, mirroring ICON's phy2dyn coupling + # (mo_interface_iconam_aes.f90). The exner/rho/theta_v trio stays EOS-consistent. + self._update_exner_and_theta_v( + rho=self.rho, + virtual_temperature=self.tv, + virtual_temperature_tendency=self._tv_tendency, + dtime=dt_seconds, + exner=prognostic.exner, + theta_v=prognostic.theta_v, + ) + + # 3. Store precip diagnostics (references; never applied to prognostic state). + self._precip_diagnostics = {name: outputs[name] for name in PRECIP_DIAGNOSTICS} + + def as_component_input(self) -> dict[str, fa.CellKField[ta.wpfloat]]: + """ + Translate to the generic Component input dict (the 10 muphys input fields). + """ + if self.rho is None or self._tracers is None: + raise RuntimeError("as_component_input called before gather_from_prognostic") + inp = {"dz": self.dz, "te": self.te, "p": self.p, "rho": self.rho} + inp.update({f"q{s}": getattr(self._tracers, f"q{s}") for s in SPECIES}) + return inp + + @property + def precip_diagnostics(self) -> dict[str, fa.CellKField[ta.wpfloat]]: + """muphys precip diagnostics keyed by name, ready for IO / plotting.""" + if self._precip_diagnostics is None: + raise RuntimeError("precip_diagnostics accessed before scatter_to_prognostic") + return self._precip_diagnostics diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/fixtures.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/fixtures.py index de9850de36..0635cbbdc8 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/fixtures.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/fixtures.py @@ -5,3 +5,15 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause + +from icon4py.model.testing.fixtures.datatest import ( + backend, + backend_like, + data_provider, + download_ser_data, + experiment, + experiment_description, + grid_savepoint, + icon_grid, + process_props, +) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_component_datatest.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_component_datatest.py new file mode 100644 index 0000000000..703f5bbc03 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_component_datatest.py @@ -0,0 +1,125 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import datetime + +import pytest +from gt4py import next as gtx + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.component import MuphysComponent +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import SPECIES +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import common, run_full_muphys +from icon4py.model.common import dimension as dims, model_backends +from icon4py.model.testing import test_utils +from icon4py.model.testing.fixtures.datatest import backend_like + +from . import utils +from .utils import download_test_data + + +_T0 = datetime.datetime(2024, 1, 1, 0, 0, 0) +_MINI = utils.MuphysExperiment(name="mini", type=utils.ExperimentType.FULL_MUPHYS) + + +class _FullDomainGrid: + """Grid stand-in for the grid-less muphys netCDF data: prognostic bounds = full domain.""" + + def __init__(self, num_cells: int, num_levels: int) -> None: + self.num_cells = num_cells + self.num_levels = num_levels + + def start_index(self, domain: object) -> gtx.int32: + return gtx.int32(0) + + def end_index(self, domain: object) -> gtx.int32: + return gtx.int32(self.num_cells) + + +@pytest.mark.uses_concat_where +@pytest.mark.datatest +@pytest.mark.level("integration") +@pytest.mark.parametrize("experiment", [_MINI], ids=lambda e: e.name) +def test_granule_matches_direct_muphys( + backend_like: model_backends.BackendLike, + experiment: utils.MuphysExperiment, +) -> None: + allocator = model_backends.get_allocator(backend_like) + graupel_input = common.GraupelInput.load(filename=experiment.input_file, allocator=allocator) + + te0 = graupel_input.t.asnumpy().copy() + q0 = {s: getattr(graupel_input, f"q{s}").asnumpy().copy() for s in SPECIES} + + muphys_program = run_full_muphys.setup_muphys( + inp=graupel_input, + dt=experiment.dt, + qnc=experiment.qnc, + backend=backend_like, + single_program=False, + ) + + granule = MuphysComponent( + grid=_FullDomainGrid(graupel_input.ncells, graupel_input.nlev), # type: ignore[arg-type] # mini data has no icon grid + dtime=datetime.timedelta(seconds=experiment.dt), + qnc=experiment.qnc, + backend=backend_like, + step=muphys_program, + ) + state = { + "dz": graupel_input.dz, + "te": graupel_input.t, + "p": graupel_input.p, + "rho": graupel_input.rho, + "qv": graupel_input.qv, + "qc": graupel_input.qc, + "qr": graupel_input.qr, + "qs": graupel_input.qs, + "qi": graupel_input.qi, + "qg": graupel_input.qg, + } + out = granule(state, _T0) + + direct = common.GraupelOutput.allocate( + allocator=allocator, + domain=gtx.domain({dims.CellDim: graupel_input.ncells, dims.KDim: graupel_input.nlev}), + ) + muphys_program( + dz=graupel_input.dz, + te=graupel_input.t, + p=graupel_input.p, + rho=graupel_input.rho, + q_in=graupel_input.q, + t_out=direct.t, + q_out=direct.q, + pflx=direct.pflx, + pr=direct.pr, + ps=direct.ps, + pi=direct.pi, + pg=direct.pg, + pre=direct.pre, + ) + + dt = experiment.dt + + # Reconstructing the updated state as ``old + tendency*dt`` is not bit-exact + assert test_utils.dallclose( + te0 + out["tend_temperature"].asnumpy() * dt, direct.t.asnumpy(), atol=1e-15 + ) + for s in SPECIES: + applied = q0[s] + out[f"tend_q{s}"].asnumpy() * dt + assert test_utils.dallclose(applied, getattr(direct, f"q{s}").asnumpy(), atol=1e-15) + + assert test_utils.dallclose(out["pflx"].asnumpy(), direct.pflx.asnumpy(), rtol=0.0, atol=0.0) + for name in ("pr", "ps", "pi", "pg", "pre"): + assert test_utils.dallclose( + out[name].asnumpy()[:, -1], + getattr(direct, name).asnumpy()[:, -1], + rtol=0.0, + atol=0.0, + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py index 564aa15db6..d756e4c6fd 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py @@ -41,6 +41,7 @@ class Experiments: @pytest.mark.uses_concat_where @pytest.mark.datatest +@pytest.mark.level("integration") @pytest.mark.parametrize( "experiment", [ diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py index d76a2a91c2..78b81d0a1d 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py @@ -48,6 +48,7 @@ class Experiments: @pytest.mark.uses_concat_where @pytest.mark.datatest +@pytest.mark.level("integration") @pytest.mark.parametrize( ("experiment", "enable_dace_hooks"), _GRAUPEL_TEST_CASES, diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_muphys_datatest.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_muphys_datatest.py new file mode 100644 index 0000000000..9a8d050030 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_muphys_datatest.py @@ -0,0 +1,147 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations + +import datetime +from typing import TYPE_CHECKING + +import pytest + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys import ( + component as muphys_component, + config as muphys_config, +) +from icon4py.model.common.states.data import QC, QG, QI, QR, QS, QV +from icon4py.model.testing import definitions, test_utils + +from ..fixtures import * # noqa: F403 + + +if TYPE_CHECKING: + import gt4py.next.typing as gtx_typing + + from icon4py.model.common.grid import icon as icon_grid_types + from icon4py.model.testing import serialbox as sb + + +# Validation of the muphys granule against Fortran ICON: the aes-graupel-init/exit +# savepoints are written around the mig block in aes_phy_main (cloud_mig = +# satad + graupel + satad), which is exactly the composition MuphysComponent runs +# (setup_muphys with single_program=False). Inputs are captured after ICON's +# dyn2phy negative-tracer clipping, so the granule sees identical state. +# +# ICON only computes the scheme on levels jks_cloudy..nlev (zmaxcloudy cutoff) +# while the granule runs the full column, hence the comparisons are restricted +# to the cloudy levels and the levels above are separately checked to produce +# (near-)zero tendencies. +# +# The exit tendencies are the prm_tend accumulators; subtracting the init +# accumulators isolates the mig contribution even if other AES processes ever +# run before mig in this experiment (today they contribute exactly zero). +# +# The granule runs with MuphysScheme.AES_GRAUPEL, the port of the ICON +# formulation (mo_aes_graupel.f90) that generates the reference data, so +# near-roundoff agreement is expected. Remaining known deviations: +# - Fortran clamps tendencies to full depletion (MAX(-q/dt), mo_cloud_mig.f90) +# while the granule reports raw (new-old)/dt; differs only where the scheme +# drives a species below zero. +# - Fortran cvd is computed as cpd - rd (1 ULP off the icon4py literal 717.60); +# likewise tfrz_hom/tfrz_het2 are computed as tmelt-37/tmelt-25 (1 ULP below +# the icon4py literals 236.15/248.15) — both only matter at exact thresholds. +# - Fortran ice_sticking carries a cia tuning factor (cloud_mig_nml, default +# 1.0 and not set in the experiment, hence a no-op here). +@pytest.mark.uses_concat_where +@pytest.mark.datatest +@pytest.mark.level("integration") +@pytest.mark.parametrize( + "experiment_description", + [definitions.Experiments.EXCLAIM_APE_AES], +) +@pytest.mark.parametrize( + "date", + ["2008-09-01T00:00:00.000", "2008-09-01T00:05:00.000", "2008-09-01T00:10:00.000"], +) +def test_muphys_granule( + date: str, + *, + data_provider: sb.IconSerialDataProvider, + icon_grid: icon_grid_types.IconGrid, + backend: gtx_typing.Backend, +) -> None: + init_savepoint = data_provider.from_savepoint_muphys_init(date=date) + exit_savepoint = data_provider.from_savepoint_muphys_exit(date=date) + + dtime = init_savepoint.dtime() + # numpy index of the first level ICON computes the scheme on (Fortran jks_cloudy is 1-based) + jks = init_savepoint.jks_cloudy() - 1 + + # MuphysConfig().qnc matches the Fortran cloud_num = 50.0e6 m^-3 (mo_cloud_mig.f90); + # the default scheme is AES_GRAUPEL, matching the Fortran that generated the data + muphys_configuration = muphys_config.MuphysConfig() + component = muphys_component.MuphysComponent( + grid=icon_grid, + dtime=datetime.timedelta(seconds=dtime), + qnc=muphys_configuration.qnc, + backend=backend, + scheme=muphys_configuration.scheme, + ) + + state = { + "dz": init_savepoint.dz(), + "te": init_savepoint.temperature(), + "p": init_savepoint.pressure(), + "rho": init_savepoint.rho(), + "qv": init_savepoint.qv(), + "qc": init_savepoint.qc(), + "qr": init_savepoint.qr(), + "qs": init_savepoint.qs(), + "qi": init_savepoint.qi(), + "qg": init_savepoint.qg(), + } + outputs = component(state, datetime.datetime.fromisoformat(date)) + + # provisional tolerances; measure on the archive + # (ICON4PY_DALLCLOSE_PRINT_INSTEAD_OF_FAIL=true) and tighten + for name, tracer_index in ( + ("tend_qv", QV), + ("tend_qc", QC), + ("tend_qr", QR), + ("tend_qs", QS), + ("tend_qi", QI), + ("tend_qg", QG), + ): + reference = ( + exit_savepoint.tend_tracer(tracer_index).asnumpy() + - init_savepoint.tend_tracer(tracer_index).asnumpy() + ) + actual = outputs[name].asnumpy() + test_utils.assert_dallclose(actual[:, jks:], reference[:, jks:], atol=1e-13) + # above the cloudy region ICON does not run the scheme; the full-column + # granule must produce (near-)zero tendencies there + test_utils.assert_dallclose(actual[:, :jks], 0.0, atol=1e-12) + + tend_ta_reference = exit_savepoint.tend_ta().asnumpy() - init_savepoint.tend_ta().asnumpy() + tend_ta_actual = outputs["tend_temperature"].asnumpy() + test_utils.assert_dallclose(tend_ta_actual[:, jks:], tend_ta_reference[:, jks:], atol=1e-10) + test_utils.assert_dallclose(tend_ta_actual[:, :jks], 0.0, atol=1e-10) + + # surface precip: the granule keeps the surface value in the last level; ICON + # only stores the aggregated prm_field diagnostics (rsfl = rain, + # ssfl = ice + snow + graupel, pr = total, ufcs = energy flux) + rain = outputs["pr"].asnumpy()[:, -1] + ice = outputs["pi"].asnumpy()[:, -1] + snow = outputs["ps"].asnumpy()[:, -1] + graupel = outputs["pg"].asnumpy()[:, -1] + energy_flux = outputs["pre"].asnumpy()[:, -1] + + test_utils.assert_dallclose(rain, exit_savepoint.rsfl().asnumpy(), atol=1e-10) + test_utils.assert_dallclose(ice + snow + graupel, exit_savepoint.ssfl().asnumpy(), atol=1e-10) + test_utils.assert_dallclose( + rain + ice + snow + graupel, exit_savepoint.pr().asnumpy(), atol=1e-10 + ) + test_utils.assert_dallclose(energy_flux, exit_savepoint.ufcs().asnumpy(), atol=1e-10) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain_aes_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain_aes_graupel.py new file mode 100644 index 0000000000..b1be20aca1 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_rain_aes_graupel.py @@ -0,0 +1,59 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx +import numpy as np +import pytest + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import ( + cloud_to_rain_aes_graupel, +) +from icon4py.model.common import dimension as dims +from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing.stencil_tests import StencilTest + + +class TestCloudToRainAesGraupel(StencilTest): + PROGRAM = cloud_to_rain_aes_graupel + OUTPUTS = ("conversion_rate",) + + @staticmethod + def reference( + connectivities: dict[gtx.Dimension, np.ndarray], + *, + t: np.ndarray, + rho: np.ndarray, + qc: np.ndarray, + qr: np.ndarray, + nc: np.ndarray, + **kwargs, + ) -> dict: + # independent numpy reference mirroring ICON mo_aes_graupel.f90 cloud_to_rain; + # the magic constants are hardcoded PARAMETERs there in the same way + au_kernel = 9.44e9 / (20.0 * 2.6e-10) * 4.0 * 6.0 / 9.0 + a_ac = [-2.155543e00, -1.148491e00, -1.882563e-02, 2.941391e-03, 5.575598e-05] + x = np.log(np.clip(rho * qr, 3.26216e-08, 6.97604e-03)) + ac_kernel = a_ac[0] + x * (a_ac[1] + x * (a_ac[2] + x * (a_ac[3] + x * a_ac[4]))) + tau = np.maximum(1.0e-30, np.minimum(1.0 - qc / (qc + qr), 0.9)) + phi = tau**0.68 + phi = 6.0e2 * phi * (1.0 - phi) ** 3 + xau = au_kernel * (qc * qc / nc) ** 2 * (1.0 + phi / (1.0 - tau) ** 2) + xac = ac_kernel * qc * qr * (tau / (tau + 5.0e-5)) ** 4 + rate = np.where((qc > 1.0e-6) & (t > 236.15), xau + xac, 0.0) + return dict(conversion_rate=rate) + + @pytest.fixture + def input_data(self, grid): + return dict( + t=data_alloc.constant_field(grid, 267.25, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=data_alloc.constant_field(grid, 0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=data_alloc.constant_field(grid, 5.52921e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + qr=data_alloc.constant_field(grid, 2.01511e-12, dims.CellDim, dims.KDim, dtype=wpfloat), + nc=100.0, + conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow_aes_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow_aes_graupel.py new file mode 100644 index 0000000000..ea369a010a --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_cloud_to_snow_aes_graupel.py @@ -0,0 +1,54 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx +import numpy as np +import pytest + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import ( + cloud_to_snow_aes_graupel, +) +from icon4py.model.common import dimension as dims +from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing.stencil_tests import StencilTest + + +class TestCloudToSnowAesGraupel(StencilTest): + PROGRAM = cloud_to_snow_aes_graupel + OUTPUTS = ("riming_snow_rate",) + + @staticmethod + def reference( + connectivities: dict[gtx.Dimension, np.ndarray], + *, + t: np.ndarray, + qc: np.ndarray, + qs: np.ndarray, + ns: np.ndarray, + lam: np.ndarray, + **kwargs, + ) -> dict: + # mirrors ICON mo_aes_graupel.f90 cloud_to_snow (riming tuning factor 3.0) + c_rim = 2.61 * 0.9 * 25.0 * 3.0 + rate = np.where( + (np.minimum(qc, qs) > 1.0e-15) & (t > 236.15), + c_rim * ns * qc * lam ** (-3.5), + 0.0, + ) + return dict(riming_snow_rate=rate) + + @pytest.fixture + def input_data(self, grid): + return dict( + t=data_alloc.constant_field(grid, 256.571, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=data_alloc.constant_field(grid, 3.31476e-05, dims.CellDim, dims.KDim, dtype=wpfloat), + qs=data_alloc.constant_field(grid, 7.47365e-06, dims.CellDim, dims.KDim, dtype=wpfloat), + ns=data_alloc.constant_field(grid, 3.37707e07, dims.CellDim, dims.KDim, dtype=wpfloat), + lam=data_alloc.constant_field(grid, 8989.78, dims.CellDim, dims.KDim, dtype=wpfloat), + riming_snow_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor_aes_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor_aes_graupel.py new file mode 100644 index 0000000000..0439b99aed --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_rain_to_vapor_aes_graupel.py @@ -0,0 +1,56 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx +import numpy as np +import pytest + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.transitions import ( + rain_to_vapor_aes_graupel, +) +from icon4py.model.common import dimension as dims +from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing.stencil_tests import StencilTest + + +class TestRainToVaporAesGraupel(StencilTest): + PROGRAM = rain_to_vapor_aes_graupel + OUTPUTS = ("conversion_rate",) + + @staticmethod + def reference( + connectivities: dict[gtx.Dimension, np.ndarray], + *, + t: np.ndarray, + rho: np.ndarray, + qc: np.ndarray, + qr: np.ndarray, + dvsw: np.ndarray, + dt: np.ndarray, + **kwargs, + ) -> dict: + # mirrors ICON mo_aes_graupel.f90 rain_to_vapor + a_ev = [-5.532194e00, 2.432848e-01, -4.145391e-02, -1.798439e-03, -1.405764e-05] + tc = t - 273.15 + evap_max = (0.61 + tc * (-0.0163 + 1.111e-4 * tc)) * (-dvsw) / dt + x = np.log(np.clip(qr * rho, 3.26216e-08, 6.97604e-03)) + evap = -np.exp(a_ev[0] + x * (a_ev[1] + x * (a_ev[2] + x * (a_ev[3] + x * a_ev[4])))) * dvsw + rate = np.where((qr > 1.0e-15) & (dvsw + qc <= 0.0), np.minimum(evap, evap_max), 0.0) + return dict(conversion_rate=rate) + + @pytest.fixture + def input_data(self, grid): + return dict( + t=data_alloc.constant_field(grid, 258.542, dims.CellDim, dims.KDim, dtype=wpfloat), + rho=data_alloc.constant_field(grid, 0.956089, dims.CellDim, dims.KDim, dtype=wpfloat), + qc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + qr=data_alloc.constant_field(grid, 3.01332e-11, dims.CellDim, dims.KDim, dtype=wpfloat), + dvsw=data_alloc.constant_field(grid, -1.0e-10, dims.CellDim, dims.KDim, dtype=wpfloat), + dt=30.0, + conversion_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_lambda_aes_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_lambda_aes_graupel.py new file mode 100644 index 0000000000..6ad98691ff --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_lambda_aes_graupel.py @@ -0,0 +1,45 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx +import numpy as np +import pytest + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import ( + snow_lambda_aes_graupel, +) +from icon4py.model.common import dimension as dims +from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing.stencil_tests import StencilTest + + +class TestSnowLambdaAesGraupel(StencilTest): + PROGRAM = snow_lambda_aes_graupel + OUTPUTS = ("riming_snow_rate",) + + @staticmethod + def reference( + connectivities: dict[gtx.Dimension, np.ndarray], + *, + rho_s: np.ndarray, + ns: np.ndarray, + **kwargs, + ) -> dict: + # mirrors ICON mo_aes_graupel.f90 snow_lambda + lam = np.where(rho_s > 1.0e-15, (2.0 * 0.069 * ns / rho_s) ** (1.0 / 3.0), 1.0e10) + return dict(riming_snow_rate=lam) + + @pytest.fixture + def input_data(self, grid): + return dict( + rho_s=data_alloc.constant_field( + grid, 1.12204 * 7.47365e-06, dims.CellDim, dims.KDim, dtype=wpfloat + ), + ns=data_alloc.constant_field(grid, 1.76669e07, dims.CellDim, dims.KDim, dtype=wpfloat), + riming_snow_rate=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number_aes_graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number_aes_graupel.py new file mode 100644 index 0000000000..3bcab844a0 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_snow_number_aes_graupel.py @@ -0,0 +1,53 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx +import numpy as np +import pytest + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import ( + snow_number_aes_graupel, +) +from icon4py.model.common import dimension as dims +from icon4py.model.common.type_alias import wpfloat +from icon4py.model.common.utils import data_allocation as data_alloc +from icon4py.model.testing.stencil_tests import StencilTest + + +class TestSnowNumberAesGraupel(StencilTest): + PROGRAM = snow_number_aes_graupel + OUTPUTS = ("number",) + + @staticmethod + def reference( + connectivities: dict[gtx.Dimension, np.ndarray], + *, + t: np.ndarray, + rho_s: np.ndarray, + **kwargs, + ) -> dict: + # mirrors ICON mo_aes_graupel.f90 snow_number + n0s1 = 13.5 * 5.65e05 + tc = np.clip(t, 233.15, 273.15) - 273.15 + alf = 10.0 ** (-1.65 + tc * (5.45e-2 + tc * 3.27e-4)) + bet = 1.42 + tc * (1.19e-2 + tc * 9.60e-5) + n0s = 13.5 * (np.maximum(rho_s, 2.0e-7) / 0.069) ** (4.0 - 3.0 * bet) / (alf**3) + y = np.exp(-0.107 * tc) + n0smn = np.maximum(0.5 * n0s1 * y, 1.0e6) + n0smx = np.minimum(1.0e2 * n0s1 * y, 1.0e9) + number = np.where(rho_s > 1.0e-15, np.minimum(n0smx, np.maximum(n0smn, n0s)), 8.00e5) + return dict(number=number) + + @pytest.fixture + def input_data(self, grid): + return dict( + t=data_alloc.constant_field(grid, 276.302, dims.CellDim, dims.KDim, dtype=wpfloat), + rho_s=data_alloc.constant_field( + grid, 1.17797 * 8.28451e-4, dims.CellDim, dims.KDim, dtype=wpfloat + ), + number=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/unit_tests/__init__.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/unit_tests/__init__.py new file mode 100644 index 0000000000..de9850de36 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/unit_tests/__init__.py @@ -0,0 +1,7 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause diff --git a/model/atmosphere/subgrid_scale_physics/physics_driver/README.md b/model/atmosphere/subgrid_scale_physics/physics_driver/README.md new file mode 100644 index 0000000000..bc2454fbdd --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/physics_driver/README.md @@ -0,0 +1,3 @@ +# icon4py-atmosphere-physics-driver + +The physics driver (`PhysicsDriver`) and its process registry (`PhysicsProcess`, `ProcessTimeControl`, `ForcingMode`) for ICON4Py, modelled on ICON-MPIM's AES physics interface (`aes_phy_main`) diff --git a/model/atmosphere/subgrid_scale_physics/physics_driver/pyproject.toml b/model/atmosphere/subgrid_scale_physics/physics_driver/pyproject.toml new file mode 100644 index 0000000000..837695f500 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/physics_driver/pyproject.toml @@ -0,0 +1,58 @@ +# -- Build system requirements (PEP 518) -- + +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools>=61.0", "wheel>=0.40.0"] + +# -- Standard project description options (PEP 621) -- +[project] +authors = [{email = "gridtools@cscs.ch", name = "ETH Zurich"}] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: BSD License", + "Operating System :: POSIX", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Atmospheric Science", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Physics" +] +dependencies = [ + # workspace members + "icon4py-common~=0.2.0", + # external dependencies + "packaging>=20.0" +] +description = "The physics driver (PhysicsDriver) and its process registry for ICON4Py." +license = {text = "BSD-3 License"} +name = "icon4py-atmosphere-physics-driver" +readme = "README.md" +requires-python = ">=3.10" +# managed by bump-my-version: +version = "0.2.0" + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +# --ruff -- +[tool.ruff] +extend = "../../../../pyproject.toml" + +[tool.ruff.lint.isort] +known-first-party = ['icon4py'] +known-third-party = ['gt4py'] + +# -- setuptools -- +[tool.setuptools.package-data] +'*' = ['*.in', '*.md', '*.rst', '*.txt', 'LICENSE', 'py.typed'] + +[tool.setuptools.packages] +find = {namespaces = true, where = ['src']} diff --git a/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/__init__.py b/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/__init__.py new file mode 100644 index 0000000000..0344c5052e --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/__init__.py @@ -0,0 +1,29 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from typing import Final + +from packaging import version as pkg_version + + +__all__ = [ + "__author__", + "__copyright__", + "__license__", + "__version__", + "__version_info__", +] + + +__author__: Final = "ETH Zurich, MeteoSwiss and individual contributors" +__copyright__: Final = "Copyright (c) 2022-2024 ETH Zurich and MeteoSwiss" +__license__: Final = "BSD-3-Clause" + + +__version__: Final = "0.2.0" +__version_info__: Final = pkg_version.parse(__version__) diff --git a/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/physics_driver.py b/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/physics_driver.py new file mode 100644 index 0000000000..e2f67e1d3d --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/physics_driver.py @@ -0,0 +1,117 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +"""The ``PhysicsDriver`` and its process / time-control types.""" + +from __future__ import annotations + +import dataclasses +import datetime +import enum +from typing import TYPE_CHECKING, Any + +from icon4py.model.atmosphere.subgrid_scale_physics.physics_driver.process_time_control import ( + ProcessTimeControl, +) +from icon4py.model.common.components.components import Component +from icon4py.model.common.components.physics_state import PhysicsState + + +if TYPE_CHECKING: + from icon4py.model.common.states import prognostic_state, tracer_states + + +class ForcingMode(enum.IntEnum): + """Per-process apply switch -- the icon4py analogue of AES ``fc_xxx``. + + Decides whether a process's computed forcing is fed back into the prognostic + state when the process runs: + + - APPLY: compute and apply it (``field += tend*dt``); the process affects the run. + - DIAGNOSTIC: compute it but do NOT apply it -- the outputs stay available for + inspection/output while the prognostic state is left unchanged ("look, don't touch"). + + This composes with ``kind`` tag (``states/model.py``): ForcingMode is + per-PROCESS, and ``kind`` is whether the field is a tendency or a diagnostic. + A process is applied only if APPLY mode, and within it only ``kind="tendency"`` fields + are added to the state. + """ + + DIAGNOSTIC = 0 + APPLY = 1 + + +@dataclasses.dataclass +class PhysicsProcess: + """A registered physics process: a component, its state adapter, and its time control. + + The component is the per-process adapter (e.g. ``MuphysComponent``); it + implements the generic ``Component`` protocol, which is how the driver types it. + The state adapter is process-specific (it translates the prognostic state to/from + *this* component's contract), so it is bundled per process rather than shared. + + ``forcing_mode`` is the per-process AES ``fc_xxx`` analogue (DIAGNOSTIC vs APPLY); + it lives here rather than on ``ProcessTimeControl`` because it is a property of the + process, not of its firing schedule. + """ + + name: str + component: Component + state: PhysicsState + time_control: ProcessTimeControl + forcing_mode: ForcingMode = ForcingMode.APPLY + + +class PhysicsDriver: + """The physics driver: runs each registered physics process in order.""" + + def __init__( + self, + processes: list[PhysicsProcess], + ) -> None: + self._processes = processes + self._recycle_cache: dict[str, dict[str, Any]] = {} + + def run( + self, + prognostic: prognostic_state.PrognosticState, + tracers: tracer_states.TracerState, + dtime: datetime.timedelta, + simulation_current_datetime: datetime.datetime, + ) -> None: + for proc in self._processes: # TODO (Yilu): rename to component + tc = proc.time_control + tc.validate_interval(dtime) + state = proc.state + state.gather_from_prognostic(prognostic, tracers) + if not tc.enable_process: + continue + if not tc.is_in_window(simulation_current_datetime): + # outside the process window: no forcing + continue + # Compute on a firing (active) step, and also on the first in-window step -- when + # there is nothing cached to recycle yet. Otherwise reuse the last computed forcing. + if tc.is_active(simulation_current_datetime) or proc.name not in self._recycle_cache: + # compute + outputs = proc.component(state.as_component_input(), simulation_current_datetime) + self._recycle_cache[proc.name] = outputs + else: + # recycle + outputs = self._recycle_cache[proc.name] + # TODO (Yilu): ForcingMode.DIAGNOSTIC (compute without applying) is not implemented yet: + # scatter_to_prognostic both applies tendencies and stores diagnostics, so a + # compute-only path needs that split first (to be done with the State-protocol + # formalization). Fail loud rather than silently apply for a DIAGNOSTIC process. + if proc.forcing_mode is not ForcingMode.APPLY: + raise NotImplementedError( + f"process '{proc.name}': only ForcingMode.APPLY is implemented; " + "DIAGNOSTIC requires splitting scatter_to_prognostic into " + "apply-tendencies vs store-diagnostics" + ) + + state.scatter_to_prognostic(prognostic, outputs, dtime) diff --git a/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/process_time_control.py b/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/process_time_control.py new file mode 100644 index 0000000000..c978c1ad63 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/physics_driver/src/icon4py/model/atmosphere/subgrid_scale_physics/physics_driver/process_time_control.py @@ -0,0 +1,75 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +"""Per-process time control for the physics driver.""" + +from __future__ import annotations + +import dataclasses +import datetime + + +@dataclasses.dataclass(frozen=True) +class ProcessTimeControl: + """icon4py analogue of the per-process time fields in AES `aes_phy_tc`. + + Mirrors `mo_aes_phy_main.f90` semantics, with one deviation: where AES + disables a process via `dt_xxx == 0`, we use an explicit `enable_process` + flag. + - `enable_process`: explicit on/off switch for the process. + - `interval` (`dt_xxx`): firing interval; must be > 0 when enabled. + - `start_date` (`sd_xxx`), `end_date` (`ed_xxx`): half-open + `[start, end)` window during which the process exists at all. + """ + + interval: datetime.timedelta + start_date: datetime.datetime + end_date: datetime.datetime + enable_process: bool = True + # TODO (Yilu): enbale_process should be removed from here, once we have a config for each component, then we can set the enable_process in the config file for each component. + # TODO (Yilu): rename enable_process to enable_component + + def is_in_window(self, simulation_current_datetime: datetime.datetime) -> bool: + return self.start_date <= simulation_current_datetime < self.end_date + + def is_active(self, simulation_current_datetime: datetime.datetime) -> bool: + """True if the process's mtime-event fires on the given step. + + Equivalent to AES `isCurrentEventActive(ev_xxx, datetime)`. Fires only + when the elapsed time is an exact integer multiple of the interval. + """ + if ( + not self.enable_process + or self.interval <= datetime.timedelta(0) + or simulation_current_datetime < self.start_date + ): + return False + elapsed = simulation_current_datetime - self.start_date + return elapsed % self.interval == datetime.timedelta(0) + + def validate_interval(self, dtime: datetime.timedelta) -> None: + """Fail loud if the firing interval cannot align with the model timestep. + + ``is_active`` fires only when the elapsed time is an exact multiple of + ``interval``. With a discrete timestep that is observed only when + ``interval`` is a positive integer multiple of ``dtime`` -- otherwise the + process fires only at common multiples of both (or never), silently. + """ + if not self.enable_process: + return + if self.interval <= datetime.timedelta(0): + raise ValueError( + f"time-control interval must be positive for an enabled process, " + f"got {self.interval}" + ) + if self.interval % dtime != datetime.timedelta(0): + raise ValueError( + f"time-control interval {self.interval} is not an integer multiple of " + f"the model timestep {dtime}: the process would fire only at common " + "multiples of both (or never)" + ) diff --git a/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/__init__.py b/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/__init__.py new file mode 100644 index 0000000000..de9850de36 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/__init__.py @@ -0,0 +1,7 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause diff --git a/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/unit_tests/__init__.py b/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/unit_tests/__init__.py new file mode 100644 index 0000000000..de9850de36 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/unit_tests/__init__.py @@ -0,0 +1,7 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause diff --git a/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/unit_tests/test_physics_driver.py b/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/unit_tests/test_physics_driver.py new file mode 100644 index 0000000000..0836849bb4 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/physics_driver/tests/physics_driver/unit_tests/test_physics_driver.py @@ -0,0 +1,402 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import dataclasses +import datetime + +import pytest + +from icon4py.model.atmosphere.subgrid_scale_physics.physics_driver.physics_driver import ( + ForcingMode, + PhysicsDriver, + PhysicsProcess, +) +from icon4py.model.atmosphere.subgrid_scale_physics.physics_driver.process_time_control import ( + ProcessTimeControl, +) +from icon4py.model.common.components.physics_state import PhysicsState +from icon4py.model.common.states.model import FieldMetaData + + +def test_field_metadata_accepts_kind() -> None: + meta: FieldMetaData = { + "standard_name": "tend_temperature", + "units": "K s-1", + "kind": "tendency", + } + assert meta["kind"] == "tendency" + + +def test_forcing_mode_values() -> None: + assert ForcingMode.DIAGNOSTIC.value == 0 + assert ForcingMode.APPLY.value == 1 + assert ForcingMode.DIAGNOSTIC is not ForcingMode.APPLY + assert len(ForcingMode) == 2 + + +_T0 = datetime.datetime(2024, 1, 1, 0, 0, 0) +_DT = datetime.timedelta(seconds=300) # 5-min physics interval + + +def _tc( + interval: datetime.timedelta = _DT, + start: datetime.datetime = _T0, + end: datetime.datetime = _T0 + datetime.timedelta(days=1), + enable_process: bool = True, +) -> ProcessTimeControl: + return ProcessTimeControl( + interval=interval, + start_date=start, + end_date=end, + enable_process=enable_process, + ) + + +class TestProcessTimeControl: + def test_enable_process_defaults_true(self) -> None: + assert _tc().enable_process is True + + def test_is_active_false_when_disabled(self) -> None: + assert _tc(enable_process=False).is_active(_T0) is False + + def test_is_active_false_when_interval_zero(self) -> None: + assert _tc(interval=datetime.timedelta(0)).is_active(_T0) is False + + def test_is_in_window_at_start_is_true(self) -> None: + assert _tc().is_in_window(_T0) is True + + def test_is_in_window_at_end_is_false(self) -> None: + end = _T0 + datetime.timedelta(hours=1) + assert _tc(end=end).is_in_window(end) is False + + def test_is_in_window_before_start_is_false(self) -> None: + assert _tc().is_in_window(_T0 - datetime.timedelta(seconds=1)) is False + + def test_is_in_window_inside_is_true(self) -> None: + assert _tc().is_in_window(_T0 + datetime.timedelta(hours=12)) is True + + def test_is_active_at_start_is_true(self) -> None: + assert _tc().is_active(_T0) is True + + def test_is_active_at_one_interval_is_true(self) -> None: + assert _tc().is_active(_T0 + _DT) is True + + def test_is_active_at_half_interval_is_false(self) -> None: + assert _tc().is_active(_T0 + _DT / 2) is False + + def test_is_active_before_start_is_false(self) -> None: + assert _tc().is_active(_T0 - datetime.timedelta(seconds=1)) is False + + def test_is_active_requires_exact_interval_multiple(self) -> None: + # Fires only at an exact integer multiple of the interval. + assert _tc().is_active(_T0 + 2 * _DT) is True + # 1 microsecond off the boundary does not fire (no tolerance). + jitter = datetime.timedelta(microseconds=1) + assert _tc().is_active(_T0 + 2 * _DT + jitter) is False + + def test_frozen_dataclass(self) -> None: + tc = _tc() + with pytest.raises(dataclasses.FrozenInstanceError): + tc.interval = datetime.timedelta(seconds=1) # type: ignore[misc] + + def test_validate_interval_accepts_integer_multiple(self) -> None: + _tc(interval=2 * _DT).validate_interval(_DT) + + def test_validate_interval_rejects_non_multiple(self) -> None: + with pytest.raises(ValueError, match="integer multiple"): + _tc(interval=1.5 * _DT).validate_interval(_DT) + + def test_validate_interval_rejects_zero_interval_when_enabled(self) -> None: + with pytest.raises(ValueError, match="positive"): + _tc(interval=datetime.timedelta(0)).validate_interval(_DT) + + def test_validate_interval_skips_disabled_process(self) -> None: + _tc(interval=1.5 * _DT, enable_process=False).validate_interval(_DT) + + +def test_physics_process_construction() -> None: + class _DummyComponent: + inputs_properties = {} + outputs_properties = {} + + def __call__(self, state, time_step): + return {} + + state = RecordingPhysicsState() + proc = PhysicsProcess( + name="muphys", + component=_DummyComponent(), + state=state, + time_control=_tc(), + ) + assert proc.name == "muphys" + assert proc.component is not None + assert proc.state is state + assert proc.time_control.enable_process + assert proc.forcing_mode is ForcingMode.APPLY + + +@dataclasses.dataclass +class RecordingComponent: + """Stub Component: records calls, returns configured outputs. + + `output_kinds` keys mirror `outputs` keys; values are 'tendency' or + 'diagnostic'. + """ + + outputs: dict[str, object] + output_kinds: dict[str, str] + call_count: int = 0 + last_state: dict | None = None + last_time: datetime.datetime | None = None + + @property + def inputs_properties(self) -> dict: + return {} + + @property + def outputs_properties(self) -> dict: + return { + k: {"standard_name": k, "units": "1", "kind": self.output_kinds[k]} + for k in self.outputs + } + + def __call__(self, state, time_step): + self.call_count += 1 + self.last_state = state + self.last_time = time_step + return dict(self.outputs) + + +@dataclasses.dataclass +class RecordingPhysicsState(PhysicsState): + """Stub PhysicsState: records refresh / scatter; returns a fixed dict + from as_component_input. Implements just enough surface for the PhysicsDriver.""" + + gather_calls: list = dataclasses.field(default_factory=list) + scatter_calls: list = dataclasses.field(default_factory=list) + + def gather_from_prognostic(self, prognostic, tracers) -> None: + self.gather_calls.append(prognostic) + + def as_component_input(self) -> dict: + return {"foo": "bar"} + + def scatter_to_prognostic(self, prognostic, outputs, dtime) -> None: + self.scatter_calls.append((prognostic, outputs, dtime)) + + +def test_recording_doubles_record_calls() -> None: + component = RecordingComponent( + outputs={"tend_temperature": "T_TEND_VALUE", "pflx": "PFLX_VALUE"}, + output_kinds={"tend_temperature": "tendency", "pflx": "diagnostic"}, + ) + state = RecordingPhysicsState() + + # Simulate what PhysicsDriver would do. + state.gather_from_prognostic("prog", "tracers") + out = component(state.as_component_input(), _T0) + state.scatter_to_prognostic("prog", out, datetime.timedelta(seconds=300)) + + assert state.gather_calls == ["prog"] + assert component.call_count == 1 + assert component.last_state == {"foo": "bar"} # what as_component_input returned + assert state.scatter_calls == [("prog", out, datetime.timedelta(seconds=300))] + + +def test_run_invokes_components_in_order() -> None: + state = RecordingPhysicsState() + comp_a = RecordingComponent( + outputs={"tend_temperature": "A"}, + output_kinds={"tend_temperature": "tendency"}, + ) + comp_b = RecordingComponent( + outputs={"tend_temperature": "B"}, + output_kinds={"tend_temperature": "tendency"}, + ) + + driver = PhysicsDriver( + processes=[ + PhysicsProcess(name="A", component=comp_a, state=state, time_control=_tc()), + PhysicsProcess(name="B", component=comp_b, state=state, time_control=_tc()), + ], + ) + + driver.run( + prognostic="prog", + tracers="tracers", + dtime=datetime.timedelta(seconds=300), + simulation_current_datetime=_T0, + ) + + assert comp_a.call_count == 1 + assert comp_b.call_count == 1 + # B's scatter must follow A's (operator-splitting ordering) + assert state.scatter_calls[0][1] == {"tend_temperature": "A"} + assert state.scatter_calls[1][1] == {"tend_temperature": "B"} + + +def test_run_raises_for_non_multiple_interval() -> None: + state = RecordingPhysicsState() + comp = RecordingComponent( + outputs={"tend_temperature": "X"}, + output_kinds={"tend_temperature": "tendency"}, + ) + driver = PhysicsDriver( + processes=[ + PhysicsProcess( + name="X", component=comp, state=state, time_control=_tc(interval=1.5 * _DT) + ), + ], + ) + + with pytest.raises(ValueError, match="integer multiple"): + driver.run( + prognostic="prog", + tracers="tracers", + dtime=_DT, + simulation_current_datetime=_T0, + ) + assert comp.call_count == 0 + + +def test_disabled_process_is_skipped() -> None: + state = RecordingPhysicsState() + comp = RecordingComponent( + outputs={"tend_temperature": "X"}, + output_kinds={"tend_temperature": "tendency"}, + ) + tc_disabled = _tc(enable_process=False) + + driver = PhysicsDriver( + processes=[ + PhysicsProcess(name="disabled", component=comp, state=state, time_control=tc_disabled) + ], + ) + + driver.run( + prognostic="prog", + tracers="tracers", + dtime=datetime.timedelta(seconds=300), + simulation_current_datetime=_T0, + ) + + assert comp.call_count == 0 + assert state.scatter_calls == [] + + +def test_out_of_window_process_does_nothing() -> None: + state = RecordingPhysicsState() + comp = RecordingComponent( + outputs={"tend_temperature": "X"}, + output_kinds={"tend_temperature": "tendency"}, + ) + # Window starts in the future — `simulation_current_datetime=_T0` is before it. + future = _T0 + datetime.timedelta(days=1) + tc = _tc(start=future, end=future + datetime.timedelta(hours=1)) + + driver = PhysicsDriver( + processes=[PhysicsProcess(name="future", component=comp, state=state, time_control=tc)], + ) + + driver.run( + prognostic="prog", + tracers="tracers", + dtime=datetime.timedelta(seconds=300), + simulation_current_datetime=_T0, + ) + + assert comp.call_count == 0 + assert state.scatter_calls == [] + + +def test_active_call_caches_outputs_and_applies_them() -> None: + state = RecordingPhysicsState() + comp = RecordingComponent( + outputs={"tend_temperature": "FRESH"}, + output_kinds={"tend_temperature": "tendency"}, + ) + driver = PhysicsDriver( + processes=[PhysicsProcess(name="p", component=comp, state=state, time_control=_tc())], + ) + + driver.run( + prognostic="prog", + tracers="tracers", + dtime=datetime.timedelta(seconds=300), + simulation_current_datetime=_T0, + ) + + assert comp.call_count == 1 + assert state.scatter_calls == [ + ("prog", {"tend_temperature": "FRESH"}, datetime.timedelta(seconds=300)) + ] + + +def test_inactive_in_window_recycles_cached_outputs() -> None: + state = RecordingPhysicsState() + # Component returns "FRESH" the first time, would return "STALE" the second + # if called — but on the recycle step it MUST NOT be called. + comp = RecordingComponent( + outputs={"tend_temperature": "FRESH"}, + output_kinds={"tend_temperature": "tendency"}, + ) + # interval = 2 * dt → process fires every other call. + interval = 2 * _DT + tc = _tc(interval=interval) + driver = PhysicsDriver( + processes=[PhysicsProcess(name="p", component=comp, state=state, time_control=tc)], + ) + + # Step 1: active (elapsed == 0), compute + cache. + driver.run( + prognostic="prog", + tracers="tracers", + dtime=_DT, + simulation_current_datetime=_T0, + ) + # Step 2: in window, but not active (elapsed == _DT, not a multiple of 2*_DT). + driver.run( + prognostic="prog", + tracers="tracers", + dtime=_DT, + simulation_current_datetime=_T0 + _DT, + ) + + # Component invoked once total (compute step only). + assert comp.call_count == 1 + # But scatter happened twice — once with the fresh tendency, once recycled. + assert len(state.scatter_calls) == 2 + assert state.scatter_calls[0][1] == {"tend_temperature": "FRESH"} + assert state.scatter_calls[1][1] == {"tend_temperature": "FRESH"} # recycled + + +def test_first_in_window_step_inactive_computes_without_keyerror() -> None: + # Regression (jcanton review): a process whose first-ever in-window step is NOT active + # (interval = 2*dt, first step lands at start + dt) used to KeyError on the empty recycle + # cache. With nothing cached to recycle yet, it must compute instead. + state = RecordingPhysicsState() + comp = RecordingComponent( + outputs={"tend_temperature": "FRESH"}, + output_kinds={"tend_temperature": "tendency"}, + ) + tc = _tc(interval=2 * _DT) # start = _T0 + driver = PhysicsDriver( + processes=[PhysicsProcess(name="p", component=comp, state=state, time_control=tc)], + ) + + # First call lands in-window but off the firing tick (elapsed == _DT, interval == 2*_DT). + driver.run( + prognostic="prog", + tracers="tracers", + dtime=_DT, + simulation_current_datetime=_T0 + _DT, + ) + + assert comp.call_count == 1 + assert state.scatter_calls == [("prog", {"tend_temperature": "FRESH"}, _DT)] diff --git a/model/common/src/icon4py/model/common/components/physics_state.py b/model/common/src/icon4py/model/common/components/physics_state.py new file mode 100644 index 0000000000..5ae2222cf4 --- /dev/null +++ b/model/common/src/icon4py/model/common/components/physics_state.py @@ -0,0 +1,41 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol + + +if TYPE_CHECKING: + import datetime + + from icon4py.model.common.states import prognostic_state, tracer_states + + +class PhysicsState(Protocol): + """ + Protocol for the physics-state adapter that the ``PhysicsDriver`` depends on. + + The concrete implementations live with their process (e.g. + ``icon4py.model.atmosphere.subgrid_scale_physics.muphys.state.State``); this + module only declares the interface the ``PhysicsDriver`` relies on, so it + stays decoupled from any specific physics state. + """ + + def gather_from_prognostic( + self, + prognostic: prognostic_state.PrognosticState, + tracers: tracer_states.TracerState, + ) -> None: ... + def as_component_input(self) -> dict[str, Any]: ... + def scatter_to_prognostic( + self, + prognostic: prognostic_state.PrognosticState, + outputs: dict[str, Any], + dtime: datetime.timedelta, + ) -> None: ... diff --git a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/update_exner_and_theta_v.py b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/update_exner_and_theta_v.py new file mode 100644 index 0000000000..a2485a0b13 --- /dev/null +++ b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/update_exner_and_theta_v.py @@ -0,0 +1,75 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx +from gt4py.next import exp, log + +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common.constants import PhysicsConstants +from icon4py.model.common.type_alias import wpfloat + + +@gtx.field_operator +def _update_exner_and_theta_v( + rho: fa.CellKField[ta.wpfloat], + virtual_temperature: fa.CellKField[ta.wpfloat], + virtual_temperature_tendency: fa.CellKField[ta.wpfloat], + dtime: ta.wpfloat, +) -> tuple[fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat]]: + """Update exner and theta_v from a physics virtual-temperature tendency. + + Mirrors ICON's physics-to-dynamics thermodynamic update in + ``mo_interface_iconam_aes.f90``: + recompute exner from the new virtual temperature via the exact equation of + state and diagnose ``theta_v = Tv / exner``, so the exner/rho/theta_v trio + stays EOS-consistent:: + + Tv_new = Tv + dtime * dTv / dt + exner_new = (rd / p0ref * rho * Tv_new) ** (rd / cpd) + theta_v = Tv_new / exner_new + + Args: + rho: air density [kg m-3] + virtual_temperature: virtual temperature before the physics update [K] + virtual_temperature_tendency: physics virtual-temperature tendency [K s-1] + dtime: time step [s] + Returns: + (new exner function, new virtual potential temperature theta_v [K]) + """ + new_virtual_temperature = virtual_temperature + virtual_temperature_tendency * dtime + new_exner = exp( + PhysicsConstants.rd_o_cpd * log(PhysicsConstants.rd_o_p0ref * rho * new_virtual_temperature) + ) + new_theta_v = new_virtual_temperature / new_exner + return new_exner, new_theta_v + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def update_exner_and_theta_v( + rho: fa.CellKField[ta.wpfloat], + virtual_temperature: fa.CellKField[ta.wpfloat], + virtual_temperature_tendency: fa.CellKField[ta.wpfloat], + dtime: ta.wpfloat, + exner: fa.CellKField[wpfloat], + theta_v: fa.CellKField[wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +) -> None: + _update_exner_and_theta_v( + rho=rho, + virtual_temperature=virtual_temperature, + virtual_temperature_tendency=virtual_temperature_tendency, + dtime=dtime, + out=(exner, theta_v), + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/src/icon4py/model/common/math/operators.py b/model/common/src/icon4py/model/common/math/operators.py index 761cae99d0..b4daf883d9 100644 --- a/model/common/src/icon4py/model/common/math/operators.py +++ b/model/common/src/icon4py/model/common/math/operators.py @@ -48,3 +48,17 @@ def _compute_difference_on_cell_k( field_b: fa.CellKField[ta.vpfloat], ) -> fa.CellKField[ta.wpfloat]: return field_a - astype(field_b, wpfloat) + + +@gtx.field_operator +def _compute_field_a_plus_coeff_times_field_b_on_cell_k( + field_a: fa.CellKField[ta.wpfloat], + coeff: ta.wpfloat, + field_b: fa.CellKField[ta.wpfloat], +) -> fa.CellKField[ta.wpfloat]: + return field_a + coeff * field_b + + +@gtx.field_operator +def _copy_field_on_cell_k(field: fa.CellKField[ta.wpfloat]) -> fa.CellKField[ta.wpfloat]: + return field diff --git a/model/common/src/icon4py/model/common/math/stencils/generic_math_operations.py b/model/common/src/icon4py/model/common/math/stencils/generic_math_operations.py index 6c2171e041..b75bfb6a85 100644 --- a/model/common/src/icon4py/model/common/math/stencils/generic_math_operations.py +++ b/model/common/src/icon4py/model/common/math/stencils/generic_math_operations.py @@ -9,7 +9,11 @@ import gt4py.next as gtx from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta -from icon4py.model.common.math.operators import _compute_difference_on_cell_k +from icon4py.model.common.math.operators import ( + _compute_difference_on_cell_k, + _compute_field_a_plus_coeff_times_field_b_on_cell_k, + _copy_field_on_cell_k, +) @gtx.program @@ -31,3 +35,45 @@ def compute_difference_on_cell_k( dims.KDim: (vertical_start, vertical_end), }, ) + + +@gtx.program +def compute_field_a_plus_coeff_times_field_b_on_cell_k( + field_a: fa.CellKField[ta.wpfloat], + coeff: ta.wpfloat, + field_b: fa.CellKField[ta.wpfloat], + output_field: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +) -> None: + _compute_field_a_plus_coeff_times_field_b_on_cell_k( + field_a, + coeff, + field_b, + out=output_field, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) + + +@gtx.program +def copy_field_on_cell_k( + field: fa.CellKField[ta.wpfloat], + output_field: fa.CellKField[ta.wpfloat], + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +) -> None: + _copy_field_on_cell_k( + field, + out=output_field, + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/common/src/icon4py/model/common/states/data.py b/model/common/src/icon4py/model/common/states/data.py index 03a84399af..2eab334a1e 100644 --- a/model/common/src/icon4py/model/common/states/data.py +++ b/model/common/src/icon4py/model/common/states/data.py @@ -144,3 +144,74 @@ icon_var_name="pres_sfc", ), ) + +# CF attributes of precipitation-flux diagnostics. +PRECIPITATION_CF_ATTRIBUTES: Final[dict[str, model.FieldMetaData]] = dict( + precipitation_flux=dict( + standard_name="precipitation_flux", + long_name="total precipitation flux (rain + snow + graupel + ice)", + units="kg m-2 s-1", + kind="diagnostic", + ), + rainfall_flux=dict( + standard_name="rainfall_flux", + long_name="rainfall flux", + units="kg m-2 s-1", + kind="diagnostic", + ), + snowfall_flux=dict( + standard_name="snowfall_flux", + long_name="snowfall flux", + units="kg m-2 s-1", + kind="diagnostic", + ), + graupelfall_flux=dict( + standard_name="graupelfall_flux", + long_name="graupelfall flux", + units="kg m-2 s-1", + kind="diagnostic", + ), + icefall_flux=dict( + standard_name="icefall_flux", + long_name="icefall flux", + units="kg m-2 s-1", + kind="diagnostic", + ), + precipitation_energy_flux=dict( + standard_name="precipitation_energy_flux", + long_name="energy flux carried by precipitation", + units="W m-2", + kind="diagnostic", + ), +) + + +def tendency_of(base: model.FieldMetaData) -> model.FieldMetaData: + """Derive generic tendency metadata for ``base`` (CF ``tendency_of_``). + + The tendency carries ``kind="tendency"`` and ``base``'s units per second; a + dimensionless base (``units="1"``) yields ``"s-1"`` rather than ``"1 s-1"``. + """ + base_units = base["units"] + units = "s-1" if base_units == "1" else f"{base_units} s-1" + tendency: model.FieldMetaData = dict( + standard_name=f"tendency_of_{base['standard_name']}", + units=units, + kind="tendency", + ) + long_name = base.get("long_name") + if long_name is not None: + tendency["long_name"] = f"tendency of {long_name}" + return tendency + + +# Generic tendencies of the temperature and tracer fields, derived from their base identities so names/units never drift. +TENDENCY_CF_ATTRIBUTES: Final[dict[str, model.FieldMetaData]] = dict( + temperature=tendency_of(DIAGNOSTIC_CF_ATTRIBUTES["temperature"]), + qv=tendency_of(COMMON_TRACER_CF_ATTRIBUTES["qv"]), + qc=tendency_of(COMMON_TRACER_CF_ATTRIBUTES["qc"]), + qi=tendency_of(COMMON_TRACER_CF_ATTRIBUTES["qi"]), + qr=tendency_of(COMMON_TRACER_CF_ATTRIBUTES["qr"]), + qs=tendency_of(COMMON_TRACER_CF_ATTRIBUTES["qs"]), + qg=tendency_of(COMMON_TRACER_CF_ATTRIBUTES["qg"]), +) diff --git a/model/common/src/icon4py/model/common/states/model.py b/model/common/src/icon4py/model/common/states/model.py index c4a01f4f7e..e23e9524e9 100644 --- a/model/common/src/icon4py/model/common/states/model.py +++ b/model/common/src/icon4py/model/common/states/model.py @@ -37,6 +37,8 @@ class OptionalMetaData(TypedDict, total=False): #: whether the vertical dimension of the field lives on interface (half) levels #: rather than full levels is_on_half_levels: bool + #: physics-component output category: "tendency" applied as field += val*dt; "diagnostic" stored unscaled + kind: Literal["tendency", "diagnostic"] class RequiredMetaData(TypedDict, total=True): diff --git a/model/common/src/icon4py/model/common/states/prognostic_state.py b/model/common/src/icon4py/model/common/states/prognostic_state.py index a5d8d1043d..a515286cec 100644 --- a/model/common/src/icon4py/model/common/states/prognostic_state.py +++ b/model/common/src/icon4py/model/common/states/prognostic_state.py @@ -39,6 +39,7 @@ class PrognosticState: tracer: TracerState = dataclasses.field( default_factory=TracerState ) # tracer concentration, one CellKField per active tracer [kg/kg] + # TODO (Anyone): tracer should not inside the PrognosticState, since it cannot be swapped after every sub-timestep, it only swaps after 1 timestep @property def w_1(self) -> fa.CellField[ta.wpfloat]: diff --git a/model/common/tests/common/components/__init__.py b/model/common/tests/common/components/__init__.py new file mode 100644 index 0000000000..de9850de36 --- /dev/null +++ b/model/common/tests/common/components/__init__.py @@ -0,0 +1,7 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause diff --git a/model/common/tests/common/components/unit_tests/__init__.py b/model/common/tests/common/components/unit_tests/__init__.py new file mode 100644 index 0000000000..de9850de36 --- /dev/null +++ b/model/common/tests/common/components/unit_tests/__init__.py @@ -0,0 +1,7 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause diff --git a/model/common/tests/common/diagnostic_calculations/unit_tests/test_diagnostic_calculations.py b/model/common/tests/common/diagnostic_calculations/unit_tests/test_diagnostic_calculations.py index 830bab3273..9f3ee7f575 100644 --- a/model/common/tests/common/diagnostic_calculations/unit_tests/test_diagnostic_calculations.py +++ b/model/common/tests/common/diagnostic_calculations/unit_tests/test_diagnostic_calculations.py @@ -10,17 +10,20 @@ from typing import TYPE_CHECKING import gt4py.next as gtx +import numpy as np import pytest import icon4py.model.common.grid.horizontal as h_grid from icon4py.model.common import dimension as dims +from icon4py.model.common.constants import PhysicsConstants from icon4py.model.common.diagnostic_calculations.stencils import ( calculate_tendency, diagnose_pressure, diagnose_surface_pressure, diagnose_temperature, + update_exner_and_theta_v, ) -from icon4py.model.common.grid import vertical as v_grid +from icon4py.model.common.grid import simple, vertical as v_grid from icon4py.model.common.interpolation.stencils import edge_2_cell_vector_rbf_interpolation as rbf from icon4py.model.common.states import diagnostic_state as diagnostics, tracer_states as tracers from icon4py.model.common.utils import data_allocation as data_alloc @@ -46,6 +49,60 @@ from icon4py.model.testing import serialbox as sb +def test_update_exner_and_theta_v(backend: gtx_typing.Backend) -> None: + """The physics coupling recomputes exner via the exact EOS and theta_v = Tv/exner. + + Mirrors ICON's mo_interface_iconam_aes.f90: + Tv_new = Tv + dt * dTv/dt + exner_new = (rd/p0ref * rho * Tv_new) ** (rd/cpd) + theta_v = Tv_new / exner_new + """ + grid = simple.simple_grid() + rho_value, tv_value, tv_tendency_value, dtime = 1.1, 280.0, 0.05, 20.0 + + rho = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=float, allocator=backend) + virtual_temperature = data_alloc.zero_field( + grid, dims.CellDim, dims.KDim, dtype=float, allocator=backend + ) + virtual_temperature_tendency = data_alloc.zero_field( + grid, dims.CellDim, dims.KDim, dtype=float, allocator=backend + ) + exner = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=float, allocator=backend) + theta_v = data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=float, allocator=backend) + rho.ndarray[...] = rho_value + virtual_temperature.ndarray[...] = tv_value + virtual_temperature_tendency.ndarray[...] = tv_tendency_value + + update_exner_and_theta_v.update_exner_and_theta_v.with_backend(backend)( + rho=rho, + virtual_temperature=virtual_temperature, + virtual_temperature_tendency=virtual_temperature_tendency, + dtime=dtime, + exner=exner, + theta_v=theta_v, + horizontal_start=0, + horizontal_end=grid.num_cells, + vertical_start=0, + vertical_end=grid.num_levels, + offset_provider={}, + ) + + new_virtual_temperature = tv_value + tv_tendency_value * dtime + expected_exner = np.exp( + PhysicsConstants.rd_o_cpd + * np.log(PhysicsConstants.rd_o_p0ref * rho_value * new_virtual_temperature) + ) + expected_theta_v = new_virtual_temperature / expected_exner + + assert test_utils.dallclose(exner.asnumpy(), expected_exner) + assert test_utils.dallclose(theta_v.asnumpy(), expected_theta_v) + # EOS consistency: by definition theta_v * exner == Tv_new + assert test_utils.dallclose( + theta_v.asnumpy() * exner.asnumpy(), + np.full_like(theta_v.asnumpy(), new_virtual_temperature), + ) + + @pytest.mark.datatest @pytest.mark.parametrize("experiment_description", [test_defs.Experiments.JW]) def test_diagnose_temperature( diff --git a/model/standalone_driver/pyproject.toml b/model/standalone_driver/pyproject.toml index 4ebf5f2a14..0c0e7d8c15 100644 --- a/model/standalone_driver/pyproject.toml +++ b/model/standalone_driver/pyproject.toml @@ -27,6 +27,8 @@ dependencies = [ # workspace members "icon4py-atmosphere-dycore~=0.3.0", "icon4py-atmosphere-diffusion~=0.3.0", + "icon4py-atmosphere-physics-driver~=0.3.0", + "icon4py-atmosphere-muphys~=0.3.0", "icon4py-common[io]~=0.3.0", # external dependencies "typer>=0.20.0", diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/config.py b/model/standalone_driver/src/icon4py/model/standalone_driver/config.py index e36e0f9d3f..10931a6426 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/config.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/config.py @@ -24,6 +24,7 @@ from icon4py.model.atmosphere.subgrid_scale_physics.microphysics import ( single_moment_six_class_gscp_graupel as graupel, ) +from icon4py.model.atmosphere.subgrid_scale_physics.muphys import config as muphys_config from icon4py.model.atmosphere.tracer_advection import tracer_advection from icon4py.model.common import ( initial_condition, @@ -324,6 +325,7 @@ class ExperimentConfig: tracer_config: tracer_states.TracerConfig | None = None tracer_advection: tracer_advection.AdvectionConfig | None = None graupel: graupel.SingleMomentSixClassIconGraupelConfig | None = None + muphys: muphys_config.MuphysConfig | None = None def __post_init__(self) -> None: # The file-based initial condition needs the clock of the driver to know which @@ -395,9 +397,19 @@ def read_experiment_config_from_fortran( do_tracer_advection = not ( "exclaim_ch_r04b09_dsl" in config_file_path.name or "exclaim_ape_R02B04" in config_file_path.name + or "exclaim_ape_aesPhys" in config_file_path.name ) - # The experiments above were run in fortran with a tracer advection scheme - # that has not been ported to ICON4Py and can not be used for testing. + # These experiments' references run with tracer transport ON, but the standalone + # driver can't run tracer advection yet: it never computes airmass (rho*ddqz) or + # supplies live mass fluxes for the advection inputs (open TODO(OngChia) "compute + # airmass" in the driver's dynamics substepping), so advection divides tracer + # density by a zero airmass -> 0/0 -> NaN. The MIURA/PPM scheme itself IS ported + + # tested; the gap is the driver's advection setup. We therefore disable it here. + # For exclaim_ape_aesPhys this makes the muphys datatest validate muphys in + # isolation -- a known transport-off mismatch vs the reference (see the + # test_standalone_driver docstring). Its moisture comes from the analytical JW IC + # (normalize_global_moisture), independent of ntracer/advection, so disabling + # advection leaves muphys' tracers intact. # TODO (jcanton): this isn't the right place to keep a special case # handling. Either fix these experiments or move the special case handling. tracer_advection_cfg = ( @@ -408,7 +420,15 @@ def read_experiment_config_from_fortran( ntracer = ( fortran_config.list_to_value(atm_dict["run_nml"]["ntracer"]) if do_tracer_advection else 0 ) - tracer_cfg = tracer_states.TracerConfig.from_ntracer(ntracer) + # AES physics implies muphys is active for the experiments we support today; the presence + # of the aes_phy_nml namelist mirrors the graupel `do_physics` check below. A robust + # dt_mig>0 check needs the raw namelist (see docs/2026-07-22-muphys-namelist-dt-mig-gate.md). + aes_physics_on = "aes_phy_nml" in atm_dict + tracer_cfg = ( + tracer_states.TracerConfig.all() + if aes_physics_on + else tracer_states.TracerConfig.from_ntracer(ntracer) + ) do_physics = "nwp_phy_nml" in atm_dict and "nwp_tuning_nml" in atm_dict # If these two namelists are missing it means that the experiment was run @@ -447,6 +467,8 @@ def read_experiment_config_from_fortran( config=dataclasses.replace(initial_condition_cfg.config, ntracer=0), ) + muphys_cfg = muphys_config.MuphysConfig() if aes_physics_on else None + return ExperimentConfig( geometry=geometry_cfg, metrics=metrics_cfg, @@ -458,6 +480,7 @@ def read_experiment_config_from_fortran( tracer_config=tracer_cfg, tracer_advection=tracer_advection_cfg, graupel=graupel_cfg, + muphys=muphys_cfg, initial_condition=initial_condition_cfg, prescribed_tendencies=prescribed_tendencies.PrescribedTendenciesConfig.from_fortran_dict( atm_dict=atm_dict, data_path=config_file_path diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py index ddaf9eca71..28f9652e48 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py @@ -20,6 +20,11 @@ from icon4py.model.atmosphere.diffusion import diffusion, diffusion_states from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as solve_nh +from icon4py.model.atmosphere.subgrid_scale_physics.muphys import ( + component as muphys_component, + state as muphys_state, +) +from icon4py.model.atmosphere.subgrid_scale_physics.physics_driver import physics_driver from icon4py.model.atmosphere.tracer_advection import tracer_advection, tracer_advection_states from icon4py.model.common import ( constants, @@ -68,6 +73,7 @@ class Granules: diffusion: diffusion.Diffusion | None = None solve_nonhydro: solve_nh.SolveNonhydro | None = None tracer_advection: tracer_advection.Advection | None = None + physics: physics_driver.PhysicsDriver | None = None def validate_granule_state_consistency( @@ -78,7 +84,7 @@ def validate_granule_state_consistency( """ Validate that enabled granules have their required states allocated. The graupel granule is currently not checked as it will be moved to the - physics interface. + physics driver. Raises: ValueError: if a granule is enabled but a state it requires is None. @@ -217,6 +223,7 @@ def initialize_granules( grid: icon_grid.IconGrid, vertical_grid: v_grid.VerticalGrid, static_field_factories: static_fields.StaticFieldFactories, + model_time_variables: driver_states.ModelTimeVariables, exchange: decomposition_defs.ExchangeRuntime, owner_mask: fa.CellField[bool], backend: gtx_typing.Backend | None, @@ -438,10 +445,32 @@ def initialize_granules( exchange=exchange, ) + physics_granule: physics_driver.PhysicsDriver | None = None + if config.muphys is not None: + muphys_process = physics_driver.PhysicsProcess( + name="muphys", + component=muphys_component.MuphysComponent( + grid=grid, + dtime=config.driver.dtime, + qnc=config.muphys.qnc, + backend=backend, + scheme=config.muphys.scheme, + ), + state=muphys_state.State(grid=grid, metrics=metrics_field_source, backend=backend), + time_control=physics_driver.ProcessTimeControl( + interval=config.driver.dtime, + start_date=config.driver.start_of_simulation, + end_date=model_time_variables.simulation_end_datetime, + enable_process=True, + ), + ) + physics_granule = physics_driver.PhysicsDriver([muphys_process]) + return Granules( solve_nonhydro=solve_nonhydro_granule, diffusion=diffusion_granule, tracer_advection=tracer_advection_granule, + physics=physics_granule, ) diff --git a/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py b/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py index aae0ab0511..20b3d6e02a 100644 --- a/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py +++ b/model/standalone_driver/src/icon4py/model/standalone_driver/standalone_driver.py @@ -70,6 +70,7 @@ def __init__( decomposition_info: decomposition_defs.DecompositionInfo, static_field_factories: static_fields.StaticFieldFactories, granules: driver_utils.Granules, + model_time_variables: driver_states.ModelTimeVariables, vertical_grid_config: v_grid.VerticalGridConfig, exchange: decomposition_defs.ExchangeRuntime, global_reductions: decomposition_defs.Reductions, @@ -84,7 +85,7 @@ def __init__( self.static_field_factories = static_field_factories self.granules = granules self.vertical_grid_config = vertical_grid_config - self.model_time_variables = driver_states.ModelTimeVariables(config=config.driver) + self.model_time_variables = model_time_variables self.timer_collection = driver_states.TimerCollection( [timer.value for timer in driver_states.DriverTimers] ) @@ -292,6 +293,14 @@ def _integrate_one_time_step( dtime=self.model_time_variables.dtime_in_seconds, ) + if self.granules.physics is not None: + self.granules.physics.run( + prognostic=prognostic_states.next, + tracers=prognostic_states.next.tracer, + dtime=self.config.driver.dtime, + simulation_current_datetime=self.model_time_variables.simulation_current_datetime, + ) + prognostic_states.swap() def _update_time_levels_for_velocity_tendencies( @@ -666,12 +675,15 @@ def initialize_driver( metrics_config=config.metrics, ) + model_time_variables = driver_states.ModelTimeVariables(config=config.driver) + log.info("initializing granules") granules = driver_utils.initialize_granules( config=config, grid=grid_manager.grid, vertical_grid=vertical_grid, static_field_factories=static_field_factories, + model_time_variables=model_time_variables, exchange=exchange, owner_mask=gtx.as_field( (dims.CellDim,), @@ -704,6 +716,7 @@ def initialize_driver( decomposition_info=decomposition_info, static_field_factories=static_field_factories, granules=granules, + model_time_variables=model_time_variables, vertical_grid_config=config.vertical_grid, exchange=exchange, global_reductions=global_reductions, diff --git a/model/standalone_driver/tests/standalone_driver/fixtures.py b/model/standalone_driver/tests/standalone_driver/fixtures.py index 2129ee43d3..dd89cad25f 100644 --- a/model/standalone_driver/tests/standalone_driver/fixtures.py +++ b/model/standalone_driver/tests/standalone_driver/fixtures.py @@ -24,6 +24,7 @@ savepoint_nonhydro_exit, savepoint_nonhydro_init, savepoint_nonhydro_step_final, + savepoint_time_step_exit, savepoint_velocity_init, step_date_exit, step_date_init, diff --git a/model/standalone_driver/tests/standalone_driver/integration_tests/test_standalone_driver.py b/model/standalone_driver/tests/standalone_driver/integration_tests/test_standalone_driver.py index a517992b80..2815fbc182 100644 --- a/model/standalone_driver/tests/standalone_driver/integration_tests/test_standalone_driver.py +++ b/model/standalone_driver/tests/standalone_driver/integration_tests/test_standalone_driver.py @@ -27,7 +27,6 @@ # Tolerances (atol, rtol) per experiment, measured across the CSCS CI backends -# (gtfn_cpu, gtfn_gpu, dace_cpu, dace_gpu). _TOLERANCES: dict[test_defs.ExperimentDescription, dict[str, tuple[float, float]]] = { test_defs.Experiments.JW: { "vn": (5.3e-7, 0.0), @@ -50,53 +49,75 @@ "theta_v": (1.2e-3, 3.6e-6), "rho": (3.5e-6, 3.7e-6), }, + test_defs.Experiments.EXCLAIM_APE_AES: { + "vn": (6e-7, 0.0), + "w": (1e-8, 0.0), + "rho": (9e-10, 0.0), + "exner": (1e-8, 0.0), + "theta_v": (0.0, 3e-8), + "qv": (2e-6, 0.0), + "qc": (1e-10, 0.0), + "qr": (1e-10, 0.0), + "qs": (1e-10, 0.0), + "qi": (1e-10, 0.0), + "qg": (1e-10, 0.0), + }, } +# Metadata selecting the MCH mid-time-step dynamics savepoints (see the MCH branch in +# the test body): solve-nonhydro exit at the corrector (istep=2) of the last substep +# (2 for MCH), and the non-initial diffusion savepoint. Only instantiated for MCH. +@pytest.fixture # type: ignore[no-redef] # deliberately shadows the fixtures.py import +def istep_exit() -> int: + return 2 + + +@pytest.fixture +def substep_exit() -> int: + return 2 + + +@pytest.fixture +def timeloop_diffusion_linit_exit() -> bool: + return False + + @pytest.mark.datatest +@pytest.mark.level("integration") @pytest.mark.embedded_remap_error @pytest.mark.parametrize( - "experiment_description, istep_exit, substep_exit, timeloop_date_init, timeloop_date_exit, step_date_exit, timeloop_diffusion_linit_init, timeloop_diffusion_linit_exit", + "experiment_description, timeloop_date_init, timeloop_date_exit, step_date_exit", [ ( test_defs.Experiments.JW, - 2, - 5, "2008-09-01T00:00:00.000", "2008-09-01T00:05:00.000", "2008-09-01T00:05:00.000", - False, - False, ), ( test_defs.Experiments.GAUSS3D, - 2, - 5, "2001-01-01T00:00:00.000", "2001-01-01T00:00:04.000", "2001-01-01T00:00:04.000", - False, - False, + ), + ( + test_defs.Experiments.EXCLAIM_APE_AES, + "2008-09-01T00:00:00.000", + "2008-09-01T00:05:00.000", + "2008-09-01T00:05:00.000", ), ( test_defs.Experiments.MCH_CH_R04B09, - 2, - 2, "2021-06-20T12:00:00.000", "2021-06-20T12:00:10.000", "2021-06-20T12:00:10.000", - True, - False, ), ( test_defs.Experiments.MCH_CH_R04B09, - 2, - 2, "2021-06-20T12:00:10.000", "2021-06-20T12:00:20.000", "2021-06-20T12:00:20.000", - False, - False, ), ], ) @@ -104,15 +125,43 @@ def test_standalone_driver( experiment_description: test_defs.ExperimentDescription, timeloop_date_init: str, timeloop_date_exit: str, - timeloop_diffusion_linit_init: bool, *, + request: pytest.FixtureRequest, tmp_path: pathlib.Path, process_props: decomp_defs.ProcessProperties, backend: gtx_typing.Backend, - savepoint_nonhydro_exit: sb.IconNonHydroExitSavepoint, - substep_exit: int, - savepoint_diffusion_exit: sb.IconDiffusionExitSavepoint, + savepoint_time_step_exit: sb.IconTimeStepExitSavepoint, ) -> None: + """End-to-end standalone-driver validation over one time step. + + Experiments validate the final prognostic state against the end-of-time-step + (``time-step-exit``) savepoint. EXCLAIM_APE_AES additionally runs muphys and also + validates the tracers. Exception: MCH_CH_R04B09 compares against the mid-time-step + dynamics savepoints, because its reference runs NWP physics + limited-area nudging + after the dynamics, which the driver does not (see the comment in the body). + Per-field tolerances live in ``_TOLERANCES``. + + muphys (EXCLAIM_APE_AES): runs ``MuphysScheme.AES_GRAUPEL`` -- the port of the exact + ICON formulation that generated the reference. Graupel is the only *physics* + parameterization active, so vn/w/rho/exner/theta_v compare tightly; the tracer + comparison carries residuals from gaps not yet ported: + + - exner / theta_v: recomputed via the exact EOS in ``scatter_to_prognostic``, mirroring + ICON's phy2dyn coupling (mo_interface_iconam_aes.f90). Measured on v6: exner ~3e-9 + (atol=1e-8), theta_v ~7e-9 relative (rtol=3e-8) -- essentially exact. + - tracer transport (KNOWN MISMATCH): the reference ran with tracer advection ON + (ltransport=.TRUE., MIURA/PPM), but the driver disables it here -- it can't yet + compute airmass (rho*ddqz) or wire live mass fluxes (TODO(OngChia)), so advection + would divide by a zero airmass. So this validates muphys in isolation, not + transport+muphys. Over one 300 s step the advective change is small: qv still passes + at atol=2e-6, qc/qr/qs/qi/qg match bit-for-bit (atol=1e-10). + TODO (Yilu): revisit once the driver computes airmass + wires the mass fluxes. + - negative tracers: ICON clips them (iqneg_d2p/iqneg_p2d); the driver does not. + - vertical extent: ICON runs graupel on jks_cloudy..nlev; muphys runs the full column. + + The muphys granule itself is validated in isolation against the aes-graupel savepoints + in test_muphys_datatest.py. + """ allocator = model_backends.get_allocator(backend) grid_file_path = grid_utils._download_grid_file(experiment_description.grid) @@ -147,20 +196,43 @@ def test_standalone_driver( ) prognostics = ds.prognostics.current + computed = { "vn": prognostics.vn, "w": prognostics.w, + "rho": prognostics.rho, "exner": prognostics.exner, "theta_v": prognostics.theta_v, - "rho": prognostics.rho, - } - references = { - "vn": savepoint_diffusion_exit.vn(), - "w": savepoint_diffusion_exit.w(), - "exner": savepoint_diffusion_exit.exner(), - "theta_v": savepoint_diffusion_exit.theta_v(), - "rho": savepoint_nonhydro_exit.rho_new(), } + if experiment_description is test_defs.Experiments.MCH_CH_R04B09: + # The MCH reference runs the full NWP physics suite (nwp_phy_nml: convection, + # radiation, SSO, graupel, satad) plus limited-area boundary nudging AFTER the + # dynamics -- none of which the driver runs -- so its end-of-step state is not + # comparable (vn differs by O(10) m/s). Validate against the mid-time-step + # dynamics savepoints instead until NWP physics is ported. + diffusion_exit = request.getfixturevalue("savepoint_diffusion_exit") + nonhydro_exit = request.getfixturevalue("savepoint_nonhydro_exit") + references = { + "vn": diffusion_exit.vn(), + "w": diffusion_exit.w(), + "rho": nonhydro_exit.rho_new(), + "exner": diffusion_exit.exner(), + "theta_v": diffusion_exit.theta_v(), + } + else: + # Nothing runs after diffusion for JW/GAUSS3D, and for EXCLAIM_APE_AES muphys + # is the only active physics: validate against the end-of-time-step savepoint. + references = { + "vn": savepoint_time_step_exit.vn(), + "w": savepoint_time_step_exit.w(), + "rho": savepoint_time_step_exit.rho(), + "exner": savepoint_time_step_exit.exner(), + "theta_v": savepoint_time_step_exit.theta_v(), + } + + for tracer in prognostics.tracers.active_fields(): + computed[tracer] = tracer.field + references[tracer] = getattr(savepoint_time_step_exit, tracer.name) tolerances = _TOLERANCES[experiment_description] diff --git a/model/testing/src/icon4py/model/testing/definitions.py b/model/testing/src/icon4py/model/testing/definitions.py index 76e2cd02cf..733170f155 100644 --- a/model/testing/src/icon4py/model/testing/definitions.py +++ b/model/testing/src/icon4py/model/testing/definitions.py @@ -156,7 +156,7 @@ class ExperimentDescription: name: str long_name: str grid: GridDescription - version: int = 5 + version: int = 6 @dataclasses.dataclass diff --git a/model/testing/src/icon4py/model/testing/fixtures/datatest.py b/model/testing/src/icon4py/model/testing/fixtures/datatest.py index ebb3a77a7c..9501fcf5fd 100644 --- a/model/testing/src/icon4py/model/testing/fixtures/datatest.py +++ b/model/testing/src/icon4py/model/testing/fixtures/datatest.py @@ -482,6 +482,21 @@ def savepoint_diffusion_exit( return sp +@pytest.fixture +def savepoint_time_step_exit( + data_provider: serialbox.IconSerialDataProvider, + step_date_exit: str, +) -> serialbox.IconTimeStepExitSavepoint: + """ + Load data from the ICON savepoint written at the end of the full time step, + after all physics tendencies have been applied. + + date of the timestamp to be selected can be set separately by overriding the + 'step_date_exit' fixture, passing 'step_date_exit=' + """ + return data_provider.from_savepoint_time_step_exit(date=step_date_exit) + + @pytest.fixture def istep_init() -> int: return 1 diff --git a/model/testing/src/icon4py/model/testing/serialbox.py b/model/testing/src/icon4py/model/testing/serialbox.py index 8d6576eff0..52801054b3 100644 --- a/model/testing/src/icon4py/model/testing/serialbox.py +++ b/model/testing/src/icon4py/model/testing/serialbox.py @@ -1962,6 +1962,118 @@ def topo_smt_c(self): return self._get_field("smooth_topography", dims.CellDim) +class IconTimeStepExitSavepoint(IconSavepoint): + """End-of-timestep prognostic state, written in perform_nh_timeloop right after + integrate_nh returns: all physics tendencies applied, time levels swapped.""" + + def vn(self): + return self._get_field("vn", dims.EdgeDim, dims.KDim) + + def w(self): + return self._get_field("w", dims.CellDim, dims.KDim) + + def rho(self): + return self._get_field("rho", dims.CellDim, dims.KDim) + + def exner(self): + return self._get_field("exner", dims.CellDim, dims.KDim) + + def theta_v(self): + return self._get_field("theta_v", dims.CellDim, dims.KDim) + + def tracer(self, ntracer: TracerIndex): + return self._get_field_component("tracers", ntracer, (dims.CellDim, dims.KDim)) + + def qv(self): + return self.tracer(QV) + + def qc(self): + return self.tracer(QC) + + def qi(self): + return self.tracer(QI) + + def qr(self): + return self.tracer(QR) + + def qs(self): + return self.tracer(QS) + + def qg(self): + return self.tracer(QG) + + +class IconMuphysSavepoint(IconSavepoint): + """Common fields of the aes-graupel-init/exit savepoints written around the mig + block (cloud_mig = satad + graupel + satad) in aes_phy_main. tend_ta/tend_tracers + are the prm_tend accumulators: exit minus init isolates the mig contribution.""" + + def temperature(self): + return self._get_field("temperature", dims.CellDim, dims.KDim) + + def tracer(self, ntracer: TracerIndex): + return self._get_field_component("tracers", ntracer, (dims.CellDim, dims.KDim)) + + def tend_ta(self): + return self._get_field("tend_ta", dims.CellDim, dims.KDim) + + def tend_tracer(self, ntracer: TracerIndex): + return self._get_field_component("tend_tracers", ntracer, (dims.CellDim, dims.KDim)) + + def qv(self): + return self.tracer(QV) + + def qc(self): + return self.tracer(QC) + + def qi(self): + return self.tracer(QI) + + def qr(self): + return self.tracer(QR) + + def qs(self): + return self.tracer(QS) + + def qg(self): + return self.tracer(QG) + + +class IconMuphysInitSavepoint(IconMuphysSavepoint): + def dz(self): + return self._get_field("dz", dims.CellDim, dims.KDim) + + def rho(self): + return self._get_field("rho", dims.CellDim, dims.KDim) + + def pressure(self): + return self._get_field("pressure", dims.CellDim, dims.KDim) + + def dtime(self): + return self.serializer.read("dtime", self.savepoint)[0] + + def jks_cloudy(self): + return int(self.serializer.read("jks_cloudy", self.savepoint)[0]) + + +class IconMuphysExitSavepoint(IconMuphysSavepoint): + def rsfl(self): + # surface rain rate + return self._get_field("rsfl", dims.CellDim) + + def ssfl(self): + # surface frozen precip rate: ice + snow + graupel + return self._get_field("ssfl", dims.CellDim) + + def pr(self): + # total surface precip rate + return self._get_field("pr", dims.CellDim) + + def ufcs(self): + # surface precip energy flux + return self._get_field("ufcs", dims.CellDim) + + class IconSerialDataProvider: def __init__( self, @@ -2253,3 +2365,21 @@ def from_savepoint_satad_exit(self, location: str, date: str) -> IconSatadExitSa return IconSatadExitSavepoint( savepoint, self.serializer, size=self.grid_size, backend=self.backend ) + + def from_savepoint_time_step_exit(self, date: str) -> IconTimeStepExitSavepoint: + savepoint = self.serializer.savepoint["time-step-exit"].id[1].date[date].as_savepoint() + return IconTimeStepExitSavepoint( + savepoint, self.serializer, size=self.grid_size, backend=self.backend + ) + + def from_savepoint_muphys_init(self, date: str) -> IconMuphysInitSavepoint: + savepoint = self.serializer.savepoint["aes-graupel-init"].id[1].date[date].as_savepoint() + return IconMuphysInitSavepoint( + savepoint, self.serializer, size=self.grid_size, backend=self.backend + ) + + def from_savepoint_muphys_exit(self, date: str) -> IconMuphysExitSavepoint: + savepoint = self.serializer.savepoint["aes-graupel-exit"].id[1].date[date].as_savepoint() + return IconMuphysExitSavepoint( + savepoint, self.serializer, size=self.grid_size, backend=self.backend + ) diff --git a/noxfile.py b/noxfile.py index 853e00da6e..51b36a0a41 100644 --- a/noxfile.py +++ b/noxfile.py @@ -55,6 +55,7 @@ class _VenvBackendKwargs(TypedDict, total=False): "atmosphere/dycore", "atmosphere/subgrid_scale_physics/microphysics", "atmosphere/subgrid_scale_physics/muphys", + "atmosphere/subgrid_scale_physics/physics_driver", "common", "standalone_driver", "testing", diff --git a/pyproject.toml b/pyproject.toml index 19437db5cc..b2f505f7e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,6 +93,7 @@ dependencies = [ "icon4py-atmosphere-dycore~=0.3.0", "icon4py-atmosphere-microphysics~=0.3.0", "icon4py-atmosphere-muphys~=0.3.0", + "icon4py-atmosphere-physics-driver~=0.3.0", "icon4py-common~=0.3.0", "icon4py-standalone_driver~=0.3.0" ] @@ -157,6 +158,10 @@ files = [ # '$MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/microphysics/tests', # '$MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/muphys/src', # '$MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/muphys/tests', + '$MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/physics_driver/src', + '$MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/state.py', + '$MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/component.py', + '$MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/data.py', # '$MYPY_CONFIG_FILE_DIR/model/common/src/', '$MYPY_CONFIG_FILE_DIR/model/common/src/icon4py/model/common/utils/', # '$MYPY_CONFIG_FILE_DIR/model/common/tests/', @@ -190,6 +195,7 @@ $MYPY_CONFIG_FILE_DIR/model/atmosphere/diffusion/src: $MYPY_CONFIG_FILE_DIR/model/atmosphere/dycore/src: $MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/microphysics/src: $MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/muphys/src: +$MYPY_CONFIG_FILE_DIR/model/atmosphere/subgrid_scale_physics/physics_driver/src: $MYPY_CONFIG_FILE_DIR/model/common/src: $MYPY_CONFIG_FILE_DIR/model/standalone_driver/src: $MYPY_CONFIG_FILE_DIR/model/standalone_driver/tests: @@ -269,6 +275,7 @@ pythonpath = [ "model/atmosphere/dycore", "model/atmosphere/subgrid_scale_physics/microphysics", "model/atmosphere/subgrid_scale_physics/muphys", + "model/atmosphere/subgrid_scale_physics/physics_driver", "model/common", "model/standalone_driver", "model/testing", @@ -414,6 +421,7 @@ icon4py-atmosphere-diffusion = {workspace = true} icon4py-atmosphere-dycore = {workspace = true} icon4py-atmosphere-microphysics = {workspace = true} icon4py-atmosphere-muphys = {workspace = true} +icon4py-atmosphere-physics-driver = {workspace = true} # dace = {index = "gridtools"} # gt4py = {git = "https://github.com/GridTools/gt4py", branch = "main"} # gt4py = {index = "test.pypi"} @@ -432,6 +440,7 @@ members = [ "model/atmosphere/dycore", "model/atmosphere/subgrid_scale_physics/microphysics", "model/atmosphere/subgrid_scale_physics/muphys", + "model/atmosphere/subgrid_scale_physics/physics_driver", "model/common", "model/standalone_driver", "model/testing", diff --git a/scripts/python/generate_ci_pipeline.py b/scripts/python/generate_ci_pipeline.py index 6b76d14b70..360a7aa865 100755 --- a/scripts/python/generate_ci_pipeline.py +++ b/scripts/python/generate_ci_pipeline.py @@ -52,6 +52,7 @@ "dycore", "microphysics", "muphys", + "physics_driver", "common", "standalone_driver", ] diff --git a/tach.toml b/tach.toml index 3b9b8746c8..4ace349188 100644 --- a/tach.toml +++ b/tach.toml @@ -6,6 +6,7 @@ source_roots = [ "model/atmosphere/dycore/src", "model/atmosphere/subgrid_scale_physics/microphysics/src", "model/atmosphere/subgrid_scale_physics/muphys/src", + "model/atmosphere/subgrid_scale_physics/physics_driver/src", "model/standalone_driver/src", "model/common/src", "model/testing/src", @@ -39,6 +40,10 @@ rename = ["serialbox:serialbox4py"] path = "icon4py.model.atmosphere.subgrid_scale_physics.muphys" depends_on = [{ path = "icon4py.model.common" }] +[[modules]] +path = "icon4py.model.atmosphere.subgrid_scale_physics.physics_driver" +depends_on = [{ path = "icon4py.model.common" }] + [[modules]] path = "icon4py.model.atmosphere.tracer_advection" depends_on = [{ path = "icon4py.model.common" }] @@ -58,6 +63,8 @@ depends_on = [{ path = "icon4py.model.common" }] [[modules]] path = "icon4py.model.standalone_driver" depends_on = [ + { path = "icon4py.model.atmosphere.subgrid_scale_physics.muphys" }, + { path = "icon4py.model.atmosphere.subgrid_scale_physics.physics_driver" }, { path = "icon4py.model.common" }, ] diff --git a/uv.lock b/uv.lock index 8d6019e778..121eb8c74f 100644 --- a/uv.lock +++ b/uv.lock @@ -30,11 +30,12 @@ version = 1 [manifest] members = [ "icon4py", - "icon4py-atmosphere-tracer_advection", "icon4py-atmosphere-diffusion", "icon4py-atmosphere-dycore", "icon4py-atmosphere-microphysics", "icon4py-atmosphere-muphys", + "icon4py-atmosphere-physics-driver", + "icon4py-atmosphere-tracer-advection", "icon4py-bindings", "icon4py-common", "icon4py-standalone-driver", @@ -1756,11 +1757,12 @@ wheels = [ [[package]] dependencies = [ - {name = "icon4py-atmosphere-tracer_advection"}, {name = "icon4py-atmosphere-diffusion"}, {name = "icon4py-atmosphere-dycore"}, {name = "icon4py-atmosphere-microphysics"}, {name = "icon4py-atmosphere-muphys"}, + {name = "icon4py-atmosphere-physics-driver"}, + {name = "icon4py-atmosphere-tracer-advection"}, {name = "icon4py-common"}, {name = "icon4py-standalone-driver"} ] @@ -1866,11 +1868,12 @@ typing-distributed = [ provides-extras = ["all", "cuda12", "cuda13", "distributed", "fortran", "io", "profiling", "rocm7", "testing"] requires-dist = [ {name = "icon4py", extras = ["distributed", "fortran", "io", "testing", "profiling"], marker = "extra == 'all'"}, - {name = "icon4py-atmosphere-tracer_advection", editable = "model/atmosphere/tracer_advection"}, {name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion"}, {name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore"}, {name = "icon4py-atmosphere-microphysics", editable = "model/atmosphere/subgrid_scale_physics/microphysics"}, {name = "icon4py-atmosphere-muphys", editable = "model/atmosphere/subgrid_scale_physics/muphys"}, + {name = "icon4py-atmosphere-physics-driver", editable = "model/atmosphere/subgrid_scale_physics/physics_driver"}, + {name = "icon4py-atmosphere-tracer-advection", editable = "model/atmosphere/tracer_advection"}, {name = "icon4py-bindings", marker = "extra == 'fortran'", editable = "bindings"}, {name = "icon4py-common", editable = "model/common"}, {name = "icon4py-common", extras = ["cuda12"], marker = "extra == 'cuda12'", editable = "model/common"}, @@ -2015,8 +2018,8 @@ dependencies = [ {name = "icon4py-common"}, {name = "packaging"} ] -name = "icon4py-atmosphere-tracer_advection" -source = {editable = "model/atmosphere/tracer_advection"} +name = "icon4py-atmosphere-diffusion" +source = {editable = "model/atmosphere/diffusion"} version = "0.3.0" [package.metadata] @@ -2032,8 +2035,8 @@ dependencies = [ {name = "icon4py-common"}, {name = "packaging"} ] -name = "icon4py-atmosphere-diffusion" -source = {editable = "model/atmosphere/diffusion"} +name = "icon4py-atmosphere-dycore" +source = {editable = "model/atmosphere/dycore"} version = "0.3.0" [package.metadata] @@ -2049,8 +2052,8 @@ dependencies = [ {name = "icon4py-common"}, {name = "packaging"} ] -name = "icon4py-atmosphere-dycore" -source = {editable = "model/atmosphere/dycore"} +name = "icon4py-atmosphere-microphysics" +source = {editable = "model/atmosphere/subgrid_scale_physics/microphysics"} version = "0.3.0" [package.metadata] @@ -2063,16 +2066,33 @@ requires-dist = [ [[package]] dependencies = [ {name = "gt4py"}, - {name = "icon4py-common"}, + {name = "icon4py-common", extra = ["io"]}, + {name = "numpy"}, {name = "packaging"} ] -name = "icon4py-atmosphere-microphysics" -source = {editable = "model/atmosphere/subgrid_scale_physics/microphysics"} +name = "icon4py-atmosphere-muphys" +source = {editable = "model/atmosphere/subgrid_scale_physics/muphys"} version = "0.3.0" [package.metadata] requires-dist = [ {name = "gt4py", specifier = "==1.1.12"}, + {name = "icon4py-common", extras = ["io"], editable = "model/common"}, + {name = "numpy", specifier = ">=1.23.3"}, + {name = "packaging", specifier = ">=20.0"} +] + +[[package]] +dependencies = [ + {name = "icon4py-common"}, + {name = "packaging"} +] +name = "icon4py-atmosphere-physics-driver" +source = {editable = "model/atmosphere/subgrid_scale_physics/physics_driver"} +version = "0.2.0" + +[package.metadata] +requires-dist = [ {name = "icon4py-common", editable = "model/common"}, {name = "packaging", specifier = ">=20.0"} ] @@ -2080,19 +2100,17 @@ requires-dist = [ [[package]] dependencies = [ {name = "gt4py"}, - {name = "icon4py-common", extra = ["io"]}, - {name = "numpy"}, + {name = "icon4py-common"}, {name = "packaging"} ] -name = "icon4py-atmosphere-muphys" -source = {editable = "model/atmosphere/subgrid_scale_physics/muphys"} +name = "icon4py-atmosphere-tracer-advection" +source = {editable = "model/atmosphere/tracer_advection"} version = "0.3.0" [package.metadata] requires-dist = [ {name = "gt4py", specifier = "==1.1.12"}, - {name = "icon4py-common", extras = ["io"], editable = "model/common"}, - {name = "numpy", specifier = ">=1.23.3"}, + {name = "icon4py-common", editable = "model/common"}, {name = "packaging", specifier = ">=20.0"} ] @@ -2235,6 +2253,8 @@ dependencies = [ {name = "gt4py"}, {name = "icon4py-atmosphere-diffusion"}, {name = "icon4py-atmosphere-dycore"}, + {name = "icon4py-atmosphere-muphys"}, + {name = "icon4py-atmosphere-physics-driver"}, {name = "icon4py-common", extra = ["io"]}, {name = "packaging"}, {name = "typer"} @@ -2249,6 +2269,8 @@ requires-dist = [ {name = "gt4py", specifier = "==1.1.12"}, {name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion"}, {name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore"}, + {name = "icon4py-atmosphere-muphys", editable = "model/atmosphere/subgrid_scale_physics/muphys"}, + {name = "icon4py-atmosphere-physics-driver", editable = "model/atmosphere/subgrid_scale_physics/physics_driver"}, {name = "icon4py-common", extras = ["io"], editable = "model/common"}, {name = "packaging", specifier = ">=20.0"}, {name = "typer", specifier = ">=0.20.0"}