xnn.common.data.atomic_data

The single data abstraction that flows through the whole library.

AtomicGraph represents one or a batch of atomic systems (molecular or periodic) as a graph. Every model in xnn.common.models consumes this object and nothing else, which is what keeps the model interface coherent.

Design notes

• Positions, cell and the integer cell_shifts are kept separately from edge displacement vectors. The displacement r_ij is recomputed inside the model as pos[dst] - pos[src] + cell_shift @ cell so that autograd can flow back to pos (forces) and cell (stress). See models.outputs.

• A batch is just several graphs concatenated along the node/edge axes with a batch vector mapping each node to its structure index. This is the same convention PyTorch Geometric uses, so interop is easy if you later want it.

• This is a plain dataclass of tensors – no heavyweight dependency. .to() moves everything to a device in one call.

Classes

AtomicGraph(pos, atomic_numbers, edge_index, ...)

Graph representation of one or a batch of atomic systems.