From 70ea601b8373a48afc5fe1f63732e0e1a32cc3bd Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 29 Nov 2024 14:11:59 -0500 Subject: [PATCH 1/3] 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 50ffba2..edfb1bf 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 dbc082d..1ac71a8 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/3] 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 1ac71a8..708d1e9 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/3] 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 708d1e9..3a65c41 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: