r/pytorch

During a research internship, I ran into a problem involving massive neighbor queries on a GPU for a large particle-dynamics simulation. This led me to experiment with and develop SquareNet, an open-source Python package for NumPy/JAX/PyTorch.

https://preview.redd.it/vb8vpyi421ih1.png?width=705&format=png&auto=webp&s=9de4e42838ca21cd08571c30c1c46c239809e8f9

Its core sorting algorithm (Cartesian sort) enables fast, greedy multidimensional reordering of raw point sets — essentially a form of gridification. Raw points, e.g. (x, y, z, ...), are mapped to unique grid multi-indices [i, j, k, ...] while trying to preserve local geometry, somewhat like a multidimensional generalization of a space-filling curve.

The collection of all multi-indices forms a grid lattice that can be processed efficiently with ML tensor-based frameworks, even when the initial dataset is an irregular point cloud.

I’m wondering whether this could be useful in contexts such as convolutional networks, non uniform fourier transform or ANN search on irregular LiDAR data.

The target use case is approximate but fast and scalable assignment preprocessing, then the grided/tensorized version of the dataset is exploited by standard tensor based frameworks, and result is converted back to the points. High-quality procedures for the assignement part already exist and are well established, such as optimal transport, but they were intractable in my context due to their O(N²/N³) complexity. Cartesian sort, by contrast, runs in O(N log N). It is specifically designed for grid assignment in a greedy setting, trading global optimality for speed and scalability.

In my practical application, involving millions of points processed in a dynamic context (Gaussian blue noise), this simple approach turned out to work well: it provided a ~100× speedup compared to exact brute-force computation of particle interactions, with negligible approximation error (I can provide more details about this experiment if useful).

One caveat is that a single gridification pass introduces a slight axis bias and can produce some distortion/outliers, which can be problematic for challenging distributions where exact geometric precision is required. If exact accuracy is critical, one possible approach would therefore be to build an ensemble of gridifications, each using a different viewpoint/rotation.

Empirically, something like 8 randomly chosen viewpoints seems to give near-perfect recovery of local geometry in a 3D test evaluated with a freud analysis (second link below). However, in the Gaussian blue noise context, where the geometry is smoother, a single viewpoint was already sufficient.

I built an interactive demo on Hugging Face (first link below) to showcase the approach. I’d really appreciate any feedback, especially on whether this idea has already been explored in related computer vision / point-cloud literature, or whether you see potential applications or obvious better alternatives that I may be missing.

interactive HF demo

3D exact nn query discussion