Physics driver with Muphys

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

@@ -0,0 +1,208 @@
+# 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 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,
+    type_alias as ta,
+)
+from icon4py.model.common.diagnostic_calculations.stencils import calculate_tendency
+from icon4py.model.common.math.stencils import generic_math_operations
+
+
+if TYPE_CHECKING:
+    import gt4py.next.typing as gtx_typing
+
+    from icon4py.model.common.states import model
+
+
+class MuphysComponent:
+    """The muphys Granule: a 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,
+        ncells: int,
+        nlev: int,
+        dtime: datetime.timedelta,
+        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
+        self._nlev = nlev
+        self._dt_seconds = dtime.total_seconds()
+        self._qnc = qnc
+        self._backend = model_options.customize_backend(program=None, backend=backend)
+
+        horizontal_sizes = {
+            "horizontal_start": gtx.int32(0),
+            "horizontal_end": gtx.int32(self._ncells),
+        }
+        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=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=horizontal_sizes,
+            vertical_sizes=vertical_sizes,
+            offset_provider={},
+        )
+
+        allocator = model_backends.get_allocator(backend)
+
+        if step is None:
+            sizes = types.SimpleNamespace(ncells=ncells, nlev=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: ncells, dims.KDim: 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 _to_tendency(
+        self,
+        old: fa.CellKField[ta.wpfloat],
+        new: fa.CellKField[ta.wpfloat],
+        out: fa.CellKField[ta.wpfloat],
+    ) -> None:
+        """``out = (new - old) / dt`` over the whole column."""
+        self._calculate_tendency(
+            dtime=self._dt_seconds,
+            old_field=old,
+            new_field=new,
+            tendency=out,
+        )
+
+    def _copy_into(self, src: fa.CellKField[ta.wpfloat], dst: fa.CellKField[ta.wpfloat]) -> None:
+        """Copy ``src`` into the granule-owned buffer ``dst``."""
+        self._copy_field(
+            field=src,
+            output_field=dst,
+        )
+
+    def __call__(
+        self, state: dict[str, model.DataField], time_step: datetime.datetime
+    ) -> 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_into(fields["te"], self._te_in)
+        for s in SPECIES:
+            self._copy_into(fields[f"q{s}"], 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._to_tendency(fields["te"], self._t_out, self._tendencies["tend_temperature"])
+        for s in SPECIES:
+            self._to_tendency(
+                fields[f"q{s}"], getattr(self._q_out, s), 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,
+            },
+        )

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 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.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/definitions.py

@@ -27,3 +27,6 @@ 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