I have few variables x, y, z
# Variables
x = np.linspace(-1.0, 1.0, num=5, endpoint=True)
y = np.linspace(-1.0, 1.0, num=5, endpoint=True)
z = np.linspace(-1.0, 1.0, num=5, endpoint=True)
# And the corresponding meshgrid
x_v, y_v, z_v = np.meshgrid(x, y, z)
Let's say I want to calculate a value of a quadratic polynomial with the coefficients C at each point of a meshgrid. For this I have a function to return a polynomial terms for a set of arguments (NOTE: here the function is generic to make it work for any number of arguments):
def pol_terms(*args):
# -- Calculate terms of the polynomial --
# Linear part (b0, b11*x1, b12*x2, ..., b1k*xk, where k=len(args))
entries = [1] + [pt for pt in args]
# Combinations part (b12*x12, b13*x13, ..., bk-1k*xk-1k)
n_args = len(args)
for i in range(n_args):
for j in range(i+1, n_args):
entries.append(args[i]*args[j])
# Quadratic part
entries += [pt**2 for pt in args]
return np.array([entries])
So, for one set of variables' values I can get the value of my polynomial as follows:
X = pol_terms(1,2,3)
X.dot(C)
I can't figure out how to make the same for the meshgrid x_v, y_v, z_v. Applying np.vectorize on the pol_terms wouldn't work because it returns an array. I don't want manually looping through all the dimensions neither because I want the solution to be generic (so it works for 2 variable as well as for 5).
Thanks in advance!
I agree you can't use np.vectorize, but not for the reason you've given.
np.vectorize works well with functions that return arrays. Just, you must say it.
ptv = np.vectorize(pol_terms, signature='(),(),()->(m)')
ptv([1,10],[2,20],[3,30])
returns
array([[ 1, 1, 2, 3, 2, 3, 6, 1, 4, 9],
[ 1, 10, 20, 30, 200, 300, 600, 100, 400, 900]])
The reasons why you can't use vectorize are, firstly, that it doesn't do much better than just looping yourself (vectorize is not for performance, as stated in documentation). And secondly, the one you can see in my example: it needs a constant number of arguments. At least if you want to specify a signature
But your code is mostly already vectorized. All what you do is additions and multiplications.
The only problem that I can see that would occur if using arrays (I mean non-0 dimension array, since scalar, in numpy, are just array with dimension 0) is the 1 term. That term implies that all other terms are of the same dimension (so dimension 0).
But that is easily solved.
Replace line
entries = [1] + [pt for pt in args]
With line
entries = [np.ones_like(args[0])] + [pt for pt in args]
(I assume there is at least one arg. If not, I let you add the test needed for the special "no arg" case)
Then
pol_terms(np.array([1,10]),np.array([2,20]),np.array([3,30]))
Returns
array([[[ 1, 1],
[ 1, 10],
[ 2, 20],
[ 3, 30],
[ 2, 200],
[ 3, 300],
[ 6, 600],
[ 1, 100],
[ 4, 400],
[ 9, 900]]])
(Note, that the extra bracket is due to your returning np.array([entries]) and not np.array(entries). I don't know if this is willingly. But that means that whatever the shape on entries, shape of returned value is 1×shape. So there is always a size-1 axis 0 in your return)
I won't copy the result of pol_terms(x_v,y_v,z_v). But (to that strange 1st axis) they seem to be what is expected to me.
For example
pt=pol_terms(x_v,y_v,z_v)
c123=pt[0,:,1,2,3]
are the terms at position (x[1],x[2],x[3]). The 0 coord is because of that extra axis I've mentionned.
So, unless I've not understood what you're trying to do, there is almost nothing to do. Your function is already able to work on vectors. But for that hard-coded scalar 1.
def pol_terms(*args):
entries = np.array([np.ones_like(args[0])] + [pt for pt in args])
return (entries[None,...]*entries[:,None,...])[np.triu_indices(len(entries))]
But it is just more compact, not faster. Your way already vectorized what really matters (the axes of the meshgrid). What this add, is just vectorization of the outest axis, the one of args itself (your 2 nested for loops, that just iterates over arguments, not over axis of those arguments)
So, for a meaningless vectorization, I pay a price: I compute both x*y and y*x, and then drop almost half of the computation (keep only x*y and forget about y*x since it is the same)
So, asympotically, this version should be about twice slower than yours (with proposed correction). It is just for the fun of the quasi-one-liner.