xnn.common.data.dataset.structure_to_graph#

xnn.common.data.dataset.structure_to_graph(s, cutoff, device=None)[source]#

Build a single-structure AtomicGraph from a dict of arrays.

Values may be array-likes or tensors; they are coerced to tensors. The neighbor list is computed via build_neighbor_list(). When cell is provided but pbc is not, full periodicity is assumed.

Parameters:

s (dict) – Structure data. Required keys: pos (N, 3) and atomic_numbers (N,). Optional keys: cell (3, 3), pbc (3,), energy (scalar), forces (N, 3), stress (3, 3), dipole (3,) (the molecular dipole moment; dipole_moment is a synonym), polarizability (3, 3), charges (N,) (reference partial charges, the targets of the HDNNP charge networks), total_charge (scalar net charge; the key charge is accepted as a synonym), fragment_charges (N,) (the net charge of the fragment each atom belongs to, for the LES charge solve), spin_multiplicity (scalar 2S + 1; multiplicity is a synonym), weight (scalar per-structure loss weight; see weighted_loss()), head (integer index of the readout head, for multi-head models; see label_head()), the per-structure term weights energy_weight / force_weight / stress_weight (scalars) and atom_force_weight (N,) (the weight of every atom’s force error).
cutoff (float) – Neighbor cutoff radius passed to the neighbor list builder.
device (torch.device, optional) – Device to build the graph on. The neighbor list is the expensive part of this function, so building it where the model already is avoids paying for it on the CPU and copying the result across. Defaults to None, which builds on the CPU as before – the right choice when the graph is being cached by a dataset rather than fed straight to a model.

Returns:

A single-structure graph (batch all zeros, n_atoms of length one), with optional target fields populated when present in s. Floating-point fields take the default dtype, except energy, which is always float64.

Return type:

AtomicGraph