Add time-step-exit and aes-graupel savepoint validation

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/component.py

@@ -15,7 +15,10 @@
 
 import gt4py.next as gtx
 
-from icon4py.model.atmosphere.subgrid_scale_physics.muphys import data as muphys_data
+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 import run_full_muphys
 from icon4py.model.common import (
@@ -50,6 +53,7 @@ def __init__(
         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 = ncells
@@ -88,6 +92,7 @@ def __init__(
                 qnc=qnc,
                 backend=backend,
                 single_program=False,
+                scheme=scheme,
             )
         self._step = step
 

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/config.py

@@ -9,10 +9,26 @@
 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 carried by
+    icon-nwp (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.0 # Cloud droplet number concentration [cm^-3], default value hardcoded
+    qnc: float = 50.0e6  # cloud droplet number concentration [m^-3], matches Fortran cloud_num
+    scheme: MuphysScheme = MuphysScheme.AES_GRAUPEL  # the driver validates against icon-nwp

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/common/constants.py

@@ -83,3 +83,40 @@ class IndexConsts(ta.wpfloat, enum.Enum):
     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 carried by icon-nwp
+    # (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 = 0.16666666666666666
+    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

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-nwp 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-nwp 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)
+    )