Skip to main content

Atomistic Learned Potential Package in JAX

Project description

apax: Atomistic learned Potentials in JAX!

Read the Docs DOI Code style: black License: MIT

apax[1] is a high-performance, extendable package for training of and inference with atomistic neural networks. It implements the Gaussian Moment Neural Network model [2, 3]. It is based on JAX and uses JaxMD as a molecular dynamics engine.

Installation

Apax is available on PyPI with a CPU version of JAX.

pip install apax

For more detailed instructions, please refer to the documentation.

CUDA Support

If you want to enable GPU support (only on Linux), please overwrite the jaxlib version:

CUDA 12:

pip install -U "jax[cuda12]"

See the Jax installation instructions for more details.

Usage

Your first apax Model

In order to train a model, you need to run

apax train config.yaml

We offer some input file templates to get new users started as quickly as possible. Simply run the following commands and add the appropriate entries in the marked fields

apax template train # use --full for a template with all input options

Please refer to the documentation for a detailed explanation of all parameters. The documentation can convenienty be accessed by running apax docs.

Molecular Dynamics

There are two ways in which apax models can be used for molecular dynamics out of the box. High performance NVT simulations using JaxMD can be started with the CLI by running

apax md config.yaml md_config.yaml

A template command for MD input files is provided as well.

The second way is to use the ASE calculator provided in apax.md.

Input File Auto-Completion

use the following command to generate JSON schemata for training and MD configuration files:

apax schema

If you are using VSCode, you can utilize them to lint and autocomplete your input files. The command creates the 2 schemata and adds them to the projects .vscode/settings.json

Authors

  • Moritz René Schäfer
  • Nico Segreto

Under the supervion of Johannes Kästner

Contributing

We are happy to receive your issues and pull requests!

Do not hesitate to contact any of the authors above if you have any further questions.

Acknowledgements

The creation of Apax was supported by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) in the framework of the priority program SPP 2363, “Utilization and Development of Machine Learning for Molecular Applications - Molecular Machine Learning” Project No. 497249646 and the Ministry of Science, Research and the Arts Baden-Württemberg in the Artificial Intelligence Software Academy (AISA). Further funding though the DFG under Germany's Excellence Strategy - EXC 2075 - 390740016 and the Stuttgart Center for Simulation Science (SimTech) was provided.

References

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

apax-0.7.0.tar.gz (114.6 kB view details)

Uploaded Source

Built Distribution

apax-0.7.0-py3-none-any.whl (143.8 kB view details)

Uploaded Python 3

File details

Details for the file apax-0.7.0.tar.gz.

File metadata

  • Download URL: apax-0.7.0.tar.gz
  • Upload date:
  • Size: 114.6 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/1.8.3 CPython/3.10.14 Linux/6.5.0-1025-azure

File hashes

Hashes for apax-0.7.0.tar.gz
Algorithm Hash digest
SHA256 423e45f6858b7f86026733888bca5800f41ea51f6f307f86e94f31739c059a56
MD5 7454d175f61e804b9d1db042919d0080
BLAKE2b-256 dd376c45f4c97a7f3f8851517b58e0e23ba0b2801c24f155ccba56eaeb3edcd6

See more details on using hashes here.

File details

Details for the file apax-0.7.0-py3-none-any.whl.

File metadata

  • Download URL: apax-0.7.0-py3-none-any.whl
  • Upload date:
  • Size: 143.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/1.8.3 CPython/3.10.14 Linux/6.5.0-1025-azure

File hashes

Hashes for apax-0.7.0-py3-none-any.whl
Algorithm Hash digest
SHA256 bdf71b2ecf43246030bb1e7d2f7050c8ed62414afbe85c45bb0af8a133d119ec
MD5 b19f8e3516f9ea6ada954b10c2549067
BLAKE2b-256 b79f6174ba9e2687240e39a369b87b494814d5f00eeb0f301b8a11f0149759f4

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page