enerzyme.models.layers.geometry.DistanceLayer#

class enerzyme.models.layers.geometry.DistanceLayer[source]#

Bases: BaseFFLayer

Compute the distance between atoms

__init__() → None[source]#
get_output(Ra: Tensor, idx_i: Tensor, idx_j: Tensor, offsets: Tensor | None = None) → Dict[str, Tensor][source]#

Compute the distance with atom pair indices

Params:#

Ra: Float tensor of atom positions, shape [N * 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]

offsets: Float tensor of distance offsets, shape [N_pair * batch_size]

Returns:#

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

with_vector_on(vij_name: str = 'vij') → None[source]#