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, energy_mask=None, forces_mask=None, stress_mask=None, dipole=None, polarizability=None, dipole_mask=None, polarizability_mask=None, charges=None, charges_mask=None, fragment_charges=None, fragment_charges_mask=None, total_charge=None, spin_multiplicity=None, weight=None, head=None, energy_weight=None, force_weight=None, stress_weight=None, atom_force_weight=None, compute_dtype=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. In a batch mixing molecular and periodic structures, a molecular structure has a zero cell and no periodic flag (its edges carry no image shift).

  • 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.

  • forces_mask (Tensor, optional) – Which structures carry force labels, bool (B,): set when a batch mixes structures with and without them (the missing ones are zeros in forces and excluded from the loss). None means all do.

  • energy_mask (Tensor, optional) – The same for energy and stress.

  • stress_mask (Tensor, optional) – The same for energy and stress.

  • dipole (Tensor, optional) – Target dipole moment per structure, of shape (B, 3) (e Angstrom in the xnn units). Present during training of the models with a dipole head (PaiNN, PhysNet, AIMNet2).

  • polarizability (Tensor, optional) – Target polarizability tensor per structure, of shape (B, 3, 3) (Angstrom^3). Present during training of PaiNN’s polarizability head.

  • dipole_mask (Tensor, optional) – The same as forces_mask for dipole / polarizability.

  • polarizability_mask (Tensor, optional) – The same as forces_mask for dipole / polarizability.

  • charges (Tensor, optional) – Target partial charge of every atom, of shape (N,) (e). Present during training of the charge networks of the 3G and 4G HDNNPs.

  • charges_mask (Tensor, optional) – The same as forces_mask for charges.

  • fragment_charges (Tensor, optional) – Net charge, in e, of the fragment each atom belongs to, shape (N,) (for the fragment constraints of the LES charge solve); None when unlabeled.

  • fragment_charges_mask (Tensor, optional) – Which structures carry fragment_charges when a batch mixes labeled and unlabeled ones, shape (B,); None means all of them.

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

  • spin_multiplicity (Tensor, optional) – Spin multiplicity 2S + 1 per structure, of shape (B,). None means closed shell (1). Read by the open-shell AIMNet2 models (aimnet2-nse).

  • 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.

  • head (Tensor, optional) – Index of the readout head each structure belongs to, of shape (B,). None means the first head. Read by MultiHead (multi-head fine-tuning) and by the loss, which weights every head’s terms separately.

  • energy_weight (Tensor, optional) – Per-structure weight of one loss term, of shape (B,), on top of weight; 0 drops that label of the structure. None means 1.

  • force_weight (Tensor, optional) – Per-structure weight of one loss term, of shape (B,), on top of weight; 0 drops that label of the structure. None means 1.

  • stress_weight (Tensor, optional) – Per-structure weight of one loss term, of shape (B,), on top of weight; 0 drops that label of the structure. None means 1.

  • atom_force_weight (Tensor, optional) – Weight of every atom’s force error, of shape (N,), on top of the structure’s weights; 0 keeps the atom as environment but drops its force label. None means 1.

  • compute_dtype (dtype | None)

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).

  • stress_mask (energy_mask, forces_mask,) – Bool (B,): the structures that carry energy / force / stress labels.

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

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

  • polarizability_mask (dipole_mask,) – Bool (B,): the structures that carry dipole / polarizability labels.

  • charges (Tensor or None) – Target partial charges (N,).

  • charges_mask (Tensor or None) – Bool (B,): the structures that carry partial-charge labels.

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

  • spin_multiplicity (Tensor or None) – Spin multiplicity per structure (B,).

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

  • head (Tensor or None) – Readout head per structure (B,).

  • stress_weight (energy_weight, force_weight,) – Per-structure term weights (B,).

  • atom_force_weight (Tensor or None) – Per-atom force weight (N,).

property model_dtype: dtype#

The dtype the model computes in (compute_dtype, or the positions’ dtype when it is not set).

Type:

torch.dtype

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

subset(keep)[source]#

The structures flagged in keep as a batch of their own.

Nodes and edges of the other structures are dropped and the indices renumbered; the kept structures stay in their original order. The positions of the result are a view-like gather of pos, so forces and stress computed on the subset flow back to the full graph.

Parameters:

keep (Tensor) – Bool tensor of shape (B,), True for the structures to keep.

Returns:

The sub-batch, with every optional field restricted the same way.

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.

The difference is formed in the dtype of the geometry and returned in model_dtype: absolute coordinates of size L in float32 leave every vector with an error of about 6e-8 L, so the deploy paths keep the positions and cell in float64 for a float32 model and only the (short) vectors are rounded.

Returns:

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

Return type:

Tensor