python - 网格索引和实际值之间的转换

Question

我有一个形状为(A1, M1, A2, M2, A3, M3, E)的网格，我使用它生成

A1, M1, A2, M2, A3, M3, E = meshgrid(Grid.aGrid, Grid.mGrid, Grid.aGrid, Grid.mGrid, Grid.aGrid, Grid.mGrid, Grid.eGrid, indexing='ij')

，其中 Grid.aGrid 是使用 linspace(aMin, aMax, nA) 生成的，其他网格也类似。

考虑一些 Z = f(A1, ...)，其中 f() 会将某些网格点标记为不相关。为了简单起见，就这样吧

Z = A1 + A2 + A3
Z[Z < 0] = NaN

考虑Z[0, 1, 2, 3, 4, 5, 6]。它包含与实际值相对应的值(aGrid[0], mGrid[1], aGrid[2], mGrid[3], aGrid[4], mGrid[5], eGrid[6]).这正是我尝试对 Z 上未用 f() 标记的所有点实现的目标:

我想创建一本字典

foo = {z1, z2, z3, ... zn}

其中 z1 等都是此类

z1 = (aGrid[0], mGrid[1],  aGrid[2], mGrid[3], aGrid[4], mGrid[5], eGrid[6])

，即Z内z1位置对应的网格值。

我想出了一些东西:

aGrid = arange(0, 10)
mGrid = arange(100, 110)
eGrid = arange(1000, 1200)

A,M,E = meshgrid(aGrid, mGrid, eGrid, indexing='ij')

# contains the grid index
Z = (A + M + E).astype(float)
Z[A < 3] = nan
# will contain the actual values, as tuples
Z2 = {}

for i, idx in enumerate(ndindex(Z.shape)):
    a = aGrid[idx[0]]
    m = mGrid[idx[1]]
    e = eGrid[idx[2]]

    if isnan(Z[idx]):
        Z2[i] = NaN
    else:
        Z2[i] = (a, m, e)

效率是关键。有没有更快/更干净的方法可以实现这一目标？除了使用字典之外还有其他方法吗？

我特别不喜欢我必须写下aGrid[idx[0]]等。是否可以使算法更通用？一些类似的事情

for i, idx in enumerate(ndindex(Z.shape)):
    # some magic happens here. What exactly?
    someMagicList = magic(aGrid, mGrid, eGrid)
    Z2[i] = someMagicList[idx]

Answer 1

使用broadcast_arrays()，结果Z2是一个形状为(20000, 3)的数组。

import numpy as np

aGrid = np.arange(0, 10, dtype=float)
mGrid = np.arange(100, 110, dtype=float)
eGrid = np.arange(1000, 1200, dtype=float)

A,M,E = np.meshgrid(aGrid, mGrid, eGrid, indexing='ij')

# contains the grid index
Z = (A + M + E).astype(float)
Z[A < 3] = np.nan

grids = [A, M, E]
grid_bc = np.broadcast_arrays(*grids)
Z2 = np.column_stack([g.ravel() for g in grid_bc])
Z2[np.isnan(Z.ravel())] = np.nan

print Z2[5900], Z2[6000]