diff --git a/.circleci/config.yml b/.circleci/config.yml index 149f4e31..3c7cce28 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -386,7 +386,7 @@ jobs: name: run tests command: | . venv/bin/activate - cd driver + cd pace GT_CACHE_ROOT=$(pwd)/.gt_cache MPIRUN_CALL="mpirun -n 6 --mca btl_vader_single_copy_mechanism none" make test_mpi - save_gt_cache: key: v1-gt_cache_driver @@ -409,9 +409,9 @@ jobs: name: run tests command: | . venv/bin/activate - cd driver + cd pace mpirun -n 6 --mca btl_vader_single_copy_mechanism none \ - python3 -m mpi4py -m pace.driver.run examples/configs/baroclinic_c12_orch_cpu.yaml + python3 -m mpi4py -m pace.run examples/configs/baroclinic_c12_orch_cpu.yaml test_main: machine: diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 47c1e8bb..4c592c6a 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -23,7 +23,7 @@ Before submitting this PR, please make sure: - [ ] You have followed the coding standards guidelines established at [Code Review Checklist](https://drive.google.com/file/d/1R0nqOxfYnzaSdoYdt8yjx5J482ETI2Ft/view?usp=sharing). - [ ] Docstrings and type hints are added to new and updated routines, as appropriate - [ ] All relevant documentation has been updated or added (e.g. README, CONTRIBUTING docs) -- [ ] For each public change and fix in `pace-util`, HISTORY has been updated +- [ ] For each public change and fix in `pace`, HISTORY has been updated - [ ] Unit tests are added or updated for non-stencil code changes Additionally, if this PR contains code authored by new contributors: diff --git a/.github/workflows/main_unit_tests.yml b/.github/workflows/main_unit_tests.yml new file mode 100644 index 00000000..79e745d6 --- /dev/null +++ b/.github/workflows/main_unit_tests.yml @@ -0,0 +1,31 @@ +name: "Main unit tests" +on: + pull_request: + types: [opened, synchronize, reopened, ready_for_review, labeled, unlabeled] + +jobs: + main_unit_tests: + runs-on: ubuntu-latest + steps: + - name: Checkout Pace repository + uses: actions/checkout@v3.5.2 + with: + submodules: 'recursive' + - name: Step Python 3.8.12 + uses: actions/setup-python@v4.6.0 + with: + python-version: '3.8.12' + - name: Install OpenMPI & Boost for gt4py + run: | + sudo apt-get install libopenmpi-dev libboost1.74-dev + - name: Install Python packages + run: | + python -m pip install --upgrade pip setuptools wheel + pip install -r requirements_dev.txt -c constraints.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: | + python -m pytest -x tests/main diff --git a/.gitignore b/.gitignore index 0a957698..5e261fa3 100644 --- a/.gitignore +++ b/.gitignore @@ -2,9 +2,9 @@ # Test outputs results.xml *.zarr -driver/examples/*.json -driver/examples/**/*.json -driver/examples/comm +examples/*.json +examples/**/*.json +examples/comm 20*-*-*-*-*-*.json *.pkl diff --git a/.jenkins/baroclinic_initialization.sh b/.jenkins/baroclinic_initialization.sh index 7d5d1c08..72a4231e 100755 --- a/.jenkins/baroclinic_initialization.sh +++ b/.jenkins/baroclinic_initialization.sh @@ -62,7 +62,7 @@ export VIRTUALENV=${JENKINS_DIR}/../venv_driver ${JENKINS_DIR}/install_virtualenv.sh ${VIRTUALENV} source ${VIRTUALENV}/bin/activate -CMD="srun python3 ${PACE_DIR}/driver/examples/baroclinic_init.py ${JENKINS_DIR}/driver_configs/${experiment}.yaml" +CMD="srun python3 ${PACE_DIR}/examples/baroclinic_init.py ${JENKINS_DIR}/driver_configs/${experiment}.yaml" run_command "${CMD}" Job${action} ${scheduler_script} $minutes if [ $? -ne 0 ] ; then exitError 1510 ${LINENO} "problem while executing script ${script}" @@ -71,7 +71,7 @@ fi module load sarus sarus pull elynnwu/pace:latest echo "####### generating figures..." -srun -C gpu --partition=debug --account=go31 --time=00:05:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/driver/examples/plot_baroclinic_init.py /work/output.zarr ${experiment} pt -1 +srun -C gpu --partition=debug --account=go31 --time=00:05:00 sarus run --mount=type=bind,source=${PACE_DIR},destination=/work elynnwu/pace:latest python /work/examples/plot_baroclinic_init.py /work/output.zarr ${experiment} pt -1 mkdir -p ${ARTIFACT_ROOT}/${experiment} echo "####### moving figures..." cp *.png ${ARTIFACT_ROOT}/${experiment}/. diff --git a/.jenkins/driver_performance.sh b/.jenkins/driver_performance.sh index f31a5ec3..22e3efac 100755 --- a/.jenkins/driver_performance.sh +++ b/.jenkins/driver_performance.sh @@ -33,7 +33,7 @@ cat << EOF > run.daint.slurm set -x export OMP_NUM_THREADS=12 export FV3_DACEMODE=BuildAndRun -srun python -m pace.driver.run ${JENKINS_DIR}/driver_configs/baroclinic_c192_6ranks.yaml +srun python -m pace.run ${JENKINS_DIR}/driver_configs/baroclinic_c192_6ranks.yaml EOF launch_job run.daint.slurm 3600 diff --git a/.jenkins/install_virtualenv.sh b/.jenkins/install_virtualenv.sh index b049b247..58ea86b6 100755 --- a/.jenkins/install_virtualenv.sh +++ b/.jenkins/install_virtualenv.sh @@ -35,7 +35,7 @@ cd ${PACE_DIR} python3 -m pip install $wheel_command -r ${PACE_DIR}/requirements_dev.txt -c ${PACE_DIR}/constraints.txt # have to be installed in non-develop mode because the directory where this gets built from # gets deleted before the tests run on daint -python3 -m pip install ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/driver ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/dsl ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/fv3core ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/physics ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/stencils ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/util ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/external/gt4py -c ${PACE_DIR}/constraints.txt +python3 -m pip install ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/driver ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/ndsl ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/pyFV3 ${PACE_VENV_INSTALL_PREFIX} ${PACE_DIR}/pySHiELD ${PACE_DIR}/external/gt4py -c ${PACE_DIR}/constraints.txt cd ${workdir} deactivate diff --git a/.jenkins/run_compare_fortran.sh b/.jenkins/run_compare_fortran.sh index 9010225e..d684ba92 100755 --- a/.jenkins/run_compare_fortran.sh +++ b/.jenkins/run_compare_fortran.sh @@ -55,7 +55,7 @@ cat << EOF > run.daint.slurm ######################################################## set -x export OMP_NUM_THREADS=12 -srun python -m pace.driver.run ${JENKINS_DIR}/driver_configs/${NAMELIST}.yaml +srun python -m pace.run ${JENKINS_DIR}/driver_configs/${NAMELIST}.yaml EOF launch_job run.daint.slurm 2400 diff --git a/.jenkins/run_diff_rank.sh b/.jenkins/run_diff_rank.sh index 75ba89bd..804294c6 100755 --- a/.jenkins/run_diff_rank.sh +++ b/.jenkins/run_diff_rank.sh @@ -29,7 +29,7 @@ cat << EOF > run.daint.slurm ######################################################## set -x export OMP_NUM_THREADS=12 -srun python -m pace.driver.run ${JENKINS_DIR}/driver_configs/baroclinic_c192_54ranks.yaml +srun python -m pace.run ${JENKINS_DIR}/driver_configs/baroclinic_c192_54ranks.yaml EOF launch_job run.daint.slurm 29000 @@ -55,7 +55,7 @@ cat << EOF > run.daint.slurm ######################################################## set -x export OMP_NUM_THREADS=12 -srun python -m pace.driver.run ${JENKINS_DIR}/driver_configs/baroclinic_c192_6ranks.yaml +srun python -m pace.run ${JENKINS_DIR}/driver_configs/baroclinic_c192_6ranks.yaml EOF launch_job run.daint.slurm 29000 diff --git a/.jenkins/test_driver.sh b/.jenkins/test_driver.sh index a2381957..dde5d114 100755 --- a/.jenkins/test_driver.sh +++ b/.jenkins/test_driver.sh @@ -3,4 +3,4 @@ JENKINS_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" PACE_IMAGE="driver_image" make -C ${JENKINS_DIR}/.. build -docker run --rm driver_image make -C /pace/driver test test_mpi +docker run --rm driver_image make -C /driver/pace/driver test test_mpi diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d5e8db5b..a6f865f7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,13 +18,10 @@ repos: rev: v0.812 hooks: - id: mypy - name: mypy-driver + name: mypy-pace + files: pace args: [--config-file, setup.cfg] - files: driver - exclude: | - (?x)^( - tests/main/driver/__init__.py | - )$ + - repo: https://github.com/pre-commit/pre-commit-hooks rev: v2.3.0 hooks: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 51dd8fd3..b45b7212 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -113,7 +113,7 @@ Turns into - Internal functions that are likely to be inlined into a larger stencil do not need this if it will just be removed in the near-term. ### GT4Py stencils -We interface to `gt4py.cartesian.gtscript.stencil` through pace.dsl.stencil, specifically the FrozenStencil, that allows us to minimize runtime overhead in calling stencils. +We interface to `gt4py.cartesian.gtscript.stencil` through ndsl.dsl.stencil, specifically the FrozenStencil, that allows us to minimize runtime overhead in calling stencils. ```python diff --git a/Makefile b/Makefile index 92771a95..71b78f08 100644 --- a/Makefile +++ b/Makefile @@ -115,7 +115,7 @@ dependencies.svg: dependencies.dot dot -Tsvg $< -o $@ .PHONY: constraints.txt -constraints.txt: driver/setup.py dsl/setup.py fv3core/setup.py physics/setup.py util/setup.py stencils/setup.py util/requirements.txt requirements_docs.txt requirements_lint.txt external/gt4py/setup.cfg requirements_dev.txt +constraints.txt: driver/setup.py ndsl/setup.py pyFV3/setup.py pySHiELD/setup.py requirements_docs.txt requirements_lint.txt external/gt4py/setup.cfg requirements_dev.txt pip-compile $^ --output-file constraints.txt sed -i.bak '/\@ git+https/d' constraints.txt rm -f constraints.txt.bak @@ -142,8 +142,8 @@ test_mpi_54rank: mpirun -n 54 $(MPIRUN_ARGS) python3 -m mpi4py -m pytest tests/mpi_54rank driver_savepoint_tests_mpi: build - TARGET=driver $(MAKE) get_test_data - $(CONTAINER_CMD) $(CONTAINER_FLAGS) bash -c "$(SAVEPOINT_SETUP) && cd $(PACE_PATH) && $(MPIRUN_CALL) python -m mpi4py -m pytest --maxfail=1 --data_path=$(EXPERIMENT_DATA_RUN)/driver/ $(TEST_ARGS) $(PHYSICS_THRESH_ARGS) -m parallel $(PACE_PATH)/physics/tests/savepoint" + TARGET=pace $(MAKE) get_test_data + $(CONTAINER_CMD) $(CONTAINER_FLAGS) bash -c "$(SAVEPOINT_SETUP) && cd $(PACE_PATH) && $(MPIRUN_CALL) python -m mpi4py -m pytest --maxfail=1 --data_path=$(EXPERIMENT_DATA_RUN)/pace/ $(TEST_ARGS) $(PHYSICS_THRESH_ARGS) -m parallel $(PACE_PATH)/physics/tests/savepoint" docs: ## generate Sphinx HTML documentation $(MAKE) -C docs html diff --git a/NDSL b/NDSL index 7c0f13fd..61c4cbe5 160000 --- a/NDSL +++ b/NDSL @@ -1 +1 @@ -Subproject commit 7c0f13fd357e276c8f6176d65d0b73b905fb3df2 +Subproject commit 61c4cbe51dd20e7d9eba50ea9f9b2427a20e0a80 diff --git a/README.md b/README.md index 959e4404..5188c318 100644 --- a/README.md +++ b/README.md @@ -77,13 +77,13 @@ Shell scripts to install Pace on specific machines such as Gaea can be found in With the environment activated, you can run an example baroclinic test case with the following command: ```shell -mpirun -n 6 python3 -m pace.driver.run driver/examples/configs/baroclinic_c12.yaml +mpirun -n 6 python3 -m pace.run examples/configs/baroclinic_c12.yaml # or with oversubscribe if you do not have at least 6 cores -mpirun -n 6 --oversubscribe python3 -m pace.driver.run driver/examples/configs/baroclinic_c12.yaml +mpirun -n 6 --oversubscribe python3 -m pace.run examples/configs/baroclinic_c12.yaml ``` -After the run completes, you will see an output direcotry `output.zarr`. An example to visualize the output is provided in `driver/examples/plot_output.py`. See the [driver example](driver/examples/README.md) section for more details. +After the run completes, you will see an output direcotry `output.zarr`. An example to visualize the output is provided in `examples/plot_output.py`. See the [driver example](examples/README.md) section for more details. ### Environment variable configuration @@ -116,7 +116,7 @@ make build ```shell make dev -mpirun --mca btl_vader_single_copy_mechanism none -n 6 python3 -m pace.driver.run /pace/driver/examples/configs/baroclinic_c12.yaml +mpirun --mca btl_vader_single_copy_mechanism none -n 6 python3 -m pace.run /examples/configs/baroclinic_c12.yaml ``` ## History diff --git a/constraints.txt b/constraints.txt index 5b20a8b7..110d3ebf 100644 --- a/constraints.txt +++ b/constraints.txt @@ -62,7 +62,6 @@ charset-normalizer==2.0.4 click==8.0.1 # via # black - # flask # gt4py cloudpickle==2.0.0 # via dask @@ -74,10 +73,12 @@ coverage==5.5 # via pytest-cov cytoolz==0.12.1 # via gt4py +dace + # via ndsl dacite==1.6.0 # via # fv3config - # pace-driver + # pace dask==2021.12.0 # via # -r requirements_dev.txt @@ -111,16 +112,19 @@ f90nml==1.3.1 # via # -r requirements_dev.txt # fv3config + # ndsl # pyfv3 # pyshield +factory-boy==3.3.0 + # via gt4py +faker==24.3.0 + # via factory-boy fasteners==0.16.3 # via zarr fastjsonschema==2.16.2 # via nbformat filelock==3.0.12 # via virtualenv -flask==2.1.2 - # via dace fparser==0.1.4 # via dace frozendict==2.3.4 @@ -143,8 +147,10 @@ google-auth-oauthlib==0.4.5 # via gcsfs gprof2dot==2021.2.21 # via pytest-profiling -gridtools-cpp +gridtools-cpp==2.3.2 # via gt4py +gt4py + # via ndsl h5netcdf==1.1.0 # via ndsl h5py==3.10.0 @@ -157,8 +163,6 @@ idna==3.2 # yarl imagesize==1.2.0 # via sphinx -importlib-metadata==4.11.3 - # via flask importlib-resources==5.10.0 # via # gt4py @@ -169,13 +173,11 @@ ipykernel==6.16.2 # via nbmake ipython==8.5.0 # via ipykernel -itsdangerous==2.1.2 - # via flask jedi==0.18.1 # via ipython jinja2==3.0.1 # via - # flask + # dace # gt4py # sphinx jsonschema==4.16.0 @@ -202,11 +204,11 @@ matplotlib-inline==0.1.6 # via # ipykernel # ipython -mpi4py==3.1.4 +mpi4py==3.1.5 # via # -r requirements_dev.txt # ndsl - # pace-driver + # pace mpmath==1.2.1 # via sympy multidict==5.1.0 @@ -234,7 +236,7 @@ netcdf4==1.6.4 # via # -r requirements_dev.txt # ndsl - # pace-driver + # pace networkx==2.6.3 # via dace ninja==1.11.1 @@ -243,7 +245,7 @@ nodeenv==1.6.0 # via pre-commit numcodecs==0.7.2 # via zarr -numpy==1.21.2 +numpy==1.23.3 # via # -r requirements_dev.txt # cftime @@ -252,7 +254,7 @@ numpy==1.21.2 # h5py # netcdf4 # numcodecs - # pace-driver + # pace # pandas # pyfv3 # pyshield @@ -286,7 +288,7 @@ pickleshare==0.7.5 # via ipython pkgutil-resolve-name==1.3.10 # via jsonschema -platformdirs==2.2.0 +platformdirs==3.9.1 # via # black # virtualenv @@ -312,7 +314,7 @@ pyasn1==0.4.8 # rsa pyasn1-modules==0.2.8 # via google-auth -pybind11==2.8.1 +pybind11==2.11.1 # via gt4py pydantic==1.7.4 # via nbmake @@ -346,6 +348,7 @@ pytest-subtests==0.5.0 # via -r requirements_dev.txt python-dateutil==2.8.2 # via + # faker # jupyter-client # pandas pytz==2021.1 @@ -357,7 +360,7 @@ pyyaml==5.4.1 # dace # dask # fv3config - # pace-driver + # pace # pre-commit # pytest-regressions pyzmq==24.0.1 @@ -368,7 +371,6 @@ recommonmark==0.7.1 # via -r requirements_docs.txt requests==2.26.0 # via - # dace # gcsfs # requests-oauthlib # sphinx @@ -454,6 +456,7 @@ typing-extensions==4.3.0 # via # aiohttp # black + # faker # gt4py # setuptools-scm urllib3==1.26.6 @@ -464,15 +467,13 @@ wcwidth==0.2.5 # via prompt-toolkit websockets==10.3 # via dace -werkzeug==2.1.2 - # via flask wheel==0.37.0 # via astunparse xarray==0.19.0 # via # -r requirements_dev.txt # ndsl - # pace-driver + # pace # pyfv3 # pyshield xxhash==2.0.2 @@ -482,11 +483,9 @@ yarl==1.6.3 zarr==2.9.2 # via # -r requirements_dev.txt - # pace-driver + # pace zipp==3.8.0 - # via - # importlib-metadata - # importlib-resources + # via importlib-resources # The following packages are considered to be unsafe in a requirements file: # setuptools diff --git a/doc_primer_orchestration.md b/doc_primer_orchestration.md index a10baf0c..f669155e 100644 --- a/doc_primer_orchestration.md +++ b/doc_primer_orchestration.md @@ -29,7 +29,7 @@ DaCe needs to be described all memory so it can interface it in the C code that File structure -------------- -`pace.dsl.dace.*` carries the structure for orchestration. +`ndsl.dsl.dace.*` carries the structure for orchestration. - `build.py`: tooling for distributed build & SDFG load. - `dace_config.py`: DaCeConfig & DaCeOrchestration enum. @@ -43,7 +43,7 @@ DaCe Config DaCe has many configuration options. When executing, it drops or reads a `dace.conf` to get/set options for execution. Because this is a performance-portable model and not a DaCe model, decision has been taken to freeze the options. -`pace.dsl.dace.dace_config` carries a set of tested options for DaCe, with doc. It also takes care of removing the `dace.conf` that will be generated automatically when using DaCe. +`ndsl.dsl.dace.dace_config` carries a set of tested options for DaCe, with doc. It also takes care of removing the `dace.conf` that will be generated automatically when using DaCe. Orchestration can be debugged by using the env var `PACE_DACE_DEBUG`. When set to `True`, this will drop a few checks: @@ -54,7 +54,7 @@ When set to `True`, this will drop a few checks: - `trace_all_outputs_at_index` drops a print on every variable at a given index to track numerical protection, - `sdfg_execution_progress`, drops a print after each kernel. Useful when encurring bad crash with no stacktrace, - insert a CUDA_ERROR_CHECK in C after each kernel. -See `dsl/pace/dsl/dace/utils.py` for details. +See `ndsl/dsl/dace/utils.py` for details. Build ----- diff --git a/docs/installation.rst b/docs/installation.rst index a4859bb6..44949a0c 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -10,7 +10,7 @@ When cloning Pace you will need to update the repository's submodules as well: .. code-block:: console - $ git clone --recursive https://github.com/ai2cm/pace.git + $ git clone --recursive https://github.com/NOAA-GFDL/pace.git or if you have already cloned the repository: diff --git a/driver/examples/README.md b/driver/examples/README.md deleted file mode 100644 index ac37af4a..00000000 --- a/driver/examples/README.md +++ /dev/null @@ -1,56 +0,0 @@ -# Driver Examples - -Here we have example scripts and configuration for running the Pace driver. -Currently this contains two examples which run the model on a baroclinic test case using the numpy backend. -You will find this runs fairly slowly, since the "compiled" code is still Python. -In the future, we will add examples for the compiled backends we support. - -The "docker" workflow here is written for the basic baroclinic test example. -`write_then_read.sh` is the second example, which shows how to configure an MPI run to write communication data to disk and then use it to repeat a single rank of that run with the same configuration. -This second example assumes you are already in an appropriate environment to run the driver, for example as documented in the "Host Machine" section below. - -Note that on the baroclinic test case example you will see significant grid imprinting in the first hour time evolution. -Rest assured that this duplicates the behavior of the original Fortran code. - -We have also included a utility to convert the zarr output of the run to netcdf, for convenience. To convert `output.zarr` to `output.nc`, you would run: - -```bash -$ python3 zarr_to_nc.py output.zarr output.nc -``` - -Another example is `baroclinic_init.py`, which initializes a barcolinic wave and writes out the grid and the initial state. To run this script with the c12 6ranks example: - -```bash -$ mpirun -n 6 python3 baroclinic_init.py ./configs/baroclinic_c12.yaml -``` -## Docker - -To run a baroclinic c12 case with Docker in a single command, run `run_docker.sh`. -This example will start from the Python 3.8 docker image, install extra dependencies and Python packages, and execute the example, leaving the output in this directory. - -To visualize the output, two example scripts are provided: -1. `plot_output.py`: To use it, you must install matplotlib (e.g. with `pip install matplotlib`). -2. `plot_cube.py`: this uses plotting tools in [fv3viz](https://github.com/ai2cm/fv3net/tree/master/external/fv3viz). Note the requirements aren't part of pace by default and need to be installed accordingly. It is recommended to use the post processing docker provided at the top level `docker/postprocessing.Dockerfile`. - -## Host Machine - -To run examples on your host machine, you will need to have an MPI library on your machine suitable for installing mpi4py. -For example, on Ubuntu 20.04 this could be the libopenmpi3 and libopenmpi-dev libraries. - -With these requirements installed, set up the virtual environment with - -```bash -$ create_venv.sh -$ . venv/bin/activate -``` - -With the environment activated, the model itself can be run with `python3 -m pace.driver.run `. -Currently this must be done using some kind of mpirun, though there are plans to enable single-tile runs without MPI. -The exact command will vary based on your MPI implementation, but you can try running - -```bash -$ mpirun -n 6 python3 -m pace.driver.run ./configs/baroclinic_c12.yaml -``` - -To run the example at C48 resolution instead of C12, you can update the value of `nx_tile` in the configuration file from 12 to 48. -Here you can also change the timestep in seconds `dt_atmos`, as well as the total run duration with `minutes`, or by adding values for `hours` or `days`. diff --git a/examples/README.md b/examples/README.md index 487badc9..1c98ba40 100644 --- a/examples/README.md +++ b/examples/README.md @@ -44,7 +44,7 @@ Within the `build_scripts` directory are a couple scripts relevant for setting u ``` $ ssh analysis -$ git clone --recursive https://github.com/ai2cm/pace.git +$ git clone --recursive https://github.com/NOAA-GFDL/pace.git $ cd pace/examples/build_scripts ``` @@ -63,3 +63,60 @@ $ jhp launch lab ``` It will take a minute or so for the server to start up. If you are on GFDL's network or have [configured your proxy settings appropriately](https://wiki.gfdl.noaa.gov/index.php/Creating_a_GFDL_SSH_Tunnel), you will then be able to navigate to the URL produced by the command in your local browser. This will take you to the JupyterLab interface, where you can open and run the example notebooks. Note that it can take some time to initially connect to the notebook kernel, so be patient. + +# Driver Examples + +Here we have example scripts and configuration for running the Pace driver. +Currently this contains two examples which run the model on a baroclinic test case using the numpy backend. +You will find this runs fairly slowly, since the "compiled" code is still Python. +In the future, we will add examples for the compiled backends we support. + +The "docker" workflow here is written for the basic baroclinic test example. +`write_then_read.sh` is the second example, which shows how to configure an MPI run to write communication data to disk and then use it to repeat a single rank of that run with the same configuration. +This second example assumes you are already in an appropriate environment to run the driver, for example as documented in the "Host Machine" section below. + +Note that on the baroclinic test case example you will see significant grid imprinting in the first hour time evolution. +Rest assured that this duplicates the behavior of the original Fortran code. + +We have also included a utility to convert the zarr output of the run to netcdf, for convenience. To convert `output.zarr` to `output.nc`, you would run: + +```bash +$ python3 zarr_to_nc.py output.zarr output.nc +``` + +Another example is `baroclinic_init.py`, which initializes a barcolinic wave and writes out the grid and the initial state. To run this script with the c12 6ranks example: + +```bash +$ mpirun -n 6 python3 baroclinic_init.py ./configs/baroclinic_c12.yaml +``` +## Docker + +To run a baroclinic c12 case with Docker in a single command, run `run_docker.sh`. +This example will start from the Python 3.8 docker image, install extra dependencies and Python packages, and execute the example, leaving the output in this directory. + +To visualize the output, two example scripts are provided: +1. `plot_output.py`: To use it, you must install matplotlib (e.g. with `pip install matplotlib`). +2. `plot_cube.py`: this uses plotting tools in [fv3viz](https://github.com/ai2cm/fv3net/tree/master/external/fv3viz). Note the requirements aren't part of pace by default and need to be installed accordingly. It is recommended to use the post processing docker provided at the top level `docker/postprocessing.Dockerfile`. + +## Host Machine + +To run examples on your host machine, you will need to have an MPI library on your machine suitable for installing mpi4py. +For example, on Ubuntu 20.04 this could be the libopenmpi3 and libopenmpi-dev libraries. + +With these requirements installed, set up the virtual environment with + +```bash +$ create_venv.sh +$ . venv/bin/activate +``` + +With the environment activated, the model itself can be run with `python3 -m pace.run `. +Currently this must be done using some kind of mpirun, though there are plans to enable single-tile runs without MPI. +The exact command will vary based on your MPI implementation, but you can try running + +```bash +$ mpirun -n 6 python3 -m pace.run examples/configs/baroclinic_c12.yaml +``` + +To run the example at C48 resolution instead of C12, you can update the value of `nx_tile` in the configuration file from 12 to 48. +Here you can also change the timestep in seconds `dt_atmos`, as well as the total run duration with `minutes`, or by adding values for `hours` or `days`. diff --git a/driver/examples/baroclinic_init.py b/examples/baroclinic_init.py similarity index 92% rename from driver/examples/baroclinic_init.py rename to examples/baroclinic_init.py index fefc0c2d..17f7f4cd 100644 --- a/driver/examples/baroclinic_init.py +++ b/examples/baroclinic_init.py @@ -2,7 +2,7 @@ import yaml -from pace.driver.run import Driver, DriverConfig +from pace import Driver, DriverConfig def parse_args(): diff --git a/driver/examples/configs/README.md b/examples/configs/README.md similarity index 100% rename from driver/examples/configs/README.md rename to examples/configs/README.md diff --git a/driver/examples/configs/analytic_test.yaml b/examples/configs/analytic_test.yaml similarity index 100% rename from driver/examples/configs/analytic_test.yaml rename to examples/configs/analytic_test.yaml diff --git a/driver/examples/configs/baroclinic_c12.yaml b/examples/configs/baroclinic_c12.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12.yaml rename to examples/configs/baroclinic_c12.yaml diff --git a/driver/examples/configs/baroclinic_c12_comm_read.yaml b/examples/configs/baroclinic_c12_comm_read.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12_comm_read.yaml rename to examples/configs/baroclinic_c12_comm_read.yaml diff --git a/driver/examples/configs/baroclinic_c12_comm_write.yaml b/examples/configs/baroclinic_c12_comm_write.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12_comm_write.yaml rename to examples/configs/baroclinic_c12_comm_write.yaml diff --git a/driver/examples/configs/baroclinic_c12_dp.yaml b/examples/configs/baroclinic_c12_dp.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12_dp.yaml rename to examples/configs/baroclinic_c12_dp.yaml diff --git a/driver/examples/configs/baroclinic_c12_explicit_physics.yaml b/examples/configs/baroclinic_c12_explicit_physics.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12_explicit_physics.yaml rename to examples/configs/baroclinic_c12_explicit_physics.yaml diff --git a/driver/examples/configs/baroclinic_c12_from_serialbox.yaml b/examples/configs/baroclinic_c12_from_serialbox.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12_from_serialbox.yaml rename to examples/configs/baroclinic_c12_from_serialbox.yaml diff --git a/driver/examples/configs/baroclinic_c12_null_comm.yaml b/examples/configs/baroclinic_c12_null_comm.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12_null_comm.yaml rename to examples/configs/baroclinic_c12_null_comm.yaml diff --git a/driver/examples/configs/baroclinic_c12_orch_cpu.yaml b/examples/configs/baroclinic_c12_orch_cpu.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12_orch_cpu.yaml rename to examples/configs/baroclinic_c12_orch_cpu.yaml diff --git a/driver/examples/configs/baroclinic_c12_read_restart_fortran.yml b/examples/configs/baroclinic_c12_read_restart_fortran.yml similarity index 100% rename from driver/examples/configs/baroclinic_c12_read_restart_fortran.yml rename to examples/configs/baroclinic_c12_read_restart_fortran.yml diff --git a/driver/examples/configs/baroclinic_c12_write_restart.yaml b/examples/configs/baroclinic_c12_write_restart.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c12_write_restart.yaml rename to examples/configs/baroclinic_c12_write_restart.yaml diff --git a/driver/examples/configs/baroclinic_c48_6ranks_serialbox_test.yaml b/examples/configs/baroclinic_c48_6ranks_serialbox_test.yaml similarity index 100% rename from driver/examples/configs/baroclinic_c48_6ranks_serialbox_test.yaml rename to examples/configs/baroclinic_c48_6ranks_serialbox_test.yaml diff --git a/driver/examples/configs/test_external_C12_1x1.yaml b/examples/configs/test_external_C12_1x1.yaml similarity index 100% rename from driver/examples/configs/test_external_C12_1x1.yaml rename to examples/configs/test_external_C12_1x1.yaml diff --git a/driver/examples/configs/test_external_C12_2x2.yaml b/examples/configs/test_external_C12_2x2.yaml similarity index 100% rename from driver/examples/configs/test_external_C12_2x2.yaml rename to examples/configs/test_external_C12_2x2.yaml diff --git a/driver/examples/configs/tropical_read_restart_fortran.yml b/examples/configs/tropical_read_restart_fortran.yml similarity index 100% rename from driver/examples/configs/tropical_read_restart_fortran.yml rename to examples/configs/tropical_read_restart_fortran.yml diff --git a/driver/examples/configs/tropicalcyclone_c128.yaml b/examples/configs/tropicalcyclone_c128.yaml similarity index 100% rename from driver/examples/configs/tropicalcyclone_c128.yaml rename to examples/configs/tropicalcyclone_c128.yaml diff --git a/driver/examples/create_venv.sh b/examples/create_venv.sh similarity index 100% rename from driver/examples/create_venv.sh rename to examples/create_venv.sh diff --git a/driver/examples/notebooks/driver_write_config.yaml b/examples/notebooks/driver_write_config.yaml similarity index 100% rename from driver/examples/notebooks/driver_write_config.yaml rename to examples/notebooks/driver_write_config.yaml diff --git a/examples/notebooks/functions.py b/examples/notebooks/functions.py index ffc37ca6..fb7ca1fa 100644 --- a/examples/notebooks/functions.py +++ b/examples/notebooks/functions.py @@ -10,12 +10,22 @@ from matplotlib import animation from units_config import units -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner +from ndsl import ( + CompilationConfig, + CubedSphereCommunicator, + CubedSpherePartitioner, + DaceConfig, + DaCeOrchestration, + GridIndexing, + Quantity, + QuantityFactory, + RunMode, + StencilConfig, + StencilFactory, + SubtileGridSizer, + TilePartitioner, +) from ndsl.constants import RADIUS -from ndsl.dsl.dace.dace_config import DaceConfig, DaCeOrchestration -from ndsl.dsl.stencil import GridIndexing, StencilConfig, StencilFactory -from ndsl.dsl.stencil_config import CompilationConfig, RunMode from ndsl.grid import ( AngleGridData, ContravariantGridData, @@ -26,12 +36,7 @@ VerticalGridData, ) from ndsl.grid.gnomonic import great_circle_distance_lon_lat -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from ndsl.quantity import Quantity -from pyFV3.stencils.fvtp2d import FiniteVolumeTransport -from pyFV3.stencils.fxadv import FiniteVolumeFluxPrep -from pyFV3.stencils.tracer_2d_1l import TracerAdvection +from pyFV3.stencils import FiniteVolumeFluxPrep, FiniteVolumeTransport, TracerAdvection class GridType(enum.Enum): diff --git a/examples/notebooks/grid_generation.ipynb b/examples/notebooks/grid_generation.ipynb index 34518f60..cb5f039e 100644 --- a/examples/notebooks/grid_generation.ipynb +++ b/examples/notebooks/grid_generation.ipynb @@ -188,7 +188,7 @@ "source": [ "## Domain decomposition\n", "\n", - "The `pace.util` package provides a number of convenience functions to facilitate domain-decomposition.\n", + "The `ndsl` package provides a number of convenience functions to facilitate domain-decomposition.\n", "\n", "- `TilePartitioner`: determines domain-decomposition on a single tile of the cubed-sphere;\n", "- `CubedSpherePartitioner`: determines domain-decomposition of the entire cubed-sphere, including information about the orientation of each tile\n", @@ -360,7 +360,7 @@ } ], "source": [ - "from ndsl.util import (\n", + "from ndsl import (\n", " CubedSphereCommunicator,\n", " CubedSpherePartitioner,\n", " Quantity,\n", @@ -451,7 +451,7 @@ "\n", "Once we have decomposed the computational domain on the cube, we still need to project the cube onto the sphere and compute the associated metric terms that are essential for numerical discretization on the cubed-sphere.\n", "\n", - "The `ndsl.util.grid` package implements helper functions which make this very easy:\n", + "The `ndsl.grid` package implements helper functions which make this very easy:\n", "- `MetricTerms`: class to compute the metric terms of the cubed-sphere grid\n", "- `GridData`: storage of basic grid data (e.g. lat, lon, area, dx, dy, ...)" ] @@ -493,7 +493,7 @@ } ], "source": [ - "from ndsl.util.grid import GridData, MetricTerms\n", + "from ndsl.grid import GridData, MetricTerms\n", "\n", "# create the object to compute metric terms\n", "metric_terms = MetricTerms(quantity_factory=quantity_factory, communicator=communicator)\n", @@ -577,7 +577,7 @@ ], "source": [ "import numpy as np\n", - "from pace.util import Quantity\n", + "from ndsl import Quantity\n", "\n", "# convert the GT4Py storages into Quantity objects\n", "lon = Quantity(\n", diff --git a/examples/notebooks/initial_condition_definition.ipynb b/examples/notebooks/initial_condition_definition.ipynb index 4b194a01..7d08eaa2 100644 --- a/examples/notebooks/initial_condition_definition.ipynb +++ b/examples/notebooks/initial_condition_definition.ipynb @@ -455,7 +455,7 @@ "metadata": {}, "outputs": [], "source": [ - "from ndsl.util.grid.gnomonic import great_circle_distance_lon_lat\n", + "from ndsl.grid.gnomonic import great_circle_distance_lon_lat\n", "\n", "tracer = quantity_factory.zeros(\n", " dims=(\"x\", \"y\", \"z\"), units=units[\"tracer\"], dtype=\"float\"\n", diff --git a/driver/examples/notebooks/serial_debugging.ipynb b/examples/notebooks/serial_debugging.ipynb similarity index 96% rename from driver/examples/notebooks/serial_debugging.ipynb rename to examples/notebooks/serial_debugging.ipynb index 6af12a00..597051a1 100644 --- a/driver/examples/notebooks/serial_debugging.ipynb +++ b/examples/notebooks/serial_debugging.ipynb @@ -27,9 +27,9 @@ "metadata": {}, "outputs": [], "source": [ - "import pace.util\n", - "import pace.driver\n", - "import pace.fv3core.stencils.d2a2c_vect\n", + "from ndsl.constants import X_DIM, Y_DIM, Z_DIM\n", + "from pace import Driver, DriverConfig\n", + "from pyFV3.stencils import DGrid2AGrid2CGridVectors\n", "import yaml\n", "import dacite\n", "import copy\n", @@ -156,7 +156,7 @@ "metadata": {}, "outputs": [], "source": [ - "driver = pace.driver.Driver(config=pace.driver.DriverConfig.from_dict(read_config))" + "driver = Driver(config=DriverConfig.from_dict(read_config))" ] }, { @@ -344,7 +344,7 @@ "metadata": {}, "outputs": [], "source": [ - "d2a2c = fv3core.stencils.d2a2c_vect.DGrid2AGrid2CGridVectors(\n", + "d2a2c = DGrid2AGrid2CGridVectors(\n", " stencil_factory=driver.stencil_factory,\n", " grid_data=driver.state.grid_data,\n", " nested=False,\n", @@ -371,8 +371,8 @@ "outputs": [], "source": [ "# need temporaries for c-grid contravariant wind multiplied by timestep\n", - "utc = driver.quantity_factory.zeros(dims=[pace.util.X_DIM, pace.util.Y_DIM, pace.util.Z_DIM], units=\"m/s\")\n", - "vtc = driver.quantity_factory.zeros(dims=[pace.util.X_DIM, pace.util.Y_DIM, pace.util.Z_DIM], units=\"m/s\")\n", + "utc = driver.quantity_factory.zeros(dims=[X_DIM, Y_DIM, Z_DIM], units=\"m/s\")\n", + "vtc = driver.quantity_factory.zeros(dims=[X_DIM, Y_DIM, Z_DIM], units=\"m/s\")\n", "d2a2c(\n", " uc=remap_state.uc,\n", " vc=remap_state.vc,\n", diff --git a/examples/notebooks/stencil_definition.ipynb b/examples/notebooks/stencil_definition.ipynb index a5efcc51..f65ed20a 100644 --- a/examples/notebooks/stencil_definition.ipynb +++ b/examples/notebooks/stencil_definition.ipynb @@ -84,7 +84,7 @@ "source": [ "## Setup helpers for building stencils\n", "\n", - "The `ndsl.dsl` package contains provides a helper class for compiling stencils. The helper class is called a stencil factory. In order to setup a stencil factory, several configuration objects have to be defined and passed in.\n", + "The `ndsl` package contains provides a helper class for compiling stencils. The helper class is called a stencil factory. In order to setup a stencil factory, several configuration objects have to be defined and passed in.\n", "\n", "- `DaceConfig`: configuration of DaCe backend\n", "- `CompilationConfig`: specification of how to compile\n", @@ -102,9 +102,7 @@ }, "outputs": [], "source": [ - "from ndsl.dsl.dace.dace_config import DaceConfig, DaCeOrchestration\n", - "from ndsl.dsl.stencil import GridIndexing, StencilConfig, StencilFactory\n", - "from ndsl.dsl.stencil_config import CompilationConfig, RunMode\n", + "from ndsl import DaceConfig, DaCeOrchestration, GridIndexing, StencilConfig, StencilFactory, CompilationConfig, RunMode\n", "\n", "dace_config = DaceConfig(\n", " communicator=None, backend=backend, orchestration=DaCeOrchestration.Python\n", @@ -155,7 +153,7 @@ "metadata": {}, "outputs": [], "source": [ - "from ndsl.dsl.stencil import StencilFactory\n", + "from ndsl import StencilFactory\n", "from ndsl.dsl.typing import FloatField\n", "\n", "\n", @@ -183,7 +181,7 @@ "source": [ "## Compiling and running the stencil\n", "\n", - "To compile and run the stencil we first have to instanciate the wrapper class. This will compile the stencil and return a callable object. Next we have to define a GT4Py data storage (field). This could also be done using the `ndsl.util.Quantity` class, but here we use plain GT4Py data storages for simplicity. Finally, we can call the stencil and pass in the field and value." + "To compile and run the stencil we first have to instanciate the wrapper class. This will compile the stencil and return a callable object. Next we have to define a GT4Py data storage (field). This could also be done using the `ndsl.quantity.Quantity` class, but here we use plain GT4Py data storages for simplicity. Finally, we can call the stencil and pass in the field and value." ] }, { @@ -582,10 +580,8 @@ "metadata": {}, "outputs": [], "source": [ - "from ndsl.dsl.dace.dace_config import DaceConfig, DaCeOrchestration\n", - "from ndsl.dsl.stencil import GridIndexing, StencilConfig, StencilFactory\n", - "from ndsl.dsl.stencil_config import CompilationConfig, RunMode\n", - "from ndsl.util.grid import AngleGridData, ContravariantGridData, DampingCoefficients, GridData, HorizontalGridData, MetricTerms, VerticalGridData\n", + "from ndsl import DaceConfig, DaCeOrchestration, GridIndexing, StencilConfig, StencilFactory, CompilationConfig, RunMode\n", + "from ndsl.grid import AngleGridData, ContravariantGridData, DampingCoefficients, GridData, HorizontalGridData, MetricTerms, VerticalGridData\n", "\n", "\n", "if nz == 1:\n", @@ -706,7 +702,7 @@ "source": [ "import copy as cp\n", "import matplotlib.pyplot as plt\n", - "from pace.fv3core.stencils.fxadv import FiniteVolumeFluxPrep\n", + "from pyFV3.stencils import FiniteVolumeFluxPrep\n", "from units_config import units\n", "\n", "fvf_prep = FiniteVolumeFluxPrep(\n", @@ -839,8 +835,7 @@ } ], "source": [ - "from pace.fv3core.stencils.fvtp2d import FiniteVolumeTransport\n", - "from pace.fv3core.stencils.tracer_2d_1l import TracerAdvection\n", + "from pyFV3.stencils import FiniteVolumeTransport, TracerAdvection\n", "\n", "grid_type = 0 # cubed-sphere\n", "hord = 6 # horizontal diffusion order\n", diff --git a/driver/examples/plot_baroclinic_init.py b/examples/plot_baroclinic_init.py similarity index 100% rename from driver/examples/plot_baroclinic_init.py rename to examples/plot_baroclinic_init.py diff --git a/driver/examples/plot_cube.py b/examples/plot_cube.py similarity index 100% rename from driver/examples/plot_cube.py rename to examples/plot_cube.py diff --git a/driver/examples/plot_output.py b/examples/plot_output.py similarity index 100% rename from driver/examples/plot_output.py rename to examples/plot_output.py diff --git a/driver/examples/plot_pcolormesh_cube.py b/examples/plot_pcolormesh_cube.py similarity index 100% rename from driver/examples/plot_pcolormesh_cube.py rename to examples/plot_pcolormesh_cube.py diff --git a/driver/examples/run_docker.sh b/examples/run_docker.sh similarity index 54% rename from driver/examples/run_docker.sh rename to examples/run_docker.sh index cf70194f..de1692ea 100755 --- a/driver/examples/run_docker.sh +++ b/examples/run_docker.sh @@ -2,4 +2,4 @@ SCRIPT_DIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )" -docker run -v ${SCRIPT_DIR}/../../:/pace -w /pace python:3.8 bash -c "apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y python3-venv python3-dev libopenmpi3 libopenmpi-dev && cd /pace/driver/examples && /pace/driver/examples/create_venv.sh && . venv/bin/activate && mpirun -n 6 --allow-run-as-root --mca btl_vader_single_copy_mechanism none python3 -m pace.driver.run /pace/driver/examples/configs/baroclinic_c12.yaml" +docker run -v ${SCRIPT_DIR}/../../:/pace -w /pace python:3.8 bash -c "apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y python3-venv python3-dev libopenmpi3 libopenmpi-dev && cd /examples && /examples/create_venv.sh && . venv/bin/activate && mpirun -n 6 --allow-run-as-root --mca btl_vader_single_copy_mechanism none python3 -m pace.run /examples/configs/baroclinic_c12.yaml" diff --git a/driver/examples/stencil_signatures.py b/examples/stencil_signatures.py similarity index 80% rename from driver/examples/stencil_signatures.py rename to examples/stencil_signatures.py index b23cef99..1ddbe474 100644 --- a/driver/examples/stencil_signatures.py +++ b/examples/stencil_signatures.py @@ -4,15 +4,14 @@ import yaml -import ndsl.dsl -import ndsl.util -import pace.driver +import pace +from ndsl import FrozenStencil def has_stencils(object): for name in dir(object): try: - stencil_found = isinstance(getattr(object, name), ndsl.dsl.FrozenStencil) + stencil_found = isinstance(getattr(object, name), FrozenStencil) except (AttributeError, RuntimeError): stencil_found = False if stencil_found: @@ -26,7 +25,7 @@ def report_stencils(obj, file: Optional[TextIO]): print(f"module {module.__name__}, class {obj.__class__.__name__}:", file=file) all_access_names = collections.defaultdict(list) for name, value in obj.__dict__.items(): - if isinstance(value, ndsl.dsl.FrozenStencil): + if isinstance(value, FrozenStencil): print(f" stencil {name}:", file=file) for arg_name, field_info in value.stencil_object.field_info.items(): if field_info is None: @@ -52,10 +51,10 @@ def report_stencils(obj, file: Optional[TextIO]): if __name__ == "__main__": with open("configs/baroclinic_c12.yaml", "r") as f: - driver_config = pace.driver.DriverConfig.from_dict(yaml.safe_load(f)) - driver_config.comm_config = pace.driver.CreatesCommSelector( - config=pace.driver.NullCommConfig(rank=0, total_ranks=6), type="null" + driver_config = pace.DriverConfig.from_dict(yaml.safe_load(f)) + driver_config.comm_config = pace.CreatesCommSelector( + config=pace.NullCommConfig(rank=0, total_ranks=6), type="null" ) - driver = pace.driver.Driver(config=driver_config) + driver = pace.Driver(config=driver_config) with open("stencil_report.txt", "w") as f: report_stencils(driver.dycore, file=f) diff --git a/driver/examples/write_then_read.sh b/examples/write_then_read.sh similarity index 57% rename from driver/examples/write_then_read.sh rename to examples/write_then_read.sh index a3e89a35..869be852 100755 --- a/driver/examples/write_then_read.sh +++ b/examples/write_then_read.sh @@ -5,5 +5,5 @@ set -e -x MPIRUN_CMD=${MPIRUN_CMD:-mpirun -n 6} -$MPIRUN_CMD python3 -m pace.driver.run configs/baroclinic_c12_comm_write.yaml --log-rank 0 -python3 -m pace.driver.run configs/baroclinic_c12_comm_read.yaml +$MPIRUN_CMD python3 -m pace.run configs/baroclinic_c12_comm_write.yaml --log-rank 0 +python3 -m pace.run configs/baroclinic_c12_comm_read.yaml diff --git a/driver/examples/zarr_to_nc.py b/examples/zarr_to_nc.py similarity index 100% rename from driver/examples/zarr_to_nc.py rename to examples/zarr_to_nc.py diff --git a/driver/Makefile b/pace/Makefile similarity index 72% rename from driver/Makefile rename to pace/Makefile index 250a0a93..435ff599 100644 --- a/driver/Makefile +++ b/pace/Makefile @@ -2,7 +2,7 @@ MPIRUN_CALL ?= mpirun -n 6 test_mpi: MPIRUN_CALL="$(MPIRUN_CALL)" pytest tests/mpi - $(MPIRUN_CALL) python3 -m mpi4py -m pace.driver.run examples/configs/baroclinic_c12.yaml + $(MPIRUN_CALL) python3 -m mpi4py -m pace.run examples/configs/baroclinic_c12.yaml cd examples && MPIRUN_CALL="$(MPIRUN_CALL)" ./write_then_read.sh clean: diff --git a/driver/README.md b/pace/README.md similarity index 81% rename from driver/README.md rename to pace/README.md index 7675eacd..86a4bca3 100644 --- a/driver/README.md +++ b/pace/README.md @@ -1,8 +1,8 @@ -# pace-driver +# pace This package provides command-line routines to run the Pace model, and utilities to write model driver scripts. -We suggest reading the code in the examples directory, or taking a look at `pace/driver/run.py` to see how the main entrypoint for this package works. +We suggest reading the code in the examples directory, or taking a look at `pace/run.py` to see how the main entrypoint for this package works. # Usage @@ -10,8 +10,8 @@ Usage examples exist in the examples directory. The command-line interface may be run in certain debugging modes in serial, but usually you will want to run it using an mpi executor such as mpirun. ```bash -$ python3 -m pace.driver.run --help -Usage: python -m pace.driver.run [OPTIONS] CONFIG_PATH +$ python3 -m pace.run --help +Usage: python -m pace.run [OPTIONS] CONFIG_PATH Run the driver. diff --git a/driver/pace/driver/__init__.py b/pace/__init__.py similarity index 74% rename from driver/pace/driver/__init__.py rename to pace/__init__.py index df00d341..d6bd5eb1 100644 --- a/driver/pace/driver/__init__.py +++ b/pace/__init__.py @@ -11,8 +11,14 @@ from .diagnostics import Diagnostics, DiagnosticsConfig from .driver import Driver, DriverConfig, RestartConfig from .grid import GeneratedGridConfig, SerialboxGridConfig -from .initialization import AnalyticInit, PredefinedStateInit, RestartInit +from .initialization import ( + AnalyticInit, + FortranRestartInit, + PredefinedStateInit, + RestartInit, +) from .registry import Registry +from .safety_checks import SafetyChecker from .state import DriverState, TendencyState diff --git a/driver/pace/driver/comm.py b/pace/comm.py similarity index 95% rename from driver/pace/driver/comm.py rename to pace/comm.py index 27b00a5b..87bfcf8e 100644 --- a/driver/pace/driver/comm.py +++ b/pace/comm.py @@ -3,12 +3,9 @@ import os from typing import Any, ClassVar, List -from ndsl.comm.caching_comm import CachingCommReader, CachingCommWriter -from ndsl.comm.comm_abc import Comm -from ndsl.comm.mpi import MPIComm -from ndsl.comm.null_comm import NullComm - -from .registry import Registry +from ndsl import MPIComm, NullComm +from ndsl.comm import CachingCommReader, CachingCommWriter, Comm +from pace.registry import Registry class CreatesComm(abc.ABC): diff --git a/driver/pace/driver/configs/__init__.py b/pace/configs/__init__.py similarity index 100% rename from driver/pace/driver/configs/__init__.py rename to pace/configs/__init__.py diff --git a/driver/pace/driver/configs/comm.py b/pace/configs/comm.py similarity index 96% rename from driver/pace/driver/configs/comm.py rename to pace/configs/comm.py index bf9c00b7..06f60acf 100644 --- a/driver/pace/driver/configs/comm.py +++ b/pace/configs/comm.py @@ -5,9 +5,8 @@ import dacite -from ndsl.comm.caching_comm import CachingCommReader, CachingCommWriter -from ndsl.comm.mpi import MPIComm -from ndsl.comm.null_comm import NullComm +from ndsl import MPIComm, NullComm +from ndsl.comm import CachingCommReader, CachingCommWriter class CreatesComm(abc.ABC): diff --git a/driver/pace/driver/diagnostics.py b/pace/diagnostics.py similarity index 97% rename from driver/pace/driver/diagnostics.py rename to pace/diagnostics.py index b5961595..b395b353 100644 --- a/driver/pace/driver/diagnostics.py +++ b/pace/diagnostics.py @@ -4,16 +4,16 @@ from datetime import datetime, timedelta from typing import List, Optional, Union -from ndsl.comm.communicator import Communicator +from ndsl import Quantity from ndsl.constants import RGRAV, Z_DIM, Z_INTERFACE_DIM from ndsl.dsl.dace.orchestration import dace_inhibitor from ndsl.filesystem import get_fs from ndsl.grid import GridData -from ndsl.monitor import Monitor, NetCDFMonitor, ZarrMonitor -from ndsl.quantity import Quantity -from pyFV3.dycore_state import DycoreState - -from .state import DriverState +from ndsl.monitor import Monitor, ZarrMonitor +from ndsl.monitor.netcdf_monitor import NetCDFMonitor +from ndsl.typing import Communicator +from pace.state import DriverState +from pyFV3 import DycoreState try: diff --git a/driver/pace/driver/driver.py b/pace/driver.py similarity index 95% rename from driver/pace/driver/driver.py rename to pace/driver.py index 2210bc96..859778c2 100644 --- a/driver/pace/driver/driver.py +++ b/pace/driver.py @@ -9,39 +9,38 @@ import dacite import yaml -import pyFV3 -import pySHiELD -from ndsl.comm.comm_abc import Comm -from ndsl.comm.communicator import ( - Communicator, +from ndsl import ( + CompilationConfig, CubedSphereCommunicator, + DaceConfig, + GridIndexing, + PerformanceCollector, + QuantityFactory, + RunMode, + StencilConfig, + StencilFactory, + SubtileGridSizer, TileCommunicator, + TilePartitioner, + ndsl_log, ) -from ndsl.comm.partitioner import TilePartitioner +from ndsl.comm import Comm from ndsl.constants import N_HALO_DEFAULT -from ndsl.dsl.dace.dace_config import DaceConfig from ndsl.dsl.dace.orchestration import dace_inhibitor, orchestrate -from ndsl.dsl.stencil import GridIndexing, StencilFactory -from ndsl.dsl.stencil_config import CompilationConfig, RunMode, StencilConfig from ndsl.dsl.typing import Float from ndsl.grid import DampingCoefficients, DriverGridData, GridData -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from ndsl.logging import ndsl_log -from ndsl.performance import PerformanceConfig -from ndsl.performance.collector import PerformanceCollector -from ndsl.performance.timer import Timer -from pace.driver.safety_checks import SafetyChecker - -# TODO: move update_atmos_state into pace.driver +from ndsl.performance import PerformanceConfig, Timer +from ndsl.typing import Communicator +from pace.comm import CreatesCommSelector +from pace.diagnostics import DiagnosticsConfig +from pace.grid import GeneratedGridConfig, GridInitializerSelector +from pace.initialization import InitializerSelector +from pace.safety_checks import SafetyChecker +from pace.state import DriverState +from pyFV3 import DynamicalCore, DynamicalCoreConfig +from pySHiELD import Physics, PhysicsConfig from pySHiELD.update import update_atmos_state -from . import diagnostics -from .comm import CreatesCommSelector -from .grid import GeneratedGridConfig, GridInitializerSelector -from .initialization import InitializerSelector -from .state import DriverState - try: import cupy as cp @@ -103,8 +102,8 @@ class DriverConfig: type="generated", config=GeneratedGridConfig() ) ) - diagnostics_config: diagnostics.DiagnosticsConfig = dataclasses.field( - default_factory=diagnostics.DiagnosticsConfig + diagnostics_config: DiagnosticsConfig = dataclasses.field( + default_factory=DiagnosticsConfig ) performance_config: PerformanceConfig = dataclasses.field( default_factory=PerformanceConfig @@ -112,12 +111,10 @@ class DriverConfig: comm_config: CreatesCommSelector = dataclasses.field( default_factory=CreatesCommSelector ) - dycore_config: pyFV3.DynamicalCoreConfig = dataclasses.field( - default_factory=pyFV3.DynamicalCoreConfig - ) - physics_config: pySHiELD.PhysicsConfig = dataclasses.field( - default_factory=pySHiELD.PhysicsConfig + dycore_config: DynamicalCoreConfig = dataclasses.field( + default_factory=DynamicalCoreConfig ) + physics_config: PhysicsConfig = dataclasses.field(default_factory=PhysicsConfig) days: int = 0 hours: int = 0 @@ -242,14 +239,14 @@ def from_dict(cls, kwargs: Dict[str, Any]) -> "DriverConfig": ) kwargs["dycore_config"] = dacite.from_dict( - data_class=pyFV3.DynamicalCoreConfig, + data_class=DynamicalCoreConfig, data=kwargs.get("dycore_config", {}), config=dacite.Config(strict=True), ) if isinstance(kwargs["physics_config"], dict): kwargs["physics_config"] = dacite.from_dict( - data_class=pySHiELD.PhysicsConfig, + data_class=PhysicsConfig, data=kwargs.get("physics_config", {}), config=dacite.Config(strict=True), ) @@ -491,7 +488,7 @@ def exit_instead_of_build(self): self._start_time = self.config.initialization.start_time ndsl_log.info("setting up dycore object started") - self.dycore = pyFV3.DynamicalCore( + self.dycore = DynamicalCore( comm=communicator, grid_data=self.state.grid_data, stencil_factory=self.stencil_factory, @@ -506,7 +503,7 @@ def exit_instead_of_build(self): ndsl_log.info("setting up physics object started") if not config.dycore_only and not config.disable_step_physics: - self.physics = pySHiELD.Physics( + self.physics = Physics( stencil_factory=self.stencil_factory, quantity_factory=self.quantity_factory, grid_data=self.state.grid_data, diff --git a/driver/pace/driver/grid.py b/pace/grid.py similarity index 97% rename from driver/pace/driver/grid.py rename to pace/grid.py index 98eaf567..0423b793 100644 --- a/driver/pace/driver/grid.py +++ b/pace/grid.py @@ -5,28 +5,23 @@ import f90nml import xarray as xr -from ndsl.comm.communicator import Communicator +from ndsl import Namelist, QuantityFactory, ndsl_log from ndsl.comm.partitioner import get_tile_index from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM from ndsl.grid import ( + AngleGridData, + ContravariantGridData, DampingCoefficients, DriverGridData, GridData, - MetricTerms, - direct_transform, -) -from ndsl.grid.helper import ( - AngleGridData, - ContravariantGridData, HorizontalGridData, + MetricTerms, VerticalGridData, ) -from ndsl.initialization.allocator import QuantityFactory -from ndsl.logging import ndsl_log -from ndsl.namelist import Namelist +from ndsl.grid.stretch_transformation import direct_transform from ndsl.stencils.testing import TranslateGrid, grid - -from .registry import Registry +from ndsl.typing import Communicator +from pace.registry import Registry class GridInitializer(abc.ABC): diff --git a/driver/pace/driver/initialization.py b/pace/initialization.py similarity index 91% rename from driver/pace/driver/initialization.py rename to pace/initialization.py index 07c59067..68a3e639 100644 --- a/driver/pace/driver/initialization.py +++ b/pace/initialization.py @@ -7,22 +7,24 @@ import f90nml -import pyFV3 import pyFV3.initialization.analytic_init as analytic_init -import pySHiELD -from ndsl.comm.communicator import Communicator +from ndsl import ( + CompilationConfig, + DaceConfig, + Namelist, + QuantityFactory, + StencilConfig, + StencilFactory, +) from ndsl.constants import X_DIM, Y_DIM -from ndsl.dsl.dace.orchestration import DaceConfig -from ndsl.dsl.stencil import StencilConfig, StencilFactory -from ndsl.dsl.stencil_config import CompilationConfig from ndsl.grid import DampingCoefficients, DriverGridData, GridData -from ndsl.initialization.allocator import QuantityFactory -from ndsl.namelist import Namelist from ndsl.stencils.testing import TranslateGrid, grid +from ndsl.typing import Communicator +from pace.registry import Registry +from pace.state import DriverState, TendencyState, _restart_driver_state +from pyFV3 import DycoreState from pyFV3.testing import TranslateFVDynamics - -from .registry import Registry -from .state import DriverState, TendencyState, _restart_driver_state +from pySHiELD import PHYSICS_PACKAGES, PhysicsState class Initializer(abc.ABC): @@ -39,7 +41,7 @@ def get_driver_state( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ) -> DriverState: ... @@ -77,7 +79,7 @@ def get_driver_state( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ) -> DriverState: return self.config.get_driver_state( quantity_factory=quantity_factory, @@ -111,7 +113,7 @@ def get_driver_state( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ) -> DriverState: dycore_state = analytic_init.init_analytic_state( analytic_init_case=self.case, @@ -122,7 +124,7 @@ def get_driver_state( moist_phys=True, comm=communicator, ) - physics_state = pySHiELD.PhysicsState.init_zeros( + physics_state = PhysicsState.init_zeros( quantity_factory=quantity_factory, schemes=schemes ) tendency_state = TendencyState.init_zeros( @@ -155,7 +157,7 @@ def get_driver_state( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ) -> DriverState: state = _restart_driver_state( self.path, @@ -206,7 +208,7 @@ def get_driver_state( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ) -> DriverState: state = _restart_driver_state( self.path, @@ -279,14 +281,14 @@ def get_driver_state( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ) -> DriverState: backend = quantity_factory.zeros( dims=[X_DIM, Y_DIM], units="unknown" ).gt4py_backend dycore_state = self._initialize_dycore_state(communicator, backend) - physics_state = pySHiELD.PhysicsState.init_zeros( + physics_state = PhysicsState.init_zeros( quantity_factory=quantity_factory, schemes=schemes, ) @@ -305,7 +307,7 @@ def _initialize_dycore_state( self, communicator: Communicator, backend: str, - ) -> pyFV3.DycoreState: + ) -> DycoreState: grid = self._get_serialized_grid(communicator=communicator, backend=backend) ser = self._serializer(communicator) @@ -342,8 +344,8 @@ class PredefinedStateInit(Initializer): used to construct the class. """ - dycore_state: pyFV3.DycoreState - physics_state: pySHiELD.PhysicsState + dycore_state: DycoreState + physics_state: PhysicsState tendency_state: TendencyState grid_data: GridData damping_coefficients: DampingCoefficients @@ -357,7 +359,7 @@ def get_driver_state( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ) -> DriverState: return DriverState( dycore_state=self.dycore_state, diff --git a/driver/pace/driver/registry.py b/pace/registry.py similarity index 98% rename from driver/pace/driver/registry.py rename to pace/registry.py index b49d5208..2254ec9f 100644 --- a/driver/pace/driver/registry.py +++ b/pace/registry.py @@ -32,7 +32,7 @@ class Registry(Generic[T]): "my_type" and then initialize it. First we import the required symbols: >>> import dataclasses - >>> from pace.driver.registry import Registry + >>> from pace.registry import Registry Then we define a registry and register a class MyConfig: diff --git a/driver/pace/driver/run.py b/pace/run.py similarity index 97% rename from driver/pace/driver/run.py rename to pace/run.py index a78a14ef..f50ff5b3 100644 --- a/driver/pace/driver/run.py +++ b/pace/run.py @@ -6,8 +6,7 @@ import yaml from ndsl.logging import AVAILABLE_LOG_LEVELS, ndsl_log - -from .driver import Driver, DriverConfig +from pace.driver import Driver, DriverConfig @click.command() diff --git a/driver/pace/driver/safety_checks.py b/pace/safety_checks.py similarity index 97% rename from driver/pace/driver/safety_checks.py rename to pace/safety_checks.py index 58d26b30..1f583a5f 100644 --- a/driver/pace/driver/safety_checks.py +++ b/pace/safety_checks.py @@ -2,8 +2,8 @@ import numpy as np -from ndsl.quantity import Quantity -from pyFV3.dycore_state import DycoreState +from ndsl import Quantity +from pyFV3 import DycoreState class VariableBounds: diff --git a/driver/pace/driver/state.py b/pace/state.py similarity index 91% rename from driver/pace/driver/state.py rename to pace/state.py index af1acc97..9dc570b6 100644 --- a/driver/pace/driver/state.py +++ b/pace/state.py @@ -5,16 +5,14 @@ import xarray as xr import ndsl.dsl.gt4py_utils as gt_utils -import pyFV3 -import pySHiELD -from ndsl.comm.communicator import Communicator +from ndsl import Quantity, QuantityFactory, SubtileGridSizer from ndsl.constants import N_HALO_DEFAULT, X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float from ndsl.filesystem import get_fs from ndsl.grid import DampingCoefficients, DriverGridData, GridData -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from ndsl.quantity import Quantity +from ndsl.typing import Communicator +from pyFV3 import DycoreState +from pySHiELD import PHYSICS_PACKAGES, PhysicsState @dataclasses.dataclass() @@ -63,8 +61,8 @@ def init_zeros(cls, quantity_factory: QuantityFactory) -> "TendencyState": @dataclasses.dataclass class DriverState: - dycore_state: pyFV3.DycoreState - physics_state: pySHiELD.PhysicsState + dycore_state: DycoreState + physics_state: PhysicsState tendency_state: TendencyState grid_data: GridData damping_coefficients: DampingCoefficients @@ -81,7 +79,7 @@ def load_state_from_restart( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ) -> "DriverState": comm = driver_config.comm_config.get_comm() communicator = Communicator.from_layout(comm=comm, layout=driver_config.layout) @@ -155,7 +153,7 @@ def save_state(self, comm, restart_path: str = "RESTART"): def _overwrite_state_from_restart( path: str, rank: int, - state: pyFV3.DycoreState, + state: DycoreState, restart_file_prefix: str, ): """ @@ -182,7 +180,7 @@ def _restart_driver_state( damping_coefficients: DampingCoefficients, driver_grid_data: DriverGridData, grid_data: GridData, - schemes: List[pySHiELD.PHYSICS_PACKAGES], + schemes: List[PHYSICS_PACKAGES], ): fs = get_fs(path) @@ -192,11 +190,11 @@ def _restart_driver_state( ) if is_fortran_restart: - dycore_state = pyFV3.DycoreState.from_fortran_restart( + dycore_state = DycoreState.from_fortran_restart( quantity_factory=quantity_factory, communicator=communicator, path=path ) else: - dycore_state = pyFV3.DycoreState.init_zeros(quantity_factory=quantity_factory) + dycore_state = DycoreState.init_zeros(quantity_factory=quantity_factory) _overwrite_state_from_restart( path, rank, @@ -204,7 +202,7 @@ def _restart_driver_state( "restart_dycore_state", ) - physics_state = pySHiELD.PhysicsState.init_zeros( + physics_state = PhysicsState.init_zeros( quantity_factory=quantity_factory, schemes=schemes ) diff --git a/pyFV3 b/pyFV3 index 80d1fc67..67f1b5c6 160000 --- a/pyFV3 +++ b/pyFV3 @@ -1 +1 @@ -Subproject commit 80d1fc67a0d3d65c9fabd5cc6e01e5fd6b12acc2 +Subproject commit 67f1b5c6e364cdb617f71c3732851681396cea55 diff --git a/pySHiELD b/pySHiELD index b54ae3d5..69fd798e 160000 --- a/pySHiELD +++ b/pySHiELD @@ -1 +1 @@ -Subproject commit b54ae3d50e12594d111a50db435722e792ef8f6c +Subproject commit 69fd798e2f01a62f9c3df46beeff073e987b9199 diff --git a/requirements_dev.txt b/requirements_dev.txt index 19693cc4..c8bbf8a0 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -17,4 +17,4 @@ numpy>=1.15 -e NDSL -e pySHiELD -e pyFV3 --e driver +-e . diff --git a/setup.cfg b/setup.cfg index 1ff1234e..7a5ea7d5 100644 --- a/setup.cfg +++ b/setup.cfg @@ -29,6 +29,6 @@ ignore_missing_imports = True follow_imports = normal namespace_packages = True strict_optional = False -mypy_path = driver:NDSL:pySHiELD:pyFV3 +mypy_path = model:NDSL:pySHiELD:pyFV3 warn_unreachable = True explicit_package_bases = True diff --git a/driver/setup.py b/setup.py similarity index 66% rename from driver/setup.py rename to setup.py index c386d623..29c9899a 100644 --- a/driver/setup.py +++ b/setup.py @@ -1,12 +1,17 @@ +import os +from pathlib import Path from typing import List from setuptools import find_namespace_packages, setup -setup_requirements: List[str] = [] +def local_pkg(name: str, relative_path: str) -> str: + """Returns an absolute path to a local package.""" + path = f"{name} @ file://{Path(os.path.abspath(__file__)).parent / relative_path}" + return path -requirements = [ - "ndsl", + +requirements: List[str] = [ "dacite", "pyyaml", "mpi4py", @@ -16,12 +21,10 @@ "zarr", ] -test_requirements: List[str] = [] - setup( author="Allen Institute for AI", - author_email="elynnw@allenai.org", + author_email="oliver.elbert@noaa.gov", python_requires=">=3.8", classifiers=[ "Development Status :: 2 - Pre-Alpha", @@ -33,11 +36,9 @@ "Programming Language :: Python :: 3.9", ], install_requires=requirements, - setup_requires=setup_requirements, - tests_require=test_requirements, - name="pace-driver", + name="pace", license="BSD license", - packages=find_namespace_packages(include=["pace.*"]), + packages=find_namespace_packages(include=["pace", "pace.*"]), include_package_data=True, url="https://github.com/NOAA-GFDL/pace", version="0.2.0", diff --git a/tests/main/driver/test_analytic_init.py b/tests/main/driver/test_analytic_init.py index 97222db4..f2d748a7 100644 --- a/tests/main/driver/test_analytic_init.py +++ b/tests/main/driver/test_analytic_init.py @@ -4,7 +4,7 @@ import pytest import yaml -import pace.driver +from pace import DriverConfig DIR = os.path.dirname(os.path.abspath(__file__)) @@ -13,7 +13,7 @@ # need to update after TESTED_CONFIGS: List[str] = [ - "../../../driver/examples/configs/analytic_test.yaml", + "../../../examples/configs/analytic_test.yaml", ] @@ -27,5 +27,5 @@ def test_analytic_init_config(tested_configs: List[str]): for config_file in tested_configs: with open(os.path.join(DIR, config_file), "r") as f: config = yaml.safe_load(f) - driver_config = pace.driver.DriverConfig.from_dict(config) + driver_config = DriverConfig.from_dict(config) assert driver_config.initialization.type == "analytic" diff --git a/tests/main/driver/test_comm_config.py b/tests/main/driver/test_comm_config.py index beea714e..59fbece2 100644 --- a/tests/main/driver/test_comm_config.py +++ b/tests/main/driver/test_comm_config.py @@ -1,7 +1,7 @@ import dataclasses import unittest.mock -from pace.driver import CreatesComm, CreatesCommSelector, WriterCommConfig +from pace import CreatesComm, CreatesCommSelector, WriterCommConfig @CreatesCommSelector.register("mock") diff --git a/tests/main/driver/test_diagnostics.py b/tests/main/driver/test_diagnostics.py index eeccab4a..a8e4d621 100644 --- a/tests/main/driver/test_diagnostics.py +++ b/tests/main/driver/test_diagnostics.py @@ -3,8 +3,8 @@ import xarray as xr import yaml -import pace.driver -from pace.driver.run import main +from pace import DiagnosticsConfig, DriverConfig, NullCommConfig +from pace.run import main DIR = os.path.dirname(os.path.abspath(__file__)) @@ -12,15 +12,15 @@ def test_diagnostics_can_be_opened(tmpdir): with open( - os.path.join(DIR, "../../../driver/examples/configs/baroclinic_c12.yaml"), "r" + os.path.join(DIR, "../../../examples/configs/baroclinic_c12.yaml"), "r" ) as f: - driver_config = pace.driver.DriverConfig.from_dict(yaml.safe_load(f)) + driver_config = DriverConfig.from_dict(yaml.safe_load(f)) diagnostics_path = os.path.join(tmpdir, "output.zarr") - driver_config.diagnostics_config = pace.driver.DiagnosticsConfig( + driver_config.diagnostics_config = DiagnosticsConfig( path=diagnostics_path, names=["u", "v", "ua", "va", "w", "delp", "pt", "qvapor"], ) - driver_config.comm_config = pace.driver.NullCommConfig(rank=0, total_ranks=6) + driver_config.comm_config = NullCommConfig(rank=0, total_ranks=6) driver_config.dt_atmos = 60 driver_config.minutes = 1 main(driver_config) diff --git a/tests/main/driver/test_diagnostics_config.py b/tests/main/driver/test_diagnostics_config.py index 5e409c88..9b96c1be 100644 --- a/tests/main/driver/test_diagnostics_config.py +++ b/tests/main/driver/test_diagnostics_config.py @@ -2,45 +2,44 @@ import pytest -import pace.driver -import pace.driver.diagnostics -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from pyFV3.dycore_state import DycoreState +from ndsl import QuantityFactory, SubtileGridSizer +from pace import DiagnosticsConfig +from pace.diagnostics import MonitorDiagnostics, NullDiagnostics, ZSelect +from pyFV3 import DycoreState def test_returns_null_diagnostics_if_no_path_given(): - config = pace.driver.DiagnosticsConfig(path=None, names=[]) + config = DiagnosticsConfig(path=None, names=[]) assert isinstance( config.diagnostics_factory(unittest.mock.MagicMock()), - pace.driver.diagnostics.NullDiagnostics, + NullDiagnostics, ) def test_returns_monitor_diagnostics_if_path_given(tmpdir): - config = pace.driver.DiagnosticsConfig( + config = DiagnosticsConfig( path=tmpdir, names=["foo"], derived_names=["bar"], - z_select=[pace.driver.diagnostics.ZSelect(level=0, names=["foo"])], + z_select=[ZSelect(level=0, names=["foo"])], ) result = config.diagnostics_factory(unittest.mock.MagicMock()) - assert isinstance(result, pace.driver.diagnostics.MonitorDiagnostics) + assert isinstance(result, MonitorDiagnostics) def test_raises_if_names_given_but_no_path(): with pytest.raises(ValueError): - pace.driver.DiagnosticsConfig(path=None, names=["foo"]) + DiagnosticsConfig(path=None, names=["foo"]) with pytest.raises(ValueError): - pace.driver.DiagnosticsConfig(path=None, derived_names=["foo"]) + DiagnosticsConfig(path=None, derived_names=["foo"]) def test_zselect_raises_error_if_not_3d(tmpdir): with pytest.raises(AssertionError): - config = pace.driver.DiagnosticsConfig( + config = DiagnosticsConfig( path=tmpdir, - z_select=[pace.driver.diagnostics.ZSelect(level=0, names=["phis"])], + z_select=[ZSelect(level=0, names=["phis"])], ) result = config.diagnostics_factory(unittest.mock.MagicMock()) quantity_factory = QuantityFactory.from_backend( @@ -53,9 +52,9 @@ def test_zselect_raises_error_if_not_3d(tmpdir): def test_zselect_raises_error_if_3rd_dim_not_z(tmpdir): with pytest.raises(ValueError): - config = pace.driver.DiagnosticsConfig( + config = DiagnosticsConfig( path=tmpdir, - z_select=[pace.driver.diagnostics.ZSelect(level=0, names=["foo"])], + z_select=[ZSelect(level=0, names=["foo"])], ) result = config.diagnostics_factory(unittest.mock.MagicMock()) quantity_factory = QuantityFactory.from_backend( diff --git a/tests/main/driver/test_docs.py b/tests/main/driver/test_docs.py index 80d4abfd..aec03fcb 100644 --- a/tests/main/driver/test_docs.py +++ b/tests/main/driver/test_docs.py @@ -1,9 +1,9 @@ import doctest -import pace.driver.registry +import pace.registry def test_registry_doc_examples(): - result = doctest.testmod(pace.driver.registry) + result = doctest.testmod(pace.registry) assert result.attempted > 0, "No doctests found" assert result.failed == 0, "doctests failed" diff --git a/tests/main/driver/test_driver.py b/tests/main/driver/test_driver.py index 602624cf..d5020809 100644 --- a/tests/main/driver/test_driver.py +++ b/tests/main/driver/test_driver.py @@ -4,15 +4,14 @@ import pytest -from ndsl.comm.null_comm import NullComm -from ndsl.dsl.stencil import StencilConfig +from ndsl import NullComm, StencilConfig from ndsl.performance.report import ( TimeReport, gather_hit_counts, gather_timing_data, get_sypd, ) -from pace.driver import CreatesCommSelector, DriverConfig, NullCommConfig +from pace import CreatesCommSelector, DriverConfig, NullCommConfig def get_driver_config( diff --git a/tests/main/driver/test_example_configs.py b/tests/main/driver/test_example_configs.py index 2e4b5d3e..d99130f5 100644 --- a/tests/main/driver/test_example_configs.py +++ b/tests/main/driver/test_example_configs.py @@ -4,12 +4,12 @@ import pytest import yaml -import pace.driver +from pace import DriverConfig dirname = os.path.dirname(os.path.abspath(__file__)) -EXAMPLE_CONFIGS_DIR = os.path.join(dirname, "../../../driver/examples/configs/") +EXAMPLE_CONFIGS_DIR = os.path.join(dirname, "../../../examples/configs/") TESTED_CONFIGS: List[str] = [ "baroclinic_c12.yaml", @@ -88,5 +88,5 @@ def test_all_configs_tested_or_excluded( def test_example_config_can_initialize(path: str, file_list: List[str]): for file_name in file_list: with open(os.path.join(path, file_name), "r") as f: - config = pace.driver.DriverConfig.from_dict(yaml.safe_load(f)) - assert isinstance(config, pace.driver.DriverConfig) + config = DriverConfig.from_dict(yaml.safe_load(f)) + assert isinstance(config, DriverConfig) diff --git a/tests/main/driver/test_restart_fortran.py b/tests/main/driver/test_restart_fortran.py index 49ef5f17..3821d20b 100644 --- a/tests/main/driver/test_restart_fortran.py +++ b/tests/main/driver/test_restart_fortran.py @@ -3,14 +3,16 @@ import numpy as np import xarray as xr -import pace.driver -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.local_comm import LocalComm -from ndsl.comm.null_comm import NullComm -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from pace.driver.initialization import FortranRestartInit +from ndsl import ( + CubedSphereCommunicator, + CubedSpherePartitioner, + LocalComm, + NullComm, + QuantityFactory, + SubtileGridSizer, + TilePartitioner, +) +from pace import FortranRestartInit, GeneratedGridConfig from pySHiELD import PHYSICS_PACKAGES @@ -42,15 +44,9 @@ def test_state_from_fortran_restart(): quantity_factory = QuantityFactory.from_backend(sizer=sizer, backend="numpy") restart_dir = os.path.join(PACE_DIR, "tests/main/data/c12_restart") - ( - damping_coefficients, - driver_grid_data, - grid_data, - ) = pace.driver.GeneratedGridConfig( + (damping_coefficients, driver_grid_data, grid_data,) = GeneratedGridConfig( restart_path=restart_dir, eta_file=restart_dir + "/fv_core.res.nc" - ).get_grid( - quantity_factory, null_communicator - ) + ).get_grid(quantity_factory, null_communicator) restart_config = FortranRestartInit(path=restart_dir) driver_state = restart_config.get_driver_state( diff --git a/tests/main/driver/test_restart_serial.py b/tests/main/driver/test_restart_serial.py index c55d9661..144abab1 100644 --- a/tests/main/driver/test_restart_serial.py +++ b/tests/main/driver/test_restart_serial.py @@ -6,16 +6,22 @@ import xarray as xr import yaml -import pace.driver -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.null_comm import NullComm -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from ndsl.quantity import Quantity -from pace.driver import CreatesComm, DriverConfig -from pace.driver.driver import RestartConfig -from pace.driver.initialization import AnalyticInit +from ndsl import ( + CubedSphereCommunicator, + CubedSpherePartitioner, + NullComm, + Quantity, + QuantityFactory, + SubtileGridSizer, + TilePartitioner, +) +from pace import ( + AnalyticInit, + CreatesComm, + DriverConfig, + GeneratedGridConfig, + RestartConfig, +) from pySHiELD import PHYSICS_PACKAGES @@ -47,7 +53,7 @@ def test_restart_save_to_disk(): with open( os.path.join( DIR, - "../../../driver/examples/configs/baroclinic_c12_write_restart.yaml", + "../../../examples/configs/baroclinic_c12_write_restart.yaml", ), "r", ) as f: @@ -69,13 +75,9 @@ def test_restart_save_to_disk(): quantity_factory = QuantityFactory.from_backend(sizer=sizer, backend=backend) eta_file = driver_config.grid_config.config.eta_file - ( - damping_coefficients, - driver_grid_data, - grid_data, - ) = pace.driver.GeneratedGridConfig(eta_file=eta_file).get_grid( - quantity_factory, communicator - ) + (damping_coefficients, driver_grid_data, grid_data,) = GeneratedGridConfig( + eta_file=eta_file + ).get_grid(quantity_factory, communicator) init = AnalyticInit() driver_state = init.get_driver_state( quantity_factory=quantity_factory, diff --git a/tests/main/driver/test_safety_checks.py b/tests/main/driver/test_safety_checks.py index 5f528c55..77560b88 100644 --- a/tests/main/driver/test_safety_checks.py +++ b/tests/main/driver/test_safety_checks.py @@ -3,8 +3,8 @@ import numpy as np import pytest -from ndsl.quantity import Quantity -from pace.driver.safety_checks import SafetyChecker +from ndsl import Quantity +from pace import SafetyChecker def test_register_variable(): diff --git a/tests/main/fv3core/test_cartesian_grid.py b/tests/main/fv3core/test_cartesian_grid.py index 528735c1..0c5b2f62 100644 --- a/tests/main/fv3core/test_cartesian_grid.py +++ b/tests/main/fv3core/test_cartesian_grid.py @@ -1,11 +1,9 @@ import numpy as np import pytest -from ndsl.comm.communicator import TileCommunicator -from ndsl.comm.null_comm import NullComm -from ndsl.comm.partitioner import TilePartitioner +from ndsl import NullComm, TileCommunicator, TilePartitioner from ndsl.constants import PI -from ndsl.grid.generation import MetricTerms +from ndsl.grid import MetricTerms @pytest.mark.parametrize("npx", [8]) diff --git a/tests/main/fv3core/test_dycore_call.py b/tests/main/fv3core/test_dycore_call.py index 9ef22548..eeb3d08b 100644 --- a/tests/main/fv3core/test_dycore_call.py +++ b/tests/main/fv3core/test_dycore_call.py @@ -4,30 +4,33 @@ from datetime import timedelta from typing import Tuple -import ndsl.dsl.stencil -import ndsl.stencils.testing -import pyFV3 import pyFV3.initialization.analytic_init as ai -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.null_comm import NullComm -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner -from ndsl.dsl.dace.dace_config import DaceConfig -from ndsl.dsl.stencil import GridIndexing +from ndsl import ( + CompilationConfig, + CubedSphereCommunicator, + CubedSpherePartitioner, + DaceConfig, + GridIndexing, + NullComm, + Quantity, + QuantityFactory, + StencilConfig, + StencilFactory, + SubtileGridSizer, + TilePartitioner, +) from ndsl.grid import DampingCoefficients, GridData, MetricTerms -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer from ndsl.performance.timer import NullTimer, Timer -from ndsl.quantity import Quantity from ndsl.stencils.testing import assert_same_temporaries, copy_temporaries -from pyFV3.dycore_state import DycoreState +from pyFV3 import DycoreState, DynamicalCore, DynamicalCoreConfig DIR = os.path.abspath(os.path.dirname(__file__)) -def setup_dycore() -> Tuple[pyFV3.DynamicalCore, pyFV3.DycoreState, Timer]: +def setup_dycore() -> Tuple[DynamicalCore, DycoreState, Timer]: backend = "numpy" - config = pyFV3.DynamicalCoreConfig( + config = DynamicalCoreConfig( layout=(1, 1), npx=13, npy=13, @@ -77,8 +80,8 @@ def setup_dycore() -> Tuple[pyFV3.DynamicalCore, pyFV3.DycoreState, Timer]: partitioner = CubedSpherePartitioner(TilePartitioner(config.layout)) communicator = CubedSphereCommunicator(mpi_comm, partitioner) dace_config = DaceConfig(communicator=communicator, backend=backend) - stencil_config = ndsl.dsl.stencil.StencilConfig( - compilation_config=ndsl.dsl.stencil.CompilationConfig( + stencil_config = StencilConfig( + compilation_config=CompilationConfig( backend=backend, rebuild=False, validate_args=True ), dace_config=dace_config, @@ -116,12 +119,12 @@ def setup_dycore() -> Tuple[pyFV3.DynamicalCore, pyFV3.DycoreState, Timer]: moist_phys=config.moist_phys, comm=communicator, ) - stencil_factory = ndsl.dsl.stencil.StencilFactory( + stencil_factory = StencilFactory( config=stencil_config, grid_indexing=grid_indexing, ) - dycore = pyFV3.DynamicalCore( + dycore = DynamicalCore( comm=communicator, grid_data=grid_data, stencil_factory=stencil_factory, diff --git a/tests/main/fv3core/test_grid.py b/tests/main/fv3core/test_grid.py index 2286ea03..0b95f3ea 100644 --- a/tests/main/fv3core/test_grid.py +++ b/tests/main/fv3core/test_grid.py @@ -4,6 +4,7 @@ import pytest from gt4py.cartesian import gtscript +from ndsl import GridIndexing from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, @@ -12,7 +13,6 @@ Z_DIM, Z_INTERFACE_DIM, ) -from ndsl.dsl.stencil import GridIndexing from ndsl.dsl.typing import Index3D diff --git a/tests/main/fv3core/test_init_from_geos.py b/tests/main/fv3core/test_init_from_geos.py index d620d9fa..6d183561 100644 --- a/tests/main/fv3core/test_init_from_geos.py +++ b/tests/main/fv3core/test_init_from_geos.py @@ -2,8 +2,9 @@ import numpy as np import pytest # noqa -import pyFV3 -from ndsl.comm.null_comm import NullComm +from ndsl import NullComm +from pyFV3 import DynamicalCore +from pyFV3.wrappers import GeosDycoreWrapper def test_geos_wrapper(): @@ -84,7 +85,7 @@ def test_geos_wrapper(): comm = NullComm(rank=0, total_ranks=6, fill_value=0.0) backend = "numpy" - wrapper = pyFV3.GeosDycoreWrapper( + wrapper = GeosDycoreWrapper( namelist=namelist, comm=comm, backend=backend, @@ -122,8 +123,8 @@ def test_geos_wrapper(): 7, ) - assert isinstance(wrapper, pyFV3.GeosDycoreWrapper) - assert isinstance(wrapper.dynamical_core, pyFV3.DynamicalCore) + assert isinstance(wrapper, GeosDycoreWrapper) + assert isinstance(wrapper.dynamical_core, DynamicalCore) u = np.ones(shape_y_interface) v = np.ones(shape_x_interface) diff --git a/tests/main/grid/test_eta.py b/tests/main/grid/test_eta.py index 200720d7..f21c8896 100755 --- a/tests/main/grid/test_eta.py +++ b/tests/main/grid/test_eta.py @@ -7,7 +7,7 @@ import xarray as xr import yaml -import pace.driver +from pace import Driver, DriverConfig, NullCommConfig """ @@ -59,9 +59,7 @@ def test_set_hybrid_pressure_coefficients_correct(km): """ dirname = os.path.dirname(os.path.abspath(__file__)) - config_file = os.path.join( - dirname, "../../../driver/examples/configs/baroclinic_c12.yaml" - ) + config_file = os.path.join(dirname, "../../../examples/configs/baroclinic_c12.yaml") with open(config_file, "r") as f: yaml_config = yaml.safe_load(f) @@ -69,9 +67,9 @@ def test_set_hybrid_pressure_coefficients_correct(km): 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) + driver_config = DriverConfig.from_dict(yaml_config) + driver_config.comm_config = NullCommConfig(rank=0, total_ranks=6) + driver = Driver(config=driver_config) p_results = driver.state.grid_data.p.data ak_results = driver.state.grid_data.ak.data @@ -103,9 +101,7 @@ def test_set_hybrid_pressure_coefficients_nofile(): """ dirname = os.path.dirname(os.path.abspath(__file__)) - config_file = os.path.join( - dirname, "../../../driver/examples/configs/baroclinic_c12.yaml" - ) + config_file = os.path.join(dirname, "../../../examples/configs/baroclinic_c12.yaml") with open(config_file, "r") as f: yaml_config = yaml.safe_load(f) @@ -113,9 +109,9 @@ def test_set_hybrid_pressure_coefficients_nofile(): del yaml_config["grid_config"]["config"]["eta_file"] try: - 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) + driver_config = DriverConfig.from_dict(yaml_config) + driver_config.comm_config = NullCommConfig(rank=0, total_ranks=6) + driver = Driver(config=driver_config) except Exception as error: if str(error) == "eta file not specified": pytest.xfail("testing eta file not specified") @@ -135,9 +131,7 @@ def test_set_hybrid_pressure_coefficients_not_mono(): """ dirname = os.path.dirname(os.path.abspath(__file__)) - config_file = os.path.join( - dirname, "../../../driver/examples/configs/baroclinic_c12.yaml" - ) + config_file = os.path.join(dirname, "../../../examples/configs/baroclinic_c12.yaml") with open(config_file, "r") as f: yaml_config = yaml.safe_load(f) @@ -148,9 +142,9 @@ def test_set_hybrid_pressure_coefficients_not_mono(): yaml_config["grid_config"]["config"]["eta_file"] = out_eta_file try: - 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) + driver_config = DriverConfig.from_dict(yaml_config) + driver_config.comm_config = NullCommConfig(rank=0, total_ranks=6) + driver = Driver(config=driver_config) except Exception as error: if os.path.isfile(out_eta_file): os.remove(out_eta_file) diff --git a/tests/main/physics/test_integration.py b/tests/main/physics/test_integration.py index 0d3d534e..dfd38af1 100644 --- a/tests/main/physics/test_integration.py +++ b/tests/main/physics/test_integration.py @@ -3,18 +3,23 @@ import numpy as np -import pySHiELD -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.null_comm import NullComm -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner -from ndsl.dsl.dace import DaceConfig -from ndsl.dsl.dace.orchestration import DaCeOrchestration -from ndsl.dsl.stencil import GridIndexing, StencilConfig, StencilFactory -from ndsl.dsl.stencil_config import CompilationConfig +from ndsl import ( + CompilationConfig, + CubedSphereCommunicator, + CubedSpherePartitioner, + DaceConfig, + DaCeOrchestration, + GridIndexing, + NullComm, + QuantityFactory, + StencilConfig, + StencilFactory, + SubtileGridSizer, + TilePartitioner, +) from ndsl.grid import GridData, MetricTerms -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer from ndsl.stencils.testing import assert_same_temporaries, copy_temporaries +from pySHiELD import PHYSICS_PACKAGES, Physics, PhysicsConfig, PhysicsState try: @@ -26,7 +31,7 @@ def setup_physics(): backend = "numpy" layout = (1, 1) - physics_config = pySHiELD.PhysicsConfig( + physics_config = PhysicsConfig( dt_atmos=225, hydrostatic=False, npx=13, npy=13, npz=79, nwat=6, do_qa=True ) mpi_comm = NullComm(rank=0, total_ranks=6 * layout[0] * layout[1], fill_value=0.0) @@ -69,17 +74,17 @@ def setup_physics(): eta_file="tests/main/input/eta79.nc", ) grid_data = GridData.new_from_metric_terms(metric_terms) - physics = pySHiELD.Physics( + physics = Physics( stencil_factory, quantity_factory, grid_data, physics_config, ) - physics_state = pySHiELD.PhysicsState.init_zeros( - quantity_factory, schemes=[pySHiELD.PHYSICS_PACKAGES["GFS_microphysics"]] + physics_state = PhysicsState.init_zeros( + quantity_factory, schemes=[PHYSICS_PACKAGES["GFS_microphysics"]] ) random = np.random.RandomState(0) - for field in fields(pySHiELD.PhysicsState): + for field in fields(PhysicsState): array = getattr(physics_state, field.name) # check that it's a storage this way, because Field is not a class if isinstance(array, (np.ndarray, cp.ndarray)): diff --git a/tests/main/test_grid_init.py b/tests/main/test_grid_init.py index d9fb6b5d..9489cd3b 100644 --- a/tests/main/test_grid_init.py +++ b/tests/main/test_grid_init.py @@ -1,13 +1,16 @@ import numpy as np import pytest -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.null_comm import NullComm -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner +from ndsl import ( + CubedSphereCommunicator, + CubedSpherePartitioner, + NullComm, + Quantity, + QuantityFactory, + SubtileGridSizer, + TilePartitioner, +) from ndsl.grid import MetricTerms -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from ndsl.quantity import Quantity def get_cube_comm(layout, rank: int): diff --git a/tests/mpi_54rank/test_ext_grid/test_external_grid.py b/tests/mpi/ext_grid/test_external_grid.py similarity index 95% rename from tests/mpi_54rank/test_ext_grid/test_external_grid.py rename to tests/mpi/ext_grid/test_external_grid.py index c5b52014..d4dbc0e5 100644 --- a/tests/mpi_54rank/test_ext_grid/test_external_grid.py +++ b/tests/mpi/ext_grid/test_external_grid.py @@ -6,19 +6,19 @@ import xarray as xr import yaml -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.mpi import MPIComm -from ndsl.comm.partitioner import ( +from ndsl import ( + CubedSphereCommunicator, CubedSpherePartitioner, + MPIComm, TilePartitioner, - get_tile_number, ) +from ndsl.comm.partitioner import get_tile_number from ndsl.constants import PI, RADIUS, X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM -from pace.driver import Driver, DriverConfig +from pace import Driver, DriverConfig DIR = os.path.dirname(os.path.abspath(__file__)) -TEST_CONFIGS_DIR = os.path.join(DIR, "../../../driver/examples/configs/") +TEST_CONFIGS_DIR = os.path.join(DIR, "../../../examples/configs/") TEST_DATA_DIR = os.path.join(DIR, "../../../test_input/") TEST_CONFIG_FILE_RANKS = [ diff --git a/driver/tests/mpi/run_save_and_load_restart.sh b/tests/mpi/restart/run_save_and_load_restart.sh similarity index 66% rename from driver/tests/mpi/run_save_and_load_restart.sh rename to tests/mpi/restart/run_save_and_load_restart.sh index 7d975142..76829be8 100755 --- a/driver/tests/mpi/run_save_and_load_restart.sh +++ b/tests/mpi/restart/run_save_and_load_restart.sh @@ -7,6 +7,6 @@ sed -i.bak 's/seconds: 225/seconds: 450/' baroclinic_c12_run_two_steps.yaml sed -i.bak 's/save_restart: true/save_restart: false/' baroclinic_c12_run_two_steps.yaml sed -i.bak 's/path: "output.zarr"/path: "run_two_steps_output.zarr"/' baroclinic_c12_run_two_steps.yaml rm *.bak -$MPIRUN_CALL python -m pace.driver.run baroclinic_c12_write_restart.yaml --log-level=ERROR -$MPIRUN_CALL python -m pace.driver.run RESTART/restart.yaml --log-level=ERROR -$MPIRUN_CALL python -m pace.driver.run baroclinic_c12_run_two_steps.yaml --log-level=ERROR +$MPIRUN_CALL python -m pace.run baroclinic_c12_write_restart.yaml --log-level=ERROR +$MPIRUN_CALL python -m pace.run RESTART/restart.yaml --log-level=ERROR +$MPIRUN_CALL python -m pace.run baroclinic_c12_run_two_steps.yaml --log-level=ERROR diff --git a/driver/tests/mpi/test_restart.py b/tests/mpi/restart/test_restart.py similarity index 89% rename from driver/tests/mpi/test_restart.py rename to tests/mpi/restart/test_restart.py index 2ac12807..2c812d87 100644 --- a/driver/tests/mpi/test_restart.py +++ b/tests/mpi/restart/test_restart.py @@ -7,12 +7,14 @@ import zarr from mpi4py import MPI -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.null_comm import NullComm -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner -from ndsl.quantity import Quantity -from pace.driver import DriverConfig -from pace.driver.state import DriverState +from ndsl import ( + CubedSphereCommunicator, + CubedSpherePartitioner, + NullComm, + Quantity, + TilePartitioner, +) +from pace import DriverConfig, DriverState from pySHiELD import PHYSICS_PACKAGES diff --git a/tests/mpi_54rank/test_grid_init.py b/tests/mpi/test_grid_init.py similarity index 95% rename from tests/mpi_54rank/test_grid_init.py rename to tests/mpi/test_grid_init.py index c5fbfaed..2b647dec 100644 --- a/tests/mpi_54rank/test_grid_init.py +++ b/tests/mpi/test_grid_init.py @@ -2,15 +2,18 @@ import numpy as np -import pyFV3 -from ndsl.comm.communicator import CubedSphereCommunicator, TileCommunicator -from ndsl.comm.mpi import MPIComm -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner -from ndsl.grid import MetricTerms -from ndsl.grid.helper import GridData -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from ndsl.quantity import Quantity +from ndsl import ( + CubedSphereCommunicator, + CubedSpherePartitioner, + MPIComm, + Quantity, + QuantityFactory, + SubtileGridSizer, + TileCommunicator, + TilePartitioner, +) +from ndsl.grid import GridData, MetricTerms +from pyFV3 import DycoreState from pyFV3.initialization.test_cases.initialize_baroclinic import init_baroclinic_state @@ -121,7 +124,7 @@ def metric_terms_to_quantity_dict(metric_terms: MetricTerms) -> Dict[str, Quanti def dycore_state_to_quantity_dict( - dycore_state: pyFV3.DycoreState, + dycore_state: DycoreState, ) -> Dict[str, Quantity]: return { "u": dycore_state.u, diff --git a/tests/savepoint/test_checkpoints.py b/tests/savepoint/test_checkpoints.py index 27aafafd..28388d0e 100644 --- a/tests/savepoint/test_checkpoints.py +++ b/tests/savepoint/test_checkpoints.py @@ -8,32 +8,31 @@ import xarray as xr import yaml -import pyFV3 -from ndsl.checkpointer import ValidationCheckpointer -from ndsl.checkpointer.thresholds import ( - SavepointThresholds, - Threshold, - ThresholdCalibrationCheckpointer, -) -from ndsl.comm.communicator import CubedSphereCommunicator -from ndsl.comm.mpi import MPIComm -from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner -from ndsl.dsl.stencil import ( +from ndsl import ( CompilationConfig, + CubedSphereCommunicator, + CubedSpherePartitioner, GridIndexing, + MPIComm, + Namelist, + Quantity, + QuantityFactory, StencilConfig, StencilFactory, + SubtileGridSizer, + TilePartitioner, +) +from ndsl.checkpointer import ( + SavepointThresholds, + Threshold, + ThresholdCalibrationCheckpointer, + ValidationCheckpointer, ) from ndsl.grid import DampingCoefficients, GridData -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from ndsl.namelist import Namelist -from ndsl.quantity import Quantity -from ndsl.stencils.testing import TranslateGrid, dataset_to_dict -from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing import Grid, TranslateGrid, dataset_to_dict from ndsl.testing import perturb -from pyFV3.initialization.dycore_state import DycoreState -from pyFV3.testing.translate_fvdynamics import TranslateFVDynamics +from pyFV3 import DycoreState, DynamicalCore, DynamicalCoreConfig +from pyFV3.testing import TranslateFVDynamics def get_grid(data_path: str, rank: int, layout: Tuple[int, int], backend: str) -> Grid: @@ -110,7 +109,7 @@ def test_fv_dynamics( ds = xr.open_dataset(os.path.join(data_path, "FVDynamics-In.nc")).sel( savepoint=0, rank=communicator.rank ) - dycore_config = pyFV3.DynamicalCoreConfig.from_namelist(namelist) + dycore_config = DynamicalCoreConfig.from_namelist(namelist) initializer = StateInitializer( ds, translate, @@ -142,7 +141,7 @@ def test_fv_dynamics( savepoint_data_path=data_path, thresholds=thresholds, rank=communicator.rank ) state, grid_data = initializer.new_state() - dycore = pyFV3.DynamicalCore( + dycore = DynamicalCore( comm=communicator, grid_data=grid_data, stencil_factory=stencil_factory, @@ -164,7 +163,7 @@ def _calibrate_thresholds( stencil_factory: StencilFactory, quantity_factory: QuantityFactory, damping_coefficients: DampingCoefficients, - dycore_config: pyFV3.DynamicalCoreConfig, + dycore_config: DynamicalCoreConfig, n_trials: int, factor: float, ): @@ -175,7 +174,7 @@ def _calibrate_thresholds( perturb(dycore_state_to_dict(trial_state)) # we need to initialize new DynamicalCore because halo updates bind # to a particular state object, currently - dycore = pyFV3.DynamicalCore( + dycore = DynamicalCore( comm=communicator, grid_data=grid_data, stencil_factory=stencil_factory, diff --git a/tests/savepoint/translate/translate_driver.py b/tests/savepoint/translate/translate_driver.py index 446ed08d..43fd345e 100644 --- a/tests/savepoint/translate/translate_driver.py +++ b/tests/savepoint/translate/translate_driver.py @@ -1,14 +1,11 @@ +from ndsl import Namelist, QuantityFactory, SubtileGridSizer from ndsl.constants import N_HALO_DEFAULT -from ndsl.initialization.allocator import QuantityFactory -from ndsl.initialization.sizer import SubtileGridSizer -from ndsl.namelist import Namelist -from pace.driver.run import Driver, DriverConfig -from pace.driver.state import TendencyState -from pyFV3._config import DynamicalCoreConfig +from pace import Driver, DriverConfig, TendencyState +from pyFV3 import DynamicalCoreConfig # TODO physics should not depend on pyFV3 # but also, driver tests should not be in physics -from pyFV3.testing.translate_fvdynamics import TranslateFVDynamics +from pyFV3.testing import TranslateFVDynamics from pyFV3.testing.validation import enable_selective_validation from pySHiELD import PHYSICS_PACKAGES, PhysicsConfig, PhysicsState