xnn.common.data.hub.spookynet_states.SpookyNetStatesBuilder#

class xnn.common.data.hub.spookynet_states.SpookyNetStatesBuilder[source]#

Bases: DatasetBuilder

Builder for the Ag3+/Ag3- and singlet/triplet CH2 sets.

load(*, split=None, cache_dir, system='ag3', n_train=1000, n_val=100, seed=0, quiet=False)[source]#

Download and read one of the two sets.

Parameters:

  • split (str or None) – None returns {"train": ..., "val": ..., "test": ...}; "train" / "val" / "test" returns that split; "all" returns every structure in file order.

  • cache_dir (pathlib.Path) – Base cache directory; the files are stored under cache_dir/"spookynet_states".

  • system (str, optional) – "ag3" (default) or "carbene".

  • n_train (int, optional) – Sizes of the training and validation splits (1000 and 100); the remaining structures are the test split.

  • n_val (int, optional) – Sizes of the training and validation splits (1000 and 100); the remaining structures are the test split.

  • seed (int, optional) – Seed of the permutation. Defaults to 0.

  • quiet (bool, optional) – Suppress download progress output.

Returns:

dict of {str – Structure dicts with pos, atomic_numbers, energy, forces, dipole, total_charge and spin_multiplicity.

Return type:

list of dict} or list of dict

Raises:

ValueError – For an unknown system or split, or splits larger than the set.