xnn.common.data.hub.ani2x.ANI2xBuilder#

class xnn.common.data.hub.ani2x.ANI2xBuilder[source]#

Bases: DatasetBuilder

Builder for the ANI-2x data set (7-element wB97X energies & forces).

See the module docstring for the dataset description and citation. The one ~3.7 GB archive is downloaded and extracted once; its atom-count groups are then parsed into xnn structure dicts.

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

Download and preprocess the ANI-2x 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/"ani2x".

  • n_atoms (int or sequence of int, optional) – Which atom-count group(s) to load (e.g. 5 or [4, 5, 6]). Defaults to every group. Fewer/smaller groups mean far less data.

  • forces (bool, optional) – Include forces (default True).

  • units (str, optional) – "eV" (default) converts energies to eV and forces to eV/A; "hartree" keeps the raw upstream values.

  • max_groups (int, optional) – Cap the number of atom-count groups read (useful for demos).

  • max_conformations (int, optional) – Cap the number of conformations kept per group.

  • 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 requested, forces (N, 3). ANI-2x is molecular, so no cell / pbc.

Return type:

list of dict} or list of dict

Raises:
On this page
Show Source