enerzyme.models.layers.electrostatics.ElectrostaticEnergyLayer

class enerzyme.models.layers.electrostatics.ElectrostaticEnergyLayer(cutoff_sr: float, cutoff_lr: float | None = None, Bohr_in_R: float = 0.5291772108, Hartree_in_E: float = 1, dielectric_constant: float = 1, cutoff_fn: Literal['polynomial', 'bump', 'smooth'] = 'smooth', flavor: Literal['PhysNet', 'SpookyNet'] = 'SpookyNet')

Bases: BaseFFLayer

__init__(cutoff_sr: float, cutoff_lr: float | None = None, Bohr_in_R: float = 0.5291772108, Hartree_in_E: float = 1, dielectric_constant: float = 1, cutoff_fn: Literal['polynomial', 'bump', 'smooth'] = 'smooth', flavor: Literal['PhysNet', 'SpookyNet'] = 'SpookyNet') → None

Calculate the electrostatic energy from distributed multipoles and atomic positions

Params:

bohr_in_Ra: the numerical value of one Bohr in the unit of atom positions.

Hartree_in_Ea: the numerical value of one Hartree in the unit of energy.

short_range_cutoff: the cutoff of short range interaction, the Coulomb’s law at long-range and a damped term at short-range to avoid the singularity at r = 0 are smoothly interpolated by phi [1]:

chi(r) = phi(2r) + 1 / sqrt{r^2 + 1} + (1 - phi(2r)) / r

long_range_cutoff: the cutoff of long range interaction, outside which the electrostatics are ignored

References:

[1] J. Chem. Theory Comput. 2019, 15, 3678−3693.

get_E_ele_a(Dij_lr: Tensor, Qa: Tensor, idx_i: Tensor, idx_j: Tensor) → Tensor

Compute the atomic electrostatic energy

Params:

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

Qa: Float tensor of atomic charges, shape [N * batch_size]

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

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

Returns:

Ea: Float tensor of atomic electrostatic energy, shape [N * batch_size]