xnn.common.data.hub.ani1.ANI1Builder#

class xnn.common.data.hub.ani1.ANI1Builder[source]#

Bases: DatasetBuilder

Builder for the ANI-1 data set (20M DFT energies, GDB-11 subsets).

See the module docstring for the dataset description and citation. The one 4.8 GB archive is downloaded and extracted once; the requested heavy-atom subsets are then parsed from their HDF5 files into xnn structure dicts.

load(*, split=None, cache_dir, heavy_atoms=None, units='eV', max_molecules=None, max_conformations=None, seed=1234, quiet=False)[source]#

Download and preprocess the ANI-1 data set.

Parameters:

split (str or None) – None returns {"all": ...}; "train" / "val" / "test" returns the paper’s per-molecule 80/10/10 partition (fixed seed, disjoint splits).
cache_dir (pathlib.Path) – Base cache directory; files live under cache_dir/"ani1".
heavy_atoms (int or sequence of int, optional) – Which heavy-atom subset(s) ani_gdb_s0X.h5 to load (X in 1-8). Defaults to all eight. Fewer subsets = far less data.
units (str, optional) – "eV" (default) converts energies to eV; "hartree" keeps the raw upstream values.
max_molecules (int, optional) – Cap the number of molecules read per subset (useful for demos).
max_conformations (int, optional) – Cap the number of conformations kept per molecule.
seed (int, optional) – Seed for the reproducible 80/10/10 split. Defaults to 1234.
quiet (bool, optional) – Suppress progress output. Defaults to False.

Returns:

dict of {str – Structure dicts with keys pos (N, 3), atomic_numbers (N,), energy (scalar) and smiles (str). ANI-1 is molecular and force-free, so no cell / pbc / forces.

Return type:

list of dict} or list of dict

Raises:

ValueError – If units or split is unrecognized, or heavy_atoms is out of the 1-8 range.
ImportError – If h5py is not installed.