空间站广场

论文

Notebooks

比赛

课程

Apps

我的主页

我的Notebooks

我的论文库

我的足迹

我的工作空间

任务

节点

文件

数据集

镜像

项目

数据库

公开

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

AI4S

Deep Learning

Tutorial

PyTorch

DDPMs

Catalysis

DFT

AI4SDeep LearningTutorialPyTorchDDPMsCatalysisDFT

段辰儒

发布于 2023-10-05

推荐镜像 :reactdiff:gpu

推荐机型 :c12_m92_1 * NVIDIA V100

简单介绍一下扩散生成模型 (DDPM)

已经玩转了图像，那试试生成分子如何？

生成分子时模型的SE(3)准变性不可或缺

建立一个LEFTNet模型

创建一个水分子作为测试体系

使用LEFTNet正向预测一个scalar和vector

随机旋转一下，再来一次

不管是scalar，还是vector，结果都是等变的

化学的核心是反应，直接生成化学反应如何？

当我们旋转化学反应的任何分子，这个反应会变化吗？

全体系SE(3)大失败 - 没有保持单个对象的对称性

考虑一个假想的“电解水反应”， H2O -> H2 + O.

正向预测

我们旋转一下反应物

再预测一次这个反应。注意这两个反应的化学实质是一样的

但是我们scalar和vector上都丢失了等变性❌

我们在化学反应中需要”对象感知“的SE(3)模型

创立一个”对象感知“的LEFTNet

再来重复一遍刚才的旋转测试

经过检验，”对象感知“的LEFTNet会保持我们想要的对称性✅

从天到秒，OA-ReactDiff加速生成过渡态结构

正确的对称性移除一切繁杂预处理的必要

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

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

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

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

意外惊喜 - 来自随机过程的馈赠

画出我们的案例反应

对于同一个反应，我们产生了多少不同的TS？

我们生成的其余过渡态一文不值？

给我一些原子，我可以生成关于它们的所有的反应

随机生成反应

分析反应结果

找出独特不重复的稳定分子

独特的反应路径

展望: 让我们”生成“未来

【OA-ReactDiff】是首个生成「3D化学反应」的扩散生成模型，它不仅『加速了1000倍』的化学反应中的3D结构搜索，还可以生成和探索『全新未知』的化学反应。

在完成这个教程后，你将能够：

学习到扩散生成模型在化学和材料的一些应用和潜在的难点
了解【OA-ReactDiff】运行的原理并应用于你的项目中

阅读和跑通该教程【最多】需 20 分钟，让我们开始吧！

文章链接 👉 OA-ReactDiff

如果大家对把OA-ReactDiff应用到自己工作的想法，但不确定是不是完全合适或者不知道如何下手，欢迎联系作者段辰儒本人，duanchenru@gmail.com

代码

文本

©️ Copyright 2023 @ 段辰儒（Chenru Duan)
日期：2023-10-XX
共享协议：本作品采用知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议进行许可。
快速开始：点击上方的开始连接按钮，推荐选择 reactdiff:gpu 镜像及 c12_m92_1*NVIDIA V100 节点配置，连接之后即可运行（其余GPU节点也可，但运行速度稍慢）。如选择CPU节点，请相应选择reactdiff:cpu-demo 镜像。
全代码运行时间：1 V100 GPU: 2分钟; 32 core CPU: 8分钟

代码

文本

简单介绍一下扩散生成模型 (DDPM)

DDPM（Diffusion Probabilistic Models）是一种基于概率统计生成和随机噪声生成数据的扩散模型。它通过模拟一个扩散过程，将嘈杂的数据转化为干净的数据样本。

这些模型专注于建模数据的分布，并模拟从一个简单分布到更复杂分布的逐步演变过程。扩散模型的基本概念是将一个简单且易于抽样的分布，通常是高斯分布，转化为更复杂的数据分布。这个转化是通过一系列可逆操作实现的。一旦模型学习了这个转化过程，它可以通过从简单分布中的一个点开始，逐渐"逆扩散"到所需的复杂数据分布，从而生成新的样本。

DDPM模型的训练需要获取扩散过程参数的知识，有效地捕捉每个转换步骤中干净和嘈杂数据之间的关系。在生成过程中，DDPM从嘈杂数据（例如，嘈杂图像）开始，迭代地逆向应用已学到的变换，以获得去噪和逼真的数据样本。

想要上手把玩一下DDPM图像生成的小伙伴们可以参考 Notebook: Diffusion Models初探：你还是不懂扩散模型的基本原理？

图一｜扩散生成模型在生成图像

拓展阅读 🔽

An Introduction to Diffusion Models for Machine Learning

Diffusion Models: A Comprehensive Survey of Methods and Applications

代码

文本

已经玩转了图像，那试试生成分子如何？

最近，DDPM的框架被逐渐应用在了化学分子相关的任务上。比如生成有机小分子的3D结构，蛋白-小分子对接，和基于蛋白质结构的药物设计。

分子的生成与上面的图像生成的核心区别在于他们的表示和对称性：

表示：对于NxN像素的彩色图像，一般我们用一个(N, N, 3)的张量表示，其中最后的3表示RGB三原色。而对于一个分子，我们需要使用（原子类别，坐标信息）来表示每一个原子。
对称性：在图像中，我们一般不要求严格的对称性。大家通常使用数据增强（augmentation）的方式使模型”记住“物体关于2D旋转的不变形。但是在分子中，因为我们直接学习三维的物体并且要求的精准度极高，数据增强往往会过于昂贵且不能满足精准度要求。于是，我们需要把对称性嵌入至模型之中。

图二｜扩散生成模型生成单一分子的过程

代码

文本

生成分子时模型的SE(3)准变性不可或缺

我们在生成3D分子是需要保持分子上物理的对称性。其中一种方式就是使用具有SE(3)准变性的图网络模型。简单理解，就是对于一个分子，我们先旋转分子(D(g))再使用模型预测(f)和先用模型预测再旋转得到的结果是一样的，即 $D (g) \circ f = f \circ D^{'} (g)$

Siyuan之前有一篇对SE(3)准变性详细介绍的notebook，希望深挖的同学们可以补充学习📚。我们这里只举一个例子给大家一个可视化的体验。

代码

文本

[ ]

# --- 导入和定义一些函数 ----

import torch

import py3Dmol

import numpy as np

from typing import Optional

from torch import tensor

from e3nn import o3

from torch_scatter import scatter_mean

from oa_reactdiff.model import LEFTNet

default_float = torch.float64

torch.set_default_dtype(default_float) # 使用双精度，测试更准确

def remove_mean_batch(

x: tensor,

indices: Optional[tensor] = None

) -> tensor:

"""将x中的每个batch的均值去掉

Args:

x (tensor): input tensor.

indices (Optional[tensor], optional): batch indices. Defaults to None.

Returns:

tensor: output tensor with batch mean as 0.

"""

if indices == None:

return x - torch.mean(x, dim=0)

mean = scatter_mean(x, indices, dim=0)

x = x - mean[indices]

return x

def draw_in_3dmol(mol: str, fmt: str = "xyz") -> py3Dmol.view:

"""画分子

Args:

mol (str): str content of molecule.

fmt (str, optional): format. Defaults to "xyz".

Returns:

py3Dmol.view: output viewer

"""

viewer = py3Dmol.view(1024, 576)

viewer.addModel(mol, fmt)

viewer.setStyle({'stick': {}, "sphere": {"radius": 0.36}})

viewer.zoomTo()

return viewer

def assemble_xyz(z: list, pos: tensor) -> str:

"""将原子序数和位置组装成xyz格式

Args:

z (list): chemical elements

pos (tensor): 3D coordinates

Returns:

str: xyz string

"""

natoms =len(z)

xyz = f"{natoms}\n\n"

for _z, _pos in zip(z, pos.numpy()):

xyz += f"{_z}\t" + "\t".join([str(x) for x in _pos]) + "\n"

return xyz

代码

文本

建立一个LEFTNet模型

为了验证SE(3)准变性，我们进行一个简单的测试。这里的模型只用于测试，所以我们只需要建立一个非常小的模型。

注：LEFTNet是一个新的SOTA级别的SE(3)图神经网络。虽然我们这里使用LEFTNet，但其所展现的性质是与模型无关的（其他SE(3)模型，如EGNN，会给出相同的结果）

代码

文本

[ ]

num_layers = 2

hidden_channels = 8

in_hidden_channels = 4

num_radial = 4

model = LEFTNet(

num_layers=num_layers,

hidden_channels=hidden_channels,

in_hidden_channels=in_hidden_channels,

num_radial=num_radial,

object_aware=False,

)

sum(p.numel() for p in model.parameters() if p.requires_grad)

代码

文本

创建一个水分子作为测试体系

代码

文本

[ ]

h = torch.rand(3, in_hidden_channels)

z = ["O", "H", "H"]

pos = tensor([

[0, 0, 0],

[1, 0, 0],

[0, 1, 0],

]).double() # 方便起见，我们这里把H-O-H的角度设为90度

edge_index = tensor([

[0, 0, 1, 1, 2, 2],

[1, 2, 0, 2, 0, 1]

]).long() # 使用全连接的方式，这里的边是无向的

代码

文本

[ ]

xyz = assemble_xyz(z, pos)

view = draw_in_3dmol(xyz, "xyz")

view # 显示分子

代码

文本

使用LEFTNet正向预测一个scalar和vector

代码

文本

[ ]

_h, _pos, __ = model.forward(

h=h,

pos=remove_mean_batch(pos),

edge_index=edge_index,

) # _h是输出的特征，_pos是输出的位置，前者为scalar，后者为vector

代码

文本

随机旋转一下，再来一次

代码

文本

[ ]

rot = o3.rand_matrix()

pos_rot = torch.matmul(pos, rot).double()

_h_rot, _pos_rot, __ = model.forward(

h=h,

pos=remove_mean_batch(pos_rot),

edge_index=edge_index,

)

代码

文本

不管是scalar，还是vector，结果都是等变的

代码

文本

[ ]

torch.max(

torch.abs(

_h - _h_rot

)

) # 旋转后的h应该不变

代码

文本

[ ]

torch.max(

torch.abs(

torch.matmul(_pos, rot).double() - _pos_rot

)

) # 旋转后的pos应该旋转

代码

文本

化学的核心是反应，直接生成化学反应如何？

从古老的炼金术开始，化学一直是一门研究和控制物质之间相互反应的学科。既然已经成功实现了生成单一分子，我们为何不把目前的扩散生成模型推广到生成一组分子，比如化学中大家都关心的「反应物，过渡态，生成物」？

图三｜化学反应图示。黑线为能量曲线，「反应物, 生成物」对应极点，「过渡态」对应鞍点

带着这样的思考，我们将扩散生成推广到多体系生成。训练时，我们同时向「反应物（蓝），过渡态（黄），生成物（红）」加入噪声，并使用一个拥有特殊对称性（后面会详细描述）的图神经网络模型预测加入的三个噪声。生成时，我们开发了两种模式。

随机生成。这种情况下，我们直接从三个高斯分布出发，生成新的化学反应，即全新”未知的“ 「反应物，过渡态，生成物」.
条件生成。在化学反应中，我们通常知道部分信息。我们这里可以使用inpainting的条件生成，比如给定反应物和生成物，直接生成过渡态。这个过程中，我们每一步都将「反应物, 生成物」中加入固定噪声，并将它们与图神经网络模型预测得到的过渡态结合。一步步迭代，直到产生「反应物(“已知”)，过渡态(“未知”)，生成物(“已知”)」.

图四｜扩散生成模型生成化学反应的过程详解

代码

文本

当我们旋转化学反应的任何分子，这个反应会变化吗？

在生成单一分子或者一个大的多分子体系时，我们一般只需要满足整个体系的SE(3)变化后（比如整体旋转）模型的输出对应转变（下图第二行）。

然而在化学反应中，这个对称性还不够强。反应中，我们还要求关于每个单一对象的SE(3)对称性（下图第三行）。一个简单的例子为，当我们旋转了反应物，模型的输出应当只旋转这个反应物，其余部分应当保持不变，而不是把它当做一个”全新的“反应。

图五｜化学反应在不同的对称性（上）下经过不同转化（左）的行为

代码

文本

全体系SE(3)大失败 - 没有保持单个对象的对称性

代码

文本

考虑一个假想的“电解水反应”， H2O -> H2 + O.

注：这个反应过程纯属虚构，只是用来帮助我们理解反应中所需的对称性的一个简化模型。同时，我们省去了过渡态结构，只考虑「反应物，生成物」的对称性关系。加入过渡态并不不会改变结论。

代码

文本

[ ]

ns = [3, ] + [2, 1] # 反应物 3个原子 (H2O)，生成物 2个原子 (H2)，1个原子 (O自由基)

ntot = np.sum(ns)

mask = tensor([0, 0, 0, 1, 1, 1]) # 用于区分反应物和生成物

z = ["O", "H", "H"] + ["H", "H", "O"]

pos_react = tensor([

[0, 0, 0],

[1, 0, 0],

[0, 1, 0],

]).double() # 方便起见，我们这里把H-O-H的角度设为90度

pos_prod = tensor([

[0, 3, -0.4],

[0, 3, 0.4],

[0, -3, 0],

]) # 将H2和O自由基分开

pos = torch.cat(

[pos_react, pos_prod],

dim=0,

) # 拼接

h = torch.rand(ntot, in_hidden_channels)

代码

文本

[ ]

xyz = assemble_xyz(z, pos)

view = draw_in_3dmol(xyz, "xyz")

view

代码

文本

在化学反应使用整体SE(3)时，我们所有的原子需要全部连接在一起，保证反应物和生成物可以通过图神经网络「相互作用」。否则的话，我们改变反应物，生成物也不会改变，问题就更大了😉

代码

文本

[ ]

from oa_reactdiff.tests.model.utils import (

generate_full_eij,

get_cut_graph_mask,

)

edge_index = generate_full_eij(ntot)

代码

文本

正向预测

代码

文本

[ ]

_h, _pos, __ = model.forward(

h=h,

pos=remove_mean_batch(pos, mask),

edge_index=edge_index,

)

代码

文本

我们旋转一下反应物

代码

文本

[ ]

rot = o3.rand_matrix()

pos_react_rot = torch.matmul(pos_react, rot).double()

pos_rot = torch.cat(

[pos_react_rot, pos_prod],

dim=0,

) # 拼接旋转过后的H2O和未旋转的H2和O自由基

xyz = assemble_xyz(z, pos_rot)

view = draw_in_3dmol(xyz, "xyz")

view

代码

文本

再预测一次这个反应。注意这两个反应的化学实质是一样的

代码

文本

[ ]

_h_rot, _pos_rot, __ = model.forward(

h=h,

pos=remove_mean_batch(pos_rot, mask),

edge_index=edge_index,

)

代码

文本

但是我们scalar和vector上都丢失了等变性❌

代码

文本

[ ]

torch.max(

torch.abs(

_h - _h_rot

)

) # 旋转后的h应该不变

代码

文本

[ ]

_pos_rot_prime = torch.cat(

[

torch.matmul(_pos[:3], rot),

_pos[3:]

]

)

torch.max(

torch.abs(

_pos_rot_prime - _pos_rot

)

) # 旋转后的pos应该旋转

代码

文本

我们在化学反应中需要”对象感知“的SE(3)模型

为了处理这个难点，从而将生成模型的框架拓展到化学反应中，我们开发了一种将SE(3)模型扩展为”对象感知“（object-aware）的SE(3)模型的普适方法，OA-ReactDiff。这个方法可以用于各种图神经网络和Transformer。

这部分比技术细节比较多，感兴趣的小伙伴可以点击我们上方的文章🔗。在这个notebook中，我们直接展示改进后的效果。

代码

文本

创立一个”对象感知“的LEFTNet

代码

文本

[ ]

model_oa = LEFTNet(

num_layers=num_layers,

hidden_channels=hidden_channels,

in_hidden_channels=in_hidden_channels,

num_radial=num_radial,

object_aware=True, # 使用object-aware模型

)

代码

文本

我们将不同原子分配到他们所属的对象中，在这里是「反应物和生成物」。

代码

文本

[ ]

subgraph_mask = get_cut_graph_mask(edge_index, 3) # 0-2是反应物的原子数

代码

文本

再来重复一遍刚才的旋转测试

代码

文本

[ ]

_h, _pos, __ = model_oa.forward(

h=h,

pos=remove_mean_batch(pos, mask),

edge_index=edge_index,

subgraph_mask=subgraph_mask,

)

代码

文本

[ ]

rot = o3.rand_matrix()

pos_react_rot = torch.matmul(pos_react, rot).double()

pos_rot = torch.cat(

[pos_react_rot, pos_prod],

dim=0,

)

_h_rot, _pos_rot, __ = model_oa.forward(

h=h,

pos=remove_mean_batch(pos_rot, mask),

edge_index=edge_index,

subgraph_mask=subgraph_mask,

)

代码

文本

经过检验，”对象感知“的LEFTNet会保持我们想要的对称性✅

代码

文本

[ ]

torch.max(

torch.abs(

_h - _h_rot

)

) # 旋转后的h应该不变

代码

文本

[ ]

_pos_rot_prime = torch.cat(

[

torch.matmul(_pos[:3], rot),

_pos[3:]

]

)

torch.max(

torch.abs(

_pos_rot_prime - _pos_rot

)

) # 旋转后的pos应该旋转

代码

文本

从天到秒，OA-ReactDiff加速生成过渡态结构

在实验中，反应物和生成物存在的时间尺度相对较长，因此相对比较好表征（NMR, 质谱等等）。但是实验上表征过渡态结构很困难。因为过渡态存在时间短暂，其能量比反应物或产物更高，使得它们难以直接分离和研究。计算上，虽然大家发展了nudged elastic band等方法，但是这些方法都比较耗时，经常需要在一个24CPU的机器上跑一天，同时只有非常惨淡（30%左右）的成功率😭。

而我们文章中开发的OA-ReactDiff架构，可以将过渡态的搜索时间从几天缩短到6秒，大大加快我们探索化学反应机理的速度，成功率，和可研究化学反应网络的复杂度。我们准备了三个例子来展示OA-ReactDiff的应用场景🎬

代码

文本

正确的对称性移除一切繁杂预处理的必要

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

而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],

)

代码

文本

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

代码

文本

RMSD低于0.1A的两个分子的区别几乎肉眼不可区分。

图六｜不同RMSD下两个分子的结构差别展示（两个分子的C原子分别有蓝色和淡黄色表示）

储存生成的过渡态并且检查它与该反应的真实过渡态的RMSD，我们的RMSD只有0.02A。

代码

文本

[ ]

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"

)

代码

文本

意外惊喜 - 来自随机过程的馈赠

由于噪声的介入，扩散生成模型有一定的随机性，这导致我们及时给定了「反应物，生成物」的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"

)

代码

文本

对于同一个反应，我们产生了多少不同的TS？

为了节省大家时间，我们已经提前使用OA-ReactDiff生成多个关于该「反应物，生产物」的过渡态结构，这个反应对应于我们文章中的图三右边的例子。

代码

文本

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

代码

文本

[ ]

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)

代码

文本

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

注：我们这里已经对RMSD做了log10处理，「-2」对应于『0.01A』，「-1」对应于『0.1A』。

代码

文本

[ ]

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())

代码

文本

直观的总结一下不同的过渡态属于哪一个聚类（id: 『0』到『5』）

注： /opt/src/demo/example-3/opt_ts/26.ts.opt.xyz所在的聚类对应于我们给定{反应物，生成物}的真实过渡态。

代码

文本

[ ]

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"

)

代码

文本

我们生成的其余过渡态一文不值？

当然不是的。任意给定一个过渡态结构，都会对应一个唯一的化学反应。但是在计算当中，受限于个人的经验或者计算的方法，每个人所能探索的化学反应也是局限的

这种限制在研究未知的复杂反应时尤为致命。它会使我们忽略到一些潜在可能发生的反应，导致会反应机理的误判，进而影响催化材料设计的思路。赵祺源博士的一篇文章深入探索了这些现象，感兴趣的小伙伴可以深入了解。

在我们这个例子中，OA-ReactDiff除了找到了我们想要的那个反应的过渡态，同时由于随机过程的特点，还探索到了相关的五个“意外的”化学反应。这个特点可以弥补现有基于化学的直觉反应探索框架😄

代码

文本

比如我们生成的下面这个反应就是意料之外的全新反应。

代码

文本

[ ]

draw_reaction("/opt/src/demo/example-3/irc", prefix="cluster", idx=1)

代码

文本

给我一些原子，我可以生成关于它们的所有的反应

除了寻找已知反应的过渡态，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),

]

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,

)

代码

文本

强烈建议大家使用V100 GPU机器执行此步（运行需30秒）。不然的话，可能会花费10分钟左右时间。大家可以给自己泡一杯🍵或者做做俯卧撑。

代码

文本

[ ]

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

代码

文本

大家可以自己罗列一下，看看这里产生的分子是不是涵盖了所有的CONH四原子的稳定分子？😁

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

代码

文本

[ ]

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。

化学材料发现的一个巨大挑战在于其巨大（ $1 0^{60}$ ）的潜在设计空间。这个空间大到我们把NVIDIA所有的显卡买下来，耗尽全球的发电量也没有办法遍历。生成式AI绕过了筛选材料空间的困难，直接生成潜在有价值的分子和材料，带给了材料发现新的可能。

扩散生成模型在化学材料方面的拓展还处在初期萌芽阶段。如果大家对把OA-ReactDiff应用到自己问题的想法，但不确定是不是完全合适或者不知道如何下手，欢迎联系作者本人duanchenru@gmail.com👏🏻

代码

文本

AI4S

Deep Learning

Tutorial

PyTorch

DDPMs

Catalysis

DFT

AI4SDeep LearningTutorialPyTorchDDPMsCatalysisDFT

已赞52

本文被以下合集收录

2023机器学习及其在化学中的应用

昌珺涵

更新于 2024-09-12

22 篇158 人关注

更高更妙的分子模拟

量子御坂

更新于 2024-09-02

9 篇6 人关注

目录

简单介绍一下扩散生成模型 (DDPM)

已经玩转了图像，那试试生成分子如何？

生成分子时模型的SE(3)准变性不可或缺

建立一个LEFTNet模型

创建一个水分子作为测试体系

使用LEFTNet正向预测一个scalar和vector

随机旋转一下，再来一次

不管是scalar，还是vector， 结果都是等变的

化学的核心是反应，直接生成化学反应如何？

当我们旋转化学反应的任何分子，这个反应会变化吗？

全体系SE(3)大失败 - 没有保持单个对象的对称性

考虑一个假想的“电解水反应”， H2O -> H2 + O.

正向预测

我们旋转一下反应物

再预测一次这个反应。注意这两个反应的化学实质是一样的

但是我们scalar和vector上都丢失了等变性❌

我们在化学反应中需要”对象感知“的SE(3)模型

创立一个”对象感知“的LEFTNet

再来重复一遍刚才的旋转测试

经过检验，”对象感知“的LEFTNet会保持我们想要的对称性✅

从天到秒，OA-ReactDiff加速生成过渡态结构

正确的对称性移除一切繁杂预处理的必要

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

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

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

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

意外惊喜 - 来自随机过程的馈赠

画出我们的案例反应

对于同一个反应，我们产生了多少不同的TS？

我们生成的其余过渡态一文不值？

给我一些原子，我可以生成关于它们的所有的反应

随机生成反应

分析反应结果

找出独特不重复的稳定分子

独特的反应路径

展望: 让我们”生成“未来

不管是scalar，还是vector，结果都是等变的