Enable starting driver from Fortran restart files for TC case

.circleci/config.yml

@@ -97,6 +97,7 @@ commands:
         description: "make target to run"
         type: enum
         enum:
+          - "test_main"
           - "test_savepoint"
           - "savepoint_tests"
           - "savepoint_tests_mpi"
@@ -395,20 +396,19 @@ jobs:
               python3 -m mpi4py -m pace.driver.run examples/configs/baroclinic_c12_orch_cpu.yaml
 
   test_main:
-    docker:
-      - image: cimg/python:3.8
+    machine:
+      image: ubuntu-2004:202111-02
+      resource_class: medium
+    environment:
+      GOOGLE_APPLICATION_CREDENTIALS: /tmp/key.json
     working_directory: ~/repo
     steps:
       - checkout
-      - run:
-          name: Install Submodules
-          command: git submodule update --init
-      - setup_environment_mpi
-      - run:
-          name: run tests
-          command: |
-            . venv/bin/activate
-            pytest tests/main
+      - make_savepoints:
+          backend: numpy
+          experiment: c12_6ranks_standard
+          target: test_main
+          num_ranks: 1
 
   test_mpi_54rank:
     docker:
@@ -564,6 +564,8 @@ workflows:
             tags:
               only: /^v.*/
       - test_main:
+          context:
+            - GCLOUD_ENCODED_KEY
           filters:
             tags:
               only: /^v.*/

.gitignore

@@ -162,3 +162,5 @@ RESTART/
 _dacegraphs
 
 **/testing/output/*
+*.png
+*.nc

driver/examples/configs/baroclinic_c12_orch_cpu.yaml

@@ -8,8 +8,7 @@ stencil_config:
 initialization:
   type: baroclinic
 performance_config:
-  performance_mode: true
-  experiment_name: c12_baroclinic
+  performance_mode: false
 nx_tile: 12
 nz: 79
 dt_atmos: 225

driver/examples/configs/baroclinic_c12_read_restart_fortran.yml

@@ -0,0 +1,93 @@
+stencil_config:
+  compilation_config:
+    backend: numpy
+    rebuild: false
+    validate_args: true
+    format_source: false
+    device_sync: false
+initialization:
+  config:
+    path: /home/ajdas/pace/restart_data/1.0
+  type: fortran_restart
+performance_config:
+  performance_mode: true
+  experiment_name: c12_baroclinic_restart
+grid_config:
+  type: generated
+nx_tile: 12
+nz: 79
+dt_atmos: 225
+minutes: 15
+layout:
+  - 1
+  - 1
+diagnostics_config:
+  path: "output.zarr"
+  names:
+    - u
+    - v
+    - ua
+    - va
+    - pt
+    - delp
+    - qvapor
+    - qliquid
+    - qice
+    - qrain
+    - qsnow
+    - qgraupel
+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

driver/examples/configs/baroclinic_c12_write_restart.yaml

@@ -18,7 +18,8 @@ seconds: 225
 layout:
   - 1
   - 1
-save_restart: true
+restart_config:
+  save_restart: true
 diagnostics_config:
   path: "output.zarr"
   names:

driver/examples/configs/tropical_read_restart_fortran.yml

@@ -0,0 +1,98 @@
+dycore_only: true
+disable_step_physics: true
+stencil_config:
+  compilation_config:
+    backend: numpy
+    rebuild: false
+    validate_args: true
+    format_source: false
+    device_sync: false
+initialization:
+  config:
+    path: gs://vcm-ml-public/dsl/initial_conditions/c128_6ranks_tc
+  type: fortran_restart
+grid_config:
+  config:
+    stretch_factor: 3.0
+    lon_target: 172.5
+    lat_target: 17.5
+    ks: 0
+    vertical_grid_from_restart: True
+  type: generated
+performance_config:
+  performance_mode: true
+  experiment_name: c128_tropical
+nx_tile: 128
+nz: 79
+dt_atmos: 60
+minutes: 20
+layout:
+  - 1
+  - 1
+diagnostics_config:
+  path: "output.zarr"
+  names:
+    - u
+    - v
+    - ua
+    - va
+    - pt
+    - delp
+    - qvapor
+  output_initial_state: true
+dycore_config:
+  a_imp: 1.0
+  beta: 0.
+  consv_te: 0.
+  d2_bg: 0.
+  d2_bg_k1: 0.2
+  d2_bg_k2: 0.15
+  d4_bg: 0.14
+  d_con: 1.0
+  d_ext: 0.0
+  dddmp: 0.5
+  delt_max: 0.002
+  do_sat_adj: false
+  do_vort_damp: true
+  fill: false
+  hord_dp: 6
+  hord_mt: 6
+  hord_tm: 6
+  hord_tr: 8
+  hord_vt: 6
+  hydrostatic: false
+  k_split: 2
+  ke_bg: 0.
+  kord_mt: 9
+  kord_tm: -9
+  kord_tr: 9
+  kord_wz: 9
+  n_split: 3
+  nord: 3
+  nwat: 6
+  p_fac: 0.1
+  rf_cutoff: 5000.
+  rf_fast: true
+  tau: 5.
+  vtdm4: 0.06
+  z_tracer: true
+  do_qa: true
+  tau_i2s: 1000.
+  tau_g2v: 900.
+  ql_gen: 0.001
+  ql_mlt: 0.001
+  qs_mlt: 0.000001
+  qi_lim: 1.0
+  dw_ocean: 0.1
+  dw_land: 0.16
+  icloud_f: 0
+  tau_l2v: 225.
+  tau_v2l: 150.
+  fv_sg_adj: 600
+  n_sponge: 48
+  p_ref: 101500.
+
+physics_config:
+  hydrostatic: false
+  nwat: 6
+  do_qa: true

driver/pace/driver/__init__.py

@@ -2,15 +2,16 @@
     CreatesComm,
     CreatesCommSelector,
     MPICommConfig,
+    NullCommConfig,
     ReaderCommConfig,
     WriterCommConfig,
 )
 from .diagnostics import Diagnostics, DiagnosticsConfig
-from .driver import Driver, DriverConfig
-from .initialization import BaroclinicConfig, PredefinedStateConfig, RestartConfig
+from .driver import Driver, DriverConfig, RestartConfig
+from .grid import GeneratedGridConfig, SerialboxGridConfig
+from .initialization import BaroclinicInit, PredefinedStateInit, RestartInit
 from .performance import PerformanceConfig
 from .registry import Registry
-from .restart import Restart
 from .state import DriverState, TendencyState
 
 

driver/pace/driver/diagnostics.py

@@ -10,7 +10,6 @@
 import pace.util
 import pace.util.grid
 from pace.dsl.dace.orchestration import dace_inhibitor
-from pace.util.quantity import QuantityMetadata
 
 from .state import DriverState
 
@@ -27,9 +26,7 @@ def store(self, time: Union[datetime, timedelta], state: DriverState):
         ...
 
     @abc.abstractmethod
-    def store_grid(
-        self, grid_data: pace.util.grid.GridData, metadata: QuantityMetadata
-    ):
+    def store_grid(self, grid_data: pace.util.grid.GridData):
         ...
 
 
@@ -133,19 +130,11 @@ def _get_derived_state(self, state: DriverState):
                     warnings.warn(f"{name} is not a supported diagnostic variable.")
         return output
 
-    def store_grid(
-        self, grid_data: pace.util.grid.GridData, metadata: QuantityMetadata
-    ):
-        zarr_grid = {}
-        for name in ["lat", "lon"]:
-            grid_quantity = pace.util.Quantity(
-                getattr(grid_data, name),
-                dims=("x_interface", "y_interface"),
-                origin=metadata.origin,
-                extent=(metadata.extent[0] + 1, metadata.extent[1] + 1),
-                units="rad",
-            )
-            zarr_grid[name] = grid_quantity
+    def store_grid(self, grid_data: pace.util.grid.GridData):
+        zarr_grid = {
+            "lat": grid_data.lat,
+            "lon": grid_data.lon,
+        }
         self.monitor.store_constant(zarr_grid)
 
 
@@ -155,9 +144,7 @@ class NullDiagnostics(Diagnostics):
     def store(self, time: Union[datetime, timedelta], state: DriverState):
         pass
 
-    def store_grid(
-        self, grid_data: pace.util.grid.GridData, metadata: QuantityMetadata
-    ):
+    def store_grid(self, grid_data: pace.util.grid.GridData):
         pass