python气象信息系统工程 第十三章

/小书匠/python

第13章 计算加速与Fortran绑定

13.1 原生代码优化

13.1.1 将代码向量化

1
2
3
4
5
6
import numpy as np
x = np.random.randn(10000)
y = np.random.randn(10000)
print(x.shape)
print(y.shape)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
x_mean, y_mean = 0, 0 # 用于保存计算中间值
cov_xy = 0 # 用于保存计算中间值
sigma_x, sigma_y = 0, 0 # 用于保存计算中间值
for i in range(10000):
   x_mean = x_mean + x[i]
   y_mean = y_mean + y[i]
x_mean = x_mean / 10000 # 计算x的均值
y_mean = y_mean / 10000 # 计算y的均值
for i in range(10000):
   cov_xy = cov_xy + (x[i] - x_mean) * (y[i] - y_mean)
   sigma_x = sigma_x + (x[i] - x_mean) ** 2
   sigma_y = sigma_y + (y[i] - y_mean) ** 2
sigma_x = sigma_x ** 0.5
sigma_y = sigma_y ** 0.5
r = cov_xy / (sigma_x * sigma_y)

1
2
3
4
5
6
7
x_mean = np.mean(x)
y_mean = np.mean(y)
cov_xy = np.sum((x - x_mean)*(y - y_mean))
sigma_x = np.sqrt(np.sum((x - x_mean)**2))
sigma_y = np.sqrt(np.sum((y - y_mean)**2))
r = cov_xy / (sigma_x * sigma_y)

13.1.2 使用Numba对循环加速

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from numba import jit
import random
@jit(nopython=True)
def r_calc(x, y):
   x_mean, y_mean = 0, 0 # 用于保存计算中间值
   cov_xy = 0 # 用于保存计算中间值
   sigma_x, sigma_y = 0, 0 # 用于保存计算中间值
   for i in range(10000):
      x_mean = x_mean + x[i]
      y_mean = y_mean + y[i]
   x_mean = x_mean / 10000 # 计算x的均值
   y_mean = y_mean / 10000 # 计算y的均值
   for i in range(10000):
      cov_xy = cov_xy + (x[i] - x_mean) * (y[i] - y_mean)
      sigma_x = sigma_x + (x[i] - x_mean) ** 2
      sigma_y = sigma_y + (y[i] - y_mean) ** 2
   sigma_x = sigma_x ** 0.5
   sigma_y = sigma_y ** 0.5
   r = cov_xy / (sigma_x * sigma_y)
   return r


13.2 独立语言绑定

13.2.1 Cython

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
cimport numpy as np
cimport cython

@cython.boundscheck(False)
cdef np.float64_t _r_calc(np.ndarray[np.float64_t, ndim=1] x, np.ndarray[np.float64_t, ndim=1] y):
   cdef np.float64_t x_mean = 0, y_mean = 0 # 用于保存计算中间值
   cdef np.float64_t cov_xy = 0 # 用于保存计算中间值
   cdef np.float64_t sigma_x = 0, sigma_y = 0 # 用于保存计算中间值
   cdef np.float64_t r
   for i in range(10000):
      x_mean = x_mean + x[i]
      y_mean = y_mean + y[i]
   x_mean = x_mean / 10000 # 计算x的均值
   y_mean = y_mean / 10000 # 计算y的均值
   for i in range(10000):
      cov_xy = cov_xy + (x[i] - x_mean) * (y[i] - y_mean)
      sigma_x = sigma_x + (x[i] - x_mean) ** 2
      sigma_y = sigma_y + (y[i] - y_mean) ** 2
   sigma_x = sigma_x ** 0.5
   sigma_y = sigma_y ** 0.5
   r = cov_xy / (sigma_x * sigma_y)
   return r

def r_calc(x, y):
   return _r_calc(x, y)

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
%%cython -a
cimport numpy as np
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
cdef np.float64_t _r_calc(np.ndarray[np.float64_t, ndim=1] x, 
                      np.ndarray[np.float64_t, ndim=1] y):
   cdef np.float64_t x_mean = 0, y_mean = 0 # 用于保存计算中间值
   cdef np.float64_t cov_xy = 0 # 用于保存计算中间值
   cdef np.float64_t sigma_x = 0, sigma_y = 0 # 用于保存计算中间值
   cdef np.float64_t r
   for i in range(10000):
      x_mean = x_mean + x[i]
      y_mean = y_mean + y[i]
   x_mean = x_mean / 10000 # 计算x的均值
   y_mean = y_mean / 10000 # 计算y的均值
   for i in range(10000):
      cov_xy = cov_xy + (x[i] - x_mean) * (y[i] - y_mean)
      sigma_x = sigma_x + (x[i] - x_mean) ** 2
      sigma_y = sigma_y + (y[i] - y_mean) ** 2
   sigma_x = sigma_x ** 0.5
   sigma_y = sigma_y ** 0.5
   r = cov_xy / (sigma_x * sigma_y)
   return r

def r_calc(x, y):
   return _r_calc(x, y)

13.2.2 Fortran

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
30
subroutine r_calc(x, nx, y, ny, r)
   implicit none
   integer :: nx, ny
   real(8) :: x(nx)
   real(8) :: y(ny)
   real(8) :: r
!f2py intent(in),depend(nx) x
!f2py intent(in),depend(ny) y
!f2py intent(out) r

   real(8) :: x_mean = 0, y_mean = 0 ! 用于保存计算中间值
   real(8) :: cov_xy = 0 ! 用于保存计算中间值
   real(8) :: sigma_x = 0, sigma_y = 0 ! 用于保存计算中间值
   integer :: i
   do i = 1, 10000
      x_mean = x_mean + x(i)
      y_mean = y_mean + y(i)
   enddo
   x_mean = x_mean / 10000 ! 计算x的均值
   y_mean = y_mean / 10000 ! 计算y的均值
   do i = 1, 10000
      cov_xy = cov_xy + (x(i) - x_mean) * (y(i) - y_mean)
      sigma_x = sigma_x + (x(i) - x_mean) ** 2
      sigma_y = sigma_y + (y(i) - y_mean) ** 2
   enddo
   sigma_x = sigma_x ** 0.5
   sigma_y = sigma_y ** 0.5
   r = cov_xy / (sigma_x * sigma_y)
end subroutine r_calc