pyproject.toml

#####################################################################
# Core Project metadata
#####################################################################
[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.11,<3.14"

license = "Apache-2.0"

classifiers = [
    "Programming Language :: Python :: 3",
    "Operating System :: OS Independent",
]
dynamic = ["version", "optional_dependencies"]

dependencies = [
    "onnx>=1.14.0",
    "warp-lang",
    "pandas",
    "nvtx",
    "treelib>=1.2.5",
    "numpy>=1.22.4",
    "torch>=2.4.0",
    "tqdm>=4.60.0",
    "requests>=2.32.2",
    "GitPython",
    "s3fs>=2023.5.0",
    "packaging>=24.2",
    "timm>=1.0.22",
    "xarray>=2025.6.1",
    "einops>=0.8.1",
    "h5py>=3.15.1",
    "cftime>=1.6.5",
    "jaxtyping>=0.3.3",
    "termcolor>=3.2.0",
    "zarr>=2.18.3",
    "hydra-core>=1.3.2",
    "tensordict>=0.10.0",
    "omegaconf>=2.3.0",
]

[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"



#####################################################################
# Flags for UV compatibility
#####################################################################

[tool.uv]
managed = true

[tool.uv.extra-build-dependencies]
torch-sparse = ["torch"]
torch-cluster = ["torch"]
torch-scatter = ["torch"]
earth2grid = ["setuptools", "torch"]

[[tool.uv.index]]
name = "nvidia"
url = "https://pypi.nvidia.com"
explicit = true  # Only use for packages that explicitly specify this index


[tool.uv.sources]
earth2grid = { url = "https://github.com/NVlabs/earth2grid/archive/main.tar.gz" }
cuml-cu13 = { index = "nvidia" }
pylibraft-cu13 = { index = "nvidia" }
cuml-cu12 = { index = "nvidia" } # TODO setup optional install groups for CUDA 12 installations
pylibraft-cu12 = { index = "nvidia" }

#####################################################################
# Flags Controlling the local build of physicsnemo
#####################################################################
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"


[tool.hatch.version]
path = "physicsnemo/__init__.py"

[tool.hatch.build.targets.wheel]
packages = ["physicsnemo"]

[tool.hatch.build.targets.sdist]
exclude = [
    "/.github",
    "/docs",
    "/examples",
    "/test",
    "/CODING_STANDARDS",
    "/benchmarks",
]


#####################################################################
# Local Development Requirements (pytest, etc)
#####################################################################

[dependency-groups]
dev = [
    "asv>=0.6.0",
    "coverage>=7.13.0",
    "import-linter>=2.7",
    "interrogate>=1.7.0",
    "onnx>=1.20.0",
    "onnxscript>=0.5.6",
    "pre-commit>=4.5.0",
    "pytest>=9.0.1",
    "pytest-dependency>=0.6.0",
    "pytest-timeout>=2.4.0",
    "ruff>=0.14.8",
    "pytest-testmon>=2.2.0",
]

#####################################################################
# Optional Dependency groups
#####################################################################

# These groups are self-referential.  So, they chain together as you
# move up the physicsnemo hierarchy.  A new extra dep in `physicsnemo/nn`
# will also be an extra dep in `physicsnemo/models`, but not the other
# direction.

[project.optional-dependencies]
utils-extras = [
    "wandb",
    "mlflow",
    "line_profiler",
    "vtk",
    "stl",
    "pylibraft-cu13", # Provides rmm
    "cupy-cuda13x",
]
mesh-extras = [
    "matplotlib>=3.10.7",
    "pyvista>=0.46.4",
]
nn-extras = [
    "cuml-cu13",
    "transformer_engine[pytorch]",
    "scipy",
    "natten",
    "nvidia-physicsnemo[utils-extras]",
]
model-extras = [
    "nvidia-physicsnemo[nn-extras]",
    "transformer_engine[pytorch]",
    "netCDF4",
    "pyvista",
    "vtk",
]
datapipes-extras = [
    "dask",
    "tfrecord",
    "nvidia-physicsnemo[model-extras]",
]

# Use case specific dependency groups.
gnns = [
    "torch_geometric",
    "torch_scatter",
    "torch_sparse",
    "torch_cluster",
    # "nvfuser",
    "nvidia-physicsnemo[model-extras]",
]
healpix = [
    "earth2grid",
]


#####################################################################
# Linting configuration
#####################################################################

[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 `E501` (line length violations),
# and `S311` (random number generators)
# and `F722` which breaks jaxtyping annotations
lint.ignore = ["E501", "S311", "F722"]

# 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"]



[project.entry-points."physicsnemo.models"]
# AFNO
AFNO = "physicsnemo.models.afno:AFNO"
DistributedAFNO = "physicsnemo.models.afno:DistributedAFNO"
ModAFNO = "physicsnemo.models.afno:ModAFNO"

# Diffusion UNets
SongUNet = "physicsnemo.models.diffusion_unets:SongUNet"
SongUNetPosEmbd = "physicsnemo.models.diffusion_unets:SongUNetPosEmbd"
SongUNetPosLtEmbd = "physicsnemo.models.diffusion_unets:SongUNetPosLtEmbd"
DhariwalUNet = "physicsnemo.models.diffusion_unets:DhariwalUNet"

# Diffusion wrappers
CorrDiffRegressionUNet = "physicsnemo.models.diffusion_unets:CorrDiffRegressionUNet"
StormCastUNet = "physicsnemo.models.diffusion_unets:StormCastUNet"

# DLWP
DLWP = "physicsnemo.models.dlwp:DLWP"

# DLWP HEALPix
HEALPixRecUNet = "physicsnemo.models.dlwp_healpix:HEALPixRecUNet"
HEALPixUNet = "physicsnemo.models.dlwp_healpix:HEALPixUNet"

# DPOT
DPOTNet = "physicsnemo.models.dpot:DPOTNet"

# Fengwu
Fengwu = "physicsnemo.models.fengwu:Fengwu"

# FIGConvNet
FIGConvUNet = "physicsnemo.models.figconvnet:FIGConvUNet"

# FNO
FNO = "physicsnemo.models.fno:FNO"

# GraphCast
GraphCastNet = "physicsnemo.models.graphcast:GraphCastNet"

# MeshGraphNets
MeshGraphNet = "physicsnemo.models.meshgraphnet:MeshGraphNet"
BiStrideMeshGraphNet = "physicsnemo.models.meshgraphnet:BiStrideMeshGraphNet"
MeshGraphKAN = "physicsnemo.models.meshgraphnet:MeshGraphKAN"
HybridMeshGraphNet = "physicsnemo.models.meshgraphnet:HybridMeshGraphNet"

# MLP
FullyConnected = "physicsnemo.models.mlp:FullyConnected"

# Pangu
Pangu = "physicsnemo.models.pangu:Pangu"

# Pix2Pix
Pix2Pix = "physicsnemo.models.pix2pix:Pix2Pix"
Pix2PixUnet = "physicsnemo.models.pix2pix:Pix2PixUnet"

# RNNs
One2ManyRNN = "physicsnemo.models.rnn:One2ManyRNN"
Seq2SeqRNN = "physicsnemo.models.rnn:Seq2SeqRNN"

# SRNNs
SRResNet = "physicsnemo.models.srrn:SRResNet"

# SwinRNN
SwinRNN = "physicsnemo.models.swinvrnn:SwinRNN"

# TopoDiff
TopoDiff = "physicsnemo.models.topodiff:TopoDiff"

# Transolver
Transolver = "physicsnemo.models.transolver:Transolver"

# VFGN
VFGN = "physicsnemo.models.vfgn:VFGNLearnedSimulator"