waxmorph.jax.graph.build_edge_index

waxmorph.jax.graph.build_edge_index#

waxmorph.jax.graph.build_edge_index(X, R, particle_count, eps_dist=0.01, max_edges=None)[source]#

Build directed COO adjacency from a detached host snapshot.

A detached float32 host snapshot includes both directions when i != j and dist(X[i], X[j]) <= R[i] + R[j] + eps_dist. Nonpositive particle_count uses all rows. max_edges=None returns [2, E] with num_edges == E. A fixed capacity returns [2, max_edges] padded by (0, 0); real edges occupy columns below num_edges. Capacity overflow raises ValueError.

Examples

>>> X = jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]])
>>> R = jnp.array([0.6, 0.6, 0.6])
>>> edge_index, num_edges = build_edge_index(X, R, 3, eps_dist=0.0)
>>> print(edge_index.tolist(), int(num_edges))
[[0, 1], [1, 0]] 2
Parameters:

  • particle_count (int)

  • eps_dist (float)

  • max_edges (int | None)

Return type:

tuple[Array, Array]