Eine Funktion ähnlich wie integrierte Karte, aber Unterstützung numpy broadcast

Question

Ich möchte eine Funktion schreiben amap, die Build-in map ähnelt aber Unterstützung numpy Broadcast und Ausgabe-Array. Es würde so funktionieren:

>>> amap(lambda x:x**2, 1) 
array(1) 
>>> amap(lambda x:x**2, [1, 2]) 
array([1, 4]) 
>>> amap(lambda x,y:y**2+x**2, 1, [1, 2]) 
array([2, 5]) 
>>> amap(lambda x,y:y**2+x**2, [1, 2], [[1], [2]]) 
array([[2, 5], [5, 8]]) 

Die übergebene Funktion arbeitet auf Skalar.

Ich habe eine Implementierung geschrieben.

def amap(func, *args): 
    '''array version of build-in map 
    amap(function, scalar/sequence[, ...]) -> array 
    ''' 
    args = np.broadcast(None, *args) 
    res = [func(*arg[1:]) for arg in args] 
    res = np.asarray(res).reshape(args.shape) 
    return res 

Es funktioniert, aber es ist irgendwie hässlich und nicht sehr effizient. Haben Sie irgendwelche guten Ideen, dies zu tun? Vor allem etwas mehr numply-stylish?

Answer 1

Verwenden Sie numpy.vectorize direkt?

>>> numpy.vectorize(lambda x:x**2)(1) 
array(1) 
>>> numpy.vectorize(lambda x:x**2)([1,2]) 
array([1, 4]) 
>>> numpy.vectorize(lambda x, y: y**2 + x**2)(1, [1,2]) 
array([2, 5]) 
>>> numpy.vectorize(lambda x, y: y**2 + x**2)([1,2], [[1], [2]]) 
array([[2, 5], 
     [5, 8]])