Broadcasting a python function on to numpy arrays

前端 未结 3 1299
鱼传尺愫
鱼传尺愫 2020-12-15 11:08

Let\'s say we have a particularly simple function like

import scipy as sp
def func(x, y):
   return x + y

This function evidently works for

3条回答
  •  独厮守ぢ
    2020-12-15 11:14

    Just to get the basic idea, you may modify your function, for example this kind of way:

    def func2(x, y):
        x, y= x[:, None], y[None, :]
        A= x+ y
        A[x<= y]= (x- y)[x<= y]
        return A
    

    Thus with your case, something like this should be a very reasonable starting point:

    In []: def func(x, y):
       ..:     x, y= x[:, None], y[None, :]
       ..:     return x+ y
       ..:
    In []: def func2(x, y):
       ..:     x, y= x[:, None], y[None, :]
       ..:     A, L= x+ y, x<= y
       ..:     A[L]= (x- y)[L]
       ..:     return A
       ..:
    In []: x, y= arange(-2, 3), arange(-2, 3)
    In []: func(x, y)
    Out[]:
    array([[-4, -3, -2, -1,  0],
           [-3, -2, -1,  0,  1],
           [-2, -1,  0,  1,  2],
           [-1,  0,  1,  2,  3],
           [ 0,  1,  2,  3,  4]])
    In []: func2(x, y)
    Out[]:
    array([[ 0, -1, -2, -3, -4],
           [-3,  0, -1, -2, -3],
           [-2, -1,  0, -1, -2],
           [-1,  0,  1,  0, -1],
           [ 0,  1,  2,  3,  0]])
    

    Although this kind of processing may seem to waste resources, it's not necessarily the case. Always measure your programs actual performance and make then (not earlier) necessary changes.

    IMHO for an additional advantage: this kind of 'vectorization' makes your code really consistent and readable eventually.

提交回复
热议问题