xnn.common.data.dataset.collate

xnn.common.data.dataset.collate(graphs)

Concatenate 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. Using batch_size=1 effectively disables batching.

A batch may mix any structures:

• molecular and periodic ones: when any structure has a cell, the batch keeps cell / pbc and a molecular structure gets a zero cell and no periodic flag (its edges carry no image shift, so every edge keeps its length, and its stress is zero);

• structures with and without energy, force, stress, dipole, polarizability or partial-charge labels: the missing ones are zeros, and energy_mask / forces_mask / stress_mask / dipole_mask / polarizability_mask / charges_mask (per structure) mark the labelled ones for the loss and the metrics;

• structures with and without total_charge (missing = neutral), spin_multiplicity (missing = 1), weight, energy_weight, force_weight, stress_weight, atom_force_weight (missing = 1) or head (missing = 0, the first head).

Parameters:

graphs (list of AtomicGraph) – Graphs (single structures or batches) to concatenate.

Returns:

One batched graph with num_graphs equal to the total number of structures.

Return type:

AtomicGraph