title: Vectorize/accelerate numpy function with two arguments of different dimensions
body: <p>I am not sure if this has been asked before. I couldn't find many relevant results on SO or Google. Anyway, here is what I am trying to do. I have a function that is created at runtime which takes in 4 parameters. </p>
  
  <pre><code>my_func(t, x, p, a)
  </code></pre>
  
  <p><code>t</code> is a scalar float.<br>
  <code>x</code> is a 1D numpy array of floats.<br>
  <code>p</code> and <code>a</code> are dictionaries.</p>
  
  <p>I have a numpy array T and a 2D numpy array X. I want to call the function with each element of T and each column of X and store the result in another numpy array. The function returns a numpy array that forms the column of the solution. 
  My current, naive, very unpythonic implementation is as follows:</p>
  
  <pre><code>for i in range(len(T)):
      _u = my_func(T[i],X[:,i],p,a)
      sol.u[:,i] = _u   # The length of u is unrelated to the other variables
  </code></pre>
  
  <p>I could've used numexpr if it had allowed custom functions. I have also looked at numpy.vectorize and tried using it like this:</p>
  
  <pre><code>f = np.vectorize(my_func)
  f.excluded.add(2)
  f.excluded.add(3)
  f(T,X,p,a)
  </code></pre>
  
  <p>This gives me the error "IndexError: invalid index to scalar variable" from somewhere inside the function. I am obviously not getting how numpy broadcasting works. </p>
  
  <p>Is it possible to do this using np.vectorize?<br>
  Or is there some way to use np.apply_along_axis?</p>
  
  <p>Another possibility I see is slicing up the arrays beforehand and forming a list of tuples with the right arguments and using them with some form of "map" functionality.</p>
  
  <p>Update: Another solution I found:</p>
  
  <pre><code>f = lambda _t, _X: my_func(_t,_X,p,a)
  u = np.array(list(map(f, T, list(X.T)))).T
  </code></pre>
  
  <p>It seems to work. But is that the most efficient way to do it?</p>
  
  <p>Update 2:
  Here is a sample function:</p>
  
  <pre><code>def my_func(_t,_X,_p,_aux):
      [x,y,v,lamX,lamY,lamV,tf,] = _X[:7]
  
      g = _aux['const']['g']
  
      theta = -2*atan((lamX*v + sqrt(g**2*lamV**2 - 2*g*lamV*lamY*v + lamX**2*v**2 + lamY**2*v**2))/(g*lamV - lamY*v))
      return theta
  </code></pre>
  
  <p>where an example of _aux would be :   </p>
  
  <pre><code>_aux = {'const': {'g':-9.81} }
  </code></pre>
  
  <p>_t and _p are not used in this case. It expects _X to be a 7-element vector. 
  In this case, the function only returns one value. But there may be cases in which it returns a list. In that case, the returned value forms a column in the output. Hope that clarifies stuff a bit.</p>