xnn.common.data.atomic_data.AtomicGraph#

class xnn.common.data.atomic_data.AtomicGraph(pos, atomic_numbers, edge_index, cell_shifts, batch, n_atoms, cell=None, pbc=None, energy=None, forces=None, stress=None, total_charge=None, weight=None)[source]#

Bases: object

Graph representation of one or a batch of atomic systems.

AtomicGraph is the single data abstraction that flows through the whole library: every model in xnn.common.models consumes this object and nothing else. It represents molecular or periodic systems as a graph where nodes are atoms and edges connect neighboring atoms within a cutoff. A batch is several graphs concatenated along the node/edge axes, with a batch vector mapping each node to its structure index (the PyTorch Geometric convention).

Positions, cell and the integer cell_shifts are stored separately from edge displacement vectors; the displacement r_ij is recomputed inside the model (see edge_vectors()) so that autograd can flow back to pos (forces) and cell (stress).

Parameters:

  • pos (Tensor) – Cartesian positions of shape (N, 3).

  • atomic_numbers (Tensor) – Integer atomic number Z per atom, of shape (N,).

  • edge_index (Tensor) – Edge list of shape (2, E) holding [src, dst] node indices; the energy of dst depends on src.

  • cell_shifts (Tensor) – Integer periodic image shift for each edge, of shape (E, 3).

  • batch (Tensor) – Structure index per atom, of shape (N,) (all zero for a single structure).

  • n_atoms (Tensor) – Number of atoms per structure, of shape (B,).

  • cell (Tensor, optional) – Lattice vectors as rows, of shape (B, 3, 3). None for molecular systems.

  • pbc (Tensor, optional) – Boolean periodicity flags, of shape (B, 3). None for molecular systems.

  • energy (Tensor, optional) – Target energy per structure, of shape (B,). Present during training.

  • forces (Tensor, optional) – Target forces, of shape (N, 3). Present during training.

  • stress (Tensor, optional) – Target stress, of shape (B, 3, 3). Present during training.

  • total_charge (Tensor, optional) – Net charge per structure, of shape (B,). None means neutral. Read by the charge-aware models (D4 dispersion, PhysNet, ReaxFF).

  • weight (Tensor, optional) – Per-structure loss weight, of shape (B,). None means every structure counts equally, which is the default and reproduces the unweighted loss exactly. Read only by weighted_loss(); no model sees it.

Variables:

  • pos (Tensor) – Cartesian positions (N, 3).

  • atomic_numbers (Tensor) – Integer atomic number Z per atom (N,).

  • edge_index (Tensor) – Edge list (2, E) as [src, dst].

  • cell_shifts (Tensor) – Integer periodic image shift per edge (E, 3).

  • batch (Tensor) – Structure index per atom (N,).

  • n_atoms (Tensor) – Atoms per structure (B,).

  • cell (Tensor or None) – Lattice vectors as rows (B, 3, 3).

  • pbc (Tensor or None) – Boolean periodicity flags (B, 3).

  • energy (Tensor or None) – Target energy (B,).

  • forces (Tensor or None) – Target forces (N, 3).

  • stress (Tensor or None) – Target stress (B, 3, 3).

  • total_charge (Tensor or None) – Net charge per structure (B,).

  • weight (Tensor or None) – Per-structure loss weight (B,).

property num_graphs: int#

Number of structures in the batch (B).

Type:

int

property num_nodes: int#

Total number of atoms (nodes) across the batch (N).

Type:

int

property num_edges: int#

Total number of edges across the batch (E).

Type:

int

to(device)[source]#

Return a copy of the graph with all tensors moved to device.

Parameters:

device (torch.device or str) – Target device to move every tensor field to.

Returns:

A new AtomicGraph whose tensor fields live on device. Non-tensor fields (e.g. None) are copied unchanged.

Return type:

AtomicGraph

edge_vectors()[source]#

Compute the displacement vector r_ij for every edge.

The displacement is pos[dst] - pos[src] plus, for periodic systems, the contribution cell_shift @ cell of the periodic image. The result is differentiable with respect to pos (forces) and cell (stress). Works for molecular (cell is None) and periodic systems alike.

Returns:

Edge displacement vectors r_ij of shape (E, 3).

Return type:

Tensor