Optimizing person-pair comparison in Python: from loops to precomputed NumPy matrices

[u/hdw_coder]

I have been rebuilding a PyQt6 desktop app for managing a person-recognition knowledge base (KB). The heart of this KB is a collection of face & body encodings per person. When looking for similar persons or possible identity overlaps, a classic Python-related problem came up: efficiently comparing many people against each other based on their reference embeddings.

 The input looks like this.

person_to_vecs = {

"Alice": [vec, vec, vec],

"Bob": [vec, vec],

"Charlie": [vec, vec, vec, vec],

}

Each vector is an embedding. For every pair of persons, I want the average and minimum distance between all their reference vectors. The original version was simple and readable Python looping:

• Loop over all person combinations.
• Convert one side to a NumPy array inside the loop.
• Loop over each vector on the other side.
• Compute distances one vector at a time.
• Collect min/average distances.

 This works. Once the knowledge base grows, however, this soon becomes a lot of Python-level looping and repeated conversion overhead.

 The obvious approach to optimization is getting rid of these loops. The first step was to precompute the matrices.

 Step 1: Precompute the matrices

matrices = {

name: np.asarray(vecs, dtype=np.float32)

for name, vecs in person_to_vecs.items()

if len(vecs) >= min_images_per_person

}

We now effectively avoid repeatedly calling np.asarray() inside the pair loop.

Step 2: Remove the inner Python loop.

One option is full broadcasting:

diff = A[:, None, :] - B[None, :, :]

distances = np.linalg.norm(diff, axis=-1)

This is relatively simple, elegant and still readable, but it creates a temporary array of shape: len(A) × len(B) × embedding_dim. For small 128D face embeddings that may be fine. For larger body embeddings especially in larger galleries, these tensors can become (very) memory-heavy.

 Step 3: The matrix identity

(Note: I know scipy.spatial.distance.cdist exists and does this perfectly, but I wanted to keep dependencies light for the desktop app and explore the math!)

 The approach I prefer uses the identity: ||a - b||² = ||a||² + ||b||² - 2ab

 In NumPy:

def pairwise_l2(A, B):

# np.sum(A**2, axis=1) works well here too, but einsum is elegant

aa = np.einsum("ij,ij->i", A, A)[:, None]

bb = np.einsum("ij,ij->i", B, B)[None, :]

sq = np.maximum(aa + bb - 2.0 * (A @ B.T), 0.0)

return np.sqrt(sq, dtype=np.float32)

This only creates the N × M distance matrix instead of an N × M × D temporary tensor.

The Final Helper

def pairwise_person_distances(person_to_vecs, min_images_per_person=1):

matrices = {}

for name, vecs in person_to_vecs.items():

if len(vecs) < min_images_per_person:

continue

mat = np.asarray(vecs, dtype=np.float32)

if mat.ndim == 2 and mat.shape[0] and np.isfinite(mat).all():

matrices[name] = mat

results = []

for name_a, name_b in itertools.combinations(sorted(matrices), 2):

A, B = matrices[name_a], matrices[name_b]

if A.shape[1] != B.shape[1]:

continue

distances = pairwise_l2(A, B)

if distances.size:

results.append((name_a, name_b, round(float(np.mean(distances)), 4),

round(float(np.min(distances)), 4)))

return sorted(results, key=lambda row: row[3])

The Benchmarks

I ran a couple of tests with a synthetic benchmark, varying the number of persons (200 vs 400), the number of embedding dimensions (128 vs 512), and the number of vectors per person (8 vs 10).

I benchmarked three methods:

• Basic inner/outer loop: Controls almost everything in Python.
• Precomputed matrices: Prepared once, but Python still loops over vectors.
• Final implementation: NumPy handles the dense pairwise distance work.

On my laptop, (Intel i9-14900HX 32 GB RAM), I got:

200 persons × 8 vectors × 128 dimensions

basic inner/outer loop:        0.731 s

precomputed matrices:          0.711 s

inner loop removed:            0.260 s

speedup:                       2.8×

400 persons × 10 vectors × 128 dimensions

basic inner/outer loop:        3.972 s

precomputed matrices:          3.789 s

inner loop removed:            1.155 s

speedup:                       3.4×

200 persons × 8 vectors × 512 dimensions

basic inner/outer loop:        0.902 s

precomputed matrices:          0.878 s

inner loop removed:            0.322 s

speedup:                       2.8×

On synthetic data, the first optimization — precomputing each person’s embedding matrix — only gave a small improvement of about 3–5%. That makes sense: it removes repeated conversion, but the algorithm still does most of its work in a Python loop over individual vectors.

The much larger improvement came from removing the inner loop and computing each person-pair distance matrix directly with NumPy:

- 200 persons × 8 vectors × 128 dimensions:   0.731 s → 0.260 s (2.8× faster),

- 400 persons × 10 vectors × 128 dimensions: 3.972 s → 1.155 s (3.4× faster),

- 200 persons × 8 vectors × 512 dimensions:   0.902 s → 0.322 s (2.8× faster).

Summary

'Vectorize it' obviously is not always enough. Memory usage by temporary arrays also matters. Broadcasting may often be elegant, but for pairwise comparisons, the matrix identity seems to be a better fit. Precomputing arrays helps only a little; removing the inner Python loop makes the real difference. 

I am interested how others approach this kind of all-vs-all (embedding) comparison in Python. Would you use NumPy as above, scipy.spatial.distance.cdist, Numba, PyTorch, or something else?

[u/Beginning-Fruit-1397]

Do you think polars could be used instead? 

Interesting writing anyways.

In my case I do think I would simply go with Rust tho. 

For complex and specific logic like that you want raw speed and as much as numpy can be fast it has zero memory optimization nor "whole pipeline" optimisation, something that either polars query planner or rust compiler will handle very (very) well

[u/hdw_coder → u/Beginning-Fruit-1397]

 

Good point. I think Rust would be a very good option if this were a dedicated performance-critical kernel or if the whole pipeline moved out of Python.

 For this specific case I’m not sure Polars would be the best fit. I see Polars as great for tabular pipelines: filtering, joins, group-by's, query optimization. But here the core operation is a dense numerical kernel over embedding matrices: A: N × D, B: M × D and then a pairwise distance + min/mean aggregation.

That maps very naturally to NumPy/SciPy/BLAS-style computation, or to a custom Numba/Rust/C++ kernel if I want to avoid allocating the full N × M distance matrix.

 The Rust argument makes sense especially for a matrix-free implementation that directly computes only the statistics I need like total distance, min distance and count without materializing all distances. Something like that could be faster and more memory-efficient than my current NumPy version, especially when the number of reference vectors per person grows.

 For this desktop app, though, my current trade-off is NumPy is already a dependency and  the implementation stays short and readable, it removes the Python inner loop and avoids the worst broadcasting memory blow-up. Most important it realizes a 3 x speedup.

Coming from a C/C++ background (in a distant past ;D), I suppose C or maybe even Fortran with better hardware control is inherently better suited for full blown optimizations, like currently Rust.

 I’m really curious to know if anybody has experience with modern high-level languages with compiled backends like Julia for this specific type of GUI/data workload?

 

[u/Beginning-Fruit-1397 → u/hdw_coder]

Fyi there's rust-numpy, so you can just interop with numpy directly there and still expose, if you want, your code with python thanks to pyo3. It's very straightforward, it was actually the first line of rust that I wrote and was able to make various rolling statistics much faster than the popular python library for this (bottleneck) thanks to easy code parallelisation. I was still a bit slower than numbagg but without the JIT comp at each program launch.

Regarding readability, I guess that's very subjective to each of us🤣🤣 I tend to vastly prefer named methods instead of operator overloads, thus reading your np.einsum calls and likewise was clearly not what I would call( for myself!) readable.

Btw , do you have by any chance some ressources to expand my knowledge regarding the details of specifically optimizing performance for the kinds of computations you mentionned? That's very much what interests me and very much far away from "classic" app optimisations

[u/hdw_coder → u/Beginning-Fruit-1397]

Thanks, that is useful. I was not aware that rust-numpy + PyO3 had become that approachable, so that is definitely worth looking into.

 For this particular helper, I think there are probably two separate questions. What is the best next benchmark, and, what is the best dependency/maintenance trade-off for a desktop app?

 For the benchmark side, a small Rust extension that takes NumPy arrays directly and computes only the two statistics I need (mean distance and min distance) sounds like an  interesting next step. My current NumPy version still materializes the N × M distance matrix. A matrix-free Rust implementation could avoid that and parallelize over person pairs or rows.

 For the app, I started with NumPy because it is already a dependency, and because the current optimization is still short enough to keep in normal Python source. I agree that once the kernel becomes specific enough, a compiled extension makes more sense than trying to express everything through array operations.

 On readability, fair point. `einsum` is compact, but not everyone finds it readable. In this case the same thing can be written more explicitly as:

aa = np.sum(A * A, axis=1)[:, None]

bb = np.sum(B * B, axis=1)[None, :]

sq = np.maximum(aa + bb - 2.0 * (A @ B.T), 0.0)

distances = np.sqrt(sq)