Best way of combining meshgrid with matrix multiplication in function

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)

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

quadraction function formula

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)

This is interpreted by Einstein convention as:

1/2 sum_ij vecx_i,... A_ij vecx_j,... - sum_i vecx_i,... b_i
Best way of combining meshgrid with matrix multiplication in function
See more ...