pyproject.toml

[project]
name = "weathergen"
version = "0.1.0"
description = "The WeatherGenerator Machine Learning Earth System Model"
readme = "README.md"
authors = [
    { name = "WeatherGenerator collaboration" }
]

requires-python = ">=3.11,<3.13"
# TODO: split the plotting dependencies into their own dep groups, they are not required.
dependencies = [
 'torch==2.6.0',
 'numpy~=2.2',
 'astropy_healpix~=1.1.2',
 'zarr~=2.17',
 'anemoi-datasets~=0.5.16',
 'pandas~=2.2',
 'pynvml',
 'tqdm',
 'matplotlib',
 'packaging',
 'wheel',
 'psutil',
 "flash-attn; sys_platform == 'linux'",
 "polars~=1.25.2",
]

[project.urls]
Homepage = "https://www.weathergenerator.eu"
Documentation = "https://readthedocs.org"
Repository = "https://github.com/ecmwf/WeatherGenerator.git"
Issues = "https://github.com/ecmwf/WeatherGenerator/issues"

[project.scripts]
train = "weathergen:train"
evaluate = "weathergen:evaluate"

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[dependency-groups]
# The development dependencies
dev = [
  "pytest",
  "ruff==0.9.7"
]


[tool.black]

# Wide rows
line-length = 100


# The linting configuration
[tool.ruff]

# Wide rows
line-length = 100

[tool.ruff.lint]
# All disabled until the code is formatted.
select = [
    # pycodestyle
    "E",
    # Pyflakes
    "F",
    # pyupgrade
    "UP",
    # flake8-bugbear
    "B",
    # flake8-simplify
    "SIM",
    # isort
    "I",
]

# These rules are sensible and should be enabled at a later stage.
ignore = [
  "E501",
  "E721",
  "E722",
  "B006",
  "B011",
  "UP008",
  "SIM115",
  "SIM117",
  "SIM118",
  "SIM102",
  "SIM210",
  "SIM212",
  "SIM401",
  "F811",
  # To ignore, not relevant for us
  "E741",
]

[tool.uv]
# This guarantees that the build is deterministic and will not be impacted 
# by future releases of dependencies or sub-dependencies. 
# See https://docs.astral.sh/uv/reference/settings/#exclude-newer
# TODO: pytorch does not publish valid release timestamps, so sadly it does not work.
# exclude-newer = "2025-03-14T00:00:00Z"

# Following the recommendations from https://docs.astral.sh/uv/guides/integration/pytorch
# The current setup is:
# linux == GPU + flashattention
# windows == GPU
# macos == CPU
[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true


[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[tool.uv.sources]
torch = [
  { index = "pytorch-cu124", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
  { index = "pytorch-cpu", marker = "sys_platform == 'macosx'"},
]
# This URL was evaluated this way:
# uv run ~/WeatherGenerator-private/hpc/hpc2020/ecmwf/get-flash-atten.sh
flash-attn = [
  { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux'"  },
]