Add generic functions for plotting power grids

src/ITCPG.jl

 module ITCPG
 
 #* Used packages
+
 using CSV
 using DataFrames
 using JLD2
 using FileIO
 using PowerModels
 using Ipopt
+using PyPlot, PyCall
+using Colors
+using LaTeXStrings
+using DataStructures
 
 #* Include other code files and export functions
 
 include("Data.jl")
-export addlocs!, getlocs
+export add_locs!, get_locs, add_tl_voltages!, get_bustypes
 
 include("PowerFlow.jl")
 export update_pf_data!, calc_init_op
 
+include("PlotPG.jl")
+export plot_pg
+
 end

src/PlotPG.jl

+#* Functions for plotting a power grid described by a network data dictionary
+#*------------------------------------------------------------------------------
+
+#= 
+Function for recursively merging two dictionaries (used for updating the dictionary containing plot settings).
+=#
+_recursive_merge(x::AbstractDict,y::AbstractDict) = merge(_recursive_merge,x,y)
+_recursive_merge(x,y) = y
+
+#=
+Returns a dictionary of default plot settings for various plotting functions.
+=#
+function _default_settings(func::Symbol)
+    if func == :plot_pg
+        defaults = Dict(
+            ### Settings for _draw_buses!
+            "Buses" => Dict(
+                "Generator" => Dict(
+                    "marker" => "s",
+                    "size" => 30,
+                    "color" => "darkorange",
+                    "alpha" => 1,
+                    "label" => "Generators",
+                    "show" => true
+                ),
+                "Load and generator" => Dict(
+                    "marker" => "o",
+                    "size" => 30,
+                    "color" => "darkorange",
+                    "alpha" => 1,
+                    "label" => "nolabel",
+                    "show" => true
+                ),
+                "Slack" => Dict(
+                    "marker" => "s",
+                    "size" => 70,
+                    "color" => "red",
+                    "alpha" => 0.6,
+                    "label" => "Slack",
+                    "show" => true
+                ),
+                "Load" => Dict(
+                    "marker" => "o",
+                    "size" => 15,
+                    "color" => "limegreen",
+                    "alpha" => 0.7,
+                    "label" => "Loads",
+                    "show" => true
+                ),
+                "Empty bus" => Dict(
+                    "marker" => "o",
+                    "size" => 2,
+                    "color" => "black",
+                    "alpha" => 0.3,
+                    "label" => "Empty buses",
+                    "show" => true
+                )
+            ),
+            ### Settings for _draw_branches!
+            "Branches" => Dict(
+                "br_status" => "active", # plot only "active" or "all" branches
+                "br_coloring" => "equal", # how to color branches (can be set to "equal", "voltage", "MW-loading", "Mvar-loading" or "MVA-loading")
+                "br_color" => "k", # default color for "equal" coloring
+                "br_lw" => 2,
+                "br_alpha" => 1
+            ),
+            "draw_ticks" => [true, true, true, true], # whether to show ticks
+            "draw_legend" => true, # whether to draw a legend
+            "xlabel" => L"Longitude in $°$",
+            "ylabel" => L"Latitude in $°$",
+        )
+    end
+
+    return defaults
+end
+
+#*------------------------------------------------------------------------------
+
+#=
+Plots the power grid described by the network data dictionary (NDD). Possible plot settings are shown in _default_settings.
+=#
+function plot_pg(
+        network_data::Dict{String,<:Any},
+        settings = Dict{String,Any}(); # dictionary containing plot settings
+        figpath::String # where to save the figure
+    )
+    ### Setup figure
+    figure::Figure, ax::PyObject = plt.subplots()::Tuple{Figure,PyObject}
+    ax.set_aspect("equal")
+    w::Float64, h::Float64 = plt.figaspect(2/3)::Vector{Float64}
+    figure.set_size_inches(1.5w, 1.5h)
+
+    ### Set plot settings
+    settings = _recursive_merge(_default_settings(:plot_pg), settings)
+    _plot_pg!(ax, network_data, settings, figpath)
+    
+    return nothing
+end
+
+#=
+Plots the power grid described by the network data dictionary (NDD) onto an already existing axes. Possible plot settings are shown in _default_settings.
+=#
+function _plot_pg!(
+        ax::PyObject, # axes to draw power grid onto
+        network_data::Dict{String,<:Any},
+        settings::Dict{String,<:Any}, # dictionary containing plot settings
+        figpath::String # where to save the figure
+    )
+
+    ### Draw power grid graph
+    nx = pyimport("networkx")
+    G::PyObject = nx.Graph() # empty graph
+    _draw_pg!(ax, G, network_data, settings) # draw power grid onto axes
+   
+    plt.savefig(figpath, bbox_inches="tight")
+
+    plt.close("all") # close figure
+    return nothing
+end
+
+#*------------------------------------------------------------------------------
+
+#=
+Draws a graph for the power grid described by the NDD with options according to the dictionary "settings" (see _default_settings for possible options).
+=#
+function _draw_pg!(
+        ax::PyObject, # axes to draw power grid onto
+        G::PyObject, # power grid graph
+        network_data::Dict{String,<:Any}, 
+        settings::Dict{String,<:Any} # dictionary containing plot settings
+    )
+
+    ### Plot buses into graph
+    G, bus_markers, bus_labels = _draw_buses!(G, network_data, settings)
+
+    ### Plot branches into graph
+    G, br_markers, br_labels, cbar = _draw_branches!(G, network_data, settings)
+
+    ### Check for a predefined area to show
+    if haskey(settings, "area")
+        area = settings["area"]
+        plt.xlim(area[1], area[2])
+        plt.ylim(area[3], area[4])
+    end
+
+    ### Axes settings
+    ax.tick_params(
+        left = settings["draw_ticks"][1],
+        bottom = settings["draw_ticks"][2], 
+        labelleft = settings["draw_ticks"][3], 
+        labelbottom = settings["draw_ticks"][4]
+    )
+    plt.xlabel(settings["xlabel"])
+    plt.ylabel(settings["ylabel"], rotation=90)
+
+    ### Draw legend, if wanted
+    if settings["draw_legend"] == true
+        all_markers = vcat(bus_markers, br_markers)
+        all_labels = vcat(bus_labels, br_labels)
+        plt.legend(all_markers, all_labels)
+    end
+
+    return ax, G
+end
+
+#*------------------------------------------------------------------------------
+
+#=
+Draws buses contained in the NDD as nodes into graph G. The nodes are displayed according to the settings dictionary (see _default_settings).
+=#
+function _draw_buses!(
+        G::PyObject, # power grid graph
+        network_data::Dict{String,<:Any},
+        settings::Dict{String,<:Any} # dictionary containing plot settings 
+    )
+
+    ### Matplotlib imports
+    nx = pyimport("networkx")
+    mlines = pyimport("matplotlib.lines")
+    
+    bustypes = get_bustypes(network_data) # types of all buses
+    pos = get_locs(network_data) # geographic bus locations
+    bus_markes = [
+        mlines.Line2D([], [], color=b["color"], marker=b["marker"], ls="None")
+        for b in collect(values(settings["Buses"]))
+        if b["label"] != "nolabel"
+    ] # markers for legend
+    bus_labels = [
+        b["label"] for b in collect(values(settings["Buses"])) 
+        if b["label"] != "nolabel"
+    ] # labels for legend
+
+    ### Draw different buses as nodes
+    for (type, buses) in bustypes
+        bus_settings = settings["Buses"][type]
+        if bus_settings["show"] == true
+            nx.draw_networkx_nodes(
+                G, pos,
+                nodelist = buses,
+                node_shape = bus_settings["marker"],
+                node_size = bus_settings["size"],
+                node_color = bus_settings["color"],
+                alpha = bus_settings["alpha"]
+            )
+        end
+    end
+
+    return G, bus_markers, bus_labels
+end
+
+#*------------------------------------------------------------------------------
+
+#=
+Draws branches contained in the NDD as edges into graph G. The branches are displayed according to the settings dictionary (see _default_settings).
+=#
+function _draw_branches!(
+        G::PyObject, # power grid graph
+        network_data::Dict{String,<:Any},
+        settings::Dict{String,<:Any} # dictionary containing plot settings 
+    )
+
+    br_settings = settings["Branches"]
+    br_coloring = br_settings["br_coloring"]
+
+    ### Draw branches according to coloring mode
+    if br_coloring == "equal"
+        G, br_markers, br_labels, cbar = _draw_br_equal!(
+            G, network_data, br_settings
+        )
+    elseif br_coloring == "voltage"
+        G, br_markers, br_labels, cbar = _draw_br_voltage!(
+            G, network_data, br_settings
+        )
+    elseif br_coloring in ["MW-loading","Mvar-loading","MVA-loading"]
+        G, br_markers, br_labels, cbar = _draw_br_branchloads!(
+            G, network_data, br_settings
+        )
+    else
+        throw(ArgumentError("Unknown branch coloring $br_coloring."))
+    end
+
+    return G, br_markers, br_labels, cbar
+end
+
+#*------------------------------------------------------------------------------
+
+#=
+Draws branches contained in the NDD with coloring mode "equal". All branches are displayed using the same color.
+=#
+function _draw_br_equal!(
+        G::PyObject, # power grid graph
+        network_data::Dict{String,<:Any},
+        br_settings::Dict{String,<:Any} # dictionary containing plot settings 
+    )
+    
+    ### Matplotlib imports
+    nx = pyimport("networkx")
+    
+    pos = get_locs(network_data) # geographic bus locations
+    branches = collect(values(network_data["branch"])) # branch dictionaries
+    
+    ### Get edges contained in the NDD
+    if br_settings["br_status"] == "active" # only plot active branches
+        edges = [(b["f_bus"],b["t_bus"]) for b in branches if b["br_status"]==1]
+    elseif br_settings["br_status"] == "all" # plot all branches
+        edges = [(b["f_bus"],b["t_bus"]) for b in branches]
+    else
+        br_status = br_settings["br_status"]
+        throw(ArgumentError("Unknown branch status $br_status."))
+    end
+
+    ### Draw edges
+    drawn_edges = nx.draw_networkx_edges(
+        G, pos, 
+        edgelist = edges,
+        width = br_settings["br_lw"],
+        edge_color = br_settings["br_color"],
+        alpha = br_settings["br_alpha"]
+    )
+
+    return G, [], [], nothing    
+end
+
+#=
+Draws branches contained in the NDD with coloring mode "MW-loading", "Mvar-loading" or "MVA-loading". The branches are colored depending on their loading (flow/capacity).
+=#
+function _draw_br_branchloads!(
+        G::PyObject, # power grid graph
+        network_data::Dict{String,<:Any},
+        br_settings::Dict{String,<:Any} # dictionary containing plot settings
+    )
+
+    ### Matplotlib imports
+    nx = pyimport("networkx")
+
+    pos = get_locs(network_data) # geographic bus locations
+    branches = collect(values(network_data["branch"])) # branch dictionaries
+    br_coloring = br_settings["br_coloring"] # what kind of loading to use
+    br_status = br_settings["br_status"]
+
+    ### Get edges contained in the NDD and their loadings
+    if br_status == "active" # only plot active branches
+        edges = [(b["f_bus"],b["t_bus"]) for b in branches if b["br_status"]==1]
+        branchloads = [b[br_coloring] for b in branches if b["br_status"]==1]
+    elseif br_status == "all" # plot all branches
+        edges = [(b["f_bus"],b["t_bus"]) for b in branches]
+        branchloads = [b[br_coloring] for b in branches]
+    else
+        throw(ArgumentError("Unknown branch status $br_status."))
+    end
+
+    ### Draw edges
+    cmap = plt.cm.inferno_r
+    vmin, vmax = 0., 1.
+    drawnedges = nx.draw_networkx_edges(
+        G, pos, 
+        edgelist = edges,
+        width = br_settings["br_lw"],
+        edge_color = branchloads,
+        edge_cmap = cmap,
+        edge_vmin = vmin,
+        edge_vmax = vmax,
+        alpha = br_settings["br_alpha"]
+    )
+
+    ### Add colorbar
+    sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin, vmax))
+    cbar = plt.colorbar(sm)
+    cbar.ax.set_ylabel(
+        "Line $br_coloring " * L"$F_{ij}/C_{ij}$", rotation=-90, va="bottom"
+    )
+
+    return G, [], [], cbar
+end
+
+#=
+Draws branches contained in the NDD with coloring mode "voltage". Transmission lines are colored according to their voltage levels.
+=#
+function _draw_br_voltage!(
+        G::PyObject, # power grid graph
+        network_data::Dict{String,<:Any},
+        br_settings::Dict{String,<:Any} # dictionary containing plot settings
+    )
+    
+    ### Matplotlib imports
+    nx = pyimport("networkx")
+    mlines = pyimport("matplotlib.lines")
+
+    pos = get_locs(network_data) # geographic bus locations
+    branches = collect(values(network_data["branch"])) # branch dictionaries
+    br_markers = Array{PyObject,1}() # markers for legend
+    br_labels = Array{String,1}() # labels for legend
+    
+    ### Get edges contained in the NDD and their voltage levels
+    if br_settings["br_status"] == "active" # only plot active branches
+        edges = [
+            (b["f_bus"],b["t_bus"]) for b in branches 
+            if b["br_status"] == 1
+        ]
+        voltages = [
+            string(b["tl_voltage"]) for b in branches 
+            if b["br_status"] == 1
+        ]
+    elseif br_settings["br_status"] == "all" # plot all branches
+        edges = [(b["f_bus"],b["t_bus"]) for b in branches]
+        voltages = [string(b["tl_voltage"]) for b in branches]
+    else
+        br_status = br_settings["br_status"]
+        throw(ArgumentError("Unknown branch status $br_status."))
+    end
+
+    ### Assign colors to voltage levels and add markers and labels for legend
+    voltages[voltages .== "0.0"] .= "k" # transformers
+    mcolors = pyimport("matplotlib.colors")
+    tableau = [
+        key for key in keys(mcolors.TABLEAU_COLORS) 
+        if key ∉ ["tab:orange", "tab:green"] # orange and green used for buses
+    ]
+    for (i, v) in enumerate(sort(unique(filter(v -> v != "k", voltages))))
+        voltages[voltages .== v] .= tableau[i]
+        push!(br_markers, mlines.Line2D([], [], color=tableau[i], ls="-"))
+        push!(br_labels, string(v) * " kV")
+    end
+
+    ### Draw edges
+    drawnedges = nx.draw_networkx_edges(
+        G, pos, 
+        edgelist = edges,
+        width = br_settings["br_lw"],
+        edge_color = voltages,
+        alpha = br_settings["br_alpha"]
+    )
+    
+    return G, br_markers, br_labels, nothing
+end
\ No newline at end of file