xnn.common.data.hub.ani1x.ANI1xBuilder#

class xnn.common.data.hub.ani1x.ANI1xBuilder[source]#

Bases: DatasetBuilder

Builder for the ANI-1x data set (active-learning DFT energies & forces).

See the module docstring for the dataset description and citation. The one ~5.6 GB ani1x-release.h5 file is downloaded once; its molecule groups are then parsed into xnn structure dicts, keeping only the conformations for which the requested level has non-NaN values.

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

Download and preprocess the ANI-1x data set.

Parameters:

split (str or None) – None returns {"all": ...}; "train" / "val" / "test" returns a per-conformation 80/10/10 partition (fixed seed, disjoint splits).
cache_dir (pathlib.Path) – Base cache directory; files live under cache_dir/"ani1x".
level (str, optional) – Level of theory: "wb97x_dz" (default, the level the ANI-1x model was trained on), "wb97x_tz", or "ccsd(t)_cbs" (energy-only).
forces (bool, optional) – Include forces (default True). Ignored for energy-only levels like "ccsd(t)_cbs"; raises if forces are requested for a level that has none.
units (str, optional) – "eV" (default) converts energies to eV and forces to eV/A; "hartree" keeps the raw upstream values.
max_molecules (int, optional) – Cap the number of molecule groups read (useful for demos).
max_conformations (int, optional) – Cap the number of (non-NaN) 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, when available and requested, forces (N, 3). ANI-1x is molecular, so no cell / pbc.

Return type:

list of dict} or list of dict

Raises:

ValueError – If units, split, or level is unrecognized, or forces are requested for an energy-only level.
ImportError – If h5py is not installed.