waxmorph.torch.graph.build_graph

waxmorph.torch.graph.build_graph(X, P, R, particle_count=0, c=None, eps_dist=0.01, device=None)

Return (node_features, edge_index, edge_features).

Shapes are [N, C], [2, E], and [E, 2]. Torch features remain differentiable, while contacts use a detached position/radius snapshot. c is required. JAX returns a 4-tuple with an edge count and optional padding.

Examples

>>> X = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]])
>>> P, R, c = torch.ones(3, 3), torch.tensor([0.6, 0.6, 0.6]), torch.ones(3)
>>> print([tuple(t.shape) for t in build_graph(X, P, R, 3, c, eps_dist=0.0)])
[(3, 1), (2, 2), (2, 2)]

Parameters:

• X (Tensor | array)

• P (Tensor | array)

• R (Tensor | array)

• particle_count (int)

• c (Tensor | array | None)

• eps_dist (float)

• device (device | str | None)

Return type:

tuple[Tensor, Tensor, Tensor]