Txgy.WFLEA/main_2.py

# -*- coding: utf-8 -*-
import sys
from PySide6.QtWidgets import QApplication, QMainWindow, QGraphicsView, QFileDialog, QColorDialog
from PySide6.QtUiTools import QUiLoader
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
from PySide6.QtWidgets import QVBoxLayout
import CREAT
import numpy as np
import matplotlib as mpl
from matplotlib.path import Path
from matplotlib.patches import Patch
from PySide6.QtGui import QPalette
import random
# plt.legend(prop={'family': 'SimHei', 'size': 15})


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.toolbar = None
        self.canvas = None
        self.file_paths = []
        self.get_color = lambda n: list(map(lambda i: "#" + "%06x" % random.randint(0, 0xFFFFFF),range(n)))
        self.color = self.get_color(100)  # 图例颜色选择值
        self.colorbar = 'jet'
        self.num = len(self.color)  # 记录选择颜色的次数
        self.oilwell_color = 'k'
        self.waterwell_color = 'k'
        self.chose = 0
        self.num = 0
        # 加载UI文件
        self.ui = QUiLoader().load('QT/main_2.ui')
        self.color_bar()
        self.ui.legend_color.clicked.connect(self.legend_color_chose)  # 选择图例颜色
        self.ui.colorbar_chose.currentTextChanged.connect(self.update_colorbar)  # 更新colorbar
        self.ui.chose_oilwell_color.clicked.connect(lambda: self.oil_well_color())  # 选择油井颜色
        self.ui.chose_oilwell_color.clicked.connect(lambda: self.plot_density())  # 更新油井颜色
        self.ui.chose_waterwell_color.clicked.connect(lambda: self.water_well_color())  # 选择水井颜色
        self.ui.chose_waterwell_color.clicked.connect(lambda: self.plot_density())  # 更新水井颜色
        # 更新图窗
        self.ui.graphicsView.repaint()
        self.ui.graphicsView.update()
        self.graphics_layout = QVBoxLayout(self.ui.graphicsView)

        self.ui.chose_datafile_button.clicked.connect(self.choose_datafile)

        self.ui.chose_wellfile_button.clicked.connect(self.choose_wellfile)

        # 绘制或清除数据和画布
        self.ui.clear_data.clicked.connect(lambda: self.clear_data())

        # 绘图设置中心比例
        self.ui.center_scale.clicked.connect(lambda: self.center_scale())
        self.ui.center_scale.clicked.connect(lambda: self.plot_density())
        self.ui.level_chose.valueChanged.connect(lambda: self.plot_density())
        self.ui.gWeight_chose.valueChanged.connect(lambda: self.plot_density())

    def self_levels(self):
        self.levels = self.ui.level_chose.value()  # 选择层参数
        self.plot_density()

    def self_gWeight(self):
        self.gWeight = self.self.ui.gWeight_chose.value()
        self.plot_density()

    def scene_fig(self):
        fig = Figure()
        # 在Figure对象中添加子图
        ax = fig.add_subplot(111)
        return fig, ax

    def canvas_adjust(self, fig):
        if self.canvas is not None:
            # 从布局中删除旧的 canvas 和 toolbar
            item = self.graphics_layout.takeAt(0)
            while item:
                widget = item.widget()
                if widget:
                    widget.setParent(None)
                item = self.graphics_layout.takeAt(0)
        self.graphics_layout.removeWidget(self.canvas)
        self.canvas = FigureCanvas(fig)
        self.toolbar = NavigationToolbar(self.canvas, self.ui.graphicsView)
        self.toolbar.setParent(self.canvas)
        self.graphics_layout.addWidget(self.canvas)
        self.adjustSize()

    def choose_datafile(self):
        file_dialog = QFileDialog(self)
        file_dialog.setFileMode(QFileDialog.ExistingFiles)
        if file_dialog.exec():
            self.file_paths += file_dialog.selectedFiles()
            self.ui.data_file_lineEdit.setText(','.join(self.file_paths))
        self.plot_density()

    def choose_wellfile(self):
        file_dialog = QFileDialog(self)
        file_dialog.setFileMode(QFileDialog.ExistingFiles)
        if file_dialog.exec():
            file_paths = file_dialog.selectedFiles()
            self.ui.well_file_lineEdit.setText(','.join(file_paths))
        self.plot_density()

    def legend_color_chose(self):
        col = QColorDialog.getColor()
        self.color[self.num] = col.name()
        self.num += 1
        self.plot_density()

    def oil_well_color(self):
        col = QColorDialog.getColor()
        self.oilwell_color = col.name()

    def water_well_color(self):
        col = QColorDialog.getColor()
        self.waterwell_color = col.name()

    def color_bar(self):
        self.ui.colorbar_chose.addItem("jet")
        self.ui.colorbar_chose.addItem("viridis")
        self.ui.colorbar_chose.addItem("coolwarm")
        self.ui.colorbar_chose.addItem('plasma')
        self.ui.colorbar_chose.addItem("magma")
        self.ui.colorbar_chose.addItem("inferno")

    def update_colorbar(self):
        colorbar_name = self.ui.colorbar_chose.currentText()
        self.colorbar = plt.get_cmap(colorbar_name)
        self.plot_density()

    def clear_data(self):
        self.file_paths = []
        self.color = self.get_color(100)
        self.num = 0
        self.ui.data_file_lineEdit.clear()
        self.ui.well_file_lineEdit.clear()
        self.canvas.figure.clf()  # 清除画布上的内容
        self.canvas.draw()

    def center_scale(self):
        if self.chose == 0:
            self.chose = 1
        else:
            self.chose = 0
        self.num = 0

    def plot_density(self):
        fig, ax = self.scene_fig()
        data_name = self.ui.data_file_lineEdit.text().split(',')  # 文件输入 =================1
        level_nums = self.ui.level_chose.value()  # 层参数 =================2
        gWeight = self.ui.gWeight_chose.value()   # 权重输入 =================3
        data = CREAT.create()
        legend_elements = []
        if self.num == len(data_name):
            self.num = 0
        for i in range(len(data_name)):
            data.data_pre(data_name[i])
            data.contours_pre(level_nums, gWeight)

            """画出密度等高线"""
            f = data.f
            x_min, x_max = data.x_range[0], data.x_range[1]
            y_min, y_max = data.y_range[0], data.y_range[1]
            levels = data.levels

            xx, yy = np.mgrid[x_min:x_max:200j, y_min:y_max:200j]
            # 填充等高线图中间的区域

            if len(data_name) > 1:
                color_i = self.color[i]  # """ 自动选择图例的颜色""" =================4
                level = [levels[0], levels[-1]]
                ff = ax.contourf(xx, yy, f, levels=level, colors=color_i, alpha=1, zorder=len(data_name) - i)
                # 生成图例
                legend_elements.append(Patch(facecolor=color_i, label='第'+f'{i+1}'+'次'))
                if i == len(data_name)-1:
                    ax.legend(handles=legend_elements, loc='upper right', prop={'family': 'SimHei', 'size': 10})
                    # ax.legend(prop={'family': 'SimHei', 'size': 15})
            else:
                # colorbars = 'jet'  # 输入选择的colorbar =================5
                cmap = mpl.colormaps.get_cmap(self.colorbar)
                colors = cmap(np.linspace(0, 1, level_nums))
                ff = ax.contourf(xx, yy, f, levels=levels, colors=colors[0:level_nums], alpha=0.5,
                                 zorder=len(data_name) - i)
                # 生成colorbar
                sm = mpl.cm.ScalarMappable(cmap=cmap, norm=mpl.colors.Normalize(vmin=ff.zmin, vmax=ff.zmax))
                sm.set_array([])
                fig.colorbar(sm, ax=ax)

        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_title('Density Distribution')

        wells_name = self.ui.well_file_lineEdit.text()  # 井文件输入 =================6
        if wells_name == '':
            self.canvas_adjust(fig)
        else:
            wells = CREAT.well_to_edge()
            wells.wells_name_and_position(wells_name)
            wells.welltoedge_distance(data.contours)

            """画井位信息"""
            typee = wells.type
            points = wells.position
            namee = wells.name
            min_distance = wells.min_distance
            contours_p = wells.welltoedge_points
            x_len = ((np.array(points)[:, 0].max()-np.array(points)[:,0].min())//100+2)*100
            y_len = ((np.array(points)[:, 1].max()-np.array(points)[:,1].min())//100+2)*100
            if self.chose == 1:
                max_x = max(np.array([(np.array(points)[:, 0].max()//100+1)*100, (np.array(points)[:, 1].max()//100+1)*100, data.x_range[1], data.y_range[1]]))
                min_x = max(np.array([(np.array(points)[:, 0].min()//100-1)*100, (np.array(points)[:, 1].min()//100-1)*100, data.x_range[0], data.y_range[0]]))
                max_x = max(max_x, -min_x)
                ax.set_xlim([-max_x, max_x])
                ax.set_ylim([-max_x, max_x])
                x_len = max(x_len, y_len)
                y_len = x_len

            for i in range(len(points)):
                is_inside = False  # 标记点是否在等高线内部
                # path = Path(contours_p)
                # if path.contains_point(points[i]):
                if np.linalg.norm(contours_p[i]) >= np.linalg.norm(points[i]):
                    is_inside = True
                if typee[i][0] == 0:
                    # oil_color = 'black'  # 输入油井颜色选择 =================8
                    ax.scatter(points[i][0], points[i][1], marker='o', edgecolors=self.oilwell_color, facecolors='none', s=70)
                    ax.scatter(points[i][0], points[i][1], marker='o', edgecolors=self.oilwell_color, facecolors='none', s=40,
                               linewidths=0.5)
                    ax.scatter(points[i][0], points[i][1], marker='o', edgecolors=self.oilwell_color, facecolors='none', s=20)
                    if wells.angle[i] <= 90:
                        ax.text(points[i][0] - 1/35*x_len, points[i][1] - 1/30*y_len, f'{namee[i][0]}',
                                fontdict={'size': '8', 'color': 'b'})  # 井名信息
                    elif wells.angle[i] > 270:
                        ax.text(points[i][0] + 3/175*x_len, points[i][1] - 1/40*y_len, f'{namee[i][0]}',
                                fontdict={'size': '8', 'color': 'b'})  # 井名信息
                    else:
                        ax.text(points[i][0] - 3/80*x_len, points[i][1] + 3/175*y_len, f'{namee[i][0]}',
                            fontdict={'size': '8', 'color': 'b'})  # 井名信息
                    if not is_inside:
                        # 绘制油井边缘最短距离点的箭头并标出距离
                        ax.quiver(points[i][0], points[i][1], contours_p[i][0] - points[i][0],
                                  contours_p[i][1] - points[i][1],
                                  angles='xy', scale=1.03,
                                  scale_units='xy', width=0.002, zorder=len(data_name) + 1)  # 绘制箭头
                        if np.linalg.norm(np.array(points[i])-np.array(contours_p[i])) > 50:
                            ax.text(points[i][0] * 1 / 2 + contours_p[i][0] * 1 / 2 + min_distance[i] / 18,
                                    points[i][1] * 1 / 2 + contours_p[i][1] * 1 / 2 - min_distance[i] / 18,
                                    f'{round(min_distance[i], 2)}m',
                                    fontdict={'size': '8', 'color': 'm'})  # 标出距离
                        else:
                            if wells.angle[i] <= 90:
                                ax.text(points[i][0]-3/70*x_len, points[i][1]+1/40*y_len, f'{round(min_distance[i], 2)}m',
                                        fontdict={'size': '8', 'color': 'm'})  # 标出距离
                            elif wells.angle[i] > 270:
                                ax.text(points[i][0] - 1/35*x_len, points[i][1] + 1/30*y_len, f'{round(min_distance[i], 2)}m',
                                        fontdict={'size': '8', 'color': 'm'})  # 标出距离
                            else:
                                ax.text(points[i][0] + 3/140*x_len, points[i][1] - 1/40*y_len, f'{round(min_distance[i], 2)}m',
                                        fontdict={'size': '8', 'color': 'm'})  # 标出距离
                else:
                    # water_color = 'black'  # 水井颜色 =================9
                    ax.scatter(points[i][0], points[i][1], marker='o', edgecolors=self.waterwell_color, facecolors=self.waterwell_color, s=70, zorder=len(data_name)+1)
                    ax.text(points[i][0] + 15, points[i][1] - 15, f'{namee[i][0]}',
                            fontdict={'size': '8', 'color': 'b'}, zorder=len(data_name) + 2)  # 井名信息
                    continue
            self.canvas_adjust(fig)


if __name__ == '__main__':
    app = QApplication([])
    stats = MainWindow()
    stats.ui.show()
    app.exec()