From 70ea601b8373a48afc5fe1f63732e0e1a32cc3bd Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 29 Nov 2024 14:11:59 -0500 Subject: [PATCH 1/5] Move RFF calculation to f64 Fix the implementation check for RayleighDamping and move it closer to the code --- pyFV3/stencils/fv_dynamics.py | 6 ------ pyFV3/stencils/ray_fast.py | 25 +++++++++++++++++++------ 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 50ffba27..edfb1bfb 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/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" diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index dbc082d4..1ac71a8a 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -10,6 +10,7 @@ log, region, sin, + f64, ) import ndsl.constants as constants @@ -18,7 +19,7 @@ from ndsl.dsl.typing import Float, FloatField, FloatFieldK -SDAY = 86400.0 +SDAY = Float(86400.0) # NOTE: The fortran version of this computes rf in the first timestep only. Then @@ -36,7 +37,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 @@ -155,14 +156,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,7 +188,7 @@ 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, }, @@ -191,7 +204,7 @@ def __call__( dt: Float, ptop: Float, ): - rf_cutoff_nudge = self._rf_cutoff + min(100.0, 10.0 * ptop) + rf_cutoff_nudge = self._rf_cutoff + min(Float(100.0), Float(10.0) * ptop) self._ray_fast_wind_compute( u, From 4c8403abef8467c5dbdc2837da1fcd3ae78d33e3 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 29 Nov 2024 14:16:03 -0500 Subject: [PATCH 2/5] lint --- pyFV3/stencils/ray_fast.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index 1ac71a8a..708d1e9f 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -5,12 +5,12 @@ FORWARD, PARALLEL, computation, + f64, horizontal, interval, log, region, sin, - f64, ) import ndsl.constants as constants From bad27a887e83b98ca2f974096c2a9b5d429ac3d8 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 2 Dec 2024 16:27:36 -0500 Subject: [PATCH 3/5] Use `ndsl.constants.SECONDS_PER_DAY` --- pyFV3/stencils/ray_fast.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index 708d1e9f..3a65c417 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -15,13 +15,10 @@ import ndsl.constants as constants from ndsl import StencilFactory, orchestrate -from ndsl.constants import X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM +from ndsl.constants import SECONDS_PER_DAY, X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField, FloatFieldK -SDAY = Float(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. @@ -77,7 +74,7 @@ def ray_fast_wind_compute( 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) + rf = compute_rff_vals(pfull, dt, rf_cutoff, tau * SECONDS_PER_DAY, ptop) with computation(FORWARD): with interval(0, 1): if pfull < rf_cutoff_nudge: From e4d65df40f45fb714b923785b6754edf774ebc8b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 5 Mar 2025 15:54:52 -0500 Subject: [PATCH 4/5] Move damping increment to own K stencils Trigger damping increment calculation exactly once --- pyFV3/stencils/ray_fast.py | 81 +++++++++++++++++++++++++++++--------- 1 file changed, 63 insertions(+), 18 deletions(-) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index 3a65c417..a28e8946 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -15,7 +15,15 @@ import ndsl.constants as constants from ndsl import StencilFactory, orchestrate -from ndsl.constants import SECONDS_PER_DAY, 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 @@ -43,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, ): """ @@ -68,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 * SECONDS_PER_DAY, ptop) with computation(FORWARD): with interval(0, 1): if pfull < rf_cutoff_nudge: @@ -191,6 +209,29 @@ def __init__( }, ) + # 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, @@ -203,13 +244,17 @@ def __call__( ): 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, ) From e077ab454340a476075396bc146f757a70071b55 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:29:08 -0400 Subject: [PATCH 5/5] Docs --- pyFV3/stencils/ray_fast.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index a28e8946..e1dea877 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -242,6 +242,16 @@ def __call__( dt: Float, ptop: Float, ): + """ + 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: