waxmorph.torch.graph.build_edge_features

waxmorph.torch.graph.build_edge_features(X, P, edge_index, particle_count, device=None)

Return geometric features for each directed edge.

\[[d_{ij},\theta_{ij}] = [\lVert x_i-x_j\rVert_2,\arccos(p_i^\top p_j)].\]

P must contain unit vectors because the feature uses supplied dot products directly. The dot product is clamped to [-1 + ANGLE_EPS, 1 - ANGLE_EPS] before acos. Torch uses the exact Euclidean norm; JAX uses sqrt(||x_i-x_j||^2 + EPS_NORM^2) and masks padded edges. Torch inputs retain gradients through positions and polarities. The output shape is [E, 2].

Examples

>>> X = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
>>> P = torch.tensor([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
>>> edge_index = torch.tensor([[0, 1], [1, 0]])
>>> print(build_edge_features(X, P, edge_index, 2).round(decimals=4).tolist())
[[1.0, 0.00139999995008111], [1.0, 0.00139999995008111]]

Parameters:

• X (Tensor | array)

• P (Tensor | array)

• edge_index (Tensor)

• particle_count (int)

• device (device | str | None)

Return type:

Tensor