[DSL] [Update] Aer Activation to NDSL == 2025.10.00

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/aer_activation.py

@@ -1,7 +1,7 @@
 from gt4py.cartesian.gtscript import PARALLEL, computation, exp, float64, interval, log, sqrt
 
 import pyMoist.constants as constants
-from ndsl import QuantityFactory, StencilFactory, orchestrate
+from ndsl import NDSLRuntime, QuantityFactory, StencilFactory
 from ndsl.constants import X_DIM, Y_DIM, Z_DIM
 from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int
 from pyMoist.field_types import FloatField_NModes
@@ -21,8 +21,8 @@
 DENSIC = float64(917.0)  # Ice crystal density in kgm-3
 
 # Default precision
-NN_MIN = 100.0e6
-NN_MAX = 1000.0e6
+NN_MIN = Float(100.0e6)
+NN_MAX = Float(1000.0e6)
 
 # ACTFRAC_Mat constants - all 64 bit
 PI = float64(3.141592653589793e00)
@@ -282,7 +282,7 @@ def aer_activation_stencil(
             nacti = NN_MAX
 
 
-class AerActivation:
+class AerActivation(NDSLRuntime):
     """
     Class for aerosol activation computation.
 
@@ -313,48 +313,28 @@ def __init__(
         NotImplementedError: If the number of modes is not equal to the expected number.
         NotImplementedError: If the neural network for aerosol is not implemented.
         """
-        orchestrate(obj=self, config=stencil_factory.config.dace_config)
+        super().__init__(dace_config=stencil_factory.config.dace_config)
 
         if constants.N_MODES != n_modes:
             raise NotImplementedError(
-                f"Coding limitation: {constants.N_MODES} modes are expected, " f"getting {n_modes}"
+                f"Coding limitation: {constants.N_MODES} modes are expected, getting {n_modes}"
             )
 
         if not USE_AERSOL_NN:
             raise NotImplementedError("Non NN Aerosol not implemented")
 
-        # Temporary buffers
+        # Locals
         quantity_factory.add_data_dimensions(
-            **{
+            {
                 "n_modes": constants.N_MODES,
             }
         )
 
-        self._nact = quantity_factory.zeros(
-            [X_DIM, Y_DIM, Z_DIM, "n_modes"],
-            units="n/a",
-            dtype=Float,
-        )
-        self._ni = quantity_factory.zeros(
-            [X_DIM, Y_DIM, Z_DIM, "n_modes"],
-            units="n/a",
-            dtype=Float,
-        )
-        self._rg = quantity_factory.zeros(
-            [X_DIM, Y_DIM, Z_DIM, "n_modes"],
-            units="n/a",
-            dtype=Float,
-        )
-        self._sig0 = quantity_factory.zeros(
-            [X_DIM, Y_DIM, Z_DIM, "n_modes"],
-            units="n/a",
-            dtype=Float,
-        )
-        self._bibar = quantity_factory.zeros(
-            [X_DIM, Y_DIM, Z_DIM, "n_modes"],
-            units="n/a",
-            dtype=Float,
-        )
+        self._nact = self.make_local(quantity_factory, [X_DIM, Y_DIM, Z_DIM, "n_modes"])
+        self._ni = self.make_local(quantity_factory, [X_DIM, Y_DIM, Z_DIM, "n_modes"])
+        self._rg = self.make_local(quantity_factory, [X_DIM, Y_DIM, Z_DIM, "n_modes"])
+        self._sig0 = self.make_local(quantity_factory, [X_DIM, Y_DIM, Z_DIM, "n_modes"])
+        self._bibar = self.make_local(quantity_factory, [X_DIM, Y_DIM, Z_DIM, "n_modes"])
 
         # Stencil
         self.aer_activation = stencil_factory.from_dims_halo(

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/savepoint/translate_aer_activation.py

@@ -1,3 +1,5 @@
+import time
+
 from f90nml import Namelist
 
 from ndsl import Quantity, StencilFactory
@@ -14,8 +16,6 @@ def __init__(self, grid, namelist: Namelist, stencil_factory: StencilFactory):
         self.quantity_factory = grid.quantity_factory
         self._grid = grid
 
-        self.nmodes_quantity_factory = AerActivation.make_nmodes_quantity_factory(self.quantity_factory)
-
         # FloatField Inputs
         self.in_vars["data_vars"] = {
             "AERO_F_DUST": {},
@@ -77,7 +77,7 @@ def make_ij_field(self, data) -> Quantity:
         return qty
 
     def make_nmodes_ijk_field(self, data) -> Quantity:
-        qty = self.nmodes_quantity_factory.empty(
+        qty = self.quantity_factory.empty(
             [X_DIM, Y_DIM, Z_DIM, "n_modes"],
             "n/a",
         )
@@ -116,7 +116,6 @@ def compute(self, inputs):
         qlcn = self.make_ijk_field(inputs["QLCN"])
         qicn = self.make_ijk_field(inputs["QICN"])
 
-        n_modes = inputs["n_modes"]
         ccn_lnd = Float(inputs["CCN_LND"])
         ccn_ocn = Float(inputs["CCN_OCN"])
 
@@ -155,8 +154,40 @@ def compute(self, inputs):
             aero_sigma=aero_sigma,
         )
 
-        return {
-            "NACTL": nactl.view[:, :, :],
-            "NACTI": nacti.view[:, :, :],
-            "NWFA": nwfa.view[:, :, :],
+        output = {
+            "NACTL": nactl.field[:, :, :].copy(),
+            "NACTI": nacti.field[:, :, :].copy(),
+            "NWFA": nwfa.field[:, :, :].copy(),
         }
+
+        # Inline benchmarking - because Aer Activation is a small enough code
+
+        s = time.perf_counter()
+
+        bench_runs = 100
+        for _ in range(0, bench_runs):
+            aer_activation(
+                aero_dgn=aero_dgn,
+                aero_num=aero_num,
+                nacti=nacti,
+                t=t,
+                plo=plmb,
+                qicn=qicn,
+                qils=qils,
+                qlcn=qlcn,
+                qlls=qlls,
+                nn_land=Float(ccn_lnd * 1.0e6),
+                frland=frland,
+                nn_ocean=Float(ccn_ocn * 1.0e6),
+                aero_hygroscopicity=aero_hygroscopicity,
+                nwfa=nwfa,
+                nactl=nactl,
+                vvel=tmp3d,
+                tke=tke,
+                aero_sigma=aero_sigma,
+            )
+
+        e = time.perf_counter()
+        print(f"Aer Activation inline bench: {e - s:.2f}s for {bench_runs} tries")
+
+        return output