pyproject.toml

[build-system]
requires = [
    "setuptools>=61.0.0",
    "setuptools-scm>=8.0",
    "wheel",
]
build-backend = "setuptools.build_meta"

[project]
name = "equitrain"
dynamic = ["version"]
description = "A training framework for machine learning interatomic potentials (MLIPs)"
authors = [
    {name = "Carlos Madariaga", email = "carlos.madariaga@bam.de"},
    {name = "Jose M. Pizarro", email = "jose.pizarro-blanco@bam.de"},
    {name = "Janosh Riebesell"},
    {name = "Philipp Benner", email = "philipp.benner@bam.de"},
]
readme = "README.md"
requires-python = ">=3.10"
keywords = ["MLIP", "machine learning", "materials"]
license = { file = "LICENSE" }
classifiers = [
    "Programming Language :: Python :: 3",
    "License :: OSI Approved :: MIT License",
    "Operating System :: OS Independent"
]
dependencies = [
    "matscipy",
    "lightning",
    "beautifulsoup4",
    "pqdm",
    "ase",
    "h5py",
    "pymatgen",
    "tqdm",
]

[project.optional-dependencies]
dev = [
  "mypy==1.0.1",
  "pre-commit",
  "ruff==0.15.11",
  "pytest",
]
docu = [
  "mkdocs",
  "mkdocs-material",
  "pymdown-extensions",
  "mkdocs-click",
]
torch = [
  "torch>=2.0.0",
  "torch_geometric>=2.4.0",
  "torch_ema>=0.3.0",
  "accelerate>=0.28.0",
]
jax = [
  "jax>=0.4",
  "jaxlib>=0.4",
  "flax",
  "jraph",
]
ani = [
  "torch>=2.0.0",
  "torch_geometric>=2.4.0",
  "torch_ema>=0.3.0",
  "accelerate>=0.28.0",
  "torchani>=2.2",  # TorchANI for ANI models
  "ase>=3.22.0",    # Atomic Simulation Environment for symbol/structure utilities
]
orb = [
  "torch>=2.0.0",
  "torch_geometric>=2.4.0",
  "torch_ema>=0.3.0",
  "accelerate>=0.28.0",
  "orb-models>=3.0",  # ORB universal interatomic potentials
  "cuml-cu11; sys_platform == 'linux'",  # Optional GPU acceleration for neighbor graphs
]
m3gnet = [
  "torch>=2.0.0",
  "torch_geometric>=2.4.0",
  "torch_ema>=0.3.0",
  "accelerate>=0.28.0",
  "matgl>=1.0.0",  # Materials Graph Library for M3GNet models
  "dgl>=1.0.0",    # Deep Graph Library required by MatGL
]
mace = [
  "torch>=2.0.0",
  "torch_geometric>=2.4.0",
  "torch_ema>=0.3.0",
  "accelerate>=0.28.0",
  "mace-torch>=0.3.6",  # Torch implementation of MACE potentials
]
mace-jax = [
  "mace-jax>=0.2.0",  # JAX implementation of MACE potentials
  "jax>=0.4",         # Core JAX runtime
  "jaxlib>=0.4",      # JAX XLA binaries
  "flax",
  "jraph",
]

[project.scripts]
equitrain            = "equitrain.scripts.equitrain:main"
equitrain-hdf5-benchmark = "equitrain.scripts.equitrain_hdf5_benchmark:main"
equitrain-evaluate   = "equitrain.scripts.equitrain_evaluate:main"
equitrain-export     = "equitrain.scripts.equitrain_export:main"
equitrain-hdf5-info  = "equitrain.scripts.equitrain_hdf5_info:main"
equitrain-inspect    = "equitrain.scripts.equitrain_inspect:main"
equitrain-predict    = "equitrain.scripts.equitrain_predict:main"
equitrain-preprocess = "equitrain.scripts.equitrain_preprocess:main"

[tool.ruff]
required-version = "==0.15.11"
# Exclude a variety of commonly ignored directories.
include = ["equitrain/**/*.py", "tests/**/*.py"]
exclude = [
    ".bzr",
    ".direnv",
    ".eggs",
    ".git",
    ".git-rewrite",
    ".hg",
    ".mypy_cache",
    ".nox",
    ".pants.d",
    ".pytype",
    ".ruff_cache",
    ".svn",
    ".tox",
    ".venv",
    "__pypackages__",
    "_build",
    "buck-out",
    "build",
    "dist",
    "node_modules",
    "venv",
    "dependencies",
]
# Same as Black.
line-length = 88
indent-width = 4

[tool.ruff.lint]
select = [
    "E", # pycodestyle
    "PL", # pylint
    "F", # Pyflakes
    "UP", # pyupgrade
    "I", # isort
]
ignore = [
    "F401", # Module imported but unused
    "E501", # Line too long ({width} > {limit} characters)
    "E701", # Multiple statements on one line (colon)
    "E731", # Do not assign a lambda expression, use a def
    "E402",  # Module level import not at top of file
    "PLR0911", # Too many return statements
    "PLR0912", # Too many branches
    "PLR0913", # Too many arguments in function definition
    "PLR0915", # Too many statements
    "PLR2004", # Magic value used instead of constant
    "PLW0603", # Using the global statement
    "PLW2901", # redefined-loop-name
    "PLR1714", # consider-using-in
    "PLR5501", # else-if-used
    "PLC0415", # import should be at top-level
]
fixable = ["ALL"]
# Allow unused variables when underscore-prefixed.
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
# this is entirely optional, you can remove this if you wish to
[tool.ruff.format]
# use single quotes for strings.
quote-style = "single"
# indent with spaces, rather than tabs.
indent-style = "space"
# Like Black, respect magic trailing commas.
skip-magic-trailing-comma = false
# Like Black, automatically detect the appropriate line ending.
line-ending = "auto"

[tool.mypy]
strict = false
ignore_missing_imports = true
follow_imports = "silent"
disable_error_code = "import, annotation-unchecked"

[tool.setuptools.packages.find]
where = ["."]
include = ["equitrain"]

[tool.equitrain]
conda_env_file = "environment.yml"

[tool.setuptools_scm]
write_to = "equitrain/_version.py"

[tool.pytest.ini_options]
filterwarnings = [
    "ignore:Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected:UserWarning",
    "ignore:`torch_geometric.distributed` has been deprecated since 2.7.0:DeprecationWarning",
    "ignore:This process \\(pid=.*\\) is multi-threaded, use of fork\\(\\):DeprecationWarning",
    "ignore:Creating NamedTuple classes using keyword arguments is deprecated:DeprecationWarning",
]