import os
import os.path as osp
from typing import Callable, List, Optional, Union
import numpy as np
import torch
from torch_geometric.data import (
Data,
InMemoryDataset,
download_url,
extract_tar,
extract_zip,
)
[docs]class MD17(InMemoryDataset):
r"""A variety of ab-initio molecular dynamics trajectories from the authors
of `sGDML <http://quantum-machine.org/gdml>`_.
This class provides access to the original MD17 datasets, their revised
versions, and the CCSD(T) trajectories.
For every trajectory, the dataset contains the Cartesian positions of atoms
(in Angstrom), their atomic numbers, as well as the total energy
(in kcal/mol) and forces (kcal/mol/Angstrom) on each atom.
The latter two are the regression targets for this collection.
.. note::
Data objects contain no edge indices as these are most commonly
constructed via the :obj:`torch_geometric.transforms.RadiusGraph`
transform, with its cut-off being a hyperparameter.
The `original MD17 dataset <https://arxiv.org/abs/1611.04678>`_ contains
ten molecule trajectories.
This version of the dataset was found to suffer from high numerical noise.
The `revised MD17 dataset <https://arxiv.org/abs/2007.09593>`_ contains the
same molecules, but the energies and forces were recalculated at the
PBE/def2-SVP level of theory using very tight SCF convergence and very
dense DFT integration grid.
The third version of the dataset contains fewer molecules, computed at the
CCSD(T) level of theory.
The benzene molecule at the DFT FHI-aims level of theory was
`released separately <https://arxiv.org/abs/1802.09238>`_.
Check the table below for detailed information on the molecule, level of
theory and number of data points contained in each dataset.
Which trajectory is loaded is determined by the :attr:`name` argument.
For the coupled cluster trajectories, the dataset comes with pre-defined
training and testing splits which are loaded separately via the
:attr:`train` argument.
+--------------------+--------------------+-------------------------------+-----------+
| Molecule | Level of Theory | Name | #Examples |
+====================+====================+===============================+===========+
| Benzene | DFT | :obj:`benzene` | 627,983 |
+--------------------+--------------------+-------------------------------+-----------+
| Uracil | DFT | :obj:`uracil` | 133,770 |
+--------------------+--------------------+-------------------------------+-----------+
| Naphthalene | DFT | :obj:`napthalene` | 326,250 |
+--------------------+--------------------+-------------------------------+-----------+
| Aspirin | DFT | :obj:`aspirin` | 211,762 |
+--------------------+--------------------+-------------------------------+-----------+
| Salicylic acid | DFT | :obj:`salicylic acid` | 320,231 |
+--------------------+--------------------+-------------------------------+-----------+
| Malonaldehyde | DFT | :obj:`malonaldehyde` | 993,237 |
+--------------------+--------------------+-------------------------------+-----------+
| Ethanol | DFT | :obj:`ethanol` | 555,092 |
+--------------------+--------------------+-------------------------------+-----------+
| Toluene | DFT | :obj:`toluene` | 442,790 |
+--------------------+--------------------+-------------------------------+-----------+
| Paracetamol | DFT | :obj:`paracetamol` | 106,490 |
+--------------------+--------------------+-------------------------------+-----------+
| Azobenzene | DFT | :obj:`azobenzene` | 99,999 |
+--------------------+--------------------+-------------------------------+-----------+
| Benzene (R) | DFT (PBE/def2-SVP) | :obj:`revised benzene` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Uracil (R) | DFT (PBE/def2-SVP) | :obj:`revised uracil` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Naphthalene (R) | DFT (PBE/def2-SVP) | :obj:`revised napthalene` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Aspirin (R) | DFT (PBE/def2-SVP) | :obj:`revised aspirin` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Salicylic acid (R) | DFT (PBE/def2-SVP) | :obj:`revised salicylic acid` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Malonaldehyde (R) | DFT (PBE/def2-SVP) | :obj:`revised malonaldehyde` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Ethanol (R) | DFT (PBE/def2-SVP) | :obj:`revised ethanol` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Toluene (R) | DFT (PBE/def2-SVP) | :obj:`revised toluene` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Paracetamol (R) | DFT (PBE/def2-SVP) | :obj:`revised paracetamol` | 100,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Azobenzene (R) | DFT (PBE/def2-SVP) | :obj:`revised azobenzene` | 99,988 |
+--------------------+--------------------+-------------------------------+-----------+
| Benzene | CCSD(T) | :obj:`benzene CCSD(T)` | 1,500 |
+--------------------+--------------------+-------------------------------+-----------+
| Aspirin | CCSD | :obj:`aspirin CCSD` | 1,500 |
+--------------------+--------------------+-------------------------------+-----------+
| Malonaldehyde | CCSD(T) | :obj:`malonaldehyde CCSD(T)` | 1,500 |
+--------------------+--------------------+-------------------------------+-----------+
| Ethanol | CCSD(T) | :obj:`ethanol CCSD(T)` | 2,000 |
+--------------------+--------------------+-------------------------------+-----------+
| Toluene | CCSD(T) | :obj:`toluene CCSD(T)` | 1,501 |
+--------------------+--------------------+-------------------------------+-----------+
| Benzene | DFT FHI-aims | :obj:`benzene FHI-aims` | 49,863 |
+--------------------+--------------------+-------------------------------+-----------+
.. warning::
It is advised to not train a model on more than 1,000 samples from the
original or revised MD17 dataset.
Args:
root (str): Root directory where the dataset should be saved.
name (str): Keyword of the trajectory that should be loaded.
train (bool, optional): Determines whether the train or test split
gets loaded for the coupled cluster trajectories.
(default: :obj:`None`)
transform (callable, optional): A function/transform that takes in an
:obj:`torch_geometric.data.Data` object and returns a transformed
version. The data object will be transformed before every access.
(default: :obj:`None`)
pre_transform (callable, optional): A function/transform that takes in
an :obj:`torch_geometric.data.Data` object and returns a
transformed version. The data object will be transformed before
being saved to disk. (default: :obj:`None`)
pre_filter (callable, optional): A function that takes in an
:obj:`torch_geometric.data.Data` object and returns a boolean
value, indicating whether the data object should be included in the
final dataset. (default: :obj:`None`)
force_reload (bool, optional): Whether to re-process the dataset.
(default: :obj:`False`)
**STATS:**
.. list-table::
:widths: 20 10 10 10 10 10
:header-rows: 1
* - Name
- #graphs
- #nodes
- #edges
- #features
- #tasks
* - Benzene
- 627,983
- 12
- 0
- 1
- 2
* - Uracil
- 133,770
- 12
- 0
- 1
- 2
* - Naphthalene
- 326,250
- 10
- 0
- 1
- 2
* - Aspirin
- 211,762
- 21
- 0
- 1
- 2
* - Salicylic acid
- 320,231
- 16
- 0
- 1
- 2
* - Malonaldehyde
- 993,237
- 9
- 0
- 1
- 2
* - Ethanol
- 555,092
- 9
- 0
- 1
- 2
* - Toluene
- 442,790
- 15
- 0
- 1
- 2
* - Paracetamol
- 106,490
- 20
- 0
- 1
- 2
* - Azobenzene
- 99,999
- 24
- 0
- 1
- 2
* - Benzene (R)
- 100,000
- 12
- 0
- 1
- 2
* - Uracil (R)
- 100,000
- 12
- 0
- 1
- 2
* - Naphthalene (R)
- 100,000
- 10
- 0
- 1
- 2
* - Aspirin (R)
- 100,000
- 21
- 0
- 1
- 2
* - Salicylic acid (R)
- 100,000
- 16
- 0
- 1
- 2
* - Malonaldehyde (R)
- 100,000
- 9
- 0
- 1
- 2
* - Ethanol (R)
- 100,000
- 9
- 0
- 1
- 2
* - Toluene (R)
- 100,000
- 15
- 0
- 1
- 2
* - Paracetamol (R)
- 100,000
- 20
- 0
- 1
- 2
* - Azobenzene (R)
- 99,988
- 24
- 0
- 1
- 2
* - Benzene CCSD-T
- 1,500
- 12
- 0
- 1
- 2
* - Aspirin CCSD-T
- 1,500
- 21
- 0
- 1
- 2
* - Malonaldehyde CCSD-T
- 1,500
- 9
- 0
- 1
- 2
* - Ethanol CCSD-T
- 2000
- 9
- 0
- 1
- 2
* - Toluene CCSD-T
- 1,501
- 15
- 0
- 1
- 2
* - Benzene FHI-aims
- 49,863
- 12
- 0
- 1
- 2
""" # noqa: E501
gdml_url = 'http://quantum-machine.org/gdml/data/npz'
revised_url = ('https://archive.materialscloud.org/record/'
'file?filename=rmd17.tar.bz2&record_id=466')
file_names = {
'benzene': 'md17_benzene2017.npz',
'uracil': 'md17_uracil.npz',
'naphtalene': 'md17_naphthalene.npz',
'aspirin': 'md17_aspirin.npz',
'salicylic acid': 'md17_salicylic.npz',
'malonaldehyde': 'md17_malonaldehyde.npz',
'ethanol': 'md17_ethanol.npz',
'toluene': 'md17_toluene.npz',
'paracetamol': 'paracetamol_dft.npz',
'azobenzene': 'azobenzene_dft.npz',
'revised benzene': 'rmd17_benzene.npz',
'revised uracil': 'rmd17_uracil.npz',
'revised naphthalene': 'rmd17_naphthalene.npz',
'revised aspirin': 'rmd17_aspirin.npz',
'revised salicylic acid': 'rmd17_salicylic.npz',
'revised malonaldehyde': 'rmd17_malonaldehyde.npz',
'revised ethanol': 'rmd17_ethanol.npz',
'revised toluene': 'rmd17_toluene.npz',
'revised paracetamol': 'rmd17_paracetamol.npz',
'revised azobenzene': 'rmd17_azobenzene.npz',
'benzene CCSD(T)': 'benzene_ccsd_t.zip',
'aspirin CCSD': 'aspirin_ccsd.zip',
'malonaldehyde CCSD(T)': 'malonaldehyde_ccsd_t.zip',
'ethanol CCSD(T)': 'ethanol_ccsd_t.zip',
'toluene CCSD(T)': 'toluene_ccsd_t.zip',
'benzene FHI-aims': 'benzene2018_dft.npz',
}
def __init__(
self,
root: str,
name: str,
train: Optional[bool] = None,
transform: Optional[Callable] = None,
pre_transform: Optional[Callable] = None,
pre_filter: Optional[Callable] = None,
force_reload: bool = False,
) -> None:
if name not in self.file_names:
raise ValueError(f"Unknown dataset name '{name}'")
self.name = name
self.revised = 'revised' in name
self.ccsd = 'CCSD' in self.name
super().__init__(root, transform, pre_transform, pre_filter,
force_reload=force_reload)
if len(self.processed_file_names) == 1 and train is not None:
raise ValueError(
f"'{self.name}' dataset does not provide pre-defined splits "
f"but the 'train' argument is set to '{train}'")
elif len(self.processed_file_names) == 2 and train is None:
raise ValueError(
f"'{self.name}' dataset does provide pre-defined splits but "
f"the 'train' argument was not specified")
idx = 0 if train is None or train else 1
self.load(self.processed_paths[idx])
def mean(self) -> float:
assert isinstance(self._data, Data)
return float(self._data.energy.mean())
@property
def raw_dir(self) -> str:
if self.revised:
return osp.join(self.root, 'raw')
return osp.join(self.root, self.name, 'raw')
@property
def processed_dir(self) -> str:
return osp.join(self.root, self.name, 'processed')
@property
def raw_file_names(self) -> Union[str, List[str]]:
name = self.file_names[self.name]
if self.revised:
return osp.join('rmd17', 'npz_data', name)
elif self.ccsd:
return [name[:-4] + '-train.npz', name[:-4] + '-test.npz']
return name
@property
def processed_file_names(self) -> List[str]:
if self.ccsd:
return ['train.pt', 'test.pt']
else:
return ['data.pt']
def download(self) -> None:
if self.revised:
path = download_url(self.revised_url, self.raw_dir)
extract_tar(path, self.raw_dir, mode='r:bz2')
os.unlink(path)
else:
url = f'{self.gdml_url}/{self.file_names[self.name]}'
path = download_url(url, self.raw_dir)
if self.ccsd:
extract_zip(path, self.raw_dir)
os.unlink(path)
def process(self) -> None:
it = zip(self.raw_paths, self.processed_paths)
for raw_path, processed_path in it:
raw_data = np.load(raw_path)
if self.revised:
z = torch.from_numpy(raw_data['nuclear_charges']).long()
pos = torch.from_numpy(raw_data['coords']).float()
energy = torch.from_numpy(raw_data['energies']).float()
force = torch.from_numpy(raw_data['forces']).float()
else:
z = torch.from_numpy(raw_data['z']).long()
pos = torch.from_numpy(raw_data['R']).float()
energy = torch.from_numpy(raw_data['E']).float()
force = torch.from_numpy(raw_data['F']).float()
data_list = []
for i in range(pos.size(0)):
data = Data(z=z, pos=pos[i], energy=energy[i], force=force[i])
if self.pre_filter is not None and not self.pre_filter(data):
continue
if self.pre_transform is not None:
data = self.pre_transform(data)
data_list.append(data)
self.save(data_list, processed_path)
def __repr__(self) -> str:
return f"{self.__class__.__name__}({len(self)}, name='{self.name}')"