Fix f32 calculation for QSat tables (#1007)

* Fix f32 calculation for QSat tables
Remove NetCDF table method (exact only remains)

* Use local `f64` cast

* Lint

* Make DQSat a proper field
Turn trigger to externals for stencil version of QSat

* Lint

---------

Co-authored-by: Florian Deconinck <florian.deconinck@gmail.com>

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/saturation/__init__.py

@@ -1,2 +1,2 @@
 from .qsat import QSat, QSat_Float, QSat_Float_Ice, QSat_Float_Liquid, QSat_FloatField
-from .types import SaturationFormulation, TableMethod
+from .types import SaturationFormulation

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/saturation/constants.py

@@ -13,7 +13,7 @@
 TMAXTBL = Float(333.0)
 TMINLQU = MAPL_TICE - Float(40.0)
 DEGSUBS = np.int32(100)
-DELTA_T = Float(1.0 / DEGSUBS)
+DELTA_T = Float(1.0) / Float(DEGSUBS)
 TABLESIZE = np.int32(TMAXTBL - TMINTBL) * DEGSUBS + 1
 TMIX = Float(-20.0)
 ESFAC = MAPL_H2OMW / MAPL_AIRMW

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/saturation/qsat.py

@@ -1,9 +1,7 @@
 import copy
-import os
 from typing import Optional
 
 import gt4py.cartesian.gtscript as gtscript
-import xarray as xr
 from gt4py.cartesian.gtscript import PARALLEL, computation, floor, interval
 
 from ndsl import QuantityFactory, StencilFactory, orchestrate
@@ -21,7 +19,8 @@
     TMINTBL,
     TMIX,
 )
-from pyMoist.saturation.types import SaturationFormulation, TableMethod
+from pyMoist.saturation.table import get_table
+from pyMoist.saturation.types import SaturationFormulation
 
 
 # FloatField with extra dimension initialized to handle table data
@@ -194,17 +193,16 @@ def QSat_FloatField(
     QSAT: FloatField,
     DQSAT: FloatField,
     RAMP: FloatField,
-    PASCALS_trigger: bool,
-    RAMP_trigger: bool,
-    DQSAT_trigger: bool,
 ):
+    from __externals__ import FILL_DQSAT, USE_PASCALS, USE_RAMP
+
     with computation(PARALLEL), interval(...):
-        if RAMP_trigger:
+        if USE_RAMP:
             URAMP = -abs(RAMP)
         else:
             URAMP = TMIX
 
-        if PASCALS_trigger:
+        if USE_PASCALS:
             PP = PL
         else:
             PP = PL * 100.0
@@ -230,12 +228,12 @@ def QSat_FloatField(
             QSAT = (TI - IT) * DQ + ese[0][IT_MINUS_1]  # type: ignore
         if PP <= QSAT:
             QSAT = MAX_MIXING_RATIO
-            if DQSAT_trigger:
+            if FILL_DQSAT:
                 DQSAT = 0.0
         else:
             DD = 1.0 / (PP - (1.0 - ESFAC) * QSAT)
             QSAT = ESFAC * QSAT * DD
-            if DQSAT_trigger:
+            if FILL_DQSAT:
                 DQSAT = ESFAC * DQ * DEGSUBS * PP * (DD * DD)
 
 
@@ -269,7 +267,9 @@ def __init__(
         stencil_factory: StencilFactory,
         quantity_factory: QuantityFactory,
         formulation: SaturationFormulation = SaturationFormulation.Staars,
-        table_method: TableMethod = TableMethod.NetCDF,
+        use_ramp: bool = False,
+        use_pascals: bool = False,
+        fill_dqsat: bool = False,
     ) -> None:
         self.extra_dim_quantity_factory = self.make_extra_dim_quantity_factory(
             quantity_factory
@@ -279,34 +279,10 @@ def __init__(
         self.esw = self.extra_dim_quantity_factory.zeros([Z_DIM, "table_axis"], "n/a")
         self.esx = self.extra_dim_quantity_factory.zeros([Z_DIM, "table_axis"], "n/a")
 
-        if table_method == TableMethod.NetCDF:
-            with xr.open_dataset(
-                os.path.join(os.path.dirname(__file__), "netCDFs", "QSat_Tables.nc")
-            ) as ds:
-                ese_array = ds.data_vars["ese"].values[0, 0, :]
-                esw_array = ds.data_vars["esw"].values[0, 0, :]
-                esx_array = ds.data_vars["esx"].values[0, 0, :]
-
-                self.ese.view[:] = ese_array
-                self.esw.view[:] = esw_array
-                self.esx.view[:] = esx_array
-        elif table_method == TableMethod.Computation:
-            raise NotImplementedError(
-                """[QSat]  Calculated tables are incorrect due to differences between
-                C and Fortran calculations.\n
-                For now the tables are being read in from a netCDF file.\n
-                See https://github.com/GEOS-ESM/SMT-Nebulae/issues/88
-                for more information."""
-            )
-            # self.table = get_table(formulation)
-
-            # self.ese.view[:] = self.table.ese
-            # self.esw.view[:] = self.table.esw
-            # self.esx.view[:] = self.table.esx
-        else:
-            raise NotImplementedError(
-                f"[QSat] Unknown {table_method} to create estimation table"
-            )
+        self.table = get_table(formulation)
+        self.ese.view[:] = self.table.ese
+        self.esw.view[:] = self.table.esw
+        self.esx.view[:] = self.table.esx
 
         self._RAMP = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
         self._DQSAT = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a")
@@ -317,6 +293,11 @@ def __init__(
         self._QSat_FloatField = stencil_factory.from_dims_halo(
             func=QSat_FloatField,
             compute_dims=[X_DIM, Y_DIM, Z_DIM],
+            externals={
+                "USE_RAMP": use_ramp,
+                "USE_PASCALS": use_pascals,
+                "FILL_DQSAT": fill_dqsat,
+            },
         )
 
     @staticmethod
@@ -333,16 +314,13 @@ def __call__(
         self,
         T: FloatField,
         PL: FloatField,
-        RAMP: Optional[FloatField] = None,
-        PASCALS: bool = False,
-        DQSAT: bool = False,
+        DQSat: Optional[FloatField] = None,
         use_table_lookup: bool = True,
     ):
-        if RAMP is None:
-            RAMP = self._RAMP
-            RAMP_trigger = False
+        if DQSat:
+            dqsat = DQSat
         else:
-            RAMP_trigger = True
+            dqsat = self._DQSAT
 
         if use_table_lookup:
             self._QSat_FloatField(
@@ -351,11 +329,8 @@ def __call__(
                 T,
                 PL,
                 self.QSat,
-                self._DQSAT,
-                RAMP,
-                PASCALS,
-                RAMP_trigger,
-                DQSAT,
+                dqsat,
+                self._RAMP,
             )
 
         if not use_table_lookup:

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/saturation/qsat_ice.py

@@ -13,42 +13,53 @@
 from pyMoist.saturation.types import SaturationFormulation
 
 
+f64 = np.float64
+
 TMINSTR = Float(-95.0)
 TMINICE = MAPL_TICE + TMINSTR
 
-
 TMINSTR = Float(-95.0)
 TSTARR1 = Float(-75.0)
 TSTARR2 = Float(-65.0)
 TSTARR3 = Float(-50.0)
 TSTARR4 = Float(-40.0)
 TMAXSTR = Float(+60.0)
 
-DI = [57518.5606e08, 2.01889049, 3.56654, 20.947031]
-CI = [9.550426, -5723.265, 3.53068, -0.00728332]
+DI = [
+    f64(Float(57518.5606e08)),
+    f64(Float(2.01889049)),
+    f64(Float(3.56654)),
+    f64(Float(20.947031)),
+]
+CI = [
+    f64(Float(9.550426)),
+    f64(Float(-5723.265)),
+    f64(Float(3.53068)),
+    f64(Float(-0.00728332)),
+]
 
 # 64-bit float in Fortran
-S16 = 0.516000335e-11 * 100.0
-S15 = 0.276961083e-8 * 100.0
-S14 = 0.623439266e-6 * 100.0
-S13 = 0.754129933e-4 * 100.0
-S12 = 0.517609116e-2 * 100.0
-S11 = 0.191372282e0 * 100.0
-S10 = 0.298152339e1 * 100.0
-S26 = 0.314296723e-10 * 100.0
-S25 = 0.132243858e-7 * 100.0
-S24 = 0.236279781e-5 * 100.0
-S23 = 0.230325039e-3 * 100.0
-S22 = 0.129690326e-1 * 100.0
-S21 = 0.401390832e0 * 100.0
-S20 = 0.535098336e1 * 100.0
-BI6 = 1.838826904e-10 * 100.0
-BI5 = 4.838803174e-8 * 100.0
-BI4 = 5.824720280e-6 * 100.0
-BI3 = 4.176223716e-4 * 100.0
-BI2 = 1.886013408e-2 * 100.0
-BI1 = 5.034698970e-1 * 100.0
-BI0 = 6.109177956e0 * 100.0
+S16 = f64(Float(0.516000335e-11) * Float(100.0))
+S15 = f64(Float(0.276961083e-8) * Float(100.0))
+S14 = f64(Float(0.623439266e-6) * Float(100.0))
+S13 = f64(Float(0.754129933e-4) * Float(100.0))
+S12 = f64(Float(0.517609116e-2) * Float(100.0))
+S11 = f64(Float(0.191372282e0) * Float(100.0))
+S10 = f64(Float(0.298152339e1) * Float(100.0))
+S26 = f64(Float(0.314296723e-10) * Float(100.0))
+S25 = f64(Float(0.132243858e-7) * Float(100.0))
+S24 = f64(Float(0.236279781e-5) * Float(100.0))
+S23 = f64(Float(0.230325039e-3) * Float(100.0))
+S22 = f64(Float(0.129690326e-1) * Float(100.0))
+S21 = f64(Float(0.401390832e0) * Float(100.0))
+S20 = f64(Float(0.535098336e1) * Float(100.0))
+BI6 = f64(Float(1.838826904e-10) * Float(100.0))
+BI5 = f64(Float(4.838803174e-8) * Float(100.0))
+BI4 = f64(Float(5.824720280e-6) * Float(100.0))
+BI3 = f64(Float(4.176223716e-4) * Float(100.0))
+BI2 = f64(Float(1.886013408e-2) * Float(100.0))
+BI1 = f64(Float(5.034698970e-1) * Float(100.0))
+BI0 = f64(Float(6.109177956e0) * Float(100.0))
 
 
 def _saturation_formulation(
@@ -69,7 +80,7 @@ def _saturation_formulation(
                 TT
                 * (TT * (TT * (TT * (TT * (TT * S16 + S15) + S14) + S13) + S12) + S11)
                 + S10
-            ) + (1.0 - W) * (
+            ) + (Float(1.0) - W) * (
                 TT
                 * (TT * (TT * (TT * (TT * (TT * S26 + S25) + S24) + S23) + S22) + S21)
                 + S20
@@ -86,7 +97,7 @@ def _saturation_formulation(
                 TT
                 * (TT * (TT * (TT * (TT * (TT * S26 + S25) + S24) + S23) + S22) + S21)
                 + S20
-            ) + (1.0 - W) * (
+            ) + (Float(1.0) - W) * (
                 TT
                 * (TT * (TT * (TT * (TT * (TT * BI6 + BI5) + BI4) + BI3) + BI2) + BI1)
                 + BI0
@@ -119,14 +130,14 @@ def qsat_ice_scalar_exact(
     else:
         TI = temperature
 
-    DX = 0.0  # only calulcated when DQ is not none
+    DX = Float(0.0)  # only calulcated when DQ is not none
     EX = _saturation_formulation(formulation, TI)
 
     if DQ is not None:
         if temperature < TMINICE:
-            DDQ = 0.0
+            DDQ = Float(0.0)
         elif temperature > MAPL_TICE:
-            DDQ = 0.0
+            DDQ = Float(0.0)
         else:
             if PL > EX:
                 DD = EX
@@ -136,16 +147,16 @@ def qsat_ice_scalar_exact(
                 EX = DD
     if PL is not None:
         if PL > EX:
-            DD = ESFAC / (PL - (1.0 - ESFAC) * EX)
+            DD = ESFAC / (PL - (Float(1.0) - ESFAC) * EX)
             EX = EX * DD
             if DQ is not None:
                 DX = DDQ * ERFAC * PL * DD * DD
         else:
             EX = MAX_MIXING_RATIO
             if DQ is not None:
-                DX = 0.0
+                DX = Float(0.0)
     else:
         if DQ is not None:
-            DX = DDQ * (1.0 / DELTA_T)
+            DX = DDQ * (Float(1.0) / DELTA_T)
 
     return EX, TI, DX