[tool.poetry]
name = "lagrangebench"
version = "0.0.2"
description = "LagrangeBench: A Lagrangian Fluid Mechanics Benchmarking Suite"
authors = [
    "Artur Toshev <artur.toshev@tum.de>",
    "Gianluca Galletti <g.galletti@tum.de>",
]
license = "MIT"
readme = "README.md"

[tool.poetry.dependencies]
python = ">=3.9,<=3.11"
cloudpickle = ">=2.2.1"
h5py = ">=3.9.0"
PyYAML = ">=6.0"
numpy = ">=1.24.4"
wandb = ">=0.15.11"
pyvista = ">=0.42.2"
jax = "0.4.18"
jaxlib = "0.4.18"
dm-haiku = ">=0.0.10"
e3nn-jax = "0.20.2"
jax-md = "0.2.8"
jmp = "^0.0.4"
jraph = "0.0.6.dev0"
optax = "0.1.7"
ott-jax = "^0.4.2"
matscipy = "^0.8.0"
torch = {version = "2.1.0+cpu", source = "torchcpu"}
wget = "^3.2"

[tool.poetry.group.dev.dependencies]
pre-commit = ">=3.3.1"
pytest = ">=7.3.1"
black = ">=23.3.0"
ruff = "0.0.265"
ipykernel = ">=6.25.1"

[tool.poetry.group.docs.dependencies]
sphinx = "^7.2.6"
sphinx-rtd-theme = "^1.3.0"

[[tool.poetry.source]]
name = "torchcpu"
url = "https://download.pytorch.org/whl/cpu"
priority = "explicit"

[tool.ruff]
ignore = ["F811", "E402"]
exclude = [
    ".git",
    ".venv",
    "venv",
]
line-length = 88

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
