用特定形状初始化Numpy数组或数组

Question

我有一个Python函数，它接受标量或向量a，执行某种对称扩展操作，例如：

import numpy


def expand(a):
    zero = numpy.zeros_like(a)
    return numpy.array([
        [zero, a, a],
        [a, zero, a],
        [a, a, zero],
        #
        [zero, -a, -a],
        [-a, zero, -a],
        [-a, -a, zero],
    ])

在这个例子中，最终的数组形状将会是 (6, 3, a.shape)。最终我需要的是形状为 (6, a.shape, 3) 或者 (3, a.shape, 6) 并且是连续的，所以我会使用 numpy.moveaxis() 进行移动。即使是 (3, 6, a.shape) 都比现在好很多。我可以使用以下代码实现：

import numpy


def expand(a):
    zero = numpy.zeros_like(a)
    return numpy.array([
        [zero, a, a, zero, -a, -a],
        [a, zero, a, -a, zero, -a],
        [a, a, zero, -a, -a, zero]
    ])

但这个版本不如第一个版本易读，特别是当转换更加复杂时。（始终存在3列。）

有没有一种方式可以直接初始化out的正确形状？（请注意，reshape()不能做到这一点，数据将以错误的顺序排列。）

- Nico Schlömer

Answer 1

有没有一种方法可以直接初始化out的正确形状？

是的，实际上有：只需使用所需的形状进行分配，然后根据需要在转置视图上进行初始化。由于视图是一个视图，所有更改将传播到原始数组，同时保持其内存布局。

import numpy
    
def expand(a):
    zero = numpy.zeros_like(a)
    out = numpy.empty((6,*a.shape,3),a.dtype)
    out.transpose(0,2,1)[...] = [
        [zero, a, a],
        [a, zero, a],
        [a, a, zero],
        #
        [zero, -a, -a],
        [-a, zero, -a],
        [-a, -a, zero],
    ]
    return out

例子：

>>> out = expand(numpy.arange(1,3))
>>> out
array([[[ 0,  1,  1],
        [ 0,  2,  2]],

       [[ 1,  0,  1],
        [ 2,  0,  2]],

       [[ 1,  1,  0],
        [ 2,  2,  0]],

       [[ 0, -1, -1],
        [ 0, -2, -2]],

       [[-1,  0, -1],
        [-2,  0, -2]],

       [[-1, -1,  0],
        [-2, -2,  0]]])
>>> out.flags
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False

- Paul Panzer

Answer 2

让我们尝试(3,6) + a.shape：

def expand(a):
    sub = 1-np.eye(3, dtype=int)
    kernel = np.hstack([sub,-sub])
    return np.multiply.outer(kernel, a)

expand(a).shape
# (3, 6, 2, 5)

现在如果您想要 (3,a.shape, 6)，您可以使用 transpose：

def expand2(a):
    sub = 1-np.eye(3, dtype=int)
    kernel = np.hstack([sub,-sub])
    return np.multiply.outer(kernel, a).transpose(0,2,3,1)

expand2(a).shape
# (3, 2, 5, 6)

- Quang Hoang

---

我的确可以工作，但我并不是特别喜欢：它对于本不必要的数字计算过于依赖，只适用于问题中的特殊情况，并且与问题中第一版相比可读性可能不如前者。 - Nico Schlömer

Answer 3

这是一个简单直接的numpy方法

>>> a = np.arange(1,3)
>>> 
>>> out = np.empty((6,*a.shape,3),a.dtype)
>>> aux = out.reshape(2,3,-1,3)
>>> aux[0] = a[:,None]
>>> aux[1] = -a[:,None]
>>> np.einsum("ijkj->ijk",aux)[...] = 0
>>> 
>>> out
array([[[ 0,  1,  1],
        [ 0,  2,  2]],

       [[ 1,  0,  1],
        [ 2,  0,  2]],

       [[ 1,  1,  0],
        [ 2,  2,  0]],

       [[ 0, -1, -1],
        [ 0, -2, -2]],

       [[-1,  0, -1],
        [-2,  0, -2]],

       [[-1, -1,  0],
        [-2, -2,  0]]])
>>> 
>>> out.flags
  C_CONTIGUOUS : True
  F_CONTIGUOUS : False
  OWNDATA : True
  WRITEABLE : True
  ALIGNED : True
  WRITEBACKIFCOPY : False
  UPDATEIFCOPY : False

- Paul Panzer