Merge pull request #4 from ecmwf-lab/feature/earthkit

Feature/earthkit

.github/workflows/python-publish.yml

@@ -15,14 +15,13 @@ jobs:
     name: Code QA
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - run: pip install black flake8 isort
-    - run: black --version
-    - run: isort --version
-    - run: flake8 --version
-    - run: isort --check .
-    - run: black --check .
-    - run: flake8 .
+    - uses: actions/checkout@v4
+    - uses: actions/setup-python@v5
+      with:
+        python-version: 3.x
+    - uses: pre-commit/[email protected]
+      env:
+        SKIP: no-commit-to-branch
 
   checks:
     if: ${{ github.event_name == 'release' }}

.gitignore

@@ -158,3 +158,9 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+
+*.onnx
+*.grib
+*.out
+_version.py

.pre-commit-config.yaml

@@ -0,0 +1,69 @@
+repos:
+
+# Empty notebookds
+- repo: local
+  hooks:
+    - id: clear-notebooks-output
+      name: clear-notebooks-output
+      files: tools/.*\.ipynb$
+      stages: [commit]
+      language: python
+      entry: jupyter nbconvert --ClearOutputPreprocessor.enabled=True --inplace
+      additional_dependencies: [jupyter]
+
+
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.4.0
+  hooks:
+  - id: check-yaml # Check YAML files for syntax errors only
+    args: [--unsafe, --allow-multiple-documents]
+  - id: debug-statements # Check for debugger imports and py37+ breakpoint()
+  - id: end-of-file-fixer # Ensure files end in a newline
+  - id: trailing-whitespace # Trailing whitespace checker
+  - id: no-commit-to-branch # Prevent committing to main / master
+  - id: check-added-large-files # Check for large files added to git
+  - id: check-merge-conflict # Check for files that contain merge conflict
+
+- repo: https://github.com/psf/black-pre-commit-mirror
+  rev: 24.1.1
+  hooks:
+    - id: black
+      args: [--line-length=120]
+
+- repo: https://github.com/pycqa/isort
+  rev: 5.13.2
+  hooks:
+  - id: isort
+    args:
+    - -l 120
+    - --force-single-line-imports
+    - --profile black
+
+
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.3.0
+  hooks:
+  - id: ruff
+    exclude: '(dev/.*|.*_)\.py$'
+    args:
+    - --line-length=120
+    - --fix
+    - --exit-non-zero-on-fix
+    - --preview
+
+- repo: https://github.com/sphinx-contrib/sphinx-lint
+  rev: v0.9.1
+  hooks:
+    - id: sphinx-lint
+
+# For now, we use it. But it does not support a lot of sphinx features
+- repo: https://github.com/dzhu/rstfmt
+  rev: v0.0.14
+  hooks:
+    - id: rstfmt
+
+- repo: https://github.com/b8raoult/pre-commit-docconvert
+  rev: "0.1.4"
+  hooks:
+  - id: docconvert
+    args: ["numpy"]

pyproject.toml

@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+# https://packaging.python.org/en/latest/guides/writing-pyproject-toml/
+
+[build-system]
+requires = [
+  "setuptools>=60",
+  "setuptools-scm>=8",
+]
+
+[project]
+name = "ai-models-panguweather"
+
+description = "An ai-models plugin to run PanguWeather"
+keywords = [
+  "ai",
+  "tools",
+]
+
+license = { file = "LICENSE" }
+authors = [
+  { name = "European Centre for Medium-Range Weather Forecasts (ECMWF)", email = "[email protected]" },
+]
+
+requires-python = ">=3.10"
+
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "License :: OSI Approved :: Apache Software License",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python :: 3 :: Only",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: Implementation :: CPython",
+  "Programming Language :: Python :: Implementation :: PyPy",
+]
+
+dynamic = [
+  "version",
+]
+
+dependencies = [
+  "ai-models>=0.7.0",
+  "onnx",
+  "onnxruntime-silicon; sys_platform == 'darwin'",
+  "onnxruntime-gpu; sys_platform == 'linux'"
+]
+
+optional-dependencies.dev = [
+  "pre-commit",
+]
+urls.Repository = "https://github.com/ecmwf-lab/ai-models-panguweather"
+entry-points."ai_models.model".panguweather = "ai_models_panguweather.model:PanguWeather"
+
+[tool.setuptools_scm]
+version_file = "src/ai_models_panguweather/_version.py"

ai_models_panguweather/model.py → src/ai_models_panguweather/model.py

@@ -18,9 +18,7 @@
 
 class PanguWeather(Model):
     # Download
-    download_url = (
-        "https://get.ecmwf.int/repository/test-data/ai-models/pangu-weather/{file}"
-    )
+    download_url = "https://get.ecmwf.int/repository/test-data/ai-models/pangu-weather/{file}"
     download_files = ["pangu_weather_24.onnx", "pangu_weather_6.onnx"]
 
     # Input

utils/pangu-gfs-input.py

@@ -11,7 +11,7 @@
 
 import sys
 
-import climetlab as cml
+import earthkit.data as ekd
 import tqdm
 
 date = sys.argv[1]
@@ -24,10 +24,9 @@
 time = int(time)
 time = f"{time:02d}"
 
-gfs = cml.load_source(
+gfs = ekd.from_source(
     "url-pattern",
-    "https://data.rda.ucar.edu/ds084.1/{year}/{date}/"
-    "gfs.0p25.{date}{time}.f{step}.grib2",
+    "https://data.rda.ucar.edu/ds084.1/{year}/{date}/" "gfs.0p25.{date}{time}.f{step}.grib2",
     date=date,
     time=time,
     step="000",
@@ -62,7 +61,7 @@
 ).order_by(param=param_sfc)
 
 print("Write", output)
-out = cml.new_grib_output(output)
+out = ekd.new_grib_output(output)
 
 G = {"gh": 9.80665}
 PARAM = {"gh": "z", "prmsl": "msl"}