https://vladfeinberg.com/2021/01/07/vectorizing-ragged-arrays.html

$$ \newcommand{\bsth}{{\boldsymbol\theta}} \newcommand{\nptime}{\
textsf{NP}} \newcommand{\ptime}{\textsf{P}} \newcommand{\disteq}{\
overset{d}{=}} %linalg \newcommand{\mat}[1]{\begin{pmatrix} #1 \end
{pmatrix}} \newcommand{\detmat}[1]{\begin{vmatrix} #1 \end{vmatrix}}
\newcommand{\spanb}[1]{\text{span}\{ #1 \}} \DeclareMathOperator{\
conv}{conv} % convex hull \DeclareMathOperator{\cone}{cone} \
DeclareMathOperator{\vectorize}{vec} \DeclareMathOperator{\matricize}
{mat} \DeclareMathOperator{\adj}{adj} \DeclareMathOperator{\diag}
{diag} \DeclareMathOperator{\tr}{tr} \DeclareMathOperator{\rank}
{rank} \DeclareMathOperator*{\argmin}{argmin} \DeclareMathOperator*{\
proj}{proj} % brackets, norms, cardinalities \newcommand{\pa}[1]{ \
left({#1}\right) } \newcommand{\ha}[1]{ \left[{#1}\right] } \
newcommand{\ca}[1]{ \left\{{#1}\right\} } \newcommand{\inner}[1]{\
left\langle #1 \right\rangle} \newcommand{\innercpy}[1]{\inner{ #1, #
1 }} \newcommand{\norm}[1]{\left\| #1 \right\|} \newcommand{\abs}[1]
{\left|{#1}\right|} \newcommand{\card}[1]{\left\vert{#1}\right\vert}
% math vectors \newcommand{\va}{\textbf{a}} \newcommand{\vb}{\textbf
{b}} \newcommand{\vc}{\textbf{c}} \newcommand{\vd}{\textbf{d}} \
newcommand{\ve}{\textbf{e}} \newcommand{\vf}{\textbf{f}} \newcommand
{\vg}{\textbf{g}} \newcommand{\vh}{\textbf{h}} \newcommand{\vi}{\
textbf{i}} \newcommand{\vj}{\textbf{j}} \newcommand{\vk}{\textbf{k}}
\newcommand{\vl}{\textbf{l}} \newcommand{\vm}{\textbf{m}} \newcommand
{\vn}{\textbf{n}} \newcommand{\vo}{\textbf{o}} \newcommand{\vp}{\
textbf{p}} \newcommand{\vq}{\textbf{q}} \newcommand{\vr}{\textbf{r}}
\newcommand{\vs}{\textbf{s}} \newcommand{\vt}{\textbf{t}} \newcommand
{\vu}{\textbf{u}} \newcommand{\vv}{\textbf{v}} \newcommand{\vw}{\
textbf{w}} \newcommand{\vx}{\textbf{x}} \newcommand{\vy}{\textbf{y}}
\newcommand{\vz}{\textbf{z}} \newcommand{\vzero}{\textbf{0}} \
newcommand{\vone}{\textbf{1}} \newcommand{\valpha}{{\boldsymbol\
alpha}} \newcommand{\vepsilon}{{\boldsymbol\epsilon}} \newcommand{\
vnu}{{\boldsymbol\nu}} \newcommand{\vpi}{{\boldsymbol\pi}} \
newcommand{\veta}{{\boldsymbol\eta}} \newcommand{\vsigma}{ {\
boldsymbol\sigma}} \newcommand{\vbeta}{ {\boldsymbol\beta}} \
newcommand{\vtheta}{ {\boldsymbol\theta}} \newcommand{\vdelta}{ {\
boldsymbol\delta}} \newcommand{\vlambda}{ {\boldsymbol\lambda}} \
newcommand{\vmu}{ {\boldsymbol\mu}} % common math sets \newcommand{\
Z}{\mathbb{Z}} \newcommand{\R}{\mathbb{R}} \newcommand{\C}{\mathbb
{C}} \newcommand{\N}{\mathbb{N}} \newcommand{\Q}{\mathbb{Q}} \
newcommand{\F}{\mathbb{F}} \newcommand{\T}{\mathbb{T}} % limits \def\
sumn{\sum_{n=0}^\infty} \def\limn{\lim_{n\rightarrow\infty}} \def\
prodn{\prod_{n=0}^\infty} % mathcal \newcommand{\mcA}{\mathcal{A}} \
newcommand{\mcB}{\mathcal{B}} \newcommand{\mcC}{\mathcal{C}} \
newcommand{\mcD}{\mathcal{D}} \newcommand{\mcE}{\mathcal{E}} \
newcommand{\mcF}{\mathcal{F}} \newcommand{\mcG}{\mathcal{G}} \
newcommand{\mcH}{\mathcal{H}} \newcommand{\mcI}{\mathcal{I}} \
newcommand{\mcJ}{\mathcal{J}} \newcommand{\mcK}{\mathcal{K}} \
newcommand{\mcL}{\mathcal{L}} \newcommand{\mcM}{\mathcal{M}} \
newcommand{\mcN}{\mathcal{N}} \newcommand{\mcO}{\mathcal{O}} \
newcommand{\mcP}{\mathcal{P}} \newcommand{\mcQ}{\mathcal{Q}} \
newcommand{\mcR}{\mathcal{R}} \newcommand{\mcS}{\mathcal{S}} \
newcommand{\mcT}{\mathcal{T}} \newcommand{\mcU}{\mathcal{U}} \
newcommand{\mcV}{\mathcal{V}} \newcommand{\mcW}{\mathcal{W}} \
newcommand{\mcX}{\mathcal{X}} \newcommand{\mcY}{\mathcal{Y}} \
newcommand{\mcZ}{\mathcal{Z}} % distribs, probability \newcommand{\
disteq}{\overset{d}{=}} \newcommand\independent{\perp \!\!\! \perp} \
DeclareMathOperator{\Laplace}{Laplace} \DeclareMathOperator{\Poisson}
{Poisson} \DeclareMathOperator{\Exponential}{Exponential} \
DeclareMathOperator{\Multinomial}{Multinomial} \DeclareMathOperator{\
Bernoulli}{Bernoulli} \DeclareMathOperator{\Categorical}{Categorical}
\DeclareMathOperator{\Uniform}{Uniform} \DeclareMathOperator{\
Binomial}{Binomial} \DeclareMathOperator{\Hypergeometric}
{Hypergeometric} \DeclareMathOperator{\GammaDist}{Gamma} \
DeclareMathOperator{\NegativeBinomial}{NegativeBinomial} \
DeclareMathOperator\sub{sub} \renewcommand{\d}[1]{\mathop{\mathrm{d}
#1 }} \newcommand{\dkl}[2]{\mathop{D_\mathrm{KL}}\left({#1}\;\middle\
|\;{#2}\right)} \newcommand{\sg}{\mathop{\mathrm{SG}}} \newcommand{\
se}{\mathop{\mathrm{SE}}} %operators \DeclareMathOperator{\power}{{\
mathcal{P}}} \DeclareMathOperator{\var}{var} \DeclareMathOperator{\
cov}{cov} \DeclareMathOperator\mathProb{\mathbb{P}} \
DeclareMathOperator\mathExp{\mathbb{E}} \DeclareMathOperator*\
mathExpUnder{\mathbb{E}} \DeclareMathOperator*\fat{fat} \renewcommand
{\P}{\mathProb} % need to overwrite stupid paragraph symbol \
newcommand{\E}{\mathExp} % need to overwrite stupid paragraph symbol
\newcommand{\set}[2]{ \left\{ #1 \,\middle|\, #2 \right\} } \
newcommand{\CE}[2]{ \mathExp\left[ #1 \,\middle|\, #2 \right] } \
renewcommand{\CP}[2]{ \mathProb\left\{ #1 \,\middle|\, #2 \right\} }
$$

Vlad Feinberg

   
Toggle navigation

  * Home
  * About
  * Categories
  * Featured
  * Search

Vectorizing Ragged Arrays (Numpy Gems, Part 4)

07 Jan 2021 >> numpy, tricks

Frequently, we run into situations where we need to deal with arrays
of varying sizes in numpy. These result in much slower code that
deals with different sizes individually. Luckily, by extracting
commutative and associative operations, we can vectorize even in such
scenarios, resulting in significant speed improvements. This is
especially pronounced when doing the same thing with deep learning
packages like torch, because vectorization matters a lot more on a
GPU.

For instance, take a typical k-means implementation, which has an
inner loop for a naive algorithm like the following.

import numpy as np
from scipy.spatial.distance import cdist

def centroids(X_nd, label_n):
    """
    Given X_nd, a 2-dimensional array of n d-dimensional points,
    and n cluster assignments label_n (a 1-d array of n labels,
    ints in range [0, k)), return (c_kd, dist_n) the k centroids c_kd and the
    squared Euclidean distances dist_n from each point to its centroid.

    Intentionally zero out any empty clusters.
    """
    n, d = X_nd.shape
    k = label_n.max() + 1
    c_kd = np.zeros((k, d))
    dist_n = np.zeros(n)
    for i in range(k):
        ilabel_n = label_n == i
        if not ilabel_n.sum():
            continue
        X_id = X_nd[ilabel_n]
        c_kd[i] = X_id.mean(axis=0)
        dist_n[ilabel_n] = cdist(c_kd[i:i+1, :], X_id, 'sqeuclidean').ravel()
    return c_kd, dist_n

We want to do the same thing (mean and compute pairwise square
distances) to each of these mixed-size X_id arrays, but the for i in
range(k) loop is difficult to vectorize.

Luckily, notice that our main reduction (np.mean) over the ragged
arrays is a composition of two operations: sum / count. Extracting
the reduction operation (the sum) into its own step will let us use
our numpy gem, np.cumsum + np.diff, to aggregate across ragged
arrays.

Then we can take adjacent differences to recover per-cluster means.
This "accumulate ragged" trick will work for any respectable ufunc
with a negation. The key to making it work is to sort such that each
cluster is contiguous.

def inverse_permutation(p):
    ip = np.empty_like(p)
    ip[p] = np.arange(len(p))
    return ip

def vcentroids(X, label):
    """
    Vectorized version of centroids.
    """
    # order points by cluster label
    ix = np.argsort(label)
    label = label[ix]
    Xz = X[ix]

    # compute pos where pos[i]:pos[i+1] is span of cluster i
    d = np.diff(label, prepend=0) # binary mask where labels change
    pos = np.flatnonzero(d) # indices where labels change
    pos = np.repeat(pos, d[pos]) # repeat for 0-length clusters
    pos = np.append(np.insert(pos, 0, 0), len(X))

    # accumulate dimension sums
    Xz = np.concatenate((np.zeros_like(Xz[0:1]), Xz), axis=0)
    Xsums = np.cumsum(Xz, axis=0)

    # reduce by taking differences of accumulations exactly at the
    # endpoints for cluster indices, using pos array
    Xsums = np.diff(Xsums[pos], axis=0)
    counts = np.diff(pos)
    c = Xsums / np.maximum(counts, 1)[:, np.newaxis]

    # re-broadcast centroids for final distance calculation
    repeated_centroids = np.repeat(c, counts, axis=0)
    aligned_centroids = repeated_centroids[inverse_permutation(ix)]
    dist = np.sum((X - aligned_centroids) ** 2, axis=1)

    return c, dist

np.random.seed(1234)

n = 10000
d = 10
k = 10000
x = np.random.randn(n, d)
label = np.random.randint(k, size=n)
c0, dists0 = centroids(x, label)
c1, dists1 = vcentroids(x, label)
np.allclose(c0, c1), np.allclose(dists0, dists1)

(True, True)

%timeit centroids(x, label)
%timeit vcentroids(x, label)

398 ms +- 3.27 ms per loop (mean +- std. dev. of 7 runs, 1 loop each)
3.16 ms +- 104 us per loop (mean +- std. dev. of 7 runs, 100 loops each)

Wow! A 100x, on a CPU no less. And all thanks to using the key
technique:

 1. Sort by ragged array membership
 2. Perform an accumulation
 3. Find boundary indices, compute adjacent differences

Thanks to my friend Ben Eisner for inspiring this post with his SO
question.

Try the notebook out yourself.

(c) Vladimir Feinberg 2021