mbePredict

class mbgdml.mbe.mbePredict(models, predict_model, use_ray=False, n_workers=1, ray_address='auto', wkr_chunk_size=None, alchemy_scalers=None, periodic_cell=None, compute_stress=False)

Predict energies and forces of structures using machine learning many-body models.

This can be parallelized with ray but needs to be initialized with the ray cli or a script using this class. Note that initializing ray tasks comes with some overhead and can make smaller computations much slower. Only GDML models can run in parallel.

Parameters:

• models (list of mbgdml.models.Model) – Machine learning model objects that contain all information to make predictions using predict_model.

• predict_model (callable) – A function that takes Z, R, entity_ids, nbody_gen, model and computes energies and forces. This will be turned into a ray remote function if use_ray is True. This can return total properties or all individual \(n\)-body energies and forces.

• use_ray (bool, default: False) – Use ray to parallelize computations.

• n_workers (int, default: 1) – Total number of workers available for ray. This is ignored if use_ray is False.

• ray_address (str, default: "auto") – Ray cluster address to connect to.

• wkr_chunk_size (int, default: None) – Number of \(n\)-body structures to assign to each spawned worker with ray. If None, it will divide up the number of predictions by n_workers.

• alchemical_scaling (list of mbeAlchemyScale, ) –

• default (None) –

  Alchemical scaling of \(n\)-body interactions of entities.

  Warning

  This has not been thoroughly tested.

• periodic_cell (mbgdml.periodic.Cell, default: None) – Use periodic boundary conditions defined by this object. If this is not None only predict() can be used.

• compute_stress (bool, default: False) –

  Danger

  This implementation is experimental and has not been verified. We do not recommend using this.

  Compute the internal virial contribution, \(\mathbf{W} \left( \mathbf{r}^N \right) / V\), of \(N\) particles at positions \(\mathbf{r}\) to the pressure stress tensor of a periodic box with volume \(V\). The kinetic contribution is not computed here.

box_scaling_type

Treatment of box scaling when mbgdml.mbe.mbePredict.only_normal_stress is True.

anisotropic (default)

No modifications are made to normal stresses which allows box vectors to scale independently.

isotropic

Average the normal stresses so box vectors will scale identically.

Type:

str

compute_nbody(Z, R, entity_ids, comp_ids, model)

Compute all \(n\)-body contributions of a single structure using a mbgdml.models.Model object.

When use_ray = True, this acts as a driver that spawns ray tasks of the predict_model function.

Parameters:

• Z (numpy.ndarray, ndim: 1) – Atomic numbers of all atoms in the system with respect to r.

• R (numpy.ndarray, shape: (len(Z), 3)) – Cartesian coordinates of a single structure.

• entity_ids (numpy.ndarray, shape: (len(Z),)) – Integers specifying which atoms belong to which entities.

• comp_ids (numpy.ndarray, shape: (len(entity_ids),)) – Relates each entity_id to a fragment label. Each item’s index is the label’s entity_id.

• model (mbgdml.models.Model) – Model that contains all information needed by model_predict.

Returns:

• float – Total \(n\)-body energy of r.

• numpy.ndarray – (shape: (len(Z), 3)) - Total \(n\)-body atomic forces of r.

• numpy.ndarray – (optional, shape: (len(Z), 3)) - The internal virial contribution to the pressure stress tensor in units of energy.

compute_nbody_decomp(Z, R, entity_ids, comp_ids, model)

Compute all \(n\)-body contributions of a single structure using a mbgdml.models.Model object.

Stores all individual entity ID combinations, energies and forces. This is more memory intensive. Structures that fall outside the descriptor cutoff will have numpy.nan as their energy and forces.

When use_ray = True, this acts as a driver that spawns ray tasks of the predict_model function.

Parameters:

• Z (numpy.ndarray, ndim: 1) – Atomic numbers of all atoms in the system with respect to r.

• r (numpy.ndarray, shape: (len(Z), 3)) – Cartesian coordinates of a single structure.

• entity_ids (numpy.ndarray, shape: (len(Z),)) – Integers specifying which atoms belong to which entities.

• comp_ids (numpy.ndarray, shape: (len(entity_ids),)) – Relates each entity_id to a fragment label. Each item’s index is the label’s entity_id.

• model (mbgdml.models.Model) – Model that contains all information needed by model_predict.

Returns:

• numpy.ndarray – (ndim: 1) - Energies of all possible \(n\)-body structures. Some elements can be numpy.nan if they are beyond the descriptor cutoff.

• numpy.ndarray – (ndim: 3) - Atomic forces of all possible \(n\)-body structure. Some elements can be numpy.nan if they are beyond the descriptor cutoff.

• numpy.ndarray – (ndim: 2) - All possible entity IDs.

finite_diff_dh

Forward and backward displacement of the cell vectors for finite differences.

Default: 1e-4

Type:

float

get_avail_entities(comp_ids_r, comp_ids_model)

Determines available entity_ids for each comp_id in a structure.

Parameters:

• comp_ids_r (numpy.ndarray, ndim: 1) – Component IDs of the structure to predict.

• comp_ids_model (numpy.ndarray, ndim: 1) – Component IDs of the model.

Returns:

(length: len(comp_ids_r))

Return type:

list

only_normal_stress

Only compute normal (xx, yy, and zz) stress. All other elements will be zero. This is recommended for MD simulations to avoid altering box angular momentum due to the antisymmetric contributions (yz, xz, and xy).

Default: False

Type:

bool

property periodic_cell

Periodic cell for MBE predictions.

Type:

mbgdml.periodic.Cell

predict(Z, R, entity_ids, comp_ids)

Predict the energies and forces of one or multiple structures.

Parameters:

• Z (numpy.ndarray, ndim: 1) – Atomic numbers of all atoms in the system with respect to R.

• R (numpy.ndarray, shape: (N, len(Z), 3)) – Cartesian coordinates of N structures to predict.

• entity_ids (numpy.ndarray, shape: (N,)) – Integers specifying which atoms belong to which entities.

• comp_ids (numpy.ndarray, shape: (N,)) – Relates each entity_id to a fragment label. Each item’s index is the label’s entity_id.

Returns:

• numpy.ndarray – (shape: (N,)) - Predicted many-body energy

• numpy.ndarray – (shape: (N, len(Z), 3)) - Predicted atomic forces.

• numpy.ndarray – (optional, shape: (N, len(Z), 3)) - The pressure stress tensor in units of energy/distance³.

predict_decomp(Z, R, entity_ids, comp_ids)

Predict the energies and forces of one or multiple structures.

Parameters:

• Z (numpy.ndarray, ndim: 1) – Atomic numbers of all atoms in the system with respect to R.

• R (numpy.ndarray, shape: (N, len(Z), 3)) – Cartesian coordinates of N structures to predict.

• entity_ids (numpy.ndarray, shape: (N,)) – Integers specifying which atoms belong to which entities.

• comp_ids (numpy.ndarray, shape: (N,)) – Relates each entity_id to a fragment label. Each item’s index is the label’s entity_id.

Returns:

• list – \(n\)-body energies of all structures in increasing order of \(n\).

• list – \(n\)-body forces of all structures in increasing order of \(n\).

• list – Structure indices and Entity IDs of all structures in increasing order of \(n\). Column 0 is the structure index in R, and the remaining columns are entity IDs.

• list – \(n\)-body orders of each returned item.

use_voigt

Convert the stress tensor to the 1D Voigt notation (xx, yy, zz, yz, xz, xy).

Default: False

Type:

bool

virial_form

The form to use from 10.1063/1.3245303 to compute the internal virial contribution to the pressure stress tensor. \(\mathbf{W} \left( \mathbf{r}^N \right)\) is the internal virial of \(N\) atoms and \(\otimes\) is the outer tensor product (i.e., np.multiply.outer).

group (default)

This computes the virial by considering contributions based on groups of atoms. The number of groups, number of atoms in each group, and number of groups is completely arbitrary. Thus, we can straightforwardly use \(n\)-body combinations as groups and get more insight into the virial contributions origin.

\[\mathbf{W} \left( \mathbf{r}^N \right) = \sum_{k \in \mathbf{0}} \sum_{w = 1}^{N_k} \mathbf{r}_w^k \otimes \mathbf{F}_w^k.\]

Here, \(k\) is a group of atoms that must be in the local cell (i.e., \(k \in \mathbf{0}\)). \(w\) is the atom index within group \(k\) (\(N_k\) is the total number of atoms in the group).

Important

This has to be specifically implemented in the relevant predictor functions.

finite_diff

Uses finite differences by perturbing the cell vectors.

Type:

str