Add aimnet2-nse model weights and card

README.md

@@ -0,0 +1,111 @@
+---
+library_name: aimnet
+license: mit
+tags:
+  - chemistry
+  - computational-chemistry
+  - molecular-dynamics
+  - interatomic-potential
+  - neural-network-potential
+  - pytorch
+  - aimnet2
+pipeline_tag: other
+---
+
+# AIMNet2 NSE (Open-Shell)
+
+AIMNet2 is a neural network interatomic potential for fast and accurate molecular simulations. This is a spin-polarized model for **open-shell chemistry** (radicals, triplet states, excited states). Accepts spin multiplicity as additional input.
+
+## 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-nse")
+
+# Single-point calculation (closed-shell)
+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
+
+# For open-shell species (e.g., methyl radical, doublet)
+results = calc(
+    {"coord": coords, "numbers": atomic_numbers, "charge": 0.0, "mult": 2.0},
+    forces=True,
+)
+```
+
+### With ASE
+
+```python
+from aimnet.calculators.aimnet2ase import AIMNet2ASE
+from ase.build import molecule
+
+atoms = molecule("H2O")
+atoms.calc = AIMNet2ASE("isayevlab/aimnet2-nse", mult=2)
+
+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
+- **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-nse",
+  "ensemble_size": 4,
+  "member_names": [
+    "aimnet2nse_0",
+    "aimnet2nse_1",
+    "aimnet2nse_2",
+    "aimnet2nse_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:c65c7e1c3dc15b97f473f544d0c6835bd8d0c3e6b294523a9311b83eb0d723b6
+size 8951304

ensemble_1.safetensors

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

ensemble_2.safetensors

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

ensemble_3.safetensors

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