From cd1bd06672f409cf61bcc15606788a7e16f0eff6 Mon Sep 17 00:00:00 2001 From: Frank <142349306+fmalatino@users.noreply.github.com> Date: Tue, 30 Jan 2024 14:07:03 -0500 Subject: [PATCH] Benchmark with DaCe cpu and gpu backends (#50) * Benchmark changes * Added benchmark configurations * Removed benchmark configs from configs to be tested in unit tests * Changed dace submodule to point to Florian fix * Dace submodule points to gcc_dies_on_dacecpu branch * Set 'rf_fast' to true in baroclinic_c384_cpu/gpu.yaml * fix mpi4py version (#51) * [Feature] Better translate test (#39) (#47) Translate test: small improvements - Parametrize perturbation upon failure - Refactor error folder to be `pwd` based - Fix GPU translate unable to dump error `nc` - Fix mixmatch precision on translate test - Update README.md Test fix: - Orchestrate YPPM for translate purposes Misc: - Fix bad logger formatting on DaCeProgress * [NASA] [Feature] Guarding against unimplemented configuration (#40) (#48) * [Feature] Guarding against unimplemented configuration (#40) Guarding against unimplemented namelists options: - a2b_ord4 - d_sw - fv_dynamics - fv_subgridz - neg_adj3 - divergence damping - xppm - yppm Misc: - Fix `netcdf_monitor` not mkdir the directory - Add `as_dict` to the dycore state to dump the dycore more easily * Unused assert * Update fv3core/pace/fv3core/stencils/yppm.py Co-authored-by: Oliver Elbert * Update fv3core/pace/fv3core/stencils/xppm.py Co-authored-by: Oliver Elbert * Change NotImplemented to ValueError for n_sponge<3 * lint --------- Co-authored-by: Oliver Elbert * Re-try of updating dace submodule to track Florian fix branch * Reverted gt4py submodule to match main checkout * Added benchmark README file * Read in ak, bk coefficients (#36) * initial changes to read in ak bk * read ak/bk * add xfail * remove input dir * further changes to unit tests * finish up test * add history * commit uncommited files * fix test comment * add input to top * read in data * read in netcdf file in eta mod * remove txt file * test * modify test and fix generate.py * remove emacs backup file * driver tests pass * fix helper.py * fix fv3core tests * fix physics test * fix grid tests * nullcommconfig * cleanup input * remove driver input * remove top level input * fix circular import problems * modify eta_file readin for test_restart_serial * comment out 91 test * rm safety checks * revert diagnostics.py * restore driver.py * revert initialization.py * restore state.py * restore analytic_init.py * restore init_utils.py and analytic_init.py * restore c_sw.py * d2a2c_vect.py * restore fv3core/stensils * restore translate_fvdynamics * restore physics/stencils * restore stencils * remove circular dependency * use pytest parametrize * cleanup generation.py * fstrinngs * add eta_file to MetricTerm init * remove eta_file argument in new_from_metric_terms and geos_wrapper * use pytest parametrize for the xfail tests * use pytest parametrize for the xfail tests * fix geos_wrapper and grid * fix tests * fstring * add test comments * fix util/HISTORY.md * fix comments * remove __init__.py from tests/main/grid * add jupyter notebooks to generate eta files * generate ak,bk,ptop on metricterm init * fix tests * exploit np.all in eta mod * remove tests/main/grid/input * update ci * test * remove input * edit ci yaml * remove push --------- Co-authored-by: mlee03 * Move Active Physics Schemes to Config (#44) * initial commit, need to adapt and run tests * revising scheme name * tests pass * update history * linting * changing typehints for physics schemes to enum instead of str * driver now works with physics config enum, tests pass * fixed tests * missed one * D-grid data input method (#42) * Testing changes reflected across branches * Undoing changes made in build_gaea_c5.sh * Testing vscode functionality, by adding a change to external_grid branch * Testing vscode functionality, by adding a change to external_grid branch * Addition of from_generated method and calc_flag to util/pace/util/grid/generation.py * Added get_grid method for external grid data to driver/pace/driver/grid.py * Preliminary xarray netcdf read in method added to driver/pace/driver/grid.py * Updating util/pace/util/grid/generation.py from_generated method * Addition of external grid data read in methods for initialization of grid. Current method uses xarray to interact with netcdf tile files. Values for longitutde, latitude, number of points in x an y, grid edge distances read in. * driver/examples/configs/test_external_C12_1x1.yaml * Preliminary unit test for external grid data read in * Current state of unit tests as of 27 Nov 2023 * External grid method and unit tests added * Re-excluding external grid data yamls from test_example_configs.py * Update driver/pace/driver/grid.py Co-authored-by: Florian Deconinck * Changed name of grid initializer function to match NetCDF dependency and class descriptor * Update util/pace/util/grid/generation.py Moved position of doc string for "from_external" MetricTerms class method Co-authored-by: Oliver Elbert * Fixed indentation error in generation.py from suggestion in PR 42 * Removal of TODO comment in grid.py, changes to method of file accessing in test_analytic_init, test_external_grid_* * Changed grid data read-in unit tests to compare data directly from file to driver grid data generated from yaml * Change to reading in lon and lat, other metric terms calculated as needed * Removed read-in of dx, dy, and area. Changed unit tests to compare calculated area to 'ideal' surface area as given by selected constants type. * Update tests/mpi_54rank/test_external_grid_1x1.py Incorrect name of test in test_external_grid_1x1.py changed to match file name Co-authored-by: Oliver Elbert * Added comparisons for read-in vs generated by driver lon, lat, dx, dy, and area data to unit tests * Added relative error calculations to unit tests for external grid data read-in * External grid data read in tests changed: relative errors printed by each rank and get_tile_number replacing get_tile_index * Removing commented out sections in test_external_grid_2x2.py Co-authored-by: Oliver Elbert * Updated external grid data read-in to take configuration and input data locations from command line, updated test description, and added documentation on grid construction to external grid data configuration selection dataclass. * Updated documentation in grid.py * Updated external grid data read in unit test to use parametrize functionality of pytest * Ammended files to reference changes to PR 36 --------- Co-authored-by: Frank Malatino Co-authored-by: Florian Deconinck Co-authored-by: Oliver Elbert --------- Co-authored-by: Oliver Elbert Co-authored-by: Florian Deconinck Co-authored-by: Oliver Elbert Co-authored-by: MiKyung Lee <58964324+mlee03@users.noreply.github.com> Co-authored-by: mlee03 Co-authored-by: Frank Malatino --- .github/workflows/main_unit_tests.yml | 4 + .gitmodules | 6 +- CONTRIBUTORS.md | 2 + benchmark_README.md | 96 ++ constraints.txt | 4 +- docs/physics/state.rst | 2 +- driver/examples/configs/baroclinic_c12.yaml | 5 + .../baroclinic_c12_explicit_physics.yaml | 98 ++ .../configs/baroclinic_c12_orch_cpu.yaml | 26 +- .../configs/baroclinic_c12_write_restart.yaml | 5 + .../examples/configs/baroclinic_c384_cpu.yaml | 9 +- .../examples/configs/baroclinic_c384_gpu.yaml | 2 +- .../configs/test_external_C12_1x1.yaml | 103 +++ .../configs/test_external_C12_2x2.yaml | 103 +++ driver/pace/driver/driver.py | 5 +- driver/pace/driver/grid.py | 97 +- driver/pace/driver/initialization.py | 16 +- driver/pace/driver/state.py | 9 +- driver/tests/mpi/test_restart.py | 2 + dsl/pace/dsl/dace/utils.py | 4 +- examples/Dockerfile | 2 +- .../notebooks/generate_eta_file_netcdf.ipynb | 148 +++ .../notebooks/generate_eta_file_xarray.ipynb | 140 +++ external/dace | 2 +- fv3core/pace/fv3core/dycore_state.py | 11 +- .../fv3core/initialization/analytic_init.py | 8 +- fv3core/pace/fv3core/stencils/a2b_ord4.py | 7 +- fv3core/pace/fv3core/stencils/d_sw.py | 38 +- .../fv3core/stencils/divergence_damping.py | 4 +- fv3core/pace/fv3core/stencils/dyn_core.py | 9 +- fv3core/pace/fv3core/stencils/fv_dynamics.py | 28 +- fv3core/pace/fv3core/stencils/fv_subgridz.py | 14 +- fv3core/pace/fv3core/stencils/neg_adj3.py | 5 +- .../pace/fv3core/stencils/remap_profile.py | 6 +- fv3core/pace/fv3core/stencils/xppm.py | 6 +- fv3core/pace/fv3core/stencils/yppm.py | 10 +- fv3core/pace/fv3core/wrappers/geos_wrapper.py | 4 +- .../savepoint/translate/translate_fvtp2d.py | 5 +- .../savepoint/translate/translate_yppm.py | 5 +- physics/pace/physics/__init__.py | 2 +- physics/pace/physics/_config.py | 20 +- physics/pace/physics/physics_state.py | 24 +- physics/pace/physics/stencils/microphysics.py | 3 +- physics/pace/physics/stencils/physics.py | 23 +- .../savepoint/translate/translate_driver.py | 5 +- .../translate/translate_gfs_physics_driver.py | 6 +- .../translate/translate_microphysics.py | 3 +- stencils/pace/stencils/testing/README.md | 10 + stencils/pace/stencils/testing/grid.py | 3 +- .../pace/stencils/testing/test_translate.py | 28 +- tests/main/driver/test_analytic_init.py | 9 +- tests/main/driver/test_example_configs.py | 7 + tests/main/driver/test_restart_fortran.py | 6 +- tests/main/driver/test_restart_serial.py | 7 +- tests/main/fv3core/test_cartesian_grid.py | 3 +- tests/main/fv3core/test_dycore_call.py | 8 +- tests/main/fv3core/test_init_from_geos.py | 3 + tests/main/grid/test_eta.py | 110 +++ tests/main/physics/test_integration.py | 4 +- tests/main/test_grid_init.py | 5 +- .../test_ext_grid/test_external_grid.py | 200 +++++ tests/mpi_54rank/test_grid_init.py | 4 +- util/HISTORY.md | 3 + util/pace/util/__init__.py | 1 + util/pace/util/grid/eta.py | 841 +----------------- util/pace/util/grid/generation.py | 85 +- util/pace/util/monitor/netcdf_monitor.py | 1 + util/pace/util/utils.py | 6 + 68 files changed, 1512 insertions(+), 968 deletions(-) create mode 100644 benchmark_README.md create mode 100644 driver/examples/configs/baroclinic_c12_explicit_physics.yaml create mode 100644 driver/examples/configs/test_external_C12_1x1.yaml create mode 100644 driver/examples/configs/test_external_C12_2x2.yaml create mode 100644 examples/notebooks/generate_eta_file_netcdf.ipynb create mode 100644 examples/notebooks/generate_eta_file_xarray.ipynb create mode 100755 tests/main/grid/test_eta.py create mode 100644 tests/mpi_54rank/test_ext_grid/test_external_grid.py diff --git a/.github/workflows/main_unit_tests.yml b/.github/workflows/main_unit_tests.yml index 5800baa6..b8a1289b 100644 --- a/.github/workflows/main_unit_tests.yml +++ b/.github/workflows/main_unit_tests.yml @@ -22,6 +22,10 @@ jobs: run: | python -m pip install --upgrade pip pip install -r requirements_dev.txt + - name: Clone datafiles + run: | + mkdir -p tests/main/input && cd tests/main/input + git clone -b store_files https://github.com/mlee03/pace.git tmp && mv tmp/*.nc . && rm -rf tmp - name: Run all main tests run: | pytest -x tests/main diff --git a/.gitmodules b/.gitmodules index 60de021d..66cd22e8 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,6 +1,8 @@ [submodule "external/gt4py"] path = external/gt4py url = https://github.com/gridtools/gt4py.git -[submodule "external/dace"] + +[submodule "dacefix"] path = external/dace - url = https://github.com/spcl/dace.git + url = https://github.com/FlorianDeconinck/dace.git + branch = fix/gcc_dies_on_dacecpu diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 0ca00ff5..72dda3d5 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -11,7 +11,9 @@ List format (alphabetical order): Surname, Name. Employer/Affiliation * Fuhrer, Oliver. Allen Institute for AI. * George, Rhea. Allen Institute for AI. * Harris, Lucas. GFDL. +* Lee, Mi Kyung. GFDL. * Kung, Chris. NASA. +* Malatino, Frank. GFDL * McGibbon, Jeremy. Allen Institute for AI. * Niedermayr, Yannick. ETH. * Savarin, Ajda. University of Washington. diff --git a/benchmark_README.md b/benchmark_README.md new file mode 100644 index 00000000..674d22fb --- /dev/null +++ b/benchmark_README.md @@ -0,0 +1,96 @@ +# Benchmarking README + +The tests contained in this archive are for benchmarking purposes only. Any +distribution beyond those personnel performing the tests need explicit approval +from NOAA/GFDL (Seth Underwood or Rusty Benson). + +## Cloning benchmark repository and generating conda environment + +Pace requires GCC > 9.2, MPI, and Python 3.8 on your system, and CUDA is required to run with a GPU backend. You will also need the headers of the boost libraries in your `$PATH` (boost itself does not need to be installed). + +```shell +cd BOOST/ROOT +wget https://boostorg.jfrog.io/artifactory/main/release/1.79.0/source/boost_1_79_0.tar.gz +tar -xzf boost_1_79_0.tar.gz +mkdir -p boost_1_79_0/include +mv boost_1_79_0/boost boost_1_79_0/include/ +export BOOST_ROOT=BOOST/ROOT/boost_1_79_0 +``` + +To clone the benchmark branch use the command: + +```shell +git clone --recursive -b benchmark git@github.com:NOAA-GFDL/pace.git +``` + +or if you have already cloned the repository: + +```shell +git submodule update --init --recursive +``` + +After cloning, change into the directory containing the clone. To generate the conda environment use the following commands: + +```shell +conda create -y --name python=3.8 +conda activate +pip3 install --upgrade pip setuptools wheel +pip3 install -r requirements_dev.txt -c constraints.txt +``` + +## Benchmarking configurations + +There are four configurations of the PACE application contained within the branch to be used for benchmarking: + +```shell +driver/examples/configs/baroclinic_c384_cpu.yaml +driver/examples/configs/baroclinic_c384_gpu.yaml +driver/examples/configs/baroclinic_c3072_cpu.yaml +driver/examples/configs/baroclinic_c3072_gpu.yaml +``` + +## Building + +To build with the DaCe backends, set the following environment variables: + +```shell +FV3_DACEMODE=Build +PACE_FLOAT_PRECISION=64 +PACE_LOGLEVEL=INFO +PYTHONOPTIMIZE=1 +OMP_NUM_THREAD=1 +``` + +Adjust the time of the configuration to be built such that the time of the build is for one timestep. For example: + +```shell +dt_atmos: 450 +seconds: 450 +``` +## Running +To build with the DaCe backends, set the following environment variables: + +```shell +FV3_DACEMODE=Run +PACE_FLOAT_PRECISION=64 +PACE_LOGLEVEL=INFO +PYTHONOPTIMIZE=1 +OMP_NUM_THREAD=1 +``` + +Adjust the time of the configuration to be run to the desired length, example: + +```shell +dt_atmos: 450 +days: 9 +``` + +The time for the build or run can be set with units of seconds, minutes, hours, or days. + +An example command to start the build or run process with MPI using the DaCe CPU backend for the c384 configuration: + +```shell +mpirun -n 1536 python3 -m pace.driver.run driver/examples/configs/baroclinic_c384_cpu.yaml +``` + +The build or run requires 1536 ranks, given that layout of 16x16 ranks per tile, and there are 6 tiles. diff --git a/constraints.txt b/constraints.txt index 73f71625..1cdfa7ce 100644 --- a/constraints.txt +++ b/constraints.txt @@ -81,7 +81,7 @@ coverage==5.5 # pytest-cov cytoolz==0.12.1 # via gt4py -dace==0.14.4 +dace==0.15.1 # via # -r requirements_dev.txt # pace-dsl @@ -184,7 +184,7 @@ googleapis-common-protos==1.53.0 # via google-api-core gprof2dot==2021.2.21 # via pytest-profiling -gridtools-cpp==2.3.0 +gridtools-cpp==2.3.1 # via gt4py h5netcdf==0.11.0 # via -r util/requirements.txt diff --git a/docs/physics/state.rst b/docs/physics/state.rst index d658171e..1b603887 100644 --- a/docs/physics/state.rst +++ b/docs/physics/state.rst @@ -38,6 +38,6 @@ You can initialize a zero-filled PhysicsState and MicrophysicsState from other P >>> quantity_factory = QuantityFactory.from_backend(sizer=sizer, backend="numpy") >>> physics_state = PhysicsState.init_zeros( - ... quantity_factory=quantity_factory, active_packages=["microphysics"] + ... quantity_factory=quantity_factory, schemes=["GFS_microphysics"] ... ) >>> microphysics_state = physics_state.microphysics diff --git a/driver/examples/configs/baroclinic_c12.yaml b/driver/examples/configs/baroclinic_c12.yaml index 1785f9ab..2b0a2e3f 100644 --- a/driver/examples/configs/baroclinic_c12.yaml +++ b/driver/examples/configs/baroclinic_c12.yaml @@ -94,3 +94,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/examples/configs/baroclinic_c12_explicit_physics.yaml b/driver/examples/configs/baroclinic_c12_explicit_physics.yaml new file mode 100644 index 00000000..ceed306b --- /dev/null +++ b/driver/examples/configs/baroclinic_c12_explicit_physics.yaml @@ -0,0 +1,98 @@ +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 +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 + schemes: + - GFS_microphysics diff --git a/driver/examples/configs/baroclinic_c12_orch_cpu.yaml b/driver/examples/configs/baroclinic_c12_orch_cpu.yaml index d70699a1..4e7b17ff 100644 --- a/driver/examples/configs/baroclinic_c12_orch_cpu.yaml +++ b/driver/examples/configs/baroclinic_c12_orch_cpu.yaml @@ -14,10 +14,30 @@ performance_config: nx_tile: 12 nz: 79 dt_atmos: 225 -minutes: 5 +seconds: 675 layout: - - 1 - - 1 + - 2 + - 2 +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. 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/examples/configs/baroclinic_c384_cpu.yaml b/driver/examples/configs/baroclinic_c384_cpu.yaml index 73916c7e..ab6a92cc 100644 --- a/driver/examples/configs/baroclinic_c384_cpu.yaml +++ b/driver/examples/configs/baroclinic_c384_cpu.yaml @@ -15,11 +15,10 @@ performance_config: nx_tile: 384 nz: 79 dt_atmos: 450 -minutes: 7 -seconds: 30 +days: 9 layout: - - 1 - - 1 + - 16 + - 16 diagnostics_config: path: output output_format: netcdf @@ -72,7 +71,7 @@ dycore_config: nwat: 6 p_fac: 0.1 rf_cutoff: 800. - rf_fast: false + rf_fast: true tau: 5. vtdm4: 0.06 z_tracer: true diff --git a/driver/examples/configs/baroclinic_c384_gpu.yaml b/driver/examples/configs/baroclinic_c384_gpu.yaml index 273d85ee..d7355db7 100644 --- a/driver/examples/configs/baroclinic_c384_gpu.yaml +++ b/driver/examples/configs/baroclinic_c384_gpu.yaml @@ -71,7 +71,7 @@ dycore_config: nwat: 6 p_fac: 0.1 rf_cutoff: 800. - rf_fast: false + rf_fast: true tau: 5. vtdm4: 0.06 z_tracer: true diff --git a/driver/examples/configs/test_external_C12_1x1.yaml b/driver/examples/configs/test_external_C12_1x1.yaml new file mode 100644 index 00000000..fb615b4e --- /dev/null +++ b/driver/examples/configs/test_external_C12_1x1.yaml @@ -0,0 +1,103 @@ +stencil_config: + compilation_config: + backend: numpy + rebuild: false + validate_args: true + format_source: false + device_sync: false +grid_config: + type: external + config: + grid_type: 0 + grid_file_path: "../../../test_input/C12.tile" + eta_file: '../../../tests/main/input/eta79.nc' +initialization: + type: analytic + config: + case: baroclinic +performance_config: + collect_performance: true + experiment_name: c12_baroclinic +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 + u_max: 355.0 + +physics_config: + hydrostatic: false + nwat: 6 + do_qa: true diff --git a/driver/examples/configs/test_external_C12_2x2.yaml b/driver/examples/configs/test_external_C12_2x2.yaml new file mode 100644 index 00000000..d07d142b --- /dev/null +++ b/driver/examples/configs/test_external_C12_2x2.yaml @@ -0,0 +1,103 @@ +stencil_config: + compilation_config: + backend: numpy + rebuild: false + validate_args: true + format_source: false + device_sync: false +grid_config: + type: external + config: + grid_type: 0 + grid_file_path: "../../../test_input/C12.tile" + eta_file: '../../../tests/main/input/eta79.nc' +initialization: + type: analytic + config: + case: baroclinic +performance_config: + collect_performance: true + experiment_name: c12_baroclinic +nx_tile: 12 +nz: 79 +dt_atmos: 225 +minutes: 15 +layout: + - 2 + - 2 +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 + u_max: 355.0 + +physics_config: + hydrostatic: false + nwat: 6 + do_qa: true diff --git a/driver/pace/driver/driver.py b/driver/pace/driver/driver.py index c8d1490f..24197621 100644 --- a/driver/pace/driver/driver.py +++ b/driver/pace/driver/driver.py @@ -231,6 +231,7 @@ def get_driver_state( damping_coefficients=damping_coefficients, driver_grid_data=driver_grid_data, grid_data=grid_data, + schemes=self.physics_config.schemes, ) @classmethod @@ -327,6 +328,9 @@ def write_for_restart( config_dict["initialization"]["type"] = "restart" config_dict["initialization"]["config"]["start_time"] = time config_dict["initialization"]["config"]["path"] = restart_path + # convert physics package enum to str + schemes = [scheme.value for scheme in config_dict["physics_config"]["schemes"]] + config_dict["physics_config"]["schemes"] = schemes # restart config doesn't have 'case' if "case" in config_dict["initialization"]["config"].keys(): del config_dict["initialization"]["config"]["case"] @@ -508,7 +512,6 @@ def exit_instead_of_build(self): quantity_factory=self.quantity_factory, grid_data=self.state.grid_data, namelist=self.config.physics_config, - active_packages=["microphysics"], ) else: # Make sure those are set to None to raise any issues diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 1cf59d07..e1c3b057 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -3,11 +3,13 @@ from typing import ClassVar, Optional, Tuple import f90nml +import xarray as xr import pace.driver import pace.dsl import pace.physics import pace.stencils +import pace.util import pace.util.grid from pace.stencils.testing import TranslateGrid from pace.util import Communicator, QuantityFactory @@ -65,7 +67,8 @@ def get_grid( communicator: Communicator, ) -> Tuple[DampingCoefficients, DriverGridData, GridData]: return self.config.get_grid( - quantity_factory=quantity_factory, communicator=communicator + quantity_factory=quantity_factory, + communicator=communicator, ) @classmethod @@ -99,6 +102,7 @@ class GeneratedGridConfig(GridInitializer): dx_const: Optional[float] = 1000.0 dy_const: Optional[float] = 1000.0 deglat: Optional[float] = 15.0 + eta_file: str = "None" def get_grid( self, @@ -112,6 +116,7 @@ def get_grid( dx_const=self.dx_const, dy_const=self.dy_const, deglat=self.deglat, + eta_file=self.eta_file, ) if self.stretch_factor != 1: # do horizontal grid transformation _transform_horizontal_grid( @@ -196,6 +201,96 @@ def get_grid( return damping_coefficients, driver_grid_data, grid_data +@GridInitializerSelector.register("external") +@dataclasses.dataclass +class ExternalNetcdfGridConfig(GridInitializer): + """ + Configuration for grid initialized from external data. + Input is from tile files generated by FRE-NCtools methods + The ExternalNetcdfGridConfig get_grid member method calls + the MetricTerms class method from_generated which generates + an object of MetricTerms to be used to generate the + damping_coefficients, driver_grid_data, and grid_data variables + We do not read in the dx, dy, or area values as there may be + inconsistencies in the constants used during calculation of the + input data and the model use. An example of two adjacent finite + volume cells in the supergrid should appear like: + + X----X----X----X----X + | | | + X X X X X + | | | + X----X----X----X----X + + The grid data must define the verticies, centroids, and mid-points + on edge of the cells contained in the computation. For more information + on grid discretization for the FV3 dynamical core please visit: + https://www.gfdl.noaa.gov/fv3/ + """ + + grid_type: Optional[int] = 0 + grid_file_path: str = "/input/tilefile" + eta_file: str = "None" + + def get_grid( + self, + quantity_factory: QuantityFactory, + communicator: Communicator, + ) -> Tuple[DampingCoefficients, DriverGridData, GridData]: + + pace_log.info("Using external grid data") + + # ToDo: refactor when grid_type is an enum + if self.grid_type <= 3: + tile_num = ( + pace.util.get_tile_index( + communicator.rank, communicator.partitioner.total_ranks + ) + + 1 + ) + tile_file = self.grid_file_path + str(tile_num) + ".nc" + else: + tile_file = self.grid_file_path + + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = communicator.partitioner.tile.subtile_slice( + rank=communicator.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npy, npx), + overlap=True, + ) + + metric_terms = MetricTerms.from_external( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + quantity_factory=quantity_factory, + communicator=communicator, + grid_type=self.grid_type, + eta_file=self.eta_file, + ) + + horizontal_data = HorizontalGridData.new_from_metric_terms(metric_terms) + vertical_data = VerticalGridData.new_from_metric_terms(metric_terms) + contravariant_data = ContravariantGridData.new_from_metric_terms(metric_terms) + angle_data = AngleGridData.new_from_metric_terms(metric_terms) + grid_data = GridData( + horizontal_data=horizontal_data, + vertical_data=vertical_data, + contravariant_data=contravariant_data, + angle_data=angle_data, + ) + + damping_coefficients = DampingCoefficients.new_from_metric_terms(metric_terms) + driver_grid_data = DriverGridData.new_from_metric_terms(metric_terms) + + return damping_coefficients, driver_grid_data, grid_data + + def _transform_horizontal_grid( metric_terms: MetricTerms, stretch_factor: float, diff --git a/driver/pace/driver/initialization.py b/driver/pace/driver/initialization.py index 3cf52376..04b08d3d 100644 --- a/driver/pace/driver/initialization.py +++ b/driver/pace/driver/initialization.py @@ -3,7 +3,7 @@ import os import pathlib from datetime import datetime -from typing import Callable, ClassVar, Type, TypeVar +from typing import Callable, ClassVar, List, Type, TypeVar import f90nml @@ -40,6 +40,7 @@ def get_driver_state( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ) -> DriverState: ... @@ -77,6 +78,7 @@ def get_driver_state( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ) -> DriverState: return self.config.get_driver_state( quantity_factory=quantity_factory, @@ -84,6 +86,7 @@ def get_driver_state( damping_coefficients=damping_coefficients, driver_grid_data=driver_grid_data, grid_data=grid_data, + schemes=schemes, ) @classmethod @@ -109,6 +112,7 @@ def get_driver_state( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ) -> DriverState: dycore_state = analytic_init.init_analytic_state( analytic_init_case=self.case, @@ -120,7 +124,7 @@ def get_driver_state( comm=communicator, ) physics_state = pace.physics.PhysicsState.init_zeros( - quantity_factory=quantity_factory, active_packages=["microphysics"] + quantity_factory=quantity_factory, schemes=schemes ) tendency_state = TendencyState.init_zeros( quantity_factory=quantity_factory, @@ -152,6 +156,7 @@ def get_driver_state( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ) -> DriverState: state = _restart_driver_state( self.path, @@ -161,6 +166,7 @@ def get_driver_state( damping_coefficients, driver_grid_data, grid_data, + schemes, ) return state @@ -201,6 +207,7 @@ def get_driver_state( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ) -> DriverState: state = _restart_driver_state( self.path, @@ -210,6 +217,7 @@ def get_driver_state( damping_coefficients, driver_grid_data, grid_data, + schemes, ) _update_fortran_restart_pe_peln(state) @@ -272,6 +280,7 @@ def get_driver_state( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ) -> DriverState: backend = quantity_factory.zeros( dims=[pace.util.X_DIM, pace.util.Y_DIM], units="unknown" @@ -280,7 +289,7 @@ def get_driver_state( dycore_state = self._initialize_dycore_state(communicator, backend) physics_state = pace.physics.PhysicsState.init_zeros( quantity_factory=quantity_factory, - active_packages=["microphysics"], + schemes=schemes, ) tendency_state = TendencyState.init_zeros(quantity_factory=quantity_factory) @@ -349,6 +358,7 @@ def get_driver_state( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ) -> DriverState: return DriverState( dycore_state=self.dycore_state, diff --git a/driver/pace/driver/state.py b/driver/pace/driver/state.py index 54241b1d..93c99b55 100644 --- a/driver/pace/driver/state.py +++ b/driver/pace/driver/state.py @@ -1,5 +1,6 @@ import dataclasses from dataclasses import fields +from typing import List import xarray as xr @@ -75,6 +76,7 @@ def load_state_from_restart( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ) -> "DriverState": comm = driver_config.comm_config.get_comm() communicator = pace.util.Communicator.from_layout( @@ -102,6 +104,7 @@ def load_state_from_restart( damping_coefficients=damping_coefficients, driver_grid_data=driver_grid_data, grid_data=grid_data, + schemes=schemes, ) return state @@ -176,6 +179,7 @@ def _restart_driver_state( damping_coefficients: pace.util.grid.DampingCoefficients, driver_grid_data: pace.util.grid.DriverGridData, grid_data: pace.util.grid.GridData, + schemes: List[pace.physics.PHYSICS_PACKAGES], ): fs = pace.util.get_fs(path) @@ -197,12 +201,11 @@ def _restart_driver_state( "restart_dycore_state", ) - active_packages = ["microphysics"] physics_state = pace.physics.PhysicsState.init_zeros( - quantity_factory=quantity_factory, active_packages=active_packages + quantity_factory=quantity_factory, schemes=schemes ) - physics_state.__post_init__(quantity_factory, active_packages) + physics_state.__post_init__(quantity_factory, schemes) tendency_state = TendencyState.init_zeros( quantity_factory=quantity_factory, ) diff --git a/driver/tests/mpi/test_restart.py b/driver/tests/mpi/test_restart.py index 5c2ccde9..92c6f72b 100644 --- a/driver/tests/mpi/test_restart.py +++ b/driver/tests/mpi/test_restart.py @@ -11,6 +11,7 @@ import pace.util from pace.driver import DriverConfig from pace.driver.state import DriverState +from pace.physics import PHYSICS_PACKAGES # The packages we import will import MPI, causing an MPI init, but we don't actually @@ -65,6 +66,7 @@ def test_restart(): damping_coefficients=damping_coefficients, driver_grid_data=driver_grid_data, grid_data=grid_data, + schemes=[PHYSICS_PACKAGES["GFS_microphysics"]], ) assert isinstance(driver_state, DriverState) diff --git a/dsl/pace/dsl/dace/utils.py b/dsl/pace/dsl/dace/utils.py index c10ad6ec..28d8abe5 100644 --- a/dsl/pace/dsl/dace/utils.py +++ b/dsl/pace/dsl/dace/utils.py @@ -31,12 +31,12 @@ def default_prefix(cls, config: DaceConfig) -> str: return f"[{config.get_orchestrate()}]" def __enter__(self): - pace_log.debug(self.prefix, f"{self.label}...") + pace_log.debug(f"{self.prefix} {self.label}...") self.start = time.time() def __exit__(self, _type, _val, _traceback): elapsed = time.time() - self.start - pace_log.debug(self.prefix, f"{self.label}...{elapsed}s.") + pace_log.debug(f"{self.prefix} {self.label}...{elapsed}s.") def _is_ref(sd: dace.sdfg.SDFG, aname: str): diff --git a/examples/Dockerfile b/examples/Dockerfile index b9822650..49fe934e 100644 --- a/examples/Dockerfile +++ b/examples/Dockerfile @@ -49,7 +49,7 @@ ENV CFLAGS="-I/usr/include -DACCEPT_USE_OF_DEPRECATED_PROJ_API_H=1" RUN python3 -m pip install \ numpy==1.21.2 \ netCDF4==1.5.7 \ - mpi4py==3.1.0 \ + mpi4py==3.1.1 \ matplotlib==3.5.2 \ ipyparallel==8.4.1 \ jupyterlab==3.4.4 \ diff --git a/examples/notebooks/generate_eta_file_netcdf.ipynb b/examples/notebooks/generate_eta_file_netcdf.ipynb new file mode 100644 index 00000000..bf9623c7 --- /dev/null +++ b/examples/notebooks/generate_eta_file_netcdf.ipynb @@ -0,0 +1,148 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "2c056479", + "metadata": { + "lines_to_next_cell": 0 + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8c96fbff", + "metadata": {}, + "outputs": [], + "source": [ + "import netCDF4 as nc\n", + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6827b1b5", + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"\n", + "This notebook uses the python netCDF4 module\n", + "to create an eta_file containg\n", + "ak and bk coefficients for km=79\n", + "\"\"\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "45d4a704", + "metadata": {}, + "outputs": [], + "source": [ + "#create a Dataset instance\n", + "coefficients = nc.Dataset(\"eta79.nc\", \"w\", format=\"NETCDF4\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b964a014", + "metadata": {}, + "outputs": [], + "source": [ + "#Set dimensionsion\n", + "km = coefficients.createDimension(\"km\", 80)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d51c395f", + "metadata": {}, + "outputs": [], + "source": [ + "#Create ak and bk variables\n", + "ak = coefficients.createVariable(\"ak\", np.float64, (\"km\"))\n", + "bk = coefficients.createVariable(\"bk\", np.float64, (\"km\"))\n", + "ak.units=\"\"\n", + "bk.units=\"\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6723352e", + "metadata": {}, + "outputs": [], + "source": [ + "#Assign and write out values of ak\n", + "ak[:] = np.array(\n", + " [ 3.000000e+02, 6.467159e+02, 1.045222e+03, 1.469188e+03, 1.897829e+03,\n", + " 2.325385e+03, 2.754396e+03, 3.191294e+03, 3.648332e+03, 4.135675e+03,\n", + " 4.668282e+03, 5.247940e+03, 5.876271e+03, 6.554716e+03, 7.284521e+03,\n", + " 8.066738e+03, 8.902188e+03, 9.791482e+03, 1.073499e+04, 1.162625e+04,\n", + " 1.237212e+04, 1.299041e+04, 1.349629e+04, 1.390277e+04, 1.422098e+04,\n", + " 1.446058e+04, 1.462993e+04, 1.473633e+04, 1.478617e+04, 1.478511e+04,\n", + " 1.473812e+04, 1.464966e+04, 1.452370e+04, 1.436382e+04, 1.417324e+04,\n", + " 1.395491e+04, 1.371148e+04, 1.344540e+04, 1.315890e+04, 1.285407e+04,\n", + " 1.253280e+04, 1.219685e+04, 1.184788e+04, 1.148739e+04, 1.111682e+04,\n", + " 1.073748e+04, 1.035062e+04, 9.957395e+03, 9.558875e+03, 9.156069e+03,\n", + " 8.749922e+03, 8.341315e+03, 7.931065e+03, 7.519942e+03, 7.108648e+03,\n", + " 6.698281e+03, 6.290007e+03, 5.884984e+03, 5.484372e+03, 5.089319e+03,\n", + " 4.700960e+03, 4.320421e+03, 3.948807e+03, 3.587201e+03, 3.236666e+03,\n", + " 2.898237e+03, 2.572912e+03, 2.261667e+03, 1.965424e+03, 1.685079e+03,\n", + " 1.421479e+03, 1.175419e+03, 9.476516e+02, 7.388688e+02, 5.497130e+02,\n", + " 3.807626e+02, 2.325417e+02, 1.054810e+02, -8.381903e-04, 0.000000e+00] )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "195c9ef5", + "metadata": {}, + "outputs": [], + "source": [ + "#Assign and write out values of bk \n", + "bk[:] = np.array(\n", + " [ 0., 0., 0., 0., 0., 0.,\n", + " 0., 0., 0., 0., 0., 0.,\n", + " 0., 0., 0., 0., 0., 0.,\n", + " 0., 0.00106595, 0.00412866, 0.00900663, 0.01554263, 0.02359921,\n", + " 0.03305481, 0.0438012, 0.05574095, 0.06878554, 0.08285347, 0.09786981,\n", + " 0.1137643, 0.130471, 0.1479275, 0.1660746, 0.1848558, 0.2042166,\n", + " 0.2241053, 0.2444716, 0.2652672, 0.286445, 0.3079604, 0.3297701,\n", + " 0.351832, 0.3741062, 0.3965532, 0.4191364, 0.4418194, 0.4645682,\n", + " 0.48735, 0.5101338, 0.5328897, 0.5555894, 0.5782067, 0.6007158,\n", + " 0.6230936, 0.6452944, 0.6672683, 0.6889648, 0.7103333, 0.7313231,\n", + " 0.7518838, 0.7719651, 0.7915173, 0.8104913, 0.828839, 0.846513,\n", + " 0.8634676, 0.8796583, 0.8950421, 0.9095779, 0.9232264, 0.9359506,\n", + " 0.9477157, 0.9584892, 0.9682413, 0.9769447, 0.9845753, 0.9911126,\n", + " 0.9965372, 1. ] )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c0f3bd9d", + "metadata": {}, + "outputs": [], + "source": [ + "#Close netcdf file\n", + "coefficients.close()" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "executable": "/usr/bin/env python3", + "main_language": "python", + "notebook_metadata_filter": "-all" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/notebooks/generate_eta_file_xarray.ipynb b/examples/notebooks/generate_eta_file_xarray.ipynb new file mode 100644 index 00000000..a47b09e6 --- /dev/null +++ b/examples/notebooks/generate_eta_file_xarray.ipynb @@ -0,0 +1,140 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "6dc5fe4c", + "metadata": { + "lines_to_next_cell": 0 + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "81be9a15", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import xarray as xr" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c74c6c07", + "metadata": {}, + "outputs": [], + "source": [ + "\"\"\"\n", + "This notebook uses the python xarray module\n", + "to create an eta_file containg\n", + "ak and bk coefficients for km=79\n", + "\"\"\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f72c5d5b", + "metadata": {}, + "outputs": [], + "source": [ + "#Assign ak data\n", + "ak=np.array(\n", + " [ 3.000000e+02, 6.467159e+02, 1.045222e+03, 1.469188e+03, 1.897829e+03,\n", + " 2.325385e+03, 2.754396e+03, 3.191294e+03, 3.648332e+03, 4.135675e+03,\n", + " 4.668282e+03, 5.247940e+03, 5.876271e+03, 6.554716e+03, 7.284521e+03,\n", + " 8.066738e+03, 8.902188e+03, 9.791482e+03, 1.073499e+04, 1.162625e+04,\n", + " 1.237212e+04, 1.299041e+04, 1.349629e+04, 1.390277e+04, 1.422098e+04,\n", + " 1.446058e+04, 1.462993e+04, 1.473633e+04, 1.478617e+04, 1.478511e+04,\n", + " 1.473812e+04, 1.464966e+04, 1.452370e+04, 1.436382e+04, 1.417324e+04,\n", + " 1.395491e+04, 1.371148e+04, 1.344540e+04, 1.315890e+04, 1.285407e+04,\n", + " 1.253280e+04, 1.219685e+04, 1.184788e+04, 1.148739e+04, 1.111682e+04,\n", + " 1.073748e+04, 1.035062e+04, 9.957395e+03, 9.558875e+03, 9.156069e+03,\n", + " 8.749922e+03, 8.341315e+03, 7.931065e+03, 7.519942e+03, 7.108648e+03,\n", + " 6.698281e+03, 6.290007e+03, 5.884984e+03, 5.484372e+03, 5.089319e+03,\n", + " 4.700960e+03, 4.320421e+03, 3.948807e+03, 3.587201e+03, 3.236666e+03,\n", + " 2.898237e+03, 2.572912e+03, 2.261667e+03, 1.965424e+03, 1.685079e+03,\n", + " 1.421479e+03, 1.175419e+03, 9.476516e+02, 7.388688e+02, 5.497130e+02,\n", + " 3.807626e+02, 2.325417e+02, 1.054810e+02, -8.381903e-04, 0.000000e+00] )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f5b85c7e", + "metadata": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "#Assign bk data\n", + "bk=np.array(\n", + " [ 0., 0., 0., 0., 0., 0.,\n", + " 0., 0., 0., 0., 0., 0.,\n", + " 0., 0., 0., 0., 0., 0.,\n", + " 0., 0.00106595, 0.00412866, 0.00900663, 0.01554263, 0.02359921,\n", + " 0.03305481, 0.0438012, 0.05574095, 0.06878554, 0.08285347, 0.09786981,\n", + " 0.1137643, 0.130471, 0.1479275, 0.1660746, 0.1848558, 0.2042166,\n", + " 0.2241053, 0.2444716, 0.2652672, 0.286445, 0.3079604, 0.3297701,\n", + " 0.351832, 0.3741062, 0.3965532, 0.4191364, 0.4418194, 0.4645682,\n", + " 0.48735, 0.5101338, 0.5328897, 0.5555894, 0.5782067, 0.6007158,\n", + " 0.6230936, 0.6452944, 0.6672683, 0.6889648, 0.7103333, 0.7313231,\n", + " 0.7518838, 0.7719651, 0.7915173, 0.8104913, 0.828839, 0.846513,\n", + " 0.8634676, 0.8796583, 0.8950421, 0.9095779, 0.9232264, 0.9359506,\n", + " 0.9477157, 0.9584892, 0.9682413, 0.9769447, 0.9845753, 0.9911126,\n", + " 0.9965372, 1. ] )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c5450f7f", + "metadata": {}, + "outputs": [], + "source": [ + "#Create a Dataset instance\n", + "coefficients = xr.Dataset(\n", + " { \"ak\": ([\"km1\"], ak),\n", + " \"bk\": ([\"km1\"], bk) \n", + " })" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a0e5487", + "metadata": {}, + "outputs": [], + "source": [ + "#Set attributes for each variable\n", + "coefficients[\"ak\"].attrs[\"units\"]=\"\"\n", + "coefficients[\"bk\"].attrs[\"units\"]=\"\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "612b0134", + "metadata": {}, + "outputs": [], + "source": [ + "#Write netcdf file\n", + "coefficients.to_netcdf(\"eta79.nc\")" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "executable": "/usr/bin/env python3", + "main_language": "python", + "notebook_metadata_filter": "-all" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/external/dace b/external/dace index b0cd25b9..22982afe 160000 --- a/external/dace +++ b/external/dace @@ -1 +1 @@ -Subproject commit b0cd25b9263a3c615ee2f3325167944628fbfde5 +Subproject commit 22982afe133bccd906d5eeee448092f5f065ff6a diff --git a/fv3core/pace/fv3core/dycore_state.py b/fv3core/pace/fv3core/dycore_state.py index 4901c799..0f73e728 100644 --- a/fv3core/pace/fv3core/dycore_state.py +++ b/fv3core/pace/fv3core/dycore_state.py @@ -1,11 +1,12 @@ -from dataclasses import dataclass, field, fields -from typing import Any, Mapping +from dataclasses import asdict, dataclass, field, fields +from typing import Any, Dict, Mapping, Union import xarray as xr import pace.dsl.gt4py_utils as gt_utils import pace.util from pace.dsl.typing import Float +from pace.util.quantity import Quantity @dataclass() @@ -450,6 +451,12 @@ def xr_dataset(self): def __getitem__(self, item): return getattr(self, item) + def as_dict(self, quantity_only=True) -> Dict[str, Union[Quantity, int]]: + if quantity_only: + return {k: v for k, v in asdict(self).items() if isinstance(v, Quantity)} + else: + return {k: v for k, v in asdict(self).items()} + TRACER_PROPERTIES = { "specific_humidity": { diff --git a/fv3core/pace/fv3core/initialization/analytic_init.py b/fv3core/pace/fv3core/initialization/analytic_init.py index e8f6b07e..544c6283 100644 --- a/fv3core/pace/fv3core/initialization/analytic_init.py +++ b/fv3core/pace/fv3core/initialization/analytic_init.py @@ -1,15 +1,11 @@ -from enum import Enum, EnumMeta +from enum import Enum import pace.util as fv3util from pace.fv3core.dycore_state import DycoreState +from pace.util import MetaEnumStr from pace.util.grid import GridData -class MetaEnumStr(EnumMeta): - def __contains__(cls, item): - return item in cls.__members__.keys() - - class Cases(Enum, metaclass=MetaEnumStr): baroclinic = "baroclinic" tropicalcyclone = "tropicalcyclone" diff --git a/fv3core/pace/fv3core/stencils/a2b_ord4.py b/fv3core/pace/fv3core/stencils/a2b_ord4.py index 95b6ab02..b4a74898 100644 --- a/fv3core/pace/fv3core/stencils/a2b_ord4.py +++ b/fv3core/pace/fv3core/stencils/a2b_ord4.py @@ -548,7 +548,11 @@ def __init__( replace: boolean, update qin to the B grid as well """ orchestrate(obj=self, config=stencil_factory.config.dace_config) - assert grid_type in [0, 4] + if grid_type != 0 and grid_type != 4: + raise RuntimeError( + "A-Grid to B-Grid 4th order (a2b_ord4):" + f" grid type {grid_type} is not implemented. 0 and 4 available." + ) self._idx: GridIndexing = stencil_factory.grid_indexing self._stencil_config = stencil_factory.config self.replace = replace @@ -720,7 +724,6 @@ def __call__(self, qin: FloatField, qout: FloatField): """ if self.grid_type < 3: - self._sw_corner_stencil( qin, qout, diff --git a/fv3core/pace/fv3core/stencils/d_sw.py b/fv3core/pace/fv3core/stencils/d_sw.py index e08af776..319054b6 100644 --- a/fv3core/pace/fv3core/stencils/d_sw.py +++ b/fv3core/pace/fv3core/stencils/d_sw.py @@ -764,19 +764,45 @@ def __init__( self.grid_indexing = stencil_factory.grid_indexing self._grid_type = config.grid_type - assert not config.inline_q, "inline_q not yet implemented" - assert ( - config.d_ext <= 0 - ), "untested d_ext > 0. need to call a2b_ord2, not yet implemented" - assert (column_namelist["damp_vt"].view[:] > dcon_threshold).all() + if config.inline_q: + raise NotImplementedError( + "D-Grid Shallow Water Lagrangian Dynamics (D_SW):" + " inline_q not implemented." + ) + if config.d_ext > 0: + raise NotImplementedError( + "D-Grid Shallow Water Lagrangian Dynamics (D_SW):" + " untested d_ext > 0. need to call a2b_ord2, not implemented." + ) # TODO: in theory, we should check if damp_vt > 1e-5 for each k-level and # only compute delnflux for k-levels where this is true - assert (column_namelist["damp_w"].view[:] > dcon_threshold).all() + all_damp_vt_above_dcon = ( + column_namelist["damp_vt"].view[:] > dcon_threshold + ).all() + if not all_damp_vt_above_dcon: + raise NotImplementedError( + "D-Grid Shallow Water Lagrangian Dynamics (D_SW):" + " damp_vt misconfiguration, some are above a d_con of" + f" {dcon_threshold}." + ) # TODO: in theory, we should check if damp_w > 1e-5 for each k-level and # only compute delnflux for k-levels where this is true + all_damp_w_above_dcon = ( + column_namelist["damp_w"].view[:] > dcon_threshold + ).all() + if not all_damp_w_above_dcon: + raise NotImplementedError( + "D-Grid Shallow Water Lagrangian Dynamics (D_SW):" + f" damp_w misconfiguration, some are above a d_con of {dcon_threshold}." + ) # only compute for k-levels where this is true self.hydrostatic = config.hydrostatic + if self.hydrostatic: + raise NotImplementedError( + "D-Grid Shallow Water Lagrangian Dynamics (D_SW):" + " Hydrostatic is not implemented" + ) def make_quantity(): return quantity_factory.zeros( diff --git a/fv3core/pace/fv3core/stencils/divergence_damping.py b/fv3core/pace/fv3core/stencils/divergence_damping.py index a3e5a32b..5c63175a 100644 --- a/fv3core/pace/fv3core/stencils/divergence_damping.py +++ b/fv3core/pace/fv3core/stencils/divergence_damping.py @@ -324,8 +324,8 @@ def __init__( config=stencil_factory.config.dace_config, ) self.grid_indexing = stencil_factory.grid_indexing - assert not nested, "nested not implemented" - # assert grid_type < 3, "Not implemented, grid_type>=3" + if nested: + raise NotImplementedError("Divergence Dampoing: nested not implemented.") # TODO: make dddmp a compile-time external, instead of runtime scalar self._dddmp = dddmp # TODO: make da_min_c a compile-time external, instead of runtime scalar diff --git a/fv3core/pace/fv3core/stencils/dyn_core.py b/fv3core/pace/fv3core/stencils/dyn_core.py index bef8f6f0..45173c80 100644 --- a/fv3core/pace/fv3core/stencils/dyn_core.py +++ b/fv3core/pace/fv3core/stencils/dyn_core.py @@ -429,9 +429,12 @@ def __init__( self.checkpointer = checkpointer grid_indexing = stencil_factory.grid_indexing self.config = config - assert config.d_ext == 0, "d_ext != 0 is not implemented" - assert config.beta == 0, "beta != 0 is not implemented" - assert not config.use_logp, "use_logp=True is not implemented" + if config.d_ext != 0: + raise RuntimeError("Acoustics (dyn_core): d_ext != 0 is not implemented") + if config.beta != 0: + raise RuntimeError("Acoustics (dyn_core): beta != 0 is not implemented") + if config.use_logp: + raise RuntimeError("Acoustics (dyn_core): use_logp=True is not implemented") self._da_min = damping_coefficients.da_min self.grid_data = grid_data self._ptop = grid_data.ptop diff --git a/fv3core/pace/fv3core/stencils/fv_dynamics.py b/fv3core/pace/fv3core/stencils/fv_dynamics.py index 8e82f549..d4ef8959 100644 --- a/fv3core/pace/fv3core/stencils/fv_dynamics.py +++ b/fv3core/pace/fv3core/stencils/fv_dynamics.py @@ -186,8 +186,17 @@ def __init__( nested = False stretched_grid = False grid_indexing = stencil_factory.grid_indexing - assert config.moist_phys, "fvsetup is only implemented for moist_phys=true" - assert config.nwat == 6, "Only nwat=6 has been implemented and tested" + if not config.moist_phys: + raise NotImplementedError( + "Dynamical core (fv_dynamics):" + " fvsetup is only implemented for moist_phys=true." + ) + if config.nwat != 6: + raise NotImplementedError( + "Dynamical core (fv_dynamics):" + f" nwat=={config.nwat} is not implemented." + " Only nwat=6 has been implemented." + ) self.comm_rank = comm.rank self.grid_data = grid_data self.grid_indexing = grid_indexing @@ -286,7 +295,10 @@ def __init__( self._cappa = self.acoustic_dynamics.cappa if not (not self.config.inline_q and constants.NQ != 0): - raise NotImplementedError("tracer_2d not implemented, turn on z_tracer") + raise NotImplementedError( + "Dynamical core (fv_dynamics):" + "tracer_2d not implemented. z_tracer available" + ) self._adjust_tracer_mixing_ratio = AdjustNegativeTracerMixingRatio( stencil_factory, quantity_factory=quantity_factory, @@ -466,16 +478,20 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool): ) if self._conserve_total_energy > 0: - raise NotImplementedError("compute total energy is not implemented") + raise NotImplementedError( + "Dynamical Core (fv_dynamics): compute total energy is not implemented" + ) if (not self.config.rf_fast) and self.config.tau != 0: raise NotImplementedError( - "Rayleigh_Super, called when rf_fast=False and tau !=0" + "Dynamical Core (fv_dynamics): Rayleigh_Super," + " called when rf_fast=False and tau !=0, is not implemented" ) if self.config.adiabatic and self.config.kord_tm > 0: raise NotImplementedError( - "unimplemented namelist options adiabatic with positive kord_tm" + "Dynamical Core (fv_dynamics): Adiabatic with positive kord_tm" + " is not implemented." ) else: if __debug__: diff --git a/fv3core/pace/fv3core/stencils/fv_subgridz.py b/fv3core/pace/fv3core/stencils/fv_subgridz.py index fded63fd..908ba296 100644 --- a/fv3core/pace/fv3core/stencils/fv_subgridz.py +++ b/fv3core/pace/fv3core/stencils/fv_subgridz.py @@ -99,7 +99,6 @@ def init( qsgs_tke: FloatField, qcld: FloatField, ): - with computation(PARALLEL), interval(...): t0 = ta u0 = ua @@ -783,12 +782,17 @@ def __init__( n_sponge: int, hydrostatic: bool, ): - assert not hydrostatic, "Hydrostatic not implemented for fv_subgridz" + if hydrostatic: + raise NotImplementedError( + "DryConvectiveAdjustment (fv_subgridz):" + " Hydrostatic is not implemented" + ) grid_indexing = stencil_factory.grid_indexing self._k_sponge = n_sponge - if self._k_sponge is not None: - if self._k_sponge < 3: - return + if self._k_sponge is not None and self._k_sponge < 3: + raise ValueError( + "DryConvectiveAdjustment (fv_subgridz): n_sponge < 3 is invalid." + ) else: self._k_sponge = grid_indexing.domain[2] if self._k_sponge < min(grid_indexing.domain[2], 24): diff --git a/fv3core/pace/fv3core/stencils/neg_adj3.py b/fv3core/pace/fv3core/stencils/neg_adj3.py index f2a6d3bb..b6aa77ee 100644 --- a/fv3core/pace/fv3core/stencils/neg_adj3.py +++ b/fv3core/pace/fv3core/stencils/neg_adj3.py @@ -356,7 +356,10 @@ def __init__( ) if hydrostatic: self._d0_vap = constants.CP_VAP - constants.C_LIQ - raise NotImplementedError("Unimplemented namelist hydrostatic=True") + raise NotImplementedError( + "Adjust Negative Tracer Mixing Ratio (neg_adj3):" + " Hydrostatic is not implemented" + ) else: self._d0_vap = constants.CV_VAP - constants.C_LIQ self._lv00 = constants.HLV - self._d0_vap * constants.TICE diff --git a/fv3core/pace/fv3core/stencils/remap_profile.py b/fv3core/pace/fv3core/stencils/remap_profile.py index fc3337b0..563e9f19 100644 --- a/fv3core/pace/fv3core/stencils/remap_profile.py +++ b/fv3core/pace/fv3core/stencils/remap_profile.py @@ -591,7 +591,11 @@ def __init__( config=stencil_factory.config.dace_config, ) - assert kord <= 10, f"kord {kord} not implemented." + if kord > 10: + raise NotImplementedError( + f"Remap Profile: kord {kord} not implemented. kord <= 10 available." + ) + self._kord = kord self._gam = quantity_factory.zeros( diff --git a/fv3core/pace/fv3core/stencils/xppm.py b/fv3core/pace/fv3core/stencils/xppm.py index 239e2d7f..afa9c2df 100644 --- a/fv3core/pace/fv3core/stencils/xppm.py +++ b/fv3core/pace/fv3core/stencils/xppm.py @@ -310,7 +310,11 @@ def __init__( # Arguments come from: # namelist.grid_type # grid.dxa - assert (grid_type < 3) or (grid_type == 4) + if grid_type == 3 or grid_type > 4: + raise NotImplementedError( + "X Piecewise Parabolic (xppm): " + f" grid type {grid_type} not implemented. <3 or 4 available." + ) self._dxa = dxa ax_offsets = stencil_factory.grid_indexing.axis_offsets(origin, domain) self._compute_flux_stencil = stencil_factory.from_origin_domain( diff --git a/fv3core/pace/fv3core/stencils/yppm.py b/fv3core/pace/fv3core/stencils/yppm.py index 69389e2b..d61b0ca3 100644 --- a/fv3core/pace/fv3core/stencils/yppm.py +++ b/fv3core/pace/fv3core/stencils/yppm.py @@ -9,6 +9,7 @@ region, ) +from pace.dsl.dace.orchestration import orchestrate from pace.dsl.stencil import StencilFactory from pace.dsl.typing import FloatField, FloatFieldIJ, Index3D from pace.fv3core.stencils import ppm @@ -295,7 +296,7 @@ def compute_y_flux( class YPiecewiseParabolic: """ - Fortran name is xppm + Fortran name is yppm """ def __init__( @@ -307,10 +308,15 @@ def __init__( origin: Index3D, domain: Index3D, ): + orchestrate(obj=self, config=stencil_factory.config.dace_config) # Arguments come from: # namelist.grid_type # grid.dya - assert (grid_type < 3) or (grid_type == 4) + if grid_type == 3 or grid_type > 4: + raise NotImplementedError( + "Y Piecewise Parabolic (xppm): " + f" grid type {grid_type} not implemented. <3 or 4 available." + ) self._dya = dya ax_offsets = stencil_factory.grid_indexing.axis_offsets(origin, domain) self._compute_flux_stencil = stencil_factory.from_origin_domain( diff --git a/fv3core/pace/fv3core/wrappers/geos_wrapper.py b/fv3core/pace/fv3core/wrappers/geos_wrapper.py index e1d1defe..9f8347f3 100644 --- a/fv3core/pace/fv3core/wrappers/geos_wrapper.py +++ b/fv3core/pace/fv3core/wrappers/geos_wrapper.py @@ -124,7 +124,9 @@ def __init__( # set up the metric terms and grid data metric_terms = pace.util.grid.MetricTerms( - quantity_factory=quantity_factory, communicator=self.communicator + quantity_factory=quantity_factory, + communicator=self.communicator, + eta_file=namelist["grid_config"]["config"]["eta_file"], ) grid_data = pace.util.grid.GridData.new_from_metric_terms(metric_terms) diff --git a/fv3core/tests/savepoint/translate/translate_fvtp2d.py b/fv3core/tests/savepoint/translate/translate_fvtp2d.py index a8315e0b..b807254d 100644 --- a/fv3core/tests/savepoint/translate/translate_fvtp2d.py +++ b/fv3core/tests/savepoint/translate/translate_fvtp2d.py @@ -1,6 +1,7 @@ import pace.dsl import pace.dsl.gt4py_utils as utils import pace.util +from pace.dsl.typing import Float from pace.fv3core.stencils.fvtp2d import FiniteVolumeTransport from pace.fv3core.testing import TranslateDycoreFortranData2Py @@ -51,11 +52,11 @@ def compute_from_storage(self, inputs): backend=self.stencil_factory.backend, ) nord_col = self.grid.quantity_factory.zeros( - dims=[pace.util.Z_DIM], units="unknown" + dims=[pace.util.Z_DIM], units="unknown", dtype=Float ) nord_col.data[:] = nord_col.np.asarray(inputs.pop("nord")) damp_c = self.grid.quantity_factory.zeros( - dims=[pace.util.Z_DIM], units="unknown" + dims=[pace.util.Z_DIM], units="unknown", dtype=Float ) damp_c.data[:] = damp_c.np.asarray(inputs.pop("damp_c")) for optional_arg in ["mass"]: diff --git a/fv3core/tests/savepoint/translate/translate_yppm.py b/fv3core/tests/savepoint/translate/translate_yppm.py index 4fc49c9c..175be9bb 100644 --- a/fv3core/tests/savepoint/translate/translate_yppm.py +++ b/fv3core/tests/savepoint/translate/translate_yppm.py @@ -1,6 +1,7 @@ import pace.dsl import pace.dsl.gt4py_utils as utils import pace.util +from pace.dsl.typing import Float from pace.fv3core.stencils import yppm from pace.fv3core.testing import TranslateDycoreFortranData2Py from pace.stencils.testing import TranslateGrid @@ -40,7 +41,9 @@ def process_inputs(self, inputs): self.ivars(inputs) self.make_storage_data_input_vars(inputs) inputs["flux"] = utils.make_storage_from_shape( - inputs["q"].shape, backend=self.stencil_factory.backend + inputs["q"].shape, + backend=self.stencil_factory.backend, + dtype=Float, ) def compute(self, inputs): diff --git a/physics/pace/physics/__init__.py b/physics/pace/physics/__init__.py index 6fdaf68d..8fa30674 100644 --- a/physics/pace/physics/__init__.py +++ b/physics/pace/physics/__init__.py @@ -1,4 +1,4 @@ -from ._config import PhysicsConfig +from ._config import PHYSICS_PACKAGES, PhysicsConfig from .physics_state import PhysicsState from .stencils.microphysics import Microphysics from .stencils.physics import Physics diff --git a/physics/pace/physics/_config.py b/physics/pace/physics/_config.py index 58d4d274..4ce3715a 100644 --- a/physics/pace/physics/_config.py +++ b/physics/pace/physics/_config.py @@ -1,13 +1,20 @@ import dataclasses -from typing import Optional, Tuple +from enum import Enum, unique +from typing import List, Optional, Tuple import f90nml -from pace.util import Namelist, NamelistDefaults +from pace.util import MetaEnumStr, Namelist, NamelistDefaults DEFAULT_INT = 0 DEFAULT_BOOL = False +DEFAULT_SCHEMES = ["GFS_microphysics"] + + +@unique +class PHYSICS_PACKAGES(Enum, metaclass=MetaEnumStr): + GFS_microphysics = "GFS_microphysics" @dataclasses.dataclass @@ -18,6 +25,7 @@ class PhysicsConfig: npy: int = DEFAULT_INT npz: int = DEFAULT_INT nwat: int = DEFAULT_INT + schemes: List = None do_qa: bool = DEFAULT_BOOL c_cracw: float = NamelistDefaults.c_cracw c_paut: float = NamelistDefaults.c_paut @@ -100,6 +108,14 @@ class PhysicsConfig: namelist_override: Optional[str] = None def __post_init__(self): + if self.schemes is None: + self.schemes = DEFAULT_SCHEMES + package_schemes = [] + for scheme in self.schemes: + if scheme not in PHYSICS_PACKAGES: + raise NotImplementedError(f"{scheme} physics scheme not implemented") + package_schemes.append(PHYSICS_PACKAGES[scheme]) + self.schemes = package_schemes if self.namelist_override is not None: try: f90_nml = f90nml.read(self.namelist_override) diff --git a/physics/pace/physics/physics_state.py b/physics/pace/physics/physics_state.py index d764ac04..83463998 100644 --- a/physics/pace/physics/physics_state.py +++ b/physics/pace/physics/physics_state.py @@ -8,6 +8,8 @@ from pace.dsl.typing import Float from pace.physics.stencils.microphysics import MicrophysicsState +from ._config import PHYSICS_PACKAGES + @dataclass() class PhysicsState: @@ -281,13 +283,15 @@ class PhysicsState: } ) quantity_factory: InitVar[pace.util.QuantityFactory] - active_packages: InitVar[List[str]] + schemes: InitVar[List[PHYSICS_PACKAGES]] def __post_init__( - self, quantity_factory: pace.util.QuantityFactory, active_packages: List[str] + self, + quantity_factory: pace.util.QuantityFactory, + schemes: List[PHYSICS_PACKAGES], ): # storage for tendency variables not in PhysicsState - if "microphysics" in active_packages: + if "GFS_microphysics" in [scheme.value for scheme in schemes]: tendency = quantity_factory.zeros( [pace.util.X_DIM, pace.util.Y_DIM, pace.util.Z_DIM], "unknown", @@ -317,7 +321,9 @@ def __post_init__( self.microphysics = None @classmethod - def init_zeros(cls, quantity_factory, active_packages: List[str]) -> "PhysicsState": + def init_zeros( + cls, quantity_factory, schemes: List[PHYSICS_PACKAGES] + ) -> "PhysicsState": initial_arrays = {} for _field in fields(cls): if "dims" in _field.metadata.keys(): @@ -329,7 +335,7 @@ def init_zeros(cls, quantity_factory, active_packages: List[str]) -> "PhysicsSta return cls( **initial_arrays, quantity_factory=quantity_factory, - active_packages=active_packages, + schemes=schemes, ) @classmethod @@ -338,7 +344,7 @@ def init_from_storages( storages: Mapping[str, Any], sizer: pace.util.GridSizer, quantity_factory: pace.util.QuantityFactory, - active_packages: List[str], + schemes: List[PHYSICS_PACKAGES], ) -> "PhysicsState": inputs: Dict[str, pace.util.Quantity] = {} for _field in fields(cls): @@ -352,15 +358,13 @@ def init_from_storages( extent=sizer.get_extent(dims), ) inputs[_field.name] = quantity - return cls( - **inputs, quantity_factory=quantity_factory, active_packages=active_packages - ) + return cls(**inputs, quantity_factory=quantity_factory, schemes=schemes) @property def xr_dataset(self): data_vars = {} for name, field_info in self.__dataclass_fields__.items(): - if name not in ["quantity_factory", "active_packages"]: + if name not in ["quantity_factory", "schemes"]: if issubclass(field_info.type, pace.util.Quantity): dims = [ f"{dim_name}_{name}" for dim_name in field_info.metadata["dims"] diff --git a/physics/pace/physics/stencils/microphysics.py b/physics/pace/physics/stencils/microphysics.py index 064bc1c7..78fbcfb6 100644 --- a/physics/pace/physics/stencils/microphysics.py +++ b/physics/pace/physics/stencils/microphysics.py @@ -14,7 +14,7 @@ import pace.physics.functions.microphysics_funcs as functions import pace.util import pace.util.constants as constants -from pace.dsl.dace.orchestration import orchestrate +from pace.dsl.dace.orchestration import dace_inhibitor, orchestrate 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 @@ -2227,6 +2227,7 @@ def setupm(self, dt_atmos: float): self._ces0 = constants.EPS * es0 self._set_timestep(dt_atmos) + @dace_inhibitor def _update_timestep_if_needed(self, timestep: float): if timestep != self._timestep: self._set_timestep(timestep=timestep) diff --git a/physics/pace/physics/stencils/physics.py b/physics/pace/physics/stencils/physics.py index a4dfdb2f..2de7d0a8 100644 --- a/physics/pace/physics/stencils/physics.py +++ b/physics/pace/physics/stencils/physics.py @@ -1,5 +1,3 @@ -from typing import List - import gt4py.cartesian.gtscript as gtscript from gt4py.cartesian.gtscript import ( BACKWARD, @@ -10,7 +8,6 @@ interval, log, ) -from typing_extensions import Literal import pace.util import pace.util.constants as constants @@ -24,10 +21,7 @@ from pace.util import X_DIM, Y_DIM, Z_DIM from pace.util.grid import GridData -from .._config import PhysicsConfig - - -PHYSICS_PACKAGES = Literal["microphysics"] +from .._config import PHYSICS_PACKAGES, PhysicsConfig def atmos_phys_driver_statein( @@ -208,8 +202,13 @@ def __init__( quantity_factory: pace.util.QuantityFactory, grid_data: GridData, namelist: PhysicsConfig, - active_packages: List[Literal[PHYSICS_PACKAGES]], ): + schemes = [scheme.value for scheme in namelist.schemes] + for scheme in schemes: + if scheme not in PHYSICS_PACKAGES: + raise NotImplementedError( + f"{scheme} is not an implemented physics parameterization" + ) orchestrate( obj=self, config=stencil_factory.config.dace_config, @@ -249,8 +248,8 @@ def make_quantity(): "pktop": self._pktop, }, ) - if "microphysics" in active_packages: - self._do_microphysics = True + if "GFS_microphysics" in schemes: + self._gfs_microphysics = True self._prepare_microphysics = stencil_factory.from_origin_domain( func=prepare_microphysics, origin=grid_indexing.origin_compute(), @@ -267,7 +266,7 @@ def make_quantity(): stencil_factory, quantity_factory, grid_data, namelist=namelist ) else: - self._do_microphysics = False + self._gfs_microphysics = False def _setup_statein(self): self._NQ = 8 # state.nq_tot - spec.namelist.dnats @@ -311,7 +310,7 @@ def __call__(self, physics_state: PhysicsState, timestep: float): physics_state.phii, physics_state.phil, ) - if self._do_microphysics: + if self._gfs_microphysics: self._prepare_microphysics( physics_state.dz, physics_state.phii, diff --git a/physics/tests/savepoint/translate/translate_driver.py b/physics/tests/savepoint/translate/translate_driver.py index 90da4883..6910df18 100644 --- a/physics/tests/savepoint/translate/translate_driver.py +++ b/physics/tests/savepoint/translate/translate_driver.py @@ -8,7 +8,7 @@ # but also, driver tests should not be in physics from pace.fv3core.testing.translate_fvdynamics import TranslateFVDynamics from pace.fv3core.testing.validation import enable_selective_validation -from pace.physics import PhysicsConfig, PhysicsState +from pace.physics import PHYSICS_PACKAGES, PhysicsConfig, PhysicsState from pace.util.namelist import Namelist @@ -43,7 +43,8 @@ def compute_parallel(self, inputs, communicator): sizer, backend=self.stencil_config.compilation_config.backend ) physics_state = PhysicsState.init_zeros( - quantity_factory=quantity_factory, active_packages=["microphysics"] + quantity_factory=quantity_factory, + schemes=[PHYSICS_PACKAGES["GFS_microphysics"]], ) tendency_state = TendencyState.init_zeros( quantity_factory=quantity_factory, diff --git a/physics/tests/savepoint/translate/translate_gfs_physics_driver.py b/physics/tests/savepoint/translate/translate_gfs_physics_driver.py index 2feb3df9..ab8f870b 100644 --- a/physics/tests/savepoint/translate/translate_gfs_physics_driver.py +++ b/physics/tests/savepoint/translate/translate_gfs_physics_driver.py @@ -2,6 +2,7 @@ import pace.dsl.gt4py_utils as utils import pace.util as util +from pace.physics import PHYSICS_PACKAGES from pace.physics.stencils.physics import Physics, PhysicsState from pace.stencils import update_atmos_state from pace.stencils.testing.translate_physics import TranslatePhysicsFortranData2Py @@ -129,18 +130,17 @@ def compute(self, inputs): quantity_factory = util.QuantityFactory.from_backend( sizer, self.stencil_factory.backend ) - active_packages = ["microphysics"] + schemes = [PHYSICS_PACKAGES["GFS_microphysics"]] physics_state = PhysicsState( **inputs, quantity_factory=quantity_factory, - active_packages=active_packages, + schemes=schemes, ) physics = Physics( self.stencil_factory, self.grid.quantity_factory, self.grid.grid_data, self.namelist, - active_packages=active_packages, ) # TODO, self.namelist doesn't have fv_sg_adj because it is PhysicsConfig # either move where GFSPhysicsDriver starts, or pass the full namelist or diff --git a/physics/tests/savepoint/translate/translate_microphysics.py b/physics/tests/savepoint/translate/translate_microphysics.py index d3259c1a..49ebc4dc 100644 --- a/physics/tests/savepoint/translate/translate_microphysics.py +++ b/physics/tests/savepoint/translate/translate_microphysics.py @@ -5,6 +5,7 @@ import pace.dsl.gt4py_utils as utils import pace.util from pace.dsl.typing import Float +from pace.physics import PHYSICS_PACKAGES from pace.physics.stencils.microphysics import Microphysics from pace.physics.stencils.physics import PhysicsState from pace.stencils.testing.translate_physics import TranslatePhysicsFortranData2Py @@ -88,7 +89,7 @@ def compute(self, inputs): inputs, sizer=sizer, quantity_factory=quantity_factory, - active_packages=["microphysics"], + schemes=[PHYSICS_PACKAGES["GFS_microphysics"]], ) microphysics = Microphysics( self.stencil_factory, quantity_factory, self.grid.grid_data, self.namelist diff --git a/stencils/pace/stencils/testing/README.md b/stencils/pace/stencils/testing/README.md index a38655d0..f8b0f04f 100644 --- a/stencils/pace/stencils/testing/README.md +++ b/stencils/pace/stencils/testing/README.md @@ -7,6 +7,7 @@ First, make sure you have followed the instruction in the top level [README](../ The unit and regression tests of pace require data generated from the Fortran reference implementation which has to be downloaded from a Google Cloud Platform storage bucket. Since the bucket is setup as "requester pays", you need a valid GCP account to download the test data. First, make sure you have configured the authentication with user credientials and configured Docker with the following commands: + ```shell gcloud auth login gcloud auth configure-docker @@ -74,3 +75,12 @@ DEV=y make savepoint_tests_mpi DEV=y make physics_savepoint_tests DEV=y make physics_savepoint_tests_mpi ``` + +## Test failure + +Test are running for each gridpoint of the domain, unless the Translate class for the test specifically restricts it. +Upon failure, the test will drop a `netCDF` faile in a `./.translate-errors` directory and named `translate-TestCase(-Rank).nc` containing input, computed output, reference and errors. + +## Environment variables + +- `PACE_TEST_N_THRESHOLD_SAMPLES`: Upon failure the system will try to pertub the output in an attempt to check for numerical instability. This means re-running the test for N samples. Default is `10`, `0` or less turns this feature off. diff --git a/stencils/pace/stencils/testing/grid.py b/stencils/pace/stencils/testing/grid.py index 4cf623b1..23d25882 100644 --- a/stencils/pace/stencils/testing/grid.py +++ b/stencils/pace/stencils/testing/grid.py @@ -6,6 +6,7 @@ import pace.util from pace.dsl import gt4py_utils as utils from pace.dsl.stencil import GridIndexing +from pace.dsl.typing import Float from pace.util.grid import ( AngleGridData, ContravariantGridData, @@ -504,7 +505,7 @@ def grid_data(self) -> "GridData": data = getattr(self, name) assert data is not None - quantity = self.quantity_factory.zeros(dims=dims, units=units) + quantity = self.quantity_factory.zeros(dims=dims, units=units, dtype=Float) if len(quantity.shape) == 3: quantity.data[:] = data[:, :, : quantity.shape[2]] elif len(quantity.shape) == 2: diff --git a/stencils/pace/stencils/testing/test_translate.py b/stencils/pace/stencils/testing/test_translate.py index 0d0141d5..796f30c1 100644 --- a/stencils/pace/stencils/testing/test_translate.py +++ b/stencils/pace/stencils/testing/test_translate.py @@ -19,7 +19,7 @@ # this only matters for manually-added print statements np.set_printoptions(threshold=4096) -OUTDIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") +OUTDIR = "./.translate-errors" GPU_MAX_ERR = 1e-10 GPU_NEAR_ZERO = 1e-15 @@ -171,21 +171,23 @@ def process_override(threshold_overrides, testobj, test_name, backend): ) -N_THRESHOLD_SAMPLES = 10 +N_THRESHOLD_SAMPLES = int(os.getenv("PACE_TEST_N_THRESHOLD_SAMPLES", 10)) def get_thresholds(testobj, input_data): - return _get_thresholds(testobj.compute, input_data) + _get_thresholds(testobj.compute, input_data) def get_thresholds_parallel(testobj, input_data, communicator): def compute(input): return testobj.compute_parallel(input, communicator) - return _get_thresholds(compute, input_data) + _get_thresholds(compute, input_data) -def _get_thresholds(compute_function, input_data): +def _get_thresholds(compute_function, input_data) -> None: + if N_THRESHOLD_SAMPLES <= 0: + return output_list = [] for _ in range(N_THRESHOLD_SAMPLES): input = copy.deepcopy(input_data) @@ -289,10 +291,14 @@ def test_sequential_savepoint( ref_data_out[varname] = [ref_data] if len(failing_names) > 0: get_thresholds(case.testobj, input_data=original_input_data) - out_filename = os.path.join(OUTDIR, f"{case.savepoint_name}.nc") + os.makedirs(OUTDIR, exist_ok=True) + out_filename = os.path.join(OUTDIR, f"translate-{case.savepoint_name}.nc") + input_data_on_host = {} + for key, _input in input_data.items(): + input_data_on_host[key] = gt_utils.asarray(_input) save_netcdf( case.testobj, - [input_data], + [input_data_on_host], [output], ref_data_out, failing_names, @@ -420,13 +426,17 @@ def test_parallel_savepoint( ) passing_names.append(failing_names.pop()) if len(failing_names) > 0: + os.makedirs(OUTDIR, exist_ok=True) out_filename = os.path.join( - OUTDIR, f"{case.savepoint_name}-{case.grid.rank}.nc" + OUTDIR, f"translate-{case.savepoint_name}-{case.grid.rank}.nc" ) try: + input_data_on_host = {} + for key, _input in input_data.items(): + input_data_on_host[key] = gt_utils.asarray(_input) save_netcdf( case.testobj, - [input_data], + [input_data_on_host], [output], ref_data, failing_names, diff --git a/tests/main/driver/test_analytic_init.py b/tests/main/driver/test_analytic_init.py index 03dda5bc..97222db4 100644 --- a/tests/main/driver/test_analytic_init.py +++ b/tests/main/driver/test_analytic_init.py @@ -7,8 +7,13 @@ import pace.driver +DIR = os.path.dirname(os.path.abspath(__file__)) + +# TODO: Location of test configurations will be changed after refactor, +# need to update after + TESTED_CONFIGS: List[str] = [ - "driver/examples/configs/analytic_test.yaml", + "../../../driver/examples/configs/analytic_test.yaml", ] @@ -20,7 +25,7 @@ ) def test_analytic_init_config(tested_configs: List[str]): for config_file in tested_configs: - with open(os.path.abspath(config_file), "r") as f: + with open(os.path.join(DIR, config_file), "r") as f: config = yaml.safe_load(f) driver_config = pace.driver.DriverConfig.from_dict(config) assert driver_config.initialization.type == "analytic" diff --git a/tests/main/driver/test_example_configs.py b/tests/main/driver/test_example_configs.py index 1fc5dec1..6cd18053 100644 --- a/tests/main/driver/test_example_configs.py +++ b/tests/main/driver/test_example_configs.py @@ -14,6 +14,7 @@ TESTED_CONFIGS: List[str] = [ "baroclinic_c12.yaml", "baroclinic_c12_dp.yaml", + "baroclinic_c12_explicit_physics.yaml", "baroclinic_c12_comm_read.yaml", "baroclinic_c12_comm_write.yaml", "baroclinic_c12_null_comm.yaml", @@ -28,6 +29,12 @@ "baroclinic_c12_orch_cpu.yaml", "tropical_read_restart_fortran.yml", "tropicalcyclone_c128.yaml", + "baroclinic_c384_cpu.yaml", + "baroclinic_c384_gpu.yaml", + "baroclinic_c3072_cpu.yaml", + "baroclinic_c3072_gpu.yaml", + "test_external_C12_1x1.yaml", + "test_external_C12_2x2.yaml", ] JENKINS_CONFIGS_DIR = os.path.join(dirname, "../../../.jenkins/driver_configs/") diff --git a/tests/main/driver/test_restart_fortran.py b/tests/main/driver/test_restart_fortran.py index 5dc70d0a..518ceed1 100644 --- a/tests/main/driver/test_restart_fortran.py +++ b/tests/main/driver/test_restart_fortran.py @@ -6,6 +6,7 @@ import pace.driver import pace.util from pace.driver.initialization import FortranRestartInit +from pace.physics import PHYSICS_PACKAGES from pace.util import ( CubedSphereCommunicator, CubedSpherePartitioner, @@ -47,7 +48,9 @@ def test_state_from_fortran_restart(): damping_coefficients, driver_grid_data, grid_data, - ) = pace.driver.GeneratedGridConfig(restart_path=restart_dir).get_grid( + ) = pace.driver.GeneratedGridConfig( + restart_path=restart_dir, eta_file=restart_dir + "/fv_core.res.nc" + ).get_grid( quantity_factory, null_communicator ) @@ -58,6 +61,7 @@ def test_state_from_fortran_restart(): damping_coefficients=damping_coefficients, driver_grid_data=driver_grid_data, grid_data=grid_data, + schemes=[PHYSICS_PACKAGES["GFS_microphysics"]], ) ds = xr.open_dataset(os.path.join(restart_dir, "fv_core.res.tile1.nc")) np.testing.assert_array_equal( diff --git a/tests/main/driver/test_restart_serial.py b/tests/main/driver/test_restart_serial.py index 3e62b863..4c4542f5 100644 --- a/tests/main/driver/test_restart_serial.py +++ b/tests/main/driver/test_restart_serial.py @@ -10,6 +10,7 @@ from pace.driver import CreatesComm, DriverConfig from pace.driver.driver import RestartConfig from pace.driver.initialization import AnalyticInit +from pace.physics import PHYSICS_PACKAGES from pace.util.null_comm import NullComm @@ -66,11 +67,14 @@ def test_restart_save_to_disk(): sizer=sizer, backend=backend ) + eta_file = driver_config.grid_config.config.eta_file ( damping_coefficients, driver_grid_data, grid_data, - ) = pace.driver.GeneratedGridConfig().get_grid(quantity_factory, communicator) + ) = pace.driver.GeneratedGridConfig(eta_file=eta_file).get_grid( + quantity_factory, communicator + ) init = AnalyticInit() driver_state = init.get_driver_state( quantity_factory=quantity_factory, @@ -78,6 +82,7 @@ def test_restart_save_to_disk(): damping_coefficients=damping_coefficients, driver_grid_data=driver_grid_data, grid_data=grid_data, + schemes=[PHYSICS_PACKAGES["GFS_microphysics"]], ) time = datetime(2016, 1, 1, 0, 0, 0) diff --git a/tests/main/fv3core/test_cartesian_grid.py b/tests/main/fv3core/test_cartesian_grid.py index db2ebe2d..986f4eaf 100644 --- a/tests/main/fv3core/test_cartesian_grid.py +++ b/tests/main/fv3core/test_cartesian_grid.py @@ -8,7 +8,7 @@ @pytest.mark.parametrize("npx", [8]) @pytest.mark.parametrize("npy", [8]) -@pytest.mark.parametrize("npz", [1]) +@pytest.mark.parametrize("npz", [79]) @pytest.mark.parametrize("dx_const", [1e2, 1e3]) @pytest.mark.parametrize("dy_const", [2e2, 3e3]) @pytest.mark.parametrize("deglat", [0.0, 15.0]) @@ -35,6 +35,7 @@ def test_cartesian_grid_generation( dx_const=dx_const, dy_const=dy_const, deglat=deglat, + eta_file="tests/main/input/eta79.nc", ) assert np.all(grid_generator.lat_agrid.data == deglat * PI / 180.0) assert np.all(grid_generator.lon_agrid.data == 0.0) diff --git a/tests/main/fv3core/test_dycore_call.py b/tests/main/fv3core/test_dycore_call.py index ac79fdb5..4fd7a471 100644 --- a/tests/main/fv3core/test_dycore_call.py +++ b/tests/main/fv3core/test_dycore_call.py @@ -19,9 +19,9 @@ DIR = os.path.abspath(os.path.dirname(__file__)) -def setup_dycore() -> Tuple[ - fv3core.DynamicalCore, fv3core.DycoreState, pace.util.Timer -]: +def setup_dycore() -> ( + Tuple[fv3core.DynamicalCore, fv3core.DycoreState, pace.util.Timer] +): backend = "numpy" config = fv3core.DynamicalCoreConfig( layout=(1, 1), @@ -97,9 +97,11 @@ def setup_dycore() -> Tuple[ quantity_factory = pace.util.QuantityFactory.from_backend( sizer=sizer, backend=backend ) + eta_file = "tests/main/input/eta79.nc" metric_terms = MetricTerms( quantity_factory=quantity_factory, communicator=communicator, + eta_file=eta_file, ) grid_data = GridData.new_from_metric_terms(metric_terms) diff --git a/tests/main/fv3core/test_init_from_geos.py b/tests/main/fv3core/test_init_from_geos.py index 16efe9e9..7e03caf0 100644 --- a/tests/main/fv3core/test_init_from_geos.py +++ b/tests/main/fv3core/test_init_from_geos.py @@ -74,6 +74,9 @@ def test_geos_wrapper(): "fv_sg_adj": 0, "n_sponge": 48, }, + "grid_config": { + "config": {"eta_file": "tests/main/input/eta91.nc"}, + }, } namelist = f90nml.namelist.Namelist(namelist_dict) diff --git a/tests/main/grid/test_eta.py b/tests/main/grid/test_eta.py new file mode 100755 index 00000000..b34e52aa --- /dev/null +++ b/tests/main/grid/test_eta.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python3 + +import os + +import numpy as np +import pytest +import xarray as xr +import yaml + +import pace.driver + + +""" +This test checks to ensure that ak and bk +values are read-in and stored properly. +In addition, this test checks to ensure that +the function set_hybrid_pressure_coefficients +fail as expected if the computed eta values +vary non-mononitically and if the eta_file +is not provided. +""" + + +def set_answers(eta_file): + + """ + Read in the expected values of ak and bk + arrays from the input eta NetCDF files. + """ + + data = xr.open_dataset(eta_file) + return data["ak"].values, data["bk"].values + + +@pytest.mark.parametrize("km", [79, 91]) +def test_set_hybrid_pressure_coefficients_correct(km): + + """This test checks to see that the ak and bk arrays + are read-in correctly and are stored as + expected. Both values of km=79 and km=91 are + tested and both tests are expected to pass + with the stored ak and bk values agreeing with the + values read-in directly from the NetCDF file. + """ + + dirname = os.path.dirname(os.path.abspath(__file__)) + config_file = os.path.join( + dirname, "../../../driver/examples/configs/baroclinic_c12.yaml" + ) + + with open(config_file, "r") as f: + yaml_config = yaml.safe_load(f) + + yaml_config["nz"] = km + yaml_config["grid_config"]["config"]["eta_file"] = f"tests/main/input/eta{km}.nc" + + driver_config = pace.driver.DriverConfig.from_dict(yaml_config) + driver_config.comm_config = pace.driver.NullCommConfig(rank=0, total_ranks=6) + driver = pace.driver.Driver(config=driver_config) + + p_results = driver.state.grid_data.p.data + ak_results = driver.state.grid_data.ak.data + bk_results = driver.state.grid_data.bk.data + ak_answers, bk_answers = set_answers(f"tests/main/input/eta{km}.nc") + + if ak_answers.size != ak_results.size: + raise ValueError("Unexpected size of bk") + if bk_answers.size != bk_results.size: + raise ValueError("Unexpected size of ak") + + if not np.array_equal(ak_answers, ak_results): + raise ValueError("Unexpected value of ak") + if not np.array_equal(bk_answers, bk_results): + raise ValueError("Unexpected value of bk") + + driver.safety_checker.clear_all_checks() + + +@pytest.mark.parametrize( + "cfile", + [ + "file_is_not_here", + "tests/main/grid/input/eta_not_mono.nc", + ], +) +@pytest.mark.xfail +def test_set_hybrid_pressure_coefficients_fail(cfile): + + """This test checks to see that the program + fails when (1) the eta_file is not specified in the yaml + configuration file; and (2), the computed eta values + increase non-monotonically. For the latter test, the eta_file + is specified in test_config_not_mono.yaml file and + the ak and bk values in the eta_file have been changed nonsensically + to result in erronenous eta values. + """ + + dirname = os.path.dirname(os.path.abspath(__file__)) + config_file = os.path.join( + dirname, "../../../driver/examples/configs/baroclinic_c12.yaml" + ) + + with open(config_file, "r") as f: + yaml_config = yaml.safe_load(f) + + yaml_config["grid_config"]["config"]["eta_file"] = cfile + + driver_config = pace.driver.DriverConfig.from_dict(yaml_config) + driver_config.comm_config = pace.driver.NullCommConfig(rank=0, total_ranks=6) + driver = pace.driver.Driver(config=driver_config) diff --git a/tests/main/physics/test_integration.py b/tests/main/physics/test_integration.py index a55d9f98..84cd6d31 100644 --- a/tests/main/physics/test_integration.py +++ b/tests/main/physics/test_integration.py @@ -64,6 +64,7 @@ def setup_physics(): metric_terms = pace.util.grid.MetricTerms( quantity_factory=quantity_factory, communicator=communicator, + eta_file="tests/main/input/eta79.nc", ) grid_data = pace.util.grid.GridData.new_from_metric_terms(metric_terms) physics = pace.physics.Physics( @@ -71,10 +72,9 @@ def setup_physics(): quantity_factory, grid_data, physics_config, - active_packages=["microphysics"], ) physics_state = pace.physics.PhysicsState.init_zeros( - quantity_factory, active_packages=["microphysics"] + quantity_factory, schemes=[pace.physics.PHYSICS_PACKAGES["GFS_microphysics"]] ) random = np.random.RandomState(0) for field in fields(pace.physics.PhysicsState): diff --git a/tests/main/test_grid_init.py b/tests/main/test_grid_init.py index 4c0e5e2b..55e00e81 100644 --- a/tests/main/test_grid_init.py +++ b/tests/main/test_grid_init.py @@ -32,18 +32,21 @@ def test_grid_init_not_decomposition_dependent(rank: int): mpi_54rank test suite. It tests only variables that are not dependent on halo updates for their values in the compute domain. """ - nx_tile, ny_tile, nz = 48, 48, 5 + nx_tile, ny_tile, nz = 48, 48, 79 + eta_file = "tests/main/input/eta79.nc" metric_terms_1by1 = MetricTerms( quantity_factory=get_quantity_factory( layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz ), communicator=get_cube_comm(rank=0, layout=(1, 1)), + eta_file=eta_file, ) metric_terms_3by3 = MetricTerms( quantity_factory=get_quantity_factory( layout=(3, 3), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz ), communicator=get_cube_comm(rank=rank, layout=(3, 3)), + eta_file=eta_file, ) partitioner = pace.util.TilePartitioner(layout=(3, 3)) assert allclose(metric_terms_1by1.grid, metric_terms_3by3.grid, partitioner, rank) diff --git a/tests/mpi_54rank/test_ext_grid/test_external_grid.py b/tests/mpi_54rank/test_ext_grid/test_external_grid.py new file mode 100644 index 00000000..03f50e65 --- /dev/null +++ b/tests/mpi_54rank/test_ext_grid/test_external_grid.py @@ -0,0 +1,200 @@ +import math +import os + +import numpy as np +import pytest +import xarray as xr +import yaml + +import pace.util +from pace.driver import Driver, DriverConfig +from pace.util.constants import PI, RADIUS +from pace.util.mpi import MPIComm + + +DIR = os.path.dirname(os.path.abspath(__file__)) +TEST_CONFIGS_DIR = os.path.join(DIR, "../../../driver/examples/configs/") +TEST_DATA_DIR = os.path.join(DIR, "../../../test_input/") + +TEST_CONFIG_FILE_RANKS = [ + ("test_external_C12_1x1.yaml", 6), + ("test_external_C12_2x2.yaml", 24), +] + +TILE_FILE_BASE_NAME = "C12.tile" + + +def get_cube_comm(layout, comm: MPIComm): + return pace.util.CubedSphereCommunicator( + comm=comm, + partitioner=pace.util.CubedSpherePartitioner( + pace.util.TilePartitioner(layout=layout) + ), + ) + + +def get_tile_num(comm: MPIComm): + return pace.util.get_tile_number(comm.rank, comm.partitioner.total_ranks) + + +# TODO: Location of test configurations and data will be changed +# after refactor, need to update after + + +@pytest.mark.parametrize("config_file_path, ranks", TEST_CONFIG_FILE_RANKS) +def test_extgrid_equals_generated(config_file_path: str, ranks: int): + """ + Test for matching values between what exists in external grid + file, and what is generated by Pace when using an "external" + grid configuration. The "external" grid configuration takes + d-grid values as input and Pace then calculates the remaining + metric terms from this data. + + External grid files are expected to be of Netcdf type + (for reference see method of generating grid data in + FRE-NCtools) and contain values for the longitudinal, + latitude, distance between longitude and latitude points, + and cell areas. + + On a run, the user should specify the correct number of ranks, + location of configuration yaml, and base name for the tile files + containing the grid data. + + ON PASS: Input and generated data will match, with zero, or + insignificant relative difference. Each rank will print a + list of maximum relative differences for each metric term + tested. + + ON FAIL: There will be a mismatch present between input and + generated grid data values. Mismatch can be determined from the + printed list of errors and the list of maximum relative differences + """ + + comm = MPIComm() + size = comm.Get_size() + if size != ranks: + pytest.skip("Number of ranks does not match amount needed for layout") + side = int(math.sqrt(size // 6)) + cube_comm = get_cube_comm(layout=(side, side), comm=comm) + + with open( + os.path.join(TEST_CONFIGS_DIR, config_file_path), + "r", + ) as ext_f: + ext_config = yaml.safe_load(ext_f) + ext_driver_config = DriverConfig.from_dict(ext_config) + + ext_driver = Driver(ext_driver_config) + + tile_num = get_tile_num(cube_comm) + tile_file = TILE_FILE_BASE_NAME + str(tile_num) + ".nc" + ds = xr.open_dataset(os.path.join(TEST_DATA_DIR, tile_file)) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + area = ds.area.values + nx = ds.nx.values.size + ny = ds.ny.values.size + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npy, npx), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], + global_extent=(npy, nx), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(ny, npx), + overlap=True, + ) + + subtile_slice_area = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_DIM], + global_extent=(ny, nx), + overlap=True, + ) + + lon_rad = lon * (PI / 180) + lat_rad = lat * (PI / 180) + + errors = [] + diffs = [] + + if not np.isclose( + ext_driver.state.grid_data.lon.view[:, :], lon_rad[subtile_slice_grid] + ).all(): + errors.append("Lon data mismatch") + + diff_lon = np.amax( + ext_driver.state.grid_data.lon.view[:, :] - lon_rad[subtile_slice_grid] + ) / np.amax(lon_rad[subtile_slice_grid]) + diffs.append(f"Lon maximum relative error = {diff_lon}") + + if not np.isclose( + ext_driver.state.grid_data.lat.view[:, :], lat_rad[subtile_slice_grid] + ).all(): + errors.append("Lat data mismatch") + + diff_lat = np.amax( + ext_driver.state.grid_data.lat.view[:, :] - lat_rad[subtile_slice_grid] + ) / np.amax(lat_rad[subtile_slice_grid]) + diffs.append(f"Lat maximum relative error = {diff_lat}") + + if not np.isclose( + ext_driver.state.grid_data.dy.view[:, :], dx[subtile_slice_dx] + ).all(): + errors.append("dx data mismatch") + + diff_dx = np.amax( + ext_driver.state.grid_data.dy.view[:, :] - dx[subtile_slice_dx] + ) / np.amax(dx[subtile_slice_dx]) + diffs.append(f"dx maximum relative error = {diff_dx}") + + if not np.isclose( + ext_driver.state.grid_data.dx.view[:, :], dy[subtile_slice_dy] + ).all(): + errors.append("dy data mismatch") + + diff_dy = np.amax( + ext_driver.state.grid_data.dx.view[:, :] - dy[subtile_slice_dy] + ) / np.amax(dy[subtile_slice_dy]) + diffs.append(f"dy maximum relative error = {diff_dy}") + + if not np.isclose( + ext_driver.state.grid_data.area.view[:, :], area[subtile_slice_area] + ).all(): + errors.append("area data mismatch") + + diff_area = np.amax( + ext_driver.state.grid_data.area.view[:, :] - area[subtile_slice_area] + ) / np.amax(area[subtile_slice_area]) + diffs.append(f"Area maximum relative error = {diff_area}") + + print(diffs) + + assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) + + surface_area_true = 4 * PI * (RADIUS ** 2) + + mpicomm = MPIComm() + + tmp = [math.fsum(row) for row in ext_driver.state.grid_data.area.view[:, :]] + rank_area_sum = math.fsum(tmp) + tile_area = mpicomm._comm.gather(rank_area_sum, root=0) + + if mpicomm.Get_rank() == 0: + total_area = math.fsum(tile_area) + assert np.isclose(total_area, surface_area_true) diff --git a/tests/mpi_54rank/test_grid_init.py b/tests/mpi_54rank/test_grid_init.py index 9404a3e0..855c539e 100644 --- a/tests/mpi_54rank/test_grid_init.py +++ b/tests/mpi_54rank/test_grid_init.py @@ -4,7 +4,9 @@ import pace.fv3core import pace.util -from pace.fv3core.initialization import init_baroclinic_state +from pace.fv3core.initialization.test_cases.initialize_baroclinic import ( + init_baroclinic_state, +) from pace.util.grid import MetricTerms from pace.util.mpi import MPIComm from pace.util.quantity import Quantity diff --git a/util/HISTORY.md b/util/HISTORY.md index 62ada765..0d54eca5 100644 --- a/util/HISTORY.md +++ b/util/HISTORY.md @@ -4,6 +4,7 @@ History latest ------ +- Added `MetaEnumStr` to utils to make enums more functional - Added `fill_for_translate_test` to MetricTerms to fill fields with NaNs only when required for testing - Added `init_cartesian` method to MetricTerms to handle grid generation for orthogonal grids - Added `from_layout` and `size` methods to TileCommunicator and Communicator @@ -13,6 +14,8 @@ latest - Added `dx_const`, `dy_const`, and `deglat` to grid generation code for doubly-periodic grids - Added f32 support to halo exchange data transformation - Use one single logger, from logging.py +- Removed hard-coded values of `ak` and `bk` arrays and added in the feature to read in `ak` and `bk` values + from a NetCDF file to compute the `eta` and `eta_v` values. v0.10.0 ------- diff --git a/util/pace/util/__init__.py b/util/pace/util/__init__.py index 8137ab77..deb3b081 100644 --- a/util/pace/util/__init__.py +++ b/util/pace/util/__init__.py @@ -70,6 +70,7 @@ from .quantity import Quantity, QuantityMetadata from .time import FMS_TO_CFTIME_TYPE, datetime64_to_datetime from .units import UnitsError, ensure_equal_units, units_are_equal +from .utils import MetaEnumStr __version__ = "0.10.0" diff --git a/util/pace/util/grid/eta.py b/util/pace/util/grid/eta.py index dc37aaa2..52fe3552 100644 --- a/util/pace/util/grid/eta.py +++ b/util/pace/util/grid/eta.py @@ -1,6 +1,8 @@ +import os from dataclasses import dataclass import numpy as np +import xarray as xr @dataclass @@ -21,7 +23,9 @@ class HybridPressureCoefficients: bk: np.ndarray -def set_hybrid_pressure_coefficients(km: int) -> HybridPressureCoefficients: +def set_hybrid_pressure_coefficients( + km: int, eta_file: str +) -> HybridPressureCoefficients: """ Sets the coefficients describing the hybrid pressure coordinates. @@ -35,820 +39,30 @@ def set_hybrid_pressure_coefficients(km: int) -> HybridPressureCoefficients: Returns: a HybridPressureCoefficients dataclass """ - if km == 79: - ak = np.array( - [ - 300, - 646.7159, - 1045.222, - 1469.188, - 1897.829, - 2325.385, - 2754.396, - 3191.294, - 3648.332, - 4135.675, - 4668.282, - 5247.94, - 5876.271, - 6554.716, - 7284.521, - 8066.738, - 8902.188, - 9791.482, - 10734.99, - 11626.25, - 12372.12, - 12990.41, - 13496.29, - 13902.77, - 14220.98, - 14460.58, - 14629.93, - 14736.33, - 14786.17, - 14785.11, - 14738.12, - 14649.66, - 14523.7, - 14363.82, - 14173.24, - 13954.91, - 13711.48, - 13445.4, - 13158.9, - 12854.07, - 12532.8, - 12196.85, - 11847.88, - 11487.39, - 11116.82, - 10737.48, - 10350.62, - 9957.395, - 9558.875, - 9156.069, - 8749.922, - 8341.315, - 7931.065, - 7519.942, - 7108.648, - 6698.281, - 6290.007, - 5884.984, - 5484.372, - 5089.319, - 4700.96, - 4320.421, - 3948.807, - 3587.201, - 3236.666, - 2898.237, - 2572.912, - 2261.667, - 1965.424, - 1685.079, - 1421.479, - 1175.419, - 947.6516, - 738.8688, - 549.713, - 380.7626, - 232.5417, - 105.481, - -0.0008381903, - 0, - ] - ) - bk = np.array( - [ - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0, - 0.001065947, - 0.004128662, - 0.009006631, - 0.01554263, - 0.02359921, - 0.03305481, - 0.0438012, - 0.05574095, - 0.06878554, - 0.08285347, - 0.09786981, - 0.1137643, - 0.130471, - 0.1479275, - 0.1660746, - 0.1848558, - 0.2042166, - 0.2241053, - 0.2444716, - 0.2652672, - 0.286445, - 0.3079604, - 0.3297701, - 0.351832, - 0.3741062, - 0.3965532, - 0.4191364, - 0.4418194, - 0.4645682, - 0.48735, - 0.5101338, - 0.5328897, - 0.5555894, - 0.5782067, - 0.6007158, - 0.6230936, - 0.6452944, - 0.6672683, - 0.6889648, - 0.7103333, - 0.7313231, - 0.7518838, - 0.7719651, - 0.7915173, - 0.8104913, - 0.828839, - 0.846513, - 0.8634676, - 0.8796583, - 0.8950421, - 0.9095779, - 0.9232264, - 0.9359506, - 0.9477157, - 0.9584892, - 0.9682413, - 0.9769447, - 0.9845753, - 0.9911126, - 0.9965372, - 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, - ] - ) + if eta_file == "None": + raise ValueError("eta file not specified") + if not os.path.isfile(eta_file): + raise ValueError("file " + eta_file + " does not exist") - 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, - ] - ) + # read file into ak, bk arrays + data = xr.open_dataset(eta_file) + ak = data["ak"].values + bk = data["bk"].values - 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, - ] - ) + # check size of ak and bk array is km+1 + if ak.size - 1 != km: + raise ValueError(f"size of ak array is not equal to km={km}") + if bk.size - 1 != km: + raise ValueError(f"size of bk array is not equal to km={km}") - 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, - ] - ) + # check that the eta values computed from ak and bk are monotonically increasing + eta, etav = check_eta(ak, bk) - elif km == 137: - ak = np.array( - [ - 1.00000000, - 1.82500005, - 3.00000000, - 4.63000011, - 6.82797718, - 9.74696636, - 13.6054239, - 18.6089306, - 24.9857178, - 32.9857101, - 42.8792419, - 54.9554634, - 69.5205765, - 86.8958817, - 107.415741, - 131.425507, - 159.279404, - 191.338562, - 227.968948, - 269.539581, - 316.420746, - 368.982361, - 427.592499, - 492.616028, - 564.413452, - 643.339905, - 729.744141, - 823.967834, - 926.344910, - 1037.20117, - 1156.85364, - 1285.61035, - 1423.77014, - 1571.62292, - 1729.44897, - 1897.51929, - 2076.09595, - 2265.43164, - 2465.77051, - 2677.34814, - 2900.39136, - 3135.11938, - 3381.74365, - 3640.46826, - 3911.49048, - 4194.93066, - 4490.81738, - 4799.14941, - 5119.89502, - 5452.99072, - 5798.34473, - 6156.07422, - 6526.94678, - 6911.87061, - 7311.86914, - 7727.41211, - 8159.35400, - 8608.52539, - 9076.40039, - 9562.68262, - 10065.9785, - 10584.6318, - 11116.6621, - 11660.0674, - 12211.5479, - 12766.8730, - 13324.6689, - 13881.3311, - 14432.1396, - 14975.6152, - 15508.2568, - 16026.1152, - 16527.3223, - 17008.7891, - 17467.6133, - 17901.6211, - 18308.4336, - 18685.7188, - 19031.2891, - 19343.5117, - 19620.0430, - 19859.3906, - 20059.9316, - 20219.6641, - 20337.8633, - 20412.3086, - 20442.0781, - 20425.7188, - 20361.8164, - 20249.5117, - 20087.0859, - 19874.0254, - 19608.5723, - 19290.2266, - 18917.4609, - 18489.7070, - 18006.9258, - 17471.8398, - 16888.6875, - 16262.0469, - 15596.6953, - 14898.4531, - 14173.3242, - 13427.7695, - 12668.2578, - 11901.3398, - 11133.3047, - 10370.1758, - 9617.51562, - 8880.45312, - 8163.37500, - 7470.34375, - 6804.42188, - 6168.53125, - 5564.38281, - 4993.79688, - 4457.37500, - 3955.96094, - 3489.23438, - 3057.26562, - 2659.14062, - 2294.24219, - 1961.50000, - 1659.47656, - 1387.54688, - 1143.25000, - 926.507812, - 734.992188, - 568.062500, - 424.414062, - 302.476562, - 202.484375, - 122.101562, - 62.7812500, - 22.8359375, - 3.75781298, - 0.00000000, - 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, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 0.00000000, - 7.00000010e-06, - 2.40000008e-05, - 5.90000018e-05, - 1.12000002e-04, - 1.99000002e-04, - 3.39999999e-04, - 5.61999972e-04, - 8.90000025e-04, - 1.35300006e-03, - 1.99200003e-03, - 2.85700010e-03, - 3.97100020e-03, - 5.37799997e-03, - 7.13300006e-03, - 9.26099997e-03, - 1.18060000e-02, - 1.48160001e-02, - 1.83179993e-02, - 2.23549996e-02, - 2.69639995e-02, - 3.21759991e-02, - 3.80260013e-02, - 4.45480011e-02, - 5.17730005e-02, - 5.97280003e-02, - 6.84479997e-02, - 7.79580027e-02, - 8.82859975e-02, - 9.94620025e-02, - 0.111505002, - 0.124448001, - 0.138312995, - 0.153125003, - 0.168909997, - 0.185689002, - 0.203491002, - 0.222332999, - 0.242244005, - 0.263242006, - 0.285353988, - 0.308598012, - 0.332938999, - 0.358253986, - 0.384362996, - 0.411125004, - 0.438391000, - 0.466003001, - 0.493800014, - 0.521619022, - 0.549301028, - 0.576691985, - 0.603648007, - 0.630035996, - 0.655736029, - 0.680643022, - 0.704668999, - 0.727738976, - 0.749796987, - 0.770798028, - 0.790717006, - 0.809535980, - 0.827256024, - 0.843881011, - 0.859431982, - 0.873929024, - 0.887408018, - 0.899900019, - 0.911448002, - 0.922096014, - 0.931881011, - 0.940859973, - 0.949064016, - 0.956550002, - 0.963352025, - 0.969512999, - 0.975077987, - 0.980072021, - 0.984542012, - 0.988499999, - 0.991984010, - 0.995002985, - 0.997630000, - 1.00000000, - ] - ) - - else: - raise NotImplementedError( - "Only grids with 72, 79, 91 or 137 vertical levels" - "have been implemented so far" - ) + if not np.all(eta[:-1] <= eta[1:]): + raise ValueError("ETA values are not monotonically increasing") + if not np.all(etav[:-1] <= etav[1:]): + raise ValueError("ETAV values are not monotonically increasing") if 0.0 in bk: ks = 0 if km == 91 else np.where(bk == 0)[0][-1] @@ -857,4 +71,11 @@ def set_hybrid_pressure_coefficients(km: int) -> HybridPressureCoefficients: raise ValueError("bk must contain at least one 0.") pressure_data = HybridPressureCoefficients(ks, ptop, ak, bk) + return pressure_data + + +def check_eta(ak, bk): + from pace.fv3core.initialization.init_utils import compute_eta + + return compute_eta(ak, bk) diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index e61afe2d..2c556576 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -17,8 +17,8 @@ ) from pace.util import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_INTERFACE_DIM from pace.util.constants import N_HALO_DEFAULT, PI, RADIUS +from pace.util.grid import eta -from .eta import set_hybrid_pressure_coefficients from .geometry import ( calc_unit_vector_south, calc_unit_vector_west, @@ -225,6 +225,8 @@ def __init__( dx_const: float = 1000.0, dy_const: float = 1000.0, deglat: float = 15.0, + extdgrid: bool = False, + eta_file: str = "None", ): self._grid_type = grid_type self._dx_const = dx_const @@ -282,10 +284,12 @@ def __init__( self._dy_center = None self._area = None self._area_c = None - self._ks = None - self._ak = None - self._bk = None - self._ptop = None + ( + self._ks, + self._ptop, + self._ak, + self._bk, + ) = self._set_hybrid_pressure_coefficients(eta_file) self._ec1 = None self._ec2 = None self._ew1 = None @@ -409,11 +413,42 @@ def __init__( self._calculate_unit_vectors_lonlat = ( self._calculate_unit_vectors_lonlat_cube_sphere ) - self._init_dgrid() - self._init_agrid() + if extdgrid is False: + self._init_dgrid() + self._init_agrid() else: raise NotImplementedError(f"Unsupported grid_type = {grid_type}") + @classmethod + def from_external( + cls, + x, + y, + quantity_factory, + communicator, + grid_type, + eta_file: str = "None", + ) -> "MetricTerms": + """ + Generates a metric terms object, using input from data contained in an + externally generated tile file + """ + terms = MetricTerms( + quantity_factory=quantity_factory, + communicator=communicator, + grid_type=grid_type, + extdgrid=True, + eta_file=eta_file, + ) + + rad_conv = PI / 180.0 + terms._grid_64.view[:, :, 0] = rad_conv * x + terms._grid_64.view[:, :, 1] = rad_conv * y + + terms._init_agrid() + + return terms + @classmethod def from_tile_sizing( cls, @@ -426,6 +461,7 @@ def from_tile_sizing( dx_const: float = 1000.0, dy_const: float = 1000.0, deglat: float = 15.0, + eta_file: str = "None", ) -> "MetricTerms": sizer = util.SubtileGridSizer.from_tile_params( nx_tile=npx - 1, @@ -447,6 +483,7 @@ def from_tile_sizing( dx_const=dx_const, dy_const=dy_const, deglat=deglat, + eta_file=eta_file, ) @property @@ -586,13 +623,6 @@ def ks(self) -> util.Quantity: """ number of levels where the vertical coordinate is purely pressure-based """ - if self._ks is None: - ( - self._ks, - self._ptop, - self._ak, - self._bk, - ) = self._set_hybrid_pressure_coefficients() return self._ks @property @@ -601,13 +631,6 @@ def ak(self) -> util.Quantity: the ak coefficient used to calculate the pressure at a given k-level: pk = ak + (bk * ps) """ - if self._ak is None: - ( - self._ks, - self._ptop, - self._ak, - self._bk, - ) = self._set_hybrid_pressure_coefficients() return self._ak @property @@ -616,13 +639,6 @@ def bk(self) -> util.Quantity: the bk coefficient used to calculate the pressure at a given k-level: pk = ak + (bk * ps) """ - if self._bk is None: - ( - self._ks, - self._ptop, - self._ak, - self._bk, - ) = self._set_hybrid_pressure_coefficients() return self._bk @property @@ -630,13 +646,6 @@ def ptop(self) -> util.Quantity: """ the pressure of the top of atmosphere level """ - if self._ptop is None: - ( - self._ks, - self._ptop, - self._ak, - self._bk, - ) = self._set_hybrid_pressure_coefficients() return self._ptop @property @@ -2128,7 +2137,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): + def _set_hybrid_pressure_coefficients(self, eta_file): ks = self.quantity_factory.zeros( [], "", @@ -2149,7 +2158,9 @@ def _set_hybrid_pressure_coefficients(self): "", dtype=Float, ) - pressure_coefficients = set_hybrid_pressure_coefficients(self._npz) + pressure_coefficients = eta.set_hybrid_pressure_coefficients( + self._npz, eta_file + ) ks = pressure_coefficients.ks ptop = pressure_coefficients.ptop ak.data[:] = asarray(pressure_coefficients.ak, type(ak.data)) diff --git a/util/pace/util/monitor/netcdf_monitor.py b/util/pace/util/monitor/netcdf_monitor.py index 0b39da60..abcb7fe3 100644 --- a/util/pace/util/monitor/netcdf_monitor.py +++ b/util/pace/util/monitor/netcdf_monitor.py @@ -94,6 +94,7 @@ def flush(self): chunk=chunk_index, tile=self._tile ) ) + Path(self._path).mkdir(exist_ok=True) if os.path.exists(chunk_path): os.remove(chunk_path) ds.to_netcdf(chunk_path, format="NETCDF4", engine="netcdf4") diff --git a/util/pace/util/utils.py b/util/pace/util/utils.py index 737a55a0..9854609d 100644 --- a/util/pace/util/utils.py +++ b/util/pace/util/utils.py @@ -1,3 +1,4 @@ +from enum import EnumMeta from typing import Iterable, Sequence, Tuple, TypeVar, Union import numpy as np @@ -21,6 +22,11 @@ T = TypeVar("T") +class MetaEnumStr(EnumMeta): + def __contains__(cls, item) -> bool: + return item in cls.__members__.keys() + + def list_by_dims( dims: Sequence[str], horizontal_list: Sequence[T], non_horizontal_value: T ) -> Tuple[T, ...]: