绘制向量函数的最 Pythonic 方式
Most Pythonic way to plot a vector function
我有一个函数 calcField
,当给定一个带有两个元素代表位置的 numpy 数组时,returns 一个数组代表该位置的电场。要求 matplotlib 为此函数绘制矢量场的最 pythonic 方法是什么?目前我可以使用这段代码,但它感觉违背了 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)
我发现它比这个替代方案更容易理解,您可能会发现它更具可读性(这是您的问题):
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([])
但是,此时我不再使用 类。有时为了更轻松地访问矢量化是值得的。在代码的第二种形式中,您基本上在 elec_fields
的第二个轴中得到每个 PointCharge.calcField()
的结果,第一个只是这些字段的 x 和 y 分量。
我有一个函数 calcField
,当给定一个带有两个元素代表位置的 numpy 数组时,returns 一个数组代表该位置的电场。要求 matplotlib 为此函数绘制矢量场的最 pythonic 方法是什么?目前我可以使用这段代码,但它感觉违背了 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)
我发现它比这个替代方案更容易理解,您可能会发现它更具可读性(这是您的问题):
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([])
但是,此时我不再使用 类。有时为了更轻松地访问矢量化是值得的。在代码的第二种形式中,您基本上在 elec_fields
的第二个轴中得到每个 PointCharge.calcField()
的结果,第一个只是这些字段的 x 和 y 分量。