File size: 3,244 Bytes
bef6930
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
---
library_name: aimnet
license: mit
tags:
  - chemistry
  - computational-chemistry
  - molecular-dynamics
  - interatomic-potential
  - neural-network-potential
  - pytorch
  - aimnet2
pipeline_tag: other
---

# AIMNet2 Pd (Palladium)

AIMNet2 is a neural network interatomic potential for fast and accurate molecular simulations. This is an extended model supporting **palladium-containing systems** for organometallic catalysis and coordination chemistry.

## Highlights

- **Fast:** Orders of magnitude faster than DFT, with optional `torch.compile` for ~5x additional speedup on GPU
- **Accurate:** Near-DFT accuracy for energies, forces, and charges
- **Broad coverage:** Supports 14 elements: H, B, C, N, O, F, Si, P, S, Cl, Se, Br, Pd, I
- **Ensemble:** 4 ensemble members for uncertainty estimation
- **Features:** Energy, forces, charges, Hessian, stress tensor, periodic boundary conditions

## Installation

```bash
pip install "aimnet[hf]"
```

## Quick Start

```python
from aimnet.calculators import AIMNet2Calculator

# Load from Hugging Face (downloads and caches automatically)
calc = AIMNet2Calculator("isayevlab/aimnet2-pd")

# Single-point calculation
results = calc(
    {"coord": coords, "numbers": atomic_numbers, "charge": 0.0},
    forces=True,
)
print(results["energy"])   # Energy in eV
print(results["forces"])   # Forces in eV/A
print(results["charges"])  # Partial charges in e
```

### With ASE

```python
from aimnet.calculators.aimnet2ase import AIMNet2ASE
from ase.build import molecule

atoms = molecule("H2O")
atoms.calc = AIMNet2ASE("isayevlab/aimnet2-pd")

energy = atoms.get_potential_energy()
forces = atoms.get_forces()
```

## Ensemble Members

This repo contains 4 ensemble members (`ensemble_0.safetensors` through `ensemble_3.safetensors`). The default loads member 0. For uncertainty estimation, load all 4 and average predictions.

| File | Description |
|------|-------------|
| `ensemble_0.safetensors` | Ensemble member 0 (default) |
| `ensemble_1.safetensors` | Ensemble member 1 |
| `ensemble_2.safetensors` | Ensemble member 2 |
| `ensemble_3.safetensors` | Ensemble member 3 |
| `config.json` | Shared model configuration and metadata |

## Model Details

- **Architecture:** AIMNet2 (Atomic Environment Vectors + message-passing MLPs)
- **Cutoff radius:** 5.0 A
- **DFT functional:** wB97M-D3
- **Dispersion:** D3BJ (externalized, handled by calculator)
- **Coulomb:** No embedded short-range Coulomb (coulomb_mode=none)
- **Format:** safetensors (converted from PyTorch state dict)

## Limitations

- **Periodic systems** are supported through the Python API / ASE interface, not through the Gradio demo
- Element coverage is limited to the 14 elements listed above; unsupported elements will raise an error
- This model does not include arsenic (As) — use `aimnet2-wb97m-d3` for As-containing systems

## Citation

```bibtex
@article{anstine2025aimnet2,
  title={AIMNet2: A Neural Network Potential to Meet your Neutral, Charged, Organic, and Elemental-Organic Needs},
  author={Anstine, Dylan and Zubatyuk, Roman and Isayev, Olexandr},
  journal={Chemical Science},
  year={2025},
  publisher={Royal Society of Chemistry},
  doi={10.1039/D4SC08572H}
}
```

## License

MIT License