代码优化实例

本文主要是针对某区域的一个月的轨迹数据进行处理，并完成代码的优化。所选择轨迹的某一行的具体格式如下表所示，其中MMSI_paragraph代表是每一条轨迹特有的id格式，Lat与Lon代表该轨迹在某一时刻的经纬度，Receivedtime（UTC+8）是历史轨迹的时间戳，大约30秒更新一次。

MMSI_paragraph	Lat	Lon	Receivedtime（UTC+8）
44236893_01	113.120742	29.425808	2020-01-28 01:21:42

1.计算每条轨迹的航行时间

首先是利用pandas读取轨迹存储至DF中，并将其时间戳转化为对应datetime的格式。

from symbol import import_stmt
import pandas as pd
from tqdm import tqdm 
from datetime import datetime
import time
DF=pd.read_csv('/home/data/AIS_data_process/trajectory_for_lunwen/异常剔除后轨迹并进行上行下行轨迹提取/直线处轨迹/trajectory_toup_01.csv')
from pandarallel import pandarallel
pandarallel.initialize(progress_bar=True)
BJS_format = "%Y-%m-%d %H:%M:%S"
DF['Receivedtime（UTC+8）'] =DF['Receivedtime（UTC+8）'].parallel_apply(lambda x: datetime.strptime(x, BJS_format)) #apply的并行
mmsi_para_list=DF['MMSI_paragraph'].unique()

下面这个是原始的计算每条轨迹航行时间的代码，该代码是将每一条轨迹单独读取出来，然后计算其在该区域的航行时间。其时间复杂度和空间复杂度都特别高。

start1=time.time()
df_m1_list=[]#存储每个id的轨迹计算后得到的dataframe
for i in range(len(mmsi_para_list)):
    dfm1=DF[DF['MMSI_paragraph']==mmsi_para_list[i]]#这样获取数据时间复杂度较高，因为涉及到每一行的判断
    total_time=(dfm1.iloc[-1]['Receivedtime（UTC+8）']-dfm1.iloc[0]['Receivedtime（UTC+8）']).total_seconds()
    dfm1['total_time']=total_time
    df_m1_list.append(dfm1)#若轨迹数量太多，则都存储在df_m1_list会占用大量的存储空间
DF1=pd.concat(df_m1_list)#concat的操作复杂度较高
end1=time.time()
print(end1-start1)

经代码优化后可以将其改写为如下形式，其利用了assign、groupby、transform等方法来对原始的轨迹进行变换。

• assign的作用是可以将total_time这一个序列追加到DF的这个表的列末。
• groupby是可以将之前自己第一版代码的for循环内置，节省了for循环单独进行每一条轨迹抽取的时间。
• transform通过内置lambda匿名函数实现表的自定义变换，并将其广播到所在的整个组。
  （标量聚合用agg，标量广播/返回同长序列transform，筛选组用filter，多列返回用apply）

start2=time.time()
DF2=DF.assign(total_time=DF.groupby("MMSI_paragraph")["Receivedtime（UTC+8）"].transform(lambda x: (x.iat[-1]-x.iat[0]).total_seconds()))
end2=time.time()
print(end2-start2)

本文所使用的DF共有643万行个轨迹点，共有26524条轨迹。经过代码优化后，计算耗时从原来的2小时21分钟缩减到现在的4秒钟。

高性能原则之一：python里面不能显式出现O(n)的写法

2 计算每条轨迹的航行长度

每一个相同的MMSI_paragraph代表是一条航行轨迹，故在本任务中需要计算相同MMSI_paragraph的每一行的轨迹之间的距离差值。如下图所示为轨迹的部分数据展示，在计算过程中需要设置函数来求其两个不同经纬度之间的距离。
该问题的难点就是如何对大规模的轨迹样本进行航行长度的求取。

在本文中最开始是先建立DF_list这个列表，然后通过for循环将DF中每一条轨迹进行抽取df，抽取后再进行for循环遍历这条轨迹的每一行数据，利用球面的经纬度距离函数求得所航行长度，最后将该长度与所抽取的df进行拼接，放到DF_list这个列表。当循环运行之后利用concat将所有计算后的df合并成新的DF。这样便得到了每一条轨迹的航行长度。
其中地球经纬度之间的距离函数如下所示：

def calcDistance(Lng_A, Lat_A, Lng_B, Lat_B):
    """
    根据两个点的经纬度求两点之间的距离
    :param Lng_A:  经度1
    :param Lat_A:   维度1
    :param Lng_B:  经度2
    :param Lat_B:   维度2
    :return:  单位米
    """
    ra = 6378.140
    rb = 6356.755
    flatten = (ra - rb) / ra
    rad_lat_A = np.radians(Lat_A)
    rad_lng_A = np.radians(Lng_A)
    rad_lat_B = np.radians(Lat_B)
    rad_lng_B = np.radians(Lng_B)
    pA = math.atan(rb / ra * np.tan(rad_lat_A))
    pB = math.atan(rb / ra * np.tan(rad_lat_B))
    xx = math.acos(np.sin(pA) * np.sin(pB) + np.cos(pA) * np.cos(pB) * np.cos(rad_lng_A - rad_lng_B))
    c1 = (np.sin(xx) - xx) * (np.sin(pA) + np.sin(pB)) ** 2 / np.cos(xx / 2) ** 2
    # 经测试当传入这两个点的经纬度一样时会返回0
    if np.sin(xx/2)==0:
        return 0
    c2 = (np.sin(xx) + xx) * (np.sin(pA) - np.sin(pB)) ** 2 / np.sin(xx / 2) ** 2
    dr = flatten / 8 * (c1 - c2)
    distance = ra * (xx + dr)
    return distance*1000

在刚才叙述的过程中，这样的操作浪费了大量的时间和空间，计算效率极低，故可以利用numba和cython进行加速处理。

#基于numba的计算加速方式
from numba import njit #使用numba加速
ra, rb = 6378.140, 6356.755 
@njit
def distance_numba(p, lng):
    '''
    计算轨迹序列的长度
    '''
    flatten = (ra - rb) / ra
    sum_value, nrows, xx = 0.0, p.shape[0], 0.0
    if nrows == 1:
        return sum_value
    for i in range(nrows-1):
        pA, pB, lng_A, lng_B = p[i], p[i+1], lng[i], lng[i+1]
        xx = sin(pA)*sin(pB)
        xx += cos(pA)*cos(pB)*cos(lng_A-lng_B)
        xx = acos(min(xx, 1))
        c1 = (sin(xx)-xx) * (sin(pA)+sin(pB)) ** 2 / cos(xx/2+1e-15) ** 2
        c2 = (sin(xx)+xx) * (sin(pA)-sin(pB)) ** 2 / sin(xx/2+1e-15) ** 2
        dr = flatten / 8 * (c1 - c2)
        dist = ra * (xx + dr) * 1000
        sum_value += dist
    return sum_value
    
DF["p"] = np.arctan(rb/ra*np.tan(np.radians(DF.Lat_d)))
agg_result = DF.groupby("MMSI_paragraph").apply(lambda x: distance_numba(x.p.values, np.radians(x.Lon_d.values)))#利用groupby对不同的轨迹分别处理
DF = DF.merge(agg_result.to_frame().rename(columns={
    
    0:"Distance"}), on="MMSI_paragraph")#将计算后的结果和DF进行merge在一起
#为了加速计算，也可以写成DF1=DF.assign(Distance=agg_result.reindex(DF.MMSI_paragraph).reset_index(drop=True)).drop('p',axis=1)
#在一些情况下可以尝试用reindex替换merge，主要原因是merge本质是O(n**2)的复杂度，数据量越大则越慢
#所以在能用reindex代替merge的时候尽可能去用reindex操作

下面是基于基于cython的计算加速方式

#启动cython
%load_ext cython

%%cython
import numpy as np
cimport numpy as np
import cython
from math import acos, sin, cos
#其中上一行的导入库也可以写为 from 

def distance_cython(double[:] p, double[:] lng):
    cdef:
        double ra = 6378.140, 
        double rb = 6356.755
        double flatten = (ra - rb) / ra
        double sum_value = 0
        double xx, pA, pB, lng_A, lng_B, c1, c2, dr
        int nrows = p.shape[0]
        int i
    if nrows == 1:
        return sum_value
    for i in range(nrows-1):
        pA, pB, lng_A, lng_B = p[i], p[i+1], lng[i], lng[i+1]
        xx = sin(pA)*sin(pB)
        xx += cos(pA)*cos(pB)*cos(lng_A-lng_B)
        xx = acos(min(xx, 1))
        if sin(xx/2)==0:
            sum_value += 0
            continue
        c1 = (sin(xx)-xx) * (sin(pA)+sin(pB)) ** 2 / (cos(xx/2) ** 2)
        c2 = (sin(xx)+xx) * (sin(pA)-sin(pB)) ** 2 / (sin(xx/2) ** 2)
        dr = flatten / 8 * (c1 - c2)
        sum_value += ra * (xx + dr) * 1000
    return sum_value
agg_result = DF.groupby("MMSI_paragraph").apply(lambda x: distance_cython(x.p.values,  np.radians(x.Lon_d.values)))
DF= DF.merge(agg_result.to_frame().rename(columns={
    
    0:"Distance"}), on="MMSI_paragraph").drop("p", axis=1)

针对26524条轨迹进行计算可以得到cython加速总共用时约为13.6秒，而numba加速后则为7.6秒。相比于原始的方法可大幅提高计算效率。

3. 获取轨迹数据的部分数据

本任务旨在随机获取每条轨迹数据的前30%到70%的数据。
原始代码如下：

DF_train=pd.DataFrame()
mmsi_para_list=DF['MMSI_paragraph'].unique()
for mmsi_para in mmsi_para_list:
    df2=DF[DF['MMSI_paragraph']==mmsi_para]
    ratio=random.uniform(0.3,0.7)
    traj_num=int(len(df2)*ratio)
    df_train=df2.iloc[:traj_num]
    DF_train=DF_train.append(df_train)

上面的写法还是无法满足高性能要求，效率较低，故可利用groupby进行的代码的优化。则可以将代码写成以下结构然后进行优化

DF.groupby("MMSI_paragraph",as_index=False, sort=False, group_keys=False).apply(lambda x: x.iloc[:int(x.shape[0]*np.random.uniform(0.3,0.7))])

在这个基础上还可以去进行代码优化。代码优化思想为：对所有MMSI_paragraph的值进行查找，找到每一个相同的MMSI_paragraph开始位置和结束位置。将位置信息存储在list中，然后根据list对数据进行抽样，以明确抽样30%-70%数据后的结束位置在哪里。接下来，根据该list在对原始的轨迹数据DF进行抽样处理。具体代码如下：

%%cython

import pandas as pd
import numpy as np
from numba import njit
import numpy as np
cimport numpy as np
import cython 


def get_count(str[:] arr):
#对所有MMSI_paragraph的值进行查找，找到每一个相同的MMSI_paragraph开始位置和结束位置
#类似于aaaabbbcc，返回数组[4,3,2]
    cdef:
        int n_groups = 26673
        str cur = arr[0]
        int[:] count = np.zeros(n_groups, dtype="int")
        int cur_g = 0
        int i
        int n = arr.shape[0]
    count[cur_g] = 1
    for i in range(1, n):
        last = cur
        cur = arr[i]
        if cur != last:
            cur_g += 1
        count[cur_g] += 1
    return np.asarray(count)

#利用idx存储 DF应该 保留的行数
n = DF.shape[0]
@njit
def generate_idx(sample, count):
    idx = np.zeros(n, dtype=np.int8)
    start = 0
    for i in range(len(count)):
        step = sample[i]
        idx[start:start+step] = 1
        start += count[i]
    return idx

count = get_count(DF.MMSI_paragraph.values)#统计MMSI_paragraph的values，然后存储至count
sample = (count * np.random.uniform(0.3,0.7,count.shape[0])).astype("int")#根据count这个list对数据进行抽样
result = DF.loc[generate_idx(sample, count).astype("bool")]#按照行数对原始数据进行抽取