xnn.common.data.dataset

Datasets and graph batching.

AtomicDataset turns a list of plain structure dicts into AtomicGraph single-structure objects (computing the neighbor list once, cached). collate concatenates several AtomicGraph`s into one batched `AtomicGraph – this is what enables optional batch training: set batch_size=1 to disable it.

Functions

collate(graphs)	Concatenate graphs into one batched graph.
energy_labelled(structures)	The structure dicts whose energy enters the loss.
structure_to_graph(s, cutoff[, device])	Build a single-structure AtomicGraph from a dict of arrays.

Classes

AtomicDataset(structures, cutoff)

PyTorch Dataset wrapping a list of structure dicts.