From 5dddd211e9132c086a5817a92490a33d28341b27 Mon Sep 17 00:00:00 2001 From: Oliver Elbert Date: Fri, 9 Dec 2022 13:55:28 -0500 Subject: [PATCH] Feature/fv3core fortran api (#395) his PR adds an API to call the Pace dycore from GEOS. Code changes: fv3core/pace/fv3core/initialization/geos_wrapper.py: Added a wrapper api to access the dynamical core from GEOS. Initialize it with a namelist, communicator, and backend, call it with numpy arrays and it returns a dictionary of numpy arrays. tests/main/fv3core/test_init_from_geos.py: A simple test that checks that the GEOS wrapper is initialized correctly, populates the dycore state correctly, steps through the dycore, and outputs arrays in the output dictionary. fv3core/pace/fv3core/init.py: Added the geos wrapper to init.py to make it importable at the base level. util/pace/util/grid/eta.py: Added code to setup atmosphere with 91 vertical levels for GEOS runs. util/pace/util/initialization/sizer.py: Revised the from_namelist method to work with our updated namelist configuration that doesn't use fv_core_nml as a key. --- fv3core/README.md | 4 + fv3core/pace/fv3core/__init__.py | 1 + .../fv3core/initialization/geos_wrapper.py | 417 ++++++++++++++++++ tests/main/fv3core/test_init_from_geos.py | 221 ++++++++++ util/pace/util/grid/eta.py | 360 ++++++++++++++- util/pace/util/initialization/sizer.py | 22 +- 6 files changed, 1018 insertions(+), 7 deletions(-) create mode 100644 fv3core/pace/fv3core/initialization/geos_wrapper.py create mode 100644 tests/main/fv3core/test_init_from_geos.py diff --git a/fv3core/README.md b/fv3core/README.md index d1842dd63..9357580b2 100644 --- a/fv3core/README.md +++ b/fv3core/README.md @@ -450,3 +450,7 @@ TEST_ARGS="-v -s --pdb" RUN_FLAGS="--rm -it" make tests ``` This can be done with any pytest target, such as `make savepoint_tests` and `make savepoint_tests_mpi`. + +### GEOS API + +The `GeosDycoreWrapper` class provides an API to run the dynamical core in a Python component of a GEOS model run. A `GeosDycoreWrapper` object is initialized with a namelist, communicator, and backend, which creates the communicators, partitioners, dycore state, and dycore object required to run the Pace dycore. A wrapper object takes numpy arrays of `u, v, w, delz, pt, delp, q, ps, pe, pk, peln, pkz, phis, q_con, omga, ua, va, uc, vc, mfxd, mfyd, cxd, cyd,` and `diss_estd` and returns a dictionary containing numpy arrays of those same variables. Wrapper objects contain a `timer` attrubite that tracks the amount of time moving input data to the dycore state, running the dynamical core, and retrieving the data from the state. diff --git a/fv3core/pace/fv3core/__init__.py b/fv3core/pace/fv3core/__init__.py index c02bf437b..5465ef279 100644 --- a/fv3core/pace/fv3core/__init__.py +++ b/fv3core/pace/fv3core/__init__.py @@ -1,5 +1,6 @@ from ._config import DynamicalCoreConfig from .initialization.dycore_state import DycoreState +from .initialization.geos_wrapper import GeosDycoreWrapper from .stencils.fv_dynamics import DynamicalCore from .stencils.fv_subgridz import DryConvectiveAdjustment diff --git a/fv3core/pace/fv3core/initialization/geos_wrapper.py b/fv3core/pace/fv3core/initialization/geos_wrapper.py new file mode 100644 index 000000000..301131c85 --- /dev/null +++ b/fv3core/pace/fv3core/initialization/geos_wrapper.py @@ -0,0 +1,417 @@ +from datetime import timedelta +from typing import Dict + +import f90nml +import numpy as np + +import pace.util +from pace import fv3core + + +class GeosDycoreWrapper: + """ + Provides an interface for the Geos model to access the Pace dycore. + Takes numpy arrays as inputs, returns a dictionary of numpy arrays as outputs + """ + + def __init__(self, namelist: f90nml.Namelist, comm: pace.util.Comm, backend: str): + self.timer = pace.util.Timer() + self.namelist = namelist + + self.dycore_config = fv3core.DynamicalCoreConfig.from_f90nml(self.namelist) + + self.layout = self.dycore_config.layout + partitioner = pace.util.CubedSpherePartitioner( + pace.util.TilePartitioner(self.layout) + ) + self.communicator = pace.util.CubedSphereCommunicator( + comm, partitioner, timer=self.timer + ) + + sizer = pace.util.SubtileGridSizer.from_namelist( + self.namelist, partitioner.tile, self.communicator.tile.rank + ) + quantity_factory = pace.util.QuantityFactory.from_backend( + sizer=sizer, backend=backend + ) + + # set up the metric terms and grid data + metric_terms = pace.util.grid.MetricTerms( + quantity_factory=quantity_factory, communicator=self.communicator + ) + grid_data = pace.util.grid.GridData.new_from_metric_terms(metric_terms) + + stencil_config = pace.dsl.stencil.StencilConfig( + compilation_config=pace.dsl.stencil.CompilationConfig( + backend=backend, rebuild=False, validate_args=False + ), + ) + + self._grid_indexing = pace.dsl.stencil.GridIndexing.from_sizer_and_communicator( + sizer=sizer, cube=self.communicator + ) + stencil_factory = pace.dsl.StencilFactory( + config=stencil_config, grid_indexing=self._grid_indexing + ) + + self.dycore_state = fv3core.DycoreState.init_zeros( + quantity_factory=quantity_factory + ) + + self.dycore_state.bdt = float(namelist["dt_atmos"]) + if "fv_core_nml" in namelist.keys(): + self.dycore_state.bdt = ( + float(namelist["dt_atmos"]) / namelist["fv_core_nml"]["k_split"] + ) + elif "dycore_config" in namelist.keys(): + self.dycore_state.bdt = ( + float(namelist["dt_atmos"]) / namelist["dycore_config"]["k_split"] + ) + else: + raise KeyError("Cannot find k_split in namelist") + + damping_coefficients = pace.util.grid.DampingCoefficients.new_from_metric_terms( + metric_terms + ) + + self.dynamical_core = fv3core.DynamicalCore( + comm=self.communicator, + grid_data=grid_data, + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + damping_coefficients=damping_coefficients, + config=self.dycore_config, + timestep=timedelta(seconds=self.dycore_config.dt_atmos), + phis=self.dycore_state.phis, + state=self.dycore_state, + ) + + self.output_dict: Dict[str, np.ndarray] = {} + self._allocate_output_dir() + + def __call__( + self, + u: np.ndarray, + v: np.ndarray, + w: np.ndarray, + delz: np.ndarray, + pt: np.ndarray, + delp: np.ndarray, + q: np.ndarray, + ps: np.ndarray, + pe: np.ndarray, + pk: np.ndarray, + peln: np.ndarray, + pkz: np.ndarray, + phis: np.ndarray, + q_con: np.ndarray, + omga: np.ndarray, + ua: np.ndarray, + va: np.ndarray, + uc: np.ndarray, + vc: np.ndarray, + mfxd: np.ndarray, + mfyd: np.ndarray, + cxd: np.ndarray, + cyd: np.ndarray, + diss_estd: np.ndarray, + ) -> Dict[str, np.ndarray]: + + with self.timer.clock("move_to_pace"): + self.dycore_state = self._put_fortran_data_in_dycore( + u, + v, + w, + delz, + pt, + delp, + q, + ps, + pe, + pk, + peln, + pkz, + phis, + q_con, + omga, + ua, + va, + uc, + vc, + mfxd, + mfyd, + cxd, + cyd, + diss_estd, + ) + + with self.timer.clock("dycore"): + self.dynamical_core.step_dynamics(state=self.dycore_state, timer=self.timer) + + with self.timer.clock("move_to_fortran"): + self.output_dict = self._prep_outputs_for_geos() + + return self.output_dict + + def _put_fortran_data_in_dycore( + self, + u: np.ndarray, + v: np.ndarray, + w: np.ndarray, + delz: np.ndarray, + pt: np.ndarray, + delp: np.ndarray, + q: np.ndarray, + ps: np.ndarray, + pe: np.ndarray, + pk: np.ndarray, + peln: np.ndarray, + pkz: np.ndarray, + phis: np.ndarray, + q_con: np.ndarray, + omga: np.ndarray, + ua: np.ndarray, + va: np.ndarray, + uc: np.ndarray, + vc: np.ndarray, + mfxd: np.ndarray, + mfyd: np.ndarray, + cxd: np.ndarray, + cyd: np.ndarray, + diss_estd: np.ndarray, + ) -> fv3core.DycoreState: + + isc = self._grid_indexing.isc + jsc = self._grid_indexing.jsc + iec = self._grid_indexing.iec + 1 + jec = self._grid_indexing.jec + 1 + + state = self.dycore_state + + # Assign compute domain: + pace.util.utils.safe_assign_array(state.u.view[:], u[isc:iec, jsc : jec + 1, :]) + pace.util.utils.safe_assign_array(state.v.view[:], v[isc : iec + 1, jsc:jec, :]) + pace.util.utils.safe_assign_array(state.w.view[:], w[isc:iec, jsc:jec, :]) + pace.util.utils.safe_assign_array(state.ua.view[:], ua[isc:iec, jsc:jec, :]) + pace.util.utils.safe_assign_array(state.va.view[:], va[isc:iec, jsc:jec, :]) + pace.util.utils.safe_assign_array( + state.uc.view[:], uc[isc : iec + 1, jsc:jec, :] + ) + pace.util.utils.safe_assign_array( + state.vc.view[:], vc[isc:iec, jsc : jec + 1, :] + ) + + pace.util.utils.safe_assign_array(state.delz.view[:], delz[isc:iec, jsc:jec, :]) + pace.util.utils.safe_assign_array(state.pt.view[:], pt[isc:iec, jsc:jec, :]) + pace.util.utils.safe_assign_array(state.delp.view[:], delp[isc:iec, jsc:jec, :]) + + pace.util.utils.safe_assign_array(state.mfxd.view[:], mfxd) + pace.util.utils.safe_assign_array(state.mfyd.view[:], mfyd) + pace.util.utils.safe_assign_array(state.cxd.view[:], cxd[:, jsc:jec, :]) + pace.util.utils.safe_assign_array(state.cyd.view[:], cyd[isc:iec, :, :]) + + pace.util.utils.safe_assign_array(state.ps.view[:], ps[isc:iec, jsc:jec]) + pace.util.utils.safe_assign_array( + state.pe.data[isc - 1 : iec + 1, jsc - 1 : jec + 1, :], pe + ) + pace.util.utils.safe_assign_array(state.pk.view[:], pk) + pace.util.utils.safe_assign_array(state.peln.view[:], peln) + pace.util.utils.safe_assign_array(state.pkz.view[:], pkz) + pace.util.utils.safe_assign_array(state.phis.view[:], phis[isc:iec, jsc:jec]) + pace.util.utils.safe_assign_array( + state.q_con.view[:], q_con[isc:iec, jsc:jec, :] + ) + pace.util.utils.safe_assign_array(state.omga.view[:], omga[isc:iec, jsc:jec, :]) + pace.util.utils.safe_assign_array( + state.diss_estd.view[:], diss_estd[isc:iec, jsc:jec, :] + ) + + # tracer quantities should be a 4d array in order: + # vapor, liquid, ice, rain, snow, graupel, cloud + pace.util.utils.safe_assign_array( + state.qvapor.view[:], q[isc:iec, jsc:jec, :, 0] + ) + pace.util.utils.safe_assign_array( + state.qliquid.view[:], q[isc:iec, jsc:jec, :, 1] + ) + pace.util.utils.safe_assign_array(state.qice.view[:], q[isc:iec, jsc:jec, :, 2]) + pace.util.utils.safe_assign_array( + state.qrain.view[:], q[isc:iec, jsc:jec, :, 3] + ) + pace.util.utils.safe_assign_array( + state.qsnow.view[:], q[isc:iec, jsc:jec, :, 4] + ) + pace.util.utils.safe_assign_array( + state.qgraupel.view[:], q[isc:iec, jsc:jec, :, 5] + ) + pace.util.utils.safe_assign_array(state.qcld.view[:], q[isc:iec, jsc:jec, :, 6]) + + return state + + def _prep_outputs_for_geos(self) -> Dict[str, np.ndarray]: + + output_dict = self.output_dict + isc = self._grid_indexing.isc + jsc = self._grid_indexing.jsc + iec = self._grid_indexing.iec + 1 + jec = self._grid_indexing.jec + 1 + + pace.util.utils.safe_assign_array( + output_dict["u"], self.dycore_state.u.data[:-1, :, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["v"], self.dycore_state.v.data[:, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["w"], self.dycore_state.w.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["ua"], self.dycore_state.ua.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["va"], self.dycore_state.va.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["uc"], self.dycore_state.uc.data[:, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["vc"], self.dycore_state.vc.data[:-1, :, :-1] + ) + + pace.util.utils.safe_assign_array( + output_dict["delz"], self.dycore_state.delz.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["pt"], self.dycore_state.pt.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["delp"], self.dycore_state.delp.data[:-1, :-1, :-1] + ) + + pace.util.utils.safe_assign_array( + output_dict["mfxd"], + self.dycore_state.mfxd.data[isc : iec + 1, jsc:jec, :-1], + ) + pace.util.utils.safe_assign_array( + output_dict["mfyd"], + self.dycore_state.mfyd.data[isc:iec, jsc : jec + 1, :-1], + ) + pace.util.utils.safe_assign_array( + output_dict["cxd"], self.dycore_state.cxd.data[isc : iec + 1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["cyd"], self.dycore_state.cyd.data[:-1, jsc : jec + 1, :-1] + ) + + pace.util.utils.safe_assign_array( + output_dict["ps"], self.dycore_state.ps.data[:-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["pe"], + self.dycore_state.pe.data[isc - 1 : iec + 1, jsc - 1 : jec + 1, :], + ) + pace.util.utils.safe_assign_array( + output_dict["pk"], self.dycore_state.pk.data[isc:iec, jsc:jec, :] + ) + pace.util.utils.safe_assign_array( + output_dict["peln"], self.dycore_state.peln.data[isc:iec, jsc:jec, :] + ) + pace.util.utils.safe_assign_array( + output_dict["pkz"], self.dycore_state.pkz.data[isc:iec, jsc:jec, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["phis"], self.dycore_state.phis.data[:-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["q_con"], self.dycore_state.q_con.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["omga"], self.dycore_state.omga.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["diss_estd"], self.dycore_state.diss_estd.data[:-1, :-1, :-1] + ) + + pace.util.utils.safe_assign_array( + output_dict["qvapor"], self.dycore_state.qvapor.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qliquid"], self.dycore_state.qliquid.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qice"], self.dycore_state.qice.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qrain"], self.dycore_state.qrain.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qsnow"], self.dycore_state.qsnow.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qgraupel"], self.dycore_state.qgraupel.data[:-1, :-1, :-1] + ) + pace.util.utils.safe_assign_array( + output_dict["qcld"], self.dycore_state.qcld.data[:-1, :-1, :-1] + ) + + return output_dict + + def _allocate_output_dir(self): + + nhalo = self._grid_indexing.n_halo + shape_centered = self._grid_indexing.domain_full(add=(0, 0, 0)) + shape_x_interface = self._grid_indexing.domain_full(add=(1, 0, 0)) + shape_y_interface = self._grid_indexing.domain_full(add=(0, 1, 0)) + shape_z_interface = self._grid_indexing.domain_full(add=(0, 0, 1)) + shape_2d = shape_centered[:-1] + + self.output_dict["u"] = np.empty((shape_y_interface)) + self.output_dict["v"] = np.empty((shape_x_interface)) + self.output_dict["w"] = np.empty((shape_centered)) + self.output_dict["ua"] = np.empty((shape_centered)) + self.output_dict["va"] = np.empty((shape_centered)) + self.output_dict["uc"] = np.empty((shape_x_interface)) + self.output_dict["vc"] = np.empty((shape_y_interface)) + + self.output_dict["delz"] = np.empty((shape_centered)) + self.output_dict["pt"] = np.empty((shape_centered)) + self.output_dict["delp"] = np.empty((shape_centered)) + + self.output_dict["mfxd"] = np.empty( + (self._grid_indexing.domain_full(add=(1 - 2 * nhalo, -2 * nhalo, 0))) + ) + self.output_dict["mfyd"] = np.empty( + (self._grid_indexing.domain_full(add=(-2 * nhalo, 1 - 2 * nhalo, 0))) + ) + self.output_dict["cxd"] = np.empty( + (self._grid_indexing.domain_full(add=(1 - 2 * nhalo, 0, 0))) + ) + self.output_dict["cyd"] = np.empty( + (self._grid_indexing.domain_full(add=(0, 1 - 2 * nhalo, 0))) + ) + + self.output_dict["ps"] = np.empty((shape_2d)) + self.output_dict["pe"] = np.empty( + (self._grid_indexing.domain_full(add=(2 - 2 * nhalo, 2 - 2 * nhalo, 1))) + ) + self.output_dict["pk"] = np.empty( + (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 1))) + ) + self.output_dict["peln"] = np.empty( + (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 1))) + ) + self.output_dict["pkz"] = np.empty( + (self._grid_indexing.domain_full(add=(-2 * nhalo, -2 * nhalo, 0))) + ) + self.output_dict["phis"] = np.empty((shape_2d)) + self.output_dict["q_con"] = np.empty((shape_centered)) + self.output_dict["omga"] = np.empty((shape_centered)) + self.output_dict["diss_estd"] = np.empty((shape_centered)) + + self.output_dict["qvapor"] = np.empty((shape_centered)) + self.output_dict["qliquid"] = np.empty((shape_centered)) + self.output_dict["qice"] = np.empty((shape_centered)) + self.output_dict["qrain"] = np.empty((shape_centered)) + self.output_dict["qsnow"] = np.empty((shape_centered)) + self.output_dict["qgraupel"] = np.empty((shape_centered)) + self.output_dict["qcld"] = np.empty((shape_centered)) diff --git a/tests/main/fv3core/test_init_from_geos.py b/tests/main/fv3core/test_init_from_geos.py new file mode 100644 index 000000000..252fe7d25 --- /dev/null +++ b/tests/main/fv3core/test_init_from_geos.py @@ -0,0 +1,221 @@ +import f90nml +import numpy as np +import pytest # noqa + +from pace import fv3core +from pace.util.null_comm import NullComm + + +def test_geos_wrapper(): + + namelist_dict = { + "stencil_config": { + "compilation_config": { + "backend": "numpy", + "rebuild": False, + "validate_args": True, + "format_source": False, + "device_sync": False, + } + }, + "initialization": {"type": "baroclinic"}, + "nx_tile": 12, + "nz": 91, + "dt_atmos": 225, + "minutes": 15, + "layout": [1, 1], + "dycore_config": { + "a_imp": 1.0, + "beta": 0.0, + "consv_te": 0.0, + "d2_bg": 0.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.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.0, + "rf_fast": True, + "tau": 10.0, + "vtdm4": 0.06, + "z_tracer": True, + "do_qa": True, + "tau_i2s": 1000.0, + "tau_g2v": 1200.0, + "ql_gen": 0.001, + "ql_mlt": 0.002, + "qs_mlt": 1e-06, + "qi_lim": 1.0, + "dw_ocean": 0.1, + "dw_land": 0.15, + "icloud_f": 0, + "tau_l2v": 300.0, + "tau_v2l": 90.0, + "fv_sg_adj": 0, + "n_sponge": 48, + }, + } + + namelist = f90nml.namelist.Namelist(namelist_dict) + + comm = NullComm(rank=0, total_ranks=6, fill_value=0.0) + backend = "numpy" + + wrapper = fv3core.GeosDycoreWrapper(namelist, comm, backend) + nhalo = 3 + shape_centered = ( + namelist["nx_tile"] + 2 * nhalo, + namelist["nx_tile"] + 2 * nhalo, + namelist["nz"], + ) + shape_x_interface = ( + namelist["nx_tile"] + 2 * nhalo + 1, + namelist["nx_tile"] + 2 * nhalo, + namelist["nz"], + ) + shape_y_interface = ( + namelist["nx_tile"] + 2 * nhalo, + namelist["nx_tile"] + 2 * nhalo + 1, + namelist["nz"], + ) + shape_z_interface = ( + namelist["nx_tile"] + 2 * nhalo, + namelist["nx_tile"] + 2 * nhalo, + namelist["nz"] + 1, + ) + shape2d = ( + namelist["nx_tile"] + 2 * nhalo, + namelist["nx_tile"] + 2 * nhalo, + ) + shape_tracers = ( + namelist["nx_tile"] + 2 * nhalo, + namelist["nx_tile"] + 2 * nhalo, + namelist["nz"], + 7, + ) + + assert isinstance(wrapper, fv3core.GeosDycoreWrapper) + assert isinstance(wrapper.dynamical_core, fv3core.DynamicalCore) + + u = np.ones(shape_y_interface) + v = np.ones(shape_x_interface) + w = np.ones(shape_centered) + delz = np.ones(shape_centered) + pt = np.ones(shape_centered) + delp = np.ones(shape_centered) + q = np.ones(shape_tracers) + ps = np.ones(shape2d) + pe = np.ones( + ( + namelist["nx_tile"] + 2, + namelist["nx_tile"] + 2, + namelist["nz"] + 1, + ) + ) + pk = np.ones( + ( + namelist["nx_tile"], + namelist["nx_tile"], + namelist["nz"] + 1, + ) + ) + peln = np.ones( + ( + namelist["nx_tile"], + namelist["nx_tile"], + namelist["nz"] + 1, + ) + ) + pkz = np.ones( + ( + namelist["nx_tile"], + namelist["nx_tile"], + namelist["nz"], + ) + ) + phis = np.ones(shape2d) + q_con = np.ones(shape_centered) + omga = np.ones(shape_centered) + ua = np.ones(shape_centered) + va = np.ones(shape_centered) + uc = np.ones(shape_x_interface) + vc = np.ones(shape_y_interface) + mfxd = np.ones( + ( + namelist["nx_tile"] + 1, + namelist["nx_tile"], + namelist["nz"], + ) + ) + mfyd = np.ones( + ( + namelist["nx_tile"], + namelist["nx_tile"] + 1, + namelist["nz"], + ) + ) + cxd = np.ones( + ( + namelist["nx_tile"] + 1, + namelist["nx_tile"] + 2 * nhalo, + namelist["nz"], + ) + ) + cyd = np.ones( + ( + namelist["nx_tile"] + 2 * nhalo, + namelist["nx_tile"] + 1, + namelist["nz"], + ) + ) + diss_estd = np.ones(shape_centered) + + output_dict = wrapper( + u, + v, + w, + delz, + pt, + delp, + q, + ps, + pe, + pk, + peln, + pkz, + phis, + q_con, + omga, + ua, + va, + uc, + vc, + mfxd, + mfyd, + cxd, + cyd, + diss_estd, + ) + + assert isinstance(output_dict["u"], np.ndarray) diff --git a/util/pace/util/grid/eta.py b/util/pace/util/grid/eta.py index 8c8063020..50afaadb3 100644 --- a/util/pace/util/grid/eta.py +++ b/util/pace/util/grid/eta.py @@ -204,14 +204,370 @@ def set_hybrid_pressure_coefficients(km: int) -> HybridPressureCoefficients: 1, ] ) + + elif km == 91: + + ak = np.array( + [ + 1.00000000, + 1.75000000, + 2.75000000, + 4.09999990, + 5.98951054, + 8.62932968, + 12.2572632, + 17.1510906, + 23.6545467, + 32.1627693, + 43.1310921, + 57.1100426, + 74.6595764, + 96.4470978, + 123.169769, + 155.601318, + 194.594009, + 241.047531, + 295.873840, + 360.046967, + 434.604828, + 520.628723, + 619.154846, + 731.296021, + 858.240906, + 1001.06561, + 1160.92859, + 1339.03992, + 1536.50012, + 1754.48938, + 1994.17834, + 2256.67407, + 2543.17139, + 2854.76392, + 3192.58569, + 3557.75366, + 3951.35107, + 4374.28662, + 4827.11084, + 5310.22168, + 5823.87793, + 6369.04248, + 6948.75244, + 7566.91992, + 8226.34277, + 8931.20996, + 9684.46191, + 10482.2725, + 11318.2793, + 12184.0771, + 13065.5674, + 13953.2207, + 14830.7285, + 15687.2617, + 16508.0645, + 17281.0996, + 17994.2988, + 18636.3223, + 19196.1797, + 19664.0723, + 20030.1914, + 20285.3691, + 20421.5254, + 20430.0684, + 20302.8730, + 20032.3711, + 19611.0664, + 19031.3848, + 18286.6426, + 17377.7930, + 16322.4639, + 15144.4033, + 13872.5674, + 12540.4785, + 11183.4170, + 9835.32715, + 8526.30664, + 7282.24512, + 6123.26074, + 5063.50684, + 4111.24902, + 3270.00122, + 2539.22729, + 1915.30762, + 1392.44995, + 963.134766, + 620.599365, + 357.989502, + 169.421387, + 51.0314941, + 2.48413086, + 0.00000000, + ] + ) + + bk = np.array( + [ + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 3.50123992e-06, + 2.81484008e-05, + 9.38666999e-05, + 2.28561999e-04, + 5.12343016e-04, + 1.04712998e-03, + 1.95625005e-03, + 3.42317997e-03, + 5.58632007e-03, + 8.65428988e-03, + 1.27844000e-02, + 1.81719996e-02, + 2.49934997e-02, + 3.34198996e-02, + 4.36249003e-02, + 5.57769015e-02, + 7.00351968e-02, + 8.65636021e-02, + 0.105520003, + 0.127051994, + 0.151319996, + 0.178477004, + 0.208675995, + 0.242069006, + 0.278813988, + 0.319043010, + 0.362558991, + 0.408596009, + 0.456384987, + 0.505111992, + 0.553902984, + 0.601903021, + 0.648333013, + 0.692534983, + 0.733981013, + 0.772292018, + 0.807236016, + 0.838724971, + 0.866774976, + 0.891497016, + 0.913065016, + 0.931702971, + 0.947658002, + 0.961175978, + 0.972495019, + 0.981844008, + 0.989410996, + 0.995342016, + 1.00000000, + ] + ) + + elif km == 72: + + ak = np.array( + [ + 1.00000000, + 2.00000024, + 3.27000046, + 4.75850105, + 6.60000134, + 8.93450165, + 11.9703016, + 15.9495029, + 21.1349030, + 27.8526058, + 36.5041084, + 47.5806084, + 61.6779099, + 79.5134125, + 101.944023, + 130.051025, + 165.079025, + 208.497040, + 262.021057, + 327.643066, + 407.657104, + 504.680115, + 621.680115, + 761.984192, + 929.294189, + 1127.69019, + 1364.34021, + 1645.71033, + 1979.16040, + 2373.04053, + 2836.78052, + 3381.00073, + 4017.54102, + 4764.39111, + 5638.79102, + 6660.34131, + 7851.23145, + 9236.57227, + 10866.3018, + 12783.7031, + 15039.3027, + 17693.0039, + 20119.2012, + 21686.5020, + 22436.3008, + 22389.8008, + 21877.5977, + 21214.9980, + 20325.8984, + 19309.6953, + 18161.8965, + 16960.8965, + 15625.9961, + 14290.9951, + 12869.5938, + 11895.8623, + 10918.1709, + 9936.52148, + 8909.99219, + 7883.42188, + 7062.19824, + 6436.26367, + 5805.32129, + 5169.61084, + 4533.90088, + 3898.20093, + 3257.08081, + 2609.20068, + 1961.31055, + 1313.48035, + 659.375244, + 4.80482578, + 0.00000000, + ] + ) + + bk = np.array( + [ + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 0.00000000, + 8.17541324e-09, + 6.96002459e-03, + 2.80100405e-02, + 6.37200624e-02, + 0.113602079, + 0.156224087, + 0.200350106, + 0.246741116, + 0.294403106, + 0.343381137, + 0.392891139, + 0.443740189, + 0.494590193, + 0.546304166, + 0.581041515, + 0.615818441, + 0.650634944, + 0.685899913, + 0.721165955, + 0.749378204, + 0.770637512, + 0.791946948, + 0.813303947, + 0.834660947, + 0.856018007, + 0.877429008, + 0.898908019, + 0.920387030, + 0.941865027, + 0.963406026, + 0.984951973, + 1.00000000, + ] + ) + else: raise NotImplementedError( - "Only a 79 vertical level grid has been implemented so far" + "Only grids with 72, 79, or 91 vertical levels have been implemented so far" ) if 0.0 in bk: - ks = np.where(bk == 0)[0][-1] + ks = 0 if km == 91 else np.where(bk == 0)[0][-1] ptop = ak[0] else: raise ValueError("bk must contain at least one 0.") + pressure_data = HybridPressureCoefficients(ks, ptop, ak, bk) return pressure_data diff --git a/util/pace/util/initialization/sizer.py b/util/pace/util/initialization/sizer.py index 10ec26e9b..e787119fe 100644 --- a/util/pace/util/initialization/sizer.py +++ b/util/pace/util/initialization/sizer.py @@ -86,11 +86,23 @@ def from_namelist( different ranks have different domain shapes. If tile_partitioner is a TilePartitioner, this argument does not matter. """ - layout = namelist["fv_core_nml"]["layout"] - # npx and npy in the namelist are cell centers, but npz is mid levels - nx_tile = namelist["fv_core_nml"]["npx"] - 1 - ny_tile = namelist["fv_core_nml"]["npy"] - 1 - nz = namelist["fv_core_nml"]["npz"] + if "fv_core_nml" in namelist.keys(): + layout = namelist["fv_core_nml"]["layout"] + # npx and npy in the namelist are cell centers, but npz is mid levels + nx_tile = namelist["fv_core_nml"]["npx"] - 1 + ny_tile = namelist["fv_core_nml"]["npy"] - 1 + nz = namelist["fv_core_nml"]["npz"] + elif "nx_tile" in namelist.keys(): + layout = namelist["layout"] + # everything is cell centered in this format + nx_tile = namelist["nx_tile"] + ny_tile = namelist["nx_tile"] + nz = namelist["nz"] + else: + raise KeyError( + "Namelist format is unrecognized, " + "expected to find nx_tile or fv_core_nml" + ) return cls.from_tile_params( nx_tile, ny_tile,