nvalchemi.models.demo.DemoModel

class nvalchemi.models.demo.DemoModel(num_atom_types=100, hidden_dim=64)

This model is a simple demo model that computes the energies and conservative forces given an atomic point cloud (not graph!).

The model is not intended to be used in production, but rather to serve as a simple example of how to implement a model in the framework, as well as useful for testing and debugging.

The following elements are design considerations:

• @dataclass(eq=False) provides a straight-forward way for users to specify arguments to construct the model.

• Type annotations are used to set expected input types and output keys, and beartype decorator is used to enforce type checking at runtime.

Parameters:

• num_atom_types (int)

• hidden_dim (int)

forward(atomic_numbers, positions, batch_indices=None, compute_forces=True)

Forward pass of the demo model.

Computes energies and forces for a given molecular system using atomic embeddings and coordinate embeddings with residual connections.

This model is not invariant or equivariant to translation/rotation.

Parameters:

• atomic_numbers (AtomicNumbers) – Atomic numbers for each atom in the system.

• positions (CoordinatesTensor) – Cartesian coordinates of atoms.

• batch_indices (BatchIndices, optional) – Batch indices for each atom.

• compute_forces (bool)

Returns:

Dictionary containing:

• energies: Predicted energy values. Shape: (batch_size, 1)

• forces: Computed forces via automatic differentiation. Shape: (batch_size, 3)

Return type:

dict[str, torch.Tensor]