pyproject.toml

# SPDX-FileCopyrightText: 2026 Samudra Authors
#
# SPDX-License-Identifier: Apache-2.0

[build-system]
build-backend = "setuptools.build_meta"
requires = [ "setuptools", "wheel" ]

[project]
name = "samudra"
version = "1.0"
description = "Train and evaluate ocean model emulators"
readme = "README.md"
keywords = [ "climate", "emulation", "machine learning", "ocean" ]
license = { file = "LICENSE" }
requires-python = ">=3.12"
classifiers = [
  "Development Status :: 4 - Beta",
  "Intended Audience :: Science/Research",
  "License :: OSI Approved :: Apache Software License",
  "Programming Language :: Python :: 3 :: Only",
  "Programming Language :: Python :: 3.12",
  "Programming Language :: Python :: 3.13",
  "Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
  "aiohttp>=3.11.13",
  "cftime>=1.6.4.post1",
  "dacite>=1.9.2",
  "dask>=2025.2",
  "einops>=0.7",
  "jaxtyping>=0.3",
  "matplotlib>=3.10.1",
  "microsoft-aurora>=1.8",
  "numpy>=1.26.4,<2",             # There is some incompatibility between a dependency and numpy 2+
  "perceiver-pytorch>=0.8.8",
  "pydantic-settings>=2.8.1",
  "pytest-xdist[psutil]>=3.6.1",
  "pyyaml>=6.0.2",
  "requests>=2.32.3",
  "s3fs>=2025.5.1",
  "scipy>=1.15.2",
  "skypilot[aws,lambda]>=0.10",
  "sqlalchemy>=2",                # skypilot requires sqlalchemy 2.0 but doesn't include that requirement
  "torch>=2.9.0a0,<2.13",         # Matches PhysicsNeMo container builds
  "torchinfo>=1.8",
  "typing-extensions>=4.15",
  "wandb>=0.19.8",
  "xarray[io]>=2025.1.2",
  "xarray-einstats[einops]>=0.8",
  "xesmf>=0.8.10",
  "zarr<3",
]
optional-dependencies.cuda = [
  "flash-attn>=2.7.4.post1,<3",
  "flash-perceiver>=0.2",
  "torchvision>=0.24.0a0",
]

[dependency-groups]
dev = [
  "beartype>=0.20",
  "cartopy>=0.24.1",
  "cmocean>=4.0.3",
  "coiled>=1.99",
  "filelock>=3.17",
  "fsspec>=2025.3",
  "hypothesis[numpy]>=6.127.2",
  "ipykernel>=6.29.5",
  "jupyter>=1.1.1",
  "memray>=1.17.1",
  "mypy>=1.15",
  "nc-time-axis>=1.4.1",
  "pandas-stubs>=2.3.0.250703",
  "py-spy>=0.4",
  "pytest>=8.3.4",
  "pytest-benchmark[histogram]>=5.1",
  "pytest-profiling>=1.8.1",
  "s3fs>=2025.3",
  "scalene>=1.5.51",
  "snakeviz>=2.2.2",
  "tenacity>=9",
  "tqdm>=4.67.1",
  "types-pyyaml>=6.0.12.20241230",
  "types-tqdm>=4.67.0.20250516",
  "xarrayutils>=2.0.1",
]
docs = [
  "mkdocs-material>=9.5",
  "mkdocstrings[python]>=0.27",
  "zensical",
]

[tool.ruff]
lint.select = [ "B006", "D", "E", "F", "I", "UP", "W" ]
# F722: https://docs.kidger.site/jaxtyping/faq/#flake8-or-ruff-are-throwing-an-error
lint.ignore = [
  "D1",
  "D200",
  "D205",
  "D212",
  "E203",
  "E402",
  "E501",  # We use ruff-format which will already wrap lines to a reasonable length.
  "F541",
  "F722",
  "UP038", # This rule is deprecated.
  "W293",
]
lint.per-file-ignores."*/__init__.py" = [ "F401" ]
lint.per-file-ignores."scripts/*" = [ "D" ]
# F811 covers the pytest pattern of importing fixtures and reusing the name as a parameter.
lint.per-file-ignores."test_*.py" = [ "D", "F811" ]
lint.isort.split-on-trailing-comma = false
lint.pydocstyle.convention = "google"

[tool.pytest]
ini_options.markers = [
  "manual: marks tests as manually run (deselect with '-m \"not manual\"').",
  "cuda: marks tests that need a CUDA-enabled device to run (deselect with '-m \"not cuda\"').",
  "only_configs: used to mark tests that only run on specific configs (not used on command line)",
  "all_configs: used to mark tests that run on all configs (not used on command line)",
]
ini_options.pythonpath = [ "src", "tests" ]
ini_options.testpaths = [ "tests" ]
ini_options.addopts = "--jaxtyping-packages=data,datasets,beartype.beartype --strict-markers"

[tool.mypy]
check_untyped_defs = true
exclude = [ "^data/" ]
overrides = [
  { module = [
    "cftime.*",
    "scipy.*",
    "fsspec.*",
    "perceiver_pytorch.*",
    "aurora.*",
  ], follow_untyped_imports = true },
  { module = [
    "flash_perceiver.*",
    "flash_attn.*",
  ], ignore_missing_imports = true, follow_untyped_imports = true },
]

[tool.uv]
extra-build-dependencies.flash-attn = [ { requirement = "torch", match-runtime = true } ]
extra-build-variables.flash-attn = { FLASH_ATTENTION_SKIP_CUDA_BUILD = "TRUE" }

[tool.uv.sources]
torch = [
  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
]
torchvision = [
  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
]

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