title: Getting argmin of multiple columns in python
body: <p>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). </p>
  
  <pre><code>    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
  
  </code></pre>
  
  <p>The for loop is taking quite some time to execute (unsurprisingly). Is there a way to speed up its execution?</p>
  
  <p>Thanks.</p>