SC.NumPy 04 | 重构地球科学数据的一千〇一种方式 100%

Introduction

上一期我们说到重构地球科学数据的两种形式，一种是索引(包含高级索引——切片)，也就是将我们所需的部分从完整的序列中截取出来；另一种就是重组，当我们无法完整读取、生产出完整序列时，如何将多个离散的数据源整合为一个统一整体。

上周我们对索引部分进行了讲解，这次我们一起来探索如何使用NumPy重组地球科学数据。

广播

在我们处理气象数据中有一个很常见的问题，当涉及到月份天数时应该如何处理？例如，通常我们下载的月尺度ERA5降水数据是单位为mm/d的日均降水量，如何将其转化为对应月份天数的总降水量？

❝

Experiments Begin

Experiment.1：最简单的方法，我直接假装它每个月都是30天，全部乘30不就完事儿了嘛？

Comment.1：可以，但精度似乎有所欠缺。

Experiment.2：那我就要循环了，把12个月挨个循环，乘上对应天数就完了。

Comment.2：可以，但没必要。而且循环的时候，时间轴不是向量化计算的会拖慢速度。（而且我们还没讲循环）

Experiment.3：想要向量化？那我使用np.full_like()创建一个数组，它和降水数据形状一模一样，然后我给它挨个赋值为对应月份的值，再把它们俩乘起来。

import numpy as np

pr_daily_arr = np.random.rand(12, 180, 360)           # 假设这是12个月的月尺度日均降水数据
day_arr = np.full_like(pr_daily_arr, np.nan)          # 假设这是对应月份的天数

days_ls = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]

for i in range(day_arr.shape[0]):
   day_arr[i] = days_ls[i]         # 给月份赋值对应天数

pr_mon_arr = pr_daily_arr * day_arr                  # 月尺度降水量

print(pr_mon_arr)

Comment.3：嗯，前面学得很好。但是这也用到了循环吧？而且，循环赋值再乘法，是不是有点多此一举了？

说了这么多方案，有没有一个很好的解决方案呢？

答案是肯定的，那就是我们的广播运算。

简而言之，广播就是这样一种机制。什么样的机制呢？如果两个数组对应相乘，而它们维度一致但形状并不一样时(假设它们都有三条轴axis=(0, 1, 2)，但仅在axis=2维度的长度不一致，其中一个长度为5，另一个为1)，如果形状不一致的维度中其中一个长度为1，那么另一个维度不为1的数组在该维度上都将根据为1的维度的值计算它所有的值。

这都什么跟什么啊？这还简而言之？这是人能理解的东西？(掀桌)

不慌，文字描述确实在考验人类的想象力，下面我们看个案例就知道了。

import numpy as np

pr_daily_arr = np.random.rand(12, 180, 360)           # 假设这是12个月的月尺度日均降水数据
day_arr = np.array([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31])          # 假设这是对应月份的天数
pr_mon_arr = pr_daily_arr * day_arr[:, np.newaxis, np.newaxis]                  # 月尺度降水量

print(pr_daily_arr[0], '\n')                    # 使用'\n'输出后换行
print(pr_mon_arr[0], '\n')

[[0.7457109  0.1999839  0.86517314 ... 0.9433673  0.44404254 0.83987611]
 [0.25508891 0.39786953 0.04070897 ... 0.81796878 0.32314888 0.43181883]
 [0.50222835 0.5400615  0.87183673 ... 0.22866725 0.25978679 0.5815957 ]
 ...
 [0.09102325 0.02775076 0.2383608  ... 0.7246745  0.30396619 0.5725769 ]
 [0.36923757 0.13091775 0.32149145 ... 0.95352614 0.21708246 0.10558568]
 [0.50097136 0.14826206 0.20927308 ... 0.67874219 0.33524133 0.25731191]] 

[[23.11703799  6.19950096 26.82036733 ... 29.24438641 13.76531887
  26.03615934]
 [ 7.90775628 12.33395546  1.26197816 ... 25.35703213 10.0176153
  13.38638364]
 [15.56907888 16.74190637 27.02693877 ...  7.08868466  8.05339061
  18.02946677]
 ...
 [ 2.82172084  0.8602737   7.38918487 ... 22.4649094   9.42295196
  17.74988382]
 [11.44636476  4.0584504   9.9662349  ... 29.55931045  6.72955629
   3.27315617]
 [15.53011226  4.59612387  6.48746543 ... 21.04100775 10.39248115
   7.97666924]] 

可以看到，不只是其中一个维度不一致可以触发广播运算。只要形状不一致的维度中，其中一个长度为1，均可在该轴上触发广播机制。

这里使用到了newaxis这个NumPy内置属性，它会在我们索引的轴不存在时给我们新添加一个轴。

？，今天为什么净是些难表述的东西。

我们再来看案例：

pr_mon_arr1 = pr_daily_arr * day_arr

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[7], line 1
----> 1 pr_mon_arr1 = pr_daily_arr * day_arr

ValueError: operands could not be broadcast together with shapes (12,180,360) (12,)

print(pr_mon_arr.shape, '\n')
print(day_arr.shape, '\n')
print(day_arr[:, np.newaxis, np.newaxis].shape, '\n')

(180, 360, 12) 

(12,) 

(12, 1, 1)

显然，在不使用newaxis时广播是失败的。因为在一开始的概念里提到了，它们的维度得一致，只是其中一个数组的维度长度可以为1。day_arr只有一个维度，而day_arr[:, np.newaxis, np.newaxis]则拥有与pr_mon_arr一致的三个维度。只是它拓展的两个新维度上长度为1。

当然，广播机制可以进阶一下。但是这种方式很容易混淆，把自己绕进去。还是建议使用newaxis方式，可以明确地看到哪条轴是被广播的。

pr_daily_arr = np.random.rand(180, 360, 12)
day_arr = np.array([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31])
pr_mon_arr = pr_daily_arr * day_arr

print(pr_daily_arr[0], '\n')
print(pr_mon_arr[0], '\n')
print(day_arr.shape, '\n')
print(pr_daily_arr.shape, '\n')

[[0.52306542 0.87345502 0.19613925 ... 0.42060492 0.82202532 0.55724046]
 [0.57119119 0.67627112 0.86564461 ... 0.00805265 0.72385797 0.78095946]
 [0.05010039 0.65136344 0.44689328 ... 0.97225033 0.13130093 0.73874389]
 ...
 [0.94381637 0.53307232 0.05288252 ... 0.24439155 0.5604889  0.59675637]
 [0.02233037 0.01284824 0.73235306 ... 0.9735771  0.74557556 0.96430204]
 [0.00130125 0.27729453 0.52910089 ... 0.3379982  0.51298472 0.02553985]] 

[[16.21502815 24.45674064  6.08031688 ... 13.03875258 24.6607597
  17.27445441]
 [17.70692678 18.93559134 26.83498294 ...  0.24963209 21.71573905
  24.20974313]
 [ 1.55311215 18.23817626 13.85369173 ... 30.1397601   3.93902787
  22.90106074]
 ...
 [29.25830742 14.92602501  1.63935811 ...  7.57613806 16.81466693
  18.49944739]
 [ 0.69224147  0.35975077 22.702945   ... 30.18089016 22.36726679
  29.89336316]
 [ 0.04033873  7.76424677 16.40212773 ... 10.47794412 15.38954152
   0.79173537]] 

(12,) 

(180, 360, 12) 

可以看见，就算我们不使用newaxis，如果两个数组对应的轴长度由内往外(也就是维度从右往左)是对应的，那么外层的轴默认触发广播。如此处day_arr的维度为12，而pr_daily_arr的维度由内往外为12→360→180，因此可以触发广播机制。(关于NumPy中轴的问题，可以移步第一期复习)

堆叠、拼接

下面回到一开始的话题，如何将多个离散数据整合为一个整体？这里主要用到数组堆叠(或拼接)，它的定义是将多个数组沿着指定轴连接在一起的操作，通常用于将具有相同形状或尺寸的多个数组组合成更大的数组。

核心函数包括了：vstack, hstack, dstack, concatenate, stack。

❝

vstack, hstack, dstack只能用于沿固定的数组轴进行拼接，分别为垂直方向(axis=0，我们常说的y轴，按行方向拼接，v即vertical)，水平方向(axis=1，我们常说的x轴，按列方向拼接，h即horizontal)和竖直方向(axis=2，我们常说的z轴，按页方向堆叠，d即deep)。

下面是案例：

arr0 = np.array([[1, 2, 3], [4, 5, 6]])
arr1 = np.array([[7, 8, 9], [10, 11, 12]])

arr2 = np.vstack((arr0, arr1))              # 垂直方向拼接
arr3 = np.hstack((arr0, arr1))              # 水平方向拼接
arr4 = np.dstack((arr0, arr1))              # 竖直方向堆叠

print(arr0, '\n')
print(arr1, '\n')
print(arr2, '\n')
print(arr3, '\n')
print(arr4, '\n')

[[1 2 3]
 [4 5 6]] 

[[ 7  8  9]
 [10 11 12]] 

[[ 1  2  3]
 [ 4  5  6]
 [ 7  8  9]
 [10 11 12]] 

[[ 1  2  3  7  8  9]
 [ 4  5  6 10 11 12]] 

[[[ 1  7]
  [ 2  8]
  [ 3  9]]

 [[ 4 10]
  [ 5 11]
  [ 6 12]]] 

既然上面三个函数这么死板，肯定不会是我们喜欢用的。一种常见的情况，我要在第4个轴要怎么叠？所以就需要用到后面两个函数了。

这里我特意把concatenate放在了stack前面，按理说它们另外四个函数应该是长得更像的。下面我们一个简单的案例告诉你为什么要按这个顺序写：

# 创建两个超多维的数组
arr0 = np.array([[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], [[[13, 14, 15], [16, 17, 18]], [[19, 20, 21], [22, 23, 24]]]])
arr1 = np.array([[[[25, 26, 27], [28, 29, 30]], [[31, 32, 33], [34, 35, 36]]], [[[37, 38, 39], [40, 41, 42]], [[43, 44, 45], [46, 47, 48]]]])

print(arr0.shape, '\n')
print(arr1.shape, '\n')
print(arr0, '\n')

(2, 2, 2, 3) 

(2, 2, 2, 3) 

[[[[ 1  2  3]
   [ 4  5  6]]

  [[ 7  8  9]
   [10 11 12]]]


 [[[13 14 15]
   [16 17 18]]

  [[19 20 21]
   [22 23 24]]]] 

arr2 = np.concatenate((arr0, arr1), axis=3)
arr3 = np.stack((arr0, arr1), axis=3)

print(arr2.shape, '\n')
print(arr3.shape, '\n')

(2, 2, 2, 6) 

(2, 2, 2, 2, 3) 

显然，concatenate是在原有轴上进行拼接的，相当于一个更自由的vstack, hstack, dstack，但是stack却会在我们指定的轴添加一个新轴，而把之前已经存在的轴移动到内层去。

当然，某种意义上，两个二维数组使用dstack与这异曲同工，但也仅限于特殊情况。

arr0 = np.array([[1, 2, 3], [4, 5, 6]])
arr1 = np.array([[7, 8, 9], [10, 11, 12]])

arr2 = np.dstack((arr0, arr1))              # 垂直方向拼接
print(arr2, '\n')
print(arr2.shape, '\n')

[[[ 1  7]
  [ 2  8]
  [ 3  9]]

 [[ 4 10]
  [ 5 11]
  [ 6 12]]] 

(2, 3, 2) 

形状转换

上面我们讲了这么多关于重新组合数组的方式，但是却都是针对多个数组而言的。但很多时候，一个数组本身的形状变换也是极其重要的。

比如，我们深度学习中经常需要降维，需要将二维空间数据展平成一维向量。又例如，两个气象数据集它们的维度都一样，均包括了纬度、经度、时间、气压场，但是它们轴的顺序却不一致，如何交换轴使它们之间能够运算？

一般而言，能够用到这里的函数都是一些较为高阶的运算和数据处理方式了。可能无法给出非常普适性的案例了，需要各位结合自己的使用场景去探索，但在底层的建模中都还算是比较常用的功能吧。

arr0 = np.array([[1, 2], [4, 5]])
arr1 = arr0.T                               # 转置
arr2 = arr0.flatten()                       # 展平
arr3 = np.repeat(arr0, 2, axis=1)           # 重复：在指定轴上重复N次(此处，N=2)，注意此处的重复是数组元素连续重复
arr4 = np.tile(arr0, (2, 2))                # 平铺：在指定轴上重复(M, N)次(此处，M=N=2)，注意此处的重复是数组元素整体复制

print(arr0, '\n')
print(arr1, '\n')
print(arr2, '\n')
print(arr3, '\n')
print(arr4, '\n')

[[1 2]
 [4 5]] 

[[1 4]
 [2 5]] 

[1 2 4 5] 

[[1 1 2 2]
 [4 4 5 5]] 

[[1 2 1 2]
 [4 5 4 5]
 [1 2 1 2]
 [4 5 4 5]] 

以上只是tile的一个二维例子，如果需要在更多轴上重复，可以继续添加轴长度参数。下面是tile一个可能的应用场景：

lon_arr = np.arange(-30, 60)                     # 生成纬度范围

lon_grid = np.tile(lon_arr, (90, 1))             # 90个经度上对应的纬度
print(lon_grid, '\n')

[[-30 -29 -28 ...  57  58  59]
 [-30 -29 -28 ...  57  58  59]
 [-30 -29 -28 ...  57  58  59]
 ...
 [-30 -29 -28 ...  57  58  59]
 [-30 -29 -28 ...  57  58  59]
 [-30 -29 -28 ...  57  58  59]] 

但是NumPy提供了更简单的网格生成方式meshgrid：

lon = np.arange(10, 120)
lat = np.arange(-30, 60)

lon_grid, lat_grid = np.meshgrid(lon, lat)
print(lon_grid, '\n')
print(lat_grid, '\n')

[[ 10  11  12 ... 117 118 119]
 [ 10  11  12 ... 117 118 119]
 [ 10  11  12 ... 117 118 119]
 ...
 [ 10  11  12 ... 117 118 119]
 [ 10  11  12 ... 117 118 119]
 [ 10  11  12 ... 117 118 119]] 

[[-30 -30 -30 ... -30 -30 -30]
 [-29 -29 -29 ... -29 -29 -29]
 [-28 -28 -28 ... -28 -28 -28]
 ...
 [ 57  57  57 ...  57  57  57]
 [ 58  58  58 ...  58  58  58]
 [ 59  59  59 ...  59  59  59]] 

当我们需要进行更复杂的数组形状重塑时，就需要用到reshape函数。通过指定各个轴的长度，重新定义数组的形状。

# 使用reshape重塑数组形状
arr0 = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
arr1 = arr0.reshape(2, 5)

print(arr0, '\n')
print(arr1, '\n')

# 如果其中一个维度的长度不想算，可以使用一次-1来表示未知的维度长度
arr2 = arr0.reshape(2, -1)
print(arr2, '\n')

[ 1  2  3  4  5  6  7  8  9 10] 

[[ 1  2  3  4  5]
 [ 6  7  8  9 10]] 

[[ 1  2  3  4  5]
 [ 6  7  8  9 10]] 

我们前面提到了当两个数据集轴的顺序不一致时，需要交换轴方可进行运算。这就需要用到transpose()函数，通过指定轴的顺序，交换数组轴。

# 交换轴
arr0 = np.random.rand(3, 4, 5)
arr1 = np.random.rand(3, 5, 4)

arr0 * arr1                             # 形状不一致，无法运算

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[51], line 5
      2 arr0 = np.random.rand(3, 4, 5)
      3 arr1 = np.random.rand(3, 5, 4)
----> 5 arr0 * arr1

ValueError: operands could not be broadcast together with shapes (3,4,5) (3,5,4)

print(arr0 * arr1.transpose(0, 2, 1))             # 将arr1的axis=1和axis=2位置交换

[[[0.27157848 0.21724557 0.78304461 0.47336294 0.24355969]
  [0.04092148 0.17370547 0.0760536  0.82288962 0.60827622]
  [0.17774185 0.16140735 0.00456923 0.18290974 0.3983372 ]
  [0.20480657 0.09730546 0.00131533 0.0537118  0.02145419]]

 [[0.00375847 0.43877457 0.12805227 0.10518511 0.11676147]
  [0.23702829 0.10692001 0.02425214 0.36935007 0.06287172]
  [0.66498854 0.11690347 0.15744815 0.09945056 0.84871872]
  [0.04651117 0.03372407 0.07167041 0.19482316 0.64568326]]

 [[0.26522254 0.03228833 0.03018882 0.0268442  0.0264654 ]
  [0.24187292 0.50955092 0.01176138 0.60136186 0.03088058]
  [0.25646249 0.35001574 0.67993698 0.1566376  0.33420856]
  [0.33259152 0.12799435 0.05562951 0.78439656 0.08188747]]]

后记

本以为这期应该不需要多少篇幅就能写完的，没想到不知不觉这么长了。但是已经说好分两期了，那就一鼓作气写完吧。

这期的内容其实有点抽象，需要一定的想象力。而且由于篇幅限制，我输出的都是数组形状，而不是数组内容。建议自行运行文中代码，输出下看看每个步骤都发生了什么。感觉还是结合实际问题理解起来会更简单一点，最重要的还是在自己的代码里用上吧。

另外，欢迎跳转和鲸社区，里面提供了每篇推文对应的Jupyter Notebook，可以省去部署环境的麻烦。

如需快速配置满足地球科学开发基本要求的Python环境，可以阅读近期的另一篇推文。

数组重构的后半部分用上的频率确实不算高，但是基本都是核心部分的代码。所以，对于期望开展更深层次、更复杂、更高阶的建模与数据处理的读者们，这些内容还是很有必要的。

另外就是我们注意到中间涉及到一些比较固定的函数用法其实是其他用法的子集的问题，说的就是你们，那几个stack。自定义参数更多的函数实际上给我们也提供了更自由的选择，赋予了更多的创造性。所以，有能力的话，还是读懂函数之后使用这些更高阶的函数吧。

最后，万恶的调休，又把我的更新计划打乱了。

那么，我们下期再见！

Manuscript: RitasCake

Proof: Philero; RitasCake

获取更多资讯，欢迎订阅微信公众号：Westerlies