performing outer addition with numpy
Sorry if this is a silly question but I am just getting started with python/numpy and I'm really not sure of the most efficient ways to go about things. I'm putting together a demo N-body simulator for some students, but for now, I am computing the force between particles by looping over the positions of those particles which is predictably as slow as molasses. Basically, given a vector x[i]
, I would like to compute:
n[i] = sum from j = 0 to n-1, j != i of (x[i]-x[j])^-2,
using numpy functions rather than looping. If there is a way to perform outer addition/multiplication:
m[i,j] = x[i]-x[j],
m[i,j] = x[i]*x[j],
I could use that to do the computation.
All universal functions that take two input arguments have an attribute outer
:
x = np.array([1, 2, 3])
np.subtract.outer(x, x)
gives:
array([[ 0, -1, -2],
[ 1, 0, -1],
[ 2, 1, 0]])
and
np.multiply.outer(x, x)
results in:
array([[1, 2, 3],
[2, 4, 6],
[3, 6, 9]])
Many of numpy's basic operators such as np.add
, np.subtract
, np.multiply
etc. are known as universal functions (ufuncs), and have (amongst other things) an .outer
method:
import numpy as np
a = np.arange(3)
b = np.arange(5)
c = np.add.outer(a, b)
print(repr(c))
# array([[0, 1, 2, 3, 4],
# [1, 2, 3, 4, 5],
# [2, 3, 4, 5, 6]])
Another very powerful technique for doing this sort of thing is broadcasting:
print(repr(a[:, None] + b[None, :]))
# array([[0, 1, 2, 3, 4],
# [1, 2, 3, 4, 5],
# [2, 3, 4, 5, 6]])
Indexing with None
(or alternatively, with np.newaxis
) inserts a new dimension, so a[:, None]
has shape (3, 1)
and b[None, :]
has shape (1, 5)
. Broadcasting "expands" the result along the singleton dimensions, so that it has shape (3, 5)
.