リスト内包表記を使わずにnumpy.arrayの行or列に関数を適用する

やりたいこと

>>> def dobule(x):
...     print(x)
...     return x * 2
...
>>> A = np.arange(6).reshape((3,2))
>>> A
array([[0, 1],
       [2, 3],
       [4, 5]])

自作関数をAの行or列単位で適用する方法を考える。

>>> np.array([double(a) for a in A])
[0 1]
[2 3]
[4 5]
array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]])

上のようにリスト内包表記を使うことで実現は可能だが、
for文を使わないで済む方法を考えたい。

numpy.vectorize

>>> import numpy as np
>>> np.vectorize(double)(A)
0
0
1
2
3
4
5
array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]])

（要素数は6つなのに何で7回printされてるんだろう？）

numpy.vectorizeだと要素ごとに関数を適用することになり、
行or列単位で適用することはできない。

numpy.apply_along_axis

>>> np.apply_along_axis(double, 0, A)
[0 2 4]
[1 3 5]
array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]])
>>> np.apply_along_axis(double, 1, A)
[0 1]
[2 3]
[4 5]
array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]])

numpy.apply_along_axisを使えば行or列単位で関数を適用することができる。