🚀[FEA]: Support 2D points in physicsnemo.nn.functional.neighbors.radius_search #2137

Is this a new feature, an improvement, or a change to existing functionality?

Improvement

How would you describe the priority of this feature request

Low (would be nice)

Please provide a clear description of problem you would like to solve.

radius_search accepts 3D points only. Its docstring and both backend signatures give shapes (N, 3) and (B, N, 3), so this is documented. But a radius graph on a planar mesh (Mesh[2, 2]) or a 2D point cloud is common. knn in the same package accepts any dimension, and PR #2112 adds a Warp knn backend for 1D to 3D points.

With 2D points, the result depends on the backend:

• The default backend is Warp, when Warp is installed. It fails with a Warp conversion error, not a ValueError about the shape. The cause is wp.from_torch(..., dtype=wp.vec3) in physicsnemo/nn/functional/neighbors/radius_search/_warp_impl.py:243-244 and 359-362.
• implementation="torch" returns correct neighbors in 2D (and in 5D) with max_points=None.
• implementation="torch" with max_points and return_points=True fails. The output buffer hardcodes 3 coordinates at radius_search/_torch_impl.py:145.

import torch
from physicsnemo.nn.functional.neighbors import radius_search

torch.manual_seed(0)
points, queries = torch.rand(50, 2), torch.rand(10, 2)
for kwargs in [{}, {"implementation": "torch"}, {"implementation": "torch", "max_points": 4, "return_points": True}]:
    try:
        out = radius_search(points, queries, 0.3, **kwargs)
        print(kwargs, "OK")
    except Exception as error:
        print(kwargs, type(error).__name__, str(error).splitlines()[0])

{} RuntimeError Could not convert PyTorch tensor with shape (50, 2) to Warp array with dtype=<class 'warp._src.types.vec3f'>, ensure that source inner shape is (3,)
{'implementation': 'torch'} OK
{'implementation': 'torch', 'max_points': 4, 'return_points': True} RuntimeError shape mismatch: value tensor of shape [40, 2] cannot be broadcast to indexing result of shape [40, 3]

The request:

1. Accept points of 1 to 3 dimensions in both backends. For Warp, padding the points with zero coordinates to 3D gives exact results.
2. Raise a clear ValueError for unsupported dimensions, before dispatch.
3. Size the torch backend's return_points buffer from points.shape[-1].

Tested with Python 3.13.8, torch 2.13.0+cu130 and warp 1.17.0 on CPU, at commit 7e646b4 (2.3.0a0).

Describe any alternatives you have considered

• Pad 2D points with a zero z coordinate before the call, torch.nn.functional.pad(points, (0, 1)). With Warp this gives the same neighbors as a brute-force cdist check. It works, but every caller must know to do it.
• Use implementation="torch". It is brute force, so it does not scale, and it still fails with return_points=True and max_points.
• Use knn with a distance cutoff. It answers a different question, because it returns a fixed number of neighbors.