2019-09-21
Best way of combining meshgrid with matrix multiplication in function
stackoverflow
Question

With `np.meshgrid` numpy provides a convenient way of plotting functions of two variables, e.g. like so:

``````def plot():
fig = plt.figure()
ax = plt.axes(projection="3d")

x = np.linspace(-6, 6, 30)
y = np.linspace(-6, 6, 30)

X, Y = np.meshgrid(x, y)
Z = f(X, Y)

ax.plot_surface(X, Y, Z)
plt.show()
``````

Unfortunately, this easy setup interferes with me defining any quadratic function

in the most natural numpy way, i.e. like so

``````def f(x, y):
vec_x = np.array([x, y])
return 1/2 * np.dot(vec_x.T, np.dot(A, vec_x)) - np.dot(b.T, vec_x)
``````

The problem is that the `meshgrid` arrays `X` and `Y` in `plot()` when calling `Z = f(X, Y)` will now be processed as arrays by the line `vec_x = np.array([x, y])` which results in `vec_x` being a `(2, 30, 30)` shape array instead of an entry-by-entry treatment which would give a `(2,)` shape array which is what I would want. Compare this to

``````def other_f(x, y):
return x + y
``````

which works perfectly in a natural way with numpy thanks to the vectorization.

I haven't used numpy and matplotlib for a while but I would really and all workarounds that I come up with feel super clumsy, so, I'd love to see a neat way to work around this.

You could use `einsum` like so:
``````np.einsum('i...,ij,j...',vec_x,A,vec_x)/2 - np.einsum('i...,i',vec_x,b)
``````1/2 sum_ij vecx_i,... A_ij vecx_j,... - sum_i vecx_i,... b_i