Source code for torch_geometric.datasets.md17

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}')"