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), total_charge (scalar net charge; the key charge is accepted as a synonym) and weight (scalar per-structure loss weight; see weighted_loss()).
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.

Return type:

AtomicGraph