add diag_manager output option

.github/workflows/main_unit_tests.yaml

@@ -50,8 +50,18 @@ jobs:
         with:
           python-version: ${{ inputs.python_version && inputs.python_version || env.python_default }}
 
-      - name: Install mpi
-        run: pip install openmpi
+      - name: Install pyFMS dependencies
+        run: |
+          sudo apt update
+          sudo apt install -y libopenmpi-dev libyaml-dev gfortran libnetcdf-dev libnetcdff-dev netcdf-bin
+
+        # need to build mpi4py from source to avoid mismatching mpich versions between os and pip
+      - name: Install mpi (MPICH flavor)
+        run: pip install mpi4py --no-cache-dir --no-binary mpi4py
+        env:
+          MPICC: mpicc
+          MPICXX: mpicxx
+          MPIFC: mpifort
 
       - name: "External trigger: Remove existing component in pace/develop"
         if: ${{ inputs.component_name }}
@@ -81,5 +91,9 @@ jobs:
 
       - name: Run baroclinic test case
         run: |
+          which mpiexec
+          which mpirun
+          mpicc --version
+          mpifort --version
           cd ${GITHUB_WORKSPACE}/pace
           mpiexec -np 6 --oversubscribe python -m pace.run examples/configs/baroclinic_c48_no_out.yaml

examples/configs/baroclinic_c12_diag_manager.yaml

@@ -0,0 +1,98 @@
+stencil_config:
+  compilation_config:
+    backend: st:numpy:cpu:IJK
+    rebuild: false
+    validate_args: true
+    format_source: false
+    device_sync: false
+initialization:
+  type: analytic
+  config:
+    case: baroclinic_instability
+performance_config:
+  collect_performance: true
+  experiment_name: c12_baroclinic_instability
+nx_tile: 12
+nz: 79
+dt_atmos: 225
+minutes: 15
+layout:
+  - 1
+  - 1
+diagnostics_config:
+  output_format: diag_manager
+  precision: Float
+  names:
+    - u
+    - v
+    - ua
+    - va
+    - pt
+    - delp
+    - qvapor
+    - qliquid
+    - qice
+    - qrain
+    - qsnow
+    - qgraupel
+    - pe
+dycore_config:
+  a_imp: 1.0
+  beta: 0.
+  consv_te: 0.
+  d2_bg: 0.
+  d2_bg_k1: 0.2
+  d2_bg_k2: 0.1
+  d4_bg: 0.15
+  d_con: 1.0
+  d_ext: 0.0
+  dddmp: 0.5
+  delt_max: 0.002
+  do_sat_adj: true
+  do_vort_damp: true
+  fill: true
+  hord_dp: 6
+  hord_mt: 6
+  hord_tm: 6
+  hord_tr: 8
+  hord_vt: 6
+  hydrostatic: false
+  k_split: 1
+  ke_bg: 0.
+  kord_mt: 9
+  kord_tm: -9
+  kord_tr: 9
+  kord_wz: 9
+  n_split: 1
+  nord: 3
+  nwat: 6
+  p_fac: 0.05
+  rf_cutoff: 3000.
+  rf_fast: true
+  tau: 10.
+  vtdm4: 0.06
+  z_tracer: true
+  do_qa: true
+  tau_i2s: 1000.
+  tau_g2v: 1200.
+  ql_gen: 0.001
+  ql_mlt: 0.002
+  qs_mlt: 0.000001
+  qi_lim: 1.0
+  dw_ocean: 0.1
+  dw_land: 0.15
+  icloud_f: 0
+  tau_l2v: 300.
+  tau_v2l: 90.
+  fv_sg_adj: 0
+  n_sponge: 48
+
+physics_config:
+  hydrostatic: false
+  nwat: 6
+  do_qa: true
+
+grid_config:
+  type: generated
+  config:
+    eta_file: 'NDSL/tests/data/eta/eta79.nc'

examples/configs/diag_table.yaml

@@ -0,0 +1,60 @@
+# diag table for use with the 'baroclinic' test case for pace
+# The var_names below must correspond to the 'names' section in the diagnostics_config section for
+# the pace config file, otherwise the data will appear as all fill values in the output netcdf.
+# The 'output_name' is what will be used in the output file if provided, and must be unique.
+title: pace_diag_manager_test
+base_date: 2000 1 1 0 0 0
+diag_files:
+  # below will result in analogous output to the 'netcdf' monitor output option
+  # it will be transposed, and contain all timesteps rather than 1 timestep per diagnostic output frequency.
+  # tile count in output filenames start at 1 instead of 0
+  - file_name: model_state_baroclinic
+    freq: 225 seconds # ie. one timestep of the model
+    time_units: seconds
+    unlimdim: time
+    module: pyfv3
+    reduction: none
+    kind: r8
+    varlist:
+      - var_name: u
+      - var_name: v
+      - var_name: ua
+      - var_name: va
+      - var_name: pt
+      - var_name: delp
+      - var_name: qvapor
+      - var_name: qliquid
+      - var_name: qice
+      - var_name: qrain
+      - var_name: qsnow
+      - var_name: qgraupel
+      - var_name: pe
+  # this tests the reduction methods, taking the min/max/avg/etc over the entire simulation length (4 timesteps)
+  - file_name: reductions_baroclinic
+    freq: 900 seconds # 4 timesteps ie. the entire simulation length
+    time_units: seconds
+    unlimdim: time
+    module: pyfv3
+    kind: r8
+    varlist:
+      - var_name: pt    # max temperature
+        output_name: max_pt
+        reduction: max
+      - var_name: qvapor # min specific humidity
+        output_name: min_qvapor
+        reduction: min
+      - var_name: ps    # sum of surface pressure
+        output_name: sum_ps
+        reduction: sum
+      - var_name: pt    # avg temperature
+        output_name: avg_pt
+        reduction: average
+      - var_name: u    # rms of x wind
+        output_name: rms_u
+        reduction: rms
+      - var_name: v    # sum of squares of y wind
+        output_name: pow2_v
+        reduction: pow2
+      - var_name: pt          # diurnal temperature range, adds an 'diurnal' dimension to the output file.
+        output_name: diurnal_pt # this is just to test the functionality, doesn't really make sense for a short simulation like this.
+        reduction: diurnal2

pace/diagnostics.py

@@ -6,12 +6,13 @@
 
 import numpy as np
 
-from ndsl import Quantity
+from ndsl import DiagManagerMonitor, Quantity
 from ndsl.constants import K_DIM, K_INTERFACE_DIM, RGRAV
 from ndsl.dsl.dace.orchestration import dace_inhibitor
 from ndsl.dsl.typing import Float
 from ndsl.grid import GridData
 from ndsl.monitor import Monitor, ZarrMonitor
+from ndsl.monitor.diag_manager_monitor import register_diag_manager_fields
 from ndsl.monitor.netcdf_monitor import NetCDFMonitor
 from ndsl.typing import Communicator
 from pace.state import DriverState
@@ -88,13 +89,15 @@ class DiagnosticsConfig:
     precision: str = "Float"
 
     def __post_init__(self):
-        if (len(self.names) > 0 or len(self.derived_names) > 0) and self.path is None:
+        if (len(self.names) > 0 or len(self.derived_names) > 0) and (
+            self.path is None and self.output_format != "diag_manager"
+        ):
             raise ValueError(
                 "DiagnosticsConfig.path must be given to enable diagnostics"
             )
-        if self.output_format not in ["zarr", "netcdf"]:
+        if self.output_format not in ["zarr", "netcdf", "diag_manager"]:
             raise ValueError(
-                f"output_format must be one of 'zarr' or 'netcdf', got {self.output_format}"
+                f"output_format must be one of 'zarr', 'netcdf', or 'diag_manager', got {self.output_format}"
             )
         if self.precision not in ["Float", "float32", "float64"]:
             raise ValueError(
@@ -109,11 +112,12 @@ def diagnostics_factory(self, communicator: Communicator) -> Diagnostics:
             communicator: provides global communication e.g. to gather state
                 or to coordinate filesystem access between ranks
         """
-        if self.path is None:
+        if self.path is None and self.output_format != "diag_manager":
             return NullDiagnostics()
 
-        if not Path(self.path).exists():
-            Path(self.path).mkdir()
+        if self.output_format != "diag_manager":
+            if not Path(self.path).exists():
+                Path(self.path).mkdir()
 
         if self.output_format == "zarr":
             store = zarr_storage.DirectoryStore(path=self.path)
@@ -135,9 +139,11 @@ def diagnostics_factory(self, communicator: Communicator) -> Diagnostics:
                 time_chunk_size=self.time_chunk_size,
                 precision=precision,
             )
+        elif self.output_format == "diag_manager":
+            monitor = DiagManagerMonitor()
         else:
             raise ValueError(
-                f"output_format must be one of 'zarr' or 'netcdf', got {self.output_format}"
+                f"output_format must be one of 'zarr', 'netcdf', or 'diag_manager', got {self.output_format}"
             )
 
         return MonitorDiagnostics(