Reorganization of fv3core analytic test case initialization (#26)

* Edits for reorganization of analytic initializations, still needs work * New updates * changed the name of analytic_test.yaml * New unit test: test_analytic_init.py * Changes to test_analytic_init.py * Changes to test_analytic_init.py * gaea changes * fixing enum in analytic_init * adding assert to test_analytic_init * comment out test * comment out test * comment out test * comment out test * Moved dycore_state and geos_wrapper * Removed commented locations for dycore and geos_wrapper * Changes as of 19 Sept 2023 * Same as before * Same as before * Same as before * Same as before * Hopefully merged local and remote changes * Removed buildenv directory * Fixed reference to dycore_state in test_diagnostics_config * Fixing test issues * Fixing test issues * Fixing test issues * Fixing test issues * Fixing test issues * Fixing test issues * Fixing test issues * Fixing test issues * Fixing test issues * Fixing test issues * baroclinic_c12.yaml reverted back to new form * Changes to baroclinic_c12.yaml to observe effect on error * test_diagnostics.py issue work start * Reup * Linting fixes * More linting * Added 'ua' back into baroclinic_c12.yaml, removed it during testing. * Update driver/pace/driver/driver.py change to del method to clear case in driver.py Co-authored-by: Oliver Elbert <[email protected]> * Update fv3core/pace/fv3core/initialization/analytic_init.py Change to closing of conditional statement to raise ValueError on else Co-authored-by: Oliver Elbert <[email protected]> * Update fv3core/pace/fv3core/initialization/test_cases/initialize_tc.py Removing pytest breakpoint Co-authored-by: Oliver Elbert <[email protected]> * Update tests/main/driver/test_analytic_init.py Removal of comment Co-authored-by: Oliver Elbert <[email protected]> * Oliver E suggestions * Oliver E suggestions pt.2 * Update fv3core/pace/fv3core/initialization/init_utils.py Co-authored-by: Florian Deconinck <[email protected]> * Reverting gt4py to the correct hash * Florian changes pre-lint * Module variables capitalized * Module variables capitalized in initialize_baroclinic.py * Linted after Florian suggestions * Update fv3core/pace/fv3core/initialization/analytic_init.py Co-authored-by: Florian Deconinck <[email protected]> * Update fv3core/pace/fv3core/initialization/analytic_init.py Co-authored-by: Florian Deconinck <[email protected]> * Update driver/pace/driver/driver.py Co-authored-by: Oliver Elbert <[email protected]> * Update fv3core/pace/fv3core/initialization/init_utils.py Co-authored-by: Oliver Elbert <[email protected]> * Update fv3core/pace/fv3core/initialization/init_utils.py Co-authored-by: Oliver Elbert <[email protected]> * Update fv3core/pace/fv3core/initialization/init_utils.py Co-authored-by: Oliver Elbert <[email protected]> * Changed variable analytic_init_str to analytic_init_case --------- Co-authored-by: Frank Malatino <[email protected]> Co-authored-by: Frank Malatino <[email protected]> Co-authored-by: Frank Malatino <[email protected]> Co-authored-by: Frank Malatino <[email protected]> Co-authored-by: Frank Malatino <[email protected]> Co-authored-by: Oliver Elbert <[email protected]> Co-authored-by: Florian Deconinck <[email protected]> Co-authored-by: Frank Malatino <[email protected]>
GEOS-ESM · Sep 29, 2023 · 1b91c76 · 1b91c76
1 parent 0cdba14
commit 1b91c76
Show file tree

Hide file tree

Showing 37 changed files with 1,172 additions and 1,014 deletions.
diff --git a/.jenkins/driver_configs/baroclinic_c192_54ranks.yaml b/.jenkins/driver_configs/baroclinic_c192_54ranks.yaml
@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: true
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 diagnostics_config:
   path: "output.zarr"
   names:

diff --git a/.jenkins/driver_configs/baroclinic_c192_6ranks.yaml b/.jenkins/driver_configs/baroclinic_c192_6ranks.yaml
@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: true
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 diagnostics_config:
   path: "output.zarr"
   names:

diff --git a/.jenkins/driver_configs/baroclinic_c48_6ranks_dycore_only.yaml b/.jenkins/driver_configs/baroclinic_c48_6ranks_dycore_only.yaml
@@ -9,7 +9,9 @@ stencil_config:
     device_sync: false
     run_mode: Run
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinicc
 performance_config:
   collect_performance: false
 nx_tile: 48

diff --git a/driver/examples/configs/analytic_test.yaml b/driver/examples/configs/analytic_test.yaml
@@ -0,0 +1,101 @@
+stencil_config:
+  compilation_config:
+    backend: numpy
+    rebuild: false
+    validate_args: true
+    format_source: false
+    device_sync: false
+initialization:
+  type: analytic
+  config:
+    case: baroclinic
+performance_config:
+  collect_performance: true
+  experiment_name: c12_baroclinic
+comm_config:
+  type: null_comm
+  config:
+    rank: 0
+    total_ranks: 6
+nx_tile: 12
+nz: 79
+dt_atmos: 225
+minutes: 15
+layout:
+  - 1
+  - 1
+diagnostics_config:
+  path: output
+  output_format: netcdf
+  names:
+    - u
+    - v
+    - ua
+    - va
+    - pt
+    - delp
+    - qvapor
+    - qliquid
+    - qice
+    - qrain
+    - qsnow
+    - qgraupel
+  z_select:
+    - level: 65
+      names:
+        - pt
+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
diff --git a/driver/examples/configs/baroclinic_c12.yaml b/driver/examples/configs/baroclinic_c12.yaml
@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: true
   experiment_name: c12_baroclinic

diff --git a/driver/examples/configs/baroclinic_c12_comm_read.yaml b/driver/examples/configs/baroclinic_c12_comm_read.yaml
@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
   experiment_name: c12_baroclinic

diff --git a/driver/examples/configs/baroclinic_c12_comm_write.yaml b/driver/examples/configs/baroclinic_c12_comm_write.yaml
@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
   experiment_name: c12_baroclinic

diff --git a/driver/examples/configs/baroclinic_c12_dp.yaml b/driver/examples/configs/baroclinic_c12_dp.yaml
@@ -13,7 +13,9 @@ grid_config:
     dy_const: 3000.0
     deglat: 10.0
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: true
   experiment_name: c12_baroclinic

diff --git a/driver/examples/configs/baroclinic_c12_null_comm.yaml b/driver/examples/configs/baroclinic_c12_null_comm.yaml
@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
   experiment_name: c12_baroclinic

diff --git a/driver/examples/configs/baroclinic_c12_orch_cpu.yaml b/driver/examples/configs/baroclinic_c12_orch_cpu.yaml
@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
 nx_tile: 12

diff --git a/driver/examples/configs/baroclinic_c12_write_restart.yaml b/driver/examples/configs/baroclinic_c12_write_restart.yaml
@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
   experiment_name: c12_baroclinic

diff --git a/driver/examples/configs/tropicalcyclone_c128.yaml b/driver/examples/configs/tropicalcyclone_c128.yaml
@@ -8,7 +8,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: tropicalcyclone
+  type: analytic
+  config:
+    case: tropicalcyclone
 performance_config:
   performance_mode: true
   experiment_name: c128_tropical

diff --git a/driver/pace/driver/__init__.py b/driver/pace/driver/__init__.py
@@ -9,7 +9,7 @@
 from .diagnostics import Diagnostics, DiagnosticsConfig
 from .driver import Driver, DriverConfig, RestartConfig
 from .grid import GeneratedGridConfig, SerialboxGridConfig
-from .initialization import BaroclinicInit, PredefinedStateInit, RestartInit
+from .initialization import AnalyticInit, PredefinedStateInit, RestartInit
 from .performance import PerformanceConfig
 from .registry import Registry
 from .state import DriverState, TendencyState

diff --git a/driver/pace/driver/diagnostics.py b/driver/pace/driver/diagnostics.py
@@ -10,7 +10,7 @@
 import pace.util
 import pace.util.grid
 from pace.dsl.dace.orchestration import dace_inhibitor
-from pace.fv3core.initialization.dycore_state import DycoreState
+from pace.fv3core.dycore_state import DycoreState
 from pace.util.constants import RGRAV
 
 from .state import DriverState

diff --git a/driver/pace/driver/driver.py b/driver/pace/driver/driver.py
@@ -322,6 +322,9 @@ def write_for_restart(
         config_dict["initialization"]["type"] = "restart"
         config_dict["initialization"]["config"]["start_time"] = time
         config_dict["initialization"]["config"]["path"] = restart_path
+        # restart config doesn't have 'case'
+        if "case" in config_dict["initialization"]["config"].keys():
+            del config_dict["initialization"]["config"]["case"]
         with open(f"{restart_path}/restart.yaml", "w") as file:
             yaml.safe_dump(config_dict, file)
 

diff --git a/driver/pace/driver/initialization.py b/driver/pace/driver/initialization.py
@@ -9,8 +9,7 @@
 
 import pace.driver
 import pace.dsl
-import pace.fv3core.initialization.baroclinic as baroclinic_init
-import pace.fv3core.initialization.tropical_cyclone as tc_init
+import pace.fv3core.initialization.analytic_init as analytic_init
 import pace.physics
 import pace.stencils
 import pace.util
@@ -93,13 +92,14 @@ def from_dict(cls, config: dict):
         return cls(config=instance, type=config["type"])
 
 
-@InitializerSelector.register("baroclinic")
+@InitializerSelector.register("analytic")
 @dataclasses.dataclass
-class BaroclinicInit(Initializer):
+class AnalyticInit(Initializer):
     """
-    Configuration for baroclinic initialization.
+    Configuration for analytic initialization.
     """
 
+    case: str = "baroclinic"
     start_time: datetime = datetime(2000, 1, 1)
 
     def get_driver_state(
@@ -110,7 +110,8 @@ def get_driver_state(
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
     ) -> DriverState:
-        dycore_state = baroclinic_init.init_baroclinic_state(
+        dycore_state = analytic_init.init_analytic_state(
+            analytic_init_case=self.case,
             grid_data=grid_data,
             quantity_factory=quantity_factory,
             adiabatic=False,
@@ -134,56 +135,6 @@ def get_driver_state(
         )
 
 
-@InitializerSelector.register("tropicalcyclone")
-@dataclasses.dataclass
-class TropicalCycloneConfig(Initializer):
-    """
-    Configuration for tropical cyclone initialization.
-    """
-
-    # TODO
-    # this can be cleaned up after grid config is separated
-
-    start_time: datetime = datetime(2000, 1, 1)
-
-    def get_driver_state(
-        self,
-        quantity_factory: pace.util.QuantityFactory,
-        communicator: pace.util.CubedSphereCommunicator,
-        damping_coefficients: pace.util.grid.DampingCoefficients,
-        driver_grid_data: pace.util.grid.DriverGridData,
-        grid_data: pace.util.grid.GridData,
-    ) -> DriverState:
-        dycore_state = tc_init.init_tc_state(
-            grid_data=grid_data,
-            quantity_factory=quantity_factory,
-            hydrostatic=False,
-            comm=communicator,
-        )
-
-        physics_state = pace.physics.PhysicsState.init_zeros(
-            quantity_factory=quantity_factory, active_packages=["microphysics"]
-        )
-        tendency_state = TendencyState.init_zeros(
-            quantity_factory=quantity_factory,
-        )
-
-        print(
-            "delp: ",
-            dycore_state.delp.data[:, :, -2].min(),
-            dycore_state.pt.data[:, :, -2].max(),
-        )
-
-        return DriverState(
-            dycore_state=dycore_state,
-            physics_state=physics_state,
-            tendency_state=tendency_state,
-            grid_data=grid_data,
-            damping_coefficients=damping_coefficients,
-            driver_grid_data=driver_grid_data,
-        )
-
-
 @InitializerSelector.register("restart")
 @dataclasses.dataclass
 class RestartInit(Initializer):

diff --git a/driver/pace/driver/safety_checks.py b/driver/pace/driver/safety_checks.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from pace.fv3core.initialization.dycore_state import DycoreState
+from pace.fv3core.dycore_state import DycoreState
 from pace.util.quantity import Quantity
 
 

diff --git a/fv3core/examples/standalone/runfile/dynamics.py b/fv3core/examples/standalone/runfile/dynamics.py
@@ -22,8 +22,8 @@
 from pace.dsl import StencilFactory
 from pace.dsl.dace.orchestration import DaceConfig
 from pace.fv3core import DynamicalCore, DynamicalCoreConfig
+from pace.fv3core.dycore_state import DycoreState
 from pace.fv3core.initialization.baroclinic import init_baroclinic_state
-from pace.fv3core.initialization.dycore_state import DycoreState
 from pace.fv3core.testing import TranslateFVDynamics
 from pace.stencils.testing import dataset_to_dict
 from pace.stencils.testing.grid import Grid

diff --git a/fv3core/pace/fv3core/__init__.py b/fv3core/pace/fv3core/__init__.py
@@ -1,8 +1,8 @@
 from ._config import DynamicalCoreConfig
-from .initialization.dycore_state import DycoreState
-from .initialization.geos_wrapper import GeosDycoreWrapper
+from .dycore_state import DycoreState
 from .stencils.fv_dynamics import DynamicalCore
 from .stencils.fv_subgridz import DryConvectiveAdjustment
+from .wrappers.geos_wrapper import GeosDycoreWrapper
 
 
 __version__ = "0.2.0"
diff --git a/...ce/fv3core/initialization/dycore_state.py → fv3core/pace/fv3core/dycore_state.py b/...ce/fv3core/initialization/dycore_state.py → fv3core/pace/fv3core/dycore_state.py
diff --git a/fv3core/pace/fv3core/initialization/__init__.py b/fv3core/pace/fv3core/initialization/__init__.py
@@ -1,2 +1 @@
-from .baroclinic import init_baroclinic_state
-from .tropical_cyclone import init_tc_state
+from .analytic_init import init_analytic_state