Merge remote-tracking branch 'romanc/dsl/fix/microphysics' into dsl/fix/microphysics

FlorianDeconinck · FlorianDeconinck · commit 28cca196ce2f · 2026-05-22T11:02:48.000-04:00
diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/setup.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/convection/GF_2020/setup.py
@@ -1285,6 +1285,11 @@ def __init__(
             },
         )
 
+        # Dev NOTE: this is an orchestration workaround. Direct call to
+        #           `self.saturation_tables.X` fails closure capture for
+        #           argument reconstruction at call time
+        self._esx = self.saturation_tables.esx
+
     def __call__(
         self,
         state: GF2020State,
@@ -1325,7 +1330,7 @@ def __call__(
             area=state.area,
             modified_area=locals.derived_state.modified_area,
             convection_fraction=state.convection_fraction,
-            esx=self.saturation_tables.esx,
+            esx=self._esx,
         )
 
         if state.seed_convection is not None:
diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/GFDL_1M.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/GFDL_1M.py
@@ -116,7 +116,6 @@ def __init__(
             quantity_factory=quantity_factory,
             config=config,
             saturation_tables=saturation_tables,
-            update_tendencies=self._update_tendencies,
         )
 
     def __call__(
diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/driver/ice_cloud.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/driver/ice_cloud.py
@@ -1851,7 +1851,7 @@ def update_precip_total(
     driver_sublimation: FloatField,
 ):
     """
-    ensure information is passed back to the rest of the model
+    Ensure information is passed back to the rest of the model.
     """
     with computation(PARALLEL), interval(...):
         sublimation = sublimation + driver_sublimation
@@ -1960,6 +1960,14 @@ def __init__(
             compute_dims=[I_DIM, J_DIM, K_DIM],
         )
 
+        # Dev NOTE: this is an orchestration workaround. Direct call to
+        #           `self.saturation_tables.X` fails closure capture for
+        #           argument reconstruction at call time
+        self._table2 = self.saturation_tables._table2
+        self._table3 = self.saturation_tables._table3
+        self._des2 = self.saturation_tables._des2
+        self._des3 = self.saturation_tables._des3
+
     def __call__(
         self,
         t: FloatField,
@@ -2032,10 +2040,10 @@ def __call__(
             ccn,
             convection_fraction,
             surface_type,
-            self.saturation_tables.table2,
-            self.saturation_tables.table3,
-            self.saturation_tables.des2,
-            self.saturation_tables.des3,
+            self._table2,
+            self._table3,
+            self._des2,
+            self._des3,
         )
 
         self._update_output(
diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/driver/warm_rain.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/driver/warm_rain.py
@@ -800,6 +800,18 @@ def __init__(
             compute_dims=[I_DIM, J_DIM, K_DIM],
         )
 
+        # Dev NOTE: this is an orchestration workaround. Direct call to
+        #           `self.saturation_tables.X` fails closure capture for
+        #           argument reconstruction at call time
+        self._table1 = (self.saturation_tables.table1,)
+        self._table2 = (self.saturation_tables.table2,)
+        self._table3 = (self.saturation_tables.table3,)
+        self._table4 = (self.saturation_tables.table4,)
+        self._des1 = (self.saturation_tables.des1,)
+        self._des2 = (self.saturation_tables.des2,)
+        self._des3 = (self.saturation_tables.des3,)
+        self._des4 = (self.saturation_tables.des4,)
+
     def __call__(
         self,
         t: FloatField,
@@ -898,14 +910,14 @@ def __call__(
             z_interface=self._locals.z_interface,
             dmass=self._locals.dmass,
             precip_fall=self._locals.precip_fall,
-            table1=self.saturation_tables.table1,
-            table2=self.saturation_tables.table2,
-            table3=self.saturation_tables.table3,
-            table4=self.saturation_tables.table4,
-            des1=self.saturation_tables.des1,
-            des2=self.saturation_tables.des2,
-            des3=self.saturation_tables.des3,
-            des4=self.saturation_tables.des4,
+            table1=self._table1,
+            table2=self._table2,
+            table3=self._table3,
+            table4=self._table4,
+            des1=self._des1,
+            des2=self._des2,
+            des3=self._des3,
+            des4=self._des4,
         )
 
         if not self.config.USE_PPM:
@@ -937,14 +949,14 @@ def __call__(
             w=w,
             rh_limited=rh_limited,
             dmass=self._locals.dmass,
-            table1=self.saturation_tables.table1,
-            table2=self.saturation_tables.table2,
-            table3=self.saturation_tables.table3,
-            table4=self.saturation_tables.table4,
-            des1=self.saturation_tables.des1,
-            des2=self.saturation_tables.des2,
-            des3=self.saturation_tables.des3,
-            des4=self.saturation_tables.des4,
+            table1=self._table1,
+            table2=self._table2,
+            table3=self._table3,
+            table4=self._table4,
+            des1=self._des1,
+            des2=self._des2,
+            des3=self._des3,
+            des4=self._des4,
         )
 
         self._update_outputs(
diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/finalize.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/finalize.py
@@ -7,6 +7,7 @@
 from pyMoist.constants import MAPL_CP, MAPL_GRAV
 from pyMoist.microphysics.GFDL_1M.config import GFDL1MConfig
 from pyMoist.microphysics.GFDL_1M.radiation_coupling import GFDL1MRadiationCoupling
+from pyMoist.microphysics.GFDL_1M.shared_stencils import update_tendencies
 from pyMoist.saturation_tables import GlobalTable_saturation_tables, SaturationVaporPressureTable, saturation_specific_humidity
 from pyMoist.shared.redistribute_clouds import redistribute_clouds
 
@@ -220,14 +221,12 @@ def __init__(
         quantity_factory: QuantityFactory,
         config: GFDL1MConfig,
         saturation_tables: SaturationVaporPressureTable,
-        update_tendencies,
     ):
         # init NDSLRuntime
         super().__init__(stencil_factory)
 
         # make the config, pre-build stencil, and saturation tables visible at runtime
         self.config = config
-        self.update_tendencies = update_tendencies
         self.saturation_tables = saturation_tables
 
         # construct stencils
@@ -268,6 +267,14 @@ def __init__(
             compute_dims=[I_DIM, J_DIM, K_DIM],
         )
 
+        self._update_tendencies = stencil_factory.from_dims_halo(
+            func=update_tendencies,
+            compute_dims=[I_DIM, J_DIM, K_DIM],
+            externals={
+                "DT_MOIST": config.DT_MOIST,
+            },
+        )
+
         self._update_rainwater_source = stencil_factory.from_dims_halo(
             func=update_rainwater_source,
             compute_dims=[I_DIM, J_DIM, K_DIM],
@@ -501,7 +508,7 @@ def __call__(
             liquid_radius=cloud_particle_effective_radius_liquid,
         )
 
-        self.update_tendencies(
+        self._update_tendencies(
             u=u,
             v=v,
             t=t,
diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/radiation_coupling.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/radiation_coupling.py
@@ -77,6 +77,11 @@ def __init__(
         if config.DO_QA:
             ndsl_log.log("[Radiation Coupling] DO_QA option implemented, but untested. " "Running untested code... proceed with caution")
 
+        # Dev NOTE: this is an orchestration workaround. Direct call to
+        #           `self.saturation_tables.X` fails closure capture for
+        #           argument reconstruction at call time
+        self._esx = self.saturation_tables.esx
+
     def __call__(
         self,
         t: Quantity,
@@ -178,5 +183,5 @@ def __call__(
                 pressure=local_p_mb,
                 vapor=mixing_ratio_vapor,
                 humidity=relative_humidity_after_pdf,
-                esx=self.saturation_tables.esx,
+                esx=self._esx,
             )
diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/setup.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/microphysics/GFDL_1M/setup.py
@@ -215,14 +215,12 @@ def compute_estimated_inversion_strength(
         lower_tropospheric_stability (FloatFieldIJ)
         estimated_inversion_strength (FloatFieldIJ)
     """
-    with computation(FORWARD), interval(0, 1):
-        from __externals__ import k_end
-
-        lower_tropospheric_stability = th700 - th.at(K=k_end)
+    with computation(FORWARD), interval(-1, None):
+        lower_tropospheric_stability = th700 - th
         lcl_height = layer_height_above_surface.at(K=lcl_level - 1)
 
         # Simplified single adiabat eq4 of https://doi.org/10.1175/JCLI3988.1
-        t850 = 0.5 * (t.at(K=k_end) + t700)
+        t850 = 0.5 * (t + t700)
         qs850, _ = saturation_specific_humidity(t=t850, p=100.0 * 850.0, esx=esx)
         gamma850 = (1.0 + (MAPL_ALHL * qs850 / (MAPL_RGAS * t850))) / (1.0 + (MAPL_ALHL * MAPL_ALHL * qs850 / (MAPL_CP * MAPL_RVAP * t850 * t850)))
         gamma850 = MAPL_GRAV / MAPL_CP * (1.0 - gamma850)
@@ -528,7 +526,7 @@ def __call__(
             mass=local_mass,
             mass_inverse=local_mass_inverse,
             t=t,
-            esx=self.saturation_tables.esx,
+            esx=self._esx,
             sat=local_saturation_specific_humidity,
             dsat=local_dsaturation_specific_humidity,
             u=u,
diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/translate_tests/microphysics/GFDL_1M/translate_GFDL_1M_Finalize.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/tests/translate_tests/microphysics/GFDL_1M/translate_GFDL_1M_Finalize.py
@@ -1,14 +1,12 @@
 from f90nml import Namelist
 from ndsl import StencilFactory
-from ndsl.constants import I_DIM, J_DIM, K_DIM
 from ndsl.stencils.testing.grid import Grid
 from ndsl.stencils.testing.savepoint import DataLoader
 from ndsl.stencils.testing.translate import TranslateFortranData2Py
 
 from pyMoist.microphysics.GFDL_1M.config import GFDL1MConfig
 from pyMoist.microphysics.GFDL_1M.finalize import GFDL1MFinalize
 from pyMoist.microphysics.GFDL_1M.locals import GFDL1MLocals
-from pyMoist.microphysics.GFDL_1M.shared_stencils import update_tendencies
 from pyMoist.microphysics.GFDL_1M.state import GFDL1MState
 from pyMoist.saturation_tables.tables.main import SaturationVaporPressureTable
 
@@ -146,19 +144,11 @@ def compute(self, inputs):
         state.tendencies.dvdt_micro.field[:] = inputs["dvdt_micro"]
         state.tendencies.dtdt_micro.field[:] = inputs["dtdt_micro"]
 
-        # construct test stencil
-        _update_tendencies = self.stencil_factory.from_dims_halo(
-            func=update_tendencies,
-            compute_dims=[I_DIM, J_DIM, K_DIM],
-            externals={"DT_MOIST": config.DT_MOIST},
-        )
-
         code = GFDL1MFinalize(
             stencil_factory=self.stencil_factory,
             quantity_factory=self.quantity_factory,
             config=config,
             saturation_tables=saturation_tables,
-            update_tendencies=_update_tendencies,
         )
         code(
             t=state.t,

Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,6 @@ def __init__(`
`116`	`116`	`quantity_factory=quantity_factory,`
`117`	`117`	`config=config,`
`118`	`118`	`saturation_tables=saturation_tables,`
`119`		`- update_tendencies=self._update_tendencies,`
`120`	`119`	`)`
`121`	`120`
`122`	`121`	`def __call__(`