Tools/tool_lib/plots.py

import sqlite3
import time
from typing import Union
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
import sys
import matplotlib.animation


class Plots():
    def __init__(self, x_axe, y_axe, z_axe=[], x_label='x', y_label='y', z_label='z', x_axe2=[], y_axe2=[], z_axe2=[], grid=True,
                 figsize=(13, 10)):
        self.x_axe = x_axe
        self.y_axe = y_axe
        self.z_axe = z_axe
        self.x_label = x_label
        self.y_label = y_label
        self.z_label = z_label
        self.x_axe2 = x_axe2
        self.y_axe2 = y_axe2
        self.z_axe2 = z_axe2
        self.grid = grid
        self.figure = plt.figure(figsize=figsize)

    def get_a_exponential_func_as_plot(self, linrange=[-5, 5], expo=2):
        x = np.linspace(linrange[0], linrange[1])
        y = x ** expo
        plt.plot(x, y, 'g')

    def plot_2D_compare_bar_chart(self, legend=[], width=0.35, title=str, path=''):
        x = np.arange(len(self.x_axe))

        if legend:
            plt.bar(x, self.y_axe, width=width, color='blue', label=legend[0])
            plt.bar(x + width, self.y_axe2, width=width, color='red', label=legend[1])
        else:
            plt.bar(x, self.y_axe, width=width, color='blue')
            plt.bar(x + width, self.y_axe2, width=width, color='red')

        if self.grid:
            plt.grid()

        plt.xlabel(self.x_label)
        plt.ylabel(self.y_label)

        if title:
            plt.title(title)

        plt.xticks(x + width / 2, self.x_axe)
        plt.legend(loc='best')

        if path:
            self.save_fig(name='plot_2D_compare_bar_chart', path=path)
        else:
            self.save_fig(name='plot_2D_compare_bar_chart')
        plt.show()

    def save_fig(self, name, path: Union[bool, str] = False, ):
        if path:
            plt.savefig(path + '{}_{}.png'.format(name, time.strftime("%Y-%m-%d_H%H-M%M")))
        else:
            plt.savefig('{}_{}.png'.format(name, time.strftime("%Y-%m-%d_H%H-M%M")))


# Constant Layout
NODE_SIZE = 200
NODE_EDGE_COLOR = 'black'
EDGE_WIDTH = 0.5
FONT_SIZE = 8
FONT_SIZE_EDGES = 1
FONT_FAMILY = 'sans-serif'
SAVE_FORMAT = 'svg'
DPI = 1200
CONNECTION_STYLE = 'arc3, rad=0.2'
ARROW_SIZE = 12
LABLE_POS = 0.35


class NetworkxPlots():
    def __init__(self, node_dict, pathing=[], color_map=[], legend=[], edges=[], edge_weightings=[], directed_graph=False):
        if directed_graph:
            self.graph = nx.DiGraph()
        else:
            self.graph = nx.Graph()

        # Nodes
        self.node_type_dict = {}
        self.node_color_dict = {}
        self.node_pos_dict = {}

        for key, value in node_dict.items():
            for ckey, cvalue in node_dict.items():
                if node_dict[key]['node_type'] == node_dict[ckey]['node_type'] and node_dict[key][
                    'node_type'] not in self.node_type_dict.keys():
                    self.node_type_dict[node_dict[key]['node_type']] = []
            self.node_type_dict[node_dict[key]['node_type']].append(key)
            self.node_color_dict[node_dict[key]['node_type']] = node_dict[key]['node_color']
            self.node_pos_dict[key] = node_dict[key]['node_name']

        # Edges can be a nxn Matrix
        self.edges = edges
        self.pathing = pathing
        self.color_map = color_map
        self.edge_weightings = edge_weightings
        self.legend = legend

    def edges_for_complete_graph(self):
        # without self_loops
        for row in range(0, len(self.edges[:, 0])):
            for column in range(0, len(self.edges[0, :])):
                if round(self.edges[row, column], 1) == 0:
                    pass
                else:
                    self.graph.add_edge(row + 1, column + 1, weight=1)  # round(self.edges[row, column], 1)

    def directions_for_directed_graph(self, pathing=[]):
        if pathing:
            for order, path in enumerate(pathing):
                self.graph.add_edge(path[0], path[1], weight=order + 1)
        else:
            for order, path in enumerate(self.pathing):
                self.graph.add_edge(path[0], path[1], weight=order + 1)

    def add_nodes_to_graph(self):
        node_numberation = [node_number for node_number in range(1, len(self.node_pos_dict.keys()) + 1)]
        for count in node_numberation:
            self.graph.add_node(count)

    def add_edges_to_graph(self):
        self.graph.add_edges_from(self.edges, weight=self.edge_weightings)

    def undirected_graph_plt_(self,
                              name_path_tupl=''):  # TODO SIMPLYFY THESE FUNCTIONS BY EXTRA FUNCTIONS WHICH SEPERATES THEM INTO SMALLER ONES
        '''
        :param name_path_tupl: (name, path) to save the pic
        :return: a showing of the undirected graph/ saves the picture
        '''
        # Setting
        plt.axis("on")
        ax = plt.gca()
        self.create_edgeless_graph_plot(ax)
        # Undirected labels
        elarge = [(u, v) for (u, v, d) in self.graph.edges(data=True) if d["weight"]]
        nx.draw_networkx_edges(self.graph, self.node_pos_dict, edgelist=elarge, width=EDGE_WIDTH, ax=ax)

        # Create Figure
        self.create_graph_settings_with_cartesic_coord(plt, ax)
        # Save
        if name_path_tupl:
            self.save_fig(name=name_path_tupl[0], path=name_path_tupl[1], format=SAVE_FORMAT, dpi=DPI)
        # plt.gca().set_aspect('equal', adjustable='box')
        plt.show()

    def directed_graph_with_path_plt_(self, name_path_tupl=''):
        '''
        :param name_path_tupl: (name, path) to save the pic
        :return: a showing of the undirected graph/ saves the picture
        '''
        # Setting
        plt.axis("on")
        ax = plt.gca()
        self.create_edgeless_graph_plot(ax)
        # Directed labels
        elarge = [(u, v) for (u, v, d) in self.graph.edges(data=True) if d["weight"]]
        nx.draw_networkx_edges(self.graph, self.node_pos_dict, edgelist=elarge, width=EDGE_WIDTH, ax=ax,
                               connectionstyle=CONNECTION_STYLE, arrowsize=ARROW_SIZE)
        edge_labels = nx.get_edge_attributes(self.graph, 'weight')
        nx.draw_networkx_edge_labels(self.graph, self.node_pos_dict, edge_labels=edge_labels, ax=ax, label_pos=LABLE_POS,
                                     font_size=FONT_SIZE_EDGES)
        # Create figure:
        self.create_graph_settings_with_cartesic_coord(plt, ax)
        # Save
        if name_path_tupl:
            self.save_fig(name=name_path_tupl[0], path=name_path_tupl[1], format=SAVE_FORMAT, dpi=DPI)
        plt.gca().set_aspect('equal', adjustable='box')
        plt.show()


    def save_fig(self, name, path: Union[bool, str] = False, format='png', dpi=1200):
        if path:
            plt.savefig(path + r'\{}_{}.'.format(name, time.strftime("%Y-%m-%d_H%H-M%M")) + format, format=format, dpi=dpi)
        else:
            plt.savefig('{}_{}'.format(name, time.strftime("%Y-%m-%d_H%H-M%M")) + format, format=format, dpi=dpi)

    def create_edgeless_graph_plot(self, ax):
        for node_type in self.node_type_dict.keys():
            nlist = self.node_type_dict[node_type]
            ncolor = self.node_color_dict[node_type]

            # draw the graph
            nx.draw_networkx_nodes(self.graph,
                                   pos=self.node_pos_dict,
                                   nodelist=nlist,
                                   ax=ax,
                                   node_color=ncolor,
                                   label=node_type,
                                   edgecolors=NODE_EDGE_COLOR,
                                   node_size=NODE_SIZE)

        nx.draw_networkx_labels(self.graph, self.node_pos_dict, font_size=FONT_SIZE, font_family=FONT_FAMILY)

    def create_graph_settings_with_cartesic_coord(self, plt, ax):
        # getting the position of the legend so that the graph is not disrupted:
        x_lim_max = max([x[0] for x in self.node_pos_dict.values()])
        x_lim_max += x_lim_max * 0.01
        y_lim_max = max([y[1] for y in self.node_pos_dict.values()])
        y_lim_max += y_lim_max * 0.05
        x_lim_min = min([x[0] for x in self.node_pos_dict.values()])
        x_lim_min -= x_lim_min * 0.01
        y_lim_min = min([y[1] for y in self.node_pos_dict.values()])
        y_lim_min -= y_lim_min * 0.05

        ax.tick_params(left=True, bottom=True, labelleft=True, labelbottom=True)
        legend_size = 16
        plt.legend(scatterpoints=1, prop={'size': legend_size})
        plt.xlim(x_lim_min - 0.5, x_lim_max + legend_size / 6)
        plt.ylim(y_lim_min - 0.5, y_lim_max + 0.5)
        plt.xlabel('x')
        plt.ylabel('y')
        plt.tight_layout()