enerzyme.models.layers.dispersion.grimme_d3.GrimmeD3EnergyLayer

class enerzyme.models.layers.dispersion.grimme_d3.GrimmeD3EnergyLayer(learnable: bool = True, Hartree_in_E: float = 1, Bohr_in_R: float = 0.5291772108, cutoff_lr: float | None = None)

Bases: BaseFFLayer

__init__(learnable: bool = True, Hartree_in_E: float = 1, Bohr_in_R: float = 0.5291772108, cutoff_lr: float | None = None) → None

DFT-D3 dispersion correction [1] with Becke-Johnson damping [2].

Params:

learnable: If True, covalent radii, are learnable.

Hartree_in_E: The numerical value of one Hartree in the unit of energy targets

Bohr_in_R: The numerical value of one Bohr in the unit of position features

References:

[1]: J. Chem. Phys. 2010, 132, 154104

[2]: J. Comput. Chem. 2011, 32(7), 1456−1465

get_E_disp_a(Za: Tensor, Dij_lr: Tensor, idx_i: Tensor, idx_j: Tensor) → Tensor

Compute DFT-D3(BJ) energy with a cutoff with the model’s parameters.

Params:

Za: Long tensor of the first atomic number in the pair, shape [N * batch_size]

Dij: Float tensor of distances shape [N_pair * batch_size]

idx_i: Long tensor of the first pair indices, shape [N_pair * batch_size]

idx_j: Long tensor of the second pair indices, shape [N_pair * batch_size]

cutoff: Cutoff for the damping function

Returns:

edisp: Float tensor of atomic dispersion energy, shape [N * batch_size]