Какой самый разумный способ сделать это в NumPy?

Question

У меня есть вопрос, на который, я думаю, может быть простой ответ.У меня есть массив с тремя измерениями - (num_users, num_dates, num_holdings).Я хотел бы инициализировать его для некоторых случайных значений теста.random.rand отлично подходит для этого, но для каждого пользователя и каждой даты третье измерение должно иметь сумму 1 (т. е. для любого пользователя и любой даты их авуары должны быть равны 1).Я могу сделать это путем итерации, как в:

num_users = 2
num_dates = 2
num_holdings = 5

test_arr = np.random.rand(num_users, num_dates, num_holdings)

for user in range(num_users):
    for date in range(num_dates):
        starting_total = np.sum(test_arr[user, date, :])
        test_arr[user, date, :] = np.divide(test_arr[user, date, :], starting_total)

# Check it works
print(np.all(np.sum(test_arr, axis=2).reshape(-1)==1))

Но если я создаю несколько массивов, он начинает работать немного медленнее.Плюс это немного неудовлетворительно.Мне было интересно, если кто-нибудь знает о лучшем способе сделать это с помощью векторной математики?

Спасибо!

Answer 1

Вы могли бы сделать

test_arr /= test_arr.sum(axis=2, keepdims=True)

Например:

In [95]: test_arr = np.random.rand(2, 2, 5)

In [96]: test_arr
Out[96]: 
array([[[0.44621493, 0.04093414, 0.30051671, 0.40939041, 0.37251939],
        [0.33997017, 0.81257008, 0.52820553, 0.55382711, 0.11720684]],

       [[0.78460482, 0.43458619, 0.07722273, 0.18181153, 0.52101088],
        [0.47933417, 0.31354249, 0.09966921, 0.59655266, 0.24816989]]])

In [97]: test_arr.sum(axis=2, keepdims=True)
Out[97]: 
array([[[1.56957558],
        [2.35177973]],

       [[1.99923614],
        [1.73726842]]])

Использование keepdims=True означает, что мы получаем результирующую фигуру (2,2,1), которая будет правильно транслироваться при делении на нее.

In [98]: test_arr /= test_arr.sum(axis=2, keepdims=True)

In [99]: test_arr.sum(axis=2)
Out[99]: 
array([[1., 1.],
       [1., 1.]])

Обратите внимание, что из-за ограниченной точности вы не получите точно 1.0 в качестве суммы, но разница незначительна:

In [100]: test_arr.sum(axis=2) - 1.0
Out[100]: 
array([[ 0.00000000e+00,  0.00000000e+00],
       [-1.11022302e-16, -1.11022302e-16]])