Updated public release (#25)

Updates the codebase to a recent commit (`2428726`). Squash-merged git history.

.github/workflows/tests.yaml

@@ -13,17 +13,41 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
+      - name: Set environment variables
+        run: |
+          echo "CURRENT_WEEK=$(date +'%Y-%U')" >> $GITHUB_ENV
       - uses: actions/setup-python@v4
         with:
           python-version: "3.10"
       - uses: actions/cache@v4
         with:
           path: ${{ env.pythonLocation }}
-          key: ${{ env.pythonLocation }}-${{ hashFiles('fme/requirements.txt') }}-${{ hashFiles('fme/dev-requirements.txt') }}-${{ hashFiles('fme/constraints.txt') }}
+          key: ${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('fme/requirements.txt') }}-${{ hashFiles('fme/dev-requirements.txt') }}-${{ hashFiles('constraints.txt') }}
       - name: Install dependencies
         run: |
-          python -m pip install uv==0.2.5
-          uv pip install --system -c constraints.txt -e fme[dev]
+          python -m pip install uv
+          uv pip install --system -c constraints.txt -e ./fme[dev]
       - name: Run pytest
         run: |
           make test
+  cpu-very-fast:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set environment variables
+        run: |
+          echo "CURRENT_WEEK=$(date +'%Y-%U')" >> $GITHUB_ENV
+      - uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+      - uses: actions/cache@v4
+        with:
+          path: ${{ env.pythonLocation }}
+          key: ${{ env.CURRENT_WEEK }}-${{ env.pythonLocation }}-${{ hashFiles('fme/requirements.txt') }}-${{ hashFiles('fme/dev-requirements.txt') }}-${{ hashFiles('constraints.txt') }}
+      - name: Install dependencies
+        run: |
+          python -m pip install uv
+          uv pip install --system -c constraints.txt -e ./fme[dev]
+      - name: Run pytest
+        run: |
+          make test_very_fast

.gitignore

@@ -136,3 +136,6 @@ dmypy.json
 
 # scratch directory for testing
 scratch/
+
+# Some in progress data pipelines get added here
+scripts/data_process/.nfs*

.pre-commit-config.yaml

@@ -1,41 +1,24 @@
 repos:
--   repo: https://github.com/psf/black
-    rev: 23.3.0
-    hooks:
-    - id: black
-      additional_dependencies: ["click==8.0.4"]
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.8.1
+  hooks:
+    - id: ruff
+      args: ["--fix", "--config", "fme/pyproject.toml"]
+    - id: ruff-format
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.3.0
+    rev: v5.0.0
     hooks:
     - id: check-added-large-files
       args: [--maxkb=250]
     - id: trailing-whitespace
-    - id: flake8
-      name: flake8
-      language_version: python3
-      exclude: "__init__.py"
-      args: [--config, setup.cfg]
-    - id: flake8
-      name: flake8 __init__.py files
-      files: "__init__.py"
-      # ignore unused import error in __init__.py files
-      args: ["--ignore=F401,E203,W503", --config, setup.cfg]
 -   repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.2.0
     hooks:
     - id: mypy
       additional_dependencies: ["types-PyYaml==5.4.3"]
-      args: [
-          --follow-imports, silent, --ignore-missing-imports
-        ]
+      args: ["--ignore-missing-imports", "--check-untyped-defs"]
       exclude: |
         (?x)^(
         .+/conf.py |
-        .+/setup.py |
         .+/conftest.py
-        )$
--   repo: https://github.com/pycqa/isort
-    rev: 5.11.5
-    hooks:
-      - id: isort
-        name: isort (python)
+        )$

Makefile

@@ -4,19 +4,26 @@ ENVIRONMENT_NAME ?= fme
 DEPLOY_TARGET ?= pypi
 
 build_docker_image:
-	docker build -f docker/Dockerfile -t $(IMAGE):$(VERSION) .
+	docker build --platform=linux/amd64 -f docker/Dockerfile -t $(IMAGE):$(VERSION) .
 
 enter_docker_image: build_docker_image
 	docker run -it --rm $(IMAGE):$(VERSION) bash
 
 # recommended to deactivate current conda environment before running this
 create_environment:
 	conda create -n $(ENVIRONMENT_NAME) python=3.10 pip
-	conda run --no-capture-output -n $(ENVIRONMENT_NAME) python -m pip install uv==0.2.5
-	conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install -c constraints.txt -e fme[dev]
+	conda run --no-capture-output -n $(ENVIRONMENT_NAME) python -m pip install uv
+	conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install -c constraints.txt -e ./fme[dev,docs]
+	conda run --no-capture-output -n $(ENVIRONMENT_NAME) uv pip install -r analysis-deps.txt
 
 test:
-	pytest --durations 20 .
+	pytest --durations 40 .
+
+test_fast:
+	pytest --durations 40 --fast .
+
+test_very_fast:
+	pytest --durations 40 --very-fast .
 
 # For maintainer use only
 # requires fme[deploy] to be installed

README.md

@@ -43,4 +43,4 @@ gs://ai2cm-public-requester-pays/2024-11-13-ai2-climate-emulator-v2-amip/data/er
 The dataset used in the [ACE2-SOM paper](https://arxiv.org/abs/2412.04418) is available at:
 ```
 gs://ai2cm-public-requester-pays/2024-12-05-ai2-climate-emulator-v2-som/SHiELD-SOM-C96
-```
+```

analysis-deps.txt

@@ -0,0 +1,13 @@
+# these are some packages which are convenient to have installed for ad-hoc analysis
+# but which are not requirements of the "fme" package. We do not relist the fme
+# dependencies here.
+beaker-py
+Bottleneck
+cartopy>=0.22.0
+dask[distributed]
+ipywidgets
+nc-time-axis
+jupyterlab
+pyproj<3.7
+seaborn
+bokeh>=3.1.0

conftest.py

@@ -0,0 +1,68 @@
+import signal
+
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--fast",
+        action="store_true",
+        default=False,
+        help="Skip slow tests",
+    )
+    parser.addoption(
+        "--very-fast",
+        action="store_true",
+        default=False,
+        help="Run only very fast tests (< 5 seconds)",
+    )
+
+
+@pytest.fixture
+def skip_slow(request, very_fast_only):
+    return very_fast_only or request.config.getoption("--fast")
+
+
+@pytest.fixture
+def very_fast_only(request):
+    return request.config.getoption("--very-fast")
+
+
+class TimeoutException(Exception):
+    pass
+
+
+def timeout_handler(signum, frame):
+    raise TimeoutException("Test took too long")
+
+
+@pytest.fixture
+def pdb_enabled(request):
+    return request.config.getoption("--pdb")
+
+
+@pytest.fixture(autouse=True, scope="function")
+def enforce_timeout(skip_slow, very_fast_only, pdb_enabled):
+    if pdb_enabled:
+        yield  # Do not enforce timeout if we are debugging
+        return
+    if very_fast_only:
+        timeout_seconds = 3
+    elif skip_slow:
+        timeout_seconds = 30
+    else:
+        timeout_seconds = 60
+    signal.signal(signal.SIGALRM, timeout_handler)
+    signal.alarm(timeout_seconds)  # Set the timeout for the test
+    try:
+        yield
+    finally:
+        signal.alarm(0)  # Disable the alarm after the test completes
+
+
+@pytest.hookimpl(tryfirst=True, hookwrapper=True)
+def pytest_runtest_call(item):
+    try:
+        yield
+    except TimeoutException:
+        pytest.fail("Test failed due to timeout")

fme/docs/_static/custom.css

@@ -0,0 +1,21 @@
+body[data-theme="dark"] {
+    --code-block-background: #202020;
+}
+
+body[data-theme="light"] {
+    --code-block-background: #f8f9fb;
+}
+
+body[data-theme="auto"] {
+    --code-block-background: #f8f9fb;
+}
+
+@media (prefers-color-scheme: dark) {
+    body[data-theme="auto"] {
+        --code-block-background: #202020;
+    }
+}
+
+div.highlight pre {
+    background: var(--code-block-background);
+}

fme/docs/builder.rst

@@ -51,7 +51,7 @@ Let's define a training configuration ``TrainConfig`` containing an ``OptimizerC
         """
         Configuration for an optimizer.
 
-        Attributes:
+        Parameters:
             optimizer_type: The type of optimizer to use.
             lr: The learning rate.
             kwargs: Additional keyword arguments to pass to the optimizer.

fme/docs/conf.py

@@ -24,7 +24,7 @@
 sys.path.insert(0, os.path.abspath(".."))
 
 import fme  # noqa
-import fme.ace
+import fme.core.registry
 
 # -- General configuration ---------------------------------------------
 
@@ -33,7 +33,7 @@
 # needs_sphinx = '1.0'
 
 # Fetch the dynamic data
-module_types = fme.ace.get_available_module_types()
+module_types = fme.core.registry.ModuleSelector.get_available_types()
 
 # Create a dynamic rst snippet that can be included in your documentation
 rst_snippet = f".. code-block:: text\n\n   {module_types}"
@@ -50,7 +50,15 @@
     "sphinx.ext.autodoc",
     "sphinx.ext.viewcode",
     "sphinx.ext.doctest",
+    "sphinx.ext.intersphinx",
+    "sphinx_autodoc_typehints",
 ]
+autodoc_typehints = "description"
+
+# Intersphinx configuration to link to other projects' documentation
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3", None),
+}
 
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]
@@ -93,6 +101,9 @@
 # The name of the Pygments (syntax highlighting) style to use.
 pygments_style = "sphinx"
 
+# Include default values when documenting parameter types.
+typehints_defaults = "comma"
+
 # If true, `todo` and `todoList` produce output, else they produce nothing.
 todo_include_todos = False
 
@@ -102,18 +113,37 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = "sphinx_rtd_theme"
+html_theme = "furo"
+html_title = f"fme v{fme.__version__}"
 
 # Theme options are theme-specific and customize the look and feel of a
 # theme further.  For a list of options available for each theme, see the
 # documentation.
 #
-# html_theme_options = {}
+html_theme_options = {
+    "light_logo": "Ai2_icon_pink_padding_RGB.png",
+    "dark_logo": "Ai2_icon_pink_padding_RGB.png",
+    "footer_icons": [
+        {
+            "name": "GitHub",
+            "url": "https://github.com/ai2cm/ace",
+            "html": """
+                <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 16 16">
+                    <path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0 0 16 8c0-4.42-3.58-8-8-8z"></path>
+                </svg>
+            """,  # noqa: E501
+            "class": "",
+        },
+    ],
+}
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = []
+html_static_path = ["_static"]
+html_css_files = ["custom.css"]
+
+html_favicon = "_static/Ai2_icon_pink_RGB.png"
 
 
 # -- Options for HTMLHelp output ---------------------------------------
@@ -176,3 +206,9 @@
         "Miscellaneous",
     ),
 ]
+
+
+# -- Options for doctest -----------------------------------------------
+doctest_global_setup = """
+import fme
+"""

fme/docs/configs/explicit-indices.yaml

@@ -0,0 +1,3 @@
+path: initial_conditions.nc
+start_indices:
+   list: [0, 3, 7]

fme/docs/configs/inference-ic-indices.yaml

@@ -0,0 +1,5 @@
+path: initial_conditions.nc
+start_indices:
+  n_initial_conditions: 3
+  first: 1
+  interval: 2

fme/docs/configs/timestamp-list.yaml

@@ -0,0 +1,5 @@
+path: initial_conditions.nc
+start_indices:
+   times:
+      - "2021-01-01T00:00:00"
+      - "2021-02-01T00:00:00"

fme/docs/index.rst

@@ -1,6 +1,9 @@
-Welcome to ACE's documentation!
+fme: Full Model Emulation
 ======================================
 
+**fme** ("full model emulation") is a python package for training and running
+climate model emulators, such as the Ai2 Climate Emulator.
+
 .. toctree::
    :maxdepth: 1
    :caption: Contents: