python - 如何以方便的数学方式输入公式来绘制矢量场？

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

我已经使用简单的 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