Home

Awesome

__ Update 6/10/24 __ We release new code, suaitable for large molecules and perioric calculations. Old code available in the old branch. Models were re-compiled and are not compatible with the new code.

AIMNet2 Calculator: Fast, Accurate Molecular Simulations

This package integrates the powerful AIMNet2 neural network potential into your simulation workflows. AIMNet2 provides fast and reliable energy, force, and property calculations for molecules containing a diverse range of elements.

Key Features:

Getting Started

1. Installation

While package is in alpha stage and repository is private, please install into your conda envoronment manually with

# install requirements
conda install -y pytorch pytorch-cuda=12.1 -c pytorch -c nvidia 
conda install -y -c pyg pytorch-cluster
conda install -y -c conda-forge openbabel ase
## pysis requirements
conda install -y -c conda-forge autograd dask distributed h5py fabric jinja2 joblib matplotlib numpy natsort psutil pyyaml rmsd scipy sympy scikit-learn
# now should not do any pip installs
pip install git+https://github.com/eljost/pysisyphus.git
# finally, this repo
git clone git@github.com:zubatyuk/aimnet2calc.git
cd aimnet2calc
python setup.py install

2. Available interfaces

ASE [https://wiki.fysik.dtu.dk/ase]

from aimnet2calc import AIMNet2ASE
calc = AIMNet2ASE('aimnet2')

To specify total molecular charge and spin multiplicity, use optional charge and mult keyword arguments, or set_charge and set_mult methods:

calc = AIMNet2ASE('aimnet2', charge=1)
atoms1.calc = calc
# calculations on atoms1 will be done with charge 1
....
atoms2.calc = calc
calc.set_charge(-2)
# calculations on atoms1 will be done with charge -2

PySisyphus [https://pysisyphus.readthedocs.io]

from aimnet2calc import AIMNet2PySis
calc = AIMNet2PySis('aimnet2')

This produces standard PySisyphus calculator.

Instead of Pysis command line utility, use aimnet2pysis. This registeres AIMNet2 calculator with PySisyphus. Example calc section for PySisyphus YAML files:

calc:
   type: aimnet              # use AIMNet2 calculator
   model: aimnet2_b973c      # use aimnet2_b973c_0.jpt model

3. Base calculator

from aimnet2calc import AIMNet2Calculator

Initialization

calc = AIMNet2Calculator('aimnet2')

will load default AIMNet2 model aimnet2_wb97m_0.jpt as defined at aimnet2calc/models.py . If file does not exist on the machine, it will be downloaded from aimnet-model-zoo repository.

calc = AIMNet2Calculator('/path/to_a/model.jpt')

will load model from the file.

Input structure

The calculator accepts a dictionary containig lists, numpy arrays, torch tensors, or anything that could be accepted by torch.as_tensor.

The input could be for a single molecule (dict keys and shapes):

coord: (B, N, 3)  # atomic coordinates in Angstrom
numbers (B, N)    # atomic numbers
charge (B,)       # molecular charge
mult (B,)         # spin multiplicity, optional

or for a concatenation of molecules:

coord: (N, 3)  # atomic coordinates in Angstrom
numbers (N,)    # atomic numbers
charge (B,)    # molecular charge
mult (B,)      # spin multiplicity, optional
mol_idx (N,)   # molecule index for each atom, should contain integers in increasing order, with (B-1) is the maximum number.

where B is the number of molecules, N is number of atoms.

Calling calculator

results = calc(data, forces=False, stress=False, hessian=False)

results would be a dictionary of PyTorch tensors containing energy, charges, and possibly forces, stress and hessian if requested.

4. Long range Coulomb model

By default, Coulomb energy is calculated in O(N^2) manner, e.g. pair interaction between every pair of atoms in system. For very large or periodic systems, O(N) Dumped-Shifted Force Coulomb model could be employed doi: 10.1063/1.2206581. With AIMNet2Calculator interface, switch between standard and DSF Coulomb implementations im AIMNet2 models:

# switch to O(N)
calc.set_lrcoulomb_method('dsf', cutoff=15.0, dsf_alpha=0.2)
# switch to O(N^2), not suitable for PBC
calc.set_lrcoulomb_method('simple')