ai2-kit (pronounced as “AI Squared Kit”) is a command line and python toolkit featured artificial intelligence × ab initio for complex chemistry systems research.
Feature Highlights
- Collection of tools and workflows to facilitate the development of automated workflows for computational chemistry research.
- Use with oh-my-batch to build your own workflow with shell script.
- Provide reusable skills for building code related to
ai2-kit.
- build-tesla: a skill for building TESLA workflows from existing examples.
2D API Design Pattern
A key innovation of ai2-kit is its 2D API design pattern, shared by the
command-line interface and Python module. The two dimensions are:
- File dimension: A single API call can operate on a collection of files.
Wildcards specify inputs, and template variables specify outputs, making
batch processing as straightforward as working with a single file.
- Operation dimension: A sequence of dependent steps can be composed into
a single expression. The CLI uses
- to separate operations, while Python
uses chained method calls. Reading, transforming, and writing data can be
combined without intermediate variables or files.
By extending both the file and operation dimensions, ai2-kit provides
powerful tools through a concise interface. For example, the following
workflow merges trajectory files, samples structures, converts their format,
and writes each result to a separate file in a single expression:
Python API
from ai2_kit.tool.ase import AseTool
(
AseTool()
.read("./lammps-*/traj.lammpstrj", format="lammps-dump-text")
.sample(100, method="random")
.write_frames("./vasp-{i:04d}/POSCAR", format="vasp")
)
Command Line Interface
ai2-kit tool ase read "./lammps-*/traj.lammpstrj" --format lammps-dump-text \
- sample 100 --method random \
- write_frames "./vasp-{i:04d}/POSCAR" --format vasp
Environment Requirements
ai2-kit requires Python 3.10 or later. On Windows, use ai2-kit through
WSL2. We recommend installing it in a Conda environment.
Installation
Install the core package with pip:
pip install ai2-kit
To install all optional features and their dependencies, including
cp2kdata, DeePMD-kit, and TensorFlow, etc, use:
pip install 'ai2-kit[all]'
Note that TensorFlow currently provides wheels for Python 3.10-3.13. The tensorflow
and all extras install it only on these Python versions.
Usage
All ai2-kit commands can be accessed via the command line interface ai2-kit.
You can run the following command to see the available commands and sub-commands:
ai2-kit
# e.g.
ai2-kit tool
ai2-kit tool dpdata
For detailed usage of each command, you can run the following command to view the help information:
ai2-kit <sub-command> --help
# e.g.
ai2-kit tool dpdata read --help
Manuals
Workflows
Workflows are built with oh-my-batch and example shell scripts, which can be easily adapted to your own research purpose.
The supplementary examples accompanying the paper provide workflows for specific research applications:
- TESLA for bulk water: Active learning for a bulk-water machine learning potential.
- TESLA for Au-CO2: Active learning for the Au-CO2 system.
- TESLA for pKa: Active learning for pKa calculations using free energy perturbation and thermodynamic integration (FEP-TI).
- TESLA for redox potentials: Active learning for redox-potential calculations using FEP-TI.
- ec-MLP: Constant-potential molecular dynamics at the Cu(111)/water interface.
- Spectroscopy: IR, Raman, and sum-frequency generation spectra at the air-water interface using machine learning molecular dynamics and Deep Wannier models.
- AI-agent-generated TESLA: A workflow built with the
build-tesla skill, using MACE, OpenMM, and VASP.
Online Apps and Notebooks
Tips
- Tips: useful tips for using
ai2-kit
Developer Guide
Install from Source
To install ai2-kit from source, run the following commands from the project root:
pip install uv
uv sync --locked
uv run ai2-kit
If you encounter connection errors, reduce uv’s download concurrency:
UV_CONCURRENT_DOWNLOADS=4 uv sync --locked
Jupyter Book is in the optional doc dependency group and is not installed by
default. To install the dependencies needed to build the documentation, run:
uv sync --locked --group doc
Local Tests
Install the runtime and development dependencies from the lockfile, then run
the lightweight suite from the project root:
uv sync --locked --extra cp2kdata
uv run --no-sync python -m pytest -q tests --ignore=tests/test_md_spectra.py -k 'not DPFFDPData'
Build Packages
uv build
Appendix: Transitioning from v1 to v2
v2 makes several significant changes to how ai2-kit workflows are built and
maintained:
- Example-driven workflows: v2 replaces v1’s configuration-driven workflow
system with workflows built from runnable examples. This approach supports a
wider range of functionality, makes workflows easier to reuse and adapt, and
makes it easier to collaborate with AI coding agents.
- Python 3.10 or later: v2 drops support for Python 3.9 and earlier. This
resolves compatibility issues with key dependencies and makes
ai2-kit easier
to install and use on newer Python versions.
- Build tooling: v2 replaces Poetry with uv for dependency management and
package builds. Use
uv sync --locked to install dependencies from uv.lock,
uv run to run commands, and uv build to build distributions.
- v1 maintenance: The v1 codebase will move to the
v1-legacy branch for
continued maintenance.
评论