'How can I manipulate a numpy array without nested loops?

If I have a MxN numpy array denoted arr, I wish to index over all elements and adjust the values like so

for m in range(arr.shape[0]):
    for n in range(arr.shape[1]):
        arr[m, n] += x**2 * np.cos(m) * np.sin(n)

Where x is a random float.

Is there a way to broadcast this over the entire array without needing to loop? Thus, speeding up the run time.



Solution 1:[1]

You can use nested generators of two-dimensional arrays:

import numpy as np
from random import random

x = random()
n, m = 10,20
arr = [[x**2 * np.cos(2*np.pi*j) * np.sin(2*np.pi*i) for j in range(m)] for i in range(n)]

Solution 2:[2]

In [156]: arr = np.ones((2, 3))

Replace the range with arange:

In [157]: m, n = np.arange(arr.shape[0]), np.arange(arr.shape[1])

And change the first array to (2,1) shape. A (2,1) array broadcasts with a (3,) to produce a (2,3) result.

In [158]: A = 0.23**2 * np.cos(m[:, None]) * np.sin(n)
In [159]: A
Out[159]: 
array([[0.        , 0.04451382, 0.04810183],
       [0.        , 0.02405092, 0.02598953]])
In [160]: arr + A
Out[160]: 
array([[1.        , 1.04451382, 1.04810183],
       [1.        , 1.02405092, 1.02598953]])

The meshgrid suggested in the accepted answer does the same thing:

In [161]: np.meshgrid(m, n, sparse=True, indexing="ij")
Out[161]: 
[array([[0],
        [1]]),
 array([[0, 1, 2]])]

This broadcasting may be clearer with:

In [162]: m, n
Out[162]: (array([0, 1]), array([0, 1, 2]))
In [163]: m[:, None] * 10 + n
Out[163]: 
array([[ 0,  1,  2],
       [10, 11, 12]])

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1 Ruslan Davletshin
Solution 2 hpaulj