`nvalchemi.data`.AtomicData#

pydantic model nvalchemi.data.AtomicData[source]#

Graph-structured container for a single atomic system.

AtomicData is the core input/output record throughout nvalchemi: a molecule, cluster, or periodic cell represented as a graph whose nodes are atoms and whose (optional) edges encode pairwise interactions. It is a pydantic.BaseModel, so every tensor field is type- and shape-validated on construction, and it mixes in DataMixin for graph conveniences (indexing, device movement, property grouping, serialization).

The only required fields are positions ([V, 3]) and atomic_numbers ([V]); everything else — neighbor lists, cell/PBC, labels such as energy/forces/stress, velocities — is optional and defaults to None. Fields are organized into node-, edge-, and system-level groups (see node_properties, edge_properties, system_properties), and custom keys can be attached at runtime via add_node_property(), add_edge_property(), and add_system_property() (extra="allow" on the model config permits these). To collate many systems into one GPU-friendly graph, pass a list of AtomicData to Batch.

Beyond the plain constructor, build instances from common toolkits with the from_atoms() (ASE Atoms) and from_structure() (pymatgen Structure/Molecule) classmethods. Several model_validator hooks run after construction and shape the resulting object: node/edge tensor counts are checked for consistency against atomic_numbers / neighbor_list; all floating-point tensors are cast to the dtype of positions (emitting a UserWarning when a cast happens); atomic_masses are auto-filled from periodictable when omitted; and all tensors are moved onto a single consistent device.

Each field below is validated by its type. Field shapes use jaxtyping axis labels: V atoms/nodes, E edges, B graphs (batch), H features.

Examples

Minimal construction requires only positions and atomic numbers:

>>> import torch
>>> from nvalchemi.data import AtomicData
>>> positions = torch.randn(4, 3)
>>> atomic_numbers = torch.tensor([1, 6, 6, 1], dtype=torch.long)
>>> data = AtomicData(positions=positions, atomic_numbers=atomic_numbers)
>>> data.num_nodes
4

Attach edges (a neighbor list of [source, target] pairs) and system-level labels for a periodic cell:

>>> neighbor_list = torch.tensor([[0, 1], [1, 0], [1, 2], [2, 1]])
>>> data = AtomicData(
...     positions=positions,
...     atomic_numbers=atomic_numbers,
...     neighbor_list=neighbor_list,
...     energy=torch.tensor([[0.5]]),
...     cell=torch.eye(3).unsqueeze(0),
...     pbc=torch.tensor([[True, True, True]]),
... )
>>> data.num_edges
4

Interoperate with ASE and batch several systems together:

>>> from ase.build import molecule
>>> from nvalchemi.data import Batch
>>> water = AtomicData.from_atoms(molecule("H2O"))
>>> batch = Batch.from_data_list([data, water])

Notes

atomic_masses is optional but never stays None: when omitted it is populated from periodictable using atomic_numbers.
Floating-point fields are coerced to the dtype of positions; passing a float64 label alongside float32 positions triggers a cast and a UserWarning. Pass matching dtypes to silence it.
validate_assignment=True means re-assigning a field re-runs validation; use add_node_property() and friends (not raw attribute assignment) to register new custom keys so they are tracked in the correct property group.

field atomic_numbers: Integer[Tensor, 'V'] [Required]#

Atomic numbers for each node [n_nodes]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field positions: Float[Tensor, 'V 3'] [Required]#

Cartesian coordinates for each atom [n_nodes, 3]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field atomic_masses: Float[Tensor, 'V'] | None = None#

Atomic masses [n_nodes]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field atom_categories: list[AtomCategory] | Integer[Tensor, 'V'] | None = None#

Atom categorical index, based on _typing.AtomCategory Enum [n_nodes]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field neighbor_list: Integer[Tensor, 'E 2'] | None = None#

Neighbor list [n_edges, 2]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field shifts: Float[Tensor, 'E 3'] | None = None#

Cartesian displacement vectors for each edge (neighbor_list_shifts @ cell) [n_edges, 3]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field neighbor_list_shifts: Num[Tensor, 'E 3'] | None = None#

Integer lattice image indices for periodic edges [n_edges, 3]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field neighbor_matrix: Integer[Tensor, 'V K'] | None = None#

Dense neighbor matrix [n_nodes, max_neighbors]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field neighbor_matrix_shifts: Num[Tensor, 'V K 3'] | None = None#

Periodic shifts for the dense neighbor matrix [n_nodes, max_neighbors, 3]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field num_neighbors: Integer[Tensor, 'V'] | None = None#

Number of valid neighbors per atom [n_nodes]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field cell: Float[Tensor, 'B 3 3'] | None = None#

Unit cell vectors [3, 3]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field pbc: Bool[Tensor, 'B 3'] | None = None#

Boolean tensor indicating periodic boundary conditions along each dimension

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field forces: Float[Tensor, 'V 3'] | None = None#

Atomic forces [n_nodes, 3]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field energy: Float[Tensor, 'B 1'] | None = None#

Total energy [1, 1]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field stress: Float[Tensor, 'B 3 3'] | None = None#

Tensile-positive Cauchy stress (eV/A^3) [1, 3, 3]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field virial: Float[Tensor, 'B 3 3'] | None = None#

Virial tensor [1, 3, 3]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field dipole: Float[Tensor, 'B 3'] | None = None#

Dipole moment of the system.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field charges: Float[Tensor, 'V'] | None = None#

Partial atomic charges [n_nodes]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field charge: Float[Tensor, 'B 1'] | None = None#

Total system charge [1, 1]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field node_attrs: Float[Tensor, 'V A'] | None = None#

Node attributes [n_nodes, n_node_attrs]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field node_alpha_spins: Float[Tensor, 'V 1'] | None = None#

Alpha spins for each atom, [n_nodes, 1]. Use this field for closed-shell spins.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field node_beta_spins: Float[Tensor, 'V 1'] | None = None#

Beta spins for each atom, [n_nodes, 1]. For restricted spin, use node_alpha_spins instead.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field spin: Float[Tensor, 'B 1'] | None = None#

Spin or multiplicity value for the system, [1, 1]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field graph_alpha_spins: Float[Tensor, 'B 1'] | None = None#

Alpha spins for the entire graph, [1, 1]

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field node_embeddings: Float[Tensor, 'V H'] | None = None#

Embeddings for each node within the batch/graph.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field edge_embeddings: Float[Tensor, 'E H'] | None = None#

Embeddings for each edge within the batch/graph.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field graph_embeddings: Float[Tensor, 'B H'] | None = None#

Embeddings for the entire graph/graphs within a batch.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field velocities: Float[Tensor, 'V 3'] | None = None#

Atomic velocities [n_nodes, 3], in units set by positions.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field momenta: Float[Tensor, 'V 3'] | None = None#

Atomic momenta [n_nodes, 3], in units set by positions.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field kinetic_energies: Float[Tensor, 'V 1'] | None = None#

Per-atom kinetic energies [n_nodes, 1], with the same units as energy.

Constraints:

func = <function _tensor_serialization at 0xeb0d9239b420>
return_type = PydanticUndefined
when_used = json

field info: dict[str, Tensor] [Optional]#: Additional unstructured information about the system.

nvalchemi.data.AtomicData#

`nvalchemi.data`.AtomicData#