dmlc--dgl
Hengrui Zhang
3317522622
* Add implementation of unsupervised model * [Doc] Update Implementor's information * [doc] add index of infograph * [Feature] QM9_v2 Dataset Support * fix a typo * move qm9dataset from data to examples * Update qm9_v2.py * Infograph -> InfoGraph * add implementation and results of semi-supervised model * Update README.md * Update README.md * fix a typo * Remove the duplicated links. * fix some typos * fix typos * update model.py * update collate fn * remove unused functions * Update model.py * add device option * Update evaluate_embedding.py * Update evaluate_embedding.py * Update unsupervised.py * Fix typos * fix bugs * Update README.md Co-authored-by: Mufei Li <mufeili1996@gmail.com>
214 行
12 KiB
Python
214 行
12 KiB
Python
import numpy as np
|
|
import os
|
|
from tqdm import tqdm
|
|
|
|
import torch as th
|
|
|
|
import dgl
|
|
from dgl.data.dgl_dataset import DGLDataset
|
|
from dgl.data.utils import download, load_graphs, _get_dgl_url, extract_archive
|
|
|
|
class QM9DatasetV2(DGLDataset):
|
|
r"""QM9 dataset for graph property prediction (regression)
|
|
This dataset consists of 130,831 molecules with 19 regression targets.
|
|
Node means atom and edge means bond.
|
|
|
|
Reference: `"MoleculeNet: A Benchmark for Molecular Machine Learning" <https://arxiv.org/abs/1703.00564>`_
|
|
Atom features come from `"Neural Message Passing for Quantum Chemistry" <https://arxiv.org/abs/1704.01212>`_
|
|
|
|
Statistics:
|
|
|
|
- Number of graphs: 130,831
|
|
- Number of regression targets: 19
|
|
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| Keys | Property | Description | Unit |
|
|
+========+==================================+===================================================================================+=============================================+
|
|
| mu | :math:`\mu` | Dipole moment | :math:`\textrm{D}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| alpha | :math:`\alpha` | Isotropic polarizability | :math:`{a_0}^3` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| homo | :math:`\epsilon_{\textrm{HOMO}}` | Highest occupied molecular orbital energy | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| lumo | :math:`\epsilon_{\textrm{LUMO}}` | Lowest unoccupied molecular orbital energy | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| gap | :math:`\Delta \epsilon` | Gap between :math:`\epsilon_{\textrm{HOMO}}` and :math:`\epsilon_{\textrm{LUMO}}` | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| r2 | :math:`\langle R^2 \rangle` | Electronic spatial extent | :math:`{a_0}^2` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| zpve | :math:`\textrm{ZPVE}` | Zero point vibrational energy | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| U0 | :math:`U_0` | Internal energy at 0K | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| U | :math:`U` | Internal energy at 298.15K | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| H | :math:`H` | Enthalpy at 298.15K | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| G | :math:`G` | Free energy at 298.15K | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| Cv | :math:`c_{\textrm{v}}` | Heat capavity at 298.15K | :math:`\frac{\textrm{cal}}{\textrm{mol K}}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| U0_atom| :math:`U_0^{\textrm{ATOM}}` | Atomization energy at 0K | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| U_atom | :math:`U^{\textrm{ATOM}}` | Atomization energy at 298.15K | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| H_atom | :math:`H^{\textrm{ATOM}}` | Atomization enthalpy at 298.15K | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| G_atom | :math:`G^{\textrm{ATOM}}` | Atomization free energy at 298.15K | :math:`\textrm{eV}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| A | :math:`A` | Rotational constant | :math:`\textrm{GHz}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| B | :math:`B` | Rotational constant | :math:`\textrm{GHz}` |
|
|
+--------+----------------------------------+-----------------------------------------------------------------------------------+---------------------------------------------+
|
|
| c | :math:`C` | Rotational constant | :math:`\textrm{GHz}` |
|
|
+--------+----------------------------------+----------------------------------------
|
|
Parameters
|
|
----------
|
|
label_keys: list
|
|
Names of the regression property, which should be a subset of the keys in the table above.
|
|
If not provided, will load all the labels.
|
|
raw_dir : str
|
|
Raw file directory to download/contains the input data directory.
|
|
Default: ~/.dgl/
|
|
force_reload : bool
|
|
Whether to reload the dataset. Default: False
|
|
verbose: bool
|
|
Whether to print out progress information. Default: True.
|
|
|
|
Attributes
|
|
----------
|
|
num_labels : int
|
|
Number of labels for each graph, i.e. number of prediction tasks
|
|
|
|
Raises
|
|
------
|
|
UserWarning
|
|
If the raw data is changed in the remote server by the author.
|
|
|
|
Examples
|
|
--------
|
|
>>> data = QM9DatasetV2(label_keys=['mu', 'alpha'])
|
|
>>> data.num_labels
|
|
|
|
|
|
>>> # make each graph dense
|
|
>>> data.to_dense()
|
|
|
|
>>> # iterate over the dataset
|
|
>>> for graph, labels in data:
|
|
... print(graph) # get information of each graph
|
|
... print(labels) # get labels of the corresponding graph
|
|
... # your code here...
|
|
>>>
|
|
|
|
"""
|
|
def __init__(self,
|
|
label_keys = None,
|
|
raw_dir=None,
|
|
force_reload=False,
|
|
verbose=True):
|
|
|
|
self.label_keys = label_keys
|
|
self._url = _get_dgl_url('dataset/qm9_ver2.zip')
|
|
super(QM9DatasetV2, self).__init__(name='qm9_v2',
|
|
url=self._url,
|
|
raw_dir=raw_dir,
|
|
force_reload=force_reload,
|
|
verbose=verbose)
|
|
|
|
def process(self):
|
|
print('begin loading dataset')
|
|
graphs, label_dict = load_graphs(os.path.join(self.raw_dir, 'qm9_v2.bin'))
|
|
self.graphs = graphs
|
|
if self.label_keys == None:
|
|
self.labels = np.stack([label_dict[key] for key in label_dict.keys()], axis=1)
|
|
else:
|
|
self.labels = np.stack([label_dict[key] for key in self.label_keys], axis=1)
|
|
|
|
def to_dense(self):
|
|
r""" Transfrom each graph to a dense graph and add additional edge attribute(distance between two atoms)
|
|
|
|
Note: This operation will deprecate graph.ndata['pos']
|
|
|
|
"""
|
|
n_graph = self.labels.shape[0]
|
|
for id in tqdm(range(n_graph), desc = 'processing graphs'):
|
|
graph = self.graphs[id]
|
|
n_nodes = graph.num_nodes()
|
|
row = th.arange(n_nodes, dtype = th.long)
|
|
col = th.arange(n_nodes, dtype = th.long)
|
|
|
|
row = row.view(-1,1).repeat(1, n_nodes).view(-1)
|
|
col = col.repeat(n_nodes)
|
|
|
|
src = graph.edges()[0]
|
|
dst = graph.edges()[1]
|
|
|
|
idx = src * n_nodes + dst
|
|
size = list(graph.edata['edge_attr'].size())
|
|
size[0] = n_nodes * n_nodes
|
|
edge_attr = graph.edata['edge_attr'].new_zeros(size)
|
|
|
|
edge_attr[idx] = graph.edata['edge_attr']
|
|
|
|
pos = graph.ndata['pos']
|
|
dist = th.norm(pos[col] - pos[row], p=2, dim=-1).view(-1, 1)
|
|
|
|
new_edge_attr = th.cat([edge_attr, dist.type_as(edge_attr)], dim = -1)
|
|
|
|
new_graph = dgl.graph((row,col))
|
|
|
|
new_graph.ndata['attr'] = graph.ndata['attr']
|
|
new_graph.edata['edge_attr'] = new_edge_attr
|
|
new_graph = new_graph.remove_self_loop()
|
|
|
|
self.graphs[id] = new_graph
|
|
|
|
def download(self):
|
|
file_path = f'{self.raw_dir}/qm9_v2.zip'
|
|
if not os.path.exists(file_path):
|
|
download(self._url, path=file_path)
|
|
extract_archive(file_path, self.raw_dir, overwrite = True)
|
|
|
|
@property
|
|
def num_labels(self):
|
|
r"""
|
|
Returns
|
|
--------
|
|
int
|
|
Number of labels for each graph, i.e. number of prediction tasks.
|
|
"""
|
|
return self.labels.shape[1]
|
|
|
|
def __getitem__(self, idx):
|
|
r""" Get graph and label by index
|
|
|
|
Parameters
|
|
----------
|
|
idx : int
|
|
Item index
|
|
|
|
Returns
|
|
-------
|
|
dgl.DGLGraph
|
|
The graph contains:
|
|
|
|
- ``ndata['pos']``: the coordinates of each atom
|
|
- ``ndata['attr']``: the atomic attributes
|
|
- ``edata['edge_attr']``: the bond attributes
|
|
|
|
Tensor
|
|
Property values of molecular graphs
|
|
"""
|
|
|
|
return self.graphs[idx], self.labels[idx]
|
|
|
|
def __len__(self):
|
|
r"""Number of graphs in the dataset.
|
|
|
|
Return
|
|
-------
|
|
int
|
|
"""
|
|
return self.labels.shape[0]
|