Auto3D.config.NNPModel

class Auto3D.config.NNPModel(*args, **kwargs)

Bases: Protocol

Protocol for Neural Network Potential models.

Custom NNP models must implement this interface to be compatible with Auto3D’s optimization engine.

coord_pad

Padding value for coordinates (typically 0).

Type:

int

species_pad

Padding value for species (typically -1).

Type:

int

Example

>>> class MyNNP(torch.nn.Module):
...     coord_pad = 0
...     species_pad = -1
...
...     def forward(self, species, coords, charges):
...         # Return energies tensor of shape (batch_size,)
...         return energies

__init__(*args, **kwargs)

Methods

__init__(*args, **kwargs)
forward(species, coords, charges)	Calculate energies for a batch of molecules.

Attributes

coord_pad	Padding value for coordinates in batched tensors.
species_pad	Padding value for atomic species in batched tensors.

coord_pad: int

Padding value for coordinates in batched tensors.

species_pad: int

Padding value for atomic species in batched tensors.

forward(species: Tensor, coords: Tensor, charges: Tensor) → Tensor

Calculate energies for a batch of molecules.

Parameters:

• species – Atomic numbers, shape (batch_size, max_atoms).

• coords – Atomic coordinates, shape (batch_size, max_atoms, 3).

• charges – Molecular charges, shape (batch_size,).

Returns:

Energies tensor of shape (batch_size,) in eV.