Python | LSTM：车辆轨迹预测

# 基于LSTM算法的车辆轨迹预测
# Developed by: gpluo


import pandas as pd
import numpy as np
import csv
import os
import matplotlib.pyplot as plt
import torch
from torch import nn
from sklearn.preprocessing import MinMaxScaler


# 创建文件夹
def dicBuild(path):
    # 去除首位空格
    path = path.strip()
    isExists = os.path.exists(path)
    if not isExists:
        os.makedirs(path)


# 创建CSV文件
def csvBuild(path, mat):
    file = open(path, mode='w', encoding='utf-8', newline='')
    csvWriter = csv.DictWriter(file, fieldnames=['Global_Time', 'Local_X', 'Local_Y'])
    csvWriter.writeheader()
    for item in mat:
        dit = {
            'Global_Time': item[0],
            'Local_X': item[1],
            'Local_Y': item[2]
        }
        csvWriter.writerow(dit)
    file.close()


# 原始数据加载
def loadOrg(vehicleID):
    dfOrg = pd.read_csv('Next_Generation_Simulation__NGSIM__Vehicle_Trajectories_and_Supporting_Data.csv')
    matOrg = np.array(dfOrg[['Vehicle_ID', 'Global_Time', 'Local_X', 'Local_Y']]).tolist()

    matTotalOrg = []
    for item in matOrg:
        if int(item[0]) == vehicleID:
            matTotalOrg.append([int(item[1]), float(item[2]), float(item[3])])

    # 对原始的时刻进行排序，使其递增
    matTotal = sorted(matTotalOrg, key=lambda ele: ele[0])

    # 在递增的条件下，找到间断点并拆分不同时间段的轨迹
    matTotal3D = []
    matFinalItem = matTotal[-1]  # 标记终点用于结束循环
    while True:
        if len(matTotal) <= 1:
            break
        matTotal3DEle = []
        for i in range(len(matTotal) - 1):
            matTotal3DEle.append([matTotal[i][0], matTotal[i][1], matTotal[i][2]])
            # 拆分：间断点或终点
            if ((matTotal[i + 1][0] - matTotal[i][0]) > 100) or (matTotal[i + 1] == matFinalItem):
                matTotal3D.append(matTotal3DEle)
                del matTotal[: i + 1]
                break
            # 去除重复数据
            elif (matTotal[i + 1][0] - matTotal[i][0]) == 0:
                i += 1

    # 每条轨迹改为从moment=0开始
    for matItem in matTotal3D:
        timeMin = matItem[0][0]
        for item in matItem:
            item[0] -= timeMin
        print('complete')

    return matTotal3D


# 数据降维：PCA算法
def methodPCA(arrayOrg):
    def plotScatter(arrayDeMean, recoveredArray):
        fig, ax = plt.subplots(figsize=(10, 10))
        ax.scatter(arrayDeMean[:, 0], arrayDeMean[:, 1])
        ax.scatter(list(recoveredArray[:, 0]), list(recoveredArray[:, 1]), c='r')
        ax.plot([arrayDeMean[:, 0], list(recoveredArray[:, 0])], [arrayDeMean[:, 1], list(recoveredArray[:, 1])])
        plt.show()

    # 去均值化
    def deMean(arrayOrg):
        arrayDeMean = (arrayOrg - np.mean(arrayOrg, axis=0))
        return arrayDeMean

    # 计算数据的协方差矩阵
    def sigmaMat(arrayDeMean):
        sigma = (arrayDeMean.T @ arrayDeMean) / arrayDeMean.shape[0]
        return sigma

    # 计算特征值、特征向量
    def usv(sigma):
        u, s, v = np.linalg.svd(sigma)
        return u, s, v

    # 对数据进行降维
    def dimReduction(arrayDeMean, u, k):
        reducedU = u[:, :k]
        z = np.dot(arrayDeMean, reducedU)
        return z

    # 数据还原
    def dataRecover(z, u, k):
        reducedU = u[:, :k]
        recoveredArray = np.dot(z, reducedU.T)
        return recoveredArray

    arrayDeMean = deMean(arrayOrg)
    sigma = sigmaMat(arrayDeMean)
    u, s, v = usv(sigma)
    z = dimReduction(arrayDeMean, u, 1)

    # 数据还原方法
    # recoveredArray = dataRecover(z, u, 1)

    # 降维前后对比
    # plotScatter(arrayDeMean, recoveredArray)

    return arrayDeMean, z, u


# 3sigma原则查找异常点并删除
def threeSigma(dfCol):  # dfCol: 数据的某一列
    rule = (dfCol.mean() - 3 * dfCol.std() > dfCol) | (dfCol.mean() + 3 * dfCol.std() < dfCol)
    index = np.arange(dfCol.shape[0])[rule]
    outrange = dfCol.iloc[index]
    return outrange


# 定义LSTM算法
class LSTMRNN(nn.Module):
    """
        Parameters：
        - input_size: feature size
        - hidden_size: number of hidden units
        - output_size: number of output
        - num_layers: layers of LSTM to stack
    """

    def __init__(self, input_size, hidden_size=1, output_size=1, num_layers=1):
        super().__init__()

        self.lstm = nn.LSTM(input_size, hidden_size, num_layers)  # utilize the LSTM model in torch.nn
        self.forwardCalculation = nn.Linear(hidden_size, output_size)

    def forward(self, _x):
        x, _ = self.lstm(_x)  # _x is input, size (seq_len, batch, input_size)
        s, b, h = x.shape  # x is output, size (seq_len, batch, hidden_size)
        x = x.view(s * b, h)
        x = self.forwardCalculation(x)
        x = x.view(s, b, -1)
        return x


# 训练函数
def trainCore(dataset):
    # 预测结果可视化
    def resultPlot():
        plt.figure()
        plt.plot(tForTraining, trainY, 'b', label='locTRN')  # 训练集原曲线
        plt.plot(t[trainLength:], testY[trainLength:], 'k', label='locTST')  # 测试集原曲线
        plt.plot(tForTest, predictiveYTest, 'm--', label='preLocTST')  # 预测曲线
        plt.plot([t[trainLength], t[trainLength]], [0, 1], 'r--', label='separation line')  # 训练集、测试集分界线
        plt.xlabel('t')
        plt.ylabel('locVar')
        plt.xlim(t[0], t[-1])
        plt.ylim(0, 1)
        plt.legend(loc='best')
        plt.grid()
        plt.show()

    # 确认GPU运算是否可行
    device = torch.device("cpu")
    if torch.cuda.is_available():
        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
        device = torch.device("cuda:0")
        print('Training of GPU.')
    else:
        print('No GPU available, training on CPU.')

    # 参数设置
    trainRatio = 0.95  # 训练集占比的约数
    batchSize = 16  # 每次抓取的样本数量
    inputFeaturesNum = 1  # 输入维度
    outputFeaturesNum = 1  # 输出维度
    hiddenSize = 64  # 隐含层的特征维度
    numLayers = 3  # 隐含层的层数(默认为1)
    maxEpochs = 10000  # 最大训练循环数
    learningRate = 1e-2  # 学习率
    lrGamma = 0.98  # 学习率衰减因子
    lossValue = 5e-5  # 均方误差允许值

    trainLength = int(len(dataset) * trainRatio - ((len(dataset) * trainRatio) % batchSize))  # 保证训练集长度为batch_size的倍数

    trainX = dataset[: trainLength, 0]
    trainY = dataset[: trainLength, 1]
    t = np.linspace(0, dataset[:, 1][-1], len(dataset))  # 时间序列间隔
    tForTraining = t[: trainLength]

    testX = dataset[:, 0]
    testY = dataset[:, 1]
    tForTest = t[:]

    # 将数组转换为1维
    trainXTensor = trainX.reshape(-1, batchSize, inputFeaturesNum)
    trainYTensor = trainY.reshape(-1, batchSize, outputFeaturesNum)

    # 把数组转化成张量(tensor)
    trainXTensor = torch.from_numpy(trainXTensor)
    trainYTensor = torch.from_numpy(trainYTensor)

    # 把张量传入device(GPU可用时，传入GPU)
    trainXTensor = trainXTensor.to(device)
    trainYTensor = trainYTensor.to(device)

    # LSTM建模
    lstmModel = LSTMRNN(input_size=inputFeaturesNum, hidden_size=hiddenSize, output_size=outputFeaturesNum,
                        num_layers=numLayers)
    lstmModel = lstmModel.train()
    lstmModel = lstmModel.to(device)
    print('LSTM model:', lstmModel)  # 打印模型参数

    # 均方误差
    lossFunction = nn.MSELoss()

    # 优化器
    optimizer = torch.optim.Adam(lstmModel.parameters(), lr=learningRate)
    torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=lrGamma)  # 学习率按指数衰减

    # 数据训练
    for epoch in range(maxEpochs):
        output = lstmModel(trainXTensor)
        loss = lossFunction(output, trainYTensor)

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if loss.item() < lossValue:
            print('Epoch [{}/{}], Loss: {:.5f}'.format(epoch + 1, maxEpochs, loss.item()))
            print("The loss value is reached")
            break
        elif (epoch + 1) % 100 == 0:
            print('Epoch: [{}/{}], Loss:{:.5f}'.format(epoch + 1, maxEpochs, loss.item()))

    # 保存模型系数
    torch.save(lstmModel.state_dict(), 'model_params.pkl')

    # 读取模型系数
    lstmModel.load_state_dict(torch.load('model_params.pkl'))

    # 趋势预测
    lstmModel = lstmModel.eval()  # 加载所有模型结构，调整至测试模型
    testXTensor = testX.reshape(-1, batchSize, inputFeaturesNum)
    testXTensor = torch.from_numpy(testXTensor)
    testXTensor = testXTensor.to(device)
    with torch.no_grad():  # 确保预测时梯度不变
        predictiveYTest = lstmModel(testXTensor).to(device)
    predictiveYTest = predictiveYTest.view(-1, outputFeaturesNum).data.cpu().numpy()

    # 结果可视化
    resultPlot()


if __name__ == '__main__':
    # 1 原始数据处理

    # 1.1 获取研究对象车辆的轨迹数据(研究对象不变时，以下数行无需运行)

    # itemID = 515  # 选取的车辆编号
    # matTotal = loadOrg(itemID)
    # dicBuild('TrajectoryFile')
    # subDic = 'TrajectoryFile\\' + str(itemID)
    # dicBuild(subDic)
    # for i in range(len(matTotal)):
    #     csvBuild(subDic + '\\TrajectoryInter' + str(i) + '.txt', matTotal[i])

    # 1.2 读取数据
    df = pd.read_csv('TrajectoryFile\\515\\TrajectoryInter1.txt')  # 对515号车来说，目前比较有代表性的是1、6号数据

    # 2 数据清洗、拆分和结构优化

    # 2.1 归一化
    scaler = MinMaxScaler()
    df['Local_X'] = scaler.fit_transform(df['Local_X'].values.reshape(-1, 1))
    df['Local_Y'] = scaler.fit_transform(df['Local_Y'].values.reshape(-1, 1))

    # 2.2 查找异常点并删除
    index_ = threeSigma(df['Local_X']).index
    df = df.drop(index_, axis=0)

    # 2.3 构建数组
    # 2.3.1 构建轨迹数组并降维
    arrayTrajectory = np.array(df[['Local_X', 'Local_Y']])  # 轨迹点数组
    arrayDeMean, z, u = methodPCA(arrayTrajectory)  # 数据降维处理
    # 2.3.2 构建时刻数组
    arrayMoment = np.array(df['Global_Time']).astype('float64')  # 时刻数组

    # 2.4 创建数据集
    formatLength = len(z) - (len(z) % 16)  # 控制数据集长度为8的倍数，方便设置batch_size
    dataset = np.zeros((formatLength, 2))  # 创建空数组以便接收数据
    dataset[:, 0] = arrayMoment[: formatLength]  # x: 时间指标
    z = scaler.fit_transform(z.reshape(-1, 1))  # 降维后轨迹数组的归一化
    zOneDim = z.reshape(-1)
    dataset[:, 1] = zOneDim[: formatLength]  # y: 轨迹指标
    dataset = dataset.astype('float32')

    # 2.5 创建差值数据集(计算差值前的数据集预测效果较好，暂不运行)
    # dataset = np.diff(dataset, axis=0)

    # 3 数据训练与结果输出
    trainCore(dataset)