OA-ReactDiff｜当扩散生成模型遇到化学反应 -- 现代“AI炼金术”

从“等价交换”的远古炼金术开始，化学一直是一门了解和控制物质间相互反应的学科。经过不断解锁和利用新的化学反应，人类在不断研发出新材料便利自己的生活的同时也在提升能量利用效率，促进可持续发展。

一个基础化学反应由反应物，过渡态，生成物三者构成。这其中过渡态对于用于理解化学反应机理和估计反应势垒高度尤为重要。然而，由于其在反应过程中存在的时间极短（飞秒量级），实验中几乎不可能分离和表征过渡态。通常情况下，大家使用量子化学的计算方法，通过已知的反应物和生成物来搜索过渡态。然而，这些计算较为昂贵，并以经常失败而“臭名远扬”。同时，受限于个人的经验直觉和计算所需的资源，每个人所能探索的化学反应也是局限的。这种限制在研究未知的复杂反应时尤为“致命”。它会使我们忽略到一些潜在可能发生的反应，导致会反应机理的误判，进而影响催化材料设计的思路。

最近，MIT的段辰儒博士等开发了一种新的扩散生成模型，OA-ReactDiff，它可以在保持所有物理上必要的对称性下快速生成新的化学反应。他们将这个生成式AI模型运用在了过渡态。他们发现OA-ReactDiff将过渡态搜索时间从原本的一天左右降到了6秒，并极大的提升了找到正确过渡态的概率。同时，借助生成式AI的随机性，OA-ReactDiff可以探索到意料之外的化学反应。这个特点弥补了现有基于化学的直觉反应探索框架，帮助建立更加完整的化学反应网络，助力研发设计新型催化材料。

---

OA-ReactDiff@Notebook

🔗 Notebook 链接

本期Notebook，我们给大家带来 OA-ReactDiff 的方法设计思路和实际场景应用，帮助大家学习扩散生成模型在化学和材料的一些潜在的应用和难点，并带领大家了解 OA-ReactDiff 运行的原理并应用于你的项目中。

OA-ReactDiff使用场景

1. 无需任何预处理，生成就是这么简单

在计算中，传统的机器学习方法通常需要将反应物和生成物的之间的原子一一对应，如果有多个反应物，反应物中的每个分子之间的几何位置还需要被精心调整。这些预处理步骤，不仅花费时间，在很多未知的反应中实际上是不可行的。

而OA-ReactDiff保证了所有化学反应中所需的对称性，所以我们不需要对反应做任何的预处理。

# --- 必要的函数导入 ---
from torch.utils.data import DataLoader

from oa_reactdiff.trainer.pl_trainer import DDPMModule


from oa_reactdiff.dataset import ProcessedTS1x
from oa_reactdiff.diffusion._schedule import DiffSchedule, PredefinedNoiseSchedule

from oa_reactdiff.diffusion._normalizer import FEATURE_MAPPING
from oa_reactdiff.analyze.rmsd import batch_rmsd

from oa_reactdiff.utils.sampling_tools import (
    assemble_sample_inputs,
    write_tmp_xyz,
)

导入预训练的模型并重新定义schedule

device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda")

ddpm_trainer = DDPMModule.load_from_checkpoint(
    checkpoint_path="/opt/src/pretrained-ts1x-diff.ckpt",
    map_location=device,
)
ddpm_trainer = ddpm_trainer.to(device)

noise_schedule: str = "polynomial_2"
timesteps: int = 150
precision: float = 1e-5

gamma_module = PredefinedNoiseSchedule(
            noise_schedule=noise_schedule,
            timesteps=timesteps,
            precision=precision,
        )
schedule = DiffSchedule(
    gamma_module=gamma_module,
    norm_values=ddpm_trainer.ddpm.norm_values
)
ddpm_trainer.ddpm.schedule = schedule
ddpm_trainer.ddpm.T = timesteps
ddpm_trainer = ddpm_trainer.to(device)

准备dataset和data loader并选取一个多分子参与的反应

dataset = ProcessedTS1x(
    npz_path="/opt/src/oa_reactdiff/data/transition1x/train.pkl",
    center=True,
    pad_fragments=0,
    device=device,
    zero_charge=False,
    remove_h=False,
    single_frag_only=False,
    swapping_react_prod=False,
    use_by_ind=True,
)
loader = DataLoader(
    dataset, 
    batch_size=1,
    shuffle=False,
    num_workers=0,
    collate_fn=dataset.collate_fn
)
itl = iter(loader)
idx = -1

len(dataset)

for _ in range(4):  # 第4个样本正好是一个多分子反应
    representations, res = next(itl)
idx += 1
n_samples = representations[0]["size"].size(0)
fragments_nodes = [
    repre["size"] for repre in representations
]
conditions = torch.tensor([[0] for _ in range(n_samples)], device=device)

故意将反应物的原子顺序相对于生成物打乱，并将生成物中的两个分子拉倒无限远（10 A）

new_order_react = torch.randperm(representations[0]["size"].item())
for k in ["pos", "one_hot", "charge"]:
    representations[0][k] = representations[0][k][new_order_react]
    
xh_fixed = [
    torch.cat(
        [repre[feature_type] for feature_type in FEATURE_MAPPING],
        dim=1,
    )
    for repre in representations
]

对这组没有任何原子排序和几何排列的「反应物，生成物」生成过渡态结构

out_samples, out_masks = ddpm_trainer.ddpm.inpaint(
    n_samples=n_samples,
    fragments_nodes=fragments_nodes,
    conditions=conditions,
    return_frames=1,
    resamplings=5,
    jump_length=5,
    timesteps=None,
    xh_fixed=xh_fixed,
    frag_fixed=[0, 2],
)

评测生成过渡态结构的质量

rmsds = batch_rmsd(
    fragments_nodes, 
    out_samples[0],
    xh_fixed,
    idx=1,
)
write_tmp_xyz(
    fragments_nodes, 
    out_samples[0], 
    idx=[0, 1, 2], 
    localpath="/opt/src/demo/inpainting"
)

rmsds = [min(1, _x) for _x in rmsds]
[(ii, round(rmsd, 2)) for ii, rmsd in enumerate(rmsds)], np.mean(rmsds), np.median(rmsds)

可以看到我们生成的过渡态原子序数和反应物原子序数是完全无需对应的，证明了OA-ReactDiff不需要原子排列的预处理

! cat /opt/src/demo/inpainting/gen_0_react.xyz

! cat /opt/src/demo/inpainting/gen_0_ts.xyz

同时，虽然生成物由多个分子构成，也无需对它们进行特定的几何排列 （我们已经在生成前把两个分子拉到了”无穷“远）

draw_in_3dmol(
    open("/opt/src/demo/inpainting/gen_0_prod.xyz", "r").read(),
    "xyz"
)

2. 随机过程的馈赠，弥补化学直觉

由于噪声的介入，扩散生成模型有一定的随机性，这导致我们及时给定了「反应物，生成物」的3D结构，每次生成出来的过渡态结构也不尽相同。

那么一个自然地问题是，这些随机生成的结构，除了和真正的过渡态一样的那个，剩余的结构有化学意义和应用价值吗❓

我们将深入探索一个案例得到这个问题的答案。

from glob import glob
import plotly.express as px

from oa_reactdiff.analyze.rmsd import xyz2pmg, pymatgen_rmsd

from pymatgen.core import Molecule
from collections import OrderedDict


def draw_reaction(react_path: str, idx: int = 0, prefix: str = "gen") -> py3Dmol.view:
    """画出反应的的{反应物，过渡态，生成物}

    Args:
        react_path (str): path to the reaction.
        idx (int, optional): index for the generated reaction. Defaults to 0.
        prefix (str, optional): prefix for distinguishing true sample and generated structure.
            Defaults to "gen".

    Returns:
        py3Dmol.view: _description_
    """
    with open(f"{react_path}/{prefix}_{idx}_react.xyz", "r") as fo:
        natoms = int(fo.readline()) * 3
    mol = f"{natoms}\n\n"
    for ii, t in enumerate(["react", "ts", "prod"]):
        pmatg_mol = xyz2pmg(f"{react_path}/{prefix}_{idx}_{t}.xyz")
        pmatg_mol_prime = Molecule(
            species=pmatg_mol.atomic_numbers,
            coords=pmatg_mol.cart_coords + 8 * ii,
        )
        mol += "\n".join(pmatg_mol_prime.to(fmt="xyz").split("\n")[2:]) + "\n"
    viewer = py3Dmol.view(1024, 576)
    viewer.addModel(mol, "xyz")
    viewer.setStyle({'stick': {}, "sphere": {"radius": 0.3}})
    viewer.zoomTo()
    return viewer

画出我们的案例反应，从左下到右上分别对应「反应物，过渡态，和生成物」

draw_reaction("/opt/src/demo/example-3/ground_truth", prefix="sample")

draw_in_3dmol(
    open("/opt/src/demo/example-3/ground_truth/sample_0_ts.xyz", "r").read(),
    "xyz"
)

将这些生成的过渡态分类排序

opt_ts_path = "/opt/src/demo/example-3/opt_ts/"
opt_ts_xyzs = glob(f"{opt_ts_path}/*ts.opt.xyz")

order_dict = {}
for xyz in opt_ts_xyzs:

    order_dict.update(
        {int(xyz.split("/")[-1].split(".")[0]): xyz}
    )
order_dict = OrderedDict(sorted(order_dict.items()))

opt_ts_xyzs = []
ind_dict = {}
for ii, v in enumerate(order_dict.values()):
    opt_ts_xyzs.append(v)
    ind_dict.update(
        {ii: v}
    )

计算这些过渡态之间的配对RMSD矩阵。如我们之前的对比图所见，RMSD低于「0.1」的分子的构型可以认为是一样的

n_ts = len(opt_ts_xyzs)
rmsd_mat = np.ones((n_ts, n_ts)) * -2.5
for ii in range(n_ts):
    for jj in range(ii+1, n_ts):
        try:
            rmsd_mat[ii, jj] = np.log10(
                pymatgen_rmsd(
                    opt_ts_xyzs[ii],
                    opt_ts_xyzs[jj],
                    ignore_chirality=True,
                )
            )
        except:
            print(ii, jj)
            pass
        rmsd_mat[jj, ii] = rmsd_mat[ii, jj]

通过K-Means将这些生成的过渡态根据结构分类

from sklearn.cluster import KMeans

def reorder_matrix(matrix, n_clusters):
    # Apply K-means clustering to rows and columns
    row_clusters = KMeans(n_clusters=n_clusters).fit_predict(matrix)
    
    # Create a permutation to reorder rows and columns
    row_permutation = np.argsort(row_clusters)
    col_permutation = np.argsort(row_clusters)

    # Apply the permutation to the matrix
    reordered_matrix = matrix[row_permutation][:, col_permutation]

    return reordered_matrix, row_permutation, row_clusters


n = n_ts  # 总体过渡态的数目
n_clusters = 6  # 我们K-Means的聚类数目

reordered_matrix, row_permutation, row_clusters = reorder_matrix(rmsd_mat, n_clusters)

展示分类结果。可以清晰的看见很多过渡态结构是比较相近的（对角线上的红色方块为一个聚类）

from IPython.display import HTML

fig = px.imshow(
    reordered_matrix, 
    color_continuous_scale="Oryel_r",
    range_color=[-2, -0.3],
)
fig.layout.font.update({"size": 18, "family": "Arial"})

fig.layout.update({"width": 650, "height": 500})
HTML(fig.to_html())

直观的总结一下不同的过渡态属于哪一个聚类

import json

cluster_dict = {}
for ii, cluster in enumerate(row_clusters):
    cluster = str(cluster)
    if cluster not in cluster_dict:
        cluster_dict[cluster] = [ind_dict[ii]]
    else:
        cluster_dict[cluster] += [ind_dict[ii]]

cluster_dict = OrderedDict(sorted(cluster_dict.items()))
cluster_dict

可以看到，在不同聚类的过渡态的几何结构完全不同

draw_in_3dmol(
    open("/opt/src/demo/example-3/generated/gen_36_ts.xyz", "r").read(),
    "xyz"
)

3. 给我一些原子，我来生成所有反应

除了寻找已知反应的过渡态，OA-ReactDiff还可以直接无条件生成新的化学反应，并用于探索反应网络

在这里，我们假想一个只有四个原子，C, N, O, H的烧杯，并使用OA-ReactDiff探索所有这四个原子可能发生的反应。

xyz_path = "/opt/src/demo/CNOH/"
n_samples = 128  # 生成的总反应数目
natm = 4  # 反应物的原子数目
fragments_nodes = [
    torch.tensor([natm] * n_samples, device=device),
    torch.tensor([natm] * n_samples, device=device),
    torch.tensor([natm] * n_samples, device=device),
]

conditions = torch.tensor([[0]] * n_samples, device=device)
h0 = assemble_sample_inputs(
    atoms=["C"] * 1 + ["O"] * 1 + ["N"] * 1 + ["H"] * 1,  # 反应物的原子种类，这里是CNOH各一个
    device=device,
    n_samples=n_samples,
    frag_type=False,
)

out_samples, out_masks = ddpm_trainer.ddpm.sample(
    n_samples=n_samples,
    fragments_nodes=fragments_nodes,
    conditions=conditions,
    return_frames=1,
    timesteps=None,
    h0=h0,
)

将生成的结果保存成xyz文件形式

write_tmp_xyz(
    fragments_nodes, 
    out_samples[0], 
    idx=[0, 1, 2], 
    ex_ind=0,
    localpath=xyz_path,
)

随机画一个生成的反应

idx = 2
assert idx < n_samples
views = draw_reaction(xyz_path, idx)
views

分析反应结果，找出独特不重复的稳定分子

from glob import glob

from pymatgen.io.xyz import XYZ
from openbabel import pybel

from oa_reactdiff.analyze.rmsd import pymatgen_rmsd


def xyz_to_smiles(fname: str) -> str:
    """将xyz格式的分子转换成smiles格式

    Args:
        fname (str): path to the xyz file.

    Returns:
        str: SMILES string.
    """
    mol = next(pybel.readfile("xyz", fname))
    smi = mol.write(format="can")
    return smi.split()[0].strip()

xyzfiles = glob(f"{xyz_path}/gen*_react.xyz") + glob(f"{xyz_path}/gen*_prod.xyz")
xyz_converter = XYZ(mol=None)
mol = xyz_converter.from_file(xyzfiles[0]).molecule
unique_mols = {xyzfiles[0]: mol}
for _xyzfile in xyzfiles:
    _mol = xyz_converter.from_file(_xyzfile).molecule
    min_rmsd = 100
    for _, mol in unique_mols.items():
        rmsd = pymatgen_rmsd(mol, _mol, ignore_chirality=True, threshold=0.5)
        min_rmsd = min(min_rmsd, rmsd)
    if min_rmsd > 0.1:  # 如果和已有的分子的rmsd都大于0.1，那么就认为是一个新的分子
        unique_mols.update({_xyzfile: _mol})
        
len(unique_mols)

我们这里只考虑稳定的（即非自由基等等）的分子

unique_idx = []
unique_smiles = []
idx = 0
for file in unique_mols:
    smi = xyz_to_smiles(file)
    if smi not in unique_smiles and not "." in smi:
        unique_smiles.append(smi)
        unique_idx.append(idx)
    idx += 1
unique_idx, unique_smiles  # 独特的分子对应的反应index和smiles

现在来画出这其中的一个分子

idx = np.random.choice(unique_idx, 1)[0]

draw_in_3dmol(open(list(unique_mols.keys())[idx], "r").read(), "xyz")

独特的反应路径：我们这里只分析稳定的分子之间的反应

unique_paths = {}
path_index = {}
for ii in range(n_samples):
    r_xyz = f"{xyz_path}/gen_{ii}_react.xyz"
    p_xyz = f"{xyz_path}/gen_{ii}_prod.xyz"
    path = set([xyz_to_smiles(r_xyz), xyz_to_smiles(p_xyz)])
    use = True
    for smi in path:
        if smi not in unique_smiles:
            use = False
    if not path in unique_paths.values() and len(path) > 1 and use:
        unique_paths[ii] = path
    if not (len(path) > 1 and use):
        continue
    sorted_smi = " & ".join(list(sorted(path)))
    if sorted_smi not in path_index:
        path_index[sorted_smi] = [ii]
    else:
        path_index[sorted_smi] += [ii]
mols_in_paths = []
for k, v in unique_paths.items():
    for _v in v:
        if not _v in mols_in_paths:
            mols_in_paths.append(_v)
            
mols_in_paths, len(mols_in_paths)

# 所有的CNOH四原子稳定反应如下：
unique_paths, len(unique_paths)

最后选取一个看一下这个反应是不是化学上合理的

idx = np.random.choice(list(unique_paths.keys()), 1)[0]
views = draw_reaction(xyz_path, idx)
views

阅读并运行完整 Notebook 请见：https://nb.bohrium.dp.tech/detail/1221842760