xnn.common.models.dispersion.edge_cell_shifts

xnn.common.models.dispersion.edge_cell_shifts(pos, edge_index, edge_vec, cell, batch)

Integer lattice shifts S of directed edges, (E, 3) int64.

Recovered from edge_vec = pos[dst] - pos[src] + S @ cell (the edge_vectors() convention) with the cell of each edge’s structure (cell (B, 3, 3), an all-zero cell marks a molecule and gives zero shifts), so no neighbor-list metadata is needed.

Parameters:

• pos (Tensor)

• edge_index (Tensor)

• edge_vec (Tensor)

• cell (Tensor)

• batch (Tensor)

Return type:

Tensor