Repository: C2SM/icon4py
Pull request #740: Clean up saturation adjustment
Author: login: OngChia
Created: 2025-05-09T13:50:21Z | Merged: 2025-05-15T14:40:09Z

Clean up saturation adjustment:

1. relocate diagnostic calculations in saturation adjustment granule to a separate module that holds all tendency calculations.
2. improve the Newton's iteration in saturation adjustment.
3. Use `FrozenNameSpace` for constants used in all diagnostic calculations instead of passing in as input arguments.


Diff:
diff --git a/model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
index 8c67b7303d..c49717e8e1 100644
--- a/model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
+++ b/model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
@@ -9,9 +9,8 @@
 from typing import Final, Optional
 
 import gt4py.next as gtx
-import numpy as np
-from gt4py.eve.utils import FrozenNamespace
-from gt4py.next import backend as gtx_backend, broadcast
+from gt4py.eve import utils as eve_utils
+from gt4py.next import backend as gtx_backend
 from gt4py.next.ffront.fbuiltins import (
     abs,
     exp,
@@ -20,66 +19,47 @@
 )
 
 import icon4py.model.common.dimension as dims
+import icon4py.model.common.utils as common_utils
 from icon4py.model.common import (
     constants as phy_const,
     field_type_aliases as fa,
     type_alias as ta,
 )
-from icon4py.model.common.diagnostic_calculations.stencils import (
-    diagnose_pressure as pressure,
-    diagnose_surface_pressure as surface_pressure,
-)
 from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid, vertical as v_grid
-from icon4py.model.common.states import (
-    diagnostic_state as diagnostics,
-    model,
-    prognostic_state as prognostics,
-    tracer_state as tracers,
-)
+from icon4py.model.common.states import model
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
+physics_constants: Final = phy_const.PhysicsConstants()
+
+
 # TODO (Chia Rui): Refactor this class when direct import is enabled for gt4py stencils
-class SaturatedPressureConstants(FrozenNamespace):
+class MicrophysicsConstants(eve_utils.FrozenNamespace[ta.wpfloat]):
     """
     Constants used for the computation of saturated pressure in saturation adjustment and microphysics.
     It was originally in mo_lookup_tables_constants.f90.
     """
 
     #: Latent heat of vaporisation for water [J kg-1]. Originally expressed as alv in ICON.
-    vaporisation_latent_heat: ta.wpfloat = 2.5008e6
+    vaporisation_latent_heat = 2.5008e6
     #: Melting temperature of ice/snow [K]
-    tmelt: ta.wpfloat = 273.15
+    melting_temperature = 273.15
 
-    #: See docstring in common/constanst.py
-    rd: ta.wpfloat = phy_const.GAS_CONSTANT_DRY_AIR
-    #: See docstring in common/constanst.py
-    rv: ta.wpfloat = phy_const.GAS_CONSTANT_WATER_VAPOR
-    #: See docstring in common/constanst.py
-    cvd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_VOLUME
-    #: See docstring in common/constanst.py
-    cpd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_PRESSURE
-
-    rv_o_rd_minus_1: ta.wpfloat = rv / rd - 1.0
-    rd_o_cpd: ta.wpfloat = phy_const.RD_O_CPD
-
-    #: Dry air heat capacity at constant pressure / water heat capacity at constant pressure - 1
-    rcpl: ta.wpfloat = 3.1733
-    #: Specific heat capacity of liquid water [J K-1 kg-1]. Originally expressed as clw in ICON.
-    spec_heat_cap_water: ta.wpfloat = (rcpl + 1.0) * cpd
+    #: Specific heat capacity of liquid water [J K-1 kg-1]. Originally expressed as clw in ICON. 4.1733 is dry air heat capacity at constant pressure / water heat capacity at constant pressure.
+    specific_heat_capacity_for_water = 4.1733 * phy_const.CPD
 
     #: p0 in Tetens formula for saturation water pressure, see eq. 5.33 in COSMO documentation. Originally expressed as c1es in ICON.
-    tetens_p0: ta.wpfloat = 610.78
+    tetens_p0 = 610.78
     #: aw in Tetens formula for saturation water pressure. Originally expressed as c3les in ICON.
-    tetens_aw: ta.wpfloat = 17.269
+    tetens_aw = 17.269
     #: bw in Tetens formula for saturation water pressure. Originally expressed as c4les in ICON.
-    tetens_bw: ta.wpfloat = 35.86
+    tetens_bw = 35.86
     #: numerator in temperature partial derivative of Tetens formula for saturation water pressure (psat tetens_der / (t - tetens_bw)^2). Originally expressed as c5les in ICON.
-    tetens_der: ta.wpfloat = tetens_aw * (tmelt - tetens_bw)
+    tetens_der = tetens_aw * (melting_temperature - tetens_bw)
 
 
 # Instantiate the class
-satpres_const: Final = SaturatedPressureConstants()
+microphy_const: Final = MicrophysicsConstants()
 
 
 @dataclasses.dataclass(frozen=True)
@@ -88,8 +68,6 @@ class SaturationAdjustmentConfig:
     max_iter: int = 10
     #: in ICON, 1.e-3 is always used for the tolerance when subroutine satad_v_3D is called.
     tolerance: ta.wpfloat = 1.0e-3
-    #: An extra step of updating the variables from new temperature is done in ICON after satad is called in at the beginning of mo_nh_interface_nwp.f90. This is a new option to update those variables.
-    diagnose_variables_from_new_temperature: bool = True
 
 
 @dataclasses.dataclass
@@ -109,36 +87,12 @@ class ConvergenceError(Exception):
         units="kg m-3",
         icon_var_name="rho",
     ),
-    exner_function=dict(
-        standard_name="dimensionless_exner_function",
-        long_name="exner function",
-        icon_var_name="exner",
-        units="1",
-    ),
     temperature=dict(
         standard_name="air_temperature",
         long_name="air temperature",
         units="K",
         icon_var_name="temp",
     ),
-    virtual_temperature=dict(
-        standard_name="air_virtual_temperature",
-        long_name="air virtual temperature",
-        units="K",
-        icon_var_name="tempv",
-    ),
-    pressure=dict(
-        standard_name="air_pressure",
-        long_name="air pressure",
-        units="Pa",
-        icon_var_name="pres",
-    ),
-    interface_pressure=dict(
-        standard_name="air_pressure_at_interface_level",
-        long_name="air pressure at model interface level",
-        units="Pa",
-        icon_var_name="pres_ifc",
-    ),
     specific_humidity=dict(
         standard_name="specific_humidity",
         long_name="ratio of water vapor mass to total moist air parcel mass",
@@ -151,60 +105,16 @@ class ConvergenceError(Exception):
         units="1",
         icon_var_name="qc",
     ),
-    specific_ice=dict(
-        standard_name="specific_ice_content",
-        long_name="ratio of cloud ice mass to total moist air parcel mass",
-        units="1",
-        icon_var_name="qi",
-    ),
-    specific_rain=dict(
-        standard_name="specific_rain_content",
-        long_name="ratio of rain mass to total moist air parcel mass",
-        units="1",
-        icon_var_name="qr",
-    ),
-    specific_snow=dict(
-        standard_name="specific_snow_content",
-        long_name="ratio of snow mass to total moist air parcel mass",
-        units="1",
-        icon_var_name="qs",
-    ),
-    specific_graupel=dict(
-        standard_name="specific_graupel_content",
-        long_name="ratio of graupel mass to total moist air parcel mass",
-        units="1",
-        icon_var_name="qg",
-    ),
 )
 
 
 #: CF attributes of saturation adjustment output variables
 _SATURATION_ADJUST_OUTPUT_ATTRIBUTES: Final[dict[str, model.FieldMetaData]] = dict(
-    tend_exner_function_due_to_satad=dict(
-        standard_name="tendency_of_dimensionless_exner_function_due_to_saturation_adjustment",
-        long_name="tendency of exner function due to saturation adjustment",
-        units="K s-1",
-    ),
     tend_temperature_due_to_satad=dict(
         standard_name="tendency_of_air_temperature_due_to_saturation_adjustment",
         long_name="tendency of air temperature due to saturation adjustment",
         units="K s-1",
     ),
-    tend_virtual_temperature_due_to_satad=dict(
-        standard_name="tendency_of_air_virtual_temperature_due_to_saturation_adjustment",
-        long_name="air virtual temperature",
-        units="K s-1",
-    ),
-    tend_pressure_due_to_satad=dict(
-        standard_name="tendency_of_air_pressure_due_to_saturation_adjustment",
-        long_name="tendency of air pressure due to saturation adjustment",
-        units="Pa s-1",
-    ),
-    tend_interface_pressure_due_to_satad=dict(
-        standard_name="tendency_of_air_pressure_at_interface_level_due_to_saturation_adjustment",
-        long_name="tendency of air pressure at model interface level due to saturation adjustment",
-        units="Pa s-1",
-    ),
     tend_specific_humidity_due_to_satad=dict(
         standard_name="tendency_of_specific_humidity_due_to_saturation_adjustment",
         long_name="tendency of ratio of water vapor mass to total moist air parcel mass due to saturation adjustment",
@@ -232,7 +142,10 @@ def __init__(
         self.grid = grid
         self.vertical_params: v_grid.VerticalGrid = vertical_params
         self.metric_state: MetricStateSaturationAdjustment = metric_state
-        self._allocate_tendencies()
+        self._allocate_local_variables()
+        self._determine_horizontal_domains()
+
+        self._xp = data_alloc.import_array_ns(self._backend)
 
         cell_domain = h_grid.domain(dims.CellDim)
         self._start_cell_nudging = self.grid.start_index(cell_domain(h_grid.Zone.NUDGING))
@@ -245,7 +158,7 @@ def input_properties(self) -> dict[str, model.FieldMetaData]:
     def output_properties(self) -> dict[str, model.FieldMetaData]:
         raise NotImplementedError
 
-    def _allocate_tendencies(self):
+    def _allocate_local_variables(self):
         #: it was originally named as tworkold in ICON. Old temperature before iteration.
         self._temperature1 = data_alloc.zero_field(
             self.grid, dims.CellDim, dims.KDim, backend=self._backend
@@ -266,83 +179,43 @@ def _allocate_tendencies(self):
         self._lwdocvd = data_alloc.zero_field(
             self.grid, dims.CellDim, dims.KDim, backend=self._backend
         )
-        self._k_field = data_alloc.index_field(
-            self.grid, dims.KDim, extend={dims.KDim: 1}, dtype=gtx.int32, backend=self._backend
-        )
-        # TODO (Chia Rui): remove local pressure and pressure_ifc when scan operator can be called along with pressure tendency computation
-        self._pressure = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        self._pressure_ifc = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, backend=self._backend
-        )
-        self._pressure_ifc = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, backend=self._backend
-        )
-        self._new_exner = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        self._new_virtual_temperature = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        # TODO (Chia Rui): remove the tendency terms below when architecture of the entire phyiscs component is ready to use.
-        self.temperature_tendency = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        self.qv_tendency = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        self.qc_tendency = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        self.virtual_temperature_tendency = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        self.exner_tendency = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        self.pressure_tendency = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, backend=self._backend
-        )
-        self.pressure_ifc_tendency = data_alloc.zero_field(
-            self.grid, dims.CellDim, dims.KDim, extend={dims.KDim: 1}, backend=self._backend
-        )
-
         self.compute_subsaturated_case_and_initialize_newton_iterations = (
             compute_subsaturated_case_and_initialize_newton_iterations.with_backend(self._backend)
         )
         self.update_temperature_by_newton_iteration = (
             update_temperature_by_newton_iteration.with_backend(self._backend)
         )
-        self.compute_newton_iteration_mask = compute_newton_iteration_mask.with_backend(
-            self._backend
-        )
-        self.copy_temperature = copy_temperature.with_backend(self._backend)
-        self.update_temperature_qv_qc_tendencies = update_temperature_qv_qc_tendencies.with_backend(
-            self._backend
-        )
-        self.compute_temperature_and_exner_tendencies_after_saturation_adjustment = (
-            compute_temperature_and_exner_tendencies_after_saturation_adjustment.with_backend(
+        self.compute_newton_iteration_mask_and_copy_temperature_on_converged_cells = (
+            compute_newton_iteration_mask_and_copy_temperature_on_converged_cells.with_backend(
                 self._backend
             )
         )
-        self.diagnose_surface_pressure = surface_pressure.diagnose_surface_pressure.with_backend(
+        self.update_temperature_qv_qc_tendencies = update_temperature_qv_qc_tendencies.with_backend(
             self._backend
         )
-        self.diagnose_pressure = pressure.diagnose_pressure.with_backend(self._backend)
-        self.compute_pressure_tendency_after_saturation_adjustment = (
-            compute_pressure_tendency_after_saturation_adjustment.with_backend(self._backend)
-        )
-        self.compute_pressure_ifc_tendency_after_saturation_adjustment = (
-            compute_pressure_ifc_tendency_after_saturation_adjustment.with_backend(self._backend)
+
+    def _determine_horizontal_domains(self):
+        cell_domain = h_grid.domain(dims.CellDim)
+        self._start_cell_nudging = self.grid.start_index(cell_domain(h_grid.Zone.NUDGING))
+        self._end_cell_local = self.grid.start_index(cell_domain(h_grid.Zone.END))
+
+    def _not_converged(self) -> bool:
+        return self._xp.any(
+            self._newton_iteration_mask.ndarray[
+                self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
+            ]
         )
 
     def run(
         self,
         dtime: ta.wpfloat,
-        prognostic_state: prognostics.PrognosticState,
-        diagnostic_state: diagnostics.DiagnosticState,
-        tracer_state: tracers.TracerState,
+        rho: fa.CellKField[ta.wpfloat],
+        temperature: fa.CellKField[ta.wpfloat],
+        qv: fa.CellKField[ta.wpfloat],
+        qc: fa.CellKField[ta.wpfloat],
+        temperature_tendency: fa.CellKField[ta.wpfloat],
+        qv_tendency: fa.CellKField[ta.wpfloat],
+        qc_tendency: fa.CellKField[ta.wpfloat],
     ):
         """
         Adjust saturation at each grid point.
@@ -356,14 +229,24 @@ def run(
         is taken, which is a common approximation and introduces only a small error.
 
         Originally inspired from satad_v_3D of ICON.
+
+        Args:
+            dtime: time step [s]
+            rho: air density [kg m-3]
+            temperature: air temperature [K]
+            qv: specific humidity [kg kg-1]
+            qc: specific cloud water content [kg kg-1]
+            temperature_tendency: air temperature tendency [K s-1]
+            qv_tendency: specific humidity tendency [s-1]
+            qc_tendency: specific cloud water content tendency [s-1]
         """
 
         self.compute_subsaturated_case_and_initialize_newton_iterations(
             self.config.tolerance,
-            diagnostic_state.temperature,
-            tracer_state.qv,
-            tracer_state.qc,
-            prognostic_state.rho,
+            temperature,
+            qv,
+            qc,
+            rho,
             self._subsaturated_mask,
             self._lwdocvd,
             self._temperature1,
@@ -376,103 +259,24 @@ def run(
             offset_provider={},
         )
 
-        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
-        temperature_list = [self._temperature1, self._temperature2]
-        ncurrent, nnext = 0, 1
-        for _ in range(self.config.max_iter):
-            if np.any(
-                self._newton_iteration_mask.ndarray[
-                    self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
-                ]
-            ):
-                self.update_temperature_by_newton_iteration(
-                    diagnostic_state.temperature,
-                    tracer_state.qv,
-                    prognostic_state.rho,
-                    self._newton_iteration_mask,
-                    self._lwdocvd,
-                    temperature_list[nnext],
-                    temperature_list[ncurrent],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
+        temperature_pair = common_utils.TimeStepPair(self._temperature1, self._temperature2)
 
-                self.compute_newton_iteration_mask(
-                    self.config.tolerance,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    self._newton_iteration_mask,
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
+        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
+        num_iter = 0
+        while self._not_converged():
+            if num_iter > self.config.max_iter:
+                raise ConvergenceError(
+                    f"Maximum iteration of saturation adjustment ({self.config.max_iter}) is not enough. The max absolute error is {self._xp.abs(self._temperature1.ndarray - self._temperature2.ndarray).max()} . Please raise max_iter"
                 )
 
-                self.copy_temperature(
-                    self._newton_iteration_mask,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-                ncurrent = (ncurrent + 1) % 2
-                nnext = (nnext + 1) % 2
-            else:
-                break
-        if np.any(
-            self._newton_iteration_mask.ndarray[
-                self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
-            ]
-        ):
-            raise ConvergenceError(
-                f"Maximum iteration of saturation adjustment ({self.config.max_iter}) is not enough. The max absolute error is {np.abs(self.new_temperature1.ndarray - self.new_temperature2.ndarray).max()} . Please raise max_iter"
-            )
-        self.update_temperature_qv_qc_tendencies(
-            dtime,
-            diagnostic_state.temperature,
-            temperature_list[ncurrent],
-            tracer_state.qv,
-            tracer_state.qc,
-            prognostic_state.rho,
-            self._subsaturated_mask,
-            self.temperature_tendency,
-            self.qv_tendency,
-            self.qc_tendency,
-            horizontal_start=self._start_cell_nudging,
-            horizontal_end=self._end_cell_local,
-            vertical_start=gtx.int32(0),
-            vertical_end=self.grid.num_levels,
-            offset_provider={},
-        )
-
-        if self.config.diagnose_variables_from_new_temperature:
-            # TODO (Chia Rui): Consider to merge the gt4py stencils below if scan operator can be called in the intermediate step
-            self.compute_temperature_and_exner_tendencies_after_saturation_adjustment(
-                dtime,
-                tracer_state.qv,
-                tracer_state.qc,
-                tracer_state.qi,
-                tracer_state.qr,
-                tracer_state.qs,
-                tracer_state.qg,
-                self.qv_tendency,
-                self.qc_tendency,
-                temperature_list[ncurrent],
-                diagnostic_state.virtual_temperature,
-                prognostic_state.exner,
-                self.virtual_temperature_tendency,
-                self.exner_tendency,
-                self._new_exner,
-                self._new_virtual_temperature,
-                self._k_field,
-                self.vertical_params.kstart_moist,
+            self.update_temperature_by_newton_iteration(
+                temperature,
+                qv,
+                rho,
+                self._newton_iteration_mask,
+                self._lwdocvd,
+                temperature_pair.next,
+                temperature_pair.current,
                 horizontal_start=self._start_cell_nudging,
                 horizontal_end=self._end_cell_local,
                 vertical_start=gtx.int32(0),
@@ -480,28 +284,11 @@ def run(
                 offset_provider={},
             )
 
-            self.diagnose_surface_pressure(
-                self._new_exner,
-                self._new_virtual_temperature,
-                self.metric_state.ddqz_z_full,
-                self._pressure_ifc,
-                phy_const.CPD_O_RD,
-                phy_const.P0REF,
-                phy_const.GRAV_O_RD,
-                horizontal_start=self._start_cell_nudging,
-                horizontal_end=self._end_cell_local,
-                vertical_start=self.grid.num_levels,
-                vertical_end=gtx.int32(self.grid.num_levels + 1),
-                offset_provider={"Koff": dims.KDim},
-            )
-
-            self.diagnose_pressure(
-                self.metric_state.ddqz_z_full,
-                self._new_virtual_temperature,
-                gtx.as_field((dims.CellDim,), self._pressure_ifc.ndarray[:, -1]),
-                self._pressure,
-                self._pressure_ifc,
-                phy_const.GRAV_O_RD,
+            self.compute_newton_iteration_mask_and_copy_temperature_on_converged_cells(
+                self.config.tolerance,
+                temperature_pair.current,
+                temperature_pair.next,
+                self._newton_iteration_mask,
                 horizontal_start=self._start_cell_nudging,
                 horizontal_end=self._end_cell_local,
                 vertical_start=gtx.int32(0),
@@ -509,29 +296,26 @@ def run(
                 offset_provider={},
             )
 
-            self.compute_pressure_tendency_after_saturation_adjustment(
-                dtime,
-                diagnostic_state.pressure,
-                self._pressure,
-                self.pressure_tendency,
-                horizontal_start=self._start_cell_nudging,
-                horizontal_end=self._end_cell_local,
-                vertical_start=gtx.int32(0),
-                vertical_end=self.grid.num_levels,
-                offset_provider={},
-            )
+            temperature_pair.swap()
+            num_iter = num_iter + 1
 
-            self.compute_pressure_ifc_tendency_after_saturation_adjustment(
-                dtime,
-                diagnostic_state.pressure_ifc,
-                self._pressure_ifc,
-                self.pressure_ifc_tendency,
-                horizontal_start=self._start_cell_nudging,
-                horizontal_end=self._end_cell_local,
-                vertical_start=gtx.int32(0),
-                vertical_end=gtx.int32(self.grid.num_levels + 1),
-                offset_provider={},
-            )
+        self.update_temperature_qv_qc_tendencies(
+            dtime,
+            temperature,
+            temperature_pair.current,
+            qv,
+            qc,
+            rho,
+            self._subsaturated_mask,
+            temperature_tendency,
+            qv_tendency,
+            qc_tendency,
+            horizontal_start=self._start_cell_nudging,
+            horizontal_end=self._end_cell_local,
+            vertical_start=gtx.int32(0),
+            vertical_end=self.grid.num_levels,
+            offset_provider={},
+        )
 
 
 @gtx.field_operator
@@ -549,9 +333,10 @@ def _latent_heat_vaporization(
         latent heat of vaporization.
     """
     return (
-        satpres_const.vaporisation_latent_heat
-        + (1850.0 - satpres_const.spec_heat_cap_water) * (t - satpres_const.tmelt)
-        - satpres_const.rv * t
+        microphy_const.vaporisation_latent_heat
+        + (1850.0 - microphy_const.specific_heat_capacity_for_water)
+        * (t - microphy_const.melting_temperature)
+        - physics_constants.rv * t
     )
 
 
@@ -566,8 +351,10 @@ def _sat_pres_water(t: fa.CellKField[ta.wpfloat]) -> fa.CellKField[ta.wpfloat]:
     Returns:
         saturation water vapour pressure.
     """
-    return satpres_const.tetens_p0 * exp(
-        satpres_const.tetens_aw * (t - satpres_const.tmelt) / (t - satpres_const.tetens_bw)
+    return microphy_const.tetens_p0 * exp(
+        microphy_const.tetens_aw
+        * (t - microphy_const.melting_temperature)
+        / (t - microphy_const.tetens_bw)
     )
 
 
@@ -587,7 +374,7 @@ def _qsat_rho(
     Returns:
         specific humidity at water saturation.
     """
-    return _sat_pres_water(t) / (rho * satpres_const.rv * t)
+    return _sat_pres_water(t) / (rho * physics_constants.rv * t)
 
 
 @gtx.field_operator
@@ -610,7 +397,7 @@ def _dqsatdT_rho(
     Returns:
         partial derivative of the specific humidity at water saturation.
     """
-    beta = satpres_const.tetens_der / (t - satpres_const.tetens_bw) ** 2 - 1.0 / t
+    beta = microphy_const.tetens_der / (t - microphy_const.tetens_bw) ** 2 - 1.0 / t
     return beta * zqsat
 
 
@@ -804,14 +591,14 @@ def _compute_subsaturated_case_and_initialize_newton_iterations(
         mask for Newton iteration case
     """
     temperature_after_all_qc_evaporated = (
-        temperature - _latent_heat_vaporization(temperature) / satpres_const.cvd * qc
+        temperature - _latent_heat_vaporization(temperature) / physics_constants.cvd * qc
     )
 
     # Check, which points will still be subsaturated even after evaporating all cloud water.
     subsaturated_mask = qv + qc <= _qsat_rho(temperature_after_all_qc_evaporated, rho)
 
     # Remains const. during iteration
-    lwdocvd = _latent_heat_vaporization(temperature) / satpres_const.cvd
+    lwdocvd = _latent_heat_vaporization(temperature) / physics_constants.cvd
 
     new_temperature1 = where(
         subsaturated_mask,
@@ -856,276 +643,47 @@ def compute_subsaturated_case_and_initialize_newton_iterations(
 
 
 @gtx.field_operator
-def _compute_newton_iteration_mask(
-    tolerance: ta.wpfloat,
-    temperature_current: fa.CellKField[ta.wpfloat],
-    temperature_next: fa.CellKField[ta.wpfloat],
-) -> fa.CellKField[bool]:
-    return where(abs(temperature_current - temperature_next) > tolerance, True, False)
-
-
-@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def compute_newton_iteration_mask(
+def _compute_newton_iteration_mask_and_copy_temperature_on_converged_cells(
     tolerance: ta.wpfloat,
     temperature_current: fa.CellKField[ta.wpfloat],
     temperature_next: fa.CellKField[ta.wpfloat],
-    newton_iteration_mask: fa.CellKField[bool],
-    horizontal_start: gtx.int32,
-    horizontal_end: gtx.int32,
-    vertical_start: gtx.int32,
-    vertical_end: gtx.int32,
-):
-    _compute_newton_iteration_mask(
-        tolerance,
-        temperature_current,
-        temperature_next,
-        out=newton_iteration_mask,
-        domain={
-            dims.CellDim: (horizontal_start, horizontal_end),
-            dims.KDim: (vertical_start, vertical_end),
-        },
-    )
-
-
-@gtx.field_operator
-def _copy_temperature(
-    newton_iteration_mask: fa.CellKField[bool],
-    temperature_current: fa.CellKField[ta.wpfloat],
-) -> fa.CellKField[ta.wpfloat]:
+) -> tuple[fa.CellKField[bool], fa.CellKField[ta.wpfloat]]:
     """
-    Copy temperature from the current to the new temperature field where the convergence criterion is already met.
+    Compute a mask for the next Newton iteration when the difference between new and old temperature is larger
+    than the tolerance.
+    Then, copy temperature from the current to the new temperature field where the convergence criterion is already met.
     Otherwise, it is zero (the value does not matter because it will be updated in next iteration).
 
     Args:
         newton_iteration_mask: mask for the next Newton iteration to be executed
-        temperature: current temperature [K]
+        temperature_current: temperature at previous Newtion iteration [K]
+        temperature_next: temperature  at current Newtion iteration [K]
     Returns:
         new temperature [K]
     """
-    return where(newton_iteration_mask, 0.0, temperature_current)
+    newton_iteration_mask = where(
+        abs(temperature_current - temperature_next) > tolerance, True, False
+    )
+    new_temperature = where(newton_iteration_mask, 0.0, temperature_current)
+    return newton_iteration_mask, new_temperature
 
 
 @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def copy_temperature(
-    newton_iteration_mask: fa.CellKField[bool],
+def compute_newton_iteration_mask_and_copy_temperature_on_converged_cells(
+    tolerance: ta.wpfloat,
     temperature_current: fa.CellKField[ta.wpfloat],
     temperature_next: fa.CellKField[ta.wpfloat],
+    newton_iteration_mask: fa.CellKField[bool],
     horizontal_start: gtx.int32,
     horizontal_end: gtx.int32,
     vertical_start: gtx.int32,
     vertical_end: gtx.int32,
 ):
-    _copy_temperature(
-        newton_iteration_mask,
+    _compute_newton_iteration_mask_and_copy_temperature_on_converged_cells(
+        tolerance,
         temperature_current,
-        out=temperature_next,
-        domain={
-            dims.CellDim: (horizontal_start, horizontal_end),
-            dims.KDim: (vertical_start, vertical_end),
-        },
-    )
-
-
-@gtx.field_operator
-def _compute_temperature_and_exner_tendencies_after_saturation_adjustment(
-    dtime: ta.wpfloat,
-    qv: fa.CellKField[ta.wpfloat],
-    qc: fa.CellKField[ta.wpfloat],
-    qi: fa.CellKField[ta.wpfloat],
-    qr: fa.CellKField[ta.wpfloat],
-    qs: fa.CellKField[ta.wpfloat],
-    qg: fa.CellKField[ta.wpfloat],
-    qv_tendency: fa.CellKField[ta.wpfloat],
-    qc_tendency: fa.CellKField[ta.wpfloat],
-    temperature: fa.CellKField[ta.wpfloat],
-    virtual_temperature: fa.CellKField[ta.wpfloat],
-    exner: fa.CellKField[ta.wpfloat],
-    k_field: fa.KField[gtx.int32],
-    kstart_moist: gtx.int32,
-) -> tuple[
-    fa.CellKField[ta.wpfloat],
-    fa.CellKField[ta.wpfloat],
-    fa.CellKField[ta.wpfloat],
-    fa.CellKField[ta.wpfloat],
-]:
-    """
-    Update virtual temperature and exner tendencies after saturation adjustment updates temperature, qv , and qc.
-
-    Args:
-        dtime: time step [s]
-        qv: specific humidity [kg kg-1]
-        qc: specific cloud water content [kg kg-1]
-        qi: specific cloud ice content [kg kg-1]
-        qr: specific rain water content [kg kg-1]
-        qs: specific snow content [kg kg-1]
-        qg: specific graupel content [kg kg-1]
-        qv_tendency: specific humidity tendency [kg kg-1 s-1]
-        qc_tendency: specific cloud water content tendency [kg kg-1 s-1]
-        temperature: air temperature [K]
-        virtual_temperature: air virtual temperature [K]
-        exner: exner function
-        k_field: k level indices
-        kstart_moist: starting level for all moisture processes
-    Returns:
-        virtual temperature tendency [K s-1], exner tendency [s-1], new exner, new virtual temperature [K]
-    """
-    qsum = where(
-        k_field < kstart_moist,
-        broadcast(0.0, (dims.CellDim, dims.KDim)),
-        qc + qc_tendency * dtime + qi + qr + qs + qg,
-    )
-
-    new_virtual_temperature = where(
-        k_field < kstart_moist,
-        virtual_temperature,
-        temperature * (1.0 + satpres_const.rv_o_rd_minus_1 * (qv + qv_tendency * dtime) - qsum),
-    )
-    new_exner = where(
-        k_field < kstart_moist,
-        exner,
-        exner
-        * (1.0 + satpres_const.rd_o_cpd * (new_virtual_temperature / virtual_temperature - 1.0)),
-    )
-
-    return (
-        (new_virtual_temperature - virtual_temperature) / dtime,
-        (new_exner - exner) / dtime,
-        new_exner,
-        new_virtual_temperature,
-    )
-
-
-@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def compute_temperature_and_exner_tendencies_after_saturation_adjustment(
-    dtime: ta.wpfloat,
-    qv: fa.CellKField[ta.wpfloat],
-    qc: fa.CellKField[ta.wpfloat],
-    qi: fa.CellKField[ta.wpfloat],
-    qr: fa.CellKField[ta.wpfloat],
-    qs: fa.CellKField[ta.wpfloat],
-    qg: fa.CellKField[ta.wpfloat],
-    qv_tendency: fa.CellKField[ta.wpfloat],
-    qc_tendency: fa.CellKField[ta.wpfloat],
-    temperature: fa.CellKField[ta.wpfloat],
-    virtual_temperature: fa.CellKField[ta.wpfloat],
-    exner: fa.CellKField[ta.wpfloat],
-    virtual_temperature_tendency: fa.CellKField[ta.wpfloat],
-    exner_tendency: fa.CellKField[ta.wpfloat],
-    new_exner: fa.CellKField[ta.wpfloat],
-    new_virtual_temperature: fa.CellKField[ta.wpfloat],
-    k_field: fa.KField[gtx.int32],
-    kstart_moist: gtx.int32,
-    horizontal_start: gtx.int32,
-    horizontal_end: gtx.int32,
-    vertical_start: gtx.int32,
-    vertical_end: gtx.int32,
-):
-    _compute_temperature_and_exner_tendencies_after_saturation_adjustment(
-        dtime,
-        qv,
-        qc,
-        qi,
-        qr,
-        qs,
-        qg,
-        qv_tendency,
-        qc_tendency,
-        temperature,
-        virtual_temperature,
-        exner,
-        k_field,
-        kstart_moist,
-        out=(
-            virtual_temperature_tendency,
-            exner_tendency,
-            new_exner,
-            new_virtual_temperature,
-        ),
-        domain={
-            dims.CellDim: (horizontal_start, horizontal_end),
-            dims.KDim: (vertical_start, vertical_end),
-        },
-    )
-
-
-@gtx.field_operator
-def _compute_pressure_tendency_after_saturation_adjustment(
-    dtime: ta.wpfloat,
-    old_pressure: fa.CellKField[ta.wpfloat],
-    new_pressure: fa.CellKField[ta.wpfloat],
-) -> fa.CellKField[ta.wpfloat]:
-    """
-    Update pressure tendency at full levels after saturation adjustment updates temperature, qv , and qc.
-
-    Args:
-        dtime: time step [s]
-        old_pressure: old air pressure at full levels [Pa]
-        new_pressure: new air pressure at full levels [Pa]
-    Returns:
-        pressure tendency [Pa s-1]
-    """
-
-    return (new_pressure - old_pressure) / dtime
-
-
-@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def compute_pressure_tendency_after_saturation_adjustment(
-    dtime: ta.wpfloat,
-    old_pressure: fa.CellKField[ta.wpfloat],
-    new_pressure: fa.CellKField[ta.wpfloat],
-    pressure_tendency: fa.CellKField[ta.wpfloat],
-    horizontal_start: gtx.int32,
-    horizontal_end: gtx.int32,
-    vertical_start: gtx.int32,
-    vertical_end: gtx.int32,
-):
-    _compute_pressure_tendency_after_saturation_adjustment(
-        dtime,
-        old_pressure,
-        new_pressure,
-        out=pressure_tendency,
-        domain={
-            dims.CellDim: (horizontal_start, horizontal_end),
-            dims.KDim: (vertical_start, vertical_end),
-        },
-    )
-
-
-@gtx.field_operator
-def _compute_pressure_ifc_tendency_after_saturation_adjustment(
-    dtime: ta.wpfloat,
-    old_pressure_ifc: fa.CellKField[ta.wpfloat],
-    new_pressure_ifc: fa.CellKField[ta.wpfloat],
-) -> fa.CellKField[ta.wpfloat]:
-    """
-    Update pressure tendency at half levels (including surface level) after saturation adjustment updates temperature, qv , and qc.
-
-    Args:
-        dtime: time step [s]
-        old_pressure_ifc: old air pressure at interface [Pa]
-        new_pressure_ifc: new air pressure at interface [Pa]
-    Returns:
-        interface pressure tendency [Pa s-1]
-    """
-    return (new_pressure_ifc - old_pressure_ifc) / dtime
-
-
-@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
-def compute_pressure_ifc_tendency_after_saturation_adjustment(
-    dtime: ta.wpfloat,
-    old_pressure_ifc: fa.CellKField[ta.wpfloat],
-    new_pressure_ifc: fa.CellKField[ta.wpfloat],
-    pressure_ifc_tendency: fa.CellKField[ta.wpfloat],
-    horizontal_start: gtx.int32,
-    horizontal_end: gtx.int32,
-    vertical_start: gtx.int32,
-    vertical_end: gtx.int32,
-):
-    _compute_pressure_ifc_tendency_after_saturation_adjustment(
-        dtime,
-        old_pressure_ifc,
-        new_pressure_ifc,
-        out=pressure_ifc_tendency,
+        temperature_next,
+        out=(newton_iteration_mask, temperature_next),
         domain={
             dims.CellDim: (horizontal_start, horizontal_end),
             dims.KDim: (vertical_start, vertical_end),
diff --git a/model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics_tests/test_saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics_tests/test_saturation_adjustment.py
index 85d045bc70..f91c70bccc 100644
--- a/model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics_tests/test_saturation_adjustment.py
+++ b/model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics_tests/test_saturation_adjustment.py
@@ -13,11 +13,6 @@
 )
 from icon4py.model.common import dimension as dims
 from icon4py.model.common.grid import vertical as v_grid
-from icon4py.model.common.states import (
-    diagnostic_state as diagnostics,
-    prognostic_state as prognostics,
-    tracer_state as tracers,
-)
 from icon4py.model.common.utils import data_allocation as data_alloc
 from icon4py.model.testing import datatest_utils as dt_utils, helpers
 
@@ -31,13 +26,10 @@
 @pytest.mark.parametrize(
     "date", ["2008-09-01T01:59:48.000", "2008-09-01T01:59:52.000", "2008-09-01T01:59:56.000"]
 )
-@pytest.mark.parametrize(
-    "location, diagnose_values", [("nwp-gscp-interface", False), ("interface-nwp", True)]
-)
+@pytest.mark.parametrize("location", ["nwp-gscp-interface", "interface-nwp"])
 def test_saturation_adjustement(
     experiment,
     location,
-    diagnose_values,
     model_top_height,
     damping_height,
     stretch_factor,
@@ -46,7 +38,6 @@ def test_saturation_adjustement(
     grid_savepoint,
     metrics_savepoint,
     icon_grid,
-    lowest_layer_thickness,
     backend,
 ):
     satad_init = data_provider.from_savepoint_satad_init(location=location, date=date)
@@ -55,7 +46,6 @@ def test_saturation_adjustement(
     config = satad.SaturationAdjustmentConfig(
         tolerance=1e-3,
         max_iter=10,
-        diagnose_variables_from_new_temperature=diagnose_values,
     )
     dtime = 2.0
 
@@ -66,6 +56,11 @@ def test_saturation_adjustement(
         vct_b=grid_savepoint.vct_b(),
         _min_index_flat_horizontal_grad_pressure=grid_savepoint.nflat_gradp(),
     )
+    temperature_tendency = data_alloc.zero_field(
+        icon_grid, dims.CellDim, dims.KDim, backend=backend
+    )
+    qv_tendency = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, backend=backend)
+    qc_tendency = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, backend=backend)
 
     metric_state = satad.MetricStateSaturationAdjustment(
         ddqz_z_full=metrics_savepoint.ddqz_z_full()
@@ -78,67 +73,27 @@ def test_saturation_adjustement(
         vertical_params=vertical_params,
         backend=backend,
     )
-    tracer_state = tracers.TracerState(
-        qv=satad_init.qv(),
-        qc=satad_init.qc(),
-        qr=satad_init.qr(),
-        qi=satad_init.qi(),
-        qs=satad_init.qs(),
-        qg=satad_init.qg(),
-    )
-    prognostic_state = prognostics.PrognosticState(
-        rho=satad_init.rho(),
-        vn=data_alloc.zero_field(icon_grid, dims.EdgeDim, dims.KDim, dtype=float),
-        w=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
-        exner=satad_init.exner(),
-        theta_v=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
-    )
-    diagnostic_state = diagnostics.DiagnosticState(
-        temperature=satad_init.temperature(),
-        virtual_temperature=satad_init.virtual_temperature(),
-        pressure=satad_init.pressure(),
-        pressure_ifc=satad_init.pressure_ifc(),
-        u=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
-        v=data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, dtype=float),
-    )
 
-    # WHEN
+    qv = satad_init.qv()
+    qc = satad_init.qc()
+    rho = satad_init.rho()
+    temperature = satad_init.temperature()
+
     # run saturation adjustment
     saturation_adjustment.run(
         dtime=dtime,
-        prognostic_state=prognostic_state,
-        diagnostic_state=diagnostic_state,
-        tracer_state=tracer_state,
+        rho=rho,
+        temperature=temperature,
+        qv=qv,
+        qc=qc,
+        temperature_tendency=temperature_tendency,
+        qv_tendency=qv_tendency,
+        qc_tendency=qc_tendency,
     )
 
-    # apply tendencies
-
-    updated_qv = tracer_state.qv.asnumpy() + saturation_adjustment.qv_tendency.asnumpy() * dtime
-    updated_qc = tracer_state.qc.asnumpy() + saturation_adjustment.qc_tendency.asnumpy() * dtime
-    updated_temperature = (
-        diagnostic_state.temperature.asnumpy()
-        + saturation_adjustment.temperature_tendency.asnumpy() * dtime
-    )
-    updated_exner = (
-        prognostic_state.exner.asnumpy() + saturation_adjustment.exner_tendency.asnumpy() * dtime
-    )
-    if diagnose_values:
-        pressure = (
-            diagnostic_state.pressure.asnumpy()
-            + saturation_adjustment.pressure_tendency.asnumpy() * dtime
-        )
-        pressure_ifc = (
-            diagnostic_state.pressure_ifc.asnumpy()
-            + saturation_adjustment.pressure_ifc_tendency.asnumpy() * dtime
-        )
-        virtual_temperature = (
-            diagnostic_state.virtual_temperature.asnumpy()
-            + saturation_adjustment.virtual_temperature_tendency.asnumpy() * dtime
-        )
-    else:
-        pressure = diagnostic_state.pressure.asnumpy()
-        pressure_ifc = diagnostic_state.pressure_ifc.asnumpy()
-        virtual_temperature = diagnostic_state.virtual_temperature.asnumpy()
+    updated_qv = qv.asnumpy() + qv_tendency.asnumpy() * dtime
+    updated_qc = qc.asnumpy() + qc_tendency.asnumpy() * dtime
+    updated_temperature = temperature.asnumpy() + temperature_tendency.asnumpy() * dtime
 
     assert helpers.dallclose(
         updated_qv,
@@ -155,23 +110,3 @@ def test_saturation_adjustement(
         satad_exit.temperature().asnumpy(),
         atol=1.0e-13,
     )
-    assert helpers.dallclose(
-        virtual_temperature,
-        satad_exit.virtual_temperature().asnumpy(),
-        atol=1.0e-13,
-    )
-    assert helpers.dallclose(
-        updated_exner,
-        satad_exit.exner().asnumpy(),
-        atol=1.0e-13,
-    )
-    assert helpers.dallclose(
-        pressure,
-        satad_exit.pressure().asnumpy(),
-        atol=1.0e-13,
-    )
-    assert helpers.dallclose(
-        pressure_ifc,
-        satad_exit.pressure_ifc().asnumpy(),
-        atol=1.0e-13,
-    )
diff --git a/model/common/src/icon4py/model/common/constants.py b/model/common/src/icon4py/model/common/constants.py
index 44db67c89c..35d5210913 100644
--- a/model/common/src/icon4py/model/common/constants.py
+++ b/model/common/src/icon4py/model/common/constants.py
@@ -10,64 +10,66 @@
 import sys
 from typing import Final
 
-from icon4py.model.common.type_alias import wpfloat
+from gt4py.eve import utils as eve_utils
+
+from icon4py.model.common import type_alias as ta
 
 
 #: Gas constant for dry air [J/K/kg], called 'rd' in ICON (mo_physical_constants.f90),
 #: see https://glossary.ametsoc.org/wiki/Gas_constant.
-GAS_CONSTANT_DRY_AIR: Final[wpfloat] = 287.04
-RD: Final[wpfloat] = GAS_CONSTANT_DRY_AIR
+GAS_CONSTANT_DRY_AIR: Final[ta.wpfloat] = 287.04
+RD: Final[ta.wpfloat] = GAS_CONSTANT_DRY_AIR
 
 #: Specific heat at constant pressure [J/K/kg]
-SPECIFIC_HEAT_CONSTANT_PRESSURE: Final[wpfloat] = 1004.64
+SPECIFIC_HEAT_CONSTANT_PRESSURE: Final[ta.wpfloat] = 1004.64
 CPD = SPECIFIC_HEAT_CONSTANT_PRESSURE
 
 #: [J/K/kg] specific heat at constant volume
-SPECIFIC_HEAT_CONSTANT_VOLUME: Final[wpfloat] = CPD - RD
-CVD: Final[wpfloat] = SPECIFIC_HEAT_CONSTANT_VOLUME
-CVD_O_RD: Final[wpfloat] = CVD / RD
-RD_O_CPD: Final[wpfloat] = RD / CPD
-CPD_O_RD: Final[wpfloat] = CPD / RD
+SPECIFIC_HEAT_CONSTANT_VOLUME: Final[ta.wpfloat] = CPD - RD
+CVD: Final[ta.wpfloat] = SPECIFIC_HEAT_CONSTANT_VOLUME
+CVD_O_RD: Final[ta.wpfloat] = CVD / RD
+RD_O_CPD: Final[ta.wpfloat] = RD / CPD
+CPD_O_RD: Final[ta.wpfloat] = CPD / RD
 
 #: Gas constant for water vapor [J/K/kg], rv in ICON.
-GAS_CONSTANT_WATER_VAPOR: Final[wpfloat] = 461.51
-RV: Final[wpfloat] = GAS_CONSTANT_WATER_VAPOR
+GAS_CONSTANT_WATER_VAPOR: Final[ta.wpfloat] = 461.51
+RV: Final[ta.wpfloat] = GAS_CONSTANT_WATER_VAPOR
 
 #: RV/RD - 1, tvmpc1 in ICON.
-RV_O_RD_MINUS_1: Final[wpfloat] = GAS_CONSTANT_WATER_VAPOR / GAS_CONSTANT_DRY_AIR - 1.0
-TVMPC1: Final[wpfloat] = RV_O_RD_MINUS_1
+RV_O_RD_MINUS_1: Final[ta.wpfloat] = GAS_CONSTANT_WATER_VAPOR / GAS_CONSTANT_DRY_AIR - 1.0
+TVMPC1: Final[ta.wpfloat] = RV_O_RD_MINUS_1
 
 #: Av. gravitational acceleration [m/s^2]
-GRAVITATIONAL_ACCELERATION: Final[wpfloat] = 9.80665
-GRAV: Final[wpfloat] = GRAVITATIONAL_ACCELERATION
-GRAV_O_RD: Final[wpfloat] = GRAV / RD
+GRAVITATIONAL_ACCELERATION: Final[ta.wpfloat] = 9.80665
+GRAV: Final[ta.wpfloat] = GRAVITATIONAL_ACCELERATION
+GRAV_O_RD: Final[ta.wpfloat] = GRAV / RD
 
 #: reference pressure for Exner function [Pa]
-REFERENCE_PRESSURE: Final[wpfloat] = 100000.0
-P0REF: Final[wpfloat] = REFERENCE_PRESSURE
+REFERENCE_PRESSURE: Final[ta.wpfloat] = 100000.0
+P0REF: Final[ta.wpfloat] = REFERENCE_PRESSURE
 
 #: sea level pressure [Pa]
-SEAL_LEVEL_PRESSURE: Final[wpfloat] = 101325.0
-P0SL_BG: Final[wpfloat] = SEAL_LEVEL_PRESSURE
+SEAL_LEVEL_PRESSURE: Final[ta.wpfloat] = 101325.0
+P0SL_BG: Final[ta.wpfloat] = SEAL_LEVEL_PRESSURE
 
 # average earth radius in [m]
 EARTH_RADIUS: Final[float] = 6.371229e6
 
 #: Earth angular velocity [rad/s]
-EARTH_ANGULAR_VELOCITY: Final[wpfloat] = 7.29212e-5
+EARTH_ANGULAR_VELOCITY: Final[ta.wpfloat] = 7.29212e-5
 
 #: sea level temperature for reference atmosphere [K]
-SEA_LEVEL_TEMPERATURE: Final[wpfloat] = 288.15
-T0SL_BG: Final[wpfloat] = SEA_LEVEL_TEMPERATURE
+SEA_LEVEL_TEMPERATURE: Final[ta.wpfloat] = 288.15
+T0SL_BG: Final[ta.wpfloat] = SEA_LEVEL_TEMPERATURE
 
 #: difference between sea level temperature and asymptotic stratospheric temperature
-DELTA_TEMPERATURE: Final[wpfloat] = 75.0
-DEL_T_BG: Final[wpfloat] = DELTA_TEMPERATURE
+DELTA_TEMPERATURE: Final[ta.wpfloat] = 75.0
+DEL_T_BG: Final[ta.wpfloat] = DELTA_TEMPERATURE
 
 #: height scale for reference atmosphere [m], defined  in mo_vertical_grid
 #: scale height [m]
 HEIGHT_SCALE_FOR_REFERENCE_ATMOSPHERE = 10000.0
-_H_SCAL_BG: Final[wpfloat] = HEIGHT_SCALE_FOR_REFERENCE_ATMOSPHERE
+_H_SCAL_BG: Final[ta.wpfloat] = HEIGHT_SCALE_FOR_REFERENCE_ATMOSPHERE
 
 # Math constants
 DBL_EPS = sys.float_info.epsilon  # EPSILON(1._wp)
@@ -78,7 +80,7 @@
 DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Final[float] = 5.0
 
 #: average earth radius in [m]
-EARTH_RADIUS: Final[wpfloat] = 6.371229e6
+EARTH_RADIUS: Final[ta.wpfloat] = 6.371229e6
 
 
 class RayleighType(enum.IntEnum):
@@ -86,3 +88,15 @@ class RayleighType(enum.IntEnum):
     CLASSIC = 1
     #: Klemp (2008) type Rayleigh damping
     KLEMP = 2
+
+
+class PhysicsConstants(eve_utils.FrozenNamespace[ta.wpfloat]):
+    rd = RD
+    rv = RV
+    cpd = CPD
+    cvd = CVD
+    rv_o_rd_minus_1 = RV_O_RD_MINUS_1
+    rd_o_cpd = RD_O_CPD
+    grav_o_rd = GRAV_O_RD
+    p0ref = P0REF
+    cpd_o_rd = CPD_O_RD
diff --git a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/calculate_tendency.py b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/calculate_tendency.py
new file mode 100644
index 0000000000..e1dd471a4e
--- /dev/null
+++ b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/calculate_tendency.py
@@ -0,0 +1,189 @@
+# 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 gt4py import next as gtx
+
+from icon4py.model.common import (
+    constants as phy_const,
+    dimension as dims,
+    field_type_aliases as fa,
+    type_alias as ta,
+)
+from icon4py.model.common.type_alias import wpfloat
+
+
+physics_constants: Final = phy_const.PhysicsConstants()
+
+
+@gtx.field_operator
+def _calculate_virtual_temperature_tendency(
+    dtime: ta.wpfloat,
+    qv: fa.CellKField[ta.wpfloat],
+    qc: fa.CellKField[ta.wpfloat],
+    qi: fa.CellKField[ta.wpfloat],
+    qr: fa.CellKField[ta.wpfloat],
+    qs: fa.CellKField[ta.wpfloat],
+    qg: fa.CellKField[ta.wpfloat],
+    temperature: fa.CellKField[ta.wpfloat],
+    virtual_temperature: fa.CellKField[ta.wpfloat],
+) -> fa.CellKField[ta.wpfloat]:
+    """
+    Update virtual temperature tendency.
+
+    Args:
+        dtime: time step [s]
+        qv: specific humidity [kg kg-1]
+        qc: specific cloud water content [kg kg-1]
+        qi: specific cloud ice content [kg kg-1]
+        qr: specific rain water content [kg kg-1]
+        qs: specific snow content [kg kg-1]
+        qg: specific graupel content [kg kg-1]
+        temperature: air temperature [K]
+        virtual_temperature: air virtual temperature [K]
+    Returns:
+        virtual temperature tendency [K s-1], exner tendency [s-1], new exner, new virtual temperature [K]
+    """
+    qsum = qc + qi + qr + qs + qg
+
+    new_virtual_temperature = temperature * (
+        wpfloat("1.0") + physics_constants.rv_o_rd_minus_1 * qv - qsum
+    )
+
+    return (new_virtual_temperature - virtual_temperature) / dtime
+
+
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
+def calculate_virtual_temperature_tendency(
+    dtime: ta.wpfloat,
+    qv: fa.CellKField[ta.wpfloat],
+    qc: fa.CellKField[ta.wpfloat],
+    qi: fa.CellKField[ta.wpfloat],
+    qr: fa.CellKField[ta.wpfloat],
+    qs: fa.CellKField[ta.wpfloat],
+    qg: fa.CellKField[ta.wpfloat],
+    temperature: fa.CellKField[ta.wpfloat],
+    virtual_temperature: fa.CellKField[ta.wpfloat],
+    virtual_temperature_tendency: fa.CellKField[ta.wpfloat],
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
+):
+    _calculate_virtual_temperature_tendency(
+        dtime,
+        qv,
+        qc,
+        qi,
+        qr,
+        qs,
+        qg,
+        temperature,
+        virtual_temperature,
+        out=virtual_temperature_tendency,
+        domain={
+            dims.CellDim: (horizontal_start, horizontal_end),
+            dims.KDim: (vertical_start, vertical_end),
+        },
+    )
+
+
+@gtx.field_operator
+def _calculate_exner_tendency(
+    dtime: ta.wpfloat,
+    virtual_temperature: fa.CellKField[ta.wpfloat],
+    virtual_temperature_tendency: fa.CellKField[ta.wpfloat],
+    exner: fa.CellKField[ta.wpfloat],
+) -> fa.CellKField[ta.wpfloat]:
+    """
+    Update exner tendency.
+
+    Args:
+        dtime: time step [s]
+        virtual_temperature: air virtual temperature [K]
+        virtual_temperature_tendency: air virtual temperature tendency [K s-1]
+        exner: exner function
+    Returns:
+        exner tendency [s-1]
+    """
+
+    new_exner = exner * (
+        wpfloat("1.0")
+        + physics_constants.rd_o_cpd * virtual_temperature_tendency / virtual_temperature * dtime
+    )
+
+    return (new_exner - exner) / dtime
+
+
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
+def calculate_exner_tendency(
+    dtime: ta.wpfloat,
+    virtual_temperature: fa.CellKField[ta.wpfloat],
+    virtual_temperature_tendency: fa.CellKField[ta.wpfloat],
+    exner: fa.CellKField[ta.wpfloat],
+    exner_tendency: fa.CellKField[ta.wpfloat],
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
+):
+    _calculate_exner_tendency(
+        dtime,
+        virtual_temperature,
+        virtual_temperature_tendency,
+        exner,
+        out=exner_tendency,
+        domain={
+            dims.CellDim: (horizontal_start, horizontal_end),
+            dims.KDim: (vertical_start, vertical_end),
+        },
+    )
+
+
+@gtx.field_operator
+def _calculate_cell_kdim_field_tendency(
+    dtime: ta.wpfloat,
+    old_field: fa.CellKField[ta.wpfloat],
+    new_field: fa.CellKField[ta.wpfloat],
+) -> fa.CellKField[ta.wpfloat]:
+    """
+    Update tendency of Cell-K Dim field.
+
+    Args:
+        dtime: time step [s]
+        old_field: any old Cell-K Dim field [unit]
+        new_field: any new Cell-K Dim field [unit]
+    Returns:
+        tendency [unit s-1]
+    """
+
+    return (new_field - old_field) / dtime
+
+
+@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
+def calculate_cell_kdim_field_tendency(
+    dtime: ta.wpfloat,
+    old_field: fa.CellKField[ta.wpfloat],
+    new_field: fa.CellKField[ta.wpfloat],
+    tendency: fa.CellKField[ta.wpfloat],
+    horizontal_start: gtx.int32,
+    horizontal_end: gtx.int32,
+    vertical_start: gtx.int32,
+    vertical_end: gtx.int32,
+):
+    _calculate_cell_kdim_field_tendency(
+        dtime,
+        old_field,
+        new_field,
+        out=tendency,
+        domain={
+            dims.CellDim: (horizontal_start, horizontal_end),
+            dims.KDim: (vertical_start, vertical_end),
+        },
+    )
diff --git a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_pressure.py b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_pressure.py
index dc19184747..d0ea3bffba 100644
--- a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_pressure.py
+++ b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_pressure.py
@@ -5,12 +5,22 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from typing import Final
+
 import gt4py.next as gtx
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program, scan_operator
 from gt4py.next.ffront.fbuiltins import exp, sqrt
 
-from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+from icon4py.model.common import (
+    constants as phy_const,
+    dimension as dims,
+    field_type_aliases as fa,
+    type_alias as ta,
+)
+
+
+physics_constants: Final = phy_const.PhysicsConstants()
 
 
 @scan_operator(axis=dims.KDim, forward=False, init=(0.0, 0.0, True))
@@ -19,12 +29,11 @@ def _scan_pressure(
     ddqz_z_full: ta.wpfloat,
     virtual_temperature: ta.wpfloat,
     surface_pressure: ta.wpfloat,
-    grav_o_rd: ta.wpfloat,
 ):
     pressure_interface = (
-        surface_pressure * exp(-grav_o_rd * ddqz_z_full / virtual_temperature)
+        surface_pressure * exp(-physics_constants.grav_o_rd * ddqz_z_full / virtual_temperature)
         if state[2]
-        else state[1] * exp(-grav_o_rd * ddqz_z_full / virtual_temperature)
+        else state[1] * exp(-physics_constants.grav_o_rd * ddqz_z_full / virtual_temperature)
     )
     pressure = (
         sqrt(surface_pressure * pressure_interface)
@@ -39,7 +48,6 @@ def _diagnose_pressure(
     ddqz_z_full: fa.CellKField[ta.wpfloat],
     virtual_temperature: fa.CellKField[ta.wpfloat],
     surface_pressure: gtx.Field[gtx.Dims[dims.CellDim], ta.wpfloat],
-    grav_o_rd: ta.wpfloat,
 ) -> tuple[fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat]]:
     """
     Update pressure by assuming hydrostatic balance (dp/dz = -rho g = p g / Rd / Tv).
@@ -49,13 +57,10 @@ def _diagnose_pressure(
         ddqz_z_full: vertical grid spacing at full levels [m]
         virtual_temperature: air virtual temperature [K]
         surface_pressure: surface air pressure [Pa]
-        grav_o_rd: gravitational constant / dry air constant [K kg m/s2/J]
     Returns:
         pressure at full levels, pressure at half levels (excluding surface level)
     """
-    pressure, pressure_ifc, _ = _scan_pressure(
-        ddqz_z_full, virtual_temperature, surface_pressure, grav_o_rd
-    )
+    pressure, pressure_ifc, _ = _scan_pressure(ddqz_z_full, virtual_temperature, surface_pressure)
     return pressure, pressure_ifc
 
 
@@ -66,7 +71,6 @@ def diagnose_pressure(
     surface_pressure: fa.CellField[ta.wpfloat],
     pressure: fa.CellKField[ta.wpfloat],
     pressure_ifc: fa.CellKField[ta.wpfloat],
-    grav_o_rd: ta.wpfloat,
     horizontal_start: gtx.int32,
     horizontal_end: gtx.int32,
     vertical_start: gtx.int32,
@@ -76,7 +80,6 @@ def diagnose_pressure(
         ddqz_z_full,
         virtual_temperature,
         surface_pressure,
-        grav_o_rd,
         out=(pressure, pressure_ifc),
         domain={
             dims.CellDim: (horizontal_start, horizontal_end),
diff --git a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py
index 4bd37c749e..109f5f562e 100644
--- a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py
+++ b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_surface_pressure.py
@@ -5,27 +5,34 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from typing import Final
+
 import gt4py.next as gtx
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
 from gt4py.next.ffront.fbuiltins import exp, log
 
-from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+from icon4py.model.common import (
+    constants as phy_const,
+    dimension as dims,
+    field_type_aliases as fa,
+    type_alias as ta,
+)
 from icon4py.model.common.dimension import Koff
 
 
+physics_constants: Final = phy_const.PhysicsConstants()
+
+
 @field_operator
 def _diagnose_surface_pressure(
     exner: fa.CellKField[ta.wpfloat],
     virtual_temperature: fa.CellKField[ta.wpfloat],
     ddqz_z_full: fa.CellKField[ta.wpfloat],
-    cpd_o_rd: ta.wpfloat,
-    p0ref: ta.wpfloat,
-    grav_o_rd: ta.wpfloat,
 ) -> fa.CellKField[ta.wpfloat]:
-    surface_pressure = p0ref * exp(
-        cpd_o_rd * log(exner(Koff[-3]))
-        + grav_o_rd
+    surface_pressure = physics_constants.p0ref * exp(
+        physics_constants.cpd_o_rd * log(exner(Koff[-3]))
+        + physics_constants.grav_o_rd
         * (
             ddqz_z_full(Koff[-1]) / virtual_temperature(Koff[-1])
             + ddqz_z_full(Koff[-2]) / virtual_temperature(Koff[-2])
@@ -41,9 +48,6 @@ def diagnose_surface_pressure(
     virtual_temperature: fa.CellKField[ta.wpfloat],
     ddqz_z_full: fa.CellKField[ta.wpfloat],
     surface_pressure: fa.CellKField[ta.wpfloat],
-    cpd_o_rd: ta.wpfloat,
-    p0ref: ta.wpfloat,
-    grav_o_rd: ta.wpfloat,
     horizontal_start: gtx.int32,
     horizontal_end: gtx.int32,
     vertical_start: gtx.int32,
@@ -53,9 +57,6 @@ def diagnose_surface_pressure(
         exner,
         virtual_temperature,
         ddqz_z_full,
-        cpd_o_rd,
-        p0ref,
-        grav_o_rd,
         out=surface_pressure,
         domain={
             dims.CellDim: (horizontal_start, horizontal_end),
diff --git a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py
index 247e23308c..efc58a6695 100644
--- a/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py
+++ b/model/common/src/icon4py/model/common/diagnostic_calculations/stencils/diagnose_temperature.py
@@ -5,9 +5,19 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from typing import Final
+
 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 import (
+    constants as phy_const,
+    dimension as dims,
+    field_type_aliases as fa,
+    type_alias as ta,
+)
+
+
+physics_constants: Final = phy_const.PhysicsConstants()
 
 
 @gtx.field_operator
@@ -20,11 +30,10 @@ def _diagnose_virtual_temperature_and_temperature(
     qg: fa.CellKField[ta.wpfloat],
     theta_v: fa.CellKField[ta.wpfloat],
     exner: fa.CellKField[ta.wpfloat],
-    rv_o_rd_minus1: ta.wpfloat,
 ) -> tuple[fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat]]:
     qsum = qc + qi + qr + qs + qg
     virtual_temperature = theta_v * exner
-    temperature = virtual_temperature / (1.0 + rv_o_rd_minus1 * qv - qsum)
+    temperature = virtual_temperature / (1.0 + physics_constants.rv_o_rd_minus_1 * qv - qsum)
     return virtual_temperature, temperature
 
 
@@ -41,7 +50,6 @@ def diagnose_virtual_temperature_and_temperature(
     exner: fa.CellKField[ta.wpfloat],
     virtual_temperature: fa.CellKField[ta.wpfloat],
     temperature: fa.CellKField[ta.wpfloat],
-    rv_o_rd_minus1: ta.wpfloat,
     horizontal_start: gtx.int32,
     horizontal_end: gtx.int32,
     vertical_start: gtx.int32,
@@ -56,7 +64,6 @@ def diagnose_virtual_temperature_and_temperature(
         qg,
         theta_v,
         exner,
-        rv_o_rd_minus1,
         out=(virtual_temperature, temperature),
         domain={
             dims.CellDim: (horizontal_start, horizontal_end),
diff --git a/model/common/tests/diagnostic_calculations_tests/test_diagnostic_calculations.py b/model/common/tests/diagnostic_calculations_tests/test_diagnostic_calculations.py
index 072b4ce227..4ae99913c6 100644
--- a/model/common/tests/diagnostic_calculations_tests/test_diagnostic_calculations.py
+++ b/model/common/tests/diagnostic_calculations_tests/test_diagnostic_calculations.py
@@ -6,18 +6,25 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import gt4py.next as gtx
 import pytest
 
 import icon4py.model.common.grid.horizontal as h_grid
-from icon4py.model.common import constants as phy_const, dimension as dims
+from icon4py.model.common import dimension as dims
 from icon4py.model.common.diagnostic_calculations.stencils import (
+    calculate_tendency,
     diagnose_pressure,
     diagnose_surface_pressure,
     diagnose_temperature,
 )
+from icon4py.model.common.grid import 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_state as tracers,
+)
 from icon4py.model.common.utils import data_allocation as data_alloc
 from icon4py.model.testing import datatest_utils as dt_utils, helpers
 
@@ -67,7 +74,6 @@ def test_diagnose_temperature(
         exner=exner,
         virtual_temperature=virtual_temperature,
         temperature=temperature,
-        rv_o_rd_minus1=phy_const.RV_O_RD_MINUS_1,
         horizontal_start=0,
         horizontal_end=icon_grid.end_index(h_grid.domain(dims.CellDim)(h_grid.Zone.END)),
         vertical_start=0,
@@ -173,9 +179,6 @@ def test_diagnose_surface_pressure(
         virtual_temperature=virtual_temperature,
         ddqz_z_full=ddqz_z_full,
         surface_pressure=surface_pressure,
-        cpd_o_rd=phy_const.CPD_O_RD,
-        p0ref=phy_const.P0REF,
-        grav_o_rd=phy_const.GRAV_O_RD,
         horizontal_start=0,
         horizontal_end=icon_grid.end_index(cell_domain(h_grid.Zone.END)),
         vertical_start=icon_grid.num_levels,
@@ -223,7 +226,6 @@ def test_diagnose_pressure(experiment, data_provider, icon_grid, backend, metric
         surface_pressure,
         pressure,
         pressure_ifc,
-        phy_const.GRAV_O_RD,
         horizontal_start=0,
         horizontal_end=icon_grid.end_index(cell_domain(h_grid.Zone.END)),
         vertical_start=0,
@@ -237,3 +239,150 @@ def test_diagnose_pressure(experiment, data_provider, icon_grid, backend, metric
         pressure_ref,
         pressure.asnumpy(),
     )
+
+
+@pytest.mark.parametrize(
+    "experiment, model_top_height, damping_height, stretch_factor",
+    [
+        (dt_utils.WEISMAN_KLEMP_EXPERIMENT, 30000.0, 8000.0, 0.85),
+    ],
+)
+@pytest.mark.parametrize(
+    "date", ["2008-09-01T01:59:48.000", "2008-09-01T01:59:52.000", "2008-09-01T01:59:56.000"]
+)
+@pytest.mark.parametrize("location", [("interface-nwp")])
+def test_diagnostic_update_after_saturation_adjustement(
+    experiment,
+    location,
+    model_top_height,
+    damping_height,
+    stretch_factor,
+    date,
+    data_provider,
+    grid_savepoint,
+    metrics_savepoint,
+    icon_grid,
+    backend,
+):
+    satad_init = data_provider.from_savepoint_satad_init(location=location, date=date)
+    satad_exit = data_provider.from_savepoint_satad_exit(location=location, date=date)
+
+    dtime = 2.0
+
+    vertical_config = v_grid.VerticalGridConfig(icon_grid.num_levels)
+    vertical_params = v_grid.VerticalGrid(
+        config=vertical_config,
+        vct_a=grid_savepoint.vct_a(),
+        vct_b=grid_savepoint.vct_b(),
+        _min_index_flat_horizontal_grad_pressure=grid_savepoint.nflat_gradp(),
+    )
+    virtual_temperature_tendency = data_alloc.zero_field(
+        icon_grid, dims.CellDim, dims.KDim, backend=backend
+    )
+    exner_tendency = data_alloc.zero_field(icon_grid, dims.CellDim, dims.KDim, backend=backend)
+
+    tracer_state = tracers.TracerState(
+        qv=satad_exit.qv(),
+        qc=satad_exit.qc(),
+        qr=satad_init.qr(),
+        qi=satad_init.qi(),
+        qs=satad_init.qs(),
+        qg=satad_init.qg(),
+    )
+    exner = satad_init.exner()
+
+    diagnostic_state = diagnostics.DiagnosticState(
+        temperature=satad_exit.temperature(),
+        virtual_temperature=satad_init.virtual_temperature(),
+        pressure=satad_init.pressure(),
+        pressure_ifc=satad_init.pressure_ifc(),
+        u=None,
+        v=None,
+    )
+
+    cell_domain = h_grid.domain(dims.CellDim)
+    start_cell_nudging = icon_grid.start_index(cell_domain(h_grid.Zone.NUDGING))
+    end_cell_local = icon_grid.start_index(cell_domain(h_grid.Zone.END))
+    calculate_tendency.calculate_virtual_temperature_tendency.with_backend(backend)(
+        dtime=dtime,
+        qv=tracer_state.qv,
+        qc=tracer_state.qc,
+        qi=tracer_state.qi,
+        qr=tracer_state.qr,
+        qs=tracer_state.qs,
+        qg=tracer_state.qg,
+        temperature=diagnostic_state.temperature,
+        virtual_temperature=diagnostic_state.virtual_temperature,
+        virtual_temperature_tendency=virtual_temperature_tendency,
+        horizontal_start=start_cell_nudging,
+        horizontal_end=end_cell_local,
+        vertical_start=vertical_params.kstart_moist,
+        vertical_end=icon_grid.num_levels,
+        offset_provider={},
+    )
+
+    updated_virtual_temperature = (
+        diagnostic_state.virtual_temperature.asnumpy()
+        + virtual_temperature_tendency.asnumpy() * dtime
+    )
+
+    calculate_tendency.calculate_exner_tendency.with_backend(backend)(
+        dtime=dtime,
+        virtual_temperature=diagnostic_state.virtual_temperature,
+        virtual_temperature_tendency=virtual_temperature_tendency,
+        exner=exner,
+        exner_tendency=exner_tendency,
+        horizontal_start=start_cell_nudging,
+        horizontal_end=end_cell_local,
+        vertical_start=vertical_params.kstart_moist,
+        vertical_end=icon_grid.num_levels,
+        offset_provider={},
+    )
+
+    updated_exner = exner.asnumpy() + exner_tendency.asnumpy() * dtime
+
+    diagnose_surface_pressure.diagnose_surface_pressure.with_backend(backend)(
+        gtx.as_field((dims.CellDim, dims.KDim), updated_exner, allocator=backend),
+        gtx.as_field((dims.CellDim, dims.KDim), updated_virtual_temperature, allocator=backend),
+        metrics_savepoint.ddqz_z_full(),
+        diagnostic_state.pressure_ifc,
+        horizontal_start=start_cell_nudging,
+        horizontal_end=end_cell_local,
+        vertical_start=icon_grid.num_levels,
+        vertical_end=gtx.int32(icon_grid.num_levels + 1),
+        offset_provider={"Koff": dims.KDim},
+    )
+
+    diagnose_pressure.diagnose_pressure.with_backend(backend)(
+        metrics_savepoint.ddqz_z_full(),
+        gtx.as_field((dims.CellDim, dims.KDim), updated_virtual_temperature, allocator=backend),
+        diagnostic_state.surface_pressure,
+        diagnostic_state.pressure,
+        diagnostic_state.pressure_ifc,
+        horizontal_start=start_cell_nudging,
+        horizontal_end=end_cell_local,
+        vertical_start=gtx.int32(0),
+        vertical_end=icon_grid.num_levels,
+        offset_provider={},
+    )
+
+    assert helpers.dallclose(
+        updated_virtual_temperature,
+        satad_exit.virtual_temperature().asnumpy(),
+        atol=1.0e-13,
+    )
+    assert helpers.dallclose(
+        updated_exner,
+        satad_exit.exner().asnumpy(),
+        atol=1.0e-13,
+    )
+    assert helpers.dallclose(
+        diagnostic_state.pressure.asnumpy(),
+        satad_exit.pressure().asnumpy(),
+        atol=1.0e-13,
+    )
+    assert helpers.dallclose(
+        diagnostic_state.pressure_ifc.asnumpy(),
+        satad_exit.pressure_ifc().asnumpy(),
+        atol=1.0e-13,
+    )
diff --git a/model/common/tests/stencil_tests/test_diagnose_pressure.py b/model/common/tests/stencil_tests/test_diagnose_pressure.py
index 1769bb33b8..cc5cc3114d 100644
--- a/model/common/tests/stencil_tests/test_diagnose_pressure.py
+++ b/model/common/tests/stencil_tests/test_diagnose_pressure.py
@@ -70,7 +70,6 @@ def input_data(self, grid: base.BaseGrid) -> dict:
             surface_pressure=surface_pressure,
             pressure=pressure,
             pressure_ifc=pressure_ifc,
-            grav_o_rd=phy_const.GRAV_O_RD,
             horizontal_start=gtx.int32(0),
             horizontal_end=gtx.int32(grid.num_cells),
             vertical_start=gtx.int32(0),
diff --git a/model/common/tests/stencil_tests/test_diagnose_surface_pressure.py b/model/common/tests/stencil_tests/test_diagnose_surface_pressure.py
index b1d21b5e94..2bbc2d1874 100644
--- a/model/common/tests/stencil_tests/test_diagnose_surface_pressure.py
+++ b/model/common/tests/stencil_tests/test_diagnose_surface_pressure.py
@@ -66,9 +66,6 @@ def input_data(self, grid: base.BaseGrid) -> dict:
             virtual_temperature=virtual_temperature,
             ddqz_z_full=ddqz_z_full,
             surface_pressure=surface_pressure,
-            cpd_o_rd=phy_const.CPD_O_RD,
-            p0ref=phy_const.P0REF,
-            grav_o_rd=phy_const.GRAV_O_RD,
             horizontal_start=gtx.int32(0),
             horizontal_end=gtx.int32(grid.num_cells),
             vertical_start=gtx.int32(grid.num_levels),
diff --git a/model/common/tests/stencil_tests/test_diagnose_temperature.py b/model/common/tests/stencil_tests/test_diagnose_temperature.py
index b968d04c5e..9e190d51eb 100644
--- a/model/common/tests/stencil_tests/test_diagnose_temperature.py
+++ b/model/common/tests/stencil_tests/test_diagnose_temperature.py
@@ -85,7 +85,6 @@ def input_data(self, grid: base.BaseGrid) -> dict:
             exner=exner,
             virtual_temperature=virtual_temperature,
             temperature=temperature,
-            rv_o_rd_minus1=phy_const.RV_O_RD_MINUS_1,
             horizontal_start=gtx.int32(0),
             horizontal_end=gtx.int32(grid.num_cells),
             vertical_start=gtx.int32(0),


Reviews:
user: login: halungge
body: I think this is a huge improvement. I have some small remarks only... 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T14:09:56Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T09:16:02Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T09:17:58Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T09:18:22Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T11:06:14Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T11:15:26Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T11:33:53Z
user: login: OngChia
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T14:53:56Z
user: login: halungge
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T11:01:48Z
user: login: halungge
body: LGTM
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T11:08:22Z

Review comments:
diff_hunk: @@ -10,64 +10,66 @@
 import sys
 from typing import Final
 
-from icon4py.model.common.type_alias import wpfloat
+from gt4py.eve.utils import FrozenNamespace
path: model/common/src/icon4py/model/common/constants.py
user: login: halungge
body: ```suggestion
from gt4py.eve import utils as eve_utils
```
created_at: 2025-05-13T09:00:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -78,11 +80,23 @@
 DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Final[float] = 5.0
 
 #: average earth radius in [m]
-EARTH_RADIUS: Final[wpfloat] = 6.371229e6
+EARTH_RADIUS: Final[ta.wpfloat] = 6.371229e6
 
 
 class RayleighType(enum.IntEnum):
     #: classical Rayleigh damping, which makes use of a reference state.
     CLASSIC = 1
     #: Klemp (2008) type Rayleigh damping
     KLEMP = 2
+
+
+class _PhysicsConstants(FrozenNamespace[ta.wpfloat]):
path: model/common/src/icon4py/model/common/constants.py
user: login: halungge
body: ```suggestion
class PhysicsConstants(eve_utils.FrozenNamespace[ta.wpfloat]):
```
I would name this class without leading `_` because you import it in other modules. `_x` is conventionally for a private object, that is something that should not be used outside of a class, module... 
created_at: 2025-05-13T12:22:38Z
author_association: CONTRIBUTOR
path: model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics_tests/test_saturation_adjustment.py
user: login: halungge
body: Would it be possible to split the test into one for `SaturationAdjustment` only and one for the diagnostic functions? Or do not have the data for that, I mean we have no data to validate what comes out of saturation adjustment without the diagnostic computation?
created_at: 2025-05-13T13:54:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -269,80 +184,36 @@ def _allocate_tendencies(self):
         self._k_field = data_alloc.index_field(
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: halungge
body: do you still need that field or is it now obsolete or can be replaced by `concat_where`?
created_at: 2025-05-13T13:58:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,66 +19,49 @@
 )
 
 import icon4py.model.common.dimension as dims
+import icon4py.model.common.utils as common_utils
 from icon4py.model.common import (
     constants as phy_const,
     field_type_aliases as fa,
     type_alias as ta,
 )
-from icon4py.model.common.diagnostic_calculations.stencils import (
-    diagnose_pressure as pressure,
-    diagnose_surface_pressure as surface_pressure,
-)
 from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid, vertical as v_grid
-from icon4py.model.common.states import (
-    diagnostic_state as diagnostics,
-    model,
-    prognostic_state as prognostics,
-    tracer_state as tracers,
-)
+from icon4py.model.common.states import model
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
+physics_constants: Final = phy_const._PhysicsConstants()
+
+
 # TODO (Chia Rui): Refactor this class when direct import is enabled for gt4py stencils
-class SaturatedPressureConstants(FrozenNamespace):
+class MicrophysicsConstants(FrozenNamespace[ta.wpfloat]):
     """
     Constants used for the computation of saturated pressure in saturation adjustment and microphysics.
     It was originally in mo_lookup_tables_constants.f90.
     """
 
     #: Latent heat of vaporisation for water [J kg-1]. Originally expressed as alv in ICON.
-    vaporisation_latent_heat: ta.wpfloat = 2.5008e6
+    vaporisation_latent_heat = 2.5008e6
     #: Melting temperature of ice/snow [K]
-    tmelt: ta.wpfloat = 273.15
-
-    #: See docstring in common/constanst.py
-    rd: ta.wpfloat = phy_const.GAS_CONSTANT_DRY_AIR
-    #: See docstring in common/constanst.py
-    rv: ta.wpfloat = phy_const.GAS_CONSTANT_WATER_VAPOR
-    #: See docstring in common/constanst.py
-    cvd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_VOLUME
-    #: See docstring in common/constanst.py
-    cpd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_PRESSURE
-
-    rv_o_rd_minus_1: ta.wpfloat = rv / rd - 1.0
-    rd_o_cpd: ta.wpfloat = phy_const.RD_O_CPD
+    tmelt = 273.15
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: halungge
body: ```suggestion
    melting_temperature = 273.15
```
?
created_at: 2025-05-13T14:01:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,66 +19,49 @@
 )
 
 import icon4py.model.common.dimension as dims
+import icon4py.model.common.utils as common_utils
 from icon4py.model.common import (
     constants as phy_const,
     field_type_aliases as fa,
     type_alias as ta,
 )
-from icon4py.model.common.diagnostic_calculations.stencils import (
-    diagnose_pressure as pressure,
-    diagnose_surface_pressure as surface_pressure,
-)
 from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid, vertical as v_grid
-from icon4py.model.common.states import (
-    diagnostic_state as diagnostics,
-    model,
-    prognostic_state as prognostics,
-    tracer_state as tracers,
-)
+from icon4py.model.common.states import model
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
+physics_constants: Final = phy_const._PhysicsConstants()
+
+
 # TODO (Chia Rui): Refactor this class when direct import is enabled for gt4py stencils
-class SaturatedPressureConstants(FrozenNamespace):
+class MicrophysicsConstants(FrozenNamespace[ta.wpfloat]):
     """
     Constants used for the computation of saturated pressure in saturation adjustment and microphysics.
     It was originally in mo_lookup_tables_constants.f90.
     """
 
     #: Latent heat of vaporisation for water [J kg-1]. Originally expressed as alv in ICON.
-    vaporisation_latent_heat: ta.wpfloat = 2.5008e6
+    vaporisation_latent_heat = 2.5008e6
     #: Melting temperature of ice/snow [K]
-    tmelt: ta.wpfloat = 273.15
-
-    #: See docstring in common/constanst.py
-    rd: ta.wpfloat = phy_const.GAS_CONSTANT_DRY_AIR
-    #: See docstring in common/constanst.py
-    rv: ta.wpfloat = phy_const.GAS_CONSTANT_WATER_VAPOR
-    #: See docstring in common/constanst.py
-    cvd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_VOLUME
-    #: See docstring in common/constanst.py
-    cpd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_PRESSURE
-
-    rv_o_rd_minus_1: ta.wpfloat = rv / rd - 1.0
-    rd_o_cpd: ta.wpfloat = phy_const.RD_O_CPD
+    tmelt = 273.15
 
     #: Dry air heat capacity at constant pressure / water heat capacity at constant pressure - 1
-    rcpl: ta.wpfloat = 3.1733
+    rcpl = 3.1733
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: halungge
body: ```suggestion
    heat_capacity_at_constant_pressure = 3.1733
```
?
created_at: 2025-05-13T14:02:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -376,162 +257,74 @@ def run(
             offset_provider={},
         )
 
-        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
-        temperature_list = [self._temperature1, self._temperature2]
-        ncurrent, nnext = 0, 1
-        for _ in range(self.config.max_iter):
-            if np.any(
-                self._newton_iteration_mask.ndarray[
-                    self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
-                ]
-            ):
-                self.update_temperature_by_newton_iteration(
-                    diagnostic_state.temperature,
-                    tracer_state.qv,
-                    prognostic_state.rho,
-                    self._newton_iteration_mask,
-                    self._lwdocvd,
-                    temperature_list[nnext],
-                    temperature_list[ncurrent],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-
-                self.compute_newton_iteration_mask(
-                    self.config.tolerance,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    self._newton_iteration_mask,
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
+        temperature_pair = common_utils.TimeStepPair(self._temperature1, self._temperature2)
 
-                self.copy_temperature(
-                    self._newton_iteration_mask,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-                ncurrent = (ncurrent + 1) % 2
-                nnext = (nnext + 1) % 2
-            else:
-                break
-        if np.any(
+        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
+        not_converged = self._xp.any(
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: halungge
body: out of curiosity: does ICON do any global reduction over the horizontal domains here? Or is this done per node in a multinode setup?
created_at: 2025-05-13T14:05:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -376,162 +257,74 @@ def run(
             offset_provider={},
         )
 
-        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
-        temperature_list = [self._temperature1, self._temperature2]
-        ncurrent, nnext = 0, 1
-        for _ in range(self.config.max_iter):
-            if np.any(
-                self._newton_iteration_mask.ndarray[
-                    self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
-                ]
-            ):
-                self.update_temperature_by_newton_iteration(
-                    diagnostic_state.temperature,
-                    tracer_state.qv,
-                    prognostic_state.rho,
-                    self._newton_iteration_mask,
-                    self._lwdocvd,
-                    temperature_list[nnext],
-                    temperature_list[ncurrent],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-
-                self.compute_newton_iteration_mask(
-                    self.config.tolerance,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    self._newton_iteration_mask,
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
+        temperature_pair = common_utils.TimeStepPair(self._temperature1, self._temperature2)
 
-                self.copy_temperature(
-                    self._newton_iteration_mask,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-                ncurrent = (ncurrent + 1) % 2
-                nnext = (nnext + 1) % 2
-            else:
-                break
-        if np.any(
+        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
+        not_converged = self._xp.any(
             self._newton_iteration_mask.ndarray[
                 self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
             ]
-        ):
-            raise ConvergenceError(
-                f"Maximum iteration of saturation adjustment ({self.config.max_iter}) is not enough. The max absolute error is {np.abs(self.new_temperature1.ndarray - self.new_temperature2.ndarray).max()} . Please raise max_iter"
-            )
-        self.update_temperature_qv_qc_tendencies(
-            dtime,
-            diagnostic_state.temperature,
-            temperature_list[ncurrent],
-            tracer_state.qv,
-            tracer_state.qc,
-            prognostic_state.rho,
-            self._subsaturated_mask,
-            self.temperature_tendency,
-            self.qv_tendency,
-            self.qc_tendency,
-            horizontal_start=self._start_cell_nudging,
-            horizontal_end=self._end_cell_local,
-            vertical_start=gtx.int32(0),
-            vertical_end=self.grid.num_levels,
-            offset_provider={},
         )
+        num_iter = 0
+        while not_converged:
+            if num_iter > self.config.max_iter:
+                raise ConvergenceError(
+                    f"Maximum iteration of saturation adjustment ({self.config.max_iter}) is not enough. The max absolute error is {self._xp.abs(self._temperature1.ndarray - self._temperature2.ndarray).max()} . Please raise max_iter"
+                )
 
-        if self.config.diagnose_variables_from_new_temperature:
-            # TODO (Chia Rui): Consider to merge the gt4py stencils below if scan operator can be called in the intermediate step
-            self.compute_temperature_and_exner_tendencies_after_saturation_adjustment(
-                dtime,
-                tracer_state.qv,
-                tracer_state.qc,
-                tracer_state.qi,
-                tracer_state.qr,
-                tracer_state.qs,
-                tracer_state.qg,
-                self.qv_tendency,
-                self.qc_tendency,
-                temperature_list[ncurrent],
-                diagnostic_state.virtual_temperature,
-                prognostic_state.exner,
-                self.virtual_temperature_tendency,
-                self.exner_tendency,
-                self._new_exner,
-                self._new_virtual_temperature,
-                self._k_field,
-                self.vertical_params.kstart_moist,
+            self.update_temperature_by_newton_iteration(
+                temperature,
+                qv,
+                rho,
+                self._newton_iteration_mask,
+                self._lwdocvd,
+                temperature_pair.next,
+                temperature_pair.current,
                 horizontal_start=self._start_cell_nudging,
                 horizontal_end=self._end_cell_local,
                 vertical_start=gtx.int32(0),
                 vertical_end=self.grid.num_levels,
                 offset_provider={},
             )
 
-            self.diagnose_surface_pressure(
-                self._new_exner,
-                self._new_virtual_temperature,
-                self.metric_state.ddqz_z_full,
-                self._pressure_ifc,
-                phy_const.CPD_O_RD,
-                phy_const.P0REF,
-                phy_const.GRAV_O_RD,
-                horizontal_start=self._start_cell_nudging,
-                horizontal_end=self._end_cell_local,
-                vertical_start=self.grid.num_levels,
-                vertical_end=gtx.int32(self.grid.num_levels + 1),
-                offset_provider={"Koff": dims.KDim},
-            )
-
-            self.diagnose_pressure(
-                self.metric_state.ddqz_z_full,
-                self._new_virtual_temperature,
-                gtx.as_field((dims.CellDim,), self._pressure_ifc.ndarray[:, -1]),
-                self._pressure,
-                self._pressure_ifc,
-                phy_const.GRAV_O_RD,
+            self.compute_newton_iteration_mask_and_copy_temperature_on_converged_cells(
+                self.config.tolerance,
+                temperature_pair.current,
+                temperature_pair.next,
+                self._newton_iteration_mask,
                 horizontal_start=self._start_cell_nudging,
                 horizontal_end=self._end_cell_local,
                 vertical_start=gtx.int32(0),
                 vertical_end=self.grid.num_levels,
                 offset_provider={},
             )
 
-            self.compute_pressure_tendency_after_saturation_adjustment(
-                dtime,
-                diagnostic_state.pressure,
-                self._pressure,
-                self.pressure_tendency,
-                horizontal_start=self._start_cell_nudging,
-                horizontal_end=self._end_cell_local,
-                vertical_start=gtx.int32(0),
-                vertical_end=self.grid.num_levels,
-                offset_provider={},
+            not_converged = self._xp.any(
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: halungge
body: you could extract this into a function `has_converged` that you call wherever needed (if I am not mistaken you run the evaluation code twice...) then when you adapt the convergence criteria you only need to do it once.
created_at: 2025-05-13T14:07:46Z
author_association: CONTRIBUTOR
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: halungge
body: +1 for this module having become much shorter... 
created_at: 2025-05-13T14:08:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -78,11 +80,23 @@
 DEFAULT_PHYSICS_DYNAMICS_TIMESTEP_RATIO: Final[float] = 5.0
 
 #: average earth radius in [m]
-EARTH_RADIUS: Final[wpfloat] = 6.371229e6
+EARTH_RADIUS: Final[ta.wpfloat] = 6.371229e6
 
 
 class RayleighType(enum.IntEnum):
     #: classical Rayleigh damping, which makes use of a reference state.
     CLASSIC = 1
     #: Klemp (2008) type Rayleigh damping
     KLEMP = 2
+
+
+class _PhysicsConstants(FrozenNamespace[ta.wpfloat]):
path: model/common/src/icon4py/model/common/constants.py
user: login: OngChia
body: Ah. Right. Thanks. 
created_at: 2025-05-14T09:16:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -269,80 +184,36 @@ def _allocate_tendencies(self):
         self._k_field = data_alloc.index_field(
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: OngChia
body: Thanks for noting that! I forgot to delete it. It is not needed anymore in saturation adjustment. 
created_at: 2025-05-14T09:17:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,66 +19,49 @@
 )
 
 import icon4py.model.common.dimension as dims
+import icon4py.model.common.utils as common_utils
 from icon4py.model.common import (
     constants as phy_const,
     field_type_aliases as fa,
     type_alias as ta,
 )
-from icon4py.model.common.diagnostic_calculations.stencils import (
-    diagnose_pressure as pressure,
-    diagnose_surface_pressure as surface_pressure,
-)
 from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid, vertical as v_grid
-from icon4py.model.common.states import (
-    diagnostic_state as diagnostics,
-    model,
-    prognostic_state as prognostics,
-    tracer_state as tracers,
-)
+from icon4py.model.common.states import model
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
+physics_constants: Final = phy_const._PhysicsConstants()
+
+
 # TODO (Chia Rui): Refactor this class when direct import is enabled for gt4py stencils
-class SaturatedPressureConstants(FrozenNamespace):
+class MicrophysicsConstants(FrozenNamespace[ta.wpfloat]):
     """
     Constants used for the computation of saturated pressure in saturation adjustment and microphysics.
     It was originally in mo_lookup_tables_constants.f90.
     """
 
     #: Latent heat of vaporisation for water [J kg-1]. Originally expressed as alv in ICON.
-    vaporisation_latent_heat: ta.wpfloat = 2.5008e6
+    vaporisation_latent_heat = 2.5008e6
     #: Melting temperature of ice/snow [K]
-    tmelt: ta.wpfloat = 273.15
-
-    #: See docstring in common/constanst.py
-    rd: ta.wpfloat = phy_const.GAS_CONSTANT_DRY_AIR
-    #: See docstring in common/constanst.py
-    rv: ta.wpfloat = phy_const.GAS_CONSTANT_WATER_VAPOR
-    #: See docstring in common/constanst.py
-    cvd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_VOLUME
-    #: See docstring in common/constanst.py
-    cpd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_PRESSURE
-
-    rv_o_rd_minus_1: ta.wpfloat = rv / rd - 1.0
-    rd_o_cpd: ta.wpfloat = phy_const.RD_O_CPD
+    tmelt = 273.15
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: OngChia
body: Good suggestion! Thanks!
created_at: 2025-05-14T09:18:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,66 +19,49 @@
 )
 
 import icon4py.model.common.dimension as dims
+import icon4py.model.common.utils as common_utils
 from icon4py.model.common import (
     constants as phy_const,
     field_type_aliases as fa,
     type_alias as ta,
 )
-from icon4py.model.common.diagnostic_calculations.stencils import (
-    diagnose_pressure as pressure,
-    diagnose_surface_pressure as surface_pressure,
-)
 from icon4py.model.common.grid import horizontal as h_grid, icon as icon_grid, vertical as v_grid
-from icon4py.model.common.states import (
-    diagnostic_state as diagnostics,
-    model,
-    prognostic_state as prognostics,
-    tracer_state as tracers,
-)
+from icon4py.model.common.states import model
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
+physics_constants: Final = phy_const._PhysicsConstants()
+
+
 # TODO (Chia Rui): Refactor this class when direct import is enabled for gt4py stencils
-class SaturatedPressureConstants(FrozenNamespace):
+class MicrophysicsConstants(FrozenNamespace[ta.wpfloat]):
     """
     Constants used for the computation of saturated pressure in saturation adjustment and microphysics.
     It was originally in mo_lookup_tables_constants.f90.
     """
 
     #: Latent heat of vaporisation for water [J kg-1]. Originally expressed as alv in ICON.
-    vaporisation_latent_heat: ta.wpfloat = 2.5008e6
+    vaporisation_latent_heat = 2.5008e6
     #: Melting temperature of ice/snow [K]
-    tmelt: ta.wpfloat = 273.15
-
-    #: See docstring in common/constanst.py
-    rd: ta.wpfloat = phy_const.GAS_CONSTANT_DRY_AIR
-    #: See docstring in common/constanst.py
-    rv: ta.wpfloat = phy_const.GAS_CONSTANT_WATER_VAPOR
-    #: See docstring in common/constanst.py
-    cvd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_VOLUME
-    #: See docstring in common/constanst.py
-    cpd: ta.wpfloat = phy_const.SPECIFIC_HEAT_CONSTANT_PRESSURE
-
-    rv_o_rd_minus_1: ta.wpfloat = rv / rd - 1.0
-    rd_o_cpd: ta.wpfloat = phy_const.RD_O_CPD
+    tmelt = 273.15
 
     #: Dry air heat capacity at constant pressure / water heat capacity at constant pressure - 1
-    rcpl: ta.wpfloat = 3.1733
+    rcpl = 3.1733
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: OngChia
body: I just realized that `rcpl` is just used for computing `specific_heat_capacity_for_water`. I checked inside ICON and it is not even used in graupel. I deleted it and directly inlined in `specific_heat_capacity_for_water`.
created_at: 2025-05-14T11:06:14Z
author_association: CONTRIBUTOR
path: model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics_tests/test_saturation_adjustment.py
user: login: OngChia
body: Yes. It is possible because variables updated in diagnostic calculations after saturation adjustment are different from those variables updated by saturation adjustment.
I split it and create a new datatest in `test_diagnostic_calculations.py`
created_at: 2025-05-14T11:15:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -376,162 +257,74 @@ def run(
             offset_provider={},
         )
 
-        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
-        temperature_list = [self._temperature1, self._temperature2]
-        ncurrent, nnext = 0, 1
-        for _ in range(self.config.max_iter):
-            if np.any(
-                self._newton_iteration_mask.ndarray[
-                    self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
-                ]
-            ):
-                self.update_temperature_by_newton_iteration(
-                    diagnostic_state.temperature,
-                    tracer_state.qv,
-                    prognostic_state.rho,
-                    self._newton_iteration_mask,
-                    self._lwdocvd,
-                    temperature_list[nnext],
-                    temperature_list[ncurrent],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-
-                self.compute_newton_iteration_mask(
-                    self.config.tolerance,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    self._newton_iteration_mask,
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
+        temperature_pair = common_utils.TimeStepPair(self._temperature1, self._temperature2)
 
-                self.copy_temperature(
-                    self._newton_iteration_mask,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-                ncurrent = (ncurrent + 1) % 2
-                nnext = (nnext + 1) % 2
-            else:
-                break
-        if np.any(
+        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
+        not_converged = self._xp.any(
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: OngChia
body: No. It is done per node even in multinode setup.

By the way, there is a gpu-version (satad_v-3D_gpu) of saturation adjustment in ICON. They put a sequential for loop inside the outer (cell, k)-loops for Newton iteration. And the innermost loop breaks when the convergence criterion is met. So, there is no reduction in the gpu version. However, GT4Py does not support loop.
created_at: 2025-05-14T11:33:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -376,162 +257,74 @@ def run(
             offset_provider={},
         )
 
-        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
-        temperature_list = [self._temperature1, self._temperature2]
-        ncurrent, nnext = 0, 1
-        for _ in range(self.config.max_iter):
-            if np.any(
-                self._newton_iteration_mask.ndarray[
-                    self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
-                ]
-            ):
-                self.update_temperature_by_newton_iteration(
-                    diagnostic_state.temperature,
-                    tracer_state.qv,
-                    prognostic_state.rho,
-                    self._newton_iteration_mask,
-                    self._lwdocvd,
-                    temperature_list[nnext],
-                    temperature_list[ncurrent],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-
-                self.compute_newton_iteration_mask(
-                    self.config.tolerance,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    self._newton_iteration_mask,
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
+        temperature_pair = common_utils.TimeStepPair(self._temperature1, self._temperature2)
 
-                self.copy_temperature(
-                    self._newton_iteration_mask,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-                ncurrent = (ncurrent + 1) % 2
-                nnext = (nnext + 1) % 2
-            else:
-                break
-        if np.any(
+        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
+        not_converged = self._xp.any(
             self._newton_iteration_mask.ndarray[
                 self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
             ]
-        ):
-            raise ConvergenceError(
-                f"Maximum iteration of saturation adjustment ({self.config.max_iter}) is not enough. The max absolute error is {np.abs(self.new_temperature1.ndarray - self.new_temperature2.ndarray).max()} . Please raise max_iter"
-            )
-        self.update_temperature_qv_qc_tendencies(
-            dtime,
-            diagnostic_state.temperature,
-            temperature_list[ncurrent],
-            tracer_state.qv,
-            tracer_state.qc,
-            prognostic_state.rho,
-            self._subsaturated_mask,
-            self.temperature_tendency,
-            self.qv_tendency,
-            self.qc_tendency,
-            horizontal_start=self._start_cell_nudging,
-            horizontal_end=self._end_cell_local,
-            vertical_start=gtx.int32(0),
-            vertical_end=self.grid.num_levels,
-            offset_provider={},
         )
+        num_iter = 0
+        while not_converged:
+            if num_iter > self.config.max_iter:
+                raise ConvergenceError(
+                    f"Maximum iteration of saturation adjustment ({self.config.max_iter}) is not enough. The max absolute error is {self._xp.abs(self._temperature1.ndarray - self._temperature2.ndarray).max()} . Please raise max_iter"
+                )
 
-        if self.config.diagnose_variables_from_new_temperature:
-            # TODO (Chia Rui): Consider to merge the gt4py stencils below if scan operator can be called in the intermediate step
-            self.compute_temperature_and_exner_tendencies_after_saturation_adjustment(
-                dtime,
-                tracer_state.qv,
-                tracer_state.qc,
-                tracer_state.qi,
-                tracer_state.qr,
-                tracer_state.qs,
-                tracer_state.qg,
-                self.qv_tendency,
-                self.qc_tendency,
-                temperature_list[ncurrent],
-                diagnostic_state.virtual_temperature,
-                prognostic_state.exner,
-                self.virtual_temperature_tendency,
-                self.exner_tendency,
-                self._new_exner,
-                self._new_virtual_temperature,
-                self._k_field,
-                self.vertical_params.kstart_moist,
+            self.update_temperature_by_newton_iteration(
+                temperature,
+                qv,
+                rho,
+                self._newton_iteration_mask,
+                self._lwdocvd,
+                temperature_pair.next,
+                temperature_pair.current,
                 horizontal_start=self._start_cell_nudging,
                 horizontal_end=self._end_cell_local,
                 vertical_start=gtx.int32(0),
                 vertical_end=self.grid.num_levels,
                 offset_provider={},
             )
 
-            self.diagnose_surface_pressure(
-                self._new_exner,
-                self._new_virtual_temperature,
-                self.metric_state.ddqz_z_full,
-                self._pressure_ifc,
-                phy_const.CPD_O_RD,
-                phy_const.P0REF,
-                phy_const.GRAV_O_RD,
-                horizontal_start=self._start_cell_nudging,
-                horizontal_end=self._end_cell_local,
-                vertical_start=self.grid.num_levels,
-                vertical_end=gtx.int32(self.grid.num_levels + 1),
-                offset_provider={"Koff": dims.KDim},
-            )
-
-            self.diagnose_pressure(
-                self.metric_state.ddqz_z_full,
-                self._new_virtual_temperature,
-                gtx.as_field((dims.CellDim,), self._pressure_ifc.ndarray[:, -1]),
-                self._pressure,
-                self._pressure_ifc,
-                phy_const.GRAV_O_RD,
+            self.compute_newton_iteration_mask_and_copy_temperature_on_converged_cells(
+                self.config.tolerance,
+                temperature_pair.current,
+                temperature_pair.next,
+                self._newton_iteration_mask,
                 horizontal_start=self._start_cell_nudging,
                 horizontal_end=self._end_cell_local,
                 vertical_start=gtx.int32(0),
                 vertical_end=self.grid.num_levels,
                 offset_provider={},
             )
 
-            self.compute_pressure_tendency_after_saturation_adjustment(
-                dtime,
-                diagnostic_state.pressure,
-                self._pressure,
-                self.pressure_tendency,
-                horizontal_start=self._start_cell_nudging,
-                horizontal_end=self._end_cell_local,
-                vertical_start=gtx.int32(0),
-                vertical_end=self.grid.num_levels,
-                offset_provider={},
+            not_converged = self._xp.any(
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: OngChia
body: Ah, you are right. And, it can be actually called just once at the beginning of the while loop. 
I still extract it into a separate private function `_is_converged()` in `SaturationAdjustment` class because it looks messy with a long condition in `while self._xp.any(self._newton_iteration_mask.ndarray[self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels]):`
created_at: 2025-05-14T14:53:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -376,162 +257,74 @@ def run(
             offset_provider={},
         )
 
-        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
-        temperature_list = [self._temperature1, self._temperature2]
-        ncurrent, nnext = 0, 1
-        for _ in range(self.config.max_iter):
-            if np.any(
-                self._newton_iteration_mask.ndarray[
-                    self._start_cell_nudging : self._end_cell_local, 0 : self.grid.num_levels
-                ]
-            ):
-                self.update_temperature_by_newton_iteration(
-                    diagnostic_state.temperature,
-                    tracer_state.qv,
-                    prognostic_state.rho,
-                    self._newton_iteration_mask,
-                    self._lwdocvd,
-                    temperature_list[nnext],
-                    temperature_list[ncurrent],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-
-                self.compute_newton_iteration_mask(
-                    self.config.tolerance,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    self._newton_iteration_mask,
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
+        temperature_pair = common_utils.TimeStepPair(self._temperature1, self._temperature2)
 
-                self.copy_temperature(
-                    self._newton_iteration_mask,
-                    temperature_list[ncurrent],
-                    temperature_list[nnext],
-                    horizontal_start=self._start_cell_nudging,
-                    horizontal_end=self._end_cell_local,
-                    vertical_start=gtx.int32(0),
-                    vertical_end=self.grid.num_levels,
-                    offset_provider={},
-                )
-                ncurrent = (ncurrent + 1) % 2
-                nnext = (nnext + 1) % 2
-            else:
-                break
-        if np.any(
+        # TODO (Chia Rui): this is inspired by the cpu version of the original ICON saturation_adjustment code. Consider to refactor this code when break and for loop features are ready in gt4py.
+        not_converged = self._xp.any(
path: model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py
user: login: halungge
body: Yes, no does it support reductions.... 
created_at: 2025-05-15T11:01:48Z
author_association: CONTRIBUTOR

PR comments:
user: login: OngChia
created_at: 2025-05-09T14:46:36Z
body: cscs-ci run default
author_association: CONTRIBUTOR
user: login: OngChia
created_at: 2025-05-10T19:30:24Z
body: cscs-ci run default
author_association: CONTRIBUTOR
user: login: OngChia
created_at: 2025-05-14T22:36:35Z
body: cscs-ci run default
author_association: CONTRIBUTOR
user: login: OngChia
created_at: 2025-05-15T05:44:13Z
body: cscs-ci run default
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-05-15T05:44:15Z
body: <!-- add-pr-comment:add-pr-comment -->

**Mandatory Tests**

Please make sure you run these tests via comment before you merge!

* `cscs-ci run default`

**Optional Tests**

To run benchmarks you can use:

* `cscs-ci run benchmark-bencher`

To run tests and benchmarks with the DaCe backend you can use:

* `cscs-ci run dace`

To run test levels ignored by the default test suite (mostly simple datatest for static fields computations) you can use:
 * `cscs-ci run extra`



For more detailed information please look at [CI in the EXCLAIM universe](https://github.com/C2SM/icon-exclaim/wiki).
author_association: NONE