How to avoid nested for loops in NumPy?-CodePudding

I have this code.

n_nodes = len(data_x)
X = np.zeros((n_nodes, n_nodes))

for i in range(n_nodes):
  for j in range(n_nodes):
    X[i, j] = data_x[i] ** j

I want to do the same task with no loops used at all. How can I do that with NumPy functions?

CodePudding user response：

I'd suggest

data_x[:,None]**np.arange(n_nodes)

A check

In [17]: data_x = np.array((3,5,7,4,6))
    ...: n_nodes = len(data_x)
    ...: X = np.zeros((n_nodes, n_nodes))
    ...: 
    ...: for i in range(n_nodes):
    ...:   for j in range(n_nodes):
    ...:     X[i, j] = data_x[i] ** j
    ...: print(X)
    ...: print('-----------')
    ...: print(data_x[:,None]**np.arange(n_nodes))
[[1.000e 00 3.000e 00 9.000e 00 2.700e 01 8.100e 01]
 [1.000e 00 5.000e 00 2.500e 01 1.250e 02 6.250e 02]
 [1.000e 00 7.000e 00 4.900e 01 3.430e 02 2.401e 03]
 [1.000e 00 4.000e 00 1.600e 01 6.400e 01 2.560e 02]
 [1.000e 00 6.000e 00 3.600e 01 2.160e 02 1.296e 03]]
-----------
[[   1    3    9   27   81]
 [   1    5   25  125  625]
 [   1    7   49  343 2401]
 [   1    4   16   64  256]
 [   1    6   36  216 1296]]

Some timing

In [18]: %timeit data_x[:,None]**np.arange(n_nodes)
2.18 µs ± 7.49 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [19]: %%timeit
    ...: for i in range(n_nodes):
    ...:     for j in range(n_nodes):
    ...:         X[i, j] = data_x[i] ** j
10.9 µs ± 107 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

CodePudding user response：

You can do this in one step with numpy.power.outer:

np.power.outer(data_x, np.arange(len(data_x)))

CodePudding user response：

If data_x is big, you will be faster using only numpy functions. You can first repeat the input array and then use np.power with a vector giving the powers. This should be calculated fully vectorised in comparison to the already given list comprehension version.

x = np.arange(10)
X = x[:,np.newaxis].repeat(x.size,axis=1)
X = np.power(X,np.arange(x.size))

If data_x is already a numpy array, you can use it directly, if not you would need to do

x = np.array(data_x)