emle.models¶

The models module contains Torch modules for creating EMLE models.

class emle.models.EMLEAEVComputer(num_species=7, hypers=None, mask=None, is_external=False, zid_map=None, device=None, dtype=None)[source]¶

Wrapper for AEVCalculator from torchani (not a subclass to make sure it works with TorchScript)

Constructor.

Parameters:

num_species (int) – Number of supported species.
hypers (dict) – Hyperparameters for the wrapped AEVComputer.
mask (torch.BoolTensor) – Mask for the features returned from wrapped AEVComputer.
is_external (bool) – Whether the features are calculated externally.
zid_map (dict or torch.tensor) – Map from zid provided here to the ones passed to AEVComputer.
device (torch.device) – The device on which to run the model.
dtype (torch.dtype) – The data type to use for the models floating point tensors.

cpu(**kwargs)[source]¶: Move all model parameters and buffers to CPU memory.

cuda(**kwargs)[source]¶: Move all model parameters and buffers to CUDA memory.

double()[source]¶: Casts all floating point model parameters and buffers to float64 precision.

float()[source]¶: Casts all floating point model parameters and buffers to float32 precision.

forward(zid, xyz)[source]¶

Evaluate the AEVs.

Parameters:

zid (torch.Tensor (N_BATCH, MAX_N_ATOMS)) – The species indices.
xyz (torch.Tensor (N_BATCH, MAX_N_ATOMS, 3)) – The atomic coordinates.

Returns:

aevs – The atomic environment vectors.

Return type:

torch.Tensor (N_BATCH, MAX_N_ATOMS, N_AEV_COMPONENTS)

to(*args, **kwargs)[source]¶: Performs Tensor dtype and/or device conversion on the model.

class emle.models.EMLEBase(params, n_ref, ref_features, q_core, emle_aev_computer=None, species=None, alpha_mode='species', device=None, dtype=None)[source]¶

Base class for the EMLE model. This is used to compute valence shell widths, core charges, valence charges, and the A_thole tensor for a batch of QM systems, which in turn can be used to compute static and induced electrostatic embedding energies using the EMLE model.

Constructor.

Parameters:

params (dict) – EMLE model parameters.
n_ref (torch.Tensor) – number of GPR references for each element in species list
ref_features (torch.Tensor) – Feature vectors for GPR references.
q_core (torch.Tensor) – Core charges for each element in species list.
alpha_mode (str) –

How atomic polarizabilities are calculated.

”species”:
one volume scaling factor is used for each species

”reference”:
scaling factors are obtained with GPR using the values learned for each reference environment
emle_aev_computer (EMLEAEVComputer) – EMLE AEV computer instance used to compute AEVs (masked and normalized).
species (List[int], Tuple[int], numpy.ndarray, torch.Tensor) – List of species (atomic numbers) supported by the EMLE model.
device (torch.device) – The device on which to run the model.
dtype (torch.dtype) – The data type to use for the models floating point tensors.

cpu(**kwargs)[source]¶: Move all model parameters and buffers to CPU memory.

cuda(**kwargs)[source]¶: Move all model parameters and buffers to CUDA memory.

double()[source]¶: Casts all floating point model parameters and buffers to float64 precision.

float()[source]¶: Casts all floating point model parameters and buffers to float32 precision.

forward(atomic_numbers, xyz_qm, q_total)[source]¶

Compute the valence widths, core charges, valence charges, and A_thole tensor for a batch of QM systems.

Parameters:

atomic_numbers (torch.Tensor (N_BATCH, N_QM_ATOMS,)) – Atomic numbers of QM atoms.
xyz_qm (torch.Tensor (N_BATCH, N_QM_ATOMS, 3)) – Positions of QM atoms in Angstrom.
q_total (torch.Tensor (N_BATCH,)) – Total charge.

Returns:

result – torch.Tensor (N_BATCH, N_QM_ATOMS,),: torch.Tensor (N_BATCH, N_QM_ATOMS,), torch.Tensor (N_BATCH, N_QM_ATOMS * 3, N_QM_ATOMS * 3,))

Valence widths, core charges, valence charges, A_thole tensor

Return type:

(torch.Tensor (N_BATCH, N_QM_ATOMS,),

static get_induced_energy(A_thole: Tensor, charges_mm: Tensor, s: Tensor, mesh_data: Tuple[Tensor, Tensor, Tensor], mask: Tensor) → Tensor[source]¶

Calculate the induced electrostatic energy.

Parameters:

A_thole (torch.Tensor (N_BATCH, MAX_QM_ATOMS * 3, MAX_QM_ATOMS * 3)) – The A matrix for induced dipoles prediction.
charges_mm (torch.Tensor (N_BATCH, MAX_MM_ATOMS,)) – MM charges.
s (torch.Tensor (N_BATCH, MAX_QM_ATOMS,)) – MBIS valence shell widths.
mesh_data (mesh_data object (output of self._get_mesh_data)) – Mesh data object.
mask (torch.Tensor (N_BATCH, MAX_QM_ATOMS)) – Mask for padded coordinates.

Returns:

result – Induced electrostatic energy.

Return type:

torch.Tensor (N_BATCH,)

static get_static_energy(q_core: Tensor, q_val: Tensor, charges_mm: Tensor, mesh_data: Tuple[Tensor, Tensor, Tensor]) → Tensor[source]¶

Calculate the static electrostatic energy.

Parameters:

q_core (torch.Tensor (N_BATCH, N_QM_ATOMS,)) – QM core charges.
q_val (torch.Tensor (N_BATCH, N_QM_ATOMS,)) – QM valence charges.
charges_mm (torch.Tensor (N_BATCH, N_MM_ATOMS,)) – MM charges.
mesh_data (mesh_data object (output of self._get_mesh_data)) – Mesh data object.

Returns:

result – Static electrostatic energy.

Return type:

torch.Tensor (N_BATCH,)

to(*args, **kwargs)[source]¶: Performs Tensor dtype and/or device conversion on the model.

class emle.models.EMLE(model=None, method='electrostatic', alpha_mode='species', atomic_numbers=None, qm_charge=0, mm_charges=None, device=None, dtype=None, create_aev_calculator=True)[source]¶

Torch model for predicting static and induced EMLE energy components.

Constructor.

Parameters:

model (str) – Path to a custom EMLE model parameter file. If None, then the default model for the specified ‘alpha_mode’ will be used.
method (str) –

The desired embedding method. Options are:

”electrostatic”:
Full ML electrostatic embedding.

”mechanical”:
ML predicted charges for the core, but zero valence charge.

”nonpol”:
Non-polarisable ML embedding. Here the induced component of the potential is zeroed.

”mm”:
MM charges are used for the core charge and valence charges are set to zero. If this option is specified then the user should also specify the MM charges for atoms in the QM region.
alpha_mode (str) –

How atomic polarizabilities are calculated.

”species”:
one volume scaling factor is used for each species

”reference”:
scaling factors are obtained with GPR using the values learned for each reference environment
atomic_numbers (List[int], Tuple[int], numpy.ndarray, torch.Tensor) – Atomic numbers for the QM region. This allows use of optimised AEV symmetry functions from the NNPOps package. Only use this option if you are using a fixed QM region, i.e. the same QM region for each evalulation of the module.
qm_charge (int) – The charge on the QM region. This can also be passed when calling the forward method. The non-default value will take precendence.
mm_charges (List[float], Tuple[Float], numpy.ndarray, torch.Tensor) – List of MM charges for atoms in the QM region in units of mod electron charge. This is required if the ‘mm’ method is specified.
device (torch.device) – The device on which to run the model.
dtype (torch.dtype) – The data type to use for the models floating point tensors.
create_aev_calculator (bool) – Whether to create an AEV calculator instance. This can be set to False for derived classes that already have an AEV calculator, e.g. ANI2xEMLE. In that case, it’s possible to hook the AEV calculator to avoid duplicating the computation.

cpu(**kwargs)[source]¶: Move all model parameters and buffers to CPU memory.

cuda(**kwargs)[source]¶: Move all model parameters and buffers to CUDA memory.

double()[source]¶: Casts all floating point model parameters and buffers to float64 precision.

float()[source]¶: Casts all floating point model parameters and buffers to float32 precision.

forward(atomic_numbers: Tensor, charges_mm: Tensor, xyz_qm: Tensor, xyz_mm: Tensor, qm_charge: int | Tensor = 0) → Tensor[source]¶

Computes the static and induced EMLE energy components.

Parameters:

atomic_numbers (torch.Tensor (N_QM_ATOMS,) or (BATCH, N_QM_ATOMS)) – Atomic numbers of QM atoms.
charges_mm (torch.Tensor (max_mm_atoms,) or (BATCH, max_mm_atoms)) – MM point charges in atomic units.
xyz_qm (torch.Tensor (N_QM_ATOMS, 3) or (BATCH, N_QM_ATOMS, 3)) – Positions of QM atoms in Angstrom.
xyz_mm (torch.Tensor (N_MM_ATOMS, 3) or (BATCH, N_MM_ATOMS, 3)) – Positions of MM atoms in Angstrom.
qm_charge (int or torch.Tensor (BATCH,)) – The charge on the QM region.

Returns:

result – The static and induced EMLE energy components in Hartree.

Return type:

torch.Tensor (2,) or (2, BATCH)

to(*args, **kwargs)[source]¶: Performs Tensor dtype and/or device conversion on the model.

class emle.models.ANI2xEMLE(emle_model=None, emle_method='electrostatic', alpha_mode='species', mm_charges=None, qm_charge=0, model_index=None, ani2x_model=None, atomic_numbers=None, device=None, dtype=None)[source]¶

Combined ANI2x and EMLE model. Predicts the in vacuo ANI2x energy along with static and induced EMLE energy components.

Constructor.

Parameters:

emle_model (str) – Path to a custom EMLE model parameter file. If None, then the default model for the specified ‘alpha_mode’ will be used.
emle_method (str) –

The desired embedding method. Options are:

”electrostatic”:
Full ML electrostatic embedding.

”mechanical”:
ML predicted charges for the core, but zero valence charge.

”nonpol”:
Non-polarisable ML embedding. Here the induced component of the potential is zeroed.

”mm”:
MM charges are used for the core charge and valence charges are set to zero.
alpha_mode (str) –

How atomic polarizabilities are calculated.

”species”:
one volume scaling factor is used for each species

”reference”:
scaling factors are obtained with GPR using the values learned for each reference environment
mm_charges (List[float], Tuple[Float], numpy.ndarray, torch.Tensor) – List of MM charges for atoms in the QM region in units of mod electron charge. This is required if the ‘mm’ method is specified.
qm_charge (int) – The charge on the QM region. This can also be passed when calling the forward method. The non-default value will take precendence.
model_index (int) – The index of the ANI2x model to use. If None, then the full 8 model ensemble will be used.
ani2x_model (torchani.models.ANI2x, NNPOPS.OptimizedTorchANI) – An existing ANI2x model to use. If None, a new ANI2x model will be created. If using an OptimizedTorchANI model, please ensure that the ANI2x model from which it derived was created using periodic_table_index=True.
atomic_numbers (List[float], Tuple[float], numpy.ndarray, torch.Tensor (N_ATOMS,)) – Atomic numbers for the QM region. This allows use of optimised AEV symmetry functions from the NNPOps package. Only use this option if you are using a fixed QM region, i.e. the same QM region for each evalulation of the module.
device (torch.device) – The device on which to run the model.
dtype (torch.dtype) – The data type to use for the models floating point tensors.

cpu(**kwargs)[source]¶: Move all model parameters and buffers to CPU memory.

cuda(**kwargs)[source]¶: Move all model parameters and buffers to CUDA memory.

double()[source]¶: Casts all model parameters and buffers to float64 precision.

float()[source]¶: Casts all model parameters and buffers to float32 precision.

forward(atomic_numbers: Tensor, charges_mm: Tensor, xyz_qm: Tensor, xyz_mm: Tensor, qm_charge: int = 0) → Tensor[source]¶

Compute the the ANI2x and static and induced EMLE energy components.

Parameters:

atomic_numbers (torch.Tensor (N_QM_ATOMS,) or (BATCH, N_QM_ATOMS)) – Atomic numbers of QM atoms. A non-existent atom is represented by -1.
charges_mm (torch.Tensor (max_mm_atoms,) or (BATCH, max_mm_atoms)) – MM point charges in atomic units.
xyz_qm (torch.Tensor (N_QM_ATOMS, 3) or (BATCH, N_QM_ATOMS, 3)) – Positions of QM atoms in Angstrom.
xyz_mm (torch.Tensor (N_MM_ATOMS, 3) or (BATCH, N_MM_ATOMS, 3)) – Positions of MM atoms in Angstrom.
qm_charge (int or torch.Tensor (BATCH,)) – The charge on the QM region.

Returns:

result – The ANI2x and static and induced EMLE energy components in Hartree.

Return type:

torch.Tensor (3,) or (3, BATCH)

to(*args, **kwargs)[source]¶: Performs Tensor dtype and/or device conversion on the model.

class emle.models.MACEEMLE(emle_model=None, emle_method='electrostatic', alpha_mode='species', mm_charges=None, qm_charge=0, mace_model=None, atomic_numbers=None, device=None, dtype=None)[source]¶

Combined MACE and EMLE model. Predicts the in vacuo MACE energy along with static and induced EMLE energy components.

Constructor.

Parameters:

emle_model (str) – Path to a custom EMLE model parameter file. If None, then the default model for the specified ‘alpha_mode’ will be used.
emle_method (str) –

The desired embedding method. Options are:

”electrostatic”:
Full ML electrostatic embedding.

”mechanical”:
ML predicted charges for the core, but zero valence charge.

”nonpol”:
Non-polarisable ML embedding. Here the induced component of the potential is zeroed.

”mm”:
MM charges are used for the core charge and valence charges are set to zero.
alpha_mode (str) –

How atomic polarizabilities are calculated.

”species”:
one volume scaling factor is used for each species

”reference”:
scaling factors are obtained with GPR using the values learned for each reference environmentw
mm_charges (List[float], Tuple[Float], numpy.ndarray, torch.Tensor) – List of MM charges for atoms in the QM region in units of mod electron charge. This is required if the ‘mm’ method is specified.
qm_charge (int) – The charge on the QM region. This can also be passed when calling the forward method. The non-default value will take precendence.
mace_model (str) – Name of the MACE-OFF23 models to use. Available models are ‘mace-off23-small’, ‘mace-off23-medium’, ‘mace-off23-large’. To use a locally trained MACE model, provide the path to the model file. If None, the MACE-OFF23(S) model will be used by default.
atomic_numbers (List[int], Tuple[int], numpy.ndarray, torch.Tensor (N_ATOMS,)) – List of atomic numbers to use in the MACE model.
device (torch.device) – The device on which to run the model.
dtype (torch.dtype) – The data type to use for the models floating point tensors.

cpu(**kwargs)[source]¶: Move all model parameters and buffers to CPU memory.

cuda(**kwargs)[source]¶: Move all model parameters and buffers to CUDA memory.

double()[source]¶: Cast all floating point model parameters and buffers to float64 precision.

float()[source]¶: Cast all floating point model parameters and buffers to float32 precision.

forward(atomic_numbers: Tensor, charges_mm: Tensor, xyz_qm: Tensor, xyz_mm: Tensor, qm_charge: int = 0) → Tensor[source]¶

Compute the the MACE and static and induced EMLE energy components.

Parameters:

atomic_numbers (torch.Tensor (N_QM_ATOMS,) or (BATCH, N_QM_ATOMS)) – Atomic numbers of QM atoms.
charges_mm (torch.Tensor (max_mm_atoms,) or (BATCH, max_mm_atoms)) – MM point charges in atomic units.
xyz_qm (torch.Tensor (N_QM_ATOMS, 3), or (BATCH, N_QM_ATOMS, 3)) – Positions of QM atoms in Angstrom.
xyz_mm (torch.Tensor (N_MM_ATOMS, 3) or (BATCH, N_MM_ATOMS, 3)) – Positions of MM atoms in Angstrom.
qm_charge (int) – The charge on the QM region.

Returns:

result – The ANI2x and static and induced EMLE energy components in Hartree.

Return type:

torch.Tensor (3,)

to(*args, **kwargs)[source]¶: Performs Tensor dtype and/or device conversion on the model.