Эффективные вычисления матрицы тройного индекса в NumPy

Я пытаюсь запустить индексирование тройной матрицы в NumPy, и хотя мой код, похоже, делает то, что я хочу, мне интересно, есть ли более эффективный способ, поскольку время выполнения слишком велико для больших матриц.

У меня есть две матрицы numpy, содержащие время в пути между парами отправления и назначения, первая из которых предназначена для этапа поездки 1 (от i до k), вторая для этапа поездки 2 (от k до j). Я хочу создать две новые матрицы:
- первая матрица, содержащая минимальное время в пути для каждой пары i-j (т.е. наименьшее время для i до k + k до j)
- вторая матрица, показывающая индекс оптимальной k промежуточной остановки, которая вернула это минимальное время в пути.

Мой код выглядит следующим образом:

nb_zones = 100
leg1 = np.random.rand(nb_zones,nb_zones)
leg2 = np.random.rand(nb_zones,nb_zones)

# initialise result matrices
total = np.zeros((nb_zones,nb_zones))
index = np.zeros((nb_zones,nb_zones))

# triple index calcs
for a in range(nb_zones):
    for b in range(nb_zones):
        # max
        total[a,b] = np.min(leg1[a] + leg2[:,b])
        # index
        index[a,b] = np.argmin(leg1[a] + leg2[:,b])

При количестве зон 800 это приводит к времени выполнения до 30 секунд, есть ли более умный способ сделать это без двойного цикла по всем ячейкам матрицы?


python numpy matrix array-broadcasting
person gcornelis    schedule 19.12.2019    source источник

Ответы (2)


arrow_upward
1
arrow_downward

Используя пустую трансляцию и указав ось на np.min и np.argmin, вы можете избежать одного цикла for:

for b in range(nb_zones):
    # max
    total[:,b] = np.min(leg1 + leg2[:,b], axis=1)
    # index
    index[:,b] = np.argmin(leg1 + leg2[:,b], axis=1)
person FBruzzesi    schedule 19.12.2019

arrow_upward
1
arrow_downward

Нет петель:

duallegs = leg1[:,:,np.newaxis] + leg2[np.newaxis, :]
index = np.argmin(duallegs, axis=1)

rowidx, colidx = np.mgrid[0:nb_zones, 0:nb_zones]
total = duallegs[(rowidx, index, colidx)]

Сроки:

import numpy as np

nb_zones = 800
leg1 = np.random.rand(nb_zones,nb_zones)
leg2 = np.random.rand(nb_zones,nb_zones)

def f1():
    # initialise result matrices
    total = np.zeros((nb_zones,nb_zones))
    index = np.zeros((nb_zones,nb_zones))

    # triple index calcs
    for a in range(nb_zones):
        for b in range(nb_zones):
            # max
            total[a,b] = np.min(leg1[a] + leg2[:,b])
            # index
            index[a,b] = np.argmin(leg1[a] + leg2[:,b])

    return total, index

def f2():
    duallegs = leg1[:,:,np.newaxis] + leg2[np.newaxis, :]
    index = np.argmin(duallegs, axis=1)

    rowidx, colidx = np.mgrid[0:nb_zones, 0:nb_zones]
    total = duallegs[(rowidx, index, colidx)]

    return total, index

%timeit f1()
%timeit f2()

8 s ± 819 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
4.85 s ± 360 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
person 9mat    schedule 19.12.2019
comment
И это, и другое предложение отлично работают, спасибо. Я был удивлен тем, что разница во времени выполнения между моим скриптом и этим, похоже, уменьшается по мере увеличения размера матрицы (количества зон), тогда как я ожидал обратного. - person gcornelis; 19.12.2019