[NASA] [Feature] Guarding against unimplemented configuration (#40) (N…

…OAA-GFDL#48)

* [Feature] Guarding against unimplemented configuration (#40)

Guarding against unimplemented namelists options:
- a2b_ord4
- d_sw
- fv_dynamics
- fv_subgridz
- neg_adj3
- divergence damping
- xppm
- yppm

Misc:
- Fix `netcdf_monitor` not mkdir the directory
- Add `as_dict` to the dycore state to dump the dycore more easily

* Unused assert

* Update fv3core/pace/fv3core/stencils/yppm.py

Co-authored-by: Oliver Elbert <[email protected]>

* Update fv3core/pace/fv3core/stencils/xppm.py

Co-authored-by: Oliver Elbert <[email protected]>

* Change NotImplemented to ValueError for n_sponge<3

* lint

---------

Co-authored-by: Oliver Elbert <[email protected]>

fv3core/pace/fv3core/dycore_state.py

@@ -1,11 +1,12 @@
-from dataclasses import dataclass, field, fields
-from typing import Any, Mapping
+from dataclasses import asdict, dataclass, field, fields
+from typing import Any, Dict, Mapping, Union
 
 import xarray as xr
 
 import pace.dsl.gt4py_utils as gt_utils
 import pace.util
 from pace.dsl.typing import Float
+from pace.util.quantity import Quantity
 
 
 @dataclass()
@@ -450,6 +451,12 @@ def xr_dataset(self):
     def __getitem__(self, item):
         return getattr(self, item)
 
+    def as_dict(self, quantity_only=True) -> Dict[str, Union[Quantity, int]]:
+        if quantity_only:
+            return {k: v for k, v in asdict(self).items() if isinstance(v, Quantity)}
+        else:
+            return {k: v for k, v in asdict(self).items()}
+
 
 TRACER_PROPERTIES = {
     "specific_humidity": {

fv3core/pace/fv3core/stencils/a2b_ord4.py

@@ -548,7 +548,11 @@ def __init__(
             replace: boolean, update qin to the B grid as well
         """
         orchestrate(obj=self, config=stencil_factory.config.dace_config)
-        assert grid_type in [0, 4]
+        if grid_type != 0 and grid_type != 4:
+            raise RuntimeError(
+                "A-Grid to B-Grid 4th order (a2b_ord4):"
+                f" grid type {grid_type} is not implemented. 0 and 4 available."
+            )
         self._idx: GridIndexing = stencil_factory.grid_indexing
         self._stencil_config = stencil_factory.config
         self.replace = replace
@@ -720,7 +724,6 @@ def __call__(self, qin: FloatField, qout: FloatField):
         """
 
         if self.grid_type < 3:
-
             self._sw_corner_stencil(
                 qin,
                 qout,

fv3core/pace/fv3core/stencils/d_sw.py

@@ -764,19 +764,45 @@ def __init__(
 
         self.grid_indexing = stencil_factory.grid_indexing
         self._grid_type = config.grid_type
-        assert not config.inline_q, "inline_q not yet implemented"
-        assert (
-            config.d_ext <= 0
-        ), "untested d_ext > 0. need to call a2b_ord2, not yet implemented"
-        assert (column_namelist["damp_vt"].view[:] > dcon_threshold).all()
+        if config.inline_q:
+            raise NotImplementedError(
+                "D-Grid Shallow Water Lagrangian Dynamics (D_SW):"
+                " inline_q not implemented."
+            )
+        if config.d_ext > 0:
+            raise NotImplementedError(
+                "D-Grid Shallow Water Lagrangian Dynamics (D_SW):"
+                " untested d_ext > 0. need to call a2b_ord2, not implemented."
+            )
         # TODO: in theory, we should check if damp_vt > 1e-5 for each k-level and
         # only compute delnflux for k-levels where this is true
-        assert (column_namelist["damp_w"].view[:] > dcon_threshold).all()
+        all_damp_vt_above_dcon = (
+            column_namelist["damp_vt"].view[:] > dcon_threshold
+        ).all()
+        if not all_damp_vt_above_dcon:
+            raise NotImplementedError(
+                "D-Grid Shallow Water Lagrangian Dynamics (D_SW):"
+                " damp_vt misconfiguration, some are above a d_con of"
+                f" {dcon_threshold}."
+            )
         # TODO: in theory, we should check if damp_w > 1e-5 for each k-level and
         # only compute delnflux for k-levels where this is true
+        all_damp_w_above_dcon = (
+            column_namelist["damp_w"].view[:] > dcon_threshold
+        ).all()
+        if not all_damp_w_above_dcon:
+            raise NotImplementedError(
+                "D-Grid Shallow Water Lagrangian Dynamics (D_SW):"
+                f" damp_w misconfiguration, some are above a d_con of {dcon_threshold}."
+            )
 
         # only compute for k-levels where this is true
         self.hydrostatic = config.hydrostatic
+        if self.hydrostatic:
+            raise NotImplementedError(
+                "D-Grid Shallow Water Lagrangian Dynamics (D_SW):"
+                " Hydrostatic is not implemented"
+            )
 
         def make_quantity():
             return quantity_factory.zeros(

fv3core/pace/fv3core/stencils/divergence_damping.py

@@ -324,8 +324,8 @@ def __init__(
             config=stencil_factory.config.dace_config,
         )
         self.grid_indexing = stencil_factory.grid_indexing
-        assert not nested, "nested not implemented"
-        # assert grid_type < 3, "Not implemented, grid_type>=3"
+        if nested:
+            raise NotImplementedError("Divergence Dampoing: nested not implemented.")
         # TODO: make dddmp a compile-time external, instead of runtime scalar
         self._dddmp = dddmp
         # TODO: make da_min_c a compile-time external, instead of runtime scalar

fv3core/pace/fv3core/stencils/dyn_core.py

@@ -429,9 +429,12 @@ def __init__(
             self.checkpointer = checkpointer
         grid_indexing = stencil_factory.grid_indexing
         self.config = config
-        assert config.d_ext == 0, "d_ext != 0 is not implemented"
-        assert config.beta == 0, "beta != 0 is not implemented"
-        assert not config.use_logp, "use_logp=True is not implemented"
+        if config.d_ext != 0:
+            raise RuntimeError("Acoustics (dyn_core): d_ext != 0 is not implemented")
+        if config.beta != 0:
+            raise RuntimeError("Acoustics (dyn_core): beta != 0 is not implemented")
+        if config.use_logp:
+            raise RuntimeError("Acoustics (dyn_core): use_logp=True is not implemented")
         self._da_min = damping_coefficients.da_min
         self.grid_data = grid_data
         self._ptop = grid_data.ptop

fv3core/pace/fv3core/stencils/fv_dynamics.py

@@ -186,8 +186,17 @@ def __init__(
         nested = False
         stretched_grid = False
         grid_indexing = stencil_factory.grid_indexing
-        assert config.moist_phys, "fvsetup is only implemented for moist_phys=true"
-        assert config.nwat == 6, "Only nwat=6 has been implemented and tested"
+        if not config.moist_phys:
+            raise NotImplementedError(
+                "Dynamical core (fv_dynamics):"
+                " fvsetup is only implemented for moist_phys=true."
+            )
+        if config.nwat != 6:
+            raise NotImplementedError(
+                "Dynamical core (fv_dynamics):"
+                f" nwat=={config.nwat} is not implemented."
+                " Only nwat=6 has been implemented."
+            )
         self.comm_rank = comm.rank
         self.grid_data = grid_data
         self.grid_indexing = grid_indexing
@@ -286,7 +295,10 @@ def __init__(
         self._cappa = self.acoustic_dynamics.cappa
 
         if not (not self.config.inline_q and constants.NQ != 0):
-            raise NotImplementedError("tracer_2d not implemented, turn on z_tracer")
+            raise NotImplementedError(
+                "Dynamical core (fv_dynamics):"
+                "tracer_2d not implemented. z_tracer available"
+            )
         self._adjust_tracer_mixing_ratio = AdjustNegativeTracerMixingRatio(
             stencil_factory,
             quantity_factory=quantity_factory,
@@ -466,16 +478,20 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool):
         )
 
         if self._conserve_total_energy > 0:
-            raise NotImplementedError("compute total energy is not implemented")
+            raise NotImplementedError(
+                "Dynamical Core (fv_dynamics): compute total energy is not implemented"
+            )
 
         if (not self.config.rf_fast) and self.config.tau != 0:
             raise NotImplementedError(
-                "Rayleigh_Super, called when rf_fast=False and tau !=0"
+                "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(
-                "unimplemented namelist options adiabatic with positive kord_tm"
+                "Dynamical Core (fv_dynamics): Adiabatic with positive kord_tm"
+                " is not implemented."
             )
         else:
             if __debug__:

fv3core/pace/fv3core/stencils/fv_subgridz.py

@@ -99,7 +99,6 @@ def init(
     qsgs_tke: FloatField,
     qcld: FloatField,
 ):
-
     with computation(PARALLEL), interval(...):
         t0 = ta
         u0 = ua
@@ -783,12 +782,17 @@ def __init__(
         n_sponge: int,
         hydrostatic: bool,
     ):
-        assert not hydrostatic, "Hydrostatic not implemented for fv_subgridz"
+        if hydrostatic:
+            raise NotImplementedError(
+                "DryConvectiveAdjustment (fv_subgridz):"
+                " Hydrostatic is not implemented"
+            )
         grid_indexing = stencil_factory.grid_indexing
         self._k_sponge = n_sponge
-        if self._k_sponge is not None:
-            if self._k_sponge < 3:
-                return
+        if self._k_sponge is not None and self._k_sponge < 3:
+            raise ValueError(
+                "DryConvectiveAdjustment (fv_subgridz): n_sponge < 3 is invalid."
+            )
         else:
             self._k_sponge = grid_indexing.domain[2]
         if self._k_sponge < min(grid_indexing.domain[2], 24):

fv3core/pace/fv3core/stencils/neg_adj3.py

@@ -356,7 +356,10 @@ def __init__(
             )
         if hydrostatic:
             self._d0_vap = constants.CP_VAP - constants.C_LIQ
-            raise NotImplementedError("Unimplemented namelist hydrostatic=True")
+            raise NotImplementedError(
+                "Adjust Negative Tracer Mixing Ratio (neg_adj3):"
+                " Hydrostatic is not implemented"
+            )
         else:
             self._d0_vap = constants.CV_VAP - constants.C_LIQ
         self._lv00 = constants.HLV - self._d0_vap * constants.TICE

fv3core/pace/fv3core/stencils/remap_profile.py

@@ -591,7 +591,11 @@ def __init__(
             config=stencil_factory.config.dace_config,
         )
 
-        assert kord <= 10, f"kord {kord} not implemented."
+        if kord > 10:
+            raise NotImplementedError(
+                f"Remap Profile: kord {kord} not implemented. kord <= 10 available."
+            )
+
         self._kord = kord
 
         self._gam = quantity_factory.zeros(

fv3core/pace/fv3core/stencils/xppm.py

@@ -310,7 +310,11 @@ def __init__(
         # Arguments come from:
         # namelist.grid_type
         # grid.dxa
-        assert (grid_type < 3) or (grid_type == 4)
+        if grid_type == 3 or grid_type > 4:
+            raise NotImplementedError(
+                "X Piecewise Parabolic (xppm): "
+                f" grid type {grid_type} not implemented. <3 or 4 available."
+            )
         self._dxa = dxa
         ax_offsets = stencil_factory.grid_indexing.axis_offsets(origin, domain)
         self._compute_flux_stencil = stencil_factory.from_origin_domain(

fv3core/pace/fv3core/stencils/yppm.py

@@ -312,7 +312,11 @@ def __init__(
         # Arguments come from:
         # namelist.grid_type
         # grid.dya
-        assert (grid_type < 3) or (grid_type == 4)
+        if grid_type == 3 or grid_type > 4:
+            raise NotImplementedError(
+                "Y Piecewise Parabolic (xppm): "
+                f" grid type {grid_type} not implemented. <3 or 4 available."
+            )
         self._dya = dya
         ax_offsets = stencil_factory.grid_indexing.axis_offsets(origin, domain)
         self._compute_flux_stencil = stencil_factory.from_origin_domain(

tests/main/fv3core/test_dycore_call.py

@@ -19,9 +19,9 @@
 DIR = os.path.abspath(os.path.dirname(__file__))
 
 
-def setup_dycore() -> Tuple[
-    fv3core.DynamicalCore, fv3core.DycoreState, pace.util.Timer
-]:
+def setup_dycore() -> (
+    Tuple[fv3core.DynamicalCore, fv3core.DycoreState, pace.util.Timer]
+):
     backend = "numpy"
     config = fv3core.DynamicalCoreConfig(
         layout=(1, 1),

util/pace/util/monitor/netcdf_monitor.py

@@ -94,6 +94,7 @@ def flush(self):
                     chunk=chunk_index, tile=self._tile
                 )
             )
+            Path(self._path).mkdir(exist_ok=True)
             if os.path.exists(chunk_path):
                 os.remove(chunk_path)
             ds.to_netcdf(chunk_path, format="NETCDF4", engine="netcdf4")