使用meshgrid和vectorize评估未知维度的晶格上的函数

时间:2014-08-26 18:13:16

标签: python numpy

如果您提前了解晶格的尺寸数,则可以直接使用meshgrid来评估网格上的函数。

from pylab import *

lattice_points = linspace(0,3,4)
xs,ys = meshgrid(lattice_points,lattice_points)
zs = xs+ys # <- stand-in function, to be replaced by something more interesting
print(zs)

可生产

[[ 0.  1.  2.  3.]
 [ 1.  2.  3.  4.]
 [ 2.  3.  4.  5.]
 [ 3.  4.  5.  6.]]

但是我希望有一个类似的版本,维度在运行时确定,或者作为参数传递。

from pylab import *

@np.vectorize
def fn(listOfVars) :
    return sum(listOfVars) # <- stand-in function, to be replaced 
                           # by something more interesting

n_vars = 2
lattice_points = linspace(0,3,4)
indices = meshgrid(*(n_vars*[lattice_points])) # this works fine

zs = fn(indices) # <-- this line is wrong, but I don't 
                 #     know what would work instead
print(zs)

制作

[[[ 0.  1.  2.  3.]
  [ 0.  1.  2.  3.]
  [ 0.  1.  2.  3.]
  [ 0.  1.  2.  3.]]

 [[ 0.  0.  0.  0.]
  [ 1.  1.  1.  1.]
  [ 2.  2.  2.  2.]
  [ 3.  3.  3.  3.]]]

但我希望它能产生与上面相同的结果。

可能有一个解决方案,你可以找到每个维度的索引,并使用itertools.product生成所有可能的索引等组合,但有没有一个很好的pythonic方式这样做?

1 个答案:

答案 0 :(得分:2)

Joe Kington和user2357112帮助我以我的方式看到错误。对于那些希望看到完整解决方案的人:

from pylab import *

## 2D "preknown case" (for testing / to compare output)
lattice_points = linspace(0,3,4)
xs,ys = meshgrid(lattice_points,lattice_points)
zs = xs+ys
print('2-D Case')
print(zs)

## 3D "preknown case" (for testing / to compare output)
lattice_points = linspace(0,3,4)
ws,xs,ys = meshgrid(lattice_points,lattice_points,lattice_points)
zs = ws+xs+ys
print('3-D Case')
print(zs)

## Solution, thanks to comments from Joe Kington and user2357112 
def fn(listOfVars) :
    return sum(listOfVars)

n_vars = 3 ## can change to 2 or 3 to compare to example cases above
lattice_points = linspace(0,3,4)
indices = meshgrid(*(n_vars*[lattice_points]))

zs = np.apply_along_axis(fn,0,indices)
print('adaptable n-D Case')
print(zs)