python处理数据——画2D、3D散点图

本文，是与我所做的数据处理相联系，前面两部分，是数据转换与计算，第三部分是2D图，与3D图的实现

第一部分 数据预处理工作，若只参考散点图，此部分可略过

import os
import pandas as pd
import sys

path="D:\\SwarmData\\datazor\\data"
def readfile(filepath):
    datafilepath = "D:\\SwarmData\\datazor\\result\\"+ filepath[31:48]+'_trans.csv'
    pickfilepath = "D:\\SwarmData\\datazor\\result250\\"+ filepath[31:48]+'_step250_simu1.csv'
    size = pd.read_csv(filepath, nrows=0)   #取出文件中的第一行数据
    swarmsize = float(((str(list(size.columns)[0]).split(" "))[0].split(":"))[1])  #取出文件中，第一行数据中的参数
    zor=float(((str(list(size.columns)[0]).split(" "))[1].split(":"))[1])

    data = pd.read_csv(filepath, header=1)  #header=1，跳过文件第一行的数据
    # data=pd.concat([data, data['step'].str.split(' ', expand=True)], axis=1, names=new_names)
    data=data['simu 01'].str.split(' ', expand=True)
    data.columns = ['step','values']
    # data = data['values'].str.split('/', expand=True)
    newname=['step','flight','number','x','y','z']
    data = pd.concat([data['step'], data['values'].str.split('/', expand=True)], axis=1)
    data.columns=newname
    outdata=data[['step','flight','x','y','z']]
    outdata['swarmsize'] = swarmsize
    outdata['zor'] = zor
    outdatastep250 = outdata.loc[outdata['step']=='250.0']
    #写入转换后的数据到文件
    outdatastep250.to_csv(datafilepath)

    #说明一下，因为转换后的数据会多一列'Unnamed: 0'，且行索引不连续
    # 所以该函数实现从文件中，读取数据，在写入，就能解决以上问题
    def pickfile(pickfilepath):
        pickdata = pd.read_csv(datafilepath)
        pickdata.drop('Unnamed: 0', axis=1, inplace=True)
        length=int((pickdata['swarmsize'].tolist())[0])
        pickdatasimu1=pickdata.iloc[0:length,:]
        # print(pickdatasimu1)
        #重新写入数据到文件
        pickdatasimu1.to_csv(pickfilepath)
    pickfile(pickfilepath)

def func(path):
    files = os.listdir(path)
    files.sort()
    for file in files:
        readfile(path+"\\"+file)

if len(sys.argv) > 1:
    filepath = sys.argv[1]
func(path)

第二部分 数据计算工作，从数据中运算得到质心，在写入文件

import os
import pandas as pd
import sys
import numpy as np

#该模块实现质心的计算，并将计算后的质心，加入到文件数据中，重新写入到文件保存
path="D:\\SwarmData\\datazor\\result250"
def readfile(filepath):
    data = pd.read_csv(filepath)
    X=data['x'].tolist()
    Y=data['y'].tolist()
    Z=data['z'].tolist()
    centroidx=np.mean(X)
    centroidy=np.mean(Y)
    centroidz=np.mean(Z)
    test=data.iloc[0:1,:]
    test['x']=centroidx
    test['y']=centroidy
    test['z']=centroidz
    # test['Unnamed: 0']='centriods'
    data= pd.concat([data, test], axis=0, ignore_index=True)
    data.to_csv("D:\\SwarmData\\datazor\\resultcentroid\\"+filepath[-35:-4]+'_centriod.csv')

def func(path):
    files = os.listdir(path)
    files.sort()
    for file in files:
        readfile(path+"\\"+file)

if len(sys.argv) > 1:
    filepath = sys.argv[1]
func(path)

第三部分 用散点图，可视化

import pandas as pd
import matplotlib.pyplot as plt
import os
import sys
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

def draw2d(filepath):
    data = pd.read_csv(filepath)
    X=data['x'].tolist()
    Y=data['y'].tolist()
    Z=data['z'].tolist()
    zor = int((data['zor'].tolist())[0])
    swarmsize=int((data['swarmsize'].tolist())[0])
    # 运算得到以质心为零点
    #质心是data的最后一行
    length = len(X) - 1
    valuex=X[length]
    for i in range(length+1):
        X[i]=X[i]-valuex
    valuey = Y[length]
    for i in range(length+1):
        Y[i] = Y[i] - valuey
    valuez = Z[length]
    for i in range(length+1):
        Z[i] = Z[i] - valuez

    fig=plt.figure()

    fig.suptitle('swarmsize:'+str(swarmsize)+'  zor:'+str(zor))
    ax1 = fig.add_subplot(2, 2, 1)
    ax2 = fig.add_subplot(2, 2, 2)
    ax3 = fig.add_subplot(2, 2, 3)

    ax1.set_title('X-Y view', fontsize=10)
    ax1.set_xlabel('X', fontsize=10)
    ax1.set_ylabel('Y', fontsize=10)
    ax1.set_xlim(-200, 200)
    ax1.set_ylim(-200, 200)

    ax2.set_title('X-Z view', fontsize=10)
    ax2.set_xlabel('X', fontsize=10)
    ax2.set_ylabel('Z', fontsize=10)
    ax2.set_xlim(-200, 200)
    ax2.set_ylim(-200, 200)

    ax3.set_title('Y-Z view', fontsize=10)
    ax3.set_xlabel('Y', fontsize=10)
    ax3.set_ylabel('Z', fontsize=10)
    ax3.set_xlim(-200, 200)
    ax3.set_ylim(-200, 200)

    # ax1 = plt.axes(projection='3d')
    ax1.scatter(X[0:length],Y[0:length], s=20, color='b')
    ax1.scatter(X[length:], Y[length:], s=50, color='r',label='centroid')

    # ax2 = plt.axes(projection='3d')
    ax2.scatter(X[0:length], Z[0:length],s=20, color='b')
    ax2.scatter(X[length:], Z[length:],s=50, color='r', label='centroid')

    ax3.scatter(Y[0:length], Z[0:length], s=20, color='b')
    ax3.scatter(Y[length:], Z[length:], s=50, color='r', label='centroid')

    ax1.legend(loc='best')
    ax2.legend(loc='best')
    ax3.legend(loc='best')
    #自动保存图片到文件
    plt.savefig("D:\\SwarmData\\datazor\\image\\" + "swarmsize" + str(swarmsize) + "zor" + str(zor) + ".png",format='png')
    plt.show()

def draw3d(filepath):
    data = pd.read_csv(filepath)
    X = data['x'].tolist()
    Y = data['y'].tolist()
    Z = data['z'].tolist()
    zor = int((data['zor'].tolist())[0])
    swarmsize = int((data['swarmsize'].tolist())[0])
    length = len(X) - 1
    valuex = X[length]
    for i in range(length + 1):
        X[i] = X[i] - valuex
    valuey = Y[length]
    for i in range(length + 1):
        Y[i] = Y[i] - valuey
    valuez = Z[length]
    for i in range(length + 1):
        Z[i] = Z[i] - valuez
    ax = plt.axes(projection='3d')
    ax.scatter3D(X[0:length-1],Y[0:length-1],Z[0:length-1], s=20, color='b')
    ax.scatter3D(X[length:], Y[length:], Z[length:], s=50, color='r', label='centriod')
    ax.legend(loc='best')
    #自动保存图片到文件
    plt.savefig("D:\\SwarmData\\datazor\\image\\" + "swarmsize" + str(swarmsize) + "zor" + str(zor) + ".png",format='png')
    plt.show()

def func(path):
    files=os.listdir(path)
    for file in files:
        draw2d(path+"\\"+file)
# filepath1 = "D:\\SwarmData\\datazor\\resultcentroid\\size30zor10zoo25z_step250_simu1_centriod.csv"
# filepath2 = "D:\\SwarmData\\datazor\\resultcentroid\\size30zor40zoo25z_step250_simu1_centriod.csv"
# filepath3 = "D:\\SwarmData\\datazor\\resultcentroid\\size200zor20zoo25_step250_simu1_centriod.csv"
# filepath4 = "D:\\SwarmData\\datazor\\resultcentroid\\size200zor50zoo55_step250_simu1_centriod.csv"
# filepath5 = "D:\\SwarmData\\datazor\\resultcentroid\\size50zor10zoo25z_step250_simu1_centriod.csv"
# filepath6 = "D:\\SwarmData\\datazor\\resultcentroid\\size50zor50zoo25z_step250_simu1_centriod.csv"
# filepath7 = "D:\\SwarmData\\datazor\\resultcentroid\\size100zor10zoo25_step250_simu1_centriod.csv"
# filepath8 = "D:\\SwarmData\\datazor\\resultcentroid\\size100zor40zoo95_step250_simu1_centriod.csv"

filepath="D:\\SwarmData\\datazor\\resultcentroid"
if len(sys.argv) > 1:
    filepath = sys.argv[1]
func(filepath)

实现效果，是在保存图片目录下的截图

2d效果图

3d效果图