From 4c72d51bf1fd8b408915d7c6d55cea565ee1ae98 Mon Sep 17 00:00:00 2001 From: mlee03 Date: Tue, 19 Dec 2023 13:02:03 -0500 Subject: [PATCH] fix circular import problems --- .../configs/baroclinic_c12_write_restart.yaml | 5 ++ driver/pace/driver/diagnostics.py | 9 ++-- driver/pace/driver/driver.py | 15 ++---- driver/pace/driver/grid.py | 12 ++--- driver/pace/driver/initialization.py | 51 +++++++++---------- driver/pace/driver/state.py | 21 ++++---- .../fv3core/initialization/analytic_init.py | 2 +- .../pace/fv3core/initialization/init_utils.py | 7 ++- fv3core/pace/fv3core/stencils/a2b_ord4.py | 2 +- fv3core/pace/fv3core/stencils/c_sw.py | 2 +- fv3core/pace/fv3core/stencils/d2a2c_vect.py | 2 +- fv3core/pace/fv3core/stencils/d_sw.py | 2 +- fv3core/pace/fv3core/stencils/del2cubed.py | 2 +- fv3core/pace/fv3core/stencils/delnflux.py | 2 +- .../fv3core/stencils/divergence_damping.py | 2 +- fv3core/pace/fv3core/stencils/dyn_core.py | 2 +- fv3core/pace/fv3core/stencils/fv_dynamics.py | 2 +- fv3core/pace/fv3core/stencils/fvtp2d.py | 2 +- fv3core/pace/fv3core/stencils/fxadv.py | 2 +- fv3core/pace/fv3core/stencils/nh_p_grad.py | 2 +- fv3core/pace/fv3core/stencils/updatedzd.py | 2 +- .../fv3core/testing/translate_fvdynamics.py | 2 +- physics/pace/physics/stencils/microphysics.py | 2 +- physics/pace/physics/stencils/physics.py | 2 +- stencils/pace/stencils/c2l_ord.py | 2 +- stencils/pace/stencils/fv_update_phys.py | 2 +- stencils/pace/stencils/testing/grid.py | 4 +- stencils/pace/stencils/update_atmos_state.py | 2 +- stencils/pace/stencils/update_dwind_phys.py | 2 +- util/pace/util/grid/__init__.py | 1 - util/pace/util/grid/eta.py | 2 +- util/pace/util/grid/generation.py | 18 +++---- util/pace/util/grid/helper.py | 4 +- util/pace/util/grid/utils.py | 23 --------- 34 files changed, 94 insertions(+), 120 deletions(-) delete mode 100644 util/pace/util/grid/utils.py diff --git a/driver/examples/configs/baroclinic_c12_write_restart.yaml b/driver/examples/configs/baroclinic_c12_write_restart.yaml index 55e5b2b1..7bca3032 100644 --- a/driver/examples/configs/baroclinic_c12_write_restart.yaml +++ b/driver/examples/configs/baroclinic_c12_write_restart.yaml @@ -92,3 +92,8 @@ physics_config: hydrostatic: false nwat: 6 do_qa: true + +grid_config: + type: generated + config: + eta_file: "tests/main/input/eta79.nc" diff --git a/driver/pace/driver/diagnostics.py b/driver/pace/driver/diagnostics.py index 36f5960a..d417f903 100644 --- a/driver/pace/driver/diagnostics.py +++ b/driver/pace/driver/diagnostics.py @@ -7,11 +7,10 @@ import pace.driver import pace.dsl import pace.stencils -import pace.util -import pace.util.grid from pace.dsl.dace.orchestration import dace_inhibitor from pace.fv3core.dycore_state import DycoreState from pace.util.constants import RGRAV +from pace.util.grid.helper import GridData from .state import DriverState @@ -28,7 +27,7 @@ def store(self, time: Union[datetime, timedelta], state: DriverState): ... @abc.abstractmethod - def store_grid(self, grid_data: pace.util.grid.GridData): + def store_grid(self, grid_data: GridData): ... @abc.abstractmethod @@ -198,7 +197,7 @@ def _get_z_select_state(self, state: DycoreState): z_select_state.update(zselect.select_data(state)) return z_select_state - def store_grid(self, grid_data: pace.util.grid.GridData): + def store_grid(self, grid_data: GridData): zarr_grid = { "lat": grid_data.lat, "lon": grid_data.lon, @@ -218,7 +217,7 @@ class NullDiagnostics(Diagnostics): def store(self, time: Union[datetime, timedelta], state: DriverState): pass - def store_grid(self, grid_data: pace.util.grid.GridData): + def store_grid(self, grid_data: GridData): pass def cleanup(self): diff --git a/driver/pace/driver/driver.py b/driver/pace/driver/driver.py index c8d1490f..3694d69d 100644 --- a/driver/pace/driver/driver.py +++ b/driver/pace/driver/driver.py @@ -13,8 +13,6 @@ import pace.dsl import pace.physics import pace.stencils -import pace.util -import pace.util.grid from pace import fv3core from pace.driver.safety_checks import SafetyChecker from pace.dsl.dace.dace_config import DaceConfig @@ -29,6 +27,7 @@ CubedSphereCommunicator, TileCommunicator, ) +from pace.util.grid.helper import DampingCoefficients, DriverGridData, GridData from pace.util.logging import pace_log from . import diagnostics @@ -165,11 +164,7 @@ def get_grid( self, communicator: pace.util.Communicator, quantity_factory: Optional[pace.util.QuantityFactory] = None, - ) -> Tuple[ - pace.util.grid.DampingCoefficients, - pace.util.grid.DriverGridData, - pace.util.grid.GridData, - ]: + ) -> Tuple[DampingCoefficients, DriverGridData, GridData]: if quantity_factory is None: sizer = pace.util.SubtileGridSizer.from_tile_params( nx_tile=self.nx_tile, @@ -193,9 +188,9 @@ def get_grid( def get_driver_state( self, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, quantity_factory: Optional[pace.util.QuantityFactory] = None, stencil_factory: Optional[pace.dsl.StencilFactory] = None, ) -> DriverState: diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 33a0b43e..90e7cd23 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -11,19 +11,17 @@ import pace.util.grid from pace.stencils.testing import TranslateGrid from pace.util import Communicator, QuantityFactory -from pace.util.grid import ( - DampingCoefficients, - DriverGridData, - GridData, - MetricTerms, - direct_transform, -) from pace.util.grid.helper import ( AngleGridData, ContravariantGridData, + DampingCoefficients, + DriverGridData, + GridData, HorizontalGridData, + MetricTerms, VerticalGridData, ) +from pace.util.grid.stretch_transformation import direct_transform from pace.util.logging import pace_log from pace.util.namelist import Namelist diff --git a/driver/pace/driver/initialization.py b/driver/pace/driver/initialization.py index 3cf52376..651da6f5 100644 --- a/driver/pace/driver/initialization.py +++ b/driver/pace/driver/initialization.py @@ -12,14 +12,13 @@ import pace.fv3core.initialization.analytic_init as analytic_init import pace.physics import pace.stencils -import pace.util -import pace.util.grid from pace import fv3core from pace.dsl.dace.orchestration import DaceConfig from pace.dsl.stencil import StencilFactory from pace.dsl.stencil_config import CompilationConfig from pace.fv3core.testing import TranslateFVDynamics from pace.stencils.testing import TranslateGrid +from pace.util.grid.helper import DampingCoefficients, DriverGridData, GridData from pace.util.namelist import Namelist from .registry import Registry @@ -37,9 +36,9 @@ def get_driver_state( self, quantity_factory: pace.util.QuantityFactory, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ) -> DriverState: ... @@ -74,9 +73,9 @@ def get_driver_state( self, quantity_factory: pace.util.QuantityFactory, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ) -> DriverState: return self.config.get_driver_state( quantity_factory=quantity_factory, @@ -106,9 +105,9 @@ def get_driver_state( self, quantity_factory: pace.util.QuantityFactory, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ) -> DriverState: dycore_state = analytic_init.init_analytic_state( analytic_init_case=self.case, @@ -149,9 +148,9 @@ def get_driver_state( self, quantity_factory: pace.util.QuantityFactory, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ) -> DriverState: state = _restart_driver_state( self.path, @@ -198,9 +197,9 @@ def get_driver_state( self, quantity_factory: pace.util.QuantityFactory, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ) -> DriverState: state = _restart_driver_state( self.path, @@ -269,9 +268,9 @@ def get_driver_state( self, quantity_factory: pace.util.QuantityFactory, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ) -> DriverState: backend = quantity_factory.zeros( dims=[pace.util.X_DIM, pace.util.Y_DIM], units="unknown" @@ -337,18 +336,18 @@ class PredefinedStateInit(Initializer): dycore_state: fv3core.DycoreState physics_state: pace.physics.PhysicsState tendency_state: TendencyState - grid_data: pace.util.grid.GridData - damping_coefficients: pace.util.grid.DampingCoefficients - driver_grid_data: pace.util.grid.DriverGridData + grid_data: GridData + damping_coefficients: DampingCoefficients + driver_grid_data: DriverGridData start_time: datetime = datetime(2016, 8, 1) def get_driver_state( self, quantity_factory: pace.util.QuantityFactory, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ) -> DriverState: return DriverState( dycore_state=self.dycore_state, diff --git a/driver/pace/driver/state.py b/driver/pace/driver/state.py index 54241b1d..79c375ef 100644 --- a/driver/pace/driver/state.py +++ b/driver/pace/driver/state.py @@ -5,10 +5,9 @@ import pace.dsl.gt4py_utils as gt_utils import pace.physics -import pace.util -import pace.util.grid from pace import fv3core from pace.dsl.typing import Float +from pace.util.grid.helper import DampingCoefficients, DriverGridData, GridData @dataclasses.dataclass() @@ -60,9 +59,9 @@ class DriverState: dycore_state: fv3core.DycoreState physics_state: pace.physics.PhysicsState tendency_state: TendencyState - grid_data: pace.util.grid.GridData - damping_coefficients: pace.util.grid.DampingCoefficients - driver_grid_data: pace.util.grid.DriverGridData + grid_data: GridData + damping_coefficients: DampingCoefficients + driver_grid_data: DriverGridData # TODO: the driver_config argument here isn't type hinted from # import due to a circular dependency. This can be fixed by refactoring @@ -72,9 +71,9 @@ def load_state_from_restart( cls, restart_path: str, driver_config, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ) -> "DriverState": comm = driver_config.comm_config.get_comm() communicator = pace.util.Communicator.from_layout( @@ -173,9 +172,9 @@ def _restart_driver_state( rank: int, quantity_factory: pace.util.QuantityFactory, communicator: pace.util.Communicator, - damping_coefficients: pace.util.grid.DampingCoefficients, - driver_grid_data: pace.util.grid.DriverGridData, - grid_data: pace.util.grid.GridData, + damping_coefficients: DampingCoefficients, + driver_grid_data: DriverGridData, + grid_data: GridData, ): fs = pace.util.get_fs(path) diff --git a/fv3core/pace/fv3core/initialization/analytic_init.py b/fv3core/pace/fv3core/initialization/analytic_init.py index e8f6b07e..4f40adbf 100644 --- a/fv3core/pace/fv3core/initialization/analytic_init.py +++ b/fv3core/pace/fv3core/initialization/analytic_init.py @@ -2,7 +2,7 @@ import pace.util as fv3util from pace.fv3core.dycore_state import DycoreState -from pace.util.grid import GridData +from pace.util.grid.helper import GridData class MetaEnumStr(EnumMeta): diff --git a/fv3core/pace/fv3core/initialization/init_utils.py b/fv3core/pace/fv3core/initialization/init_utils.py index 42252e87..bfb55f71 100644 --- a/fv3core/pace/fv3core/initialization/init_utils.py +++ b/fv3core/pace/fv3core/initialization/init_utils.py @@ -8,8 +8,11 @@ import pace.util.constants as constants from pace.dsl.typing import Float from pace.fv3core.dycore_state import DycoreState -from pace.util.grid import lon_lat_midpoint -from pace.util.grid.gnomonic import get_lonlat_vect, get_unit_vector_direction +from pace.util.grid.gnomonic import ( + get_lonlat_vect, + get_unit_vector_direction, + lon_lat_midpoint, +) # maximum windspeed amplitude - close to windspeed of zonal-mean time-mean diff --git a/fv3core/pace/fv3core/stencils/a2b_ord4.py b/fv3core/pace/fv3core/stencils/a2b_ord4.py index 95b6ab02..9dba2001 100644 --- a/fv3core/pace/fv3core/stencils/a2b_ord4.py +++ b/fv3core/pace/fv3core/stencils/a2b_ord4.py @@ -17,7 +17,7 @@ from pace.dsl.typing import Float, FloatField, FloatFieldI, FloatFieldIJ from pace.fv3core.stencils.basic_operations import copy_defn from pace.util import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from pace.util.grid import GridData +from pace.util.grid.helper import GridData # comact 4-pt cubic interpolation diff --git a/fv3core/pace/fv3core/stencils/c_sw.py b/fv3core/pace/fv3core/stencils/c_sw.py index e74b9319..c5ce83e9 100644 --- a/fv3core/pace/fv3core/stencils/c_sw.py +++ b/fv3core/pace/fv3core/stencils/c_sw.py @@ -14,7 +14,7 @@ from pace.fv3core.stencils.d2a2c_vect import DGrid2AGrid2CGridVectors from pace.stencils import corners from pace.util import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from pace.util.grid import GridData +from pace.util.grid.helper import GridData def zero_delpc_ptc(delpc: FloatField, ptc: FloatField): diff --git a/fv3core/pace/fv3core/stencils/d2a2c_vect.py b/fv3core/pace/fv3core/stencils/d2a2c_vect.py index 1b3e4d33..a01a7503 100644 --- a/fv3core/pace/fv3core/stencils/d2a2c_vect.py +++ b/fv3core/pace/fv3core/stencils/d2a2c_vect.py @@ -8,7 +8,7 @@ from pace.fv3core.stencils.a2b_ord4 import a1, a2, lagrange_x_func, lagrange_y_func from pace.stencils import corners from pace.util import X_DIM, Y_DIM, Z_DIM -from pace.util.grid import GridData +from pace.util.grid.helper import GridData c1 = -2.0 / 14.0 diff --git a/fv3core/pace/fv3core/stencils/d_sw.py b/fv3core/pace/fv3core/stencils/d_sw.py index e08af776..301dca78 100644 --- a/fv3core/pace/fv3core/stencils/d_sw.py +++ b/fv3core/pace/fv3core/stencils/d_sw.py @@ -24,7 +24,7 @@ from pace.fv3core.stencils.xtp_u import advect_u_along_x from pace.fv3core.stencils.ytp_v import advect_v_along_y from pace.util import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from pace.util.grid import DampingCoefficients, GridData +from pace.util.grid.helper import DampingCoefficients, GridData dcon_threshold = 1e-5 diff --git a/fv3core/pace/fv3core/stencils/del2cubed.py b/fv3core/pace/fv3core/stencils/del2cubed.py index d4676eac..d06b44c8 100644 --- a/fv3core/pace/fv3core/stencils/del2cubed.py +++ b/fv3core/pace/fv3core/stencils/del2cubed.py @@ -7,7 +7,7 @@ from pace.dsl.typing import Float, FloatField, FloatFieldIJ, cast_to_index3d from pace.fv3core.stencils.basic_operations import copy_defn from pace.util import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from pace.util.grid import DampingCoefficients +from pace.util.grid.helper import DampingCoefficients # diff --git a/fv3core/pace/fv3core/stencils/delnflux.py b/fv3core/pace/fv3core/stencils/delnflux.py index 898a8a7f..9c76e695 100644 --- a/fv3core/pace/fv3core/stencils/delnflux.py +++ b/fv3core/pace/fv3core/stencils/delnflux.py @@ -15,7 +15,7 @@ from pace.dsl.stencil import StencilFactory, get_stencils_with_varied_bounds from pace.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK from pace.util import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from pace.util.grid import DampingCoefficients +from pace.util.grid.helper import DampingCoefficients def calc_damp( diff --git a/fv3core/pace/fv3core/stencils/divergence_damping.py b/fv3core/pace/fv3core/stencils/divergence_damping.py index a3e5a32b..ba03dd86 100644 --- a/fv3core/pace/fv3core/stencils/divergence_damping.py +++ b/fv3core/pace/fv3core/stencils/divergence_damping.py @@ -21,7 +21,7 @@ ) from pace.fv3core.stencils.d2a2c_vect import contravariant from pace.util import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from pace.util.grid import DampingCoefficients, GridData +from pace.util.grid.helper import DampingCoefficients, GridData @gtscript.function diff --git a/fv3core/pace/fv3core/stencils/dyn_core.py b/fv3core/pace/fv3core/stencils/dyn_core.py index bef8f6f0..6769a9b9 100644 --- a/fv3core/pace/fv3core/stencils/dyn_core.py +++ b/fv3core/pace/fv3core/stencils/dyn_core.py @@ -42,7 +42,7 @@ Z_DIM, Z_INTERFACE_DIM, ) -from pace.util.grid import DampingCoefficients, GridData +from pace.util.grid.helper import DampingCoefficients, GridData HUGE_R = 1.0e40 diff --git a/fv3core/pace/fv3core/stencils/fv_dynamics.py b/fv3core/pace/fv3core/stencils/fv_dynamics.py index 8e82f549..ba7225e6 100644 --- a/fv3core/pace/fv3core/stencils/fv_dynamics.py +++ b/fv3core/pace/fv3core/stencils/fv_dynamics.py @@ -21,7 +21,7 @@ from pace.fv3core.stencils.remapping import LagrangianToEulerian from pace.stencils.c2l_ord import CubedToLatLon from pace.util import X_DIM, Y_DIM, Z_INTERFACE_DIM, Timer, constants -from pace.util.grid import DampingCoefficients, GridData +from pace.util.grid.helper import DampingCoefficients, GridData from pace.util.logging import pace_log from pace.util.mpi import MPI diff --git a/fv3core/pace/fv3core/stencils/fvtp2d.py b/fv3core/pace/fv3core/stencils/fvtp2d.py index 3369bcab..6e30bea3 100644 --- a/fv3core/pace/fv3core/stencils/fvtp2d.py +++ b/fv3core/pace/fv3core/stencils/fvtp2d.py @@ -12,7 +12,7 @@ from pace.fv3core.stencils.xppm import XPiecewiseParabolic from pace.fv3core.stencils.yppm import YPiecewiseParabolic from pace.util import X_DIM, Y_DIM, Z_DIM -from pace.util.grid import DampingCoefficients, GridData +from pace.util.grid.helper import DampingCoefficients, GridData @gtscript.function diff --git a/fv3core/pace/fv3core/stencils/fxadv.py b/fv3core/pace/fv3core/stencils/fxadv.py index 1527c898..6b854b2b 100644 --- a/fv3core/pace/fv3core/stencils/fxadv.py +++ b/fv3core/pace/fv3core/stencils/fxadv.py @@ -11,7 +11,7 @@ from pace.dsl.stencil import StencilFactory from pace.dsl.typing import Float, FloatField, FloatFieldIJ from pace.fv3core.stencils.d2a2c_vect import contravariant -from pace.util.grid import GridData +from pace.util.grid.helper import GridData def main_uc_vc_contra( diff --git a/fv3core/pace/fv3core/stencils/nh_p_grad.py b/fv3core/pace/fv3core/stencils/nh_p_grad.py index 5504ba2b..4b56b847 100644 --- a/fv3core/pace/fv3core/stencils/nh_p_grad.py +++ b/fv3core/pace/fv3core/stencils/nh_p_grad.py @@ -6,7 +6,7 @@ from pace.dsl.typing import Float, FloatField, FloatFieldIJ from pace.fv3core.stencils.a2b_ord4 import AGrid2BGridFourthOrder from pace.util import X_DIM, Y_DIM, Z_INTERFACE_DIM -from pace.util.grid import GridData +from pace.util.grid.helper import GridData def set_k0_and_calc_wk( diff --git a/fv3core/pace/fv3core/stencils/updatedzd.py b/fv3core/pace/fv3core/stencils/updatedzd.py index 92dd1c20..dbc74667 100644 --- a/fv3core/pace/fv3core/stencils/updatedzd.py +++ b/fv3core/pace/fv3core/stencils/updatedzd.py @@ -18,7 +18,7 @@ Z_DIM, Z_INTERFACE_DIM, ) -from pace.util.grid import DampingCoefficients, GridData +from pace.util.grid.helper import DampingCoefficients, GridData DZ_MIN = constants.DZ_MIN diff --git a/fv3core/pace/fv3core/testing/translate_fvdynamics.py b/fv3core/pace/fv3core/testing/translate_fvdynamics.py index cdd773f7..3561d36c 100644 --- a/fv3core/pace/fv3core/testing/translate_fvdynamics.py +++ b/fv3core/pace/fv3core/testing/translate_fvdynamics.py @@ -12,7 +12,7 @@ from pace.fv3core.dycore_state import DycoreState from pace.stencils.testing import ParallelTranslateBaseSlicing from pace.stencils.testing.translate import TranslateFortranData2Py -from pace.util.grid import GridData +from pace.util.grid.helper import GridData class TranslateDycoreFortranData2Py(TranslateFortranData2Py): diff --git a/physics/pace/physics/stencils/microphysics.py b/physics/pace/physics/stencils/microphysics.py index 064bc1c7..5da2676a 100644 --- a/physics/pace/physics/stencils/microphysics.py +++ b/physics/pace/physics/stencils/microphysics.py @@ -18,7 +18,7 @@ from pace.dsl.stencil import StencilFactory from pace.dsl.typing import Float, FloatField, FloatFieldIJ, Int from pace.util import X_DIM, Y_DIM, Z_DIM -from pace.util.grid import GridData +from pace.util.grid.helper import GridData from .._config import PhysicsConfig diff --git a/physics/pace/physics/stencils/physics.py b/physics/pace/physics/stencils/physics.py index a4dfdb2f..cfc8f9f0 100644 --- a/physics/pace/physics/stencils/physics.py +++ b/physics/pace/physics/stencils/physics.py @@ -22,7 +22,7 @@ from pace.physics.stencils.get_prs_fv3 import get_prs_fv3 from pace.physics.stencils.microphysics import Microphysics from pace.util import X_DIM, Y_DIM, Z_DIM -from pace.util.grid import GridData +from pace.util.grid.helper import GridData from .._config import PhysicsConfig diff --git a/stencils/pace/stencils/c2l_ord.py b/stencils/pace/stencils/c2l_ord.py index e4610b69..ed0d5834 100644 --- a/stencils/pace/stencils/c2l_ord.py +++ b/stencils/pace/stencils/c2l_ord.py @@ -14,7 +14,7 @@ from pace.dsl.stencil import StencilFactory from pace.dsl.typing import Float, FloatField, FloatFieldIJ from pace.util.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from pace.util.grid import GridData +from pace.util.grid.helper import GridData A1 = 0.5625 diff --git a/stencils/pace/stencils/fv_update_phys.py b/stencils/pace/stencils/fv_update_phys.py index fe027cd0..2fe28b2e 100644 --- a/stencils/pace/stencils/fv_update_phys.py +++ b/stencils/pace/stencils/fv_update_phys.py @@ -11,7 +11,7 @@ from pace.stencils.c2l_ord import CubedToLatLon from pace.stencils.update_dwind_phys import AGrid2DGridPhysics from pace.util import X_DIM, Y_DIM -from pace.util.grid import DriverGridData, GridData +from pace.util.grid.helper import DriverGridData, GridData # TODO: This is the same as moist_cv.py in fv3core, should move to integration dir diff --git a/stencils/pace/stencils/testing/grid.py b/stencils/pace/stencils/testing/grid.py index 4cf623b1..e8d7919d 100644 --- a/stencils/pace/stencils/testing/grid.py +++ b/stencils/pace/stencils/testing/grid.py @@ -6,13 +6,13 @@ import pace.util from pace.dsl import gt4py_utils as utils from pace.dsl.stencil import GridIndexing -from pace.util.grid import ( +from pace.util.grid.generation import GridDefinitions +from pace.util.grid.helper import ( AngleGridData, ContravariantGridData, DampingCoefficients, DriverGridData, GridData, - GridDefinitions, HorizontalGridData, MetricTerms, VerticalGridData, diff --git a/stencils/pace/stencils/update_atmos_state.py b/stencils/pace/stencils/update_atmos_state.py index 789e40ea..d9b777ad 100644 --- a/stencils/pace/stencils/update_atmos_state.py +++ b/stencils/pace/stencils/update_atmos_state.py @@ -9,7 +9,7 @@ from pace.dsl.stencil import StencilFactory from pace.dsl.typing import Float, FloatField from pace.stencils.fv_update_phys import ApplyPhysicsToDycore -from pace.util.grid import DriverGridData, GridData +from pace.util.grid.helper import DriverGridData, GridData # TODO: when this file is not importable from physics or fv3core, import diff --git a/stencils/pace/stencils/update_dwind_phys.py b/stencils/pace/stencils/update_dwind_phys.py index f5d3242d..5823a989 100644 --- a/stencils/pace/stencils/update_dwind_phys.py +++ b/stencils/pace/stencils/update_dwind_phys.py @@ -5,7 +5,7 @@ from pace.dsl.stencil import StencilFactory from pace.dsl.typing import FloatField, FloatFieldI, FloatFieldIJ from pace.util import X_DIM, Y_DIM, Z_DIM -from pace.util.grid import DriverGridData +from pace.util.grid.helper import DriverGridData def set_winds_zero( diff --git a/util/pace/util/grid/__init__.py b/util/pace/util/grid/__init__.py index 2a755d0c..5e488743 100644 --- a/util/pace/util/grid/__init__.py +++ b/util/pace/util/grid/__init__.py @@ -21,4 +21,3 @@ VerticalGridData, ) from .stretch_transformation import direct_transform -from .utils import compute_eta, vertical_coordinate diff --git a/util/pace/util/grid/eta.py b/util/pace/util/grid/eta.py index 6f476461..ff459e68 100644 --- a/util/pace/util/grid/eta.py +++ b/util/pace/util/grid/eta.py @@ -4,7 +4,7 @@ import numpy as np import xarray as xr -from .utils import compute_eta +from pace.fv3core.initialization.init_utils import compute_eta @dataclass diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index a7a2258a..a0ba5d54 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -581,12 +581,12 @@ def dyc(self) -> util.Quantity: self._dx_center, self._dy_center = self._compute_dxdy_center() return self._dy_center - def ks(self, eta_file="") -> util.Quantity: + def ks(self, eta_file="None") -> util.Quantity: """ number of levels where the vertical coordinate is purely pressure-based """ if self._ks is None: - if not eta_file == "": + if not eta_file == "None": ( self._ks, self._ptop, @@ -597,13 +597,13 @@ def ks(self, eta_file="") -> util.Quantity: raise IOError("eta file is not specified") return self._ks - def ak(self, eta_file="") -> util.Quantity: + def ak(self, eta_file="None") -> util.Quantity: """ the ak coefficient used to calculate the pressure at a given k-level: pk = ak + (bk * ps) """ if self._ak is None: - if not eta_file == "": + if not eta_file == "None": ( self._ks, self._ptop, @@ -614,13 +614,13 @@ def ak(self, eta_file="") -> util.Quantity: raise IOError("eta file is not specified") return self._ak - def bk(self, eta_file="") -> util.Quantity: + def bk(self, eta_file="None") -> util.Quantity: """ the bk coefficient used to calculate the pressure at a given k-level: pk = ak + (bk * ps) """ if self._bk is None: - if not eta_file == "": + if not eta_file == "None": ( self._ks, self._ptop, @@ -631,12 +631,12 @@ def bk(self, eta_file="") -> util.Quantity: raise IOError("eta file is not specified") return self._bk - def ptop(self, eta_file="") -> util.Quantity: + def ptop(self, eta_file="None") -> util.Quantity: """ the pressure of the top of atmosphere level """ if self._ptop is None: - if not eta_file == "": + if not eta_file == "None": ( self._ks, self._ptop, @@ -2136,7 +2136,7 @@ def _compute_area_c_cartesian(self): area_cgrid_64.data[:, :] = self._dx_const * self._dy_const return quantity_cast_to_model_float(self.quantity_factory, area_cgrid_64) - def _set_hybrid_pressure_coefficients(self, eta_file=""): + def _set_hybrid_pressure_coefficients(self, eta_file="None"): ks = self.quantity_factory.zeros( [], "", diff --git a/util/pace/util/grid/helper.py b/util/pace/util/grid/helper.py index 9b7b3d3f..321c4def 100644 --- a/util/pace/util/grid/helper.py +++ b/util/pace/util/grid/helper.py @@ -144,7 +144,7 @@ def __post_init__(self): @classmethod def new_from_metric_terms( - cls, metric_terms: MetricTerms, eta_file="" + cls, metric_terms: MetricTerms, eta_file="None" ) -> "VerticalGridData": return cls( ak=metric_terms.ak(eta_file), @@ -323,7 +323,7 @@ def __init__( self._fC_agrid = None @classmethod - def new_from_metric_terms(cls, metric_terms: MetricTerms, eta_file=""): + def new_from_metric_terms(cls, metric_terms: MetricTerms, eta_file="None"): horizontal_data = HorizontalGridData.new_from_metric_terms(metric_terms) vertical_data = VerticalGridData.new_from_metric_terms( metric_terms, eta_file=eta_file diff --git a/util/pace/util/grid/utils.py b/util/pace/util/grid/utils.py deleted file mode 100644 index 17a15568..00000000 --- a/util/pace/util/grid/utils.py +++ /dev/null @@ -1,23 +0,0 @@ -import math - - -eta_0 = 0.252 -surface_pressure = 1.0e5 # units of (Pa), from Table VI of DCMIP2016 - - -def vertical_coordinate(eta_value): - """ - Equation (1) JRMS2006 - computes eta_v, the auxiliary variable vertical coordinate - """ - return (eta_value - eta_0) * math.pi * 0.5 - - -def compute_eta(ak, bk): - """ - Equation (1) JRMS2006 - eta is the vertical coordinate and eta_v is an auxiliary vertical coordinate - """ - eta = 0.5 * ((ak[:-1] + ak[1:]) / surface_pressure + bk[:-1] + bk[1:]) - eta_v = vertical_coordinate(eta) - return eta, eta_v