Accelerating Machine Learning MD Workflows with MACE and FBTK

Machine Learning Interatomic Potentials (MLIPs, NNPs) have revolutionized computational chemistry. Universal models like MACE now allow us to perform large-scale Molecular Dynamics (MD) simulations with near-DFT accuracy. However, even as computation speeds up, the "Initial Structure Construction" remains a significant bottleneck in the preprocessing phase.

The Workflow: Build with FBTK, Solve with MACE

FBTK (Forblaze ToolKit) drastically shortens this preprocessing stage using high-performance Rust-based algorithms. This post introduces an efficient workflow: using FBTK to build geometrically sound, overlap-free structures at blistering speeds, and then passing that system to MACE for high-accuracy simulation.

Simple Environment Setup in Just Two Lines

Setting up this workflow is remarkably straightforward. In any standard Python environment, you can prepare everything with just two commands:

pip install mace-torch ase
pip install fbtk

A key advantage is that FBTK has no external dependencies. Even when introduced into complex existing environments where tools like RadonPy, MACE, or PyTorch are already active, there is no risk of causing dependency conflicts. You can extend your capabilities with a guaranteed standalone installation without worrying about breaking your research environment.

1. Solution Building: Manual Grid vs. FBTK Automatic Packing

For example, when creating a system of 8 Benzene molecules (96 atoms), traditional methods using only ASE require manual grid calculations and rotation/translation logic.

With FBTK, you simply declare your target density and let the tool handle the packing:

import fbtk
# 1. Define molecule from SMILES
mol = fbtk.Molecule.from_smiles("c1ccccc1", name="Benzene")

# 2. Pack 8 molecules at a density of 0.8 g/cm3
builder = fbtk.Builder(density=0.8)
builder.add_molecule(mol, count=8)
system = builder.build()

# 3. Rapid relaxation with FBTK's built-in UFF logic
system.relax()

# 4. Convert to ASE and pass to MACE
atoms = system.to_ase()

FBTK’s built-in relaxation (system.relax()) is implemented in Rust and completes in less than 1 second for a system of this size.

2. Advanced Polymer Control: Atactic Polystyrene

Stereoregularity (tacticity) is crucial in polymer simulations. FBTK provides native support for generating polymer chains directly from SMILES strings with specific tacticity control.

# Define Styrene monomer and generate an atactic chain
monomer = fbtk.Molecule.from_smiles("*CC(*)(c1ccccc1)", name="Styrene")
chain = fbtk.Molecule.from_polymer(monomer, degree=20, tacticity="atactic")

# Pack chains at 1.05 g/cm3
builder = fbtk.Builder(density=1.05)
builder.add_molecule(chain, count=3)
system = builder.build()

You can build complex polymer systems using nothing but SMILES strings, without relying on heavy external dependencies like RDKit.

3. High-Accuracy MD with MACE

Once FBTK provides an overlap-free arrangement as an ASE Atoms object, you can immediately start a high-fidelity simulation using the MACE-OFF23 model.

from mace.calculators import mace_off
from ase.optimize import BFGS
from ase.md.langevin import Langevin

# Setup MACE calculator
calc = mace_off(model="small", device="cuda")
atoms.calc = calc

# 3. Geometry Optimization (MACE) - Stabilization before MD
print("Relaxing with MACE...")
opt = BFGS(atoms)
opt.run(fmax=0.05, steps=100)

# 4. Run MD Simulation
# ... (See full script below)

Because FBTK prepares an overlap-free arrangement beforehand, you minimize the risk of atoms colliding and causing numerical instability during the initial phases of the MACE MD.

Full Execution Script

The following self-contained script demonstrates the entire workflow from benzene solution construction with FBTK to MD simulation with MACE.

import torch
import fbtk
from mace.calculators import mace_off
from ase.optimize import BFGS
from ase.md.langevin import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase.io import Trajectory
from ase.units import fs

# 1. Build System (FBTK)
mol = fbtk.Molecule.from_smiles("c1ccccc1", name="Benzene")
builder = fbtk.Builder(density=0.8)
builder.add_molecule(mol, count=8)
system = builder.build()
system.relax() # Fast relaxation with internal UFF (Remove overlaps)
atoms = system.to_ase()

# 2. Setup Calculation (MACE)
device = "cuda" if torch.cuda.is_available() else "cpu"
calc = mace_off(model="small", device=device)
atoms.calc = calc

# 3. Geometry Optimization (MACE) - Stabilization before MD
print("Relaxing with MACE...")
opt = BFGS(atoms)
opt.run(fmax=0.05, steps=100)

# 4. Run MD Simulation
MaxwellBoltzmannDistribution(atoms, temperature_K=300)
dyn = Langevin(atoms, 1.0 * fs, temperature_K=300, friction=0.02)
traj = Trajectory("mace_md.traj", 'w', atoms)
dyn.attach(traj.write, interval=10)

# Print progress
def print_status():
    temp = atoms.get_temperature()
    epot = atoms.get_potential_energy()
    print(f"Step {dyn.nsteps}: T = {temp:.1f} K, Epot = {epot:.2f} eV")

dyn.attach(print_status, interval=100)
dyn.run(1000)

Summary

FBTK is a library designed to streamline computational chemistry workflows by significantly reducing the overhead of initial structure construction and large-scale analysis. It can be integrated into existing environments without conflict and is highly effective for MD simulations using the latest machine learning interatomic potentials.

• FBTK Official Documentation: https://fbtk.forblaze-works.com/