waxmorph.torch.graph.build_edge_index

waxmorph.torch.graph.build_edge_index(X, R, particle_count, eps_dist=0.01, device=None)

Return unpadded directed COO contacts with shape [2, E].

A detached float32 CPU snapshot includes both orientations when \(i\ne j\) and \(\lVert x_i-x_j\rVert_2\le R_i+R_j+\varepsilon\), where \(\varepsilon=\mathtt{eps\_dist}\). Autodiff treats this topology as fixed. Nonpositive particle_count uses all rows; device otherwise defaults from the inputs. JAX can pad to a fixed edge capacity and returns an additional edge count.

Examples

>>> X = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]])
>>> R = torch.tensor([0.6, 0.6, 0.6])
>>> print(build_edge_index(X, R, 3, eps_dist=0.0).tolist())
[[0, 1], [1, 0]]

Parameters:

• X (Tensor | array)

• R (Tensor | array)

• particle_count (int)

• eps_dist (float)

• device (device | str | None)

Return type:

Tensor