updated physics calls

Author: oelbert

pace/driver.py

@@ -37,10 +37,9 @@
 from pace.initialization import InitializerSelector
 from pace.safety_checks import SafetyChecker
 from pace.state import DriverState
-from pyfv3 import DynamicalCore, DynamicalCoreConfig
-from pyfv3.initialization.analytic_init import AnalyticCase
-from pyshield import Physics, PhysicsConfig
-from pyshield.update import update_atmos_state
+from pyFV3 import DynamicalCore, DynamicalCoreConfig
+from pySHiELD import Physics, PhysicsConfig
+from pySHiELD.update import update_atmos_state
 
 
 try:
@@ -116,7 +115,6 @@ class DriverConfig:
         default_factory=DynamicalCoreConfig
     )
     physics_config: PhysicsConfig = dataclasses.field(default_factory=PhysicsConfig)
-
     days: int = 0
     hours: int = 0
     minutes: int = 0
@@ -267,6 +265,9 @@ def from_dict(cls, kwargs: Dict[str, Any]) -> "DriverConfig":
         kwargs["comm_config"] = CreatesCommSelector.from_dict(
             kwargs.get("comm_config", {})
         )
+        kwargs["initialization"] = InitializerSelector.from_dict(
+            kwargs["initialization"]
+        )
         if "grid_config" in kwargs:
             kwargs["grid_config"] = GridInitializerSelector.from_dict(
                 kwargs["grid_config"]
@@ -278,17 +279,6 @@ def from_dict(cls, kwargs: Dict[str, Any]) -> "DriverConfig":
                 if grid_type > 3:
                     kwargs["dycore_config"].ntiles = 1
 
-        analytic_hooks = {}
-        if kwargs["initialization"]["type"] == "analytic":
-            kwargs["initialization"]["config"]["dycore_config"] = kwargs[
-                "dycore_config"
-            ]
-            analytic_hooks[AnalyticCase] = AnalyticCase
-        kwargs["initialization"] = InitializerSelector.from_dict(
-            kwargs["initialization"],
-            hooks=analytic_hooks,
-        )
-
         if (
             isinstance(kwargs["stencil_config"], dict)
             and "dace_config" in kwargs["stencil_config"].keys()
@@ -331,9 +321,6 @@ def write_for_restart(
             config_dict["dycore_config"].pop(field, None)
             config_dict["physics_config"].pop(field, None)
         config_dict["initialization"]["type"] = "restart"
-        # Remove existing initialization config and repopulate
-        config_dict["initialization"].pop("config")
-        config_dict["initialization"]["config"] = {}
         config_dict["initialization"]["config"]["start_time"] = time
         config_dict["initialization"]["config"]["path"] = restart_path
         # convert physics package enum to str
@@ -663,7 +650,13 @@ def _critical_path_step_all(
                         timestep=dt,
                     )
                     if not self.config.dycore_only:
-                        self.physics(self.state.physics_state, timestep=dt)
+                        self.physics(
+                            self.state.physics_state,
+                            self.state.radiation_state,
+                            self.state.surface_state,
+                            self.time,
+                            timestep=dt,
+                        )
                     self.end_of_step_update(
                         dycore_state=self.state.dycore_state,
                         phy_state=self.state.physics_state,

pace/state.py

@@ -12,7 +12,7 @@
 from ndsl.grid import DampingCoefficients, DriverGridData, GridData
 from ndsl.typing import Communicator
 from pyfv3 import DycoreState
-from pyshield import PHYSICS_PACKAGES, PhysicsState
+from pyshield import PHYSICS_PACKAGES, PhysicsState, RadiationState, SurfaceState
 
 
 @dataclasses.dataclass()
@@ -63,6 +63,8 @@ def init_zeros(cls, quantity_factory: QuantityFactory) -> "TendencyState":
 class DriverState:
     dycore_state: DycoreState
     physics_state: PhysicsState
+    surface_state: SurfaceState
+    radiation_state: RadiationState
     tendency_state: TendencyState
     grid_data: GridData
     damping_coefficients: DampingCoefficients
@@ -120,6 +122,12 @@ def save_state(self, comm, restart_path: str = "RESTART"):
         self.physics_state.xr_dataset.to_netcdf(
             f"{restart_path}/restart_physics_state_{current_rank}.nc"
         )
+        self.radiation_state.xr_dataset.to_netcdf(
+            f"{restart_path}/restart_radiation_state_{current_rank}.nc"
+        )
+        self.surface_state.xr_dataset.to_netcdf(
+            f"{restart_path}/restart_surface_state_{current_rank}.nc"
+        )
         # we can also convert the state to Fortran's restart format using
         # code similar to this commented code. We don't need this feature right
         # now so we haven't implemented it, but this is a good starter.