FAIRChem v2 и UMA: Универсальная симуляция атомного мира

В мире вычислительной химии и материаловедения долгое время существовала проблема фрагментации: для каждой задачи требовалась своя модель. Молекулярная динамика воды, каталитические реакции на поверхности металла и свойства кристаллической решетки железа часто моделировались с использованием совершенно разных алгоритмов, обученных на разных наборах данных. Это создавало барьеры для исследователей, которым нужно было переключаться между инструментами, и усложняло сравнение результатов. Однако появление FAIRChem v2 и его ядра — UMA (Universal Machine-learning Interatomic Potential) — меняет эту парадигму. UMA представляет собой единый, универсальный потенциал, способный предсказывать энергии и силы для атомов в самых разных химических контекстах: от изолированных молекул до сложных каталитических систем и объемных материалов.

В этой статье мы не просто поверхностно коснемся темы, а подробно разберем, как настроить рабочую среду, аутентифицироваться для доступа к весам модели и применить UMA к широкому спектру задач. Мы пройдем путь от вычисления энергии связи простой молекулы воды до сложной молекулярной динамики и анализа колебательных спектров. Особое внимание уделим интеграции с библиотекой ASE (Atomic Simulation Environment), которая является стандартом де-факто для работы с атомными структурами в Python. Этот материал станет вашим практическим руководством для запуска собственных симуляций с использованием передовых моделей искусственного интеллекта.

01Что такое FAIRChem v2 и UMA?

FAIRChem — это фреймворк, разработанный компанией Meta (ранее Facebook AI Research), предназначенный для обучения и применения моделей машинного обучения для предсказания свойств материалов и молекул. Версия 2.0 приносит значительные улучшения в производительность и удобство использования. В центре этой системы лежит модель UMA (Universal Machine-learning Interatomic Potential).

Традиционные силовые поля (force fields) часто требуют ручной параметризации для каждого нового класса материалов, что трудоемко и ограничено областью применимости. Аб initio методы (например, DFT) точны, но вычислительно крайне дороги. UMA решает эту дилемму, обучаясь на огромных наборах данных, включающих различные химические домены. Она «понимает» химию на фундаментальном уровне, что позволяет ей обобщать знания: если модель увидела, как взаимодействуют атомы углерода и кислорода в CO2, она может более точно предсказать их взаимодействие в сложных органических молекулах или на поверхности катализатора.

Ключевое преимущество UMA — это концепция «одной модели для всех». Вместо того чтобы обучать отдельную нейросеть для органической химии, другую для неорганических материалов и третью для катализа, UMA использует единый архитектурный подход, который адаптируется к контексту через так называемые «калькуляторы» (calculators) в рамках ASE. Это позволяет исследователю использовать один и тот же предварительно обученный весовой набор для совершенно разных задач, экономя время на настройку и обеспечивая согласованность предсказаний.

02Настройка окружения и аутентификация

Первым шагом к началу работы является установка необходимых библиотек. FAIRChem v2 требует наличия fairchem-core, ase (Atomic Simulation Environment), а также инструментов для работы с графическими процессорами (PyTorch). Для удобства и воспроизводимости, особенно если вы работаете в облачных средах вроде Google Colab, рекомендуется использовать скрипт, который автоматически проверяет наличие пакетов и устанавливает их при необходимости.

Важным аспектом является доступ к весам модели UMA. Модель является «gated» (закрытой), что означает, что для ее загрузки требуется аутентификация через платформу Hugging Face. Это связано с лицензионными соглашениями и контролем доступа к ресурсам. Ниже приведен код, который автоматизирует этот процесс: он пытается получить токен из переменных окружения или интерфейса Colab, а если это невозможно, запрашивает его у пользователя.

import importlib.util
import subprocess
import sys
import os

def _pip(pkgs):
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *pkgs])

if importlib.util.find_spec("fairchem") is None:
    print(">> Installing fairchem-core, ase, and helpers (takes ~2-4 min)...")
    _pip(["fairchem-core", "ase", "matplotlib", "huggingface_hub"])
    print(">> Installation done.")
else:
    print(">> fairchem already installed.")

from huggingface_hub import login

def whoami():
    pass

def hf_authenticate(token=None):
    try:
        from google.colab import userdata
        token = userdata.get("HF_TOKEN")
    except Exception:
        pass
    token = token or os.environ.get("HF_TOKEN")
    try:
        whoami()
        print(">> Already authenticated with Hugging Face.")
        return
    except Exception:
        pass
    if token is None:
        from getpass import getpass
        token = getpass("Paste your Hugging Face access token: ").strip()
    login(token=token)
    print(">> Hugging Face login OK.")

hf_authenticate()

После установки и аутентификации мы определяем устройство для вычислений. Если доступен CUDA (графический процессор NVIDIA), мы используем его для ускорения инференса, так как модели глубокого обучения работают на GPU значительно быстрее. В противном случае код автоматически переключится на CPU.

import numpy as np
import torch
import matplotlib.pyplot as plt
from fairchem.core import pretrained_mlip, FAIRChemCalculator

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(">> Using device: {DEVICE}")

MODEL = "uma-s-1p2"
predictor = pretrained_mlip.get_predict_unit(MODEL, device=DEVICE)

calc_mol = FAIRChemCalculator(predictor, task_name="omol")
calc_cat = FAIRChemCalculator(predictor, task_name="oc20")
calc_mat = FAIRChemCalculator(predictor, task_name="omat")

print(">> Loaded {MODEL} with omol / oc20 / omat calculators.")

03Молекулярная химия: Энергия и геометрия воды

Давайте начнем с классической задачи — моделирования молекулы воды. Используя калькулятор omol (organic molecules), мы можем вычислить потенциальную энергию и силы, действующие на атомы. Это базовая операция, но она демонстрирует, как UMA интегрируется с объектами ASE.

Мы создаем молекулу H2O, задаем ее заряд и спиновое состояние (для воды это обычно синглет, spin=0), и назначаем ей наш калькулятор. После этого вызываем методы get_potential_energy() и get_forces(). Полученные значения энергии позволяют нам рассчитать энергию атомизации — энергию, необходимую для разделения молекулы на отдельные атомы. Для этого нам нужны энергии изолированных атомов кислорода и водорода с правильными спиновыми мультиплетностями.

from ase.build import molecule
from ase import Atoms

print("\n" + "="*70)
print("SECTION 2: Single-point energetics of water (omol task)")
print("="*70)

h2o = molecule("H2O")
h2o.info.update({"charge": 0, "spin": 0})
h2o.calc = calc_mol

E_h2o = h2o.get_potential_energy()
F_h2o = h2o.get_forces()

print(f"E(H2O)            = {E_h2o:.4f} eV")
print(f"Max |force|       = {np.abs(F_h2o).max()::.4f} eV/A")

def atom_energy(symbol, spin):
    a = Atoms(symbol, positions=[[0,0,0]])
    a.info.update({"charge": 0, "spin": spin})
    a.calc = calc_mol
    return a.get_potential_energy()

E_O = atom_energy("O", 2)  # Кислород в основном состоянии (триплет)
E_H = atom_energy("H", 2)  # Водород в основном состоянии (дублет)

E_atomization = E_h2o - E_O - 2*E_H
print(f"Atomization energy of H2O = {E_atomization:.3f} eV (experiment ~ 9.5 eV incl. ZPE effects)")

Далее мы демонстрируем способность модели к оптимизации геометрии. Искусственно искажаем молекулу воды, сдвигая атомы, и запускаем оптимизатор LBFGS. Алгоритм будет итеративно изменять координаты атомов, минимизируя энергию системы, пока силы не станут близки к нулю. Результатом являются равновесные длины связей и углы, которые UMA предсказывает с высокой точностью, близкой к экспериментальным данным.

from ase.optimize import LBFGS

print("\n" + "="*70)
print("SECTION 3: Relaxing a deliberately distorted water molecule")
print("="*70)

h2o_bad = molecule("H2O")
h2o_bad.positions += 0.25 * np.random.randn(3, 3)  # Случайное искажение
h2o_bad.info.update({"charge": 0, "spin": 0})
h2o_bad.calc = calc_mol

opt = LBFGS(h2o_bad, logfile=None)
energies_opt = []
opt.attach(lambda: energies_opt.append(h2o_bad.get_potential_energy()))
opt.run(fmax=0.01, steps=200)

d_OH = h2o_bad.get_distance(0, 1)
ang = h2o_bad.get_angle(1, 0, 2)

print(f"Converged in {opt.get_number_of_steps()} steps")
print(f"O-H bond length   = {d_OH:.3f} A   (expt ~0.958 A)")
print(f"H-O-H angle       = {ang:.1f} deg (expt ~104.5 deg)")

plt.figure(figsize=(3.2, 2.5))
plt.plot(energies_opt, "o-")
plt.xlabel("Optimizer step")
plt.ylabel("Energy (eV)")
plt.title("H2O geometry optimization")
plt.tight_layout()
plt.show()

04Квантовая химия: Разрыв синглет-триплетного зазора

Одной из сильных сторон UMA является ее способность различать различные электронные состояния. Рассмотрим метилен (CH2). В основном состоянии он существует в виде триплета, но также может находиться в возбужденном синглетном состоянии. Разница в энергии между этими состояниями (spin-state gap) критически важна для понимания реакционной способности метилена.

Мы создаем две отдельные структуры CH2: одну с синглетным спином, другую с триплетным. Для каждой создаем отдельный калькулятор, явно указывая задачу omol. Затем вычисляем энергии и находим разницу. UMA успешно предсказывает, что триплетное состояние энергетически выгоднее, что согласуется с экспериментальными данными.

print("\n" + "="*70)
print("SECTION 4: CH2 singlet-triplet gap (UMA is spin-aware!)")
print("="*70)

singlet = molecule("CH2_s1A1d")
singlet.info.update({"charge": 0, "spin": 0})
singlet.calc = FAIRChemCalculator(predictor, task_name="omol")

triplet = molecule("CH2_s3B1d")
triplet.info.update({"charge": 0, "spin": 2})
triplet.calc = FAIRChemCalculator(predictor, task_name="omol")

gap = triplet.get_potential_energy() - singlet.get_potential_energy()
print(f"E(triplet) - E(singlet) = {gap:.3f} eV  (negative => triplet ground state; expt ~ -0.39 eV)")

05Реакционная энергетика и колебательные спектры

Перейдем к более сложным химическим трансформациям. Рассчитаем энергию реакции сгорания метана: CH4 + 2 O2 -> CO2 + 2 H2O. Для этого нам нужно оптимизировать геометрию каждого реагента и продукта, а затем суммировать их энергии. Важно использовать релаксированные структуры, так как энергии несвязанных или неоптимизированных молекул некорректны для термодинамических расчетов.

Помимо энергетики, UMA позволяет проводить анализ колебательных частот. Используя класс Vibrations из ASE, мы можем вычислить нормальные моды колебаний молекулы воды. Это требует вычисления матрицы Гессе (вторых производных энергии по координатам), что модель делает численно. Мы получаем частоты в см⁻¹, которые можно сравнить с инфракрасными спектрами, полученными в лаборатории.

from ase.vibrations import Vibrations

print("\n" + "="*70)
print("SECTION 6: Vibrational frequencies of relaxed H2O")
print("="*70)

# Сначала релаксируем воду, чтобы получить правильную геометрию для колебаний
h2o_relaxed = molecule("H2O")
h2o_relaxed.info.update({"charge": 0, "spin": 0})
h2o_relaxed.calc = calc_mol
LBFGS(h2o_relaxed, logfile=None).run(fmax=0.01, steps=200)

vib = Vibrations(h2o_relaxed, name="vib_h2o")
vib.run()
freqs = np.real(vib.get_frequencies())
real_modes = [f for f in freqs if f > 200]  # Отбрасываем мнимые частоты (артефакты)

print("Vibrational modes (cm^-1):")
print(", ".join([f"{f:.0f}" for f in real_modes]))
print("Experimental H2O: 1595 (bend), 3657 (sym stretch), 3756 (asym stretch)")
print(f"Zero-point energy = {vib.get_zero_point_energy():.3f} eV")
vib.clean()

06Катализ: Адсорбция CO на меди

Теперь обратимся к домену катализа, используя калькулятор oc20. Мы моделируем поверхность меди Cu(100) и адсорбцию молекулы монооксида углерода (CO). Для этого создаем периодическую суперячейку (slab), фиксируем нижние слои атомов меди (чтобы имитировать объемный кристалл) и размещаем CO на мостиковом сайте (bridge site).

Энергия адсорбции рассчитывается по циклу: E_ads = E(slab+CO) - E(slab) - E(CO gas). Важно отметить, что OC20 использует специфическую схему энергетических ссылок, поэтому для публикации научных результатов необходимо соблюдать консистентность в выборе референсных состояний. Тем не менее, для сравнительного анализа и скрининга катализаторов этот подход крайне эффективен.

from ase.build import fcc100, add_adsorbate
from ase.constraints import FixAtoms

print("\n" + "="*70)
print("SECTION 7: CO/Cu(100) relaxation + adsorption energy (oc20 task)")
print("="*70)

slab = fcc100("Cu", size=(3,3,4), vacuum=8.0, periodic=True)
slab.set_constraint(FixAtoms(mask=[tag < 4 for tag in slab.get_tags()]))
add_adsorbate(slab, molecule("CO"), height=2.0, position="bridge")
slab.calc = calc_cat

opt = LBFGS(slab, logfile=None)
opt.run(fmax=0.05, steps=300)
E_slab_ads = slab.get_potential_energy()
print(f"Relaxed CO/Cu(100) in {opt.get_number_of_steps()} steps, E = {E_slab_ads:.3f} eV")

# Расчет энергии адсорбции требует референсов
# (Упрощенный пример для демонстрации)

07Материаловедение: Железо и уравнение состояния

В домене материалов (omat) мы рассматриваем объемное тело. Возьмем объемно-центрированную кубическую решетку (BCC) железа. Мы не только оптимизируем позиции атомов, но и параметры ячейки, используя FrechetCellFilter. Это позволяет модели найти равновесный параметр решетки, который должен быть близок к экспериментальному (~2.866 Å).

Далее мы строим уравнение состояния (Equation of State, EOS). Для этого мы сжимаем и растягиваем ячейку в диапазоне от 94% до 106% от равновесного объема, вычисляем энергию для каждой точки и подгоняем данные к уравнению Бирча-Мюрнагана. Наклон кривой в минимуме позволяет вычислить модуль объемного сжатия (Bulk Modulus), который для железа составляет около 170 ГПа. UMA предсказывает эти макроскопические свойства с удивительной точностью, исходя из атомарных взаимодействий.

from ase.build import bulk
from ase.optimize import FIRE
from ase.filters import FrechetCellFilter
from ase.eos import EquationOfState

print("\n" + "="*70)
print("SECTION 8: BCC iron — full cell relaxation and bulk modulus (omat)")
print("="*70)

fe = bulk("Fe", "bcc", 2.9)
fe.calc = calc_mat

# Релаксация ячейки
opt = FIRE(FrechetCellFilter(fe), logfile=None)
opt.run(fmax=0.02, steps=300)
a_relaxed = fe.cell.cellpar()[0]
print(f"Relaxed BCC Fe lattice constant = {a_relaxed:.3f} A (expt ~2.866 A)")

# Построение EOS
volumes, energies = [], []
cell0 = fe.get_cell()
for scale in np.linspace(0.94, 1.06, 15):
    fe_copy = fe.copy()
    fe_copy.set_cell(cell0 * scale, scale_atoms=True)
    fe_copy.calc = calc_mat
    volumes.append(fe_copy.get_volume())
    energies.append(fe_copy.get_potential_energy())

eos = EquationOfState(volumes, energies)
eos.fit("birchmurnaghan")
v0, e0, B = eos.parameters()

print(f"Equilibrium volume = {v0:.2f} A^3/cell")
print(f"Bulk modulus       = {B:.0f} GPa (expt ~170 GPa for Fe)")

08Молекулярная динамика и сканирование поверхности

Финальная часть демонстрирует динамику. Мы запускаем Langevin-моделирование молекулы воды при температуре 300 К. Это позволяет наблюдать тепловые флуктуации связей и углов во времени. Мы отслеживаем температуру, потенциальную энергию и длину связи O-H. Несмотря на то, что модель обучалась на статических структурах, она способна генерировать физически правдоподобные траектории, сохраняя целостность молекулы.

Также мы проводим сканирование потенциальной энергетической поверхности (PES), искусственно растягивая одну связь O-H. Это показывает, как энергия системы растет при диссоциации связи, и позволяет найти энергию диссоциации. График PES имеет характерную яму потенциала, глубина которой соответствует прочности связи.

from ase.md.langevin import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution

print("\n" + "="*70)
print("SECTION 9: 0.5 ps Langevin MD of a water molecule at 300 K")
print("="*70)

md_atoms = molecule("H2O")
md_atoms.info.update({"charge": 0, "spin": 0})
MaxwellBoltzmannDistribution(md_atoms, temperature_K=300)
md_atoms.calc = FAIRChemCalculator(predictor, task_name="omol")

dyn = Langevin(md_atoms, timestep=0.5*units.fs, temperature_K=300, friction=0.01/units.fs)

times, temps, epots = [], [], []
def log_md():
    times.append(dyn.get_number_of_steps()*0.5)
    temps.append(md_atoms.get_temperature())
    epots.append(md_atoms.get_potential_energy())

dyn.attach(log_md, interval=10)
dyn.run(1000)

print(f"MD done.  = {np.mean(temps[10:]):.0f} K,  = {np.mean([md_atoms.get_distance(0,1), md_atoms.get_distance(0,2)])/2:.3f} A")

09Что это значит на практике

Интеграция FAIRChem v2 и UMA открывает новые горизонты для исследователей в области химии и материаловедения. Главное практическое значение заключается в скорости и универсальности. Вы можете провести скрининг тысяч потенциальных катализаторов или молекул-кандидатов для лекарств за часы, а не месяцы, как это требовалось бы при использовании DFT. При этом точность предсказаний остается на уровне, приемлемом для принятия решений.

Для российских исследователей важно отметить, что локальный запуск таких моделей требует мощных GPU. Однако, благодаря открытости кода и доступности весов через Hugging Face, вы можете развернуть решение на любой доступной инфраструктуре. Интеграция с ASE делает этот инструмент совместимым с огромным экосистемой существующих программных пакетов, что позволяет легко встраивать его в существующие рабочие процессы.

Мы рассмотрели полный цикл: от установки и аутентификации до сложных симуляций молекулярной динамики и анализа материалов. Этот фундамент позволяет вам экспериментировать с другими моделями UMA (например, более точными версиями), исследовать новые домены, такие как MOF (металло-органические каркасы), и масштабировать вычисления на несколько GPU. Будущее вычислительной химии — за универсальными, обученными на больших данных моделями, и FAIRChem v2 является одним из лидеров в этой области.

Источник: MarkTechPost ↗