Convert storage temporaries in fv3core to Quantity

driver/examples/configs/tropical_read_restart_fortran.yml

@@ -90,7 +90,6 @@ dycore_config:
   tau_v2l: 150.
   fv_sg_adj: 600
   n_sponge: 48
-  p_ref: 101500.
 
 physics_config:
   hydrostatic: false

driver/pace/driver/diagnostics.py

@@ -113,7 +113,7 @@ def diagnostics_factory(self, communicator: pace.util.Communicator) -> Diagnosti
         else:
             fs = pace.util.get_fs(self.path)
             if not fs.exists(self.path):
-                fs.makedirs(self.path)
+                fs.makedirs(self.path, exist_ok=True)
             if self.output_format == "zarr":
                 store = zarr_storage.DirectoryStore(path=self.path)
                 monitor: pace.util.Monitor = pace.util.ZarrMonitor(

driver/pace/driver/driver.py

@@ -471,6 +471,7 @@ def exit_instead_of_build(self):
                 comm=communicator,
                 grid_data=self.state.grid_data,
                 stencil_factory=self.stencil_factory,
+                quantity_factory=self.quantity_factory,
                 damping_coefficients=self.state.damping_coefficients,
                 config=self.config.dycore_config,
                 timestep=self.config.timestep,
@@ -491,6 +492,7 @@ def exit_instead_of_build(self):
             if not config.disable_step_physics:
                 self.dycore_to_physics = update_atmos_state.DycoreToPhysics(
                     stencil_factory=self.stencil_factory,
+                    quantity_factory=self.quantity_factory,
                     dycore_config=self.config.dycore_config,
                     do_dry_convective_adjust=config.do_dry_convective_adjustment,
                     dycore_only=self.config.dycore_only,

driver/pace/driver/initialization.py

@@ -376,7 +376,7 @@ def _update_fortran_restart_pe_peln(state: DriverState) -> None:
     and updates the driver state with values.
     """
 
-    ptop = state.grid_data._vertical_data.ak.view[0]
+    ptop = state.grid_data.ak.view[0]
     pe = state.dycore_state.pe
     peln = state.dycore_state.peln
     delp = state.dycore_state.delp

dsl/pace/dsl/stencil.py

@@ -79,7 +79,7 @@ def report_diff(arg: np.ndarray, numpy_arg: np.ndarray, label) -> str:
     failures_8 = n_points - np.sum(
         np.logical_or(
             nans_match,
-            metric_err < 1e-10,
+            metric_err < 1e-8,
         )
     )
     greatest_error = np.max(metric_err[~np.isnan(metric_err)])

fv3core/examples/standalone/runfile/acoustics.py

@@ -179,7 +179,6 @@ def driver(
             nested,
             stretched_grid,
             dycore_config.acoustic_dynamics,
-            input_data["pfull"],
             input_data["phis"],
             input_data["wsd"],
             state,

fv3core/pace/fv3core/_config.py

@@ -267,9 +267,6 @@ class DynamicalCoreConfig:
     breed_vortex_inline: bool = NamelistDefaults.breed_vortex_inline
     use_old_omega: bool = NamelistDefaults.use_old_omega
     rf_fast: bool = NamelistDefaults.rf_fast
-    p_ref: float = (
-        NamelistDefaults.p_ref
-    )  # Surface pressure used to construct a horizontally-uniform reference
     adiabatic: bool = NamelistDefaults.adiabatic
     nf_omega: int = NamelistDefaults.nf_omega
     fv_sg_adj: int = NamelistDefaults.fv_sg_adj
@@ -375,7 +372,6 @@ def from_namelist(cls, namelist: Namelist) -> "DynamicalCoreConfig":
             breed_vortex_inline=namelist.breed_vortex_inline,
             use_old_omega=namelist.use_old_omega,
             rf_fast=namelist.rf_fast,
-            p_ref=namelist.p_ref,
             adiabatic=namelist.adiabatic,
             nf_omega=namelist.nf_omega,
             fv_sg_adj=namelist.fv_sg_adj,

fv3core/pace/fv3core/stencils/a2b_ord4.py

@@ -11,7 +11,7 @@
     sqrt,
 )
 
-import pace.dsl.gt4py_utils as utils
+import pace.util
 from pace.dsl.dace.orchestration import orchestrate
 from pace.dsl.stencil import GridIndexing, StencilFactory
 from pace.dsl.typing import FloatField, FloatFieldI, FloatFieldIJ
@@ -532,6 +532,7 @@ class AGrid2BGridFourthOrder:
     def __init__(
         self,
         stencil_factory: StencilFactory,
+        quantity_factory: pace.util.QuantityFactory,
         grid_data: GridData,
         grid_type,
         z_dim=Z_DIM,
@@ -564,15 +565,18 @@ def __init__(
 
         self.replace = replace
 
-        def make_storage():
-            return utils.make_storage_from_shape(
-                self._idx.max_shape,
-                backend=self._stencil_config.compilation_config.backend,
-            )
+        self._tmp_qx = quantity_factory.zeros(
+            dims=[X_INTERFACE_DIM, Y_DIM, z_dim], units="unknown"
+        )
+        self._tmp_qy = quantity_factory.zeros(
+            dims=[X_DIM, Y_INTERFACE_DIM, z_dim], units="unknown"
+        )
+        # TODO: the dimensions of tmp_qout_edges may not be correct, verify
+        # with Lucas and either update the code or remove this comment
+        self._tmp_qout_edges = quantity_factory.zeros(
+            dims=[X_DIM, Y_DIM, z_dim], units="unknown"
+        )
 
-        self._tmp_qx = make_storage()
-        self._tmp_qy = make_storage()
-        self._tmp_qout_edges = make_storage()
         _, (z_domain,) = self._idx.get_origin_domain([z_dim])
         corner_domain = (1, 1, z_domain)
 

fv3core/pace/fv3core/stencils/basic_operations.py

@@ -1,8 +1,7 @@
 import gt4py.gtscript as gtscript
-from gt4py.gtscript import FORWARD, PARALLEL, computation, cos, interval, sin
+from gt4py.gtscript import PARALLEL, computation, interval
 
 from pace.dsl.typing import FloatField, FloatFieldIJ
-from pace.util.constants import OMEGA
 
 
 def copy_defn(q_in: FloatField, q_out: FloatField):
@@ -31,17 +30,6 @@ def adjust_divide_stencil(adjustment: FloatField, q_out: FloatField):
         q_out = q_out / adjustment
 
 
-def compute_coriolis_parameter_defn(
-    f: FloatFieldIJ, lon: FloatFieldIJ, lat: FloatFieldIJ, alpha: float
-):
-    with computation(FORWARD), interval(0, 1):
-        f = (
-            2.0
-            * OMEGA
-            * (-1.0 * cos(lon) * cos(lat) * sin(alpha) + sin(lat) * cos(alpha))
-        )
-
-
 @gtscript.function
 def sign(a, b):
     asignb = abs(a)

fv3core/pace/fv3core/stencils/c_sw.py

@@ -7,11 +7,10 @@
     region,
 )
 
-import pace.dsl.gt4py_utils as utils
+import pace.util
 from pace.dsl.dace.orchestration import orchestrate
 from pace.dsl.stencil import StencilFactory
 from pace.dsl.typing import FloatField, FloatFieldIJ
-from pace.fv3core.stencils.basic_operations import compute_coriolis_parameter_defn
 from pace.fv3core.stencils.d2a2c_vect import DGrid2AGrid2CGridVectors
 from pace.stencils import corners
 from pace.util import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM
@@ -481,22 +480,6 @@ def update_y_velocity(
         velocity_c = velocity_c - tmp_flux * flux + rdyc * (ke[0, -1, 0] - ke)
 
 
-def compute_fC(
-    stencil_factory: StencilFactory, lon: FloatFieldIJ, lat: FloatFieldIJ, backend: str
-):
-    """
-    Compute the coriolis parameter on the C-grid
-    """
-    fC = utils.make_storage_from_shape(lon.shape, backend=backend)
-    fC_stencil = stencil_factory.from_dims_halo(
-        compute_coriolis_parameter_defn,
-        compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM],
-        compute_halos=(3, 3),
-    )
-    fC_stencil(fC, lon, lat, 0.0)
-    return fC
-
-
 class CGridShallowWaterDynamics:
     """
     Fortran name is c_sw
@@ -505,35 +488,23 @@ class CGridShallowWaterDynamics:
     def __init__(
         self,
         stencil_factory: StencilFactory,
+        quantity_factory: pace.util.QuantityFactory,
         grid_data: GridData,
         nested: bool,
         grid_type: int,
         nord: int,
     ):
         orchestrate(obj=self, config=stencil_factory.config.dace_config)
-        grid_indexing = stencil_factory.grid_indexing
         self.grid_data = grid_data
         self._dord4 = True
-        self._fC = compute_fC(
-            stencil_factory,
-            self.grid_data.lon,
-            self.grid_data.lat,
-            backend=stencil_factory.backend,
-        )
-        origin_halo1 = (grid_indexing.isc - 1, grid_indexing.jsc - 1, 0)
-        self.delpc = utils.make_storage_from_shape(
-            grid_indexing.max_shape,
-            origin=origin_halo1,
-            backend=stencil_factory.backend,
-        )
+        self._fC = self.grid_data.fC
+        # TODO: double-check the dimensions on these, they may be incorrect
+        # as they are only documentation and not used by the code
+        self.delpc = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], units="unknown")
         """
         pressure thickness on c-grid forward step
         """
-        self.ptc = utils.make_storage_from_shape(
-            grid_indexing.max_shape,
-            origin=origin_halo1,
-            backend=stencil_factory.backend,
-        )
+        self.ptc = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], units="unknown")
         """
         potential temperature on c-grid forward step
         """
@@ -545,23 +516,23 @@ def __init__(
 
         self._D2A2CGrid_Vectors = DGrid2AGrid2CGridVectors(
             stencil_factory,
-            grid_data,
-            nested,
-            grid_type,
-            self._dord4,
+            quantity_factory=quantity_factory,
+            grid_data=grid_data,
+            nested=nested,
+            grid_type=grid_type,
+            dord4=self._dord4,
         )
 
-        def make_storage():
-            return utils.make_storage_from_shape(
-                grid_indexing.max_shape,
-                backend=stencil_factory.backend,
-            )
+        def make_quantity() -> pace.util.Quantity:
+            return quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], units="unknown")
 
-        self._tmp_ke = make_storage()
-        self._tmp_vort = make_storage()
-        self._tmp_fx = make_storage()
-        self._tmp_fx1 = make_storage()
-        self._tmp_fx2 = make_storage()
+        # TODO: double-check the dimensions on these, they may be incorrect
+        # as they are only documentation and not used by the code
+        self._tmp_ke = make_quantity()
+        self._tmp_vort = make_quantity()
+        self._tmp_fx = make_quantity()
+        self._tmp_fx1 = make_quantity()
+        self._tmp_fx2 = make_quantity()
 
         if nord > 0:
             self._divergence_corner = stencil_factory.from_dims_halo(

fv3core/pace/fv3core/stencils/d2a2c_vect.py

@@ -1,7 +1,7 @@
 import gt4py.gtscript as gtscript
 from gt4py.gtscript import PARALLEL, computation, horizontal, interval, region
 
-import pace.dsl.gt4py_utils as utils
+import pace.util
 from pace.dsl.dace.orchestration import orchestrate
 from pace.dsl.stencil import StencilFactory
 from pace.dsl.typing import FloatField, FloatFieldIJ
@@ -385,6 +385,7 @@ class DGrid2AGrid2CGridVectors:
     def __init__(
         self,
         stencil_factory: StencilFactory,
+        quantity_factory: pace.util.QuantityFactory,
         grid_data: GridData,
         nested: bool,
         grid_type: int,
@@ -418,16 +419,8 @@ def __init__(
         npt = 4 if not nested else 0
         if npt > grid_indexing.domain[0] - 1 or npt > grid_indexing.domain[1] - 1:
             npt = 0
-        self._utmp = utils.make_storage_from_shape(
-            grid_indexing.max_shape,
-            grid_indexing.origin_full(),
-            backend=stencil_factory.backend,
-        )
-        self._vtmp = utils.make_storage_from_shape(
-            grid_indexing.max_shape,
-            grid_indexing.origin_full(),
-            backend=stencil_factory.backend,
-        )
+        self._utmp = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], units="m/s")
+        self._vtmp = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], units="m/s")
 
         js1 = npt + OFFSET if grid_indexing.south_edge else grid_indexing.jsc - 1
         je1 = ny - npt if grid_indexing.north_edge else grid_indexing.jec + 1