在numpy中初始化向量场

Question

我想创建一个numpy数组，用来表示一个二维的向量场，这个向量场是在一个100 x 100的点网格上定义的：

import numpy as np
dx = dy = 0.1
nx = ny = 100
x, y =  np.meshgrid(np.arange(0,nx*dx,dx), np.arange(0,ny*dy,dy))

这个向量场是围绕点(cx, cy)以恒定速度旋转的，我可以用普通的Python循环来初始化它：

v = np.empty((nx, ny, 2))
cx, cy = 5, 5
s = 2
for i in range(nx):
    for j in range(ny):
        rx, ry = i*dx - cx, j*dy - cy
        r = np.hypot(rx, ry)
        if r == 0:
            v[i,j] = 0,0
            continue
        # (-ry/r, rx/r): the unit vector tangent to the circle centred at (cx,cy), radius r
        v[i,j] = (s * -ry/r, s * rx/r)

但是在用numpy进行向量化时遇到了问题。我目前得到的最接近的结果是：

v = np.array([s * -(y-cy) / np.hypot(x-cx, y-cy), s * (x-cx) / np.hypot(x-cx, y-cy)])
v = np.rollaxis(v, 1, 0)
v = np.rollaxis(v, 2, 1)
v[np.isinf(v)] = 0

但这个结果不对，跟我想要的结果不一样。用numpy初始化向量场的正确方法是什么呢？

补充说明：好吧，现在我有点困惑，按照下面的建议，我尝试了：

vx = s * -(y-cy) / np.hypot(x-cx, y-cy)
vy = s * (x-cx) / np.hypot(x-cx, y-cy)
v = np.dstack((vx, vy))
v[np.isnan(v)] = 0

但得到的数组完全不同……

Answer 1

根据你最初的设置：

import numpy as np

dx = dy = 0.1
nx = ny = 100
x, y = np.meshgrid(np.arange(0, nx * dx, dx),
                   np.arange(0, ny * dy, dy))
cx = cy = 5
s = 2

你可以这样计算 v：

rx, ry = y - cx, x - cy
r = np.hypot(rx, ry)
v2 = s * np.dstack((-ry, rx)) / r[..., None]
v2[np.isnan(v2)] = 0

如果你想要更高级一点，可以把 yx 创建成一个三维数组，然后对它进行所有操作：

# we make these [2,] arrays to broadcast over the last output dimension
c = np.array([5, 5])
s = np.array([-2, 2])

# this creates a [100, 100, 2] mesh, where the last dimension corresponds
# to (y, x)
yx = np.mgrid[0:nx * dx:dx, 0:ny * dy:dy].T

yxdiff = yx - c[None, None, :]
r = np.hypot(yxdiff[..., 0], yxdiff[..., 1])[..., None]
v3 = s[None, None, :] * yxdiff / r
v3[np.isnan(v3)] = 0

检查一下这两种方法是否和你原来的代码得到的结果一样：

print np.all(v == v2), np.all(v == v3)
# True, True

编辑

为什么用 rx, ry = y - cx, x - cy 而不是 rx, ry = x - cx, y - cy？我也觉得这很不直观——我之所以这样做只是为了和你原来的代码的输出保持一致。

问题在于，在你的网格中，连续的 x 值实际上是在 x 的连续 列 中，而连续的 y 值是在 y 的连续 行 中，也就是说 x[:, j] 是第 j 个 x 值，而 y[i, :] 是第 i 个 y 值。然而，在你的内部循环中，你把 dx 乘以 i（这是你的 行 索引），把 dy 乘以 j（这是你的 列 索引）。所以你实际上是把输出的 x 和 y 维度搞反了。