NumPy 快速入门#

前提条件#

你需要了解一些 Python.为了复习,请参阅 Python 教程.

要运行这些示例,除了 NumPy 之外,您还需要安装 matplotlib.

Learner profile

这是对NumPy中数组的快速概述.它演示了如何表示和操作n维(\(n>=2\))数组.特别是,如果你不知道如何将常用函数应用于n维数组(不使用for循环),或者如果你想了解n维数组的轴和形状属性,这篇文章可能会有所帮助.

Learning Objectives

阅读后,你应该能够:

  • 理解NumPy中的一维、二维和n维数组的区别;

  • 了解如何在不使用for循环的情况下对n维数组应用一些线性代数操作;

  • 理解n维数组的轴和形状属性.

基础知识#

NumPy 的主要对象是同质的多维数组.它是一个元素表(通常是数字),所有元素类型相同,通过非负整数元组索引.在 NumPy 中,维度称为 .

例如,3D空间中一个点的坐标数组 [1, 2, 1] 有一个轴.该轴有3个元素,所以我们说它的长度是3.在下图中,数组有2个轴.第一个轴的长度是2,第二个轴的长度是3.

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

NumPy 的数组类被称为 ndarray.它也被称为别名 array.注意,``numpy.array`` 与标准 Python 库类 array.array 不同,后者仅处理一维数组并提供较少的功能.一个 ndarray 对象的更重要属性包括:

ndarray.ndim

数组的轴(维度)数量.

ndarray.shape

数组的维度.这是一个整数元组,表示数组在每个维度的大小.对于一个有 n 行和 m 列的矩阵,``shape`` 将是 (n,m).``shape`` 元组的长度因此是轴的数量,``ndim``.

ndarray.size

数组中元素的总数.这等于 shape 元素的乘积.

ndarray.dtype

描述数组中元素类型的对象.可以使用标准的 Python 类型创建或指定 dtype.此外,NumPy 还提供了自己的类型.numpy.int32、numpy.int16 和 numpy.float64 是一些例子.

ndarray.itemsize

数组中每个元素的字节大小.例如,类型为 float64 的数组具有 itemsize 8 (=64/8),而类型为 complex32 的数组具有 itemsize 4 (=32/8).它等同于 ndarray.dtype.itemsize.

ndarray.data

包含数组实际元素的缓冲区.通常,我们不需要使用这个属性,因为我们通过索引功能访问数组中的元素.

一个例子#

>>> import numpy as np
>>> a = np.arange(15).reshape(3, 5)
>>> a
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])
>>> a.shape
(3, 5)
>>> a.ndim
2
>>> a.dtype.name
'int64'
>>> a.itemsize
8
>>> a.size
15
>>> type(a)
<class 'numpy.ndarray'>
>>> b = np.array([6, 7, 8])
>>> b
array([6, 7, 8])
>>> type(b)
<class 'numpy.ndarray'>

数组创建#

有几种方法可以创建数组.

例如,你可以使用 array 函数从常规的 Python 列表或元组创建一个数组.结果数组的类型是从序列中元素的类型推断出来的.

>>> import numpy as np
>>> a = np.array([2, 3, 4])
>>> a
array([2, 3, 4])
>>> a.dtype
dtype('int64')
>>> b = np.array([1.2, 3.5, 5.1])
>>> b.dtype
dtype('float64')

一个常见的错误是使用多个参数调用 array,而不是提供单个序列作为参数.

>>> a = np.array(1, 2, 3, 4)    # WRONG
Traceback (most recent call last):
  ...
TypeError: array() takes from 1 to 2 positional arguments but 4 were given
>>> a = np.array([1, 2, 3, 4])  # RIGHT

array 将序列的序列转换为二维数组,将序列的序列的序列转换为三维数组,依此类推.

>>> b = np.array([(1.5, 2, 3), (4, 5, 6)])
>>> b
array([[1.5, 2. , 3. ],
       [4. , 5. , 6. ]])

数组的类型也可以在创建时显式指定:

>>> c = np.array([[1, 2], [3, 4]], dtype=complex)
>>> c
array([[1.+0.j, 2.+0.j],
       [3.+0.j, 4.+0.j]])

通常,数组的元素最初是未知的,但其大小是已知的.因此,NumPy 提供了几个函数来创建具有初始占位符内容的数组.这些函数减少了增长数组的必要性,这是一个昂贵的操作.

函数 zeros 创建一个充满零的数组,函数 ones 创建一个充满一的数组,而函数 empty 创建一个初始内容随机且取决于内存状态的数组.默认情况下,创建的数组的 dtype 是 float64,但可以通过关键字参数 dtype 指定.

>>> np.zeros((3, 4))
array([[0., 0., 0., 0.],
       [0., 0., 0., 0.],
       [0., 0., 0., 0.]])
>>> np.ones((2, 3, 4), dtype=np.int16)
array([[[1, 1, 1, 1],
        [1, 1, 1, 1],
        [1, 1, 1, 1]],

       [[1, 1, 1, 1],
        [1, 1, 1, 1],
        [1, 1, 1, 1]]], dtype=int16)
>>> np.empty((2, 3)) 
array([[3.73603959e-262, 6.02658058e-154, 6.55490914e-260],  # may vary
       [5.30498948e-313, 3.14673309e-307, 1.00000000e+000]])

要创建数字序列,NumPy 提供了 arange 函数,该函数类似于 Python 内置的 range,但返回一个数组.

>>> np.arange(10, 30, 5)
array([10, 15, 20, 25])
>>> np.arange(0, 2, 0.3)  # it accepts float arguments
array([0. , 0.3, 0.6, 0.9, 1.2, 1.5, 1.8])

arange 与浮点数参数一起使用时,由于有限的浮点精度,通常无法预测获得的元素数量.出于这个原因,通常使用函数 linspace 会更好,该函数接收我们想要的元素数量作为参数,而不是步长:

>>> from numpy import pi
>>> np.linspace(0, 2, 9)                   # 9 numbers from 0 to 2
array([0.  , 0.25, 0.5 , 0.75, 1.  , 1.25, 1.5 , 1.75, 2.  ])
>>> x = np.linspace(0, 2 * pi, 100)        # useful to evaluate function at lots of points
>>> f = np.sin(x)

打印数组#

当你打印一个数组时,NumPy 以类似于嵌套列表的方式显示它,但具有以下布局:

  • 最后一个轴从左到右打印,

  • 倒数第二个从上到下打印,

  • 其余部分也从上到下打印,每个切片之间用空行分隔.

一维数组随后打印为行,二维数组打印为矩阵,三维数组打印为矩阵列表.

>>> a = np.arange(6)                    # 1d array
>>> print(a)
[0 1 2 3 4 5]
>>>
>>> b = np.arange(12).reshape(4, 3)     # 2d array
>>> print(b)
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]
>>>
>>> c = np.arange(24).reshape(2, 3, 4)  # 3d array
>>> print(c)
[[[ 0  1  2  3]
  [ 4  5  6  7]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [16 17 18 19]
  [20 21 22 23]]]

请参阅 下方 以获取更多关于 reshape 的详细信息.

如果数组太大而无法打印,NumPy 会自动跳过数组的中间部分,只打印角落部分:

>>> print(np.arange(10000))
[   0    1    2 ... 9997 9998 9999]
>>>
>>> print(np.arange(10000).reshape(100, 100))
[[   0    1    2 ...   97   98   99]
 [ 100  101  102 ...  197  198  199]
 [ 200  201  202 ...  297  298  299]
 ...
 [9700 9701 9702 ... 9797 9798 9799]
 [9800 9801 9802 ... 9897 9898 9899]
 [9900 9901 9902 ... 9997 9998 9999]]

要禁用此行为并强制 NumPy 打印整个数组,可以使用 set_printoptions 更改打印选项.

>>> np.set_printoptions(threshold=sys.maxsize)  # sys module should be imported

基本操作#

数组上的算术运算符应用 逐元素 .创建一个新数组并用结果填充.

>>> a = np.array([20, 30, 40, 50])
>>> b = np.arange(4)
>>> b
array([0, 1, 2, 3])
>>> c = a - b
>>> c
array([20, 29, 38, 47])
>>> b**2
array([0, 1, 4, 9])
>>> 10 * np.sin(a)
array([ 9.12945251, -9.88031624,  7.4511316 , -2.62374854])
>>> a < 35
array([ True,  True, False, False])

与许多矩阵语言不同,乘积运算符 * 在 NumPy 数组中是按元素操作的.可以使用 @ 运算符(在 python >=3.5 中)或 dot 函数或方法来执行矩阵乘积:

>>> A = np.array([[1, 1],
...               [0, 1]])
>>> B = np.array([[2, 0],
...               [3, 4]])
>>> A * B     # elementwise product
array([[2, 0],
       [0, 4]])
>>> A @ B     # matrix product
array([[5, 4],
       [3, 4]])
>>> A.dot(B)  # another matrix product
array([[5, 4],
       [3, 4]])

一些操作,例如 +=*= ,会就地修改现有数组而不是创建一个新数组.

>>> rg = np.random.default_rng(1)  # create instance of default random number generator
>>> a = np.ones((2, 3), dtype=int)
>>> b = rg.random((2, 3))
>>> a *= 3
>>> a
array([[3, 3, 3],
       [3, 3, 3]])
>>> b += a
>>> b
array([[3.51182162, 3.9504637 , 3.14415961],
       [3.94864945, 3.31183145, 3.42332645]])
>>> a += b  # b is not automatically converted to integer type
Traceback (most recent call last):
    ...
numpy._core._exceptions._UFuncOutputCastingError: Cannot cast ufunc 'add' output from dtype('float64') to dtype('int64') with casting rule 'same_kind'

当操作不同类型的数组时,结果数组的类型对应于更通用或更精确的类型(这种行为称为向上转换).

>>> a = np.ones(3, dtype=np.int32)
>>> b = np.linspace(0, pi, 3)
>>> b.dtype.name
'float64'
>>> c = a + b
>>> c
array([1.        , 2.57079633, 4.14159265])
>>> c.dtype.name
'float64'
>>> d = np.exp(c * 1j)
>>> d
array([ 0.54030231+0.84147098j, -0.84147098+0.54030231j,
       -0.54030231-0.84147098j])
>>> d.dtype.name
'complex128'

许多一元操作,例如计算数组中所有元素的总和,都是作为 ndarray 类的方法实现的.

>>> a = rg.random((2, 3))
>>> a
array([[0.82770259, 0.40919914, 0.54959369],
       [0.02755911, 0.75351311, 0.53814331]])
>>> a.sum()
3.1057109529998157
>>> a.min()
0.027559113243068367
>>> a.max()
0.8277025938204418

默认情况下,这些操作将应用于数组,就好像它是一个数字列表,而不管其形状如何.然而,通过指定 axis 参数,您可以沿着数组的指定轴应用操作:

>>> b = np.arange(12).reshape(3, 4)
>>> b
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>>
>>> b.sum(axis=0)     # sum of each column
array([12, 15, 18, 21])
>>>
>>> b.min(axis=1)     # min of each row
array([0, 4, 8])
>>>
>>> b.cumsum(axis=1)  # cumulative sum along each row
array([[ 0,  1,  3,  6],
       [ 4,  9, 15, 22],
       [ 8, 17, 27, 38]])

通用函数#

NumPy 提供了熟悉的数学函数,如 sin、cos 和 exp.在 NumPy 中,这些被称为”通用函数”(ufunc).在 NumPy 中,这些函数对数组进行逐元素操作,并生成一个数组作为输出.

>>> B = np.arange(3)
>>> B
array([0, 1, 2])
>>> np.exp(B)
array([1.        , 2.71828183, 7.3890561 ])
>>> np.sqrt(B)
array([0.        , 1.        , 1.41421356])
>>> C = np.array([2., -1., 4.])
>>> np.add(B, C)
array([2., 0., 6.])

索引、切片和迭代#

一维 数组可以被索引、切片和迭代,非常类似于 列表 和其他 Python 序列.

>>> a = np.arange(10)**3
>>> a
array([  0,   1,   8,  27,  64, 125, 216, 343, 512, 729])
>>> a[2]
8
>>> a[2:5]
array([ 8, 27, 64])
>>> # equivalent to a[0:6:2] = 1000;
>>> # from start to position 6, exclusive, set every 2nd element to 1000
>>> a[:6:2] = 1000
>>> a
array([1000,    1, 1000,   27, 1000,  125,  216,  343,  512,  729])
>>> a[::-1]  # reversed a
array([ 729,  512,  343,  216,  125, 1000,   27, 1000,    1, 1000])
>>> for i in a:
...     print(i**(1 / 3.))
...
9.999999999999998  # may vary
1.0
9.999999999999998
3.0
9.999999999999998
4.999999999999999
5.999999999999999
6.999999999999999
7.999999999999999
8.999999999999998

多维 数组可以有一个轴对应一个索引.这些索引在一个用逗号分隔的元组中给出:

>>> def f(x, y):
...     return 10 * x + y
...
>>> b = np.fromfunction(f, (5, 4), dtype=int)
>>> b
array([[ 0,  1,  2,  3],
       [10, 11, 12, 13],
       [20, 21, 22, 23],
       [30, 31, 32, 33],
       [40, 41, 42, 43]])
>>> b[2, 3]
23
>>> b[0:5, 1]  # each row in the second column of b
array([ 1, 11, 21, 31, 41])
>>> b[:, 1]    # equivalent to the previous example
array([ 1, 11, 21, 31, 41])
>>> b[1:3, :]  # each column in the second and third row of b
array([[10, 11, 12, 13],
       [20, 21, 22, 23]])

当提供的索引少于轴的数量时,缺失的索引被视为完整切片 :

>>> b[-1]   # the last row. Equivalent to b[-1, :]
array([40, 41, 42, 43])

b[i] 中括号内的表达式被视为一个 i ,后面跟着尽可能多的 : 实例以表示剩余的轴.NumPy 也允许你使用点来写成 b[i, ...].

(...) 代表所需的冒号数量,以生成一个完整的索引元组.例如,如果 x 是一个具有5个轴的数组,那么

  • x[1, 2, ...] 等同于 x[1, 2, :, :, :]

  • x[..., 3]x[:, :, :, :, 3] 并且

  • x[4, ..., 5, :]x[4, :, :, 5, :].

>>> c = np.array([[[  0,  1,  2],  # a 3D array (two stacked 2D arrays)
...                [ 10, 12, 13]],
...               [[100, 101, 102],
...                [110, 112, 113]]])
>>> c.shape
(2, 2, 3)
>>> c[1, ...]  # same as c[1, :, :] or c[1]
array([[100, 101, 102],
       [110, 112, 113]])
>>> c[..., 2]  # same as c[:, :, 2]
array([[  2,  13],
       [102, 113]])

遍历 多维数组是根据第一个轴进行的:

>>> for row in b:
...     print(row)
...
[0 1 2 3]
[10 11 12 13]
[20 21 22 23]
[30 31 32 33]
[40 41 42 43]

然而,如果想要对数组中的每个元素执行操作,可以使用 flat 属性,这是一个遍历数组所有元素的 迭代器:

>>> for element in b.flat:
...     print(element)
...
0
1
2
3
10
11
12
13
20
21
22
23
30
31
32
33
40
41
42
43

形状操作#

改变数组的形状#

数组有一个由每个轴上的元素数量给出的形状:

>>> a = np.floor(10 * rg.random((3, 4)))
>>> a
array([[3., 7., 3., 4.],
       [1., 4., 2., 2.],
       [7., 2., 4., 9.]])
>>> a.shape
(3, 4)

可以通过各种命令改变数组的形状.请注意,以下三个命令都会返回一个修改后的数组,但不会改变原始数组:

>>> a.ravel()  # returns the array, flattened
array([3., 7., 3., 4., 1., 4., 2., 2., 7., 2., 4., 9.])
>>> a.reshape(6, 2)  # returns the array with a modified shape
array([[3., 7.],
       [3., 4.],
       [1., 4.],
       [2., 2.],
       [7., 2.],
       [4., 9.]])
>>> a.T  # returns the array, transposed
array([[3., 1., 7.],
       [7., 4., 2.],
       [3., 2., 4.],
       [4., 2., 9.]])
>>> a.T.shape
(4, 3)
>>> a.shape
(3, 4)

ravel 产生的数组中的元素顺序通常是”C风格”,即最右边的索引”变化最快”,所以 a[0, 0] 之后的元素是 a[0, 1].如果数组被重塑为其他形状,数组仍然被视为”C风格”.NumPy 通常以这种顺序创建数组,所以 ravel 通常不需要复制其参数,但如果数组是通过对另一个数组进行切片创建的,或者使用了不寻常的选项创建的,则可能需要复制.函数 ravelreshape 也可以通过一个可选参数指示使用 FORTRAN 风格的数组,其中最左边的索引变化最快.

reshape 函数返回其参数的修改形状,而 ndarray.resize 方法修改数组本身:

>>> a
array([[3., 7., 3., 4.],
       [1., 4., 2., 2.],
       [7., 2., 4., 9.]])
>>> a.resize((2, 6))
>>> a
array([[3., 7., 3., 4., 1., 4.],
       [2., 2., 7., 2., 4., 9.]])

如果在重塑操作中给定一个维度为 -1 ,则其他维度会自动计算:

>>> a.reshape(3, -1)
array([[3., 7., 3., 4.],
       [1., 4., 2., 2.],
       [7., 2., 4., 9.]])

堆叠不同的数组#

几个数组可以沿着不同的轴堆叠在一起:

>>> a = np.floor(10 * rg.random((2, 2)))
>>> a
array([[9., 7.],
       [5., 2.]])
>>> b = np.floor(10 * rg.random((2, 2)))
>>> b
array([[1., 9.],
       [5., 1.]])
>>> np.vstack((a, b))
array([[9., 7.],
       [5., 2.],
       [1., 9.],
       [5., 1.]])
>>> np.hstack((a, b))
array([[9., 7., 1., 9.],
       [5., 2., 5., 1.]])

函数 column_stack 将一维数组作为列堆叠成二维数组.它仅对于二维数组等同于 hstack:

>>> from numpy import newaxis
>>> np.column_stack((a, b))  # with 2D arrays
array([[9., 7., 1., 9.],
       [5., 2., 5., 1.]])
>>> a = np.array([4., 2.])
>>> b = np.array([3., 8.])
>>> np.column_stack((a, b))  # returns a 2D array
array([[4., 3.],
       [2., 8.]])
>>> np.hstack((a, b))        # the result is different
array([4., 2., 3., 8.])
>>> a[:, newaxis]  # view `a` as a 2D column vector
array([[4.],
       [2.]])
>>> np.column_stack((a[:, newaxis], b[:, newaxis]))
array([[4., 3.],
       [2., 8.]])
>>> np.hstack((a[:, newaxis], b[:, newaxis]))  # the result is the same
array([[4., 3.],
       [2., 8.]])

一般来说,对于维度超过两维的数组,`hstack` 沿着它们的第二个轴堆叠,`vstack` 沿着它们的第一轴堆叠,而 concatenate 允许提供一个可选参数,指定沿着哪个轴进行堆叠.

Note

在复杂情况下,`r_` 和 c_ 对于通过沿一个轴堆叠数字来创建数组非常有用.它们允许使用范围字面量 :.:

>>> np.r_[1:4, 0, 4]
array([1, 2, 3, 0, 4])

当与数组作为参数一起使用时,`r_` 和 c_ 在默认行为上类似于 vstackhstack,但允许一个可选参数,给出沿哪个轴进行连接.

将一个数组拆分为几个较小的数组#

使用 hsplit,你可以沿着数组的水平轴分割数组,可以通过指定要返回的等形状数组的数量,或者通过指定分割应发生的列来实现:

>>> a = np.floor(10 * rg.random((2, 12)))
>>> a
array([[6., 7., 6., 9., 0., 5., 4., 0., 6., 8., 5., 2.],
       [8., 5., 5., 7., 1., 8., 6., 7., 1., 8., 1., 0.]])
>>> # Split `a` into 3
>>> np.hsplit(a, 3)
[array([[6., 7., 6., 9.],
       [8., 5., 5., 7.]]), array([[0., 5., 4., 0.],
       [1., 8., 6., 7.]]), array([[6., 8., 5., 2.],
       [1., 8., 1., 0.]])]
>>> # Split `a` after the third and the fourth column
>>> np.hsplit(a, (3, 4))
[array([[6., 7., 6.],
       [8., 5., 5.]]), array([[9.],
       [7.]]), array([[0., 5., 4., 0., 6., 8., 5., 2.],
       [1., 8., 6., 7., 1., 8., 1., 0.]])]

vsplit 沿垂直轴分割,而 array_split 允许指定沿哪个轴分割.

副本和视图#

在操作和处理数组时,它们的数据有时会被复制到一个新的数组中,有时则不会.这常常是初学者感到困惑的原因.有三种情况:

完全没有复制#

简单的赋值不会复制对象或其数据.

>>> a = np.array([[ 0,  1,  2,  3],
...               [ 4,  5,  6,  7],
...               [ 8,  9, 10, 11]])
>>> b = a            # no new object is created
>>> b is a           # a and b are two names for the same ndarray object
True

Python 传递可变对象作为引用,因此函数调用不会进行复制.

>>> def f(x):
...     print(id(x))
...
>>> id(a)  # id is a unique identifier of an object 
148293216  # may vary
>>> f(a)   
148293216  # may vary

查看或浅拷贝#

不同的数组对象可以共享相同的数据.``view`` 方法创建一个查看相同数据的新数组对象.

>>> c = a.view()
>>> c is a
False
>>> c.base is a            # c is a view of the data owned by a
True
>>> c.flags.owndata
False
>>>
>>> c = c.reshape((2, 6))  # a's shape doesn't change
>>> a.shape
(3, 4)
>>> c[0, 4] = 1234         # a's data changes
>>> a
array([[   0,    1,    2,    3],
       [1234,    5,    6,    7],
       [   8,    9,   10,   11]])

切片数组会返回其视图:

>>> s = a[:, 1:3]
>>> s[:] = 10  # s[:] is a view of s. Note the difference between s = 10 and s[:] = 10
>>> a
array([[   0,   10,   10,    3],
       [1234,   10,   10,    7],
       [   8,   10,   10,   11]])

深度复制#

copy 方法生成数组及其数据的完整副本.

>>> d = a.copy()  # a new array object with new data is created
>>> d is a
False
>>> d.base is a  # d doesn't share anything with a
False
>>> d[0, 0] = 9999
>>> a
array([[   0,   10,   10,    3],
       [1234,   10,   10,    7],
       [   8,   10,   10,   11]])

有时,如果不再需要原始数组,应在切片后调用 copy.例如,假设 a 是一个巨大的中间结果,最终结果 b 仅包含 a 的一小部分,那么在用切片构造 b 时应进行深拷贝:

>>> a = np.arange(int(1e8))
>>> b = a[:100].copy()
>>> del a  # the memory of ``a`` can be released.

如果使用 b = a[:100],``a`` 将被 b 引用,即使执行 del a,``a`` 仍将保留在内存中.

函数和方法概述#

以下是一些有用的 NumPy 函数和方法名称,按类别排序.查看 按主题分类的例程和对象 获取完整列表.

数组创建

arange, array, copy, empty, empty_like, eye, fromfile, fromfunction, identity, linspace, logspace, mgrid, ogrid, ones, ones_like, r_, zeros, zeros_like

转换

ndarray.astype, atleast_1d, atleast_2d, atleast_3d, mat

操作

array_split, column_stack, concatenate, diagonal, dsplit, dstack, hsplit, hstack, ndarray.item, newaxis, ravel, repeat, reshape, resize, squeeze, swapaxes, take, transpose, vsplit, vstack

问题

all, any, nonzero, where

排序

argmax, argmin, argsort, max, min, ptp, searchsorted, sort

操作

choose, compress, cumprod, cumsum, inner, ndarray.fill, imag, prod, put, putmask, real, sum

基本统计

cov, mean, std, var

基本线性代数

cross, dot, outer, linalg.svd, vdot

较不基本#

广播规则#

广播允许通用函数以有意义的方式处理不完全相同形状的输入.

广播的第一条规则是,如果所有输入数组没有相同的维数,则会在较小数组的形状前反复添加”1”,直到所有数组具有相同的维数.

广播的第二条规则确保在一个特定维度上大小为1的数组表现得好像它们在该维度上具有最大形状的数组的大小.数组元素的值被假定为沿该维度对”广播”数组是相同的.

在应用广播规则之后,所有数组的大小必须匹配.更多细节可以在 广播 中找到.

高级索引和索引技巧#

NumPy 提供的索引功能比常规的 Python 序列更多.除了像之前看到的那样通过整数和切片进行索引外,数组还可以通过整数数组和布尔数组进行索引.

使用索引数组进行索引#

>>> a = np.arange(12)**2  # the first 12 square numbers
>>> i = np.array([1, 1, 3, 8, 5])  # an array of indices
>>> a[i]  # the elements of `a` at the positions `i`
array([ 1,  1,  9, 64, 25])
>>>
>>> j = np.array([[3, 4], [9, 7]])  # a bidimensional array of indices
>>> a[j]  # the same shape as `j`
array([[ 9, 16],
       [81, 49]])

当索引数组 a 是多维的时,一个单一的索引数组指的是 a 的第一个维度.以下示例通过使用调色板将标签图像转换为彩色图像来展示这种行为.

>>> palette = np.array([[0, 0, 0],         # black
...                     [255, 0, 0],       # red
...                     [0, 255, 0],       # green
...                     [0, 0, 255],       # blue
...                     [255, 255, 255]])  # white
>>> image = np.array([[0, 1, 2, 0],  # each value corresponds to a color in the palette
...                   [0, 3, 4, 0]])
>>> palette[image]  # the (2, 4, 3) color image
array([[[  0,   0,   0],
        [255,   0,   0],
        [  0, 255,   0],
        [  0,   0,   0]],

       [[  0,   0,   0],
        [  0,   0, 255],
        [255, 255, 255],
        [  0,   0,   0]]])

我们也可以为多个维度提供索引.每个维度的索引数组必须具有相同的形状.

>>> a = np.arange(12).reshape(3, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> i = np.array([[0, 1],  # indices for the first dim of `a`
...               [1, 2]])
>>> j = np.array([[2, 1],  # indices for the second dim
...               [3, 3]])
>>>
>>> a[i, j]  # i and j must have equal shape
array([[ 2,  5],
       [ 7, 11]])
>>>
>>> a[i, 2]
array([[ 2,  6],
       [ 6, 10]])
>>>
>>> a[:, j]
array([[[ 2,  1],
        [ 3,  3]],

       [[ 6,  5],
        [ 7,  7]],

       [[10,  9],
        [11, 11]]])

在Python中,``arr[i, j]`` 完全等同于 arr[(i, j)]—因此我们可以将 ij 放入一个 tuple 中,然后用那个进行索引.

>>> l = (i, j)
>>> # equivalent to a[i, j]
>>> a[l]
array([[ 2,  5],
       [ 7, 11]])

然而,我们不能通过将 ij 放入一个数组来实现这一点,因为这个数组将被解释为索引 a 的第一个维度.

>>> s = np.array([i, j])
>>> # not what we want
>>> a[s]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: index 3 is out of bounds for axis 0 with size 3
>>> # same as `a[i, j]`
>>> a[tuple(s)]
array([[ 2,  5],
       [ 7, 11]])

使用数组进行索引的另一个常见用途是搜索时间相关序列的最大值:

>>> time = np.linspace(20, 145, 5)  # time scale
>>> data = np.sin(np.arange(20)).reshape(5, 4)  # 4 time-dependent series
>>> time
array([ 20.  ,  51.25,  82.5 , 113.75, 145.  ])
>>> data
array([[ 0.        ,  0.84147098,  0.90929743,  0.14112001],
       [-0.7568025 , -0.95892427, -0.2794155 ,  0.6569866 ],
       [ 0.98935825,  0.41211849, -0.54402111, -0.99999021],
       [-0.53657292,  0.42016704,  0.99060736,  0.65028784],
       [-0.28790332, -0.96139749, -0.75098725,  0.14987721]])
>>> # index of the maxima for each series
>>> ind = data.argmax(axis=0)
>>> ind
array([2, 0, 3, 1])
>>> # times corresponding to the maxima
>>> time_max = time[ind]
>>>
>>> data_max = data[ind, range(data.shape[1])]  # => data[ind[0], 0], data[ind[1], 1]...
>>> time_max
array([ 82.5 ,  20.  , 113.75,  51.25])
>>> data_max
array([0.98935825, 0.84147098, 0.99060736, 0.6569866 ])
>>> np.all(data_max == data.max(axis=0))
True

你也可以使用数组作为目标进行索引赋值:

>>> a = np.arange(5)
>>> a
array([0, 1, 2, 3, 4])
>>> a[[1, 3, 4]] = 0
>>> a
array([0, 0, 2, 0, 0])

然而,当索引列表包含重复项时,分配会进行多次,只留下最后一个值:

>>> a = np.arange(5)
>>> a[[0, 0, 2]] = [1, 2, 3]
>>> a
array([2, 1, 3, 3, 4])

这是相当合理的,但如果你想要使用 Python 的 += 结构,请注意它可能不会做你期望的事情:

>>> a = np.arange(5)
>>> a[[0, 0, 2]] += 1
>>> a
array([1, 1, 3, 3, 4])

尽管在索引列表中0出现了两次,第0个元素只增加了一次.这是因为Python要求 a += 1 等同于 a = a + 1.

使用布尔数组进行索引#

当我们用(整数)索引数组索引数组时,我们提供了要选择的索引列表.使用布尔索引的方法不同;我们明确选择数组中我们想要的项和我们不想要的项.

对于布尔索引,人们能想到的最自然的方法是使用与原始数组*具有相同形状*的布尔数组:

>>> a = np.arange(12).reshape(3, 4)
>>> b = a > 4
>>> b  # `b` is a boolean with `a`'s shape
array([[False, False, False, False],
       [False,  True,  True,  True],
       [ True,  True,  True,  True]])
>>> a[b]  # 1d array with the selected elements
array([ 5,  6,  7,  8,  9, 10, 11])

这个属性在赋值中非常有用:

>>> a[b] = 0  # All elements of `a` higher than 4 become 0
>>> a
array([[0, 1, 2, 3],
       [4, 0, 0, 0],
       [0, 0, 0, 0]])

你可以查看以下示例,了解如何使用布尔索引来生成 Mandelbrot 集 的图像:

>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> def mandelbrot(h, w, maxit=20, r=2):
...     """Returns an image of the Mandelbrot fractal of size (h,w)."""
...     x = np.linspace(-2.5, 1.5, 4*h+1)
...     y = np.linspace(-1.5, 1.5, 3*w+1)
...     A, B = np.meshgrid(x, y)
...     C = A + B*1j
...     z = np.zeros_like(C)
...     divtime = maxit + np.zeros(z.shape, dtype=int)
...
...     for i in range(maxit):
...         z = z**2 + C
...         diverge = abs(z) > r                    # who is diverging
...         div_now = diverge & (divtime == maxit)  # who is diverging now
...         divtime[div_now] = i                    # note when
...         z[diverge] = r                          # avoid diverging too much
...
...     return divtime
>>> plt.clf()
>>> plt.imshow(mandelbrot(400, 400))
../_images/quickstart-1.png

第二种使用布尔值索引的方法更类似于整数索引;对于数组的每个维度,我们给出一个一维布尔数组来选择我们想要的切片:

>>> a = np.arange(12).reshape(3, 4)
>>> b1 = np.array([False, True, True])         # first dim selection
>>> b2 = np.array([True, False, True, False])  # second dim selection
>>>
>>> a[b1, :]                                   # selecting rows
array([[ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>>
>>> a[b1]                                      # same thing
array([[ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>>
>>> a[:, b2]                                   # selecting columns
array([[ 0,  2],
       [ 4,  6],
       [ 8, 10]])
>>>
>>> a[b1, b2]                                  # a weird thing to do
array([ 4, 10])

请注意,一维布尔数组的长度必须与你想要切片的维度(或轴)的长度一致.在前面的例子中,``b1`` 的长度为 3(即 a 中的 数),而 b2 的长度为 4,适合索引 a 的第 2 轴(列).

ix_() 函数#

ix_ 函数可以用来组合不同的向量,以便为每个 n-元组获得结果.例如,如果你想计算所有从每个向量 a、b 和 c 中取出的三元组的 a+b*c:

>>> a = np.array([2, 3, 4, 5])
>>> b = np.array([8, 5, 4])
>>> c = np.array([5, 4, 6, 8, 3])
>>> ax, bx, cx = np.ix_(a, b, c)
>>> ax
array([[[2]],

       [[3]],

       [[4]],

       [[5]]])
>>> bx
array([[[8],
        [5],
        [4]]])
>>> cx
array([[[5, 4, 6, 8, 3]]])
>>> ax.shape, bx.shape, cx.shape
((4, 1, 1), (1, 3, 1), (1, 1, 5))
>>> result = ax + bx * cx
>>> result
array([[[42, 34, 50, 66, 26],
        [27, 22, 32, 42, 17],
        [22, 18, 26, 34, 14]],

       [[43, 35, 51, 67, 27],
        [28, 23, 33, 43, 18],
        [23, 19, 27, 35, 15]],

       [[44, 36, 52, 68, 28],
        [29, 24, 34, 44, 19],
        [24, 20, 28, 36, 16]],

       [[45, 37, 53, 69, 29],
        [30, 25, 35, 45, 20],
        [25, 21, 29, 37, 17]]])
>>> result[3, 2, 4]
17
>>> a[3] + b[2] * c[4]
17

你也可以如下实现 reduce:

>>> def ufunc_reduce(ufct, *vectors):
...    vs = np.ix_(*vectors)
...    r = ufct.identity
...    for v in vs:
...        r = ufct(r, v)
...    return r

然后像这样使用它:

>>> ufunc_reduce(np.add, a, b, c)
array([[[15, 14, 16, 18, 13],
        [12, 11, 13, 15, 10],
        [11, 10, 12, 14,  9]],

       [[16, 15, 17, 19, 14],
        [13, 12, 14, 16, 11],
        [12, 11, 13, 15, 10]],

       [[17, 16, 18, 20, 15],
        [14, 13, 15, 17, 12],
        [13, 12, 14, 16, 11]],

       [[18, 17, 19, 21, 16],
        [15, 14, 16, 18, 13],
        [14, 13, 15, 17, 12]]])

与普通的 ufunc.reduce 相比,这个版本的 reduce 的优势在于它利用了 广播规则 以避免创建一个大小为输出乘以向量数量的参数数组.

使用字符串进行索引#

请参见 结构化数组.

技巧和窍门#

这里我们提供一个简短且有用的提示列表.

“自动”重塑#

要改变数组的尺寸,你可以省略其中一个尺寸,该尺寸将自动推导:

>>> a = np.arange(30)
>>> b = a.reshape((2, -1, 3))  # -1 means "whatever is needed"
>>> b.shape
(2, 5, 3)
>>> b
array([[[ 0,  1,  2],
        [ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11],
        [12, 13, 14]],

       [[15, 16, 17],
        [18, 19, 20],
        [21, 22, 23],
        [24, 25, 26],
        [27, 28, 29]]])

向量堆叠#

我们如何从一组等大小的行向量构建一个二维数组?在 MATLAB 中,这非常简单:如果 xy 是两个长度相同的向量,你只需要做 m=[x;y].在 NumPy 中,这可以通过 column_stackdstackhstackvstack 函数来完成,具体取决于在哪个维度上进行堆叠.例如:

>>> x = np.arange(0, 10, 2)
>>> y = np.arange(5)
>>> m = np.vstack([x, y])
>>> m
array([[0, 2, 4, 6, 8],
       [0, 1, 2, 3, 4]])
>>> xy = np.hstack([x, y])
>>> xy
array([0, 2, 4, 6, 8, 0, 1, 2, 3, 4])

这些函数在超过两个维度背后的逻辑可能会很奇怪.

直方图#

NumPy 的 histogram 函数应用于数组时返回一对向量:数组的直方图和分箱边缘的向量.注意:matplotlib 也有一个构建直方图的函数(称为 hist,类似于 Matlab),它与 NumPy 中的函数不同.主要区别在于 pylab.hist 自动绘制直方图,而 numpy.histogram 仅生成数据.

>>> import numpy as np
>>> rg = np.random.default_rng(1)
>>> import matplotlib.pyplot as plt
>>> # Build a vector of 10000 normal deviates with variance 0.5^2 and mean 2
>>> mu, sigma = 2, 0.5
>>> v = rg.normal(mu, sigma, 10000)
>>> # Plot a normalized histogram with 50 bins
>>> plt.hist(v, bins=50, density=True)       # matplotlib version (plot)
(array...)
>>> # Compute the histogram with numpy and then plot it
>>> (n, bins) = np.histogram(v, bins=50, density=True)  # NumPy version (no plot)
>>> plt.plot(.5 * (bins[1:] + bins[:-1]), n) 
../_images/quickstart-2.png

使用 Matplotlib >=3.4 你也可以使用 plt.stairs(n, bins).

进一步阅读#