[WIP] Rayleigh Damping mixed precision fix

pyFV3/stencils/fv_dynamics.py

@@ -483,12 +483,6 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool):
                 "Dynamical Core (fv_dynamics): compute total energy is not implemented"
             )
 
-        if (not self.config.rf_fast) and self.config.tau != 0:
-            raise NotImplementedError(
-                "Dynamical Core (fv_dynamics): Rayleigh_Super,"
-                " called when rf_fast=False and tau !=0, is not implemented"
-            )
-
         if self.config.adiabatic and self.config.kord_tm > 0:
             raise NotImplementedError(
                 "Dynamical Core (fv_dynamics): Adiabatic with positive kord_tm"

pyFV3/stencils/ray_fast.py

@@ -5,6 +5,7 @@
     FORWARD,
     PARALLEL,
     computation,
+    f64,
     horizontal,
     interval,
     log,
@@ -14,13 +15,18 @@
 
 import ndsl.constants as constants
 from ndsl import StencilFactory, orchestrate
-from ndsl.constants import X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM
+from ndsl.boilerplate import get_factories_single_tile
+from ndsl.constants import (
+    SECONDS_PER_DAY,
+    X_INTERFACE_DIM,
+    Y_INTERFACE_DIM,
+    Z_DIM,
+    X_DIM,
+    Y_DIM,
+)
 from ndsl.dsl.typing import Float, FloatField, FloatFieldK
 
 
-SDAY = 86400.0
-
-
 # NOTE: The fortran version of this computes rf in the first timestep only. Then
 # rf_initialized let's you know you can skip it. Here we calculate it every
 # time.
@@ -36,7 +42,7 @@ def compute_rf_vals(pfull, bdt, rf_cutoff, tau0, ptop):
 @gtscript.function
 def compute_rff_vals(pfull, dt, rf_cutoff, tau0, ptop):
     rffvals = compute_rf_vals(pfull, dt, rf_cutoff, tau0, ptop)
-    rffvals = 1.0 / (1.0 + rffvals)
+    rffvals = f64(1.0) / (f64(1.0) + rffvals)
     return rffvals
 
 
@@ -45,14 +51,31 @@ def dm_layer(rf, dp, wind):
     return (1.0 - rf) * dp * wind
 
 
+def ray_fast_damping_increment(
+    pfull: FloatFieldK,  # type:ignore
+    dt: Float,  # type:ignore
+    ptop: Float,  # type:ignore
+    rf: FloatField,  # type:ignore
+):
+    """rf is rayleigh damping increment, fraction of vertical velocity
+    left after doing rayleigh damping (w -> w * rf)
+    """
+    from __externals__ import rf_cutoff, tau
+
+    with computation(PARALLEL), interval(...):
+        if pfull < rf_cutoff:
+            # rf is rayleigh damping increment, fraction of vertical velocity
+            # left after doing rayleigh damping (w -> w * rf)
+            rf = compute_rff_vals(pfull, dt, rf_cutoff, tau * SECONDS_PER_DAY, ptop)
+
+
 def ray_fast_wind_compute(
     u: FloatField,
     v: FloatField,
     w: FloatField,
     delta_p_ref: FloatFieldK,  # reference delta pressure
     pfull: FloatFieldK,  # input layer pressure reference?
-    dt: Float,
-    ptop: Float,
+    rf: FloatFieldK,
     rf_cutoff_nudge: Float,
 ):
     """
@@ -70,13 +93,6 @@ def ray_fast_wind_compute(
     from __externals__ import hydrostatic, local_ie, local_je, rf_cutoff, tau
 
     # dm_stencil
-    with computation(PARALLEL), interval(...):
-        # TODO -- in the fortran model rf is only computed once, repeating
-        # the computation every time ray_fast is run is inefficient
-        if pfull < rf_cutoff:
-            # rf is rayleigh damping increment, fraction of vertical velocity
-            # left after doing rayleigh damping (w -> w * rf)
-            rf = compute_rff_vals(pfull, dt, rf_cutoff, tau * SDAY, ptop)
     with computation(FORWARD):
         with interval(0, 1):
             if pfull < rf_cutoff_nudge:
@@ -155,14 +171,26 @@ class RayleighDamping:
     Fortran name: ray_fast.
     """
 
-    def __init__(self, stencil_factory: StencilFactory, rf_cutoff, tau, hydrostatic):
+    def __init__(
+        self,
+        stencil_factory: StencilFactory,
+        rf_cutoff: Float,
+        tau: Float,
+        hydrostatic: bool,
+    ):
         orchestrate(obj=self, config=stencil_factory.config.dace_config)
         grid_indexing = stencil_factory.grid_indexing
-        self._rf_cutoff = rf_cutoff
+        self._rf_cutoff = Float(rf_cutoff)
         origin, domain = grid_indexing.get_origin_domain(
             [X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM]
         )
 
+        if tau == 0:
+            raise NotImplementedError(
+                "Dynamical Core (fv_dynamics): RayleighDamping,"
+                " with tau <= 0, is not implemented"
+            )
+
         ax_offsets = grid_indexing.axis_offsets(origin, domain)
         local_axis_offsets = {}
         for axis_offset_name, axis_offset_value in ax_offsets.items():
@@ -175,12 +203,35 @@ def __init__(self, stencil_factory: StencilFactory, rf_cutoff, tau, hydrostatic)
             domain=domain,
             externals={
                 "hydrostatic": hydrostatic,
-                "rf_cutoff": rf_cutoff,
+                "rf_cutoff": self._rf_cutoff,
                 "tau": tau,
                 **local_axis_offsets,
             },
         )
 
+        # We compute the damping increment once using a trick to write a
+        # FloatFieldK as a (1, 1, K) 3D writable Field
+        K_stencil_factory, K_quantity_factory = get_factories_single_tile(
+            1,
+            1,
+            domain[2],
+            0,
+            stencil_factory.backend,
+        )
+        self._ray_fast_damping_increment = K_stencil_factory.from_origin_domain(
+            ray_fast_damping_increment,
+            origin=(0, 0, origin[2]),
+            domain=(1, 1, domain[2]),
+            externals={
+                "rf_cutoff": self._rf_cutoff,
+                "tau": tau,
+            },
+        )
+        self._damping_increment = K_quantity_factory.ones(
+            [X_DIM, Y_DIM, Z_DIM], units="n/a"
+        )
+        self._initialize_damping_increment = False
+
     def __call__(
         self,
         u: FloatField,
@@ -191,15 +242,29 @@ def __call__(
         dt: Float,
         ptop: Float,
     ):
-        rf_cutoff_nudge = self._rf_cutoff + min(100.0, 10.0 * ptop)
+        """
+        Args:
+            u (inout)
+            v (inout)
+            w (inout)
+            dp (in)
+            pfull (in)
+            dt (in)
+            ptop (in)
+        """
+        rf_cutoff_nudge = self._rf_cutoff + min(Float(100.0), Float(10.0) * ptop)
 
+        if not self._initialize_damping_increment:
+            self._ray_fast_damping_increment(
+                pfull=pfull, dt=dt, ptop=ptop, rf=self._damping_increment
+            )
+            self._initialize_damping_increment = True
         self._ray_fast_wind_compute(
-            u,
-            v,
-            w,
-            dp,
-            pfull,
-            dt,
-            ptop,
-            rf_cutoff_nudge,
+            u=u,
+            v=v,
+            w=w,
+            delta_p_ref=dp,
+            pfull=pfull,
+            rf=self._damping_increment.view[0, 0, :],
+            rf_cutoff_nudge=rf_cutoff_nudge,
         )