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

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,) 

如果我不按Shift+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)

- Julia

---

你尝试使用NumPy和NumPy.linalg的向量操作来精确复制该公式吗？ - Lutz Lehmann

这个不起作用的伪代码正试图使用那些函数（dot和abs）。 - Julia

@LutzL：我看到了问题，但我不知道如何规避它。我无法想象任何熟练使用matplotlib的用户都会像我在示例中所做的那样，通过编写所有坐标来输入向量场的公式。一定有一种方法可以使其更易读、易写、减少错误，并且更符合数学角度的要求。 - Julia

我需要阅读文档，但我认为在numpy中有一个vectorize函数，可以完美地实现这些目的。 - Lutz Lehmann

不要使用 np.vectorize 来复杂化问题。它适用于接受标量值的函数。B_diapole 已经可以处理 X 和 Y 数组了。 - hpaulj

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 进行计算：

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)

- hpaulj

---

很有趣的是，看看如何使用 sympy 简化我想要的语法。 - Julia

Answer 2

我使用简单的计算机代数系统简化了你的表达式。

--- 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

- gboffi