python - 基于数组的 Numpy 3d 数组赋值

Question

取一个 2D numpy.array，假设:

mat = numpy.random.rand(3,3)

In [153]: mat
Out[153]: 
array([[ 0.16716156,  0.90822617,  0.83888038],
       [ 0.89771815,  0.62627978,  0.34992542],
       [ 0.11097042,  0.80858005,  0.0437299 ]])

将索引更改为 numpy.nan 非常简单

以下其中一项效果很好:

In [154]: diag = numpy.diag_indices(mat.shape[0], ndim = 2)
In [155]: mat[diag] = numpy.nan

或

In [156]: numpy.fill_diagonal(mat, numpy.nan)

但假设我有一个 3D 数组，我希望在第 3 维的每个维度上完全相同的过程。

mat = numpy.random.rand(3, 5, 5)

In [158]: mat

 Out[158]: 
 array([[[ 0.65000325,  0.71059547,  0.31880388,  0.24818623,  0.57722849],
         [ 0.26908326,  0.41962004,  0.78642476,  0.25711662,  0.8662998 ],
         [ 0.15332566,  0.12633147,  0.54032977,  0.17322095,  0.17210078],
         [ 0.81952873,  0.20751669,  0.73514815,  0.00884358,  0.89222687],
         [ 0.62775839,  0.53657471,  0.99611842,  0.75051645,  0.59328044]],

        [[ 0.28718216,  0.84982865,  0.27830082,  0.90604492,  0.43119512],
         [ 0.43039373,  0.76557782,  0.58089787,  0.81135684,  0.39151152],
         [ 0.70592711,  0.30625204,  0.9753166 ,  0.32806864,  0.21947731],
         [ 0.74600317,  0.33711673,  0.16203076,  0.6002213 ,  0.74996638],
         [ 0.63555715,  0.71719058,  0.81420001,  0.28968442,  0.01368163]],

        [[ 0.06474027,  0.51966572,  0.006429  ,  0.98590784,  0.35708074],
         [ 0.44977222,  0.63719921,  0.88325451,  0.53820139,  0.51526687],
         [ 0.98529117,  0.46219441,  0.09349748,  0.11406291,  0.47697128],
         [ 0.77446136,  0.87423445,  0.71810465,  0.39019846,  0.94070077],
         [ 0.09154989,  0.36295161,  0.19740833,  0.17803146,  0.6498038 ]]])

一个合乎逻辑的方法(我认为)是:

mat[:, diag] = numpy.nan  # doesn't do it

事实上，要完成这个，我需要:

In [190]: rng = numpy.arange(5)
In [191]: for i in numpy.arange(mat.shape[0]):
   .....:     mat[i, rng, rng] = numpy.nan
   .....:     

In [192]: mat
Out[192]: 
array([[[        nan,  0.4040426 ,  0.89449522,  0.63593736,  0.94922036],
        [ 0.40682651,         nan,  0.30812181,  0.01726625,  0.75655994],
        [ 0.23925763,  0.41476223,         nan,  0.91590111,  0.18391644],
        [ 0.99784977,  0.71636554,  0.21252766,         nan,  0.24195636],
        [ 0.41137357,  0.84705055,  0.60086461,  0.16403918,         nan]],

       [[        nan,  0.26183712,  0.77621913,  0.5479058 ,  0.17142263],
        [ 0.17969373,         nan,  0.89742863,  0.65698339,  0.95817106],
        [ 0.79048886,  0.16365168,         nan,  0.97394435,  0.80612441],
        [ 0.94169129,  0.10895737,  0.92614597,         nan,  0.08689534],
        [ 0.20324943,  0.91402716,  0.23112819,  0.2556875 ,         nan]],

       [[        nan,  0.43177039,  0.76901587,  0.82069345,  0.64351534],
        [ 0.14148584,         nan,  0.35820379,  0.17434688,  0.78884305],
        [ 0.85232784,  0.93526843,         nan,  0.80981366,  0.57326785],
        [ 0.82104636,  0.63453196,  0.5872653 ,         nan,  0.96214559],
        [ 0.69959383,  0.70257404,  0.92471502,  0.50077728,         nan]]])

它适用于速度至关重要的应用程序，因此如果没有基于数组的以下实现，我将在 Cython 中执行 for 循环/赋值

Answer 1

这似乎可行:

diag = numpy.diag_indices(mat.shape[1], ndim = 2)
mat[:, diag[0], diag[1]] = numpy.nan

问题是 diag 是一个 2 元素元组，因此在 3D 索引中按原样使用它是行不通的，而且不幸的是，使用 *diag 我们的语法无效。但是，您也可以这样做:

diag = (Ellipsis, *numpy.diag_indices(mat.shape[-1], ndim = 2))
mat[diag] = numpy.nan

在这种情况下，diag 是您需要将其用作索引的三元素元组。 Ellipsis 是表示 的对象: 在索引中根据需要重复多次。此版本适用于 >2 的任意维数，其中最后两个表示您想要的方阵。