Question

如果我在matplotlib中绘制一个矢量场，我通常会明确地为每个组件写下公式，以避免问题，例如形状和广播。然而，在稍微复杂的公式中，代码变得混乱，无论是写还是读。

考虑以下示例，我想绘制由此公式定义的矢量字段：

有没有方便的方法来输入更符合向量操作的数学公式，如下面我的（不工作）伪代码？

# Run with ipython3 notebook
%matplotlib inline
from pylab import *

## The following works, but the mathematical formula is a complete mess to red
def B_dipole(m, a, x,y):
    return (3*(x - a[0])*(m[0]*(x - a[0]) + m[1]*(y-a[1]))/((x - a[0])**2 + (y-a[1])**2)**(5/2.0) -m[0]/((x - a[0])**2 + (y-a[1])**2)**(3/2.0),3*(y - a[1])*(m[0]*(x - a[0]) + m[1]*(y-a[1]))/((x - a[0])**2 + (y-a[1])**2)**(5/2.0) -m[1]/((x - a[0])**2 + (y-a[1])**2)**(3/2.0))

## I want something like (but doesn't work)
#def B_dipole(m, a, x,y):
#    r = array([x,y])
#    rs = r - a ## shifted r
#    mrs = dot(m,rs) ## dot product of m and rs
#    RS = dot(rs,rs)**(0.5) ## euclidian norm of rs
#    ret = 3*mrs*rs/RS**5 - m/RS**3 ## vector/array to return
#    return ret

x0, x1=-10,10
y0, y1=-10,10

X=linspace(x0,x1,55)
Y=linspace(y0,y1,55)
X,Y=meshgrid(X, Y)

m = [1,2]
a = [3,4]

Bx,By = B_dipole(m,a,X,Y)

fig = figure(figsize=(10,10))
ax = fig.add_subplot(1, 1, 1)
ax.streamplot(X, Y, Bx, By,color='black',linewidth=1,density=2)
#ax.quiver(X,Y,Bx,By,color='black',minshaft=2)
show()

输出：

修改我的非工作代码的错误消息：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-2-43b4694cc590> in <module>()
     26 a = [3,4]
     27 
---> 28 Bx,By = B_dipole(m,a,X,Y)
     29 
     30 fig = figure(figsize=(10,10))

<ipython-input-2-43b4694cc590> in B_dipole(m, a, x, y)
     10 def B_dipole(m, a, x,y):
     11     r = array([x,y])
---> 12     rs = r - a ## shifted r
     13     mrs = dot(m,rs) ## dot product of m and rs
     14     RS = dot(rs,rs)**0.5 ## euclidian norm of rs

ValueError: operands could not be broadcast together with shapes (2,55,55) (2,)

如果我不移动r：

，则会显示错误消息

--
ValueError                                Traceback (most recent call last)
<ipython-input-4-e0a352fa4178> in <module>()
     23 a = [3,4]
     24 
---> 25 Bx,By = B_dipole(m,a,X,Y)
     26 
     27 fig = figure(figsize=(10,10))

<ipython-input-4-e0a352fa4178> in B_dipole(m, a, x, y)
      8     r = array([x,y])
      9     rs = r# - a ## not shifted r
---> 10     mrs = dot(m,rs) ## dot product of m and rs
     11     RS = dot(rs,rs)**0.5 ## euclidian norm of rs
     12     ret = 3*mrs*rs/RS**5 - m/RS**3 ## vector/array to return

ValueError: shapes (2,) and (2,55,55) not aligned: 2 (dim 0) != 55 (dim 1)

Answer 1

我想我应该从你的公式开始，但我会尝试更紧凑地表达工作B_dipole：

def B_dipole(m, a, x,y):
    return (3*(x - a[0])*(m[0]*(x - a[0]) + m[1]*(y-a[1]))/((x - a[0])**2 + (y-a[1])**2)**(5/2.0) -m[0]/((x - a[0])**2 + (y-a[1])**2)**(3/2.0),3*(y - a[1])*(m[0]*(x - a[0]) + m[1]*(y-a[1]))/((x - a[0])**2 + (y-a[1])**2)**(5/2.0) -m[1]/((x - a[0])**2 + (y-a[1])**2)**(3/2.0))

def B_dipole(m, a, x,y):
    x0 = x - a[0]
    y1 = y - a[1]
    return (3*x0*(m[0]*x0 + m[1]*y1)/(x0**2 + y1**2)**(5/2.0) -m[0]/(x0**2 + y1**2)**(3/2.0),3*y1*(m[0]*x0 + m[1]*y1)/(x0**2 + y1**2)**(5/2.0) -m[1]/(x0**2 + y1**2)**(3/2.0))

我可能删除了太多（）。但我看到了其他重复的模式，例如。

(x0**2 + y1**2)
(m[0]*x0 + m[1]*y1)

sympy可能是将公式转换为numpy表达式的有用工具。我自己并没有太多使用它，但已经帮助了一些SO问题。

 r_abs = np.sqrt(x0**2 + y1**2))
 mr = m[0]*x0 + m[1]*y1

 (3*x0*(mr)/(r_abs)**(5) -m[0]/(r_abs)**(3), 3*y1*(mr)/(r_abs)**(5) -m[1]/(r_abs)**(3))

但是让我们用数组来表达这个：

In [21]: m = np.array([1,2]); a = np.array([3,4])

In [45]: X,Y = np.meshgrid(x,y,indexing='xy')
In [46]: X0 = X-a[0]; Y1 = Y-a[1]
In [47]: r_abs = (X0**2 + Y1**2)**.5
In [48]: mr = m[0]*X0 + m[1]*Y1
In [49]: Bx = 3*X0*mr/r_abs**5 - m[0]/r_abs**3
In [50]: By = 3*Y1*mr/r_abs**5 - m[1]/r_abs**3
In [51]: pyplot.streamplot(X,Y,By,Bx)

与你的情节相同。

让我们尝试将X和Y合并为一个数组并使用dot s：

In [52]: XY=np.stack([X,Y])
In [53]: XY.shape
Out[53]: (2, 55, 55)
In [54]: XYa = XY - a[:,None,None]
# dot doesn't work with 3d array; use einsum instea
In [55]: mr = np.dot(m,XYa)
...
ValueError: shapes (2,) and (2,55,55) not aligned: 2 (dim 0) != 55 (dim 1)
In [71]: mr = np.einsum('i,i...',m,XYa)
In [72]: r_abs = (XYa**2).sum(axis=0)**.5
In [73]: B = 3*XYa*mr/r_abs**5 - m[:,None,None]/r_abs**3
In [74]: B.shape
Out[74]: (2, 55, 55)
In [75]: pyplot.streamplot(XY[0],XY[1],B[0],B[1])
Out[75]: <matplotlib.streamplot.StreamplotSet at 0xab71feac>

可以将X和Y变量组合到更高维数组中，从而将计算减少到R2向量计算，但我不确定它是否会使事情变得更简单。

同一件事的复杂版本：

In [76]: XYj=X+1j*Y
In [77]: XYja = XYj-(3+4j)
In [98]: r_abs = np.abs(XYja)
In [103]: m_r = (XYja*(1-2j)).real   # right values, but?
In [107]: Ba = 3*XYja*m_r/r_abs**5 - (1+2j)/r_abs**3
In [108]: pyplot.streamplot(XYj.real,XYj.imag,Ba.real,Ba.imag)

Answer 2

我使用简单的CAS

简化了你的表达

--- Emacs Calculator Mode ---
    3 (m0*(x - a0) + m1*(y - a1)) (x - a0)               m0                  3 (m0*(x - a0) + m1*(y - a1)) (y - a1)               m1
4:  -------------------------------------- - -------------------------- + i*(-------------------------------------- - --------------------------)
                   2           2 2.5                  2           2 1.5                     2           2 2.5                  2           2 1.5
          ((x - a0)  + (y - a1) )            ((x - a0)  + (y - a1) )               ((x - a0)  + (y - a1) )            ((x - a0)  + (y - a1) )

3:  [X = x - a0, Y = y - a1]

    3 X*(X m0 + Y m1)        m0           3 Y*(X m0 + Y m1)        m1
2:  ----------------- - ------------ + i*(----------------- - ------------)
        2    2 2.5        2    2 1.5          2    2 2.5        2    2 1.5
      (X  + Y )         (X  + Y )           (X  + Y )         (X  + Y )

    3 X*(X m0 + Y m1)   m0       3 Y*(X m0 + Y m1)   m1
1:  ----------------- - --- + i*(----------------- - ---)
            5.           3.              5.           3.
           R            R               R            R

我将该字段的两个组成部分表示为复数的实部和虚部。

从最后一个表达开始，可能是写

x, y = np.meshgrid(...)
X, Y = x-a[0], y-a[1]
R = np.sqrt(X*X+Y*Y)
H = X*m[0]+Y*m[1]
Fx = 3*X*H/R**5-m[0]/R**3
Fy = 3*Y*H/R**5-m[1]/R**3

如何以快速数学方式输入公式来绘制矢量场？

2 个答案: