python - Function of meshgrid numpy

Question

I have a problem where I want to plot the determinant of a matrix as a function of parameters, so I have a script

def f(x, y):

    DD = np.matrix([[0., 0.],[0., 0.]]) + 0.j
    omega = x + 1.j * y

    # set up dispersion matrix
    DD[0,0] = 1 + omega
    DD[1,0] = omega
    DD[0,1] = omega
    DD[1,1] = 1 - omega

    metric = np.linalg.det(DD)

    return metric

xx = np.arange(1., 2., 0.1)
yy = np.arange(1., 2., 0.1)

x, y = np.meshgrid(xx, yy)


FPlot = f(x, y)

plt.contourf(x, y, FPlot)
plt.show()

There is a type error where, since omega is a meshgrid, I can't stick it into the matrix or compute the determinant -- numpy demands scalars in matrices. What's the best way around this to get the right mesh and evaluated determinants?

Answer 1

如果您拥有全新的 numpy 1.8，您可以使用其中一项很酷的新功能：线性代数 gufuncs。基本上，当您调用 时np.linalg.det，它将计算您传入的任何形状数组的最后两个维度的行列式：

xx = np.arange(1., 2., 0.1)
yy = np.arange(1., 2., 0.1)

x, y = np.meshgrid(xx, yy)
omega = x + 1.j * y

dispersion_matrices = np.empty(omega.shape + (2, 2), dtype=omega.dtype)
dispersion_matrices[..., 0, 0] = 1 + omega
dispersion_matrices[..., 1, 0] = omega
dispersion_matrices[..., 0, 1] = omega
dispersion_matrices[..., 1, 1] = 1 - omega
FPlot = np.linalg.det(dispersion_matrices)

plt.contourf(x, y, FPlot.imag) # the part missing in unutbus's plot!
plt.show()

Answer 2

你可以使用np.frompyfunc：

import numpy as np
import matplotlib.pyplot as plt

def f(x, y):

    DD = np.matrix([[0., 0.],[0., 0.]]) + 0.j
    omega = x + 1.j * y

    # set up dispersion matrix
    DD[0,0] = 1 + omega
    DD[1,0] = omega
    DD[0,1] = omega
    DD[1,1] = 1 - omega

    metric = np.linalg.det(DD)

    return metric
f = np.frompyfunc(f, 2, 1)

xx = np.arange(1., 2., 0.1)
yy = np.arange(1., 2., 0.1)
x, y = np.meshgrid(xx, yy)

FPlot = f(x, y)

plt.contourf(x, y, FPlot)  # Note that this is only using the real part of FPlot
plt.show()