waxmorph.jax.graph

JAX graph construction with detached topology and differentiable features.

Positions and radii are copied to the host for adjacency. The gradient contract covers JAX node and edge features built from live state arrays.

Functions

build_edge_features(X, P, edge_index, ...)	Return float32 [distance, polarity angle] for each directed edge.
build_edge_index(X, R, particle_count[, ...])	Build directed COO adjacency from a detached host snapshot.
build_graph(X, P, R[, particle_count, c, ...])	Build JAX node features, directed edges, edge features, and real-edge count.
build_node_features(c, particle_count)	Return concentrations as float32 [N, C] node features.