Metadata-Version: 2.4
Name: aimnet
Version: 0.1.0
Summary: AIMNet Machine Learned Interatomic Potential
Project-URL: Homepage, https://github.com/isayevlab/aimnetcentral
Project-URL: Documentation, https://isayevlab.github.io/aimnetcentral/
Project-URL: Repository, https://github.com/isayevlab/aimnetcentral
Author-email: Roman Zubatyuk <zubatyuk@gmail.com>
License: MIT License
        
        Copyright (c) 2024, Roman Zubatyuk
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Keywords: computational chemistry,deep learning,interatomic potential,machine learning,molecular dynamics
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.11
Requires-Dist: click>=8.1.7
Requires-Dist: h5py>=3.12.1
Requires-Dist: jinja2>=3.1.6
Requires-Dist: numpy
Requires-Dist: nvalchemi-toolkit-ops>=0.2
Requires-Dist: pyyaml>=6.0.2
Requires-Dist: requests>=2.32.3
Requires-Dist: torch>=2.4
Requires-Dist: warp-lang>=1.11
Provides-Extra: ase
Requires-Dist: ase==3.27.0; extra == 'ase'
Provides-Extra: pysis
Requires-Dist: pysisyphus; extra == 'pysis'
Provides-Extra: train
Requires-Dist: omegaconf>=2.3.0; extra == 'train'
Requires-Dist: pytorch-ignite>=0.5.1; extra == 'train'
Requires-Dist: wandb>=0.18.5; extra == 'train'
Description-Content-Type: text/markdown

[![Release](https://img.shields.io/github/v/release/isayevlab/aimnetcentral)](https://github.com/isayevlab/aimnetcentral/releases)
[![Python](https://img.shields.io/badge/python-3.11%20%7C%203.12-blue)](https://www.python.org/)
[![Build status](https://img.shields.io/github/actions/workflow/status/isayevlab/aimnetcentral/main.yml?branch=main)](https://github.com/isayevlab/aimnetcentral/actions/workflows/main.yml?query=branch%3Amain)
[![codecov](https://codecov.io/gh/isayevlab/aimnetcentral/branch/main/graph/badge.svg)](https://codecov.io/gh/isayevlab/aimnetcentral)
[![License](https://img.shields.io/github/license/isayevlab/aimnetcentral)](https://github.com/isayevlab/aimnetcentral/blob/main/LICENSE)

- **Github repository**: <https://github.com/isayevlab/aimnetcentral/>
- **Documentation** <https://isayevlab.github.io/aimnetcentral/>

# AIMNet2 : ML interatomic potential for fast and accurate atomistic simulations

## Key Features

- Accurate and Versatile: AIMNet2 excels at modeling neutral, charged, organic, and elemental-organic systems.
- Flexible Interfaces: Use AIMNet2 through convenient calculators for popular simulation packages like ASE and PySisyphus.
- Flexible Long-Range Interactions: Optionally employ the Damped-Shifted Force (DSF) or Ewald summation Coulomb models for accurate calculations in large or periodic systems.

## Requirements

### Python Version

AIMNet2 requires **Python 3.11 or 3.12**.

### GPU Support (Optional)

AIMNet2 works on CPU out of the box. For GPU acceleration:

- **CUDA GPU**: Install PyTorch with CUDA support from [pytorch.org](https://pytorch.org/get-started/locally/)
- **compile_mode**: Requires CUDA for ~5x MD speedup (see Performance Optimization)

Example PyTorch installation with CUDA 12.4:

```bash
pip install torch --index-url https://download.pytorch.org/whl/cu124
```

## Available Models

| Model           | Elements                                      | Description                                   |
| --------------- | --------------------------------------------- | --------------------------------------------- |
| `aimnet2`       | H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I | wB97M-D3 (default)                            |
| `aimnet2_b973c` | H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I | B97-3c functional                             |
| `aimnet2_2025`  | H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I | B97-3c + improved intermolecular interactions |
| `aimnet2nse`    | H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I | Open-shell chemistry                          |
| `aimnet2pd`     | H, B, C, N, O, F, Si, P, S, Cl, Se, Br, Pd, I | Palladium-containing systems                  |

_Each model has ensemble members (append \_0 to \_3). Ensemble averaging recommended for production use._

## Installation

### Basic Installation

Install from PyPI:

```bash
pip install aimnet
```

### Optional Features

AIMNet2 provides optional extras for different use cases:

**ASE Calculator** (for atomistic simulations with ASE):

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

**PySisyphus Calculator** (for reaction path calculations):

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

**Training** (for model training and development):

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

**All Features**:

```bash
pip install "aimnet[ase,pysis,train]"
```

### Development Installation

For contributors, use [uv](https://docs.astral.sh/uv/) for fast dependency management:

```bash
git clone https://github.com/isayevlab/aimnetcentral.git
cd aimnetcentral
make install
source .venv/bin/activate
```

## Quick Start

### Basic Usage (Core)

```python
from aimnet.calculators import AIMNet2Calculator

# Load a pre-trained model
calc = AIMNet2Calculator("aimnet2")

# Prepare input
data = {
    "coord": coordinates,  # Nx3 array
    "numbers": atomic_numbers,  # N array
    "charge": 0.0,
}

# Run inference
results = calc(data, forces=True)
print(results["energy"], results["forces"])
```

### Output Data

The calculator returns a dictionary with the following keys:

| Key       | Shape                   | Description                          |
| --------- | ----------------------- | ------------------------------------ |
| `energy`  | `(,)` or `(B,)`         | Total energy in eV                   |
| `charges` | `(N,)` or `(B, N)`      | Atomic partial charges in e          |
| `forces`  | `(N, 3)` or `(B, N, 3)` | Atomic forces in eV/A (if requested) |
| `hessian` | `(N, 3, N, 3)`          | Second derivatives (if requested)    |
| `stress`  | `(3, 3)`                | Stress tensor for PBC (if requested) |

_B = batch size, N = number of atoms_

### ASE Integration

With `aimnet[ase]` installed:

```python
from ase.io import read
from aimnet.calculators import AIMNet2ASE

atoms = read("molecule.xyz")
atoms.calc = AIMNet2ASE("aimnet2")

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

### Periodic Boundary Conditions

For periodic systems, provide a unit cell:

```python
data = {
    "coord": coordinates,
    "numbers": atomic_numbers,
    "charge": 0.0,
    "cell": cell_vectors,  # 3x3 array in Angstrom
}
results = calc(data, forces=True, stress=True)
```

### Long-Range Coulomb Methods

Configure electrostatic interactions for large or periodic systems:

```python
# Damped-Shifted Force (DSF) - recommended for periodic systems
calc.set_lrcoulomb_method("dsf", cutoff=15.0, dsf_alpha=0.2)

# Ewald summation - for accurate periodic electrostatics
calc.set_lrcoulomb_method("ewald", ewald_accuracy=1e-8)
```

### Performance Optimization

For molecular dynamics simulations, use `compile_mode` for ~5x speedup:

```python
calc = AIMNet2Calculator("aimnet2", compile_mode=True)
```

Requirements:

- CUDA GPU required
- Not compatible with periodic boundary conditions
- Best for repeated inference on similar-sized systems

### Training

With `aimnet[train]` installed:

```bash
aimnet train --config my_config.yaml --model aimnet2.yaml
```

## Technical Details

### Batching and Neighbor Lists

The `AIMNet2Calculator` automatically selects the optimal strategy based on system size (`nb_threshold`, default 120 atoms) and hardware:

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.
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.

### Adaptive Neighbor List

In sparse mode, AIMNet2 uses an `AdaptiveNeighborList` that automatically resizes its buffer to maintain efficient utilization (~75%) while preventing overflows.

- **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`.
- **Padding**: Rows with fewer neighbors than `max_neighbors` are padded with the index `N_total` (a dummy atom index).
- **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.

## Development

Common development tasks using `make`:

```bash
make check       # Run linters and code quality checks
make test        # Run tests with coverage
make docs        # Build and serve documentation
make build       # Build distribution packages
```

## Citation

If you use AIMNet2 in your research, please cite the appropriate paper:

**AIMNet2 (main model):**

```bibtex
@article{aimnet2,
  title={AIMNet2: A Neural Network Potential to Meet Your Neutral, Charged, Organic, and Elemental-Organic Needs},
  author={Anstine, Dylan M and Zubatyuk, Roman and Isayev, Olexandr},
  journal={Chemical Science},
  volume={16},
  pages={10228--10244},
  year={2025},
  doi={10.1039/D4SC08572H}
}
```

**AIMNet2-NSE:** [ChemRxiv preprint](https://chemrxiv.org/engage/chemrxiv/article-details/692d304c65a54c2d4a7ab3c7)

**AIMNet2-Pd:** [ChemRxiv preprint](https://chemrxiv.org/engage/chemrxiv/article-details/67d7b7f7fa469535b97c021a)

## License

See [LICENSE](LICENSE) file for details.
