python编程练习：提取Visual MODFLOW水均衡数据(.ZOT)文件至表格

一、功能介绍

1、将水均衡数据(.ZOT)文件中的各水均衡组分提取到表格（.xls）文件

• 按关键字进行数据分类，关键字优先级依次为：层、补\排、各均衡组分
• 合并每层中补给或排泄的越流项，并将原文件中各均衡组分的英文名称翻译为中文
• 数据导出至表格文件中的三个工作表“All”、“In”、“Out”中
  

2、绘制各层补给、排泄项水均衡组分扇形图

• 不绘制数值为0的水均衡项
• 单个扇形图中，占比最高的一项会突出显示
  

二、代码

import re
import pandas as pd
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 显示中文


def read_zot_file(in_zot):
    """ 寻找.ZOT文件中含有等号的行，将其中的空格删去，并将 :param in_zot: MODFLOW模拟后生成的水均衡文件(.ZOT) :return: 元素为字符串的列表，格式为["均衡组分名a,数值1","均衡组分名b,数值2",......] """
    with open(in_zot) as f:
        lines = [line.strip().replace(' ', "").replace('=', ',') for line in f if line.find("=") != -1]
    return lines


def class_by_layer(datas):
    """ :param datas: 格式为["均衡组分名a,数值1","均衡组分名b,数值2",......],由方法read_zot_file()方法 :return: in_comp：[[layer1补给项],[layer2补给项],[layer3补给项]] out_comp：[[layer1排泄项],[layer2排泄项],[layer3排泄项]] """
    # 各层级的关键词
    sep_layer = "PercentDiscrepancy"  # 层分隔标志符
    sep_in = "TotalIN"  # IN标志符
    sep_out = "TotalOUT"  # OUT标志符
    # 各级关键字所在的位置
    in_index = []
    out_index = []
    layer_index = []
    for i, v in enumerate(datas):
        if sep_in in v:
            in_index.append(i)
        elif sep_out in v:
            out_index.append(i)
        elif sep_layer in v:
            layer_index.append(i)
    # 分离各层
    index_s = 0
    in_comp = []
    out_comp = []
    for i in range(len(layer_index)):
        in_comp.append(datas[index_s:in_index[i]])
        out_comp.append(datas[in_index[i] + 1:out_index[i]])
        index_s = layer_index[i] + 1
    return in_comp, out_comp


def translate(eg_words, dict_trans):
    zh_words = {
   }
    # 补给项、排泄项的翻译字典，默认如下
    for index, words in enumerate(eg_words):
        list_temp = {
   }
        temp = 0
        for word in words:
            p = word.split(',')
            if p[0] in dict_trans.keys():
                # 如果字典存在该英文单词，翻译成对应中文
                list_temp[dict_trans[p[0]]] = p[1]
            elif re.match(r'Zone\dto\d+', p[0]):
                # 如果该项原格式为'Zone\dto\d+'，则记录其值至temp
                temp += float(word.split(',')[1])
            else:
                # 如果字典不存在该英文单词，不翻译
                list_temp[p[0]] = p[1]
        list_temp["越流"] = temp
        zh_words["layer{}".format(index+1)] = list_temp
    return zh_words


def draw_pieplots(df1, df2):
    # 绘制饼状图组图
    size = (2, df1.shape[1])
    plt.subplots(size[0], size[1])
    plt.figure(figsize=(15, 10))
    for i in range(size[0] * size[1]):
        plt.subplot(size[0], size[1], i + 1)
        target = "layer{}".format(i % size[1] + 1)
        if i // size[1] == 0:
            df = df1[df1[target] != 0][target]
        else:
            df = df2[df2[target] != 0][target]
        highlight = list(df).index(max(df))
        explode = [0.05 if j == highlight else 0 for j in range(len(df.index))]
        plt.pie(x=df, explode=explode, labels=df.index, autopct="%.3f%%")
        plt.title("{}".format(target))
        plt.legend()


if __name__ == "__main__":
    dict_in = {
   'CONSTANTHEAD': "定水头",
                "WELLS": "井",
                "RIVERLEAKAGE": "河流入渗",
                "ET": "潜水蒸发",
                "RECHARGE": "降雨入渗",
                "STORAGE": "储存量",
                }
    dict_out = {
   'CONSTANTHEAD': "定水头",
                "WELLS": "井",
                "RIVERLEAKAGE": "河流排泄",
                "ET": "潜水蒸发",
                "RECHARGE": "降雨入渗",
                "STORAGE": "释放量",
                }
    # 输入zot文件，并读取
    # ！！！在此更改要转换的（.ZOT）文件路径，例如C:\sim_2\XXX.ZOT
    in_path = "zot3.ZOT"
    zufen = read_zot_file(in_path)
    # 分离补给项、排泄项
    supply, consume = class_by_layer(zufen)
    # 将其中的英语单词翻译成中文
    df_in = pd.DataFrame(translate(supply, dict_in), dtype='float')
    df_out = pd.DataFrame(translate(consume, dict_out), dtype='float')
    # 绘制饼状图组图
    draw_pieplots(df1=df_in, df2=df_out)
    # 计算补给项和排泄项的总和
    df_in.loc["总补给"] = df_in.sum()
    df_out.loc["总排泄"] = df_out.sum()
    df_all = pd.concat([df_in, df_out], axis=0)
    df_all.loc["补排差"] = df_all.loc["总补给"] - df_all.loc["总排泄"]
    # 导出至excel
    # ！！！在此更改要转换的（.ZOT）文件路径，此处为zot2.xls
    with pd.ExcelWriter(r'zot2.xls') as writer:
        df_all.to_excel(writer, sheet_name="All")
        df_in.to_excel(writer, sheet_name="In")
        df_out.to_excel(writer, sheet_name="Out")
    # 绘图，可选
    plt.show()