xnn.common.data.hub.rmd17.RMD17Builder#

class xnn.common.data.hub.rmd17.RMD17Builder[source]#

Bases: DatasetBuilder

Builder for the revised MD17 dataset.

See the module docstring for the dataset description and citation. Only the requested molecule’s .npz (and, for the standard splits, the two index CSVs) are downloaded, so a single molecule costs far less than the 1 GB bundle.

load(*, split=None, cache_dir, molecule=None, fold=1, units='eV', n_train=None, n_test=None, quiet=False)[source]#

Download and preprocess one rMD17 molecule.

Parameters:

split (str or None) – None returns {"train": ..., "test": ...} for the chosen fold; "train" / "test" returns that split; "all" returns every conformation (no split filtering), for custom splits.
cache_dir (pathlib.Path) – Base cache directory; files are stored under cache_dir/"rmd17".
molecule (str) – Which molecule to load (required). One of aspirin, azobenzene, benzene, ethanol, malonaldehyde, naphthalene, paracetamol, salicylic, toluene, uracil.
fold (int, optional) – Which of the five official splits (1-5) to use. Defaults to 1. Ignored when split="all".
units (str, optional) – "eV" (default) converts energies to eV and forces to eV/A; "kcal/mol" keeps the upstream units.
n_train (int, optional) – Truncate the train/test split to the first n structures. The official splits hold 1000 each; the authors warn against training on more than 1000 samples.
n_test (int, optional) – Truncate the train/test split to the first n structures. The official splits hold 1000 each; the authors warn against training on more than 1000 samples.
quiet (bool, optional) – Suppress download progress output. Defaults to False.

Returns:

dict of {str – Structure dicts with keys pos (N, 3), atomic_numbers (N,), energy (scalar) and forces (N, 3). rMD17 is molecular, so no cell / pbc is set.

Return type:

list of dict} or list of dict

Raises:

ValueError – If molecule is missing/unknown, fold is not in 1-5, or units / split is unrecognized.