Metatensor and metatomic: foundational libraries for interoperable atomistic machine learning

This paper introduces metatensor and metatomic, two foundational software libraries designed to bridge the gap between traditional atomistic modeling tools and modern machine learning frameworks by providing interoperable, portable data structures and model interfaces that facilitate the adoption of ML in atomic-scale simulations.

The Gist

Imagine the world of computer simulations for atoms and molecules as a massive, bustling international city. In this city, there are two very different groups of people trying to work together:

1. The Old Guard (Traditional Simulation Engines): These are the veteran engineers who have been building bridges and roads for decades. They speak "Fortran," "C," and "C++." They are incredibly fast, reliable, and built for heavy lifting, but they are a bit rigid and don't like to change their blueprints easily.
2. The New Wave (Machine Learning Models): These are the brilliant, creative data scientists who speak "Python," "Julia," and "PyTorch." They are building amazing new tools using AI to predict how atoms behave. They are fast to innovate but often build things that only work in their own specific neighborhood.

The Problem:
For years, these two groups couldn't talk to each other. If a scientist wanted to use a new AI model to predict how a molecule moves, they had to build a custom, one-off bridge between the AI and the simulation engine. It was like trying to plug a USB-C device into a 1980s cassette player—you needed a weird adapter, it was hard to make, and if you changed the device, the adapter broke. This made it slow, expensive, and frustrating to combine the best of both worlds.

The Solution: metatensor and metatomic
The authors of this paper introduced two new "universal adapters" and a "standardized language" to fix this. Think of them as the USB-C port and the Universal Power Adapter for the atomic world.

1. metatensor: The "Smart Filing Cabinet"

Imagine you have a box of data. In the old days, you just dumped numbers into a box. But in science, the context is just as important as the numbers.

• The Analogy: metatensor is like a smart filing cabinet. Instead of just throwing papers in a drawer, it puts every sheet in a folder that has a label saying exactly what it is, where it came from, and how it relates to other papers.
• Why it matters: It handles "sparse" data (data where most spots are empty, like a sparse spreadsheet) very efficiently. It also keeps the "gradients" (which are like the derivative or the "slope" of the data, crucial for physics) right next to the data itself, so they never get lost.
• The Magic: It speaks every language. Whether you are writing code in Python, C++, or Rust, metatensor translates the data so everyone understands it perfectly. It's the common language that lets the AI and the simulation engine exchange complex information without confusion.

2. metatomic: The "Universal Model Passport"

Once you have the data, you need to share the model itself (the AI brain).

• The Analogy: metatomic is like a universal passport and instruction manual for an AI model. Usually, an AI model is like a custom-built robot that only works in one specific factory. metatomic puts that robot in a standardized shipping container.
• How it works: It wraps the AI model, its "weights" (the learned knowledge), and a list of what it can do (e.g., "I can calculate energy" or "I can predict forces") into a single, portable file.
• The Magic: Now, a simulation engine (like LAMMPS, which runs on supercomputers) can look at this passport, say, "Ah, this model can calculate energy? Great, let's use it!" without needing to know how the robot was built inside. It turns the complex task of connecting an AI to a simulator from a custom engineering project into a simple "plug-and-play" action.

The Ecosystem: A City of Tools

The paper doesn't just stop at the adapters; it shows a whole new city built around them:

• metatrain: A factory that builds these AI models using the new standards.
• featomic & torch-spex: Tools that create the "descriptors" (the features) the AI needs to understand atoms, like a translator turning raw atomic positions into a language the AI understands.
• FlashMD: A super-fast AI that predicts the future movement of atoms directly, skipping the slow step-by-step calculations of traditional physics.
• Chemiscope & PLUMED: Tools that let scientists visualize and explore these new models easily, like a map and compass for the new city.

The Result: A Seamless Future

The authors tested this by running simulations of water and complex molecules. They found that using these new tools was almost as fast as the old, custom-built methods, but with a massive advantage: flexibility.

• Before: If you wanted to switch from one AI model to another, you had to rewrite your simulation code.
• Now: You just swap the "passport" (the metatomic file), and the simulation engine keeps running without a hitch.

In Summary:
This paper introduces a new infrastructure that allows the "Old Guard" of high-performance physics simulations and the "New Wave" of Machine Learning to finally hold hands. By creating a standard way to store data (metatensor) and a standard way to package AI models (metatomic), they have removed the friction that was slowing down scientific discovery. Now, researchers can focus on solving big problems in chemistry and materials science, rather than wasting time building custom bridges between software that should have been talking to each other all along.

Technical Summary

1. Problem Statement

The integration of Machine Learning (ML) into atomistic simulations has revolutionized materials science and chemistry, offering high accuracy at reduced computational costs. However, the field suffers from severe fragmentation and interoperability issues:

• Diverse Ecosystems: Traditional simulation engines (e.g., LAMMPS, GROMACS, Quantum ESPRESSO) are often written in Fortran, C, or C++, while modern ML frameworks (PyTorch, JAX, NumPy) are predominantly Python-based.
• Data Structure Mismatch: ML models require specific data structures to handle gradients, sparsity, and metadata (e.g., atomic positions, forces, tensor symmetries), which existing libraries (like NumPy or Pandas) do not natively support efficiently for physical sciences.
• Integration Overhead: Connecting a specific ML model to a specific simulation engine requires custom, time-consuming interfaces. This creates an $O(M \times N)$ complexity problem (where $M$ is the number of models and $N$ is the number of engines), limiting reproducibility and flexibility.
• Model Portability: ML models consist of both weights and code (architecture). Sharing them across different languages and environments is difficult without a standardized container format.

2. Methodology

The authors introduce two foundational libraries, metatensor and metatomic, designed to act as a universal "hourglass" interface between ML frameworks and simulation engines.

A. metatensor: A Self-Describing Data Format

metatensor provides a multi-platform, multi-language storage format for arrays with potentially sparse indices, specifically designed for atomistic ML.

• Core Objects:
  1. Labels: Named, multi-dimensional indices (metadata) stored as integer arrays. They define dimensions like "system," "atom," "components" (e.g., x, y, z), and "properties."
  2. TensorBlock: A dense floating-point data array decorated with Labels. It stores values and their associated gradients (e.g., forces as gradients of energy) recursively.
  3. TensorMap: A key/value map acting as a block-sparse storage format. It groups multiple TensorBlocks based on specific symmetry patterns (e.g., irreducible representations of the O(3) group) to avoid storing zero values.
• Key Features:
  • Gradient Support: Native handling of gradients with respect to atomic positions and strain, essential for calculating forces and virials.
  • Sparsity: Efficient block-sparse representation for equivariant data (common in atomistic ML).
  • Serialization: Uses the language-agnostic npz (NumPy) format for robust, long-term storage.
  • Implementation: Core library in Rust with C API bindings for C++, Python, and TorchScript.

B. metatomic: A Standardized Model Interface

metatomic defines a unified interface for atomistic ML models, enabling them to be used across different simulation engines.

• Architecture: It wraps a model (code + weights) and metadata (capabilities, authors, inputs/outputs).
• Workflow:
  1. Declaration: The model declares what outputs it can compute (e.g., energy, forces).
  2. Request: The simulation engine queries the model for required inputs (e.g., neighbor lists, atomic positions).
  3. Execution: The engine prepares data and calls the model, receiving outputs as TensorMap objects.
• Benefit: Reduces integration complexity from $O(M \times N)$ to $O(M + N)$. Once a model implements the metatomic interface, it works with any compatible engine.

C. The Ecosystem

The paper describes a modular ecosystem built on these foundations:

• metatrain: A CLI tool for training, evaluating, and exporting models (using metatomic format) for various architectures (GAP, Behler-Parrinello, GNNs).
• featomic & torch-spex: Libraries for computing atomic descriptors (e.g., SOAP, ACE) with high performance and GPU support.
• torch-pme: A library for efficient long-range electrostatic interactions (PME) with automatic differentiation.
• vesin: A fast neighbor list calculator.
• Integrations: Interfaces with LAMMPS, i-PI, ASE, PLUMED, eOn, and chemiscope.

3. Key Contributions

1. Standardized Data Format (metatensor): Introduced a metadata-rich, gradient-friendly, block-sparse tensor format that bridges the gap between Python ML and C/Fortran simulation codes.
2. Unified Model Interface (metatomic): Established a "hourglass" design pattern allowing any ML model to interface with any simulation engine without custom code for each pair.
3. Performance Optimization: Demonstrated that the overhead of the metatomic interface is negligible (~2 µs/atom) compared to model execution time, even on GPUs.
4. Modular Toolchain: Released a suite of interoperable tools (metatrain, featomic, torch-spex, etc.) that allow users to mix and match components for custom workflows.
5. Demonstrated Interoperability: Showcased successful integration of complex workflows, including:
   • Training universal potentials (PET-MAD) and exporting them to LAMMPS and ASE.
   • Running Path Integral Molecular Dynamics (PIMD) with quantum nuclear effects using i-PI.
   • Computing chemical shifts for NMR crystallography (ShiftML).
   • Accelerated molecular dynamics via FlashMD (direct trajectory prediction).

4. Results

• Performance Benchmarks:
  • Overhead: In LAMMPS simulations of liquid water, metatomic achieved 18.3 timesteps/s vs. 18.7 timesteps/s for a native integration (MACE-OFF24), showing <2% overhead.
  • Descriptor Calculation: featomic outperformed existing libraries (QUIP, librascal, DScribe) in both speed and memory efficiency, particularly when computing gradients for sparse data (e.g., 8 GiB peak memory vs. 30 GiB for competitors).
  • Scaling: The PET-MAD potential showed significant speedups on GPUs when using the KOKKOS-enabled LAMMPS interface compared to CPU-only ASE implementations.
• Scientific Applications:
  • PET-MAD: A universal interatomic potential trained on 85 elements, demonstrating high accuracy and generalization.
  • ShiftML: Successfully predicted NMR chemical shifts for organic crystals and liquids, handling complex thermal distortions.
  • FlashMD: Achieved 1-2 orders of magnitude acceleration in MD by directly predicting trajectories rather than integrating forces.
  • Free Energy Surfaces: Used metatomic models within PLUMED to define custom collective variables for exploring the free energy landscape of Lennard-Jones clusters.

5. Significance

This work addresses a critical bottleneck in computational chemistry: the inability to easily combine state-of-the-art ML models with established simulation infrastructure.

• Interoperability: By decoupling model development from simulation execution, it allows researchers to focus on algorithm design rather than writing custom interfaces.
• Reproducibility: The standardized metatomic format ensures that models can be shared, validated, and reused across the community, adhering to FAIR principles.
• Future-Proofing: The design supports future integration with emerging frameworks (JAX, Julia, Fortran) and allows for the evolution of ML architectures without breaking existing simulation pipelines.
• Democratization: Tools like metatrain and the "Atomistic Cookbook" lower the barrier to entry, enabling non-ML experts to utilize advanced ML potentials in their research.

In conclusion, metatensor and metatomic serve as the foundational infrastructure for a new generation of interoperable, high-performance atomistic machine learning, bridging the gap between traditional physics-based simulations and modern data-driven modeling.