Radius Graph

[alexsoleg]

Dear devs,

Is it feasible to update the repository to allow a radius graph neighbour list using GPU, instead of only knn?

Best Regards

[dodgebc]

Interesting! Are you talking about something like radius_neighbors_graph from sci-kit learn?

If you can afford to store it as a dense matrix you can always do jnp.linalg.norm(points - points[:,None], axis=-1) < radius. But assuming you want to use a sparse matrix for a larger number of points it has the same issue as SciPy's query_ball_point. The size of the output is dynamically determined by the input values so it's not possible to write a jit-compatible function for it.

One option could be to do a fixed-k neighbor search and then mask the points that are not within your radius... can you tell me anything more about your intended use case?

[alexsoleg]

Thanks for the quick answer!

Yes, for the moment I am using jnp.linalg.norm(points - points[:,None], axis=-1) < radius but the main problem is that even with jit compilation is our main bottleneck in applications such as molecular dynamics where a graph creation is needed at each time step. Also, the scalability is problematic since it scales O(N^2) and for large proteins this still remains challenging.

In Pytorch Geometric, they have a built-in function that is already optimized for GPU: Pool Radius . However, I haven't seen anything similar in JAX; it could be interesting to have something similar, but GPU-optimized.

The option you mentioned can be interesting and can also be used to filter by batch. But one question comes to my mind, do you think the speed can be increased if instead of jnp.linalg.norm(points - points[:,None], axis=-1) < radius, we first build the tree and then query points for each time step with your implementation?

For the moment we are using this:

from typing import Optional, Tuple
from flax import linen as nn
import jax.numpy as jnp
import jax 

radius_graph = lambda **kwargs: RadiusGraph(**kwargs)

class RadiusGraph(nn.Module):
    """
    Build a (directed) radius-graph with double-vmap, requiring `batch`.

    Attributes
    ----------
    r_max : float
      maximum edge length
    size  : int
      how many edges to output (pads/truncates to exactly this length)
    loop  : bool
      if False exclude self-loops (i→i)
    fill  : int
      what to fill src/dst slots with when padding/truncating
    """
    r_max: float
    size:  int
    fill:  int
    loop:  bool = False

    @nn.compact
    def __call__(
        self,
        pos:   jnp.ndarray,   # [N, D]
        batch: jnp.ndarray    # [N], ints in 0..B-1
    ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
        # 1) two‐level vmap to compute pairwise norms → [N, N]
        def row_dists(x):
            return jax.vmap(lambda y: jnp.linalg.norm(x - y))(pos)
        r = jax.vmap(row_dists)(pos)   # [N, N]


        

        # 2) build mask: same-batch & within-radius & (no self if desired)
        mask = (batch[:, None] == batch[None, :])    # same batch
        
        if self.loop:
            mask &= (r <= self.r_max)
        else:
            mask &= (r <= self.r_max) & (r > 0.0)
        
        
        # 3) static‐size nonzero → two [size] arrays, padding with `fill`
        src, dst = jnp.where(mask, size=self.size, fill_value=self.fill)


        # 4) distances for those edges
        d_ij = r[src, dst]    # [size]

        return src, dst, d_ij

[dodgebc]

Cool! If neighbor queries are the performance bottleneck, the pure-JAX version might not be for you. But if you use the CUDA extension, doing a complete jk.build_and_query can be much faster than forming the pairwise distance matrix for sufficiently large N and small k, not to mention the memory benefits. Take a look at the readme for more details and note that cudaKDTree is currently even faster. If you have any installation issues please do let me know since it's work-in-progress!

Looking at the pytorch function, it seems they have both a radius and max neighbors argument. It might be a good idea to have an r_max argument in query_neighbors and return a fixed array size with a mask, but for now I'd guess first querying max neighbors and then masking isn't going to be much slower unless you make max neighbors much larger than the typical count (in which case I can only suggest cudaKDTree's max heap feature to handle large-k searches). Does that make sense?

[alexsoleg]

I tried this but still slower than the double vmap on a RTX4090 :(

Average over 100 runs:
KDTree: 0.0274 s ± 0.0005
Double-vmap: 0.0002 s ± 0.0000

Any suggestions?

import time
from typing import Tuple
import jax
import jax.numpy as jnp
from flax import linen as nn
import jaxkd as jk

# Module 1: KDTree-based implementation
class RadiusGraph_KDTree(nn.Module):
    r_max:      float
    max_neigh:  int
    size:       int
    fill:       int = -1
    loop:       bool = False

    @nn.compact
    def __call__(self, pos: jnp.ndarray, batch: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
        N = pos.shape[0]
        nbrs, dists = jk.build_and_query(pos, pos, k=self.max_neigh, cuda=True)
        src = jnp.repeat(jnp.arange(N), self.max_neigh)
        dst = nbrs.reshape(-1)
        dij = dists.reshape(-1)

        b_src = jnp.repeat(batch, self.max_neigh)
        b_dst = batch[dst]
        valid = (b_src == b_dst) & (dij <= self.r_max)
        valid = jnp.where(self.loop, valid, valid & (src != dst))

        masked_dij = jnp.where(valid, dij, jnp.inf)
        src_fill = jnp.where(valid, src, self.fill)
        dst_fill = jnp.where(valid, dst, self.fill)

        _, idx = jax.lax.top_k(-masked_dij, self.size)
        src_out = src_fill[idx]
        dst_out = dst_fill[idx]
        dij_out = masked_dij[idx]
        dij_out = jnp.where(dij_out == jnp.inf, self.fill, dij_out)
        return src_out, dst_out, dists

# Module 2: Double-vmap-based implementation
class RadiusGraph(nn.Module):
    r_max: float
    size:  int
    fill:  int
    loop:  bool = False

    @nn.compact
    def __call__(self, pos: jnp.ndarray, batch: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
        # compute full pairwise distances
        def row_dists(x):
            return jax.vmap(lambda y: jnp.linalg.norm(x - y))(pos)
        r = jax.vmap(row_dists)(pos)

        mask = (batch[:, None] == batch[None, :])
        if self.loop:
            mask &= (r <= self.r_max)
        else:
            mask &= (r <= self.r_max) & (r > 0.0)

        src, dst = jnp.where(mask, size=self.size, fill_value=self.fill)
        d_ij = r[src, dst]
        return src, dst, d_ij

# Benchmark with init, apply, and averaging over iterations
if __name__ == '__main__':
    # Parameters
    N, D, B = 500, 3, 10
    num_iters = 100
    key = jax.random.PRNGKey(42)
    pos = jax.random.uniform(key, (N, D))
    batch = jax.random.randint(key, (N,), 0, B)

    # Dynamically set max_neigh to N
    max_neigh = N//2
    # Common template
    params_kd = dict(r_max=0.5, max_neigh=max_neigh, size=1024, fill=-1, loop=False)
    params_vm = dict(r_max=0.5, size=1024, fill=-1, loop=False)

    # Initialize models
    kdt_model = RadiusGraph_KDTree(**params_kd)
    vm_model = RadiusGraph(**params_vm)
    rng1, rng2 = jax.random.split(key)
    vars_kd = kdt_model.init(rng1, pos, batch)
    vars_vm = vm_model.init(rng2, pos, batch)

    # JIT-compiled apply functions
    run_kd = jax.jit(lambda p, b: kdt_model.apply(vars_kd, p, b))
    run_vm = jax.jit(lambda p, b: vm_model.apply(vars_vm, p, b))

    # Warm-up
    _ = run_kd(pos, batch)
    _ = run_vm(pos, batch)

    # Measure KDTree over multiple iterations
    kd_times = []
    for _ in range(num_iters):
        t0 = time.perf_counter()
        out = run_kd(pos, batch)
        jax.block_until_ready(out)
        kd_times.append(time.perf_counter() - t0)
    avg_kd = sum(kd_times) / num_iters

    # Measure vmap over multiple iterations
    vm_times = []
    for _ in range(num_iters):
        t0 = time.perf_counter()
        out = run_vm(pos, batch)
        jax.block_until_ready(out)
        vm_times.append(time.perf_counter() - t0)
    avg_vm = sum(vm_times) / num_iters

    print(f"Average over {num_iters} runs:")
    print(f"  KDTree: {avg_kd:.4f} s ± {jnp.std(jnp.array(kd_times)):.4f}")
    print(f"  Double-vmap: {avg_vm:.4f} s ± {jnp.std(jnp.array(vm_times)):.4f}")

[dodgebc]

Yes, this is where we get to the part about "sufficiently large N and small k." If k is a significant fraction of N, tree search just doesn't make sense since the whole point is to be able to prune away most of the points. In terms of absolute numbers, k=250 is probably too big for this implementation, and N=500 is way too small to worry about any of this, since it won't even saturate the GPU. I'd suggest you time the two approaches for various N and k to get a better feel for it.

points = jr.normal(key, (10_000, 3))
jk.build_and_query(points, points, k=8, cuda=True)
jk.extras.query_neighbors_pairwise(points, points, k=8)

By the way, none of this changes the fact that having a r_max argument for queries is a useful feature and I should get on it!