pythonのforを使わない高速な書き方を教えてください。

Question

以下のプログラムのコメントで書かれている該当部分をforを使わず記述したいです。
何か別の処理が遅くなさそうな記述方法を教えてください！

import numpy as np
def test(value):
    return value+1, value+2, value+3, value+4

A = np.random.rand(300, 900)
B = np.zeros((300,900,4))

#　以下をfor文を使わない別の書き方を教えて欲しいです！
for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        B[i][j] = test(A[i][j])

Answer 1

numpy使うなら、配列同士の和のほうが楽かなと
shapeが(x, y)のランダム配列aのn繰り返しがA
1からnまでの配列bの(x,y)繰り返しがB
結果はそれを足しています。

import numpy as np
n = 4  # test内でvalueに足す最大の数
x = 2  # 元は300
y = 3  # 元は900

a = np.random.rand(y, x)
A = np.tile(a, (n, 1, 1)).T
b = np.array([i+1 for i in range(n)])
B = np.tile(b, (x, y, 1))
print(A + B)
# [[[1.75150523 2.75150523 3.75150523 4.75150523]
#   [1.21935891 2.21935891 3.21935891 4.21935891]
#   [1.16490574 2.16490574 3.16490574 4.16490574]]

#  [[1.73834419 2.73834419 3.73834419 4.73834419]
#   [1.58224169 2.58224169 3.58224169 4.58224169]
#   [1.48666611 2.48666611 3.48666611 4.48666611]]]

Answer 2

testという関数を使いたいなら以下です。

import numpy as np
def test(value):
    return value+1, value+2, value+3, value+4

A = np.random.rand(300, 900)
B = np.zeros((300,900,4))

B[:,:, 0], B[:,:, 1], B[:, :, 2], B[:, :, 3] = test(A)