python 可视化plotly 等高线热力图

python 可视化plotly 等高线热力图

CSV文件数据如下，要求将第1列转速设置为x轴，第2列扭矩设置为y轴，第3列效率为等高线。

思路分析：
1)在控制端输入文件路径，读取csv文件，该代码只对三列数据进行处理，多于或少于三列都会报错退出，且数据行只有第一行，且其余是数据行。对数据进行空数据行删除处理。
2）数据的预处理，x是第1列数据，y是第2列数据，画等高线Z则是一个数组，其数据格式应该如下，即相同扭矩下不同转速的效率组成的数组。也就是CSV文件格式要转化为以下格式。首先要把所有转速和扭矩下的效率全部显示出来。需要用到笛卡尔积， 即x(1000-16000一个16个值) 与y(20-300 15个值)生成一个只含有X和Y两列的数据集，再与原数据合并生成新的数据集，没有效率值则自动生成NAN值。这时再对NAN值进行赋予0值。再遍历y从20-300，获取每个y下面不同的效率值是一个list,同时又是列表index_list1的值。
3）将2中获取的x,y, index_list1传入画图的函数中，进行plotly作图。

# -*- coding:utf8 -*-
import numpy as np
import pandas as pd
import plotly.graph_objs as go
import plotly.offline as py  #设置离线画图


""" 打开原始CSV文件数据，并且检查输入查询的列名是否存在于文件中，
    若存在，则提取这些列的数据，返回新的dataframe,
    若不存在，则报错不执行程序
    返回值： new_Df
"""
def check_file_column():
    input_file = input('请输入csv文件路径:')
    filename = input_file.strip()
    with open(filename, 'rb') as f:
        csv_data = pd.read_csv(f, encoding='gb2312')
    if len(csv_data.columns) != 3:
        print('文件列名不等于3列，请修改文件!')
        exit()
    correct_list = list(csv_data.columns.values)
    # 删除空的数据行
    new_Df = csv_data.dropna()
    # 将导入的各个类型变成正确的类型
    new_Df[new_Df.columns.values[0]] = new_Df[new_Df.columns.values[0]].apply(int)
    new_Df[new_Df.columns.values[2]] = new_Df[new_Df.columns.values[2]].apply(float)
    new_Df[new_Df.columns.values[1]] = new_Df[new_Df.columns.values[1]].apply(float)
    # 将new_Df进行排序
    new_Df.sort_values(new_Df.columns.values[0], inplace=True)
    return new_Df, correct_list


""" 数据的处理：
    传入参数x 是第一列去重列表，参数y是第二列去重列表，csv_data是提取相应列名的数据集
    将x,y调用函数笛卡尔积生成新的dataframe,并且与csv_data合并，给空值赋0
    将所有转矩的效率提取到一个二维数组index_list1
    返回： index_list1 
"""
def handle_data(x,y,csv_data):
    new_pd = itertool_func(x, y,csv_data)
    new_pd[csv_data.columns.values[0]] = new_pd[csv_data.columns.values[0]].apply(int)
    new_pd[csv_data.columns.values[1]] = new_pd[csv_data.columns.values[1]].apply(int)
    # 合并两个dataframe，与原来的数据集为基础
    df3 = pd.merge(csv_data, new_pd, how='right', on=list(new_pd.columns.values))
    # 给空值赋予0
    df4 = df3.fillna(value=0)
    index_list1 = []
    for y1 in y:
        index_list1.append(df4.loc[df4[csv_data.columns.values[1]] == y1][csv_data.columns.values[2]].values)
    return index_list1


""" 传入参数A是x数组， 参数B是y数组
将两个数组进行笛卡尔积操作生成新的数据集"""
def itertool_func(A,B,csv_data):
    lenB = len(B)
#     print(lenB)
    new_pd = pd.DataFrame(columns=csv_data.columns.values[:2])
    for a in A:
        curA = np.array([a]*lenB)
        curA.shape = (lenB,1)
        curB = np.array(B)
        curB.shape = (lenB,1)
        join_h = np.hstack((curA, curB))
        new_pd = new_pd.append(pd.DataFrame(join_h, columns=csv_data.columns.values[:2]), ignore_index=True)
    return new_pd


"""传入参数x，参数y,参数z是二维数组，反映效率
画等高线图
"""
def plot_color(x, y, z):
    fig = go.Figure(
        data=go.Contour(
            z=z,
            x=x,  # horizontal axis
            y=y,  # vertical axis
            # colorscale=[[0, 'white'], [0.5, 'red'], [1.0, 'rgb(0, 0, 255)']],
            colorscale=[[0, 'white'], [0.5, 'red'], [1.0, 'darkred']],
            contours=dict(
                start=60,    # colorbar的最小刻度值
                end=100,   #colorbar的最大刻度值
                size=4,      #即每个4画一个等高线
                showlabels=True,  # show labels on contours
                labelfont=dict(  # label font properties
                    size=12,
                    color='white',
                ),
            ),

        ))
    fig.update_layout(  # all "layout" attributes: /python/reference/#layout
        title=dict(text="Motor Efficiency Presented by matplotlib", font=dict(family='Arial', size=24, color='black')),
        xaxis=dict(  # all "layout's" "xaxis" attributes: /python/reference/#layout-xaxis
            title="Speed",  # more about "layout's" "xaxis's" "title": /python/reference/#layout-xaxis-title
            titlefont=dict(color='black', size=24)  # 设置坐标轴标签的字体及颜色
        ),
        yaxis=dict(
            title='Torque',
            titlefont=dict(color='black', size=24)  # 设置坐标轴标签的字体及颜色

        ),
    )
    fig.show()
    py.plot(fig, filename='pictures/contour.html',  # 会生成一个网页文件
            image='png', )  # 设置保存的文件类型，不会在本地有个png的文件，需要在生成的网页打开另存为png的文件


def main():
    csv_data, correct_list = check_file_column()
    # 选取第一列的所有值并去重并排序，设置为X坐标，
    col_1 = csv_data.iloc[:,0]
    x = sorted(list(set(col_1.values)))
    # 选取第2列所有值并去重并排序，设置为Y坐标
    col_2 = csv_data.iloc[:,1]
    y = sorted(list(set(col_2.values)))
    # 处理数据
    index_list1 = handle_data(x,y,csv_data)
    # 画图
    plot_color(x, y, index_list1)


if __name__ == '__main__':
    main()

最后生成一个网页，鼠标放在图上会显示x,y,z的相应值，而且网页右上方有照相功能，点击之后浏览器能够自动下载png图片。上面的代码 py.plot有本地离线下载网页的功能。