xnn.common.data.dataset.collate#

xnn.common.data.dataset.collate(graphs)[source]#

Concatenate single-structure graphs into one batched graph.

Nodes and edges are concatenated along their respective axes, edge_index entries are offset by the running node count, and a batch vector mapping each node to its structure index is built. Optional fields (cell/pbc, energy, forces, stress, total_charge, weight) are only included when present in every input graph. Using batch_size=1 effectively disables batching.

Parameters:

graphs (list of AtomicGraph) – Single-structure graphs to concatenate.

Returns:

One batched graph with num_graphs equal to len(graphs).

Return type:

AtomicGraph