waxmorph.jax.graph.build_graph

Contents

waxmorph.jax.graph.build_graph#

waxmorph.jax.graph.build_graph(X, P, R, particle_count=0, c=None, eps_dist=0.01, max_edges=None)[source]#

Build JAX node features, directed edges, edge features, and real-edge count.

c is required despite its compatibility default. Shapes are [N, C], [2, Q], [Q, 2], and scalar, where Q is the real edge count or max_edges when padded. JAX’s fourth value and optional padding differ from PyTorch’s unpadded three-tuple.

Examples

>>> X = jnp.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]])
>>> P, R, c = jnp.ones((3, 3)), jnp.array([0.6, 0.6, 0.6]), jnp.ones(3)
>>> out = build_graph(X, P, R, 3, c, eps_dist=0.0)
>>> print([tuple(a.shape) for a in out[:3]], int(out[3]))
[(3, 1), (2, 2), (2, 2)] 2
Parameters:

  • particle_count (int)

  • eps_dist (float)

  • max_edges (int | None)

Return type:

tuple[Array, Array, Array, Array]

Contents