Add aimnet2-wb97m-d3 model weights and card

README.md

@@ -0,0 +1,105 @@
+---
+library_name: aimnet
+license: mit
+tags:
+  - chemistry
+  - computational-chemistry
+  - molecular-dynamics
+  - interatomic-potential
+  - neural-network-potential
+  - pytorch
+  - aimnet2
+pipeline_tag: other
+---
+
+# AIMNet2 wB97M-D3
+
+AIMNet2 is a neural network interatomic potential for fast and accurate molecular simulations. This model is trained at the **wB97M-D3** level of theory.
+
+## 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, As, Se, Br, 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-wb97m-d3")
+
+# 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-wb97m-d3")
+
+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:** Short-range embedded + external long-range (DSF/Ewald)
+- **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
+
+## 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

config.json

@@ -0,0 +1,38 @@
+{
+  "config_schema_version": 1,
+  "family_name": "aimnet2-wb97m-d3",
+  "ensemble_size": 4,
+  "member_names": [
+    "aimnet2_wb97m_d3_0",
+    "aimnet2_wb97m_d3_1",
+    "aimnet2_wb97m_d3_2",
+    "aimnet2_wb97m_d3_3"
+  ],
+  "cutoff": 5.0,
+  "needs_coulomb": true,
+  "needs_dispersion": true,
+  "coulomb_mode": "sr_embedded",
+  "implemented_species": [
+    1,
+    5,
+    6,
+    7,
+    8,
+    9,
+    14,
+    15,
+    16,
+    17,
+    33,
+    34,
+    35,
+    53
+  ],
+  "library_name": "aimnet",
+  "tags": [
+    "chemistry",
+    "molecular-dynamics",
+    "force-field",
+    "aimnet2"
+  ]
+}

ensemble_0.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:930f3f32a14202ddb3d7c90d5ab0e4f22da880e33e5dcd144b13503aa27de4f5
+size 8825656

ensemble_1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4d45344c96fefee4067597ccc33103c60398cfbbcf4052008ace8eb6d0bcd3f
+size 8825656

ensemble_2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff9b2b8f5b7c3be9ca02d1a8a96127a94e2988a2cc18d02afd854f5061915157
+size 8825656

ensemble_3.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3475a7c49186a41e244af899aa69b9d4e4253c8b0ae3189c9ec11f5fd7ec0263
+size 8825656