'''
by wufeil
DeepChem-使用图卷积神经网络对化学分子建模
在深度学习中，图卷积神经网络特别适合处理化学分子数据，也经常被用于处理分子数据。
这是因为化学分子本身就是一个图数据。但是如何将分子转化为图神经网络可以使用的数据格式呢？
接下来就做简单的介绍。
首先，标准的卷积网络，最简单最传统的就是CNN，convolutional neural network，是处理图片最常用的神经网络模型。
CNN的输入是像素矩阵，每一个像素都是一个含有三个元素的向量，代表红黄蓝的数值。所以CNN的输入形状是；(batch_size, lenght, widht, color(3) )
数据会逐步通过一系列的卷积层。每一层将来自像素及其相邻像素的数据组合在一起，以生成该像素的新数据向量。
较早的层检测到小规模的局部模式，而较后的层则检测到更大，更抽象的模式。
卷积层通常与执行某些操作（例如在局部区域上的最大值或最小值）的池化层交替。
图卷积神经网络也是相似的。
但是图卷积神经网络是操作在图上面的。操作的对象是其周围的节点和边。
所谓节点就是化学分子中的原子，边则是源自之间的连接的化学键。
所以图神经网络输入的是图中节点的特征矩阵，在卷积和池化的过程中，不断的聚合周围节点和边的信息，生成新的节点新的边特征。
（如果是图层面的分类模型，有可能节点和边的数量都会发生变化）。
关于图神经网络的知识，有一本书《深入浅出图神经网络：GNN原理解析》写的很不错，有精力的同学可以学习参考。
'''
'''
接下来的关键就是如何使用图神经网络处理为我们的化学分子建模呢？
我们这里要使用到的数据集是MolNet里面的Tox21数据集。
之前的博客有写道：这些数据集可以很方便的额调用，只需要featurizer=‘GraphConv’就可以加载适合图神经网络的数据集
所以，在这里如何将化学分子处理成适合图神经网络的数据集就不做重点介绍,大家可以看我之前的博客 [《DeepChem的内置数据集及使用方法》](https://editor.csdn.net/md/?articleId=110817827)。
这里仅介绍图神经网络的使用方法
'''
'''
导入数据
'''
import deepchem as dc
tasks, datasets, transformers = dc.molnet.load_tox21(featurizer='GraphConv', splitter='random')
train_set, valid_set, test_set = datasets
'''
使用deepchem内置的图神经网络模型对数据集进行训练
（该模型很简单，并不一定适合真实的数据集！！！真实的数据集需要特殊的数据处理和特殊的模型结构，特殊的损失函数的！！！
这些就是做AI的人在干的事情）
如果真的要完成某些任务，那就要自己设计模型，编写自定义的模型，那么才能有效果。
那么要在Deepchem中使用自定义模型，但是正如前一篇博客[《DeepChem-使用自定义Keras\Torch深度学习模型示例》](https://editor.csdn.net/md/?articleId=111171039)所说的，不建议。
最好是直接使用pytorch，tensorflow来写，来训练，预测。
因为deepchem的API接口不友好，部分功能还缺失.另外，再运行的时候，有好多的报警，警告，
感觉deepchem调用深度学习网络进行训练的时候不稳定，很有可能崩溃。
更为重要的是，torch, tensorflow现在已经很成熟，很简单了。python有个一两年基础，基本上可以掌握的。
'''
model = dc.models.GraphConvModel(n_tasks=len(tasks), mode='classification')
model.fit(train_set, nb_epoch=50)
'''
评估模型
'''
metric = dc.metrics.Metric(dc.metrics.roc_auc_score)
print('train roc suc score:', model.evaluate(train_set, [metric], transformers))
print('valid roc suc score:', model.evaluate(valid_set, [metric], transformers))
print('test roc suc score:', model.evaluate(test_set, [metric], transformers))
'''
要在Deepchem中使用自定义模型，但是正如前一篇博客《DeepChem-使用自定义Keras\Torch深度学习模型示例》所说的，不建议。
但是在这里，还是介绍一些如何使用自定义的图神经网络模型。
可以结合之前的博客《DeepChem-使用自定义Keras\Torch深度学习模型示例》
'''
'''
首先来自定义一个图神经网络模型，使用tensorflow.keras
继承tf.keras.Model类，定义__init__()和call()函数即可
'''
from deepchem.models.layers import GraphConv, GraphPool, GraphGather
import tensorflow as tf
import tensorflow.keras.layers as layers
batch_size = 100
n_tasks=len(tasks)
class MyGraphConvModel(tf.keras.Model):
def __init__(self):
super(MyGraphConvModel, self).__init__()
self.gc1 = GraphConv(128, activation_fn=tf.nn.tanh)
self.batch_norm1 = layers.BatchNormalization()
self.gp1 = GraphPool()
self.gc2 = GraphConv(128, activation_fn=tf.nn.tanh)
self.batch_norm2 = layers.BatchNormalization()
self.gp2 = GraphPool()
self.dense1 = layers.Dense(256, activation=tf.nn.tanh)
self.batch_norm3 = layers.BatchNormalization()
self.readout = GraphGather(batch_size=batch_size, activation_fn=tf.nn.tanh)
self.dense2 = layers.Dense(n_tasks*2)
self.logits = layers.Reshape((n_tasks, 2))
self.softmax = layers.Softmax()
def call(self, inputs):
gc1_output = self.gc1(inputs)
batch_norm1_output = self.batch_norm1(gc1_output)
gp1_output = self.gp1([batch_norm1_output] + inputs[1:])
gc2_output = self.gc2([gp1_output] + inputs[1:])
batch_norm2_output = self.batch_norm1(gc2_output)
gp2_output = self.gp2([batch_norm2_output] + inputs[1:])
dense1_output = self.dense1(gp2_output)
batch_norm3_output = self.batch_norm3(dense1_output)
readout_output = self.readout([batch_norm3_output] + inputs[1:])
logits_output = self.logits(self.dense2(readout_output))
return self.softmax(logits_output)
'''
接下来就是使用自定义的图神经网络，实例化一个deepchem的模型。
需要传递自定义的图神经网络和损失函数
'''
model = dc.models.KerasModel(MyGraphConvModel(),loss=dc.models.losses.CategoricalCrossEntropy())
'''
train_set虽然已经使用了featurizer='GraphConv'，但是还不能直接作为图神经网络的输入。因为图神经网络的输入需要的是numpy数组。
而featurizer='GraphConv'返回的对象是ConvMol，因此还需要将ConvMol转换为numpy数组。
（在ConvMol中，并不是直接包含图神经网络可以使用的邻接矩阵，度矩阵等，而是：
atom_features (np.ndarray) – Has shape (n_atoms, n_feat)，
adj_list (list) – List of length n_atoms, with neighor indices of each atom.
max_deg (int, optional) – Maximum degree of any atom.
min_deg (int, optional) – Minimum degree of any atom.
详见：https://deepchem.readthedocs.io/en/latest/api_reference/dataclasses.html#graph-convolutions）
详情请见：文档中关于分子特征化部分https://deepchem.readthedocs.io/en/latest/api_reference/featurizers.html#
因此，先创建一个数据加载转换函数data_generator，其输入是dataset
'''
from deepchem.metrics import to_one_hot
from deepchem.feat.mol_graphs import ConvMol
import numpy as np
import tensorflow as tf
def data_generator(dataset, epochs=1):
for ind, (X_b, y_b, w_b, ids_b) in enumerate(dataset.iterbatches(batch_size, epochs,
deterministic=False, pad_batches=True)):
multiConvMol = ConvMol.agglomerate_mols(X_b)
inputs = [multiConvMol.get_atom_features(), multiConvMol.deg_slice, np.array(multiConvMol.membership)]
for i in range(1, len(multiConvMol.get_deg_adjacency_lists())):
inputs.append(multiConvMol.get_deg_adjacency_lists()[i])
labels = [to_one_hot(y_b.flatten(), 2).reshape(-1, n_tasks, 2)]
weights = [w_b]
yield (inputs, labels, weights)
'''
使用data_generator转换数据以后直接输入到model
然后评估验证集和测试集
注意这里使用的是model.fit_generator(), 而不是model.fit()
'''
model.fit_generator(data_generator(train_set, epochs=50))
print('Training set score:', model.evaluate_generator(data_generator(train_set), [metric], transformers))
print('Test set score:', model.evaluate_generator(data_generator(test_set), [metric], transformers))
'''
最后总结一下：
1.deepchem里面有简单的图神经网络模型，dc.models.GraphConvModel()，其输入的是：ConvMol数据类型，因此能与deepchem内置的数据集
和featurizer='GraphConv'无缝对接；
2.deepchem也支持自定义的图神经网络，仅需要继承tf.keras.Model/tf.torch.Model，定义__init__()和call()函数。
但是自定义的图神经网络需要的是numpy的数据类型，因此需要定义一个data_generator，将ConvMol数据类型改变成numpy数据类型，
返回一个数据generator,因此要使用model.fit_generator()
显然，如果要真的使用这个图神经网络，那么真的需要有一些图神经网络的知识，知道什么是度，边，邻接矩阵等基础概念。
当然也要知道DeepChem中关于图神经网络的生成方法。
'''
