初始化numpy中的矢量场

时间:2014-08-13 02:05:51

标签: python arrays numpy vector vectorization

我想初始化一个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

但得到一个完全不同的数组...

1 个答案:

答案 0 :(得分:0)

从初始设置开始:

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创建为3D数组,并通过它广播所有操作:

# 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, :]是< em> i y值。但是,在内部循环中,您将dx乘以i,这是索引,而dy乘以j,这是{{1}}您的索引。因此,您正在翻转输出的x和y维度。