绘制向量函数的最具python风格的方法

绘制向量函数的最具python风格的方法,python,numpy,matplotlib,Python,Numpy,Matplotlib,我有一个函数calcField,当给定一个包含两个表示位置的元素的numpy数组时,它返回一个表示该位置电场的数组。要求matplotlib为此函数绘制向量场的最简单方式是什么?目前,我已经有了这段代码,但它感觉与numpy的精神背道而驰,并且相对不可读 Y, X = np.mgrid[-3:3:100j, -3:3:100j] vectors = np.array([[field.calcField(r) for r in row] for row in

我有一个函数
calcField
,当给定一个包含两个表示位置的元素的numpy数组时,它返回一个表示该位置电场的数组。要求matplotlib为此函数绘制向量场的最简单方式是什么?目前,我已经有了这段代码,但它感觉与numpy的精神背道而驰,并且相对不可读

Y, X = np.mgrid[-3:3:100j, -3:3:100j]

vectors = np.array([[field.calcField(r) for r in row] 
                  for row in [zip(a, b) for a, b in zip(X, Y)]])
U = np.array([[vector[0] for vector in row] for row in vectors])
V = np.array([[vector[1] for vector in row] for row in vectors])

plt.streamplot(X, Y, U, V, color=U, linewidth=2, cmap=plt.cm.autumn)
编辑:根据要求,calcField的代码:

import constants
import numpy as np
import numpy.linalg as l
class Field:
    def __init__(self, charges = []):
       self.charges = charges
    def addCharge(self, charge):
        self.charges = self.charges + [charge]
    def calcField(self, point):
        point = np.array(point)
        return sum([charge.calcField(point) for charge in self.charges])

class PointCharge:
    def __init__(self, q, position):
        self.q = q
        self.position = np.array(position)
    def calcField(self, point):
        return constants.k * self.q * (point - self.position) / l.norm (point - self.position)**3

使用流线绘制一组点电荷电场的矢量化代码可能如下所示:

num_charges = 4
charges = np.random.random_integers(-5,5,num_charges)
charges[charges==0] = 5
charges_positions = np.random.random((num_charges, 2))

y,x = np.mgrid[0:1:40j, 0:1:40j]
xdist = x - charges_positions[:,0].reshape(-1,1,1)
ydist = y - charges_positions[:,1].reshape(-1,1,1)

denom = ((xdist**2 + ydist**2)**1.5)
# Ignoring Coulomb's constant here...
Ex = (charges.reshape(-1,1,1) * xdist / denom).sum(axis=0)
Ey = (charges.reshape(-1,1,1) * ydist / denom).sum(axis=0)
我觉得这比这个选择更容易理解,你可能会觉得更可读(这是你的问题):

你可以很容易地画出来。我将继续使用上一个代码块中的格式(如果使用第一个表单,则需要交换一些索引):

pos\u费用=费用>0
负电荷=电荷<0
f、 (ax,ax1)=plt.子批次(1,2)
ax.绘图(位置[0,位置费用],[1,位置费用],“ro”)
ax.图(位置[0,负电荷],[1,负电荷],“bo”)
ax.streamplot(x,y,Ex,Ey,color='k')
ax.set_aspect('equal',可调='box'))
ax.设置标题(“电场”)
ax.set_xticks([])
ax.set_-yticks([])
然而,现在我不再使用类了。有时,享受更轻松的矢量化是值得的。在代码的第二种形式中,基本上在
电气场的第二个轴上有每个
PointCharge.calcField()
的结果,第一个只是这些场的x和y分量


尝试矢量化函数
calcField
。一旦你掌握了这个窍门,你就可以在这个代码摘录中去掉所有其他for循环。为了补充奥利弗所说的,你可以编写
U,V=calcField(X,Y)
。只有在值得付出努力的情况下才这样做,即此代码非常慢,或者您知道您将来将经常使用
calcField()
。我已经养成了将N空间中的向量表示为长度N numpy数组的习惯,因此,我不必跟踪单个组件,而只能处理向量操作。例如,点电荷场的表达式就是
k*q*(pos1-pos2)/np.linalg.norm(pos1-pos2)**3
。在同时对多个输入进行“矢量化”时,有没有办法保持数学的可读性?我们可以看看
calcField
num_charges = 4
charges = np.random.random_integers(-5,5,(num_charges,1,1))
charges[charges==0] = 5  # only for clarity
positions = np.random.random((2, num_charges,1,1))

y,x = np.mgrid[0:1:40j, 0:1:40j]
M,N = y.shape
xy = np.array([x,y]).reshape(2,1, M,N)
rad_dist = xy - positions
denom = np.linalg.norm(rad_dist, axis=(0))**3

elec_fields = charges * rad_dist / denom
Ex, Ey = elec_fields.sum(axis=1)
pos_charges = charges > 0
neg_charges = charges < 0
f,(ax,ax1) = plt.subplots(1,2)
ax.plot(positions[0, pos_charges], positions[1, pos_charges], 'ro ')
ax.plot(positions[0, neg_charges], positions[1, neg_charges], 'bo ')
ax.streamplot(x,y, Ex, Ey, color='k')
ax.set_aspect('equal', adjustable='box')
ax.set_title('Electric field')
ax.set_xticks([])
ax.set_yticks([])