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 single-structure graphs into one batched graph.

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.