FrankenPotential

class franken.rf.model.FrankenPotential(gnn_config, rf_config, jac_chunk_size='auto', scale_by_Z=True, num_species=1, atomic_energies=None)

Bases: Module

Maps atomic configurations into high-dimensional vectors through a neural network and random features.

This class provides both functions to map atom configurations with GNN + random features through methods feature_map() and grad_feature_map(), and to perform inference on the potential function of atom configurations given a learned linear model (energy_and_forces()).

Parameters:

gnn_config (BackboneConfig) – Configuration object for the GNN. Sets the ID and other backbone parameters.
rf_config (RFConfig) – Configuration object for the random features (kernel-approximation). Sets all relevant kernel parameters, as well as the number of random features.
jac_chunk_size (int | Literal['auto']) – Force calculation requires computing Jacobians through the GNN. Since this is memory intensive, we can do this in batches across atoms of each configuration. When dealing with large systems, the chunk size becomes important to ensure no out-of-memory errors occur. This can either be set manually or be set automatically (by passing “auto”) which will attempt to determine the largest chunk that fits in memory. Defaults to “auto”.
scale_by_Z (bool) – Whether features should be scaled indipendently for each different species. Defaults to True.
num_species (int) – The number of species that this model will be trained on. Defaults to 1.
atomic_energies (Mapping[int, Tensor | float] | None) – A precomputed dictionary mapping atomic numbers to the average energy of that species. Can be safely left to its default of None in most cases.

gnn

The graph neural network which is used for feature extraction, loaded from a pretrained checkpoint.

Type:: torch.nn.Module

rf

Random-features module franken.rf.heads.RandomFeaturesHead.

Type:: torch.nn.Module

Note

The automatic Jacobian chunking is known to be error-prone. If you encounter out-of-memory errors with this option active, try manually setting the jac_chunk_size parameter.

energy(weights, data)

Computes the energy of a configuration, using a linear model.

if weights is not provided, the weights stored in the FrankenPotential.rf random features object will be used instead.

This function returns the energy of the configuration scaled by the number of atoms in the configuration itself

\[\text{out} = \text{num atoms} \times E(\text{configuration})\]

Parameters:

weights (Tensor | None) – The linear coefficients of the energy model.
configuration – The molecular configuration whose energy to compute.

Returns:

A tensor of size [n_models, n_structures].: n_models denotes the number of models present in the weights attribute; n_structures the number of different structures present in the input data.

Return type:

energies

energy_and_forces(data, weights=None, forces_mode='torch.autograd', add_energy_shift=True)

Infer energy and forces of an atomic configuration using a learned random-features model.

The parameter weights can be used to specified the model’s weights. Otherwise the weights stored in FrankenPotential.rf.weights will be used instead.

The different values of forces_mode correspond to different ways of differentiating through the model to obtain the forces acting on the atoms:

"torch.func" is best for when weights contains multiple linear models on which to perform inference at the same time (in that case weights should be a matrix of shape [n_linear_models, model_size]).
"torch.autograd" is best for when a single linear model is used (i.e. when weights is a vector of shape [model_size])
"no_forces" can be used if forces are not required.

Parameters:

weights (Tensor | None) – weights of the random feature model. Defaults to None, in which case the weights set in FrankenPotential.rf will be used instead.
forces_mode (str) – how to compute the model’s forces. Defaults to "torch.autograd".
add_energy_shift (bool) – whether to add the energy shift to the energy.

Returns:

A tuple containing a tensor representing the potential energy (this is scalar, unless doing inference with multiple models when it can be vector-valued), and another optional tensor representing the forces acting on each atom of the given configuration. If multiple models are given, the forces will have shape [n_linear_models, n_atoms, 3], otherwise they will have shape [n_atoms, 3].

Return type:

Tuple[Tensor, Tensor | None]

energy_and_forces_from_fmaps(data, energy_fmap, forces_fmap, weights=None, add_energy_shift=True)

Compute energies and forces from pre-computed featuremaps. This function does not require calling the underlying GNN.

Parameters:

data (Configuration) – Configuration object describing one or more atomic structures
energy_fmap (Tensor) – Tensor of size [n_structures, num_random_features] containing the original feature map
forces_fmap (Tensor) – Tensor of size [num_random_features, num_atoms * 3] containing the gradients of the original feature map
weights (Tensor | None) – Optional tensor of size [num_models, num_random_features]. This contains the weights of a trained RF model. If no weights are provided, the weights contained in the RF model attached to this class will be used. Energies and forces can be computed for multiple models simulatenously by passing in multiple weight vectors (arranged in 2D).

Returns:

Tensor of size [n_models, n_structures] forces : Tensor of size [n_models, n_atoms, 3]

Return type:

energies

feature_map(data)

Obtain an embedding of each atom, and map it through random features.

The RF mapping computes an average so the final feature map is per-structure, instead of per-atom. In case data contains multiple structures, multiple feature maps are computed.

Returns:: Tensor of size [n_structures, n_random_features]
Return type:: feature_map

forward(data, weights=None, add_energy_shift=True, compute_forces=True)

See docstring of energy_and_forces().

This function defaults to using the ‘torch.autograd’ strategy which allows the model to be jit-compiled.

grad_energy_autograd(weights, data)

Computes the gradient of the energy() acting on a configuration.

The gradient is equivalent to the negative force acting on the configuration.

if weights is not provided, the weights stored in the FrankenPotential.rf random features object will be used instead.

This function returns a tuple: energy_gradient, energy. See grad_energy_func() for a discussion on the performance characteristics of the two implementations.

Parameters:

weights (Tensor or None) – the linear coefficients to compute the energy
data (Configuration) – the molecular configuration whose energy and gradient to compute.

grad_energy_func(weights, data)

Computes the gradient of the energy() acting on a configuration.

The gradient is equivalent to the negative force acting on the configuration.

if weights is not provided, the weights stored in the FrankenPotential.rf random features object will be used instead.

This function returns a tuple: energy_gradient, energy.

Note

This function uses the torch.func package to compute the gradient, which is particularly useful when computing the energy gradient with multiple linear models. In this case weights can be a matrix whose first dimension is the number of linear models. When computing the gradient for a single linear model, use the grad_energy_autograd() method instead for better performance.

Parameters:

weights (Tensor or None) – the linear coefficients to compute the energy
data (Configuration) – the molecular configuration whose energy and gradient to compute.