Getting argmin of multiple columns in python

Question

I'm currently running some code that finds the angles that most accurately describe a function by comparing the result of a function to a `look-up table' using np.linalg.norm. At the moment I'm doing this for each time step of my simulation using a for loop, like so (I've also put in what it's comparing to, to highlight that it's a tensor that one must compare to at each time step).

    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

The for loop is taking quite some time to execute (unsurprisingly). Is there a way to speed up its execution?

Thanks.

Could you please add some comments to your code describing the current flow? Thank you. — Suthiro, May 09 '20 at 14:40
`argmin` has an `axis` argument which applies it to a certain axis. For use on each column you would use `axis=0`. — Paddy Harrison, May 09 '20 at 16:32
Except that's not the issue - the issue is one needs to subtract the value from the tensor at each time step, but in some "quick" broadcasted away as opposed to the current for loop implementation. — user212883, May 09 '20 at 16:35
If you can provide some example data then we may be able to help. — Paddy Harrison, May 09 '20 at 16:39
matrix3pol has dimensions of 180*360*3 (i.e. 3 detectors, the angular dependence of detector intensity) and then Int3det is 3 * N_timestep (i.e. intensity seen at each detector as a function of time step). Thus at the moment one gets what angles lead to the intensity at each time step by taking the argmin of the difference between the 3 detectors at each time step and what we would expect to see for each angle. This is done using a for loop, which I would like to speed up. — user212883, May 09 '20 at 18:06

score 0 · Answer 1 · answered May 10 '20 at 09:23

I believe it is not an answer, but it is too long for a comment. Since you are trying to find the values that differs the least (root finding), there is no point of using np.linalg.norm() at all, you can try to compare the unnormalized sum (squared or absolute):

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)

Another approach could be to use Scipy.optimize, namely constrained L-BFGS-B. Despite being more complex than brute force root finding, an optimization routine often computes faster.

Also, it is worth to try true table lookup, given that thetarange and phirange are limited and known beforehand. The idea is simple: calculate an array containing Intensities(theta,phi) once and then use calculated values to find the arguments, without subtracting. So the intensity value is used later as an array indice to find corresponding theta and phi. I often do very similar thing in C++.

Also, please, consider using Numba to speed up things. However, for perfomance-critical routines I usually write a library in C++ and call it from Python.

Getting argmin of multiple columns in python

1 Answers1