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.12,<3.13"
# TODO: split the plotting dependencies into their own dep groups, they are not required.
dependencies = [
 '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',
 "polars~=1.25.2",
 "omegaconf~=2.3.0",
 "dask~=2025.5.1",
 "hatchling",
 "numexpr>=2.11.0",
 "weathergen-common",
 "weathergen-evaluate",
 "weathergen-readers-extra",
]


[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.run_train:train"
train_continue = "weathergen.run_train:train_continue"
inference = "weathergen.run_train:inference"
evaluate = "weathergen.evaluate.run_evaluation:evaluate"
plot_train = "weathergen.utils.plot_training:plot_train"
export = "weathergen.evaluate.export.export_inference:export"

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

[tool.hatch.build.targets.wheel]
packages = ["src/weathergen"]

[dependency-groups]
# The development dependencies
dev = [
 "ipykernel>=6.30.0",
 "jupytext>=1.17.2",
 "pytest~=8.3.5",
 "pytest-mock>=3.14.1",
 "ruff==0.9.7",
 "tensorboard>=2.20.0",
 "pdbpp>=0.11.7",
 "pyrefly==0.36.0",
]


# Torch listed as optional dependencies.
# uv and python can only filter dependencies by platform, not by capability.
# Following the recommendations from https://docs.astral.sh/uv/guides/integration/pytorch
# We need to support:
# x86_64: cpu (unit tests) + gpu 
# aarch64: gpu
[project.optional-dependencies]

cpu = [
  'torch==2.6.0',
]

gpu = [
  'torch==2.6.0+cu126',
  # flash-attn also has a torch dependency.
  "flash-attn",
]


[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",
    # Banned imports
    "TID",
    # Naming conventions
    "N",
    # print
    "T201"
]

# These rules are sensible and should be enabled at a later stage.
ignore = [
  # "B006",
  "B011",
  "UP008",
  "SIM117",
  "SIM118",
  "SIM102",
  "SIM401",
  # To ignore, not relevant for us
  "SIM108", # in case additional norm layer supports are added in future
  "N817", # we use heavy acronyms, e.g., allowing 'import LongModuleName as LMN' (LMN is accepted)
  "E731", # overly restrictive and less readable code
  "N812", # prevents us following the convention for importing torch.nn.functional as F
]

[tool.ruff.lint.flake8-tidy-imports.banned-api]
"numpy.ndarray".msg = "Do not use 'ndarray' to describe a numpy array type, it is a function. Use numpy.typing.NDArray or numpy.typing.NDArray[np.float32] for example"

[tool.ruff.format]
# Use Unix `\n` line endings for all files
line-ending = "lf"




[tool.uv]
# Most work is done a distributed filesystem, where hardlink is not always possible.
# Also, trying to resolve some permissions issue, see 44.
link-mode = "symlink"
# 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"

# The minimum version of uv required.
# It is tightly controlled because the format of uv.lock has changed 
# over revisions, causing reformats to happen without reason.
# Also, relatively recent versions are required to support workspaces.
required-version = ">=0.7.0"

# The supported environments
# TODO: add macos and windows (CPU only, for running tests)
environments = [
    "sys_platform == 'linux' and platform_machine == 'aarch64'",
    "sys_platform == 'linux' and platform_machine == 'x86_64'",
#    "sys_platform == 'darwin'",
]

# One can only have cpu or gpu.
conflicts = [
  [
    { extra = "cpu" },
    { extra = "gpu" },
  ],
]


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


[tool.pyrefly]
project-includes = ["src/"]
project-excludes = [
]

[tool.pyrefly.errors]
bad-argument-type = false
unsupported-operation = false
missing-attribute = false
no-matching-overload = false
bad-context-manager = false

# To do:
bad-assignment = false
bad-return = false
index-error = false
not-iterable = false
not-callable = false



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

[tool.uv.sources]
weathergen-common = { workspace = true }
weathergen-evaluate = { workspace = true }
weathergen-metrics = { workspace = true }
weathergen-readers-extra = { workspace = true }


flash-attn = [
# The build of Cathal O'Brien is not compatible with the libc build on santis.
# Hardcode the reference to the swiss cluster for the time being.
# TODO: open issue
#  { url = "https://github.com/cathalobrien/get-flash-attn/releases/download/v0.1-alpha/flash_attn-2.7.4+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
# This version was rebuilt locally on santis and uploaded.
  { url = "https://object-store.os-api.cci1.ecmwf.int/weathergenerator-dev/wheels/flash_attn-2.7.3-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
  { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiTRUE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux'  and platform_machine == 'x86_64'" },
#  { index = "pytorch-cpu", marker = "sys_platform == 'darwin'"},
]


torch = [
# Explicit pin for GPU
  { url = "https://download.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp312-cp312-linux_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" },
  { url = "https://download.pytorch.org/whl/cu126/torch-2.6.0%2Bcu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux"  and platform_machine == "x86_64"', extra="gpu" },
# Use the public repo for CPU versions.
  { index = "pytorch-cpu", marker = "sys_platform == 'linux'", extra="cpu"},
]

[tool.pytest.ini_options]
log_cli = true
log_cli_level = "INFO"
log_cli_format = "%(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s)"
log_cli_date_format = "%Y-%m-%d %H:%M:%S"


[tool.uv.workspace]
members = [
   "packages/common",
   "packages/evaluate",
   "packages/metrics",
   "packages/readers_extra",
]