Muphys bug fix: drivers run_graupel_only and run_full_muphys

model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
   # workspace members
   "icon4py-common[io]>=0.0.6",
   # external dependencies
+  "numpy>=1.23.3",
   "gt4py==1.1.2",
   "packaging>=20.0"
 ]

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py

@@ -0,0 +1,20 @@
+# 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 NamedTuple
+
+from icon4py.model.common import field_type_aliases as fa, type_alias as ta
+
+
+class Q(NamedTuple):
+    v: fa.CellKField[ta.wpfloat]  # Specific humidity
+    c: fa.CellKField[ta.wpfloat]  # Specific cloud water content
+    r: fa.CellKField[ta.wpfloat]  # Specific rain water
+    s: fa.CellKField[ta.wpfloat]  # Specific snow water
+    i: fa.CellKField[ta.wpfloat]  # Specific ice water content
+    g: fa.CellKField[ta.wpfloat]  # Specific graupel water content

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment.py

@@ -0,0 +1,92 @@
+# 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 maximum, where
+
+from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.frozen import g_ct, t_d
+from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import Q
+from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import (
+    _newton_raphson,
+    _qsat_rho,
+)
+from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+
+
+@gtx.field_operator
+def _saturation_adjustment(
+    te: fa.CellKField[ta.wpfloat], rho: fa.CellKField[ta.wpfloat], q_in: Q
+) -> tuple[
+    fa.CellKField[ta.wpfloat],
+    fa.CellKField[ta.wpfloat],
+    fa.CellKField[ta.wpfloat],
+]:
+    """
+    Compute the saturation adjustment which revises internal energy and water contents
+
+    Args:
+        Tx:                    Temperature
+        rho:                   Density containing dry air and water constituents
+        Q:                     Class with humidity, cloud, rain, snow, ice and graupel water
+
+    Result:                    Tuple containing
+                               - Revised temperature
+                               - Revised specific cloud water content
+                               - Revised specific vapor content
+    """
+    qti = q_in.s + q_in.i + q_in.g
+    qt = q_in.v + q_in.c + q_in.r + qti
+    cvc = t_d.cvd * (1.0 - qt) + t_d.clw * q_in.r + g_ct.ci * qti
+    cv = cvc + t_d.cvv * q_in.v + t_d.clw * q_in.c
+    ue = cv * te - q_in.c * g_ct.lvc
+    Tx_hold = ue / (cv + q_in.c * (t_d.cvv - t_d.clw))
+    qx_hold = _qsat_rho(Tx_hold, rho)
+
+    Tx = te
+    # Newton-Raphson iteration: 6 times the same operations
+    Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue)
+    Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue)
+    Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue)
+    Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue)
+    Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue)
+    Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue)
+
+    # At this point we hope Tx has converged
+    qx = _qsat_rho(Tx, rho)
+
+    # Is it possible to unify the where for all three outputs??
+    mask = q_in.v + q_in.c <= qx_hold
+    te = where(mask, Tx_hold, Tx)
+    qce = where(mask, 0.0, maximum(q_in.v + q_in.c - qx, 0.0))
+    qve = where(mask, q_in.v + q_in.c, qx)
+
+    return te, qve, qce
+
+
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
+def saturation_adjustment(
+    te: fa.CellKField[ta.wpfloat],  # Temperature
+    rho: fa.CellKField[ta.wpfloat],  # Density containing dry air and water constituents
+    q_in: Q,  # Class with humidity, cloud, rain, snow, ice and graupel water
+    te_out: fa.CellKField[ta.wpfloat],  # Temperature
+    qve_out: fa.CellKField[ta.wpfloat],  # Specific humidity
+    qce_out: fa.CellKField[ta.wpfloat],  # Specific cloud water content
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
+):
+    _saturation_adjustment(
+        te,
+        rho,
+        q_in,
+        out=(te_out, qve_out, qce_out),
+        domain={
+            dims.CellDim: (horizontal_start, horizontal_end),
+            dims.KDim: (vertical_start, vertical_end),
+        },
+    )

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py

@@ -6,7 +6,7 @@
 # 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, where
+from gt4py.next import exp
 
 from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.frozen import g_ct, t_d
 from icon4py.model.common import field_type_aliases as fa, type_alias as ta
@@ -338,78 +338,3 @@ def _newton_raphson(
     dux = cv + dqx * (g_ct.lvc + (t_d.cvv - t_d.clw) * Tx)
     Tx = Tx - (ux - ue) / dux
     return Tx
-
-
-@gtx.field_operator
-def _saturation_adjustment(
-    te: fa.CellKField[ta.wpfloat],
-    qve: fa.CellKField[ta.wpfloat],
-    qce: fa.CellKField[ta.wpfloat],
-    qre: fa.CellKField[ta.wpfloat],
-    qti: fa.CellKField[ta.wpfloat],
-    rho: fa.CellKField[ta.wpfloat],
-) -> tuple[
-    fa.CellKField[ta.wpfloat],
-    fa.CellKField[ta.wpfloat],
-    fa.CellKField[ta.wpfloat],
-    fa.CellKField[bool],
-]:
-    """
-    Compute the saturation adjustment which revises internal energy and water contents
-
-    Args:
-        Tx:                    Temperature
-        qve:                   Specific humidity
-        qce:                   Specific cloud water content
-        qre:                   Specific rain water
-        qti:                   Specific mass of all ice species (total-ice)
-        rho:                   Density containing dry air and water constituents
-
-    Result:                    Tuple containing
-                               - Revised temperature
-                               - Revised specific cloud water content
-                               - Revised specific vapor content
-                               - Mask specifying where qce+qve less than holding capacity
-    """
-    qt = qve + qce + qre + qti
-    cvc = t_d.cvd * (1.0 - qt) + t_d.clw * qre + g_ct.ci * qti
-    cv = cvc + t_d.cvv * qve + t_d.clw * qce
-    ue = cv * te - qce * g_ct.lvc
-    Tx_hold = ue / (cv + qce * (t_d.cvv - t_d.clw))
-    qx_hold = _qsat_rho(Tx_hold, rho)
-
-    Tx = te
-    # Newton-Raphson iteration: 6 times the same operations
-    Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue)
-    Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue)
-    Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue)
-    Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue)
-    Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue)
-    Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue)
-
-    # At this point we hope Tx has converged
-    qx = _qsat_rho(Tx, rho)
-
-    # Is it possible to unify the where for all three outputs??
-    mask = qve + qce <= qx_hold
-    te = where((qve + qce <= qx_hold), Tx_hold, Tx)
-    qce = where((qve + qce <= qx_hold), 0.0, maximum(qve + qce - qx, 0.0))
-    qve = where((qve + qce <= qx_hold), qve + qce, qx)
-
-    return te, qve, qce, mask
-
-
-@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def saturation_adjustment(
-    te: fa.CellKField[ta.wpfloat],  # Temperature
-    qve: fa.CellKField[ta.wpfloat],  # Specific humidity
-    qce: fa.CellKField[ta.wpfloat],  # Specific cloud water content
-    qre: fa.CellKField[ta.wpfloat],  # Specific rain water
-    qti: fa.CellKField[ta.wpfloat],  # Specific mass of all ice species (total-ice)
-    rho: fa.CellKField[ta.wpfloat],  # Density containing dry air and water constituents
-    te_out: fa.CellKField[ta.wpfloat],  # Temperature
-    qve_out: fa.CellKField[ta.wpfloat],  # Specific humidity
-    qce_out: fa.CellKField[ta.wpfloat],  # Specific cloud water content
-    mask_out: fa.CellKField[bool],  # Specific cloud water content
-):
-    _saturation_adjustment(te, qve, qce, qre, qti, rho, out=(te_out, qve_out, qce_out, mask_out))

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/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