bijective gridification for point clouds

[ArmanddeCacqueray]

Hi everyone,

I’m currently working on point cloud analysis and recently started using freud analytics, which I find very useful — especially for computing structure factors.

In my own work, I developed a data structure called SquareNet, which reorganizes point clouds onto a regular grid, similarly to voxelization. The key difference is that this mapping is bijective: each point is assigned to exactly one cell, and each cell contains exactly one point. The method work in any dimension and scales to datasets with millions of points.

This allows very efficient spatial queries through simple tensor slicing instead of traditional neighbor search structures (KD-trees, etc.) or looping on irregular voxels data.

I share a visual example here of how my mapping works:
https://github.com/ArmanddeCacqueray/SquareNet/blob/main/gridification_paris.gif

With this approach, I can, for example, compute a ground state configurations for hundreds of thousands of particles in ~20 seconds on a laptop CPU:
https://github.com/For-a-few-DPPs-more/hyperuniform-samplers/blob/main/notebooks/blue_noise.ipynb

I maid a python package out of it:
pip install squarenet

I was wondering whether this kind of structure could be relevant for speeding up operations like freud.density.LocalDensity or similar spatial queries.

I’d be very interested in feedback from people working with freud or similar toolkits.

GitHub:
https://github.com/ArmanddeCacqueray/SquareNet

[janbridley]

Thank you for the contribution, this is a very interesting approach. That being said, I'm not certain if this makes sense in Freud -- while constructing the spatial sort is useful, the query can make no guarantees on whether all physically close points are close in index space. I'm able to achieve 100% recall with very large values of wr, but doing so results orders of magnitude more neighbors than required.

If I've missed anything, you have a way to guarantee complete neighbor queries, or have a prototype showing how squarenet could fit into Freud, please feel free to let me know!

# minimal working example
import warnings
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections as mc
import freud
from squarenet import SquareNet

r_max = 1.05

# Generate noisy BCC lattice
box, pts = freud.data.UnitCell.bcc().generate_system(
    num_replicas=(8, 8, 8), sigma_noise=0.05, seed=0
)
N = len(pts)
grid = (N // 64, 8, 8)

# Ground-truth neighbors
nl = (
    freud.locality.AABBQuery(box, pts)
    .query(pts, dict(r_max=r_max, exclude_ii=True))
    .toNeighborList()
)
# Filter bonds resulting from periodicity
freud_pairs = {
    (min(int(i), int(j)), max(int(i), int(j)))
    for i, j in zip(nl.query_point_indices, nl.point_indices)
    if np.linalg.norm(pts[int(i)] - pts[int(j)]) <= r_max
}

# SquareNet window neighbors (wr=10)
sn = SquareNet(IJ_=grid, max_iter=800, warnings_=False)
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    sn.fit(pts)

ncloud = sn.invert_map(sn.views(np.arange(N), wr=10)).reshape(N, -1)
sn_pairs = set()
for i in range(N):
    for j in ncloud[i]:
        if j != i:
            sn_pairs.add((min(i, int(j)), max(i, int(j))))

found = freud_pairs & sn_pairs
missed = freud_pairs - sn_pairs

# Plot
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6), sharex=True, sharey=True)

for ax, title in [
    (ax1, f"Ground truth (freud)\n{len(freud_pairs)} pairs"),
    (
        ax2,
        f"SquareNet (wr=10)\n{len(sn_pairs)} pairs — "
        f"recall = {len(found) / len(freud_pairs):.0%}",
    ),
]:
    ax.add_collection(
        mc.LineCollection(
            [[pts[i, :2], pts[j, :2]] for i, j in found],
            colors="#00C9A4",
            linewidths=0.8,
            alpha=0.5,
            label="found",
            linestyle="dashed",
        )
    )

    if ax is ax2:
        ax.add_collection(
            mc.LineCollection(
                [[pts[i, :2], pts[j, :2]] for i, j in missed],
                colors="#925A4F",
                linewidths=1.8,
                alpha=0.8,
                label=f"missed ({len(missed)})",
            )
        )

    ax.scatter(pts[:, 0], pts[:, 1], s=12, c="k", zorder=3)
    ax.set_title(title, fontsize=11)
    ax.set_aspect("equal")
    ax.legend(loc="upper right", fontsize=9)
    ax.margins(0.05)

fig.suptitle(f"BCC lattice (N={N}): SquareNet neighbor recall", fontsize=13, fontweight="bold")
fig.tight_layout()
fig.savefig("demo_recall.png", dpi=150, bbox_inches="tight")
print(
    f"  freud: {len(freud_pairs)}  SN: {len(sn_pairs)}  "
    f"found: {len(found)}  missed: {len(missed)}  recall: {len(found) / len(freud_pairs):.0%}"
)
plt.show()

[ArmanddeCacqueray]

Hi Jenna,

Thank you for giving this example and I'm really sorry for the long time it took me to answer.

I also tryed to use squarenet for approximate or exact nearest neighbor on my side
(basically, I am currently try to see how it goes on this https://github.com/erikbern/ann-benchmarks/),
and after some experiments I found that the best solution is to use a pool of nets, instead
of just a single net with a window size getting bigger and bigger.

Each net individually will make mistakes, but together they can achive 100% recall with
a very small window radius, and thus a very reasonable amount of candidates.

Here, on the 5000 true neighbors, the SquareAssembly achieve:

• 86% recall with only 16000 candidates
• 98% recall with 24000
• and perfect recall with 43000

43000 total candidates means that only 43 candidates are evaluated per query, out of
the full 1024 sized dataset. On 1024 points it may not sound very impressive,
but as the number of points increases I think gridification could become a pertinent tool
for problems involving neighborhood. As an exemple, I use gridification myself in an other project
(https://github.com/For-a-few-DPPs-more/rgbn) to perform particle simulation with
millions of interacting particles in the context of gaussian blue noise, where KD-tree
or other dynamic data structure adapted for local interraction would make the code much
more complexe, especially for GPU. With squarenet, particle interaction is just a matter of numpy.roll or jax.roll.

#to make sure the latest version is used:
pip install --upgrade squarenet

# minimal working example, now with a pool of SquareNets
import warnings
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.collections as mc
import freud
from squarenet import SquareNet
import math

r_max = 1.05

# Generate noisy BCC lattice
box, pts = freud.data.UnitCell.bcc().generate_system(
    num_replicas=(8, 8, 8), sigma_noise=0.05, seed=0
)
N = len(pts)


# Ground-truth neighbors
nl = (
    freud.locality.AABBQuery(box, pts)
    .query(pts, dict(r_max=r_max, exclude_ii=True))
    .toNeighborList()
)
# Filter bonds resulting from periodicity
freud_pairs = {
    (min(int(i), int(j)), max(int(i), int(j)))
    for i, j in zip(nl.query_point_indices, nl.point_indices)
    if np.linalg.norm(pts[int(i)] - pts[int(j)]) <= r_max
}

class SquareAssembly():
    def __init__(self, gridshape, nexperts):
        """
        A square assembly contains n experts.
        Each expert is a SquareNet, fitted on a different rotation of the dataset.
        The more experts, the more accurate the assembly.
        """
        self.gridshape = np.array(gridshape)
        self.nets = []
        self.rotations = []
        self.shifts = np.array([
            [0, 0, 0],
            [-1, 0, 0],
            [1, 0, 0],
            [0, -1, 0],
            [0, 1, 0],
            [0, 0, -1],
            [0, 0, 1]
            
        ])
        for _ in range(nexperts):
            self.nets.append(SquareNet(gridshape=gridshape, warnings_ = False, verbose = 0))
            Q, _ = np.linalg.qr(np.random.randn(3, 3))
            self.rotations.append(Q)
        self.unpaddings = []  
    def fit(self, pts):
        for net, Q in zip(self.nets, self.rotations):
            #=================================================
            #each net will see a different rotation of the data,
            #allowing the full assembly to capture structure
            #from different perspectives.
            rotated = pts @ Q
            #=================================================
            #padding allow to fit N points with grishape > N...
            #... net.pad wil just add as many empty cells as neaded
            padded = net.pad(rotated) 
            #=================================================
            true_pts = np.isfinite(padded).all(axis = -1)
            unpadding = np.full(len(padded), -1, dtype = int)
            unpadding[true_pts] = np.arange(len(pts))
            net.fit(padded)
            self.unpaddings.append(unpadding)
    def query(self, pt):
        candidates = []
        for net, Q, unpad in zip(self.nets, self.rotations, self.unpaddings):
            rotated_pt = pt @ Q 
            for side in ["left", "right"]:
                #=================================================
                #search_sorted is basically dichotomy on multi indexes
                # it naturaly generalises numpy 1D search sorted
                grid_idx = net.search_sorted(rotated_pt, side = side)
                #=================================================
                for shift in self.shifts:
                    padded_idx = net.invert_mapidx((grid_idx + shift).clip(0, self.gridshape - 1))
                    idx = unpad[padded_idx]
                    if idx >= 0:
                      candidates.append(idx)
        return np.unique(np.array(candidates))

def try_params(gridshape, nexperts):
    sn = SquareAssembly(gridshape, nexperts)
    sn.fit(pts)
    sn_pairs = set()
    for i in range(N):
        for j in sn.query(pts[i]):
            if j != i:
              sn_pairs.add((min(i, j), max(i, j)))

    found = freud_pairs & sn_pairs
    missed = freud_pairs - sn_pairs

    # Plot
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6), sharex=True, sharey=True)

    for ax, title in [
        (ax1, f"Ground truth (freud)\n{len(freud_pairs)} pairs"),
        (
            ax2,
            f"SquareNet (params {gridshape, nexperts}) \n{len(sn_pairs)} pairs — "
            f"recall = {len(found) / len(freud_pairs):.0%}",
        ),
    ]:
        ax.add_collection(
            mc.LineCollection(
                [[pts[i, :2], pts[j, :2]] for i, j in found],
                colors="#00C9A4",
                linewidths=0.8,
                alpha=0.5,
                label="found",
                linestyle="dashed",
            )
        )

        if ax is ax2:
            ax.add_collection(
                mc.LineCollection(
                    [[pts[i, :2], pts[j, :2]] for i, j in missed],
                    colors="#925A4F",
                    linewidths=1.8,
                    alpha=0.8,
                    label=f"missed ({len(missed)})",
                )
            )

        ax.scatter(pts[:, 0], pts[:, 1], s=12, c="k", zorder=3)
        ax.set_title(title, fontsize=11)
        ax.set_aspect("equal")
        ax.legend(loc="upper right", fontsize=9)
        ax.margins(0.05)

    fig.suptitle(f"BCC lattice (N={N}): SquareNet neighbor recall", fontsize=13, fontweight="bold")
    fig.tight_layout()
    fig.savefig("demo_recall.png", dpi=150, bbox_inches="tight")
    print(
        f"  freud: {len(freud_pairs)}  SN: {len(sn_pairs)}  "
        f"found: {len(found)}  missed: {len(missed)}  recall: {len(found) / len(freud_pairs):.0%}"
    )
    plt.show()

DN = math.ceil(N**(1/3))
gridshape = (DN,)*3
try_params(gridshape, 2)
try_params(gridshape, 4)
try_params(gridshape, 8)
try_params(gridshape, 16)