jaxip 0.7.0

JAX-based Interatomic Potential

JAX-based Interatomic Potential

Description

Jaxip is an optimized Python library on basis of Google JAX that enables development of emerging machine learning interatomic potentials for use in computational physics, chemistry, material science. These potentials are necessary for conducting large-scale molecular dynamics simulations of complex materials with ab initio accuracy.

See documentation for more information.

Features

• The design of Jaxip is simple and flexible, which makes it easy to incorporate atomic descriptors and potentials.

• It uses autograd to make defining new descriptors straightforward.

• Jaxip is written purely in Python and optimized with just-in-time (JIT) compilation.

• It also supports GPU-accelerated computing, which can significantly speed up preprocessing and model training.

Warning

This package is under heavy development and the current focus is on the implementation of high-dimensional neural network potential (HDNNP) proposed by Behler et al. (2007).

Installation

To install Jaxip, run this command in your terminal:

$ pip install jaxip

For machines with an NVIDIA GPU please follow the installation instruction on the documentation.

Examples

Defining an atomic descriptor

This script demonstrates the process of evaluating an array of atomic-centered symmetry functions (ACSF) for a specific element, which can be utilized to evaluate the descriptor values for any structure. The resulting values can then be used to construct a machine learning potential.

from jaxip.datasets import RunnerDataset
from jaxip.descriptors import ACSF
from jaxip.descriptors.acsf import CutoffFunction, G2, G3

# Read atomic structure dataset (e.g. water molecules)
structures = RunnerDataset('input.data')
structure = structures[0]

# Define ACSF descriptor for hydrogen element
descriptor = ACSF(element='H')

# Add radial and angular symmetry functions
cfn = CutoffFunction(r_cutoff=12.0, cutoff_type='tanh')
descriptor.add(G2(cfn, eta=0.5, r_shift=0.0), 'H')
descriptor.add(G3(cfn, eta=0.001, zeta=2.0, lambda0=1.0, r_shift=12.0), 'H', 'O')
print(descriptor)

values = descriptor(structure)
print("Descriptor values:\n", values)

gradient = descriptor.grad(structure, atom_index=0)
print("Descriptor gradient:\n", gradient)

Training a machine learning potential

This example illustrates how to quickly create a high-dimensional neural network potential (HDNNP) instance from an in input setting files and train it on input structures. The trained potential can then be used to evaluate the energy and force components for new structures.

from jaxip.datasets import RunnerDataset
from jaxip.potentials import NeuralNetworkPotential

# Read atomic data
structures = RunnerDataset("input.data")
structure = structures[0]

# Instantiate potential from input settings file
nnp = NeuralNetworkPotential.from_file("input.nn")

# Fit descriptor scaler and model weights
nnp.fit_scaler(structures)
nnp.fit_model(structures)
nnp.save()

# Or loading from files
#nnp.load()

# Total energy
nnp(structure)

# Force components
nnp.compute_force(structure)

Example files: input.data and input.nn

License

This project is licensed under the GNU General Public License (GPL) version 3 - see the LICENSE file for details.