nequix/pyproject.toml at main · atomicarchitects/nequix · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
[project]
name = "nequix"
version = "0.4.5"
description = "Nequix source code"
readme = "README.md"
authors = [
    { name = "Teddy Koker", email = "teddy.koker@gmail.com" }
]
requires-python = ">=3.10"
license = { file = "LICENSE" }

dependencies = [
    "jax[cuda12]>=0.4.34; sys_platform == 'linux'",
    "jax>=0.4.34; sys_platform == 'darwin'",
    "e3nn-jax>=0.20.8",
    "jraph>=0.0.6.dev0",
    "equinox>=0.11.11",
    "optax>=0.2.5",
    "ase>=3.24.0",
    "tqdm>=4.67.1",
    "wandb>=0.19.11",
    "pyyaml>=6.0.2",
    "matscipy>=1.1.1",
    "h5py>=3.14.0",
    "wandb-osh>=1.2.2",
    "cloudpickle>=3.1.1",
    "ase-db-backends",
]

[project.urls]
Homepage = "https://pypi.org/project/nequix/"
Repository = "https://github.com/atomicarchitects/nequix"


[project.scripts]
nequix_train = "nequix.train:main"

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"


[tool.hatch.build]
include = ["nequix/*"]

[dependency-groups]
dev = [
    "ipython>=8.36.0",
    "jax-md",
    "pandas>=2.3.1",
    "pre-commit>=4.1.0",
    "pytest>=8.3.5",
    "pytest-benchmark>=5.1.0",
]

[tool.ruff]
line-length = 100

[tool.uv.sources]
ase-db-backends = { git = "https://gitlab.com/ase/ase-db-backends.git" }
jax-md = { git = "https://github.com/jax-md/jax-md.git" }

[project.optional-dependencies]
torch = [
  "e3nn>=0.5.8",
  "openequivariance>=0.6.4",
  "torch>=2.7.0",
  "torch-geometric>=2.6.1",
  "setuptools",
]
torch-sim = [
  "nequix[torch]",
  "torch-sim-atomistic>=0.5.2; python_version >= '3.12'",
  "torch>=2.8.0; python_version >= '3.12'",
]
oeq = [
  # You need to install openequivariance_extjax separately
  # uv pip install openequivariance_extjax --no-build-isolation
  "openequivariance[jax]>=0.6.4",
]
pft = [
    "phonopy>=2.43.1",
]