Suppose we have observations $\{y_i\}_{i=1}^{n}$. I would like to compute the average pairwise distance defined by: $$ D = \frac{1}{n^2}\sum_{i=1}^{n}\sum_{j=1}^{n} |y_i - y_j| $$ This computation is done repeatedly in a loop, so optimizing it for speed is important. In my application, each $y_i$ is an array of shape (n, p, k), where the final result of $D$ should be averaged across the first dimension - resulting in an array of shape (p, k). Simplest way to do this is a naive for-loop: Y # Array of shape (n, p, k) D = 0 for i in range(n): for j in range(n): D += np.abs(Y[i, :, :] - Y[j, :, :]) D = D / n**2 However this is of course very slow. My current solution uses shape broadcasting/vectorization and is extremely fast: D = np.mean( np.abs( np.expand_dims(Y, axis=0) - np.expand_dims(Y, axis=1) ), axis=[0,1] ) However a problem I have with this is that I am quickly hitting out-of-memory errors, since it requires storing two intermediate results of shapes (n, n, p, k). For example if n=100, p=100, k=300, this is already two arrays of over 6M points. Any suggestions on how to compute $D$ in a way that is more memory-efficient but also fast? Could we rewrite $D$ somehow in terms of sample statistics? Continue reading...
