Initial infrastructure

.circleci/config.yml

@@ -0,0 +1,28 @@
+version: 2.1
+orbs:
+  gcp-gcr: circleci/[email protected]
+  gh: circleci/[email protected]
+jobs:
+  build:
+    machine:
+      image: ubuntu-2004:202111-02
+    environment:
+      FSSPEC_GS_REQUESTER_PAYS: vcm-ml
+      GOOGLE_APPLICATION_CREDENTIALS: /tmp/key.json
+    steps:
+      - checkout
+      - run: make update_submodules
+      - run:
+          name: "gcloud auth"
+          command: |
+            echo $ENCODED_GOOGLE_CREDENTIALS | base64 -d > $GOOGLE_APPLICATION_CREDENTIALS
+            echo "export GCLOUD_SERVICE_KEY=\$(echo \$ENCODED_GOOGLE_CREDENTIALS | base64 --decode)" >> $BASH_ENV
+      - gcp-gcr/gcr-auth
+      - run: make build
+      - run: make test
+
+workflows:
+  version: 2
+  build:
+    jobs:
+      - build

.gitmodules

@@ -0,0 +1,18 @@
+[submodule "submodules/SHiELD_build"]
+	path = submodules/SHiELD_build
+	url = https://github.com/NOAA-GFDL/SHiELD_build.git
+[submodule "submodules/SHiELD_SRC/FMS"]
+	path = submodules/SHiELD_SRC/FMS
+	url = https://github.com/NOAA-GFDL/FMS.git
+[submodule "submodules/SHiELD_SRC/FMSCoupler"]
+	path = submodules/SHiELD_SRC/FMSCoupler
+	url = https://github.com/NOAA-GFDL/FMSCoupler.git
+[submodule "submodules/SHiELD_SRC/GFDL_atmos_cubed_sphere"]
+	path = submodules/SHiELD_SRC/GFDL_atmos_cubed_sphere
+	url = https://github.com/NOAA-GFDL/GFDL_atmos_cubed_sphere.git
+[submodule "submodules/SHiELD_SRC/SHiELD_physics"]
+	path = submodules/SHiELD_SRC/SHiELD_physics
+	url = https://github.com/NOAA-GFDL/SHiELD_physics.git
+[submodule "submodules/SHiELD_SRC/atmos_drivers"]
+	path = submodules/SHiELD_SRC/atmos_drivers
+	url = https://github.com/NOAA-GFDL/atmos_drivers.git

Makefile

@@ -0,0 +1,25 @@
+VERSION ?= $(shell git rev-parse HEAD)
+
+build:
+	docker build -f docker/Dockerfile -t shield-wrapper:$(VERSION) .
+
+test:
+	docker/docker_run.sh \
+		-v $(shell pwd)/wrapper/tests:/wrapper/tests \
+		shield-wrapper:$(VERSION) make -C /wrapper test
+
+test_regtest_reset:
+	docker/docker_run.sh \
+		-v $(shell pwd)/wrapper/tests:/wrapper/tests \
+		shield-wrapper:$(VERSION) make -C /wrapper test_regtest_reset
+
+enter:
+	docker/docker_run.sh \
+		-v $(shell pwd)/wrapper/tests:/wrapper/tests \
+		-it shield-wrapper:$(VERSION) bash
+
+lock_pip:
+	pip-compile requirements.in
+
+update_submodules:
+	git submodule update --init --recursive

README.md

@@ -0,0 +1,51 @@
+# SHiELD-wrapper
+
+This repository contains code for a Python-wrapped version of the public
+version of GFDL's SHiELD model.  It aims to be a drop-in replacement for
+[AI2's existing Python-wrapped version of
+FV3GFS](https://github.com/ai2cm/fv3gfs-fortran/tree/master/FV3/wrapper), which
+was described in [McGibbon et
+al. (2021)](https://gmd.copernicus.org/articles/14/4401/2021/gmd-14-4401-2021.html),
+and used in several published hybrid machine learning projects.  Much of the
+code is identical or nearly identical to this original wrapper; the main
+difference is that the wrapper in this repository wraps the latest public
+version of GFDL's SHiELD model instead of an old fork of FV3GFS.  At this
+moment some features lag behind, because they require changes to the fortran
+code of SHiELD.  These features will be added incrementally as the required
+changes are made to SHiELD.
+
+## Development
+
+Development and testing of SHiELD-wrapper is currently only supported within a
+Docker container.  Similar to how things are done in AI2's fv3net repository,
+tests are set up to be run from within the Docker image defined in the
+Dockerfile within this repository.  To build the docker image locally, make
+sure to first initialize all the submodules, and then run `make build`:
+
+```
+$ make update_submodules
+$ make build
+```
+
+Once the Docker image has been built, you can enter it with the wrapper testing
+code editable from outside the container with:
+
+```
+$ make enter
+```
+
+To run the tests, you can use the following:
+
+```
+$ make test
+```
+
+To reset the regression test checksums, use:
+
+```
+$ make test_regtest_reset
+```
+
+The Python dependencies used in building and testing the wrapper are defined in
+the `requirements.in` file.  The full set of pinned dependencies these resolve
+into are defined in `requirements.txt`, which is produced using `pip-compile`.

docker/Dockerfile

@@ -0,0 +1,99 @@
+# Use identical base image (20.04) to what we use in the fv3net prognostic run:
+# https://github.com/ai2cm/fv3net/blob/382409f76fd26d70457b349a48fa0e63f4f26303/Makefile#L12
+#
+# This is convenient for a couple reasons:
+#
+# * It installs Python 3.8 as the system Python, which should allow us to
+#   seamlessly install the additional Python packages used in the existing
+#   prognostic_run image, to their exact versions.
+#
+# * It installs a version of MPICH that does not enable link time optimization
+#   (LTO) by default.  Newer Ubuntu LTS base images, e.g. 22.04, install
+#   versions of MPICH that enable link time optimization by default in their
+#   compiler wrappers (i.e. mpif90 and the like; see the presence of -flto=auto
+#   and other LTO-related flags in the output of `mpif90 -show`).  LTO can lead
+#   to differences in answers between applications that were linked differently
+#   (e.g. the fortran executable and the Python-wrapped model). Compiling
+#   without LTO allows the fortran executable and the Python wrapper to
+#   robustly produce bitwise identical results.  See
+#   https://github.com/ai2cm/SHiELD-minimal/pull/6 for some details on MPICH in
+#   the Ubuntu LTS 22.04 image.  This can be addressed by manually adding
+#   `-fno-lto` compiler flags to Makefile templates, but it is easiest not to 
+#   need to worry about this.
+FROM ubuntu@sha256:9101220a875cee98b016668342c489ff0674f247f6ca20dfc91b91c0f28581ae
+
+# Install and use GNU version 10 compiliers instead of the default version 9,
+# since the latest version of FMS fails to build with GNU version 9
+# compilers.
+RUN apt-get update && \
+    DEBIAN_FRONTEND="noninteractive" apt-get -y install \
+    cmake \
+    gcc-10 \
+    gfortran-10 \
+    g++-10 \
+    git \
+    libmpich-dev \
+    libnetcdf-dev \
+    libnetcdff-dev \
+    m4 \
+    make \
+    mpich \
+    pkg-config \
+    python3 \
+    python3-dev \
+    python3-pip \
+    python3-setuptools
+
+# GNU version 9 compilers are unavoidably installed by default; using
+# update-alternatives ensures that gcc, gfortran, and g++ point to GNU version
+# 10 compilers, including within MPI-wrapped compilers.
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 1
+RUN update-alternatives --install /usr/bin/gfortran gfortran /usr/bin/gfortran-10 1
+RUN update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-10 1
+
+ENV SHiELD_FC=mpif90
+ENV SHiELD_CC=mpicc
+ENV SHiELD_CXX=mpicxx
+ENV SHiELD_LD=mpif90
+
+ENV SUBMODULE_DIR=/SHiELD
+COPY submodules ${SUBMODULE_DIR}/
+
+# Build FMS, NCEPlibs, and SHiELD using SHiELD_build
+# FMS_CPPDEFS needed to address https://github.com/NOAA-GFDL/FMS/issues/426 and
+# to build FMS with the -fPIC flag, which is required for the wrapper.
+RUN cd ${SUBMODULE_DIR}/SHiELD_build/Build && \
+    FMS_CPPDEFS="-DHAVE_GETTID -fPIC" \
+    FC=${SHiELD_FC} \
+    CC=${SHiELD_CC} \
+    CXX=${SHiELD_CXX} \
+    LD=${SHiELD_LD} \
+    TEMPLATE=site/gnu.mk \
+    AVX_LEVEL=-march=native \
+    ./COMPILE shield_wrapper 64bit gnu
+
+COPY requirements.txt /tmp
+RUN \
+    FC=${SHiELD_FC} \
+    CC=${SHiELD_CC} \
+    CXX=${SHiELD_CXX} \
+    LD=${SHiELD_LD} \
+    python3 -m pip install -r /tmp/requirements.txt
+
+# Prescribe LDSHARED as well to ensure that the wrapper is built with the GNU
+# 10 compilers.  See https://github.com/ai2cm/fv3gfs-fortran/issues/330 for
+# more comprehensive discussion.  This may not be absolutely necessary (we do
+# not worry about this in fv3net), but it feels cleaner not to mix compilers if
+# we can avoid it.
+COPY wrapper /wrapper
+RUN make -C wrapper \
+    FC=${SHiELD_FC} \
+    CC=${SHiELD_CC} \
+    CXX=${SHiELD_CXX} \
+    LD=${SHiELD_LD} \
+    LDSHARED="${SHiELD_CC} -shared" \
+    FPIC=Y \
+    OPENMP=Y \
+    AVX=Y \
+    build install
+RUN python3 -c "import shield.wrapper"

docker/docker_run.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# Taken largely from fv3net:
+# https://github.com/ai2cm/fv3net/blob/master/tools/docker-run
+
+set -e
+dockerArgs=()
+
+# Google authentication
+if [[ -n $GOOGLE_APPLICATION_CREDENTIALS ]] && \
+   [[ "$(jq -r .type $GOOGLE_APPLICATION_CREDENTIALS)" == "service_account" ]]
+then
+    >&2 echo "Using service account credentials"
+    dockerArgs+=(-v "${GOOGLE_APPLICATION_CREDENTIALS}:/tmp/key.json")
+    dockerArgs+=(-e "GOOGLE_APPLICATION_CREDENTIALS=/tmp/key.json")
+else
+    >&2 echo "Using user credentials"
+    dockerArgs=(-v ~/.config/gcloud:/root/.config/gcloud)
+fi
+
+docker run -e FSSPEC_GS_REQUESTER_PAYS=vcm-ml "${dockerArgs[@]}" $@

requirements.in

@@ -0,0 +1,14 @@
+cftime>=1.2.1
+cython
+jinja2
+fv3config
+mpi4py>=3.0.3
+netCDF4>=1.4.2
+numpy>=1.16
+pace-util>=0.7.0
+pkgconfig
+pytest
+pytest-regtest
+pyyaml>=5
+xarray>=0.15.1
+wheel

requirements.txt

@@ -0,0 +1,65 @@
+#
+# This file is autogenerated by pip-compile
+# To update, run:
+#
+#    pip-compile requirements.in
+#
+aiohttp==3.8.5            # via gcsfs
+aiosignal==1.3.1          # via aiohttp
+appdirs==1.4.4            # via fv3config
+async-timeout==4.0.3      # via aiohttp
+attrs==23.1.0             # via aiohttp
+cachetools==5.3.1         # via google-auth
+certifi==2023.7.22        # via netcdf4, requests
+cftime==1.6.2             # via -r requirements.in, netcdf4, pace-util
+charset-normalizer==3.2.0  # via aiohttp, requests
+cython==3.0.2             # via -r requirements.in
+dacite==1.8.1             # via fv3config
+decorator==5.1.1          # via gcsfs
+exceptiongroup==1.1.3     # via pytest
+f90nml==1.4.3             # via fv3config, pace-util
+frozenlist==1.4.0         # via aiohttp, aiosignal
+fsspec==2023.9.1          # via fv3config, gcsfs, pace-util
+fv3config==0.9.0          # via -r requirements.in
+gcsfs==2023.9.1           # via fv3config
+google-api-core==2.11.1   # via google-cloud-core, google-cloud-storage
+google-auth-oauthlib==1.1.0  # via gcsfs
+google-auth==2.23.0       # via gcsfs, google-api-core, google-auth-oauthlib, google-cloud-core, google-cloud-storage
+google-cloud-core==2.3.3  # via google-cloud-storage
+google-cloud-storage==2.10.0  # via gcsfs
+google-crc32c==1.5.0      # via google-resumable-media
+google-resumable-media==2.6.0  # via google-cloud-storage
+googleapis-common-protos==1.60.0  # via google-api-core
+idna==3.4                 # via requests, yarl
+iniconfig==2.0.0          # via pytest
+jinja2==3.1.2             # via -r requirements.in
+markupsafe==2.1.3         # via jinja2
+mpi4py==3.1.4             # via -r requirements.in
+multidict==6.0.4          # via aiohttp, yarl
+netcdf4==1.6.4            # via -r requirements.in
+numpy==1.24.4             # via -r requirements.in, cftime, netcdf4, pace-util, pandas, xarray
+oauthlib==3.2.2           # via requests-oauthlib
+pace-util==0.9.0          # via -r requirements.in
+packaging==23.1           # via pytest, xarray
+pandas==2.0.3             # via xarray
+pkgconfig==1.5.5          # via -r requirements.in
+pluggy==1.3.0             # via pytest
+protobuf==4.24.3          # via google-api-core
+pyasn1-modules==0.3.0     # via google-auth
+pyasn1==0.5.0             # via pyasn1-modules, rsa
+pytest-regtest==1.5.1     # via -r requirements.in
+pytest==7.4.2             # via -r requirements.in, pytest-regtest
+python-dateutil==2.8.2    # via pandas
+pytz==2023.3.post1        # via pandas
+pyyaml==6.0.1             # via -r requirements.in, fv3config
+requests-oauthlib==1.3.1  # via google-auth-oauthlib
+requests==2.31.0          # via gcsfs, google-api-core, google-cloud-storage, requests-oauthlib
+rsa==4.9                  # via google-auth
+six==1.16.0               # via python-dateutil
+tomli==2.0.1              # via pytest
+typing-extensions==4.8.0  # via pace-util
+tzdata==2023.3            # via pandas
+urllib3==1.26.16          # via google-auth, requests
+wheel==0.41.2             # via -r requirements.in
+xarray==2023.1.0          # via -r requirements.in
+yarl==1.9.2               # via aiohttp

wrapper/LICENSE

@@ -0,0 +1,32 @@
+
+
+BSD License
+
+Copyright (c) 2019, Vulcan Technologies LLC
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice, this
+  list of conditions and the following disclaimer in the documentation and/or
+  other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from this
+  software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
+IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
+INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
+OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
+OF THE POSSIBILITY OF SUCH DAMAGE.
+

wrapper/MANIFEST.in

@@ -0,0 +1,16 @@
+include LICENSE
+
+recursive-include shield *.json *.yml
+recursive-include tests *.py *.sh
+recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif
+recursive-include lib *
+recursive-include shield *
+
+recursive-exclude * __pycache__
+recursive-exclude * *.py[co]
+recursive-exclude * *.mod *.o
+recursive-exclude * *.log
+
+prune tests/*workdir
+prune tests/logs/*
+prune lib/external