Skip to content

Commit 3573805

Browse files
brianhennbrianhenn
committed
squash linear history
1 parent 0c36153 commit 3573805

File tree

401 files changed

+65395
-196
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

401 files changed

+65395
-196
lines changed

‎.github/pull_request_template.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,6 @@ Changes:
55
- Can group multiple related symbols on a single bullet
66

77
- [ ] Tests added
8+
- [ ] If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated
89

910
Resolves #<github issues> (delete if none)

‎.github/workflows/docs.yaml‎

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,30 +8,56 @@ on:
88
branches:
99
- main
1010

11+
# Add a button to manually trigger the workflow
12+
workflow_dispatch:
13+
14+
# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
15+
permissions:
16+
contents: read
17+
pages: write
18+
id-token: write
19+
20+
# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued.
21+
# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete.
22+
concurrency:
23+
group: "pages"
24+
cancel-in-progress: false
25+
1126
jobs:
1227
build:
1328
runs-on: ubuntu-latest
1429
steps:
1530
- uses: actions/checkout@v3
31+
- name: Set environment variables
32+
run: |
33+
echo "CURRENT_WEEK=$(date +'%Y-%U')" >> $GITHUB_ENV
1634
- uses: actions/setup-python@v4
1735
with:
18-
python-version: "3.10"
36+
python-version: "3.11"
1937
- uses: actions/cache@v4
2038
with:
2139
path: ${{ env.pythonLocation }}
22-
key: ${{ env.pythonLocation }}-${{ hashFiles('fme/requirements.txt') }}-${{ hashFiles('fme/docs/requirements.txt') }}-${{ hashFiles('fme/constraints.txt') }}
40+
key: v1-${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('docs/requirements.txt') }}-${{ hashFiles('requirements-healpix.txt') }}-${{ hashFiles('constraints.txt') }}
2341
- name: Install dependencies
2442
run: |
25-
python -m pip install uv==0.2.5
26-
uv pip install --system -c constraints.txt -e fme[docs]
43+
python -m pip install uv
44+
# can't use --system here because it doesn't work with --no-build-isolation and a healpix dependency
45+
uv venv --python 3.11
46+
uv pip install -c constraints.txt -e .[docs]
47+
uv pip install --no-build-isolation -c constraints.txt -e .[docs,healpix]
2748
- name: Build docs
2849
run: |
29-
cd fme/docs && make doctest html
30-
- name: Deploy to GitHub Pages
31-
uses: peaceiris/actions-gh-pages@v3
32-
if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
50+
source .venv/bin/activate
51+
cd docs && make doctest html
52+
- name: Upload to GitHub Pages
53+
uses: actions/upload-pages-artifact@v3
3354
with:
34-
publish_branch: gh-pages
35-
github_token: ${{ secrets.GITHUB_TOKEN }}
36-
publish_dir: fme/docs/_build/
37-
force_orphan: true
55+
path: docs/_build/html
56+
deploy:
57+
runs-on: ubuntu-latest
58+
# Only deploy on manual trigger or push to main
59+
if: ${{ github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main') }}
60+
needs: build
61+
steps:
62+
- name: Deploy to GitHub Pages
63+
uses: actions/deploy-pages@v4

‎.github/workflows/pre-commit.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ jobs:
1818
- name: Set up Python
1919
uses: actions/setup-python@v2
2020
with:
21-
python-version: "3.10"
21+
python-version: "3.11"
2222

2323
- name: Install dependencies
2424
run: |

‎.github/workflows/tests.yaml‎

Lines changed: 45 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,23 @@ jobs:
1818
echo "CURRENT_WEEK=$(date +'%Y-%U')" >> $GITHUB_ENV
1919
- uses: actions/setup-python@v4
2020
with:
21-
python-version: "3.10"
21+
python-version: "3.11"
2222
- uses: actions/cache@v4
2323
with:
2424
path: ${{ env.pythonLocation }}
25-
key: ${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('fme/requirements.txt') }}-${{ hashFiles('fme/dev-requirements.txt') }}-${{ hashFiles('constraints.txt') }}
25+
key: v1-${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}-${{ hashFiles('requirements-healpix.txt') }}-${{ hashFiles('constraints.txt') }}
2626
- name: Install dependencies
2727
run: |
2828
python -m pip install uv
29-
uv pip install --system -c constraints.txt -e ./fme[dev]
30-
- name: Run pytest
29+
# can't use --system here because it doesn't work with --no-build-isolation and a healpix dependency
30+
uv venv --python 3.11
31+
uv pip install -c constraints.txt -e .[dev]
32+
uv pip install --no-build-isolation -c constraints.txt -e .[dev,healpix]
33+
- name: Run pytest with code coverage
3134
run: |
32-
make test
33-
cpu-very-fast:
35+
source .venv/bin/activate
36+
make test_cov
37+
cpu-very-fast-no-healpix:
3438
runs-on: ubuntu-latest
3539
steps:
3640
- uses: actions/checkout@v3
@@ -39,15 +43,47 @@ jobs:
3943
echo "CURRENT_WEEK=$(date +'%Y-%U')" >> $GITHUB_ENV
4044
- uses: actions/setup-python@v4
4145
with:
42-
python-version: "3.10"
46+
python-version: "3.11"
4347
- uses: actions/cache@v4
4448
with:
4549
path: ${{ env.pythonLocation }}
46-
key: ${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('fme/requirements.txt') }}-${{ hashFiles('fme/dev-requirements.txt') }}-${{ hashFiles('constraints.txt') }}
50+
key: v1-${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}-${{ hashFiles('constraints.txt') }}
4751
- name: Install dependencies
4852
run: |
4953
python -m pip install uv
50-
uv pip install --system -c constraints.txt -e ./fme[dev]
54+
uv pip install --system -c constraints.txt -e .[dev]
5155
- name: Run pytest
5256
run: |
5357
make test_very_fast
58+
gpu:
59+
# These steps run on our self-hosted k8s GPU runners. See below for setup notes
60+
# https://github.com/ai2cm/long-lived-infrastructure#installing-the-github-actions-runner-controller
61+
runs-on: arc-runner-set
62+
steps:
63+
- uses: actions/checkout@v3
64+
- name: Set environment variables
65+
run: |
66+
echo "CURRENT_WEEK=$(date +'%Y-%U')" >> $GITHUB_ENV
67+
- uses: actions/setup-python@v4
68+
with:
69+
python-version: "3.11"
70+
- uses: actions/cache@v4
71+
with:
72+
path: ${{ env.pythonLocation }}
73+
key: ${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('requirements.txt') }}-${{ hashFiles('requirements-dev.txt') }}-${{ hashFiles('requirements-healpix.txt') }}-${{ hashFiles('constraints.txt') }}
74+
- name: Install dependencies
75+
run: |
76+
sudo apt-get update
77+
sudo apt-get install -y make git build-essential
78+
python -m pip install uv
79+
# can't use --system here because it doesn't work with --no-build-isolation and a healpix dependency
80+
uv venv --python 3.11
81+
uv pip install -c constraints.txt -e .[dev]
82+
uv pip install --no-build-isolation -c constraints.txt -e .[dev,healpix]
83+
- name: Run CUDA check and pytest
84+
run: |
85+
source .venv/bin/activate
86+
export PATH=/usr/local/nvidia/bin:${PATH}
87+
export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:${LD_LIBRARY_PATH}
88+
python3 fme/require_gpu.py
89+
make test

‎.gitignore‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1-
fme/docs/available_modules.rst
2-
fme/docs/_build
1+
docs/available_modules.rst
2+
docs/_build
3+
filtered_repo/
34

45
.vscode
56

‎.pre-commit-config.yaml‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,22 +3,30 @@ repos:
33
rev: v0.8.1
44
hooks:
55
- id: ruff
6-
args: ["--fix", "--config", "fme/pyproject.toml"]
6+
args: ["--fix", "--config", "pyproject.toml"]
77
- id: ruff-format
88
- repo: https://github.com/pre-commit/pre-commit-hooks
99
rev: v5.0.0
1010
hooks:
1111
- id: check-added-large-files
1212
args: [--maxkb=250]
1313
- id: trailing-whitespace
14+
- id: file-contents-sorter
15+
files: |
16+
(?x)^(
17+
^.*requirements.*\.txt|
18+
constraints.txt|
19+
analysis-deps.txt
20+
)$
1421
- repo: https://github.com/pre-commit/mirrors-mypy
15-
rev: v1.2.0
22+
rev: v1.15.0
1623
hooks:
1724
- id: mypy
18-
additional_dependencies: ["types-PyYaml==5.4.3"]
25+
additional_dependencies: ["types-PyYaml==5.4.3", "wandb==0.19.0"]
1926
args: ["--ignore-missing-imports", "--check-untyped-defs"]
2027
exclude: |
2128
(?x)^(
2229
.+/conf.py |
23-
.+/conftest.py
30+
.+/conftest.py |
31+
fme/ace/models/makani_fcn2/.*
2432
)$

‎Makefile‎

Lines changed: 41 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,59 @@
11
VERSION ?= $(shell git rev-parse --short HEAD)
22
IMAGE ?= fme
3+
REGISTRY ?= registry.nersc.gov/m4492/ai2cm
34
ENVIRONMENT_NAME ?= fme
5+
USERNAME ?= $(shell beaker account whoami --format=json | jq -r '.[0].name')
46
DEPLOY_TARGET ?= pypi
7+
BEAKER_WORKSPACE = ai2/ace
8+
9+
ifeq ($(shell uname), Linux)
10+
CONDA_PACKAGES=gxx_linux-64 pip
11+
else
12+
CONDA_PACKAGES=pip
13+
endif
514

615
build_docker_image:
7-
docker build --platform=linux/amd64 -f docker/Dockerfile -t $(IMAGE):$(VERSION) .
16+
DOCKER_BUILDKIT=1 docker build --platform=linux/amd64 -f docker/Dockerfile -t $(IMAGE):$(VERSION) --target production .
17+
18+
push_shifter_image: build_docker_image
19+
docker tag $(IMAGE):$(VERSION) $(REGISTRY)/$(IMAGE):$(VERSION)
20+
docker push $(REGISTRY)/$(IMAGE):$(VERSION)
21+
22+
build_beaker_image: build_docker_image
23+
beaker image create --name $(IMAGE)-$(VERSION) $(IMAGE):$(VERSION)
24+
25+
build_podman_image:
26+
podman-hpc build -f docker/Dockerfile -t $(IMAGE):$(VERSION) .
27+
28+
migrate_podman_image: build_podman_image
29+
podman-hpc migrate $(IMAGE):$(VERSION)
830

931
enter_docker_image: build_docker_image
1032
docker run -it --rm $(IMAGE):$(VERSION) bash
1133

34+
launch_beaker_session:
35+
./launch-beaker-session.sh $(USERNAME)/$(IMAGE)-$(VERSION)
36+
37+
build_deps_only_image:
38+
DOCKER_BUILDKIT=1 docker build --platform=linux/amd64 -f docker/Dockerfile -t $(IMAGE)-deps-only:$(VERSION) --target deps-only .
39+
beaker image create $(IMAGE)-deps-only:$(VERSION) --name $(IMAGE)-deps-only-$(VERSION) --workspace ai2/ace-ci-tests
40+
41+
1242
# recommended to deactivate current conda environment before running this
1343
create_environment:
14-
conda create -n $(ENVIRONMENT_NAME) python=3.10 pip
44+
conda create -n $(ENVIRONMENT_NAME) python=3.11 $(CONDA_PACKAGES)
1545
conda run --no-capture-output -n $(ENVIRONMENT_NAME) python -m pip install uv
16-
conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install -c constraints.txt -e ./fme[dev,docs]
46+
conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install -c constraints.txt -e .[dev,docs]
47+
conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install --no-build-isolation -c constraints.txt -e .[dev,docs,healpix]
1748
conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install -r analysis-deps.txt
1849

1950
test:
2051
pytest --durations 40 .
2152

53+
# --cov must come after pytest args to use the sources defined by config
54+
test_cov:
55+
pytest --durations 40 --cov --cov-report=term-missing:skip-covered --cov-config=pyproject.toml .
56+
2257
test_fast:
2358
pytest --durations 40 --fast .
2459

@@ -29,11 +64,11 @@ test_very_fast:
2964
# requires fme[deploy] to be installed
3065

3166
build_pypi:
32-
rm -rf fme/dist
33-
cd fme && python -m build
67+
rm -rf dist
68+
python -m build
3469

3570
deploy_pypi: build_pypi
36-
cd fme && twine upload --repository $(DEPLOY_TARGET) dist/*
71+
twine upload --repository $(DEPLOY_TARGET) dist/*
3772

3873
deploy_test_pypi: DEPLOY_TARGET = testpypi
3974
deploy_test_pypi: deploy_pypi

‎PACKAGE_README.md‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
[![Docs](https://readthedocs.org/projects/ai2-climate-emulator/badge/?version=latest)](https://ai2-climate-emulator.readthedocs.io/en/latest/)
2+
3+
# FME: Weather/Climate Model Emulation
4+
This package contains code to train and evaluate weather/climate model emulators as seen in
5+
"ACE: A fast, skillful learned global atmospheric model for climate prediction" ([arxiv:2310.02074](https://arxiv.org/abs/2310.02074))
6+
and "Application of the Ai2 Climate Emulator to E3SMv2's global atmosphere model, with a focus on precipitation fidelity"
7+
([JGR-ML](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2024JH000136)).
8+
9+
10+
## Installation
11+
12+
The package can be installed via PyPI using:
13+
14+
```
15+
pip install fme
16+
```
17+
18+
## Quickstart
19+
20+
A quickstart guide may be found [here](https://ai2-climate-emulator.readthedocs.io/en/latest/quickstart.html).
21+
22+
## Documentation
23+
24+
See complete documentation [here](https://ai2-climate-emulator.readthedocs.io/en/latest/).

‎README.md‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,12 @@
77

88
Ai2 Climate Emulator (ACE) is a fast machine learning model that simulates global atmospheric variability in a changing climate over time scales ranging from hours to centuries.
99

10-
This repo contains code accompanying four papers describing ACE models:
10+
This repo contains code accompanying five papers describing ACE models:
1111
- "ACE: A fast, skillful learned global atmospheric model for climate prediction" ([link](https://arxiv.org/abs/2310.02074))
1212
- "Application of the Ai2 Climate Emulator to E3SMv2's global atmosphere model, with a focus on precipitation fidelity" ([link](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2024JH000136))
13-
- "ACE2: Accurately learning subseasonal to decadal atmospheric variability and forced responses" ([link](https://arxiv.org/abs/2411.11268))
13+
- "ACE2: Accurately learning subseasonal to decadal atmospheric variability and forced responses" ([link](https://www.nature.com/articles/s41612-025-01090-0))
1414
- "ACE2-SOM: Coupling to a slab ocean and learning the sensitivity of climate to changes in CO2" ([link](https://arxiv.org/abs/2412.04418))
15+
- "Applying the ACE2 Emulator to SST Green's Functions for the E3SMv3 Global Atmosphere Model" ([link](https://arxiv.org/abs/2505.08742))
1516

1617
## Installation
1718

‎analysis-deps.txt‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
1-
# these are some packages which are convenient to have installed for ad-hoc analysis
2-
# but which are not requirements of the "fme" package. We do not relist the fme
3-
# dependencies here.
4-
beaker-py
1+
# these are convenient to have installed for ad-hoc analysis but are not requirements of the "fme" package
52
Bottleneck
6-
cartopy>=0.22.0
3+
beaker-gantry>=1.12.0
4+
beaker-py
5+
bokeh>=3.1.0
76
dask[distributed]
87
ipywidgets
9-
nc-time-axis
108
jupyterlab
9+
nc-time-axis
1110
pyproj<3.7
1211
seaborn
13-
bokeh>=3.1.0

0 commit comments

Comments
 (0)