enerzyme.models.layers.electrostatics.ChargeConservationLayer

class enerzyme.models.layers.electrostatics.ChargeConservationLayer

Bases: BaseFFLayer

__init__() → None

Correct the atomic charges to make their summation equal to the total charge by [1]

q^{corrected}_i = q_i - 1 / N (sum_{j=1}^N q_j - Q)

References:

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

get_output(Za: Tensor, Qa: Tensor, Q: Tensor | None = None, batch_seg: Tensor | None = None) → Dict[Literal['Qa', 'Q'], Tensor]

Correct the atomic charge

Params:

Za: Long tensor of atomic numbers, shape [N * batch_size]

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

Q: Float tensor of total charges, shape [batch_size]

batch_seg: Long tensor of batch indices, shape [N * batch_size]

Returns:

Qa_corrected: Float tensor of corrected atomic charge, shape [N * batch_size]

raw_Q: Float tensor of total atomic charge before correction, shape [batch_size]