CODEWITHSUNDEEP
pythonnumpy
kevinkayaks
kevinkayaks

Reputation: 2726

mapping over higher dimensional numpy arrays

If I have a 1d array a and want to map a function f over each element, I could do 

>>> import numpy as np 
>>> a = np.arange(5)
>>> def f(x):
...     return 3*x + x**2 #whatever 
>>> np.fromiter(map(f,a),float)
array([ 0.,  4., 10., 18., 28.])

I'd like to do something analogous with more complex arrays. One example calculation is this: compose paired 3x3 arrays with matrix multiplication 

>>> a = np.arange(5*2*3**2).reshape(5,2,3,3)
>>> def f(x):
...     return np.matmul(x[0],x[1])
# is there a smarter way? 
>>> np.array([f(x) for x in a])
array([[[   42,    45,    48],
        [  150,   162,   174],
        [  258,   279,   300]],

       [[ 1716,  1773,  1830],
        [ 1986,  2052,  2118],
        [ 2256,  2331,  2406]],

       [[ 5334,  5445,  5556],
        [ 5766,  5886,  6006],
        [ 6198,  6327,  6456]],

       [[10896, 11061, 11226],
        [11490, 11664, 11838],
        [12084, 12267, 12450]],

       [[18402, 18621, 18840],
        [19158, 19386, 19614],
        [19914, 20151, 20388]]])

Another example calculation would be transform every vector in an array of vectors by matrix multiplication

>>> a = np.arange(3*5).reshape(5,3)
>>> def f(x):
...     M = np.arange(3*3).reshape(3,3)
...     return np.dot(M,x)
>>> np.array([f(x) for x in a])
array([[  5,  14,  23],
       [ 14,  50,  86],
       [ 23,  86, 149],
       [ 32, 122, 212],
       [ 41, 158, 275]])

Is there a nice way to do such computations with an np.fromiter like approach? What is the most pythonic way to do these operations with numpy? Is there an approach which handles every example problem here as np.specialnumpything(map(f,a))?

Upvotes: 0

Views: 162

Answers (2)

Daniel F
Daniel F

Reputation: 14399

This is just as easily implemented with broadcasting. Namely:

a = np.arange(5)
a*3 + a**2

array([ 0,  4, 10, 18, 28])

a = np.arange(5*2*3**2).reshape(5,2,3,3)

a[:, 0] @ a[:, 1]

array([[[   42,    45,    48],
        [  150,   162,   174],
        [  258,   279,   300]],

       [[ 1716,  1773,  1830],
        [ 1986,  2052,  2118],
        [ 2256,  2331,  2406]],

       [[ 5334,  5445,  5556],
        [ 5766,  5886,  6006],
        [ 6198,  6327,  6456]],

       [[10896, 11061, 11226],
        [11490, 11664, 11838],
        [12084, 12267, 12450]],

       [[18402, 18621, 18840],
        [19158, 19386, 19614],
        [19914, 20151, 20388]]])


a = np.arange(3*5).reshape(5,3)
M = np.arange(3*3).reshape(3,3)

M.dot(a.T).T

array([[  5,  14,  23],
       [ 14,  50,  86],
       [ 23,  86, 149],
       [ 32, 122, 212],
       [ 41, 158, 275]])

np.einsum('kj, ij -> ik', M, a)

array([[  5,  14,  23],
       [ 14,  50,  86],
       [ 23,  86, 149],
       [ 32, 122, 212],
       [ 41, 158, 275]])

Upvotes: 1

meetaig
meetaig

Reputation: 915

I would use the builtin numpy.nditer, which could be what you're looking for:

https://docs.scipy.org/doc/numpy/reference/arrays.nditer.html

from the examples:

>>> a = np.arange(6).reshape(2,3)
>>> a
>>> array([[0, 1, 2],
           [3, 4, 5]])
>>> with np.nditer(a, op_flags=['readwrite']) as it:
    ...    for x in it:
    ...        x[...] = 2 * x
    ...
>>> a
    array([[ 0,  2,  4],
           [ 6,  8, 10]])

Upvotes: 1

Related Questions