[build-system]
requires = ["setuptools", "setuptools-scm<9.0.0"]
build-backend = "setuptools.build_meta"

[project]
name = "nvidia-physicsnemo"
authors = [
  { name="NVIDIA PhysicsNeMo Team"},
]
description = "A deep learning framework for AI-driven multi-physics systems"
readme = "README.md"
requires-python = ">=3.10"
license = "Apache-2.0"
dependencies = [
    "certifi>=2023.7.22",
    "fsspec>=2023.1.0",
    "numpy>=1.22.4",
    "onnx>=1.14.0",
    "packaging>=24.2",
    "s3fs>=2023.5.0",
    "setuptools>=77.0.3",
    "timm>=1.0.0",
    "torch>=2.4.0",
    "tqdm>=4.60.0",
    "treelib>=1.2.5",
    "xarray>=2023.1.0",
    "zarr>=2.14.2",
]
classifiers = [
    "Programming Language :: Python :: 3",
    "Operating System :: OS Independent",
]
dynamic = ["version"]

[project.urls]
Homepage = "https://github.com/NVIDIA/physicsnemo"
Documentation = "https://docs.nvidia.com/physicsnemo/index.html#core"
Issues = "https://github.com/NVIDIA/physicsnemo/issues"
Changelog = "https://github.com/NVIDIA/physicsnemo/blob/main/CHANGELOG.md"

[project.optional-dependencies]
launch = [
    "hydra-core>=1.2.0",
    "termcolor>=2.1.1",
    "wandb>=0.13.7",
    "mlflow>=2.1.1",
    "pydantic>=2.4.2",
    "imageio>=2.28.1",
    "moviepy>=1.0.3",
]

dev = [
    "pytest>=6.0.0",
    "pyyaml>=6.0",
    "interrogate==1.5.0",
    "coverage==6.5.0",
    "ruff==0.12.5",
    "moto[s3]>=5.0.28",
    "pre-commit>=4.0.0"
]

makani = [
    # TODO(akamenev): PyPI does not allow direct URL deps, update once Makani is in PyPI
    # "makani @ git+https://github.com/NVIDIA/modulus-makani.git@v0.1.0",
    "torch-harmonics>=0.6.5,<0.7.1",
    "tensorly>=0.8.1",
    "tensorly-torch>=0.4.0",
]

fignet = [
    "jaxtyping>=0.2",
    "torch_scatter>=2.1",
    "torchinfo>=1.8",
    "warp-lang>=1.0",
    "webdataset>=0.2",
]

storage = [
    "multi-storage-client[boto3]>=0.33.0",
]

shardtensor = [
    "wrapt>=1.15.0",
]

natten = [
    "natten",
    "einops",
]

all = [
    "nvidia_dali_cuda120>=1.35.0",
    "h5py>=3.7.0",
    "netcdf4>=1.6.3",
    "ruamel.yaml>=0.17.22",
    "scikit-learn>=1.0.2",
    "scikit-image>=0.24.0",
    "warp-lang>=1.0",
    "vtk>=9.2.6",
    "pyvista>=0.40.1",
    "cftime>=1.6.2",
    "einops>=0.7.0",
    "pyspng>=0.1.0",
    "shapely>=2.0.6",
    "pytz>=2023.3",
    "nvtx>=0.2.8",
    "nvidia-physicsnemo[launch]",
    "nvidia-physicsnemo[dev]",
    "nvidia-physicsnemo[makani]",
    "nvidia-physicsnemo[fignet]",
    "nvidia-physicsnemo[storage]",
]


[tool.setuptools.dynamic]
version = {attr = "physicsnemo.__version__"}

[tool.setuptools.packages.find]
include = ["physicsnemo", "physicsnemo.*"]

[tool.ruff]
# Enable flake8/pycodestyle (`E`), Pyflakes (`F`), flake8-bandit (`S`),
# isort (`I`), and performance 'PERF' rules.
lint.select = ["E", "F", "S", "I", "PERF"]
lint.fixable = ["I"]

# Never enforce `E402`, `E501` (line length violations),
# and `S311` (random number generators)
lint.ignore = ["E501", "S311"]

# Exclude the docs and experimental folders (this applies to both lint and format)
exclude = ["docs", "physicsnemo/experimental"]

[tool.ruff.lint.per-file-ignores]
# Ignore `F401` (import violations) in all `__init__.py` files, and in `docs/*.py`.
"__init__.py" = ["F401"]
"docs/*.py" = ["F401"]

# Ignore `S101` (assertions) in all `test` files.
"test/*.py" = ["S101"]

# ==== UV configuration ====
[tool.uv]
no-build-isolation-package = ["torch_scatter"]
managed = false

[project.entry-points."physicsnemo.models"]
AFNO = "physicsnemo.models.afno:AFNO"
DLWP = "physicsnemo.models.dlwp:DLWP"
FNO = "physicsnemo.models.fno:FNO"
GraphCastNet = "physicsnemo.models.graphcast:GraphCastNet"
MeshGraphNet = "physicsnemo.models.meshgraphnet:MeshGraphNet"
FullyConnected = "physicsnemo.models.mlp:FullyConnected"
Pix2Pix = "physicsnemo.models.pix2pix:Pix2Pix"
One2ManyRNN = "physicsnemo.models.rnn:One2ManyRNN"
SRResNet = "physicsnemo.models.srrn:SRResNet"
Pangu = "physicsnemo.models.pangu:Pangu"
Fengwu = "physicsnemo.models.fengwu:Fengwu"
SwinRNN = "physicsnemo.models.swinvrnn:SwinRNN"
EDMPrecondSR = "physicsnemo.models.diffusion:EDMPrecondSR"
UNet = "physicsnemo.models.diffusion:UNet"