以矢量化方式创建给定大小的 NumPy 矩阵

Question

如果有内置的 NumPy 函数或任何 向量化 方法来生成以下二维 NumPy 矩阵 n > 1 的值？让我给出一些 所需矩阵 的示例，因为 n 等于 2、3 和 4。

一种方法是在下三角部分构造一个全为 1 的矩阵，然后简单地乘以一个向量 np.arange(1, n+1)。还有其他选择吗？

import numpy as np

对于 n = 2

n = 2
arr = np.array([[1, 0], [2, 2]])

# array([[1, 0],
#        [2, 2]])

对于 n = 3

n = 3
arr = np.array([[1, 0, 0], [2, 2, 0], [3, 3, 3]])

# array([[1, 0, 0],
#        [2, 2, 0],
#        [3, 3, 3]])

对于 n = 4

n = 4
arr = np.array([[1, 0, 0, 0], [2, 2, 0, 0], [3, 3, 3, 0], [4, 4, 4, 4]])

# array([[1, 0, 0, 0],
#        [2, 2, 0, 0],
#        [3, 3, 3, 0],
#        [4, 4, 4, 4]])

我的尝试（n = 4 的示例）使用嵌套 for 循环。它适用于 n > 1:

的任何值

import numpy as np

n = 4
arr = np.zeros((n, n))

for i in range(n):
    for j in range(0, i+1):
        arr[i][j] = i+1 

# array([[1., 0., 0., 0.],
#        [2., 2., 0., 0.],
#        [3., 3., 3., 0.],
#        [4., 4., 4., 4.]])

Answer 1

您可以使用 np.tril 和 np.broadcast_to 创建初始 arange 变体。

import numpy as np

n = 4
np.tril(np.broadcast_to(np.arange(1, n+1)[:, None], (n, n)))

array([[1, 0, 0, 0],
       [2, 2, 0, 0],
       [3, 3, 3, 0],
       [4, 4, 4, 4]])

Answer 2

这是一个使用 np.where 和广播的解决方案：

x = np.arange(1,n+1)
np.where(x[:,None]>=x,x[:,None],0)
# array([[1, 0, 0, 0],
#        [2, 2, 0, 0],
#        [3, 3, 3, 0],
#        [4, 4, 4, 4]])

Answer 3

我们可以使用范围数组乘以 lower-triangular 掩码 -

In [44]: n = 5

In [45]: np.arange(1,n+1)[:,None]*np.tri(n,dtype=bool)
Out[45]: 
array([[1, 0, 0, 0, 0],
       [2, 2, 0, 0, 0],
       [3, 3, 3, 0, 0],
       [4, 4, 4, 4, 0],
       [5, 5, 5, 5, 5]])

考虑到 arithmetic-oriented 的性质，它很容易移植到 numexpr 以利用 multi-cores 处理大数据 -

import numexpr as ne

ne.evaluate('A*B',{'A':np.arange(1,n+1)[:,None],'B':np.tri(n,dtype=bool)})

基准测试

包括此处发布的所有解决方案。

基准测试脚本 -

import numpy as np
import perfplot

import numexpr as ne

def numexpr_range_broadcast(n):
    return ne.evaluate('A*B',{'A':np.arange(1,n+1)[:,None],'B':np.tri(n,dtype=bool)})

def where_method(n):
    x = np.arange(1,n+1)
    return np.where(x[:,None]>=x,x[:,None],0)

perfplot.show(
    setup=lambda n: n,
    kernels=[
        lambda n: where_method(n),
        lambda n: np.tril(np.broadcast_to(np.arange(1, n+1)[:, None], (n, n))),
        lambda n: np.arange(1,n+1)[:,None]*np.tri(n,dtype=bool),
        lambda n: numexpr_range_broadcast(n),
        ],
    labels=['where','tril_broadacast','range_broadcast','numexpr_range_broadcast'],
    n_range=[10, 20, 50, 80, 100, 200, 500, 800, 1000, 2000, 5000],
    xlabel='n',
    logx=True,
    logy=True,
    )

输出-

因此，在 100 之前的较低 n's，基于 np.where 的一方获胜，在较大的 n's，numexpr 闪耀。

以矢量化方式创建给定大小的 NumPy 矩阵

Creating NumPy matrix of a given size in a vectorized manner

python

numpy

vectorization

基准测试