xnn.common.data.dataset.AtomicDataset#

class xnn.common.data.dataset.AtomicDataset(structures, cutoff)[source]#

Bases: Dataset

PyTorch Dataset wrapping a list of structure dicts.

Each structure is converted to a single-structure AtomicGraph via structure_to_graph() on first access and cached, so the neighbor list for a given index is only computed once.

Parameters:

  • structures (list of dict) – Structure dicts, each in the format accepted by structure_to_graph().

  • cutoff (float) – Neighbor cutoff radius used when building each graph.

Variables:

  • structures (list of dict) – The wrapped structure dicts.

  • cutoff (float) – Neighbor cutoff radius.

classmethod from_file(path, cutoff, index=':', **target_keys)[source]#

Build a dataset from a structure file readable by ASE.

Any ASE-readable format works (.xyz / .extxyz / .cif / VASP / …); frames are converted via load_structures(), picking up energy / forces / stress targets when the file carries them. Requires the ase extra.

Parameters:

  • path (str) – Path to the structure file.

  • cutoff (float) – Neighbor cutoff radius used when building each graph.

  • index (str, optional) – Frame selection passed to ase.io.read(); the default ":" loads all frames.

  • **target_keys – energy_key / forces_key / stress_key overrides for files that store targets under non-standard names (e.g. energy_key="REF_energy"); see atoms_to_structure().

Returns:

Dataset over all selected frames.

Return type:

AtomicDataset

classmethod from_atoms(atoms, cutoff, **target_keys)[source]#

Build a dataset from ASE Atoms object(s) already in memory.

Parameters:

  • atoms (ase.Atoms or list of ase.Atoms) – Structure(s) to convert, via atoms_to_structure().

  • cutoff (float) – Neighbor cutoff radius used when building each graph.

  • **target_keys – energy_key / forces_key / stress_key overrides; see atoms_to_structure().

Returns:

Dataset over the given structure(s).

Return type:

AtomicDataset