| 129 | ) |
| 130 | |
| 131 | def process(self): |
| 132 | npz_path = f"{self.raw_dir}/qm9_eV.npz" |
| 133 | data_dict = np.load(npz_path, allow_pickle=True) |
| 134 | # data_dict['N'] contains the number of atoms in each molecule. |
| 135 | # Atomic properties (Z and R) of all molecules are concatenated as single tensors, |
| 136 | # so you need this value to select the correct atoms for each molecule. |
| 137 | self.N = data_dict["N"] |
| 138 | self.R = data_dict["R"] |
| 139 | self.Z = data_dict["Z"] |
| 140 | self.label = np.stack( |
| 141 | [data_dict[key] for key in self.label_keys], axis=1 |
| 142 | ) |
| 143 | self.N_cumsum = np.concatenate([[0], np.cumsum(self.N)]) |
| 144 | |
| 145 | def download(self): |
| 146 | file_path = f"{self.raw_dir}/qm9_eV.npz" |