aimnet 0.0.1__tar.gz → 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- aimnet-0.1.0/.gitignore +171 -0
- aimnet-0.1.0/PKG-INFO +308 -0
- aimnet-0.1.0/README.md +249 -0
- aimnet-0.1.0/aimnet/__init__.py +7 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/base.py +24 -8
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/calculators/__init__.py +4 -4
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/calculators/aimnet2ase.py +19 -6
- aimnet-0.1.0/aimnet/calculators/calculator.py +1080 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/calculators/model_registry.py +2 -5
- aimnet-0.1.0/aimnet/calculators/model_registry.yaml +71 -0
- aimnet-0.1.0/aimnet/cli.py +78 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/config.py +8 -9
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/data/sgdataset.py +23 -22
- aimnet-0.1.0/aimnet/kernels/__init__.py +66 -0
- aimnet-0.1.0/aimnet/kernels/conv_sv_2d_sp_wp.py +478 -0
- aimnet-0.1.0/aimnet/models/__init__.py +14 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/models/aimnet2.py +19 -22
- aimnet-0.1.0/aimnet/models/base.py +219 -0
- aimnet-0.1.0/aimnet/models/convert.py +30 -0
- aimnet-0.1.0/aimnet/models/utils.py +735 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/modules/__init__.py +1 -1
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/modules/aev.py +49 -48
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/modules/core.py +14 -13
- aimnet-0.1.0/aimnet/modules/lr.py +648 -0
- aimnet-0.1.0/aimnet/modules/ops.py +537 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/nbops.py +105 -15
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/ops.py +90 -18
- aimnet-0.1.0/aimnet/train/export_model.py +226 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/train/loss.py +7 -7
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/train/metrics.py +5 -6
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/train/train.py +4 -1
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/train/utils.py +42 -13
- aimnet-0.1.0/pyproject.toml +218 -0
- aimnet-0.0.1/PKG-INFO +0 -78
- aimnet-0.0.1/README.md +0 -49
- aimnet-0.0.1/aimnet/calculators/calculator.py +0 -320
- aimnet-0.0.1/aimnet/calculators/model_registry.yaml +0 -33
- aimnet-0.0.1/aimnet/calculators/nb_kernel_cpu.py +0 -222
- aimnet-0.0.1/aimnet/calculators/nb_kernel_cuda.py +0 -217
- aimnet-0.0.1/aimnet/calculators/nbmat.py +0 -220
- aimnet-0.0.1/aimnet/cli.py +0 -22
- aimnet-0.0.1/aimnet/models/__init__.py +0 -2
- aimnet-0.0.1/aimnet/models/base.py +0 -51
- aimnet-0.0.1/aimnet/modules/lr.py +0 -243
- aimnet-0.0.1/aimnet/train/__init__.py +0 -0
- aimnet-0.0.1/aimnet/train/pt2jpt.py +0 -81
- aimnet-0.0.1/pyproject.toml +0 -143
- {aimnet-0.0.1 → aimnet-0.1.0}/LICENSE +0 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/calculators/aimnet2pysis.py +0 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/constants.py +0 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/data/__init__.py +0 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/dftd3_data.pt +0 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/models/aimnet2.yaml +0 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/models/aimnet2_dftd3_wb97m.yaml +0 -0
- {aimnet-0.0.1/aimnet → aimnet-0.1.0/aimnet/train}/__init__.py +0 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/train/calc_sae.py +0 -0
- {aimnet-0.0.1 → aimnet-0.1.0}/aimnet/train/default_train.yaml +0 -0
aimnet-0.1.0/.gitignore
ADDED
|
@@ -0,0 +1,171 @@
|
|
|
1
|
+
docs/source
|
|
2
|
+
aimnet/calculators/assets/
|
|
3
|
+
|
|
4
|
+
scripts/compile/*.pt
|
|
5
|
+
scripts/compile/*.jpt
|
|
6
|
+
scripts/compile/*.sae
|
|
7
|
+
scripts/compile/aimnet2_b973c_d3.yaml
|
|
8
|
+
scripts/compile/aimnet2_wb97m_d3.yaml
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
# From https://raw.githubusercontent.com/github/gitignore/main/Python.gitignore
|
|
12
|
+
|
|
13
|
+
# Byte-compiled / optimized / DLL files
|
|
14
|
+
__pycache__/
|
|
15
|
+
*.py[cod]
|
|
16
|
+
*$py.class
|
|
17
|
+
|
|
18
|
+
# C extensions
|
|
19
|
+
*.so
|
|
20
|
+
|
|
21
|
+
# Distribution / packaging
|
|
22
|
+
.Python
|
|
23
|
+
build/
|
|
24
|
+
develop-eggs/
|
|
25
|
+
dist/
|
|
26
|
+
downloads/
|
|
27
|
+
eggs/
|
|
28
|
+
.eggs/
|
|
29
|
+
lib/
|
|
30
|
+
lib64/
|
|
31
|
+
parts/
|
|
32
|
+
sdist/
|
|
33
|
+
var/
|
|
34
|
+
wheels/
|
|
35
|
+
share/python-wheels/
|
|
36
|
+
*.egg-info/
|
|
37
|
+
.installed.cfg
|
|
38
|
+
*.egg
|
|
39
|
+
MANIFEST
|
|
40
|
+
|
|
41
|
+
# PyInstaller
|
|
42
|
+
# Usually these files are written by a python script from a template
|
|
43
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
44
|
+
*.manifest
|
|
45
|
+
*.spec
|
|
46
|
+
|
|
47
|
+
# Installer logs
|
|
48
|
+
pip-log.txt
|
|
49
|
+
pip-delete-this-directory.txt
|
|
50
|
+
|
|
51
|
+
# Unit test / coverage reports
|
|
52
|
+
htmlcov/
|
|
53
|
+
.tox/
|
|
54
|
+
.nox/
|
|
55
|
+
.coverage
|
|
56
|
+
.coverage.*
|
|
57
|
+
.cache
|
|
58
|
+
nosetests.xml
|
|
59
|
+
coverage.xml
|
|
60
|
+
*.cover
|
|
61
|
+
*.py,cover
|
|
62
|
+
.hypothesis/
|
|
63
|
+
.pytest_cache/
|
|
64
|
+
cover/
|
|
65
|
+
|
|
66
|
+
# Translations
|
|
67
|
+
*.mo
|
|
68
|
+
*.pot
|
|
69
|
+
|
|
70
|
+
# Django stuff:
|
|
71
|
+
*.log
|
|
72
|
+
local_settings.py
|
|
73
|
+
db.sqlite3
|
|
74
|
+
db.sqlite3-journal
|
|
75
|
+
|
|
76
|
+
# Flask stuff:
|
|
77
|
+
instance/
|
|
78
|
+
.webassets-cache
|
|
79
|
+
|
|
80
|
+
# Scrapy stuff:
|
|
81
|
+
.scrapy
|
|
82
|
+
|
|
83
|
+
# Sphinx documentation
|
|
84
|
+
docs/_build/
|
|
85
|
+
|
|
86
|
+
# PyBuilder
|
|
87
|
+
.pybuilder/
|
|
88
|
+
target/
|
|
89
|
+
|
|
90
|
+
# Jupyter Notebook
|
|
91
|
+
.ipynb_checkpoints
|
|
92
|
+
|
|
93
|
+
# IPython
|
|
94
|
+
profile_default/
|
|
95
|
+
ipython_config.py
|
|
96
|
+
|
|
97
|
+
# pyenv
|
|
98
|
+
# For a library or package, you might want to ignore these files since the code is
|
|
99
|
+
# intended to run in multiple environments; otherwise, check them in:
|
|
100
|
+
# .python-version
|
|
101
|
+
|
|
102
|
+
# pipenv
|
|
103
|
+
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
|
104
|
+
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
|
105
|
+
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
|
106
|
+
# install all needed dependencies.
|
|
107
|
+
#Pipfile.lock
|
|
108
|
+
|
|
109
|
+
.pdm.toml
|
|
110
|
+
|
|
111
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
|
112
|
+
__pypackages__/
|
|
113
|
+
|
|
114
|
+
# Celery stuff
|
|
115
|
+
celerybeat-schedule
|
|
116
|
+
celerybeat.pid
|
|
117
|
+
|
|
118
|
+
# SageMath parsed files
|
|
119
|
+
*.sage.py
|
|
120
|
+
|
|
121
|
+
# Environments
|
|
122
|
+
.env
|
|
123
|
+
.venv
|
|
124
|
+
env/
|
|
125
|
+
venv/
|
|
126
|
+
ENV/
|
|
127
|
+
env.bak/
|
|
128
|
+
venv.bak/
|
|
129
|
+
|
|
130
|
+
# Spyder project settings
|
|
131
|
+
.spyderproject
|
|
132
|
+
.spyproject
|
|
133
|
+
|
|
134
|
+
# Rope project settings
|
|
135
|
+
.ropeproject
|
|
136
|
+
|
|
137
|
+
# mkdocs documentation
|
|
138
|
+
/site
|
|
139
|
+
|
|
140
|
+
# mypy
|
|
141
|
+
.mypy_cache/
|
|
142
|
+
.dmypy.json
|
|
143
|
+
dmypy.json
|
|
144
|
+
|
|
145
|
+
# Pyre type checker
|
|
146
|
+
.pyre/
|
|
147
|
+
|
|
148
|
+
# pytype static type analyzer
|
|
149
|
+
.pytype/
|
|
150
|
+
|
|
151
|
+
# Cython debug symbols
|
|
152
|
+
cython_debug/
|
|
153
|
+
|
|
154
|
+
# Vscode config files
|
|
155
|
+
.vscode/
|
|
156
|
+
|
|
157
|
+
# Vscode history
|
|
158
|
+
.history/
|
|
159
|
+
|
|
160
|
+
# Cursor IDE
|
|
161
|
+
.cursor/
|
|
162
|
+
|
|
163
|
+
# Claude Code
|
|
164
|
+
.claude/
|
|
165
|
+
|
|
166
|
+
# PyCharm
|
|
167
|
+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
|
168
|
+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
|
169
|
+
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
|
170
|
+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
|
171
|
+
.idea/
|
aimnet-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,308 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: aimnet
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: AIMNet Machine Learned Interatomic Potential
|
|
5
|
+
Project-URL: Homepage, https://github.com/isayevlab/aimnetcentral
|
|
6
|
+
Project-URL: Documentation, https://isayevlab.github.io/aimnetcentral/
|
|
7
|
+
Project-URL: Repository, https://github.com/isayevlab/aimnetcentral
|
|
8
|
+
Author-email: Roman Zubatyuk <zubatyuk@gmail.com>
|
|
9
|
+
License: MIT License
|
|
10
|
+
|
|
11
|
+
Copyright (c) 2024, Roman Zubatyuk
|
|
12
|
+
|
|
13
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
14
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
15
|
+
in the Software without restriction, including without limitation the rights
|
|
16
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
17
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
18
|
+
furnished to do so, subject to the following conditions:
|
|
19
|
+
|
|
20
|
+
The above copyright notice and this permission notice shall be included in all
|
|
21
|
+
copies or substantial portions of the Software.
|
|
22
|
+
|
|
23
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
24
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
25
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
26
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
27
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
28
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
29
|
+
SOFTWARE.
|
|
30
|
+
License-File: LICENSE
|
|
31
|
+
Keywords: computational chemistry,deep learning,interatomic potential,machine learning,molecular dynamics
|
|
32
|
+
Classifier: Intended Audience :: Science/Research
|
|
33
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
34
|
+
Classifier: Operating System :: OS Independent
|
|
35
|
+
Classifier: Programming Language :: Python :: 3
|
|
36
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
37
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
38
|
+
Classifier: Topic :: Scientific/Engineering :: Chemistry
|
|
39
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
40
|
+
Requires-Python: >=3.11
|
|
41
|
+
Requires-Dist: click>=8.1.7
|
|
42
|
+
Requires-Dist: h5py>=3.12.1
|
|
43
|
+
Requires-Dist: jinja2>=3.1.6
|
|
44
|
+
Requires-Dist: numpy
|
|
45
|
+
Requires-Dist: nvalchemi-toolkit-ops>=0.2
|
|
46
|
+
Requires-Dist: pyyaml>=6.0.2
|
|
47
|
+
Requires-Dist: requests>=2.32.3
|
|
48
|
+
Requires-Dist: torch>=2.4
|
|
49
|
+
Requires-Dist: warp-lang>=1.11
|
|
50
|
+
Provides-Extra: ase
|
|
51
|
+
Requires-Dist: ase==3.27.0; extra == 'ase'
|
|
52
|
+
Provides-Extra: pysis
|
|
53
|
+
Requires-Dist: pysisyphus; extra == 'pysis'
|
|
54
|
+
Provides-Extra: train
|
|
55
|
+
Requires-Dist: omegaconf>=2.3.0; extra == 'train'
|
|
56
|
+
Requires-Dist: pytorch-ignite>=0.5.1; extra == 'train'
|
|
57
|
+
Requires-Dist: wandb>=0.18.5; extra == 'train'
|
|
58
|
+
Description-Content-Type: text/markdown
|
|
59
|
+
|
|
60
|
+
[](https://github.com/isayevlab/aimnetcentral/releases)
|
|
61
|
+
[](https://www.python.org/)
|
|
62
|
+
[](https://github.com/isayevlab/aimnetcentral/actions/workflows/main.yml?query=branch%3Amain)
|
|
63
|
+
[](https://codecov.io/gh/isayevlab/aimnetcentral)
|
|
64
|
+
[](https://github.com/isayevlab/aimnetcentral/blob/main/LICENSE)
|
|
65
|
+
|
|
66
|
+
- **Github repository**: <https://github.com/isayevlab/aimnetcentral/>
|
|
67
|
+
- **Documentation** <https://isayevlab.github.io/aimnetcentral/>
|
|
68
|
+
|
|
69
|
+
# AIMNet2 : ML interatomic potential for fast and accurate atomistic simulations
|
|
70
|
+
|
|
71
|
+
## Key Features
|
|
72
|
+
|
|
73
|
+
- Accurate and Versatile: AIMNet2 excels at modeling neutral, charged, organic, and elemental-organic systems.
|
|
74
|
+
- Flexible Interfaces: Use AIMNet2 through convenient calculators for popular simulation packages like ASE and PySisyphus.
|
|
75
|
+
- Flexible Long-Range Interactions: Optionally employ the Damped-Shifted Force (DSF) or Ewald summation Coulomb models for accurate calculations in large or periodic systems.
|
|
76
|
+
|
|
77
|
+
## Requirements
|
|
78
|
+
|
|
79
|
+
### Python Version
|
|
80
|
+
|
|
81
|
+
AIMNet2 requires **Python 3.11 or 3.12**.
|
|
82
|
+
|
|
83
|
+
### GPU Support (Optional)
|
|
84
|
+
|
|
85
|
+
AIMNet2 works on CPU out of the box. For GPU acceleration:
|
|
86
|
+
|
|
87
|
+
- **CUDA GPU**: Install PyTorch with CUDA support from [pytorch.org](https://pytorch.org/get-started/locally/)
|
|
88
|
+
- **compile_mode**: Requires CUDA for ~5x MD speedup (see Performance Optimization)
|
|
89
|
+
|
|
90
|
+
Example PyTorch installation with CUDA 12.4:
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
pip install torch --index-url https://download.pytorch.org/whl/cu124
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
## Available Models
|
|
97
|
+
|
|
98
|
+
| Model | Elements | Description |
|
|
99
|
+
| --------------- | --------------------------------------------- | --------------------------------------------- |
|
|
100
|
+
| `aimnet2` | H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I | wB97M-D3 (default) |
|
|
101
|
+
| `aimnet2_b973c` | H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I | B97-3c functional |
|
|
102
|
+
| `aimnet2_2025` | H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I | B97-3c + improved intermolecular interactions |
|
|
103
|
+
| `aimnet2nse` | H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I | Open-shell chemistry |
|
|
104
|
+
| `aimnet2pd` | H, B, C, N, O, F, Si, P, S, Cl, Se, Br, Pd, I | Palladium-containing systems |
|
|
105
|
+
|
|
106
|
+
_Each model has ensemble members (append \_0 to \_3). Ensemble averaging recommended for production use._
|
|
107
|
+
|
|
108
|
+
## Installation
|
|
109
|
+
|
|
110
|
+
### Basic Installation
|
|
111
|
+
|
|
112
|
+
Install from PyPI:
|
|
113
|
+
|
|
114
|
+
```bash
|
|
115
|
+
pip install aimnet
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
### Optional Features
|
|
119
|
+
|
|
120
|
+
AIMNet2 provides optional extras for different use cases:
|
|
121
|
+
|
|
122
|
+
**ASE Calculator** (for atomistic simulations with ASE):
|
|
123
|
+
|
|
124
|
+
```bash
|
|
125
|
+
pip install "aimnet[ase]"
|
|
126
|
+
```
|
|
127
|
+
|
|
128
|
+
**PySisyphus Calculator** (for reaction path calculations):
|
|
129
|
+
|
|
130
|
+
```bash
|
|
131
|
+
pip install "aimnet[pysis]"
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
**Training** (for model training and development):
|
|
135
|
+
|
|
136
|
+
```bash
|
|
137
|
+
pip install "aimnet[train]"
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
**All Features**:
|
|
141
|
+
|
|
142
|
+
```bash
|
|
143
|
+
pip install "aimnet[ase,pysis,train]"
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
### Development Installation
|
|
147
|
+
|
|
148
|
+
For contributors, use [uv](https://docs.astral.sh/uv/) for fast dependency management:
|
|
149
|
+
|
|
150
|
+
```bash
|
|
151
|
+
git clone https://github.com/isayevlab/aimnetcentral.git
|
|
152
|
+
cd aimnetcentral
|
|
153
|
+
make install
|
|
154
|
+
source .venv/bin/activate
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
## Quick Start
|
|
158
|
+
|
|
159
|
+
### Basic Usage (Core)
|
|
160
|
+
|
|
161
|
+
```python
|
|
162
|
+
from aimnet.calculators import AIMNet2Calculator
|
|
163
|
+
|
|
164
|
+
# Load a pre-trained model
|
|
165
|
+
calc = AIMNet2Calculator("aimnet2")
|
|
166
|
+
|
|
167
|
+
# Prepare input
|
|
168
|
+
data = {
|
|
169
|
+
"coord": coordinates, # Nx3 array
|
|
170
|
+
"numbers": atomic_numbers, # N array
|
|
171
|
+
"charge": 0.0,
|
|
172
|
+
}
|
|
173
|
+
|
|
174
|
+
# Run inference
|
|
175
|
+
results = calc(data, forces=True)
|
|
176
|
+
print(results["energy"], results["forces"])
|
|
177
|
+
```
|
|
178
|
+
|
|
179
|
+
### Output Data
|
|
180
|
+
|
|
181
|
+
The calculator returns a dictionary with the following keys:
|
|
182
|
+
|
|
183
|
+
| Key | Shape | Description |
|
|
184
|
+
| --------- | ----------------------- | ------------------------------------ |
|
|
185
|
+
| `energy` | `(,)` or `(B,)` | Total energy in eV |
|
|
186
|
+
| `charges` | `(N,)` or `(B, N)` | Atomic partial charges in e |
|
|
187
|
+
| `forces` | `(N, 3)` or `(B, N, 3)` | Atomic forces in eV/A (if requested) |
|
|
188
|
+
| `hessian` | `(N, 3, N, 3)` | Second derivatives (if requested) |
|
|
189
|
+
| `stress` | `(3, 3)` | Stress tensor for PBC (if requested) |
|
|
190
|
+
|
|
191
|
+
_B = batch size, N = number of atoms_
|
|
192
|
+
|
|
193
|
+
### ASE Integration
|
|
194
|
+
|
|
195
|
+
With `aimnet[ase]` installed:
|
|
196
|
+
|
|
197
|
+
```python
|
|
198
|
+
from ase.io import read
|
|
199
|
+
from aimnet.calculators import AIMNet2ASE
|
|
200
|
+
|
|
201
|
+
atoms = read("molecule.xyz")
|
|
202
|
+
atoms.calc = AIMNet2ASE("aimnet2")
|
|
203
|
+
|
|
204
|
+
energy = atoms.get_potential_energy()
|
|
205
|
+
forces = atoms.get_forces()
|
|
206
|
+
```
|
|
207
|
+
|
|
208
|
+
### Periodic Boundary Conditions
|
|
209
|
+
|
|
210
|
+
For periodic systems, provide a unit cell:
|
|
211
|
+
|
|
212
|
+
```python
|
|
213
|
+
data = {
|
|
214
|
+
"coord": coordinates,
|
|
215
|
+
"numbers": atomic_numbers,
|
|
216
|
+
"charge": 0.0,
|
|
217
|
+
"cell": cell_vectors, # 3x3 array in Angstrom
|
|
218
|
+
}
|
|
219
|
+
results = calc(data, forces=True, stress=True)
|
|
220
|
+
```
|
|
221
|
+
|
|
222
|
+
### Long-Range Coulomb Methods
|
|
223
|
+
|
|
224
|
+
Configure electrostatic interactions for large or periodic systems:
|
|
225
|
+
|
|
226
|
+
```python
|
|
227
|
+
# Damped-Shifted Force (DSF) - recommended for periodic systems
|
|
228
|
+
calc.set_lrcoulomb_method("dsf", cutoff=15.0, dsf_alpha=0.2)
|
|
229
|
+
|
|
230
|
+
# Ewald summation - for accurate periodic electrostatics
|
|
231
|
+
calc.set_lrcoulomb_method("ewald", ewald_accuracy=1e-8)
|
|
232
|
+
```
|
|
233
|
+
|
|
234
|
+
### Performance Optimization
|
|
235
|
+
|
|
236
|
+
For molecular dynamics simulations, use `compile_mode` for ~5x speedup:
|
|
237
|
+
|
|
238
|
+
```python
|
|
239
|
+
calc = AIMNet2Calculator("aimnet2", compile_mode=True)
|
|
240
|
+
```
|
|
241
|
+
|
|
242
|
+
Requirements:
|
|
243
|
+
|
|
244
|
+
- CUDA GPU required
|
|
245
|
+
- Not compatible with periodic boundary conditions
|
|
246
|
+
- Best for repeated inference on similar-sized systems
|
|
247
|
+
|
|
248
|
+
### Training
|
|
249
|
+
|
|
250
|
+
With `aimnet[train]` installed:
|
|
251
|
+
|
|
252
|
+
```bash
|
|
253
|
+
aimnet train --config my_config.yaml --model aimnet2.yaml
|
|
254
|
+
```
|
|
255
|
+
|
|
256
|
+
## Technical Details
|
|
257
|
+
|
|
258
|
+
### Batching and Neighbor Lists
|
|
259
|
+
|
|
260
|
+
The `AIMNet2Calculator` automatically selects the optimal strategy based on system size (`nb_threshold`, default 120 atoms) and hardware:
|
|
261
|
+
|
|
262
|
+
1. **Dense Mode (O(N²))**: Used for small molecules on GPU. Input is kept in 3D batched format `(B, N, 3)`. No neighbor list is computed; the model uses a fully connected graph for maximum parallelism.
|
|
263
|
+
2. **Sparse Mode (O(N))**: Used for large systems or CPU execution. Input is flattened to 2D `(N_total, 3)` with an adaptive neighbor list. This ensures linear memory scaling.
|
|
264
|
+
|
|
265
|
+
### Adaptive Neighbor List
|
|
266
|
+
|
|
267
|
+
In sparse mode, AIMNet2 uses an `AdaptiveNeighborList` that automatically resizes its buffer to maintain efficient utilization (~75%) while preventing overflows.
|
|
268
|
+
|
|
269
|
+
- **Format**: The neighbor list is stored as a 2D integer matrix `nbmat` of shape `(N_total, max_neighbors)`. Each row `i` contains the indices of atoms neighboring atom `i`.
|
|
270
|
+
- **Padding**: Rows with fewer neighbors than `max_neighbors` are padded with the index `N_total` (a dummy atom index).
|
|
271
|
+
- **Buffer Management**: The buffer size `max_neighbors` is always a multiple of 16 for memory alignment. It dynamically expands (by 1.5x) on overflow and shrinks if utilization drops significantly below the target, ensuring robust performance during MD simulations where density fluctuates.
|
|
272
|
+
|
|
273
|
+
## Development
|
|
274
|
+
|
|
275
|
+
Common development tasks using `make`:
|
|
276
|
+
|
|
277
|
+
```bash
|
|
278
|
+
make check # Run linters and code quality checks
|
|
279
|
+
make test # Run tests with coverage
|
|
280
|
+
make docs # Build and serve documentation
|
|
281
|
+
make build # Build distribution packages
|
|
282
|
+
```
|
|
283
|
+
|
|
284
|
+
## Citation
|
|
285
|
+
|
|
286
|
+
If you use AIMNet2 in your research, please cite the appropriate paper:
|
|
287
|
+
|
|
288
|
+
**AIMNet2 (main model):**
|
|
289
|
+
|
|
290
|
+
```bibtex
|
|
291
|
+
@article{aimnet2,
|
|
292
|
+
title={AIMNet2: A Neural Network Potential to Meet Your Neutral, Charged, Organic, and Elemental-Organic Needs},
|
|
293
|
+
author={Anstine, Dylan M and Zubatyuk, Roman and Isayev, Olexandr},
|
|
294
|
+
journal={Chemical Science},
|
|
295
|
+
volume={16},
|
|
296
|
+
pages={10228--10244},
|
|
297
|
+
year={2025},
|
|
298
|
+
doi={10.1039/D4SC08572H}
|
|
299
|
+
}
|
|
300
|
+
```
|
|
301
|
+
|
|
302
|
+
**AIMNet2-NSE:** [ChemRxiv preprint](https://chemrxiv.org/engage/chemrxiv/article-details/692d304c65a54c2d4a7ab3c7)
|
|
303
|
+
|
|
304
|
+
**AIMNet2-Pd:** [ChemRxiv preprint](https://chemrxiv.org/engage/chemrxiv/article-details/67d7b7f7fa469535b97c021a)
|
|
305
|
+
|
|
306
|
+
## License
|
|
307
|
+
|
|
308
|
+
See [LICENSE](LICENSE) file for details.
|