waxmorph.jax.graph.build_edge_features

waxmorph.jax.graph.build_edge_features#

waxmorph.jax.graph.build_edge_features(X, P, edge_index, particle_count)[source]#

Return float32 [distance, polarity angle] for each directed edge.

\[\begin{split}d_{ij} &= \sqrt{\lVert x_i-x_j\rVert_2^2 + \mathrm{EPS\_NORM}^2}, \\ \theta_{ij} &= \arccos\!\left(\operatorname{clip} (p_i^\top p_j,-1+\mathrm{ANGLE\_EPS},1-\mathrm{ANGLE\_EPS})\right).\end{split}\]

JAX regularizes every real-edge distance by EPS_NORM; PyTorch uses the exact norm. P must contain unit vectors because the feature uses supplied dot products directly. Clamping keeps arccos gradients finite. Self-edge padding is masked to zero in both columns.

Examples

>>> X = jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
>>> P = jnp.array([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
>>> edge_index = jnp.array([[0, 1], [1, 0]])
>>> print(jnp.round(build_edge_features(X, P, edge_index, 2), 4).tolist())
[[1.0, 0.00139999995008111], [1.0, 0.00139999995008111]]
Parameters:

  • edge_index (Array)

  • particle_count (int)

Return type:

Array