Получение argmin нескольких столбцов в python

Question

В настоящее время я запускаю код, который находит углы, которые наиболее точно описывают функцию, сравнивая результат функции с «таблицей поиска» с использованием np.linalg.norm. На данный момент я делаю это для каждого временного шага моей симуляции, используя для l oop, вот так (я также добавил то, с чем он сравнивается, чтобы подчеркнуть, что это тензор, с которым нужно сравнивать на каждом шаг времени).

    Int3det = TMake.intfromcoordrodswvl(coordinates, avgintensity, wavelength, NAobj, NAcond, alpha3, ndetectors=3) # make 3 by N_timesteps vector
    xpdens = 180 # number of points in x angle 
    ypdens = 360 # number of points in y angle


    thetarange = np.linspace(0, np.pi, xpdens) # range of theta angles
    phirange = np.linspace(-np.pi, np.pi, ypdens) # range of phi angles
    matrix3pol = np.zeros([xpdens, ypdens, 3]) # initialise matrix of intensities for 3 detectors
    for j in enumerate(thetarange): # for angles in theta
        tempmatrix = IntGen.intensitydistribution(j[1], phirange, wavelength, NAobj, NAcond, alpha3, 3) # make intensities for theta, phirange angles for 3 detectors
        matrix3pol[j[0], :, 0] = tempmatrix[:, 0] # put detector 1 in axis 1
        matrix3pol[j[0], :, 1] = tempmatrix[:, 1] # put detector 2 in axis 2
        matrix3pol[j[0], :, 2] = tempmatrix[:, 2] # put detector 3 in axis 3


    for tstep in np.arange(0, len(Int3det)): # for timestep in time trace
         indices3 = np.unravel_index(np.argmin(np.linalg.norm(np.subtract(matrix3pol, Int3det.values[tstep, :]), axis=-1)), matrix3pol[:, :, 0].shape) # get what angle most likely for time step
         Int3res[tstep, 0] = thetarange[indices3[0]] # put theta in result matrix for time step
         Int3res[tstep, 1] = phirange[indices3[1]] # put phi in result matrix for time step

Выполнение for l oop занимает довольно много времени (неудивительно). Есть ли способ ускорить его выполнение?

Спасибо.

Answer 1

Я считаю, что это не ответ, но он слишком длинный для комментария. Поскольку вы пытаетесь найти значения, которые меньше всего отличаются (root поиск), нет никакого смысла использовать np.linalg.norm(), вы можете попытаться сравнить ненормализованную сумму (в квадрате или абсолютную):

difference = np.sum(np.abs(np.subtract(matrix3pol, Int3det.values[tstep, :])))
# or, whichever is faster
difference = np.sum((np.subtract(matrix3pol, Int3det.values[tstep, :]))**2)
# then
np.unravel_index(np.argmin(difference, axis=-1)), matrix3pol[:, :, 0].shape)

Другой подход может заключаться в использовании Scipy.optimize , а именно ограниченного L-BFGS-B . Несмотря на то, что процедура оптимизации сложнее, чем поиск root методом грубой силы, она часто вычисляется быстрее.

Кроме того, стоит попробовать настоящий поиск по таблице, учитывая, что thetarange и phirange ограничены и известны заранее . Идея проста: один раз вычислить массив, содержащий Intensities(theta,phi), а затем использовать вычисленные значения для поиска аргументов без вычитания. Таким образом, значение интенсивности используется позже как индекс массива для нахождения соответствующих тета и фи. Я часто делаю очень похожие вещи в C ++ .

Также, пожалуйста, подумайте об использовании Numba для ускорения работы. Однако для подпрограмм, критичных к производительности, я обычно пишу библиотеку на C ++ и вызываю ее из Python.