Physics driver with microphysics

model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/saturation_adjustment.py

@@ -15,14 +15,17 @@
 from icon4py.model.atmosphere.subgrid_scale_physics.microphysics.stencils import (
     saturation_adjustment_stencils as satad_stencils,
 )
-from icon4py.model.common import field_type_aliases as fa, model_options, type_alias as ta
+from icon4py.model.common import (
+    field_type_aliases as fa,
+    model_backends,
+    model_options,
+    type_alias as ta,
+)
 from icon4py.model.common.grid import horizontal as h_grid
 from icon4py.model.common.utils import data_allocation as data_alloc
 
 
 if TYPE_CHECKING:
-    import gt4py.next.typing as gtx_typing
-
     from icon4py.model.common.grid import icon as icon_grid, vertical as v_grid
     from icon4py.model.common.states import model
 
@@ -100,14 +103,15 @@ def __init__(
         grid: icon_grid.IconGrid,
         vertical_params: v_grid.VerticalGrid,
         metric_state: MetricStateSaturationAdjustment,
-        backend: gtx_typing.Backend | None,
+        backend: model_backends.BackendLike,
     ):
         self._backend = backend
+        self._allocator = model_backends.get_allocator(backend)
         self.config = config
         self._grid = grid
         self._vertical_params: v_grid.VerticalGrid = vertical_params
         self._metric_state: MetricStateSaturationAdjustment = metric_state
-        self._xp = data_alloc.import_array_ns(self._backend)
+        self._xp = data_alloc.import_array_ns(self._allocator)
 
         self._allocate_local_variables()
         self._determine_horizontal_domains()
@@ -123,23 +127,23 @@ def output_properties(self) -> dict[str, model.FieldMetaData]:
     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, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         #: it was originally named as twork in ICON. New temperature before iteration.
         self._temperature2 = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         #: A mask that indicates whether the grid cell is subsaturated or not.
         self._subsaturated_mask = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=bool, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=bool, allocator=self._allocator
         )
         #: A mask that indicates whether next Newton iteration is required.
         self._newton_iteration_mask = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=bool, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=bool, allocator=self._allocator
         )
         #: latent heat vaporization / dry air heat capacity at constant volume
         self._lwdocvd = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
 
     def _initialize_gt4py_programs(self):
@@ -203,7 +207,7 @@ def _initialize_gt4py_programs(self):
     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))
+        self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL))
 
     def _not_converged(self) -> bool:
         return self._xp.any(
@@ -215,14 +219,14 @@ def _not_converged(self) -> bool:
 
     def run(
         self,
-        dtime: ta.wpfloat,
+        temperature_tendency: fa.CellKField[ta.wpfloat],
+        qv_tendency: fa.CellKField[ta.wpfloat],
+        qc_tendency: fa.CellKField[ta.wpfloat],
         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],
+        dtime: ta.wpfloat,
     ):
         """
         Adjust saturation at each grid point.
@@ -238,14 +242,14 @@ def run(
         Originally inspired from satad_v_3D of ICON.
 
         Args:
-            dtime: time step [s]
+            temperature_tendency: air temperature tendency [K s-1]
+            qv_tendency: specific humidity tendency [s-1]
+            qc_tendency: specific cloud water content tendency [s-1]
             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]
+            dtime: time step [s]
         """
 
         temperature_pair = common_utils.TimeStepPair(self._temperature1, self._temperature2)

model/atmosphere/subgrid_scale_physics/microphysics/src/icon4py/model/atmosphere/subgrid_scale_physics/microphysics/single_moment_six_class_gscp_graupel.py

@@ -22,6 +22,7 @@
     constants as physics_constants,
     dimension as dims,
     field_type_aliases as fa,
+    model_backends,
     model_options,
     type_alias as ta,
 )
@@ -30,8 +31,6 @@
 
 
 if TYPE_CHECKING:
-    import gt4py.next.typing as gtx_typing
-
     from icon4py.model.common.grid import icon as icon_grid, vertical as v_grid
 
 
@@ -86,6 +85,8 @@ class SingleMomentSixClassIconGraupelConfig:
     rain_n0: ta.wpfloat = 1.0
     #: coefficient for snow-graupel conversion by riming. Originally defined as csg in mo_nwp_tuning_config.f90 in ICON.
     snow2graupel_riming_coeff: ta.wpfloat = 0.5
+    #: cloud number concentration [1/m3]
+    cloud_number_concentration: ta.wpfloat = 100.0e6
 
 
 @dataclasses.dataclass
@@ -96,18 +97,19 @@ class MetricStateIconGraupel:
 class SingleMomentSixClassIconGraupel:
     def __init__(
         self,
-        graupel_config: SingleMomentSixClassIconGraupelConfig,
+        config: SingleMomentSixClassIconGraupelConfig,
         grid: icon_grid.IconGrid,
         metric_state: MetricStateIconGraupel,
         vertical_params: v_grid.VerticalGrid,
-        backend: gtx_typing.Backend | None,
+        backend: model_backends.BackendLike,
     ):
-        self.config = graupel_config
+        self.config = config
         self._initialize_configurable_parameters()
         self._grid = grid
         self._metric_state = metric_state
         self.vertical_params = vertical_params
         self._backend = backend
+        self._allocator = model_backends.get_allocator(backend)
 
         self._initialize_local_fields()
         self._determine_horizontal_domains()
@@ -218,50 +220,53 @@ def _initialize_configurable_parameters(self):
         )
 
     def _initialize_local_fields(self):
+        self.qnc = data_alloc.zero_field(
+            self._grid, dims.CellDim, dtype=ta.wpfloat, allocator=self._allocator
+        )
         self.rhoqrv_old_kup = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.rhoqsv_old_kup = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.rhoqgv_old_kup = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.rhoqiv_old_kup = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.vnew_r = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.vnew_s = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.vnew_g = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.vnew_i = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.rain_precipitation_flux = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.snow_precipitation_flux = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.graupel_precipitation_flux = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.ice_precipitation_flux = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
         self.total_precipitation_flux = data_alloc.zero_field(
-            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._backend
+            self._grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=self._allocator
         )
 
     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))
+        self._end_cell_local = self._grid.end_index(cell_domain(h_grid.Zone.LOCAL))
 
     def _initialize_gt4py_programs(self):
         self._icon_graupel = model_options.setup_program(
@@ -346,7 +351,6 @@ def run(
         qi: fa.CellKField[ta.wpfloat],
         qs: fa.CellKField[ta.wpfloat],
         qg: fa.CellKField[ta.wpfloat],
-        qnc: fa.CellField[ta.wpfloat],
         temperature_tendency: fa.CellKField[ta.wpfloat],
         qv_tendency: fa.CellKField[ta.wpfloat],
         qc_tendency: fa.CellKField[ta.wpfloat],
@@ -389,7 +393,7 @@ def run(
             qr=qr,
             qs=qs,
             qg=qg,
-            qnc=qnc,
+            qnc=self.qnc,
             temperature_tendency=temperature_tendency,
             qv_tendency=qv_tendency,
             qc_tendency=qc_tendency,

model/atmosphere/subgrid_scale_physics/microphysics/tests/microphysics/integration_tests/test_single_moment_six_class_gscp_graupel.py

@@ -15,7 +15,7 @@
     microphysics_options as mphys_options,
     single_moment_six_class_gscp_graupel as graupel,
 )
-from icon4py.model.common import dimension as dims, type_alias as ta
+from icon4py.model.common import dimension as dims, model_backends, type_alias as ta
 from icon4py.model.common.grid import vertical as v_grid
 from icon4py.model.common.states import (
     diagnostic_state as diagnostics,
@@ -94,11 +94,15 @@ def test_graupel(
         temperature=entry_savepoint.temperature(),
         virtual_temperature=None,
         pressure=entry_savepoint.pressure(),
-        pressure_ifc=None,
+        pressure_at_half_levels=None,
         u=None,
         v=None,
     )
 
+    allocator = model_backends.get_allocator(backend)
+    xp = data_alloc.import_array_ns(allocator)
+    qnc = xp.mean(entry_savepoint.qnc())
+
     graupel_config = graupel.SingleMomentSixClassIconGraupelConfig(
         liquid_autoconversion_option=mphys_options.LiquidAutoConversionType.SEIFERT_BEHENG,
         ice_stickeff_min=0.01,
@@ -108,18 +112,17 @@ def test_graupel(
         rain_mu=0.0,
         rain_n0=1.0,
         snow2graupel_riming_coeff=0.5,
+        cloud_number_concentration=qnc,
     )
 
     graupel_microphysics = graupel.SingleMomentSixClassIconGraupel(
-        graupel_config=graupel_config,
+        config=graupel_config,
         grid=icon_grid,
         metric_state=metric_state,
         vertical_params=vertical_params,
         backend=backend,
     )
 
-    qnc = entry_savepoint.qnc()
-
     temperature_tendency = data_alloc.zero_field(
         icon_grid, dims.CellDim, dims.KDim, dtype=ta.wpfloat, allocator=backend
     )
@@ -153,7 +156,6 @@ def test_graupel(
         tracer_state.qi,
         tracer_state.qs,
         tracer_state.qg,
-        qnc,
         temperature_tendency,
         qv_tendency,
         qc_tendency,