Пример #1
0
    def highlight_plot(self, param):
        color_select = []
        for i in range(len(self.config_all)):
            if self.colors[i] != 'none' and param[1][i]:
                color_select.append(self.colors[i])
            else:
                pass

        config_ = []
        config_select = []
        for i in range(len(self.config_all)):
            config = []
            if param[0][i] and param[1][i]:
                current_row = param[2][i]
                for j in range(len(self.config_all[i])):
                    if self.config_all[i][j]:
                        config.append(self.config_all[i][j][current_row])
                config_.append(config)
                config_select_ = list(map(list, zip(*config_)))
                config_select.append(config_select_)

        detect = Detection(self.obj, config_select)
        patches, paths = detect.plot()
        for i in range(len(patches)):
            p = PatchCollection(patches[i], color=color_select[i], edgecolor=color_select[i], alpha=0.6)
            p.set_zorder(6)
            self.ax.add_collection(p)
        for m in range(len(paths)):
            for n in range(len(paths[m])):
                patch_polygon = PathPatch(paths[m][n], facecolor=color_select[m], edgecolor=color_select[i],
                                          alpha=0.6)
                patch_polygon.set_zorder(6)
                self.ax.add_patch(patch_polygon)

        self.redraw()
    def draw_nodes(G,
                   pos,
                   nodelist=None,
                   node_size=300,
                   node_color='r',
                   node_shape='o',
                   alpha=1.0,
                   cmap=None,
                   vmin=None,
                   vmax=None,
                   ax=None,
                   linewidths=None,
                   label=None,
                   **kwds):

        try:
            import matplotlib.pyplot as plt
            import numpy
        except ImportError:
            raise ImportError("Matplotlib required for draw()")
        except RuntimeError:
            print("Matplotlib unable to open display")
            raise

        if ax is None:
            ax = plt.gca()

        if nodelist is None:
            nodelist = list(G.nodes())

        if not nodelist or len(nodelist) == 0:  # empty nodelist, no drawing
            return None

        try:
            xy = numpy.asarray([pos[v] for v in nodelist])
        except KeyError as e:
            raise nx.NetworkXError('Node %s has no position.' % e)
        except ValueError:
            raise nx.NetworkXError('Bad value in node positions.')

        radius = 7

        circles = []
        for v in nodelist:
            circ = Circle(pos[v], radius=radius, fill=True)
            circles.append(circ)

        node_collection = PatchCollection(circles,
                                          facecolors='r',
                                          edgecolors='k',
                                          linewidths=2,
                                          alpha=0.8)
        node_collection.set_zorder(3)
        ax.add_collection(node_collection)
        return node_collection
Пример #3
0
def _create_node_collection(nodes, coords, size=5, patch_type="circle", color=None, picker=False,
                            infos=None, **kwargs):
    """
    Creates a collection with patches for the given nodes. Can be used generically for different \
    types of nodes (bus in pandapower network, but also other nodes, e.g. in a networkx graph).

    :param nodes: indices of the nodes to plot
    :type nodes: iterable
    :param coords: list of node coordinates (shape (2, N))
    :type coords: iterable
    :param size: size of the patches (handed over to patch creation function)
    :type size: float
    :param patch_type: type of patches that chall be created for the nodes - can be one of\
        - "circle" for a circle\
        - "rect" for a rectangle\
        - "poly<n>" for a polygon with n edges
    :type patch_type: str, default "circle"
    :param color: colors or color of the node patches
    :type color: iterable, float
    :param picker: picker argument passed to the patch collection
    :type picker: bool, default False
    :param infos: list of infos belonging to each of the patches (can be displayed when hovering \
        over the elements)
    :type infos: list, default None
    :param kwargs: keyword arguments are passed to the patch maker and patch collection
    :type kwargs:
    :return: pc - patch collection for the nodes
    """
    if len(coords) == 0:
        return None

    infos = list(infos) if infos is not None else []
    patches = node_patches(coords, size, patch_type, color, **kwargs)
    pc = PatchCollection(patches, match_original=True, picker=picker)
    pc.node_indices = np.array(nodes)

    pc.patch_type = patch_type
    pc.size = size
    if 'orientation' in kwargs:
        pc.orientation = kwargs['orientation']
    if "zorder" in kwargs:
        pc.set_zorder(kwargs["zorder"])
    pc.info = infos
    pc.patch_type = patch_type
    pc.size = size
    if "zorder" in kwargs:
        pc.set_zorder(kwargs["zorder"])
    if 'orientation' in kwargs:
        pc.orientation = kwargs['orientation']

    return pc
Пример #4
0
    def zoneLines(self, edgecolour='black'):        #was 'red'
        """
        Set boundary colour for defined ESRI shapes
        edgecolour -- HTML colour name for boundary
        """

        pc2 = PatchCollection(self.patches, match_original=True)
        pc2.set_facecolor('none')
        pc2.set_edgecolor(edgecolour)
        pc2.set_alpha(0.5) #5.0
        pc2.set_linewidth(0.5)
        pc2.set_zorder(25) # 500

        sq2 = self.ax.add_collection(pc2)
Пример #5
0
def plotss(data_feats,ax):
    """
    Plots secondary structures

    :param data_feats: pandas dataframe with position wise features
    :param ax: axes object
    """

#     plt.figure(figsize=(40, 1),dpi=500)
#     ax=plt.subplot(111)
    ax.set_axis_off()
    xlim=len(data_feats)

    #no sec struct
    for aai,rowi in data_feats.iterrows():
        if not pd.isnull(rowi.loc["Secondary structure"]):
            ini=aai+1
            break
    for aai,rowi in data_feats.iloc[::-1].iterrows():
        if not pd.isnull(rowi.loc["Secondary structure"]):
            end=aai+1
            break
    x,y = np.array([[ini, end], [0, 0]])
    line = mlines.Line2D(x, y, lw=5,color="black")
    line.set_zorder(0)
    ax.add_line(line)

    patches = []
    patches_colors=[]
    ssi_ini=""
    aai_ini=0
    for aai,rowi in data_feats.fillna("").iterrows():
        aai=aai-1
        ssi=rowi.loc["Secondary structure"]
        if ssi!=ssi_ini:
            patches,patches_colors,aai_ini,ssi_ini,ax=plotsspatches(ssi,ssi_ini,aai,aai_ini,patches,patches_colors,ax)
    collection = PatchCollection(patches, cmap=plt.cm.Accent, alpha=1)
    colors = np.linspace(0, 1, len(patches))
    collection.set_array(np.array(patches_colors))
    collection.set_zorder(20)
    ax.add_collection(collection)
    # ax.set_xticklabels(range(xlim))
    ax.set_xlim((0-0.5,xlim+0.5))
    ax.set_ylim((-1.1,1.1))
    ax.text(xlim+1,-0.5,"Secondary structure",fontdict={'size': 20})
#     plt.show()
    return ax
Пример #6
0
    def plotStand(self, param, flag, stand_param):
        self.fig.clear()
        self.ax = self.fig.add_subplot(111) #添加子图
        self.fig.subplots_adjust(left=0.05, right=0.85, top=0.95, bottom=0.05) #设置边距
        # self.ax.clear()
        self.config_all = self.init_config(param) #获取参数

        if param[0][0] != '-- Choose ODD please --':
            self.obj = self.odd_import(self.ax, param[0])

        config_select = []
        color_select = []
        for i in range(len(self.config_all)):
            if self.colors[i] != 'none':
                config_select.append(self.config_all[i])
                color_select.append(self.colors[i])
            else:
                pass
        detect = Detection(self.obj, config_select)     #opencv检测
        patches, paths = detect.plot()  # 路径点坐标


        if flag:
            self.loadYaml(stand_param)
        else:
            pass
        # for sensorType in self.yFile.yamlData:
        #     s=self.loadYaml(sensorType)
        #     s.set_zorder(1)  # 设置图层
        #     self.ax.add_collection(s)

        for i in range(len(patches)):
            p = PatchCollection(patches[i], color=color_select[i], edgecolor=color_select[i], alpha=0.15)
            p.set_zorder(6)  #设置图层
            self.ax.add_collection(p)  #将PatchCollection添加进axes对象
        for m in range(len(paths)):
            for n in range(len(paths[m])):
                patch_polygon = PathPatch(paths[m][n], facecolor=color_select[m], edgecolor=color_select[m],
                                          alpha=0.15)
                patch_polygon.set_zorder(6)
                self.ax.add_patch(patch_polygon)
        self.in_view = detect.postdetection_list
        self.in_scale_view = detect.detection_list
        self.draw_ego_vehicle()
Пример #7
0
    def zoneColour(self, colours):
        """
        Set display colours for defined ESRI shapes
        colours -- list containing HTML colour names
        """

        self.colours = colours

        if not (isinstance(self.colours, list)):
            raise Exception('Invalid list of zone colours')

        pc = PatchCollection(self.patches, match_original=True)
        pc.set_facecolor(self.colours)
        pc.set_edgecolor('none')
        pc.set_alpha(0.5)
        pc.set_linewidth(0.5)
        pc.set_zorder(20)

        sq = self.ax.add_collection(pc)
Пример #8
0
 def plotGeometry(self, ax = None):
     if (ax is None):
         fig = plt.figure()
         ax = fig.add_subplot(111)
     ax.set_aspect(1)
     
     self.blockPatch.set_facecolor('#F0F0F0')
     self.blockPatch.set_edgecolor('k')
     ax.add_patch(self.blockPatch)    
     
     pc = PatchCollection(self.channelPatches, match_original = True)
     pc.set_facecolor('w')
     pc.set_edgecolor('k')
     pc.set_zorder(2)
     ax.add_collection(pc)
        
     self.plotDimensions(ax)
     ax.set_xlim(self.xMin, self.xMax)
     ax.set_ylim(self.yMin, self.yMax)
     plt.show()
Пример #9
0
def plotss(data_feats,ax,refi_lims,fontsize=20,ticks=False):
    """
    Plots secondary structures

    :param data_feats: pandas dataframe with position wise features
    :param ax: axes object
    """
    if not ticks:
        ax.set_axis_off()
    #no sec struct
    for aai,rowi in data_feats.iterrows():
        if not pd.isnull(rowi.loc["Secondary structure"]):
            ini=aai+1
            break
    for aai,rowi in data_feats.iloc[::-1].iterrows():
        if not pd.isnull(rowi.loc["Secondary structure"]):
            end=aai+1
            break
    x,y = np.array([[ini, end], [0, 0]])
    line = mlines.Line2D(x, y, lw=5,color="black")
    line.set_zorder(0)
    ax.add_line(line)

    patches = []
    patches_colors=[]
    ssi_ini=""
    aai_ini=0
    for aai,rowi in data_feats.fillna("").iterrows():
        aai=aai-1
        ssi=rowi.loc["Secondary structure"]
        if ssi!=ssi_ini:
            patches,patches_colors,aai_ini,ssi_ini,ax=plotsspatches(ssi,ssi_ini,aai,aai_ini,patches,patches_colors,ax)
    collection = PatchCollection(patches, cmap=plt.cm.Accent, alpha=1)
    colors = np.linspace(0, 1, len(patches))
    collection.set_array(np.array(patches_colors))
    collection.set_zorder(20)
    ax.add_collection(collection)
    ax.set_xlim((refi_lims[0]-0.5,refi_lims[1]+0.5))
    ax.set_ylim((-1.1,1.1))
    ax.text(refi_lims[1]+1,-0.5,"%sSecondary structure" % (r'$\leftarrow$'),fontdict={'size': fontsize})
    return ax
Пример #10
0
    def odd_import(self, scenario):
        try:
            lines1 = scenario['lines']['boundary_lines']
            lines2 = scenario['lines']['solid_lines']
            lines3 = scenario['lines']['dashed_lines']
            lines4 = scenario['lines']['thick_lines']
            objects = scenario['objects']
            common_rotation = scenario['common_rotation']
            common_offset = scenario['common_offset']
        except Exception:
            lines1 = []
            lines2 = []
            lines3 = []
            lines4 = []
            objects = []
            common_rotation = 0.0
            common_offset = [0, 0]
            QMessageBox.warning(self, 'Warning', 'Scenario file format error!')

        boundary_lines = self.line_rotation(lines1, common_rotation)
        solid_lines = self.line_rotation(lines2, common_rotation)
        dashed_lines = self.line_rotation(lines3, common_rotation)
        thick_lines = self.line_rotation(lines4, common_rotation)

        for l in boundary_lines:
            for m in l:
                m[0] = (m[0] + common_offset[0]) * self.zoom_factor
                m[1] = (m[1] + common_offset[1]) * self.zoom_factor

        for l in solid_lines:
            for m in l:
                m[0] = (m[0] + common_offset[0]) * self.zoom_factor
                m[1] = (m[1] + common_offset[1]) * self.zoom_factor

        for l in dashed_lines:
            for m in l:
                m[0] = (m[0] + common_offset[0]) * self.zoom_factor
                m[1] = (m[1] + common_offset[1]) * self.zoom_factor

        for l in thick_lines:
            for m in l:
                m[0] = (m[0] + common_offset[0]) * self.zoom_factor
                m[1] = (m[1] + common_offset[1]) * self.zoom_factor

        line_segments1 = LineCollection(boundary_lines, linewidths=2, colors=(0, 0, 0), linestyles='solid')
        line_segments1.set_zorder(0)
        line_segments2 = LineCollection(solid_lines, linewidths=1, colors=(0.1, 0.1, 0.1), linestyles='solid')
        line_segments2.set_zorder(0)
        line_segments3 = LineCollection(dashed_lines, linewidths=1, colors=(0.2, 0.2, 0.2), linestyles='dashed')
        line_segments3.set_zorder(0)
        line_segments4 = LineCollection(thick_lines, linewidths=4.5, colors=(0.8, 0.8, 0.8), linestyles='solid')
        line_segments4.set_zorder(0)

        self.ax.add_collection(line_segments1)
        self.ax.add_collection(line_segments2)
        self.ax.add_collection(line_segments3)
        self.ax.add_collection(line_segments4)

        patch = []
        for obj in range(len(objects)):
            objects_xy = self.object_rotation(objects[obj][1], common_rotation)
            objects_xy[0] = objects_xy[0] + common_offset[0]
            objects_xy[1] = objects_xy[1] + common_offset[1]
            objects_r = objects[obj][2] + common_rotation
            patch.append(self.add_object(objects[obj][0], objects_xy, objects_r))
            self.ax.text((objects_xy[0]) * self.zoom_factor, (objects_xy[1] + 2.5) * self.zoom_factor,
                         'obj' + str(obj + 1), fontsize=10)
        patch_segment = PatchCollection(patch, facecolors='gray')
        patch_segment.set_zorder(4)
        self.ax.add_collection(patch_segment)
        return patch
def plot_eigvect_excitation_haldane(xy,
                                    fig,
                                    dos_ax,
                                    eig_ax,
                                    eigval,
                                    eigvect,
                                    en,
                                    marker_num=0,
                                    black_t0lines=False,
                                    mark_t0=True,
                                    title='auto',
                                    normalization=1.,
                                    alpha=0.6,
                                    lw=1,
                                    zorder=10):
    """Draws normal mode ellipsoids on axis eig_ax.
    If black_t0lines is true, draws the black line from pinning site to positions. The difference from
    hlatpfns.construct_haldane_eigvect_DOS_plot() is doesn't draw lattice.

    Parameters
    ----------
    xy: array 2N x 3
        Equilibrium position of the gyroscopes
    fig :
        figure with lattice and DOS drawn
    dos_ax: matplotlib axis instance or None
        axis for the DOS plot. If None, ignores this input
    eig_ax : matplotlib axis instance
        axis for the eigenvalue plot
    eigval : array of dimension 2nx1
        Eigenvalues of matrix for system
    eigvect : array of dimension 2nx2n
        Eigenvectors of matrix for system.
        Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like:
        x0, y0, x1, y1, ... xNP, yNP.
    en: int
        Number of the eigenvalue you are plotting
    marker_num : int in (0, 80)
        where in the phase (0 to 80) to call t=t0. This sets "now" for drawing where in the normal mode to draw
    black_t0lines : bool
        Draw black lines extending from the pinning site to the current site (where 'current' is determined by
        marker_num)

    Returns
    ----------
    fig : matplotlib figure instance
        completed figure for normal mode
    [scat_fg, pp, f_mark, lines12_st] :
        things to be cleared before next normal mode is drawn
    """
    s = leplt.absolute_sizer()

    ev = eigval[en]
    ev1 = ev

    # Show where current eigenvalue is in DOS plot
    if dos_ax is not None:
        (f_mark, ) = dos_ax.plot([np.real(ev), np.real(ev)], dos_ax.get_ylim(),
                                 '-r')

    NP = len(xy)

    im1 = np.real(ev)
    plt.sca(eig_ax)

    if title == 'auto':
        eig_ax.set_title('$\omega = %0.6f$' % im1)
    elif title is not None and title not in ['', 'none']:
        eig_ax.set_title(title)

    # Preallocate ellipsoid plot vars
    angles_arr = np.zeros(NP, dtype=float)

    patch = []
    colors = np.zeros(NP)
    x0s = np.zeros(NP, dtype=float)
    y0s = np.zeros(NP, dtype=float)
    mag1 = eigvect[en]

    # Pick a series of times to draw out the ellipsoid
    time_arr = np.arange(81) * 2 * np.pi / (np.abs(ev1) * 80)
    exp1 = np.exp(1j * ev1 * time_arr)

    # Normalization for the ellipsoids
    mag1 /= np.max(np.abs(mag1))
    mag1 *= normalization

    if black_t0lines:
        lines_1 = []
    else:
        lines_12_st = []

    for i in range(NP):
        x_disps = 0.5 * (exp1 * mag1[i]).real
        y_disps = 0.5 * (exp1 * mag1[i]).imag
        x_vals = xy[i, 0] + x_disps
        y_vals = xy[i, 1] + y_disps

        poly_points = np.array([x_vals, y_vals]).T
        polygon = Polygon(poly_points, True)

        # x0 is the marker_num^th element of x_disps
        x0 = x_disps[marker_num]
        y0 = y_disps[marker_num]

        x0s[i] = x_vals[marker_num]
        y0s[i] = y_vals[marker_num]

        if black_t0lines:
            # These are the black lines protruding from pivot point to current position
            lines_1.append([[xy[i, 0], x_vals[marker_num]],
                            [xy[i, 1], y_vals[marker_num]]])

        mag = np.sqrt(x0**2 + y0**2)
        if mag > 0:
            anglez = np.arccos(x0 / mag)
        else:
            anglez = 0

        if y0 < 0:
            anglez = 2 * np.pi - anglez

        angles_arr[i] = anglez
        patch.append(polygon)
        colors[i] = anglez

    # this is the part that puts a dot a t=0 point
    if mark_t0:
        scat_fg = eig_ax.scatter(x0s, y0s, s=s(.02), c='k')
    else:
        scat_fg = []

    pp = PatchCollection(patch, cmap='hsv', lw=lw, alpha=alpha, zorder=zorder)

    pp.set_array(np.array(colors))
    pp.set_clim([0, 2 * np.pi])
    pp.set_zorder(1)

    eig_ax.add_collection(pp)

    if black_t0lines:
        lines_12 = [zip(x, y) for x, y in lines_1]
        lines_12_st = LineCollection(lines_12, linewidth=0.8)
        lines_12_st.set_color('k')
        eig_ax.add_collection(lines_12_st)

    eig_ax.set_aspect('equal')

    return fig, [scat_fg, pp, f_mark, lines_12_st]
Пример #12
0
    def onselect(self, verts):
        """
        Defines what happens when a polygon is completed

        :param verts: polygon vertices automatically returned by the Polygon Selector
        :return:
        """

        if len(self.vals
               ) == 0:  # If the user forgets to press enter to initiate value
            self.vals.append(0)
            print("Press enter in the text box")

        # When a polygon is completed or modified after completion, the `onselect`
        path = Path(verts)
        # function is called and passed a list of the vertices as ``(xdata, ydata)`` tuples.

        self.ind = np.nonzero(path.contains_points(self.xys))[0]  #

        # xycenter = self.xys[self.ind]
        # xycenter.data
        self.index.append([self.vals[-1], self.ind, verts])
        # OP : Value, Index #, XY Coordinates selected points, Polygon vtx

        # Lists all values assigned to each polygon,
        current_vals = [self.index[i][0] for i in range(len(self.index))]
        # stored in the 'index' list
        vals_cb = list(dict.fromkeys(current_vals))  # Removes duplicates

        lsps = find_norm(current_vals)  # Get colors

        cc = [self.cmap(v) for v in lsps]  # Colors used for the points/blocks

        ccb = lsps.copy()  # Copy list
        ccb = list(dict.fromkeys(ccb))  # Removes duplicates
        ccb.sort()
        # List of colors for the colorbar
        cols = colors.ListedColormap([self.cmap(v) for v in ccb])

        if (
                len(vals_cb) > 1
        ):  # Creating a colormap fn to automatically rescale the colors of selected zones

            try:
                vals_cb.sort()
                # Adding one bound more to have enough intervals
                bounds = vals_cb + [max(vals_cb) * 1.1]
                ticks = vals_cb  # Ticks labels
                cbnorm = colors.BoundaryNorm(bounds, cols.N)
                mcb = colorbar.ColorbarBase(
                    self.axcb,
                    cmap=cols,
                    norm=cbnorm,
                    boundaries=bounds,
                    ticks=ticks,
                    ticklocation="right",
                    orientation="vertical",
                )
            except:
                pass

        else:
            pass

        if (not self.blocks.any()
            ):  # If no blocks, simply coloring the points accordingly
            for i in range(len(self.index)):
                self.fc[self.index[i][1]] = self.cmap(lsps[i])
            self.collection.set_facecolors(self.fc)  # updates points color

        else:
            if len(self.index) > 0:
                for i in range(len(self.index)):  # If blocks
                    # Concerned mesh blocks
                    bid = self.blocks[self.index[i][1]]
                    patches = []
                    for b in bid:
                        polygon = Polygon(b, closed=True)
                        patches.append(polygon)
                    p = PatchCollection(patches, alpha=1, facecolors=cc[i])
                    p.set_zorder(1)
                    self.ax.add_collection(p)
            else:
                pass

        # self.collection.set_facecolors(self.fc)  # updates points color

        self.canvas.draw()

        # Necessary to copy to not mess with the original array
        idx = np.copy(self.index)

        # What you see is what you get

        # Removes duplicates from older selections, scanning for the latest
        for i in range(len(idx) - 1, 0, -1):
            # polygons first to the last.
            indc = idx[i][1]
            indp = idx[i - 1][1]
            upd = np.array(list(set(indp) - set(indc)))
            # Updates the oldest polygons values, removing duplicates index values
            self.index[i - 1][1] = upd

        # final_results = np.ones(len(self.points))*self.bck  # Multiplying by the background value

        for item in self.index:
            for i in item[1]:
                self.final_results[i] = item[0]

        # Now save the results
        if self.model_name:
            print("Saving file...")
            # The final output is a 1-column file with the values assigned,
            with open(self.model_name, "w") as mm:
                # in the same order as the blocks provided. The file is opened and saved every time a modification is
                # made.
                mm.write(str(int(self.Npts)) + "\n")
                [mm.write(str(fr) + "\n") for fr in self.final_results]
                mm.close()
Пример #13
0
def draw_networkx_edges(G, pos,
                        edgelist=None,
                        width=1.0,
                        edge_color='k',
                        style='solid',
                        alpha=None,
                        edge_cmap=None,
                        edge_vmin=None,
                        edge_vmax=None,
                        ax=None,
                        arrows=True,
                        label=None,
                        **kwds):
    """Draw the edges of the graph G.

    This draws only the edges of the graph G.

    Parameters
    ----------
    G : graph
       A networkx graph

    pos : dictionary
       A dictionary with nodes as keys and positions as values.
       Positions should be sequences of length 2.

    edgelist : collection of edge tuples
       Draw only specified edges(default=G.edges())

    width : float
       Line width of edges (default =1.0)

    edge_color : color string, or array of floats
       Edge color. Can be a single color format string (default='r'),
       or a sequence of colors with the same length as edgelist.
       If numeric values are specified they will be mapped to
       colors using the edge_cmap and edge_vmin,edge_vmax parameters.

    style : string
       Edge line style (default='solid') (solid|dashed|dotted,dashdot)

    alpha : float
       The edge transparency (default=1.0)

    edge_ cmap : Matplotlib colormap
       Colormap for mapping intensities of edges (default=None)

    edge_vmin,edge_vmax : floats
       Minimum and maximum for edge colormap scaling (default=None)

    ax : Matplotlib Axes object, optional
       Draw the graph in the specified Matplotlib axes.

    arrows : bool, optional (default=True)
       For directed graphs, if True draw arrowheads.

    label : [None| string]
       Label for legend

    Returns
    -------
    matplotlib.collection.LineCollection
        `LineCollection` of the edges

    Notes
    -----
    For directed graphs, "arrows" (actually just thicker stubs) are drawn
    at the head end.  Arrows can be turned off with keyword arrows=False.
    Yes, it is ugly but drawing proper arrows with Matplotlib this
    way is tricky.

    Examples
    --------
    >>> G=nx.dodecahedral_graph()
    >>> edges=nx.draw_networkx_edges(G,pos=nx.spring_layout(G))

    Also see the NetworkX drawing examples at
    http://networkx.lanl.gov/gallery.html

    See Also
    --------
    draw()
    draw_networkx()
    draw_networkx_nodes()
    draw_networkx_labels()
    draw_networkx_edge_labels()
    """
    try:
        import matplotlib
        import matplotlib.pyplot as plt
        import matplotlib.patches as patches
        import matplotlib.cbook as cb
        from matplotlib.colors import colorConverter, Colormap
        from matplotlib.collections import LineCollection
        from matplotlib.collections import PatchCollection
        import numpy
    except ImportError:
        raise ImportError("Matplotlib required for draw()")
    except RuntimeError:
        print("Matplotlib unable to open display")
        raise

    if ax is None:
        ax = plt.gca()

    if edgelist is None:
        edgelist = G.edges()

    if not edgelist or len(edgelist) == 0:  # no edges!
        return None

    # set edge positions
    edge_pos = numpy.asarray([(pos[e[0]], pos[e[1]]) for e in edgelist])

    if not cb.iterable(width):
        lw = (width,)
    else:
        lw = width

    if not cb.is_string_like(edge_color) \
           and cb.iterable(edge_color) \
           and len(edge_color) == len(edge_pos):
        if numpy.alltrue([cb.is_string_like(c)
                         for c in edge_color]):
            # (should check ALL elements)
            # list of color letters such as ['k','r','k',...]
            edge_colors = tuple([colorConverter.to_rgba(c, alpha)
                                 for c in edge_color])
        elif numpy.alltrue([not cb.is_string_like(c)
                           for c in edge_color]):
            # If color specs are given as (rgb) or (rgba) tuples, we're OK
            if numpy.alltrue([cb.iterable(c) and len(c) in (3, 4)
                             for c in edge_color]):
                edge_colors = tuple(edge_color)
            else:
                # numbers (which are going to be mapped with a colormap)
                edge_colors = None
        else:
            raise ValueError('edge_color must consist of either color names or numbers')
    else:
        if cb.is_string_like(edge_color) or len(edge_color) == 1:
            edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
        else:
            raise ValueError('edge_color must be a single color or list of exactly m colors where m is the number or edges')

    edge_collection = None

    if not arrows:
        edge_collection = LineCollection(edge_pos,
                                         colors=edge_colors,
                                         linewidths=lw,
                                         antialiaseds=(1,),
                                         linestyle=style,
                                         transOffset = ax.transData,
                                         )

        edge_collection.set_zorder(1)  # edges go behind nodes
        edge_collection.set_label(label)
        ax.add_collection(edge_collection)

        # Note: there was a bug in mpl regarding the handling of alpha values for
        # each line in a LineCollection.  It was fixed in matplotlib in r7184 and
        # r7189 (June 6 2009).  We should then not set the alpha value globally,
        # since the user can instead provide per-edge alphas now.  Only set it
        # globally if provided as a scalar.
        if cb.is_numlike(alpha):
            edge_collection.set_alpha(alpha)

        if edge_colors is None:
            if edge_cmap is not None:
                assert(isinstance(edge_cmap, Colormap))
            edge_collection.set_array(numpy.asarray(edge_color))
            edge_collection.set_cmap(edge_cmap)
            if edge_vmin is not None or edge_vmax is not None:
                edge_collection.set_clim(edge_vmin, edge_vmax)
            else:
                edge_collection.autoscale()

    arrow_collection = None

    if G.is_directed() and arrows:

        # a directed graph hack
        # draw thick line segments at head end of edge
        # waiting for someone else to implement arrows that will work
        arrow_colors = edge_colors
        arrow_patches = []
        for src, dst in edge_pos:
            x1, y1 = src
            x2, y2 = dst
            dx = x2-x1   # x offset
            dy = y2-y1   # y offset
            arrow_width = width / 20.0
            line_width = arrow_width / 20.0
            arrow_patch = patches.FancyArrow(x1, y1, dx, dy, width=line_width,
                                             length_includes_head=True, head_width=arrow_width,
                                             head_length=arrow_width)
            arrow_patches.append(arrow_patch)

        arrow_collection = PatchCollection(arrow_patches,
                                           edgecolors=arrow_colors,
                                           facecolors=arrow_colors,
                                           linewidths=lw,
                                           linestyle=style,
                                           antialiaseds=(1,),
                                           transOffset = ax.transData,
                                          )

        arrow_collection.set_zorder(1)  # edges go behind nodes
        arrow_collection.set_label(label)
        ax.add_collection(arrow_collection)

    # update view
    minx = numpy.amin(numpy.ravel(edge_pos[:, :, 0]))
    maxx = numpy.amax(numpy.ravel(edge_pos[:, :, 0]))
    miny = numpy.amin(numpy.ravel(edge_pos[:, :, 1]))
    maxy = numpy.amax(numpy.ravel(edge_pos[:, :, 1]))

    w = maxx-minx
    h = maxy-miny
    padx,  pady = 0.05*w, 0.05*h
    corners = (minx-padx, miny-pady), (maxx+padx, maxy+pady)
    ax.update_datalim(corners)
    ax.autoscale_view()

    if arrow_collection:
        return arrow_collection
    else:
        return edge_collection
Пример #14
0
def create_bus_collection(net, buses=None, size=5, marker="o", patch_type="circle", colors=None,
                          cmap=None, norm=None, infofunc=None, **kwargs):
    """
    Creates a matplotlib patch collection of pandapower buses.
    
    Input:

        **net** (PandapowerNet) - The pandapower network
               
    Optional:
    
        **buses** (list, None) - The buses for which the collections are created. If None, all buses in the network are considered.

        **size** (int, 5) - patch size

        **marker** (str, "o") - patch marker

        **patch_type** (str, "circle") - patch type, can be
        
                - "circle" for a circle
                - "rect" for a rectanlge
                - "poly<n>" for a polygon with n edges 
        
        **infofunc** (function, None) - infofunction for the patch element
        
        **colors** (list, None) - list of colors for every element
        
        **cmap** - colormap for the patch colors
        
        **picker** - picker argument passed to the patch collection
        
        **kwargs - key word arguments are passed to the patch function
        
    """
    buses = net.bus.index.tolist() if buses is None else list(buses)
    patches = []
    infos = []
    def figmaker(x, y, i):
        if patch_type=="circle":
            if colors:
                fig = Circle((x, y), size, color=colors[i], **kwargs)
            else:
                fig = Circle((x, y), size, **kwargs)
        elif patch_type=="rect":
            if colors:
                fig = Rectangle([x - size, y - size], 2*size, 2*size, color=colors[i], **kwargs)
            else:
                fig = Rectangle([x - size, y - size], 2*size, 2*size, **kwargs)
        elif patch_type.startswith("poly"):
            edges = int(patch_type[4:])
            if colors:
                fig = RegularPolygon([x, y], numVertices=edges, radius=size, color=colors[i],
                                     **kwargs)
            else:
                fig = RegularPolygon([x, y], numVertices=edges, radius=size, **kwargs)
        if infofunc:
            infos.append(infofunc(buses[i]))
        return fig
    patches = [figmaker(x, y, i)
               for i, (x, y) in enumerate(zip(net.bus_geodata.loc[buses].x.values,
                                              net.bus_geodata.loc[buses].y.values))
               if x != -1 and x != np.nan]
    pc = PatchCollection(patches, match_original=True)
    if cmap:
        pc.set_cmap(cmap)
        pc.set_norm(norm)
        pc.set_array(net.res_bus.vm_pu.loc[buses])
        pc.has_colormap = True
        pc.cbar_title = "Bus Voltage [pu]"

    pc.patch_type = patch_type
    pc.size = size
    if "zorder" in kwargs:
        pc.set_zorder(kwargs["zorder"])
    pc.info = infos
    return pc
Пример #15
0
def draw_frontline(x,
                   y,
                   ax,
                   line_style="-",
                   line_color='k',
                   line_width=1,
                   line_alpha=1,
                   marker_shape="polygon",
                   marker_num_sides=3,
                   marker_angle=0,
                   marker_size=12,
                   marker_interval=24,
                   marker_offset=0,
                   marker_face_color='k',
                   marker_edge_color='k',
                   marker_edge_width=1,
                   marker_aspect_ratio=1.,
                   marker_alternate_sides=False,
                   marker_theta1=0,
                   marker_theta2=180,
                   marker_arrow_shape="full",
                   marker_arrow_overhang=0,
                   marker_arrow_length_includes_head=False,
                   marker_arrow_head_starts_at_zero=False,
                   marker_alpha=1,
                   zorder=None,
                   **kwargs):
    """
	Draw a "frontline" in matplotlib,
	this is a line decorated with a pattern that is repeated at regular
	distances (independent of location of points making up the line),
	and keeping the same orientation with respect to that line.
	Examples are faults on geological maps, weather fronts, topographic
	highs and lows, etc.

	The easiest way to accomplish this is using markers or patches.
	However, the drawback of this solution is that the aspect ratio of
	the matplotlib axes should not be changed after the frontline is
	drawn. Doing so modifies the line curvature, resulting in a mismatch
	with the angles at which the markers have been drawn previously.
	Simple translation does not have this problem.

	:param x:
		array-like, X coordinates of points defining the line.
	:param y:
		array-like, Y coordinates of points defining the line.
	:param axes:
		axes instance in which frontline should be drawn.
	:param line_style:
		matplotlib linestyle definition.
		(default: "-")
	:param line_color:
		matplotlib color definition for line.
		(default: 'k')
	:param line_width:
		matplotlib line width in points.
		(default: 1)
	:param line_alpha:
		float, alpha transparency for line.
		(default: 1)
	:param marker_shape:
		str, style of front marker,
		one of ("polygon", "star", "asterisk", "circle", "arc",
		"arrow", "ellipse", "rectangle").
		Polygons, stars, asterisks and circles are implemented as
		matplotlib markers; arcs, arrows, ellipses and rectangles as
		matplotlib patches.
		marker_shape may also be a matplotlib patch instance, in that
		case, it is recommended to set lower left coordinate at (0, 0).
		(default: "polygon")
	:param marker_num_sides:
		int, number of sides of front marker. Only applies to polygons, stars,
		and asterisks.
		(default: 3)
	:param marker_angle:
		float, angle in degrees of front marker with respect to line.
		Does not apply to circles.
		(default: 0)
	:param marker_size:
		int, size of front marker in points.
		(default: 12)
	:param marker_interval:
		int, str or array-like, interval of front marker along line.
		- positive int: interval between markers in points
		- negative int (or zero): number of markers along line
		- str: marker interval will be rounded to closest value that
			results in even spacing along the line
		- array-like: distances along line, in range [0,1].
		(default: 24)
	:param marker_offset:
		int, offset of front marker with respect to line.
		(default: 0)
	:param marker_face_color:
		matplotlib color definition for fill color of front marker.
		(default: 'k')
	:param marker_edge_color:
		matplotlib color definition for line color of front marker.
		(default: 'k')
	:param marker_edge_width:
		int, line width of front marker in points.
		(default: 1)
	:param marker_aspect_ratio:
		float, aspect ratio of width (direction perpendicular to line)
		to length (direction along line) of front marker. Only applies
		to arcs, arrows, ellipses and rectangles.
		(default: 1.)
	:param marker_alternate_sides:
		bool, whether or not front markers should alternate between
		two sides of line.
		(default: False)
	:param marker_theta1:
		float, starting angle in degrees. Only applies to arcs.
		(default: 0)
	:param marker_theta2:
		float, ending angle in degrees. Only applies to arcs.
		(default: 180)
	:param marker_arrow_shape:
		str, arrow shape, one of ["full", "left", "right"].
		Only applies to arrows.
		(default: "full")
	:param marker_arrow_overhang:
		float, fraction that the arrow is swept back (0 means triangular shape).
		Can be negative or greater than one. Only applies to arrows.
		(default: 0.)
	:param marker_arrow_length_includes_head:
		bool, True if head is to be counted in calculating arrow length.
		Only applies to arrows.
		(default: False)
	:param marker_arrow_head_starts_at_zero:
		bool, if True, arrow head starts being drawn at coordinate 0 instead of
		ending at coordinate 0. Only applies to arrows.
		(default: False)
	:param marker_alpha:
		float, alpha transparency for front marker.
		(default: 1)
	:param zorder:
		int, sets the zorder for both the line and the front markers.
		(default: None)
	:param kwargs:
		remaining keyword arguments for matplotlib plot command

	:return:
		instance of :class:`FrontlineLegendHandler`, which can be passed
		(in handler_map dictionary) to matplotlib legend function
	"""
    dpi = ax.get_figure().dpi

    #legend_artist = None
    legend_handler = None

    ## Plot line first
    if not line_style in ("None", None):
        ax.plot(x,
                y,
                lw=line_width,
                color=line_color,
                ls=line_style,
                alpha=line_alpha,
                zorder=zorder,
                **kwargs)

    ## Transforms between data domain and display domain
    pt_to_pixel_transform = ax.get_figure().dpi_scale_trans.frozen().scale(1. /
                                                                           72)
    forward_transform = ax.transData
    inverse_transform = ax.transData.inverted()

    ## Compute cumulative distance along line and angles in pixel domain
    display_data_coords = forward_transform.transform(list(zip(x, y)))
    display_data_x, display_data_y = zip(*display_data_coords)
    display_data_angles = np.arctan2(np.diff(display_data_y),
                                     np.diff(display_data_x))
    ## mlab.distances_along_curve function was deprecated in mpl 2.2, and removed in 3.1
    #display_distance = mlab.distances_along_curve(display_data_coords)
    display_distance = (np.sum(np.diff(display_data_coords, axis=0)**2,
                               axis=1))**(1. / 2)
    display_cum_distance = np.concatenate([[0.], np.cumsum(display_distance)])

    if marker_shape == "asterisk" and marker_num_sides <= 2:
        marker_length = 0
    else:
        marker_length = marker_size
    ## Make sure marker_offset is 2D (along line / perpendicular to line)
    if isinstance(marker_offset, (int, float)):
        marker_offset = np.array([0, marker_offset])
    else:
        marker_offset = np.asarray(marker_offset[:2])

    ## Correct for a slight mismatch with dash lengths
    pt_corr_factor = 1.004

    ## Convert points to pixels (1 pt = 1/72 inch)
    marker_length_px = pt2pixel(marker_length, dpi, pt_corr_factor)
    marker_size_px = pt2pixel(marker_size, dpi, pt_corr_factor)
    marker_offset_px = pt2pixel(marker_offset, dpi, pt_corr_factor)

    ## Compute marker coordinates in data units based on
    ## interpolation of cumulative distance along line in display units
    start_distance = marker_length_px / 2. + marker_offset_px[0]
    end_distance = display_cum_distance[-1] - marker_length_px / 2.
    if isinstance(marker_interval, (list, np.ndarray)):
        marker_interval_px = np.asarray(marker_interval)
        display_marker_distances = np.interp(marker_interval_px, [0., 1.],
                                             [start_distance, end_distance])
    elif isinstance(marker_interval, basestring):
        marker_interval_px = pt2pixel(float(marker_interval), dpi,
                                      pt_corr_factor)
        num_markers = np.abs(
            np.round((end_distance - start_distance) / marker_interval_px)) + 1
        display_marker_distances = np.linspace(start_distance, end_distance,
                                               num_markers)
    elif marker_interval > 0:
        marker_interval_px = pt2pixel(marker_interval, dpi, pt_corr_factor)
        display_marker_distances = np.arange(start_distance, end_distance,
                                             marker_interval_px)
    elif marker_interval <= 0:
        display_marker_distances = np.linspace(start_distance, end_distance,
                                               np.abs(marker_interval))
    display_marker_x = np.interp(display_marker_distances,
                                 display_cum_distance, display_data_x)
    display_marker_y = np.interp(display_marker_distances,
                                 display_cum_distance, display_data_y)
    display_marker_angles = np.interp(display_marker_distances,
                                      display_cum_distance[1:],
                                      display_data_angles)
    if marker_alternate_sides:
        display_marker_angles[1::2] += np.pi
    display_marker_x += (marker_offset_px[1] *
                         np.cos(display_marker_angles + np.pi / 2))
    display_marker_y += (marker_offset_px[1] *
                         np.sin(display_marker_angles + np.pi / 2))
    marker_coords = inverse_transform.transform(
        list(zip(display_marker_x, display_marker_y)))

    patch_list = []

    ## Regular markers
    if marker_shape in ["polygon", "star", "asterisk", "circle"]:
        marker_shape_code = {
            "polygon": 0,
            "star": 1,
            "asterisk": 2,
            "circle": 3
        }[marker_shape]
        if marker_num_sides == 0:
            marker_shape_code = 3
        if marker_shape_code in (0, 1, 3):
            ## Convert marker size in diamter to size in area
            # TODO: find correct conversion for different shapes
            radius = marker_size / 2.
            marker_size = radius * (radius**2) / (marker_size)
        for i, (marker_x, marker_y) in enumerate(marker_coords):
            angle = np.degrees(display_marker_angles[i]) + marker_angle
            marker = (marker_num_sides, marker_shape_code, angle)
            ax.plot(marker_x,
                    marker_y,
                    linestyle='None',
                    marker=marker,
                    markersize=marker_size,
                    mec=marker_edge_color,
                    mfc=marker_face_color,
                    mew=marker_edge_width,
                    alpha=marker_alpha,
                    zorder=zorder)

        def legend_artist(self, legend, orig_handle, fontsize, handlebox):
            ## Note: In matplotlib versions 1.3.X and maybe also 1.4.X,
            ## the 'self' property needs to be removed...
            # TODO: marker_alternate_sides!
            marker = (marker_num_sides, marker_shape_code, marker_angle)
            x0, y0 = handlebox.xdescent, handlebox.ydescent
            width, height = handlebox.width, handlebox.height

            x = np.array([x0, x0 + width])
            y = np.ones_like(x) * (y0 + height / 2. - marker_size / 4.)
            line1 = mpl.lines.Line2D(x,
                                     y,
                                     lw=line_width,
                                     color=line_color,
                                     ls=line_style,
                                     alpha=line_alpha)
            x = np.arange(x0, x0 + width, marker_interval_px)
            y = np.ones_like(x) * (y0 + height / 2. - marker_size / 4.)
            line2 = mpl.lines.Line2D(x,
                                     y,
                                     linestyle='None',
                                     marker=marker,
                                     markersize=marker_size,
                                     mec=marker_edge_color,
                                     mfc=marker_face_color,
                                     mew=marker_edge_width,
                                     alpha=marker_alpha)

            handlebox.add_artist(line1)
            handlebox.add_artist(line2)
            return (line1, line2)

        if mpl.__version__ <= "1.3.1":
            legend_handler = legend_artist
        else:
            #legend_handler = type('FrontlineLegendHandler',  (), {'legend_artist': classmethod(legend_artist)})()
            legend_handler = FrontlineLegendHandler()
            if PY2:
                legend_handler.legend_artist = MethodType(
                    legend_artist, legend_handler, FrontlineLegendHandler)
            else:
                legend_handler.legend_artist = MethodType(
                    legend_artist, legend_handler)

    ## Patches
    elif marker_shape in ["arc", "arrow", "ellipse", "rectangle"
                          ] or isinstance(marker_shape, mpl.patches.Patch):
        for i in range(len(display_marker_x)):
            dmx, dmy = display_marker_x[i], display_marker_y[i]
            angle = np.degrees(display_marker_angles[i]) + marker_angle
            if marker_shape == "arc":
                patch = mpl.patches.Arc((0, 0),
                                        marker_size_px,
                                        marker_size_px * marker_aspect_ratio,
                                        theta1=marker_theta1,
                                        theta2=marker_theta2,
                                        angle=angle,
                                        fill=False)
                ## Workaround because in some mpl versions, theta1,2 params are ignored...
                patch._path = mpl.patches.Path.arc(marker_theta1,
                                                   marker_theta2)
            elif marker_shape == "arrow":
                dx = np.cos(np.radians(angle)) * marker_length_px / 2.
                dy = np.sin(np.radians(angle)) * marker_length_px / 2.
                patch = mpl.patches.FancyArrow(
                    0,
                    0,
                    dx,
                    dy,
                    width=marker_edge_width,
                    head_width=marker_size_px / 2 * marker_aspect_ratio,
                    head_length=marker_size_px / 2,
                    shape=marker_arrow_shape,
                    overhang=marker_arrow_overhang,
                    length_includes_head=marker_arrow_length_includes_head,
                    head_starts_at_zero=marker_arrow_head_starts_at_zero)
            elif marker_shape == "ellipse":
                patch = mpl.patches.Ellipse(
                    (0, 0),
                    marker_size_px,
                    marker_size_px * marker_aspect_ratio,
                    angle=angle)
            elif marker_shape == "rectangle":
                patch = mpl.patches.Rectangle(
                    (0, 0),
                    marker_size_px,
                    marker_size_px * marker_aspect_ratio,
                    angle=angle)
            elif isinstance(marker_shape, mpl.patches.Patch):
                patch = marker_shape
                tf = pt_to_pixel_transform + mpl.transforms.Affine2D(
                ).rotate_deg(angle)
                patch.set_transform(tf)
            ## Shift patch back half its width if needed
            if marker_shape in ("arrow", "rectangle") or isinstance(
                    marker_shape, mpl.patches.Patch):
                dmx += (-np.cos(np.radians(angle)) * marker_length_px / 2.)
                dmy += (-np.sin(np.radians(angle)) * marker_length_px / 2.)
            tf = mpl.transforms.Affine2D().translate(dmx, dmy)

            ## Copy tranformed vertices to a new patch, otherwise patches
            ## do not move with curve if plot is translated
            path = patch.get_path()
            path = path.cleaned(transform=(patch.get_transform() + tf +
                                           inverse_transform))
            vertices, codes = path.vertices, path.codes
            if marker_shape == "arc":
                ## Modify codes to make filling of arc patches possible
                ## Alternatively, see: https://stackoverflow.com/questions/30642391/how-to-draw-a-filled-arc-in-matplotlib
                if codes[-1] != mpl.path.Path.CLOSEPOLY:
                    codes = np.concatenate(
                        [codes[:-1], [mpl.path.Path.CLOSEPOLY]])

            # Note: the following lines work too, but in the case of arcs,
            # some points seem to disappear
            #patch.set_transform(tf + inverse_transform)
            #vertices = patch.get_verts()
            #codes = [mpl.path.Path.MOVETO] + [mpl.path.Path.LINETO] * (len(vertices) - 2) + [mpl.path.Path.CLOSEPOLY]

            ## Copy patch (and apply color and linewidth options)
            path = mpl.path.Path(vertices, codes)
            #patch = mpl.patches.PathPatch(path, ec=marker_edge_color, fc=marker_face_color, lw=marker_edge_width, alpha=marker_alpha)
            #patch.set_zorder(zorder)
            #ax.add_patch(patch)
            patch = mpl.patches.PathPatch(path)
            patch_list.append(patch)

        ## Plot patches collectively, and apply color and linewidth options
        col = PatchCollection(patches=patch_list,
                              facecolors=marker_face_color,
                              edgecolors=marker_edge_color,
                              linewidths=marker_edge_width,
                              alpha=marker_alpha)
        col.set_zorder(zorder)
        ax.add_collection(col)

        # TODO: custom legend handler for patches

    return legend_handler
Пример #16
0
def plot(network, margin=0.05, ax=None, geomap=True, projection=None,
         bus_colors='b', line_colors={'Line':'g', 'Link':'cyan'}, bus_sizes=10,
         line_widths={'Line':2, 'Link':2},
         flow=None, title="", line_cmap=None, bus_cmap=None, boundaries=None,
         geometry=False, branch_components=['Line', 'Link'], jitter=None,
         basemap=None, basemap_parameters=None, color_geomap=None):
    """
    Plot the network buses and lines using matplotlib and Basemap.

    Parameters
    ----------
    margin : float
        Margin at the sides as proportion of distance between max/min x,y
    ax : matplotlib ax, defaults to plt.gca()
        Axis to which to plot the network
    geomap: bool/str, default True
        Switch to use Basemap or Cartopy (depends on what is installed).
        If string is passed, it will be used as a resolution argument.
        For Basemap users 'c' (crude), 'l' (low), 'i' (intermediate),
        'h' (high), 'f' (full) are valid resolutions options.
        For Cartopy users '10m', '50m', '110m' are valid resolutions options.
    projection: cartopy.crs.Projection, defaults to None
        Define the projection of your geomap, only valid if cartopy is
        installed. If None (default) is passed the projection for cartopy
        is set to cartopy.crs.PlateCarree
    bus_colors : dict/pandas.Series
        Colors for the buses, defaults to "b"
    bus_sizes : dict/pandas.Series
        Sizes of bus points, defaults to 10
    line_colors : dict/pandas.Series
        Colors for the lines, defaults to "g" for Lines and "cyan" for
        Links. Colors for branches other than Lines can be
        specified using a pandas Series with a MultiIndex.
    line_widths : dict/pandas.Series
        Widths of lines, defaults to 2. Widths for branches other
        than Lines can be specified using a pandas Series with a
        MultiIndex.
    flow : snapshot/pandas.Series/function/string
        Flow to be displayed in the plot, defaults to None. If an element of
        network.snapshots is given, the flow at this timestamp will be
        displayed. If an aggregation function is given, is will be applied
        to the total network flow via pandas.DataFrame.agg (accepts also
        function names). Otherwise flows can be specified by passing a pandas
        Series with MultiIndex including all necessary branch components.
        Use the line_widths argument to additionally adjust the size of the
        flow arrows.
    title : string
        Graph title
    line_cmap : plt.cm.ColorMap/str|dict
        If line_colors are floats, this color map will assign the colors.
        Use a dict to specify colormaps for more than one branch type.
    bus_cmap : plt.cm.ColorMap/str
        If bus_colors are floats, this color map will assign the colors
    boundaries : list of four floats
        Boundaries of the plot in format [x1,x2,y1,y2]
    branch_components : list of str
        Branch components to be plotted, defaults to Line and Link.
    jitter : None|float
        Amount of random noise to add to bus positions to distinguish
        overlapping buses
    basemap_parameters : dict
        Specify a dict with additional constructor parameters for the
        Basemap. Will disable Cartopy.
        Use this feature to set a custom projection.
        (e.g. `{'projection': 'tmerc', 'lon_0':10.0, 'lat_0':50.0}`)
    color_geomap : dict or bool
        Specify colors to paint land and sea areas in.
        If True, it defaults to `{'ocean': 'lightblue', 'land': 'whitesmoke'}`.
        If no dictionary is provided, colors are white.

    Returns
    -------
    bus_collection, branch_collection1, ... : tuple of Collections
        Collections for buses and branches.
    """
    defaults_for_branches = pd.Series({
        'Link': dict(color="cyan", width=2),
        'Line': dict(color="b", width=2),
        'Transformer': dict(color='green', width=2)
    }).rename_axis('component')

    if not plt_present:
        logger.error("Matplotlib is not present, so plotting won't work.")
        return

    if basemap is not None:
        logger.warning("argument `basemap` is deprecated, "
                       "use `geomap` instead.")
        geomap = basemap

    if geomap:
        if not (cartopy_present or basemap_present):
            # Not suggesting Basemap since it is being deprecated
            logger.warning("Cartopy needs to be installed to use `geomap=True`.")
            geomap = False

        # Use cartopy by default, fall back on basemap
        use_basemap = False
        use_cartopy = cartopy_present
        if not use_cartopy:
            use_basemap = basemap_present

        # If the user specifies basemap parameters, they prefer
        # basemap over cartopy.
        # (This means that you can force the use of basemap by
        # setting `basemap_parameters={}`)
        if basemap_present:
            if basemap_parameters is not None:
                logger.warning("Basemap is being deprecated, consider "
                               "switching to Cartopy.")
                use_basemap = True
                use_cartopy = False

        if use_cartopy:
            if projection is None:
                projection = get_projection_from_crs(network.srid)

            if ax is None:
                ax = plt.gca(projection=projection)
            else:
                assert isinstance(ax, cartopy.mpl.geoaxes.GeoAxesSubplot), (
                        'The passed axis is not a GeoAxesSubplot. You can '
                        'create one with: \nimport cartopy.crs as ccrs \n'
                        'fig, ax = plt.subplots('
                        'subplot_kw={"projection":ccrs.PlateCarree()})')
    elif ax is None:
        ax = plt.gca()

    x, y = network.buses["x"],  network.buses["y"]

    axis_transform = ax.transData

    if geomap:
        if use_cartopy:
            axis_transform = draw_map_cartopy(network, x, y, ax,
                    boundaries, margin, geomap, color_geomap)
            new_coords = pd.DataFrame(
                    ax.projection.transform_points(axis_transform,
                                                   x.values, y.values),
                       index=network.buses.index, columns=['x', 'y', 'z'])
            x, y = new_coords['x'], new_coords['y']
        elif use_basemap:
            basemap_transform = draw_map_basemap(network, x, y, ax,
                    boundaries, margin, geomap, basemap_parameters, color_geomap)

            # A non-standard projection might be used; the easiest way to
            # support this is to tranform the bus coordinates.
            x, y = basemap_transform(x.values, y.values)
            x = pd.Series(x, network.buses.index)
            y = pd.Series(y, network.buses.index)

    if jitter is not None:
        x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
        y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))

    if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index, pd.MultiIndex):
        # We are drawing pies to show all the different shares
        assert len(bus_sizes.index.levels[0].difference(network.buses.index)) == 0, \
            "The first MultiIndex level of bus_sizes must contain buses"
        assert (isinstance(bus_colors, dict) and
                set(bus_colors).issuperset(bus_sizes.index.levels[1])), \
            "bus_colors must be a dictionary defining a color for each element " \
            "in the second MultiIndex level of bus_sizes"

        bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)
        if geomap:
            bus_sizes *= projected_area_factor(ax, network.srid)**2

        patches = []
        for b_i in bus_sizes.index.levels[0]:
            s = bus_sizes.loc[b_i]
            radius = s.sum()**0.5
            if radius == 0.0:
                ratios = s
            else:
                ratios = s/s.sum()

            start = 0.25
            for i, ratio in ratios.iteritems():
                patches.append(Wedge((x.at[b_i], y.at[b_i]), radius,
                                     360*start, 360*(start+ratio),
                                     facecolor=bus_colors[i]))
                start += ratio
        bus_collection = PatchCollection(patches, match_original=True)
        ax.add_collection(bus_collection)
    else:
        c = pd.Series(bus_colors, index=network.buses.index)
        s = pd.Series(bus_sizes, index=network.buses.index, dtype="float").fillna(10)
        bus_collection = ax.scatter(x, y, c=c, s=s, cmap=bus_cmap, edgecolor='face')

    def as_branch_series(ser):
        # ensure that this function always return a multiindexed series
        if isinstance(ser, dict) and set(ser).issubset(branch_components):
            return pd.concat(
                    {c.name: pd.Series(s, index=c.df.index) for c, s in
                         zip(network.iterate_components(ser.keys()), ser.values())},
                    names=['component', 'name'])
        elif isinstance(ser, pd.Series) and isinstance(ser.index, pd.MultiIndex):
            return ser.rename_axis(index=['component', 'name'])
        else:
            ser =  pd.Series(ser, network.lines.index)
            return pd.concat([ser], axis=0, keys=['Line'],
                             names=['component', 'name']).fillna(0)

    line_colors = as_branch_series(line_colors)
    line_widths = as_branch_series(line_widths)

    if not isinstance(line_cmap, dict):
        line_cmap = {'Line': line_cmap}

    branch_collections = []

    if flow is not None:
        flow = (_flow_ds_from_arg(flow, network, branch_components)
                .pipe(as_branch_series)
                .div(sum(len(t.df) for t in
                         network.iterate_components(branch_components)) + 100))
        flow = flow.mul(line_widths[flow.index], fill_value=1)
        # update the line width, allows to set line widths separately from flows
        line_widths.update((5 * flow.abs()).pipe(np.sqrt))
        arrows = directed_flow(network, flow, x=x, y=y, ax=ax, geomap=geomap,
                               branch_colors=line_colors,
                               branch_comps=branch_components,
                               cmap=line_cmap['Line'])
        branch_collections.append(arrows)


    for c in network.iterate_components(branch_components):
        l_defaults = defaults_for_branches[c.name]
        l_widths = line_widths.get(c.name, l_defaults['width'])
        l_nums = None
        l_colors = line_colors.get(c.name, l_defaults['color'])

        if isinstance(l_colors, pd.Series):
            if issubclass(l_colors.dtype.type, np.number):
                l_nums = l_colors
                l_colors = None
            else:
                l_colors.fillna(l_defaults['color'], inplace=True)

        if not geometry:
            segments = (np.asarray(((c.df.bus0.map(x),
                                     c.df.bus0.map(y)),
                                    (c.df.bus1.map(x),
                                     c.df.bus1.map(y))))
                        .transpose(2, 0, 1))
        else:
            from shapely.wkt import loads
            from shapely.geometry import LineString
            linestrings = c.df.geometry[lambda ds: ds != ''].map(loads)
            assert all(isinstance(ls, LineString) for ls in linestrings), (
                "The WKT-encoded geometry in the 'geometry' column must be "
                "composed of LineStrings")
            segments = np.asarray(list(linestrings.map(np.asarray)))

        l_collection = LineCollection(segments,
                                      linewidths=l_widths,
                                      antialiaseds=(1,),
                                      colors=l_colors,
                                      transOffset=ax.transData)

        if l_nums is not None:
            l_collection.set_array(np.asarray(l_nums))
            l_collection.set_cmap(line_cmap.get(c.name, None))
            l_collection.autoscale()

        ax.add_collection(l_collection)
        l_collection.set_zorder(3)

        branch_collections.append(l_collection)

    bus_collection.set_zorder(4)

    ax.update_datalim(compute_bbox_with_margins(margin, x, y))
    ax.autoscale_view()

    if geomap:
        if use_cartopy:
            ax.outline_patch.set_visible(False)
        ax.axis('off')

    ax.set_title(title)

    return (bus_collection,) + tuple(branch_collections)
Пример #17
0
def plot(network, margin=0.05, ax=None, basemap=True, bus_colors='b',
         line_colors='g', bus_sizes=10, line_widths=2, title="",
         line_cmap=None, bus_cmap=None, boundaries=None,
         geometry=False, branch_components=['Line', 'Link'], jitter=None):
    """
    Plot the network buses and lines using matplotlib and Basemap.

    Parameters
    ----------
    margin : float
        Margin at the sides as proportion of distance between max/min x,y
    ax : matplotlib ax, defaults to plt.gca()
        Axis to which to plot the network
    basemap : bool, default True
        Switch to use Basemap
    bus_colors : dict/pandas.Series
        Colors for the buses, defaults to "b"
    bus_sizes : dict/pandas.Series
        Sizes of bus points, defaults to 10
    line_colors : dict/pandas.Series
        Colors for the lines, defaults to "g" for Lines and "cyan" for
        Links. Colors for branches other than Lines can be
        specified using a pandas Series with a MultiIndex.
    line_widths : dict/pandas.Series
        Widths of lines, defaults to 2. Widths for branches other
        than Lines can be specified using a pandas Series with a
        MultiIndex.
    title : string
        Graph title
    line_cmap : plt.cm.ColorMap/str|dict
        If line_colors are floats, this color map will assign the colors.
        Use a dict to specify colormaps for more than one branch type.
    bus_cmap : plt.cm.ColorMap/str
        If bus_colors are floats, this color map will assign the colors
    boundaries : list of four floats
        Boundaries of the plot in format [x1,x2,y1,y2]
    branch_components : list of str
        Branch components to be plotted, defaults to Line and Link.
    jitter : None|float
        Amount of random noise to add to bus positions to distinguish
        overlapping buses

    Returns
    -------
    bus_collection, branch_collection1, ... : tuple of Collections
        Collections for buses and branches.
    """

    defaults_for_branches = {
        'Link': dict(color="cyan", width=2),
        'Line': dict(color="b", width=2),
        'Transformer': dict(color='green', width=2)
    }

    if not plt_present:
        logger.error("Matplotlib is not present, so plotting won't work.")
        return

    if ax is None:
        ax = plt.gca()

    def compute_bbox_with_margins(margin, x, y):
        #set margins
        pos = np.asarray((x, y))
        minxy, maxxy = pos.min(axis=1), pos.max(axis=1)
        xy1 = minxy - margin*(maxxy - minxy)
        xy2 = maxxy + margin*(maxxy - minxy)
        return tuple(xy1), tuple(xy2)

    x = network.buses["x"]
    y = network.buses["y"]

    if jitter is not None:
        x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
        y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))

    if basemap and basemap_present:
        if boundaries is None:
            (x1, y1), (x2, y2) = compute_bbox_with_margins(margin, x, y)
        else:
            x1, x2, y1, y2 = boundaries
        bmap = Basemap(resolution='l', epsg=network.srid,
                       llcrnrlat=y1, urcrnrlat=y2, llcrnrlon=x1,
                       urcrnrlon=x2, ax=ax)
        bmap.drawcountries()
        bmap.drawcoastlines()

        x, y = bmap(x.values, y.values)
        x = pd.Series(x, network.buses.index)
        y = pd.Series(y, network.buses.index)

    if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index, pd.MultiIndex):
        # We are drawing pies to show all the different shares
        assert len(network.buses.index.difference(bus_sizes.index.levels[0])) == 0, \
            "The first MultiIndex level of bus_sizes must contain buses"
        assert isinstance(bus_colors, dict) and set(bus_colors).issuperset(bus_sizes.index.levels[1]), \
            "bus_colors must be a dictionary defining a color for each element " \
            "in the second MultiIndex level of bus_sizes"

        bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)

        patches = []
        for b_i in bus_sizes.index.levels[0]:
            s = bus_sizes.loc[b_i]
            radius = s.sum()**0.5
            ratios = s/s.sum()

            start = 0.25
            for i, ratio in ratios.iteritems():
                patches.append(Wedge((x.at[b_i], y.at[b_i]), radius,
                                     360*start, 360*(start+ratio),
                                     facecolor=bus_colors[i]))
                start += ratio
        bus_collection = PatchCollection(patches, match_original=True)
        ax.add_collection(bus_collection)
    else:
        c = pd.Series(bus_colors, index=network.buses.index)
        if c.dtype == np.dtype('O'):
            c.fillna("b", inplace=True)
            c = list(c.values)
        s = pd.Series(bus_sizes, index=network.buses.index, dtype="float").fillna(10)
        bus_collection = ax.scatter(x, y, c=c, s=s, cmap=bus_cmap)

    def as_branch_series(ser):
        if isinstance(ser, dict) and set(ser).issubset(branch_components):
            return pd.Series(ser)
        elif isinstance(ser, pd.Series):
            if isinstance(ser.index, pd.MultiIndex):
                return ser
            index = ser.index
            ser = ser.values
        else:
            index = network.lines.index
        return pd.Series(ser,
                         index=pd.MultiIndex(levels=(["Line"], index),
                                             labels=(np.zeros(len(index)),
                                                     np.arange(len(index)))))

    line_colors = as_branch_series(line_colors)
    line_widths = as_branch_series(line_widths)
    if not isinstance(line_cmap, dict):
        line_cmap = {'Line': line_cmap}

    branch_collections = []
    for c in network.iterate_components(branch_components):
        l_defaults = defaults_for_branches[c.name]
        l_widths = line_widths.get(c.name, l_defaults['width'])
        l_nums = None
        l_colors = line_colors.get(c.name, l_defaults['color'])

        if isinstance(l_colors, pd.Series):
            if issubclass(l_colors.dtype.type, np.number):
                l_nums = l_colors
                l_colors = None
            else:
                l_colors.fillna(l_defaults['color'], inplace=True)

        if not geometry:
            segments = (np.asarray(((c.df.bus0.map(x),
                                     c.df.bus0.map(y)),
                                    (c.df.bus1.map(x),
                                     c.df.bus1.map(y))))
                        .transpose(2, 0, 1))
        else:
            from shapely.wkt import loads
            from shapely.geometry import LineString
            linestrings = c.df.geometry.map(loads)
            assert all(isinstance(ls, LineString) for ls in linestrings), \
                "The WKT-encoded geometry in the 'geometry' column must be composed of LineStrings"
            segments = np.asarray(list(linestrings.map(np.asarray)))
            if basemap and basemap_present:
                segments = np.transpose(bmap(*np.transpose(segments, (2, 0, 1))), (1, 2, 0))

        l_collection = LineCollection(segments,
                                      linewidths=l_widths,
                                      antialiaseds=(1,),
                                      colors=l_colors,
                                      transOffset=ax.transData)

        if l_nums is not None:
            l_collection.set_array(np.asarray(l_nums))
            l_collection.set_cmap(line_cmap.get(c.name, None))
            l_collection.autoscale()

        ax.add_collection(l_collection)
        l_collection.set_zorder(1)

        branch_collections.append(l_collection)

    bus_collection.set_zorder(2)

    ax.update_datalim(compute_bbox_with_margins(margin, x, y))
    ax.autoscale_view()

    ax.set_title(title)

    return (bus_collection,) + tuple(branch_collections)
Пример #18
0
    def odd_import(self, ax, odd):
        if odd[1] == -1:
            odd[1] = 0
        self.odd_tool.update_odd()
        file_name = self.odd_tool.odd_list[odd[0]]
        odd_path = 'scenarios/' + str(odd[0]) + '/' + file_name[odd[1]] + '.json'

        with open(odd_path, 'r') as load_f:
            load_dict = load(load_f)
        list_dict = list(load_dict.keys())
        scenario = load_dict[list_dict[0]]

        try:
            lines1 = scenario['lines']['boundary_lines']
            lines2 = scenario['lines']['solid_lines']
            lines3 = scenario['lines']['dashed_lines']
            lines4 = scenario['lines']['thick_lines']
            objects = scenario['objects']
            common_rotation = scenario['common_rotation']
            common_offset = scenario['common_offset']
        except Exception:
            lines1 = []
            lines2 = []
            lines3 = []
            lines4 = []
            objects = []
            common_rotation = 0.0
            common_offset = [0, 0]
            QMessageBox.warning(self, 'Warning', 'Scenario file format error!')

        boundary_lines = self.line_rotation(lines1, common_rotation)
        solid_lines = self.line_rotation(lines2, common_rotation)
        dashed_lines = self.line_rotation(lines3, common_rotation)
        thick_lines = self.line_rotation(lines4, common_rotation)

        for l in boundary_lines:
            for m in l:
                m[0] = (m[0] + common_offset[0]) * self.zoom_factor
                m[1] = (m[1] + common_offset[1]) * self.zoom_factor

        for l in solid_lines:
            for m in l:
                m[0] = (m[0] + common_offset[0]) * self.zoom_factor
                m[1] = (m[1] + common_offset[1]) * self.zoom_factor

        for l in dashed_lines:
            for m in l:
                m[0] = (m[0] + common_offset[0]) * self.zoom_factor
                m[1] = (m[1] + common_offset[1]) * self.zoom_factor

        for l in thick_lines:
            for m in l:
                m[0] = (m[0] + common_offset[0]) * self.zoom_factor
                m[1] = (m[1] + common_offset[1]) * self.zoom_factor

        line_segments1 = LineCollection(boundary_lines, linewidths=2, colors=(0, 0, 0), linestyles='solid')
        line_segments1.set_zorder(0)
        line_segments2 = LineCollection(solid_lines, linewidths=1, colors=(0.1, 0.1, 0.1), linestyles='solid')
        line_segments2.set_zorder(0)
        line_segments3 = LineCollection(dashed_lines, linewidths=1, colors=(0.2, 0.2, 0.2), linestyles=(0, (20, 20)))
        line_segments3.set_zorder(0)
        line_segments4 = LineCollection(thick_lines, linewidths=4.5, colors=(0.8, 0.8, 0.8), linestyles='solid')
        line_segments4.set_zorder(0)

        self.ax.add_collection(line_segments1)
        self.ax.add_collection(line_segments2)
        self.ax.add_collection(line_segments3)
        self.ax.add_collection(line_segments4)

        patch = []
        for obj in range(len(objects)):
            objects_xy = self.object_rotation(objects[obj][1], common_rotation)
            objects_xy[0] = objects_xy[0] + common_offset[0]
            objects_xy[1] = objects_xy[1] + common_offset[1]
            objects_r = objects[obj][2] + common_rotation
            patch.append(self.add_object(objects[obj][0], objects_xy, objects_r))
            self.ax.text((objects_xy[0]) * self.zoom_factor, (objects_xy[1] + 2.5) * self.zoom_factor,
                         'obj' + str(obj + 1), fontsize=10)
        patch_segment = PatchCollection(patch, facecolors='gray')
        patch_segment.set_zorder(2)
        ax.add_collection(patch_segment)
        return patch
Пример #19
0
#7 substract .1 subtract .05 equals -.15

print(X3)
ax3.set_ylabel('cov ' + betastr)
#fig.set_title(betastr + ' for reg 1st level beta ~ cov score')
ax3.set_xticks(X3)
ax3data = [data_dict[reg][0] for reg in regs1]
ax3dots = [data_dict[reg][1] for reg in regs1]
colors1 = ['purple', 'peachpuff', 'darkred', 'yellow', 'darkblue', 'lightskyblue',
          'darkolivegreen', 'lightgreen']
patches = []
for i, dot in enumerate(ax3dots):
    circ = mpatches.Ellipse(xy = (X3[i] + .125, ax3dots[i]), width = .05, height = .05, color = 'black')
    patches.append(circ)
patches = PatchCollection(patches, match_original=True, zorder=10)
patches.set_zorder(20)
ax3.set_xticks([x + .125 for x in X3])
for i in range(8):
    ax3.bar(X3[i], ax3data[i], color=colors1[i], width=0.25, alpha = .8, zorder=2)
ax3.set_xticklabels(labs1)
ax3.set_ylim(min(ax3data) - .25, max(ax3data) + .25)
line_xs = ax3.get_xlim()
line_ys = [0, 0]
ax3.add_line(lines.Line2D(line_xs, line_ys, linewidth=1, alpha=.5, color='black'))
ax3.add_collection(patches)
ax3.set_xlim(X3[0] - .5, X3[-1] + .7)

X2= []
for i, x in enumerate(X3[::2]):
    X2.append(np.mean([X3[2*i], X3[2*i + 1]]))
X2 = np.arange(4)
Пример #20
0
    def __init__(
            self,
            model_name=None,
            centerxy=np.array([]),
            blocks=np.array([]),
            values=np.array([]),
            values_log=0,
            bck=1,
    ):
        """
        :param model_name: path to the file to be created containing the output
        :param centerxy: x-y coordinates of the center of the cells
        :param blocks: coordinates of the corners of the different blocks
        :param values: array containing the value assigned to each block
        :param values_log: flag indicating if values should be log transformed or not
        :param bck: background value
        """

        if centerxy.any(
        ):  # If block center coordinates are provided, plot them.
            xs = centerxy[:, 0]
            ys = centerxy[:, 1]
            fig, ax = plt.subplots()
            pts = ax.scatter(xs, ys, c="black", alpha=0.5)
        else:  # Else, if only blocks are provided, the centers are computed based on the mean of the coordinates of
            # their corners
            try:
                centerxy = np.array([np.mean(b, axis=0)
                                     for b in blocks])  # Mean computed
                xs = centerxy[:, 0]
                ys = centerxy[:, 1]
                fig, ax = plt.subplots()
                pts = ax.scatter(xs, ys, c="black", alpha=0.5)
            except Exception as e:
                print(e)
                exit()

        # I use this option to put the points in first plan, as to always see them.
        pts.set_zorder(2)

        self.points = centerxy

        self.bck = bck

        self.model_name = model_name

        self.ax = ax  # Axes object

        self.blocks = blocks

        self.ax.set_title("Now setting value: 0.0")  # Initial title

        self.ax.set_facecolor((0.86, 0.86, 0.86))  # Set gray background
        # Adjusts plt dimensions to insert colorbar, textbox...
        plt.subplots_adjust(bottom=0.2, right=0.8)

        self.vals = (
            []
        )  # List that will contain the values assigned to the different polygons

        self.cmap = cm.get_cmap("jet")  # Color map function to be later used

        # It is necessary to define a different axes for the box (normalized)
        axbox = plt.axes([0.4, 0.07, 0.1, 0.07])
        # 4 - tuple of floats
        # rect = [left, bottom, width, height]. A new axes is added
        # with dimensions rect in normalized (0, 1) units using ~.Figure.add_axes on the current figure.

        # Location of colorbar for the user-defined polygons
        self.axcb = plt.axes([0.82, 0.07, 0.015, 0.83])
        bounds = [0] + [1]  # Adding one bound more to have enough intervals
        ticks = [0]  # Ticks - default value
        cols = colors.ListedColormap([self.cmap(v) for v in ticks])
        cbnorm = colors.BoundaryNorm(bounds, cols.N)
        mcb = colorbar.ColorbarBase(
            self.axcb,
            cmap=cols,
            norm=cbnorm,
            boundaries=bounds,
            ticks=ticks,
            ticklocation="right",
            orientation="vertical",
        )
        #
        #         textstr = """Select points in the figure by enclosing them within a polygon.
        # Press the 'esc' key to start a new polygon.
        # Try holding the 'shift' key to move all of the vertices.
        # Try holding the 'ctrl' key to move a single vertex."""
        #
        #         #axtxt = plt.axes([0.15, 0.0, 0.2, 0.15])
        #         props = dict(boxstyle='round', facecolor='green', alpha=0.5)
        #         ax.text(0, -0.2, textstr, transform=ax.transAxes, fontsize=10, bbox=props)

        # Text box to input the value to input
        self.vinput = TextBox(axbox, label=None, initial="0")
        # What happens when pressing enter
        self.vinput.on_submit(self.button_submit)

        self.index = []  # List that will contain the final results !

        # Creates a canvas from the ax of the scatter plot
        self.canvas = self.ax.figure.canvas

        self.collection = pts  # 'collection' is the scatter plot
        # Necessary for later to define if points enclosed by polygon - basically
        self.xys = pts.get_offsets()
        # equals to the - x-y coordinates of the different points.
        self.Npts = len(self.xys)  # Number of points

        # Ensure that we have separate colors for each object
        self.fc = pts.get_facecolors()  # Gets the rgb of the points

        if not values.any():
            facecolors = "gray"
            alpha = 0.35  # Opacity of the polygons - if no values assigned - soft gray
        else:
            cmap2 = cm.get_cmap("coolwarm")
            if values_log:
                # Making a nice linear space
                itv = 10**np.linspace(np.log10(min(values)),
                                      np.log10(max(values)), 12)
                # out ouf log values to represent some ticks on the colorbar
                norm2 = colors.LogNorm(
                    vmin=min(values),
                    vmax=max(values))  # Log norm for color bar
                # Necessary arg to produce the color scale
                formatter = LogFormatter(10, labelOnlyBase=False)
            else:
                itv = np.linspace(min(values), max(values), 8)  # Linear space
                norm2 = colors.Normalize(vmin=min(values), vmax=max(values))
                formatter = None
            # Individual color of each polygon
            facecolors = [cmap2(norm2(v)) for v in values]
            alpha = 0.6  # Opacity of the polygons

            # Colorbar if initial values present - plotting a nice color bar
            ticks2 = [round(v, 1) for v in itv]
            plt.subplots_adjust(left=0.2, bottom=0.2, right=0.8)
            axcb1 = plt.axes([0.15, 0.07, 0.015, 0.83])
            cb1 = colorbar.ColorbarBase(
                axcb1,
                cmap=cmap2,
                norm=norm2,
                ticks=ticks2,
                boundaries=None,
                ticklocation="left",
                format=formatter,
                orientation="vertical",
            )
            cb1.set_ticklabels(ticks2)  # Setting the proper labels

        if (
                self.blocks.any()
        ):  # If blocks are provided. I should change this as the direction this code is going is
            # to provide blocks by default

            xs = self.blocks[:, :, 0]  # x-coordinates blocks corners
            ys = self.blocks[:, :, 1]  # y-coordinates blocks corners

            patches = [
            ]  # Coloring the blocks, in gray or with different colors
            for b in blocks:
                polygon = Polygon(b, closed=True)
                patches.append(polygon)
            p = PatchCollection(patches,
                                alpha=alpha,
                                facecolors=facecolors,
                                edgecolors="black")
            p.set_zorder(0)
            self.ax.add_collection(p)

            # 5% padding in x-direction for visualization
            padx = (xs.max() - xs.min()) * 0.05
            pady = (ys.max() - ys.min()) * 0.05  # 5% padding
            self.ax.set_xlim(xs.min() - padx, xs.max() + padx)
            self.ax.set_ylim(ys.min() - pady, ys.max() + pady)

        self.collection.set_facecolors(facecolors)  # Coloring the points

        # Polygon selector object
        self.poly = PolygonSelector(self.ax, self.onselect)
        self.ind = []  # Initiates the ind list for a new polygon!

        def handle_close(evt):  # Function that disconnects the PolygonSelector
            self.disconnect()

        # When closing window, finishes the job
        self.canvas.mpl_connect("close_event", handle_close)

        # Final results, array filled with background value.
        self.final_results = np.ones(len(self.points)) * self.bck
Пример #21
0
    if zone == 'R':
        color = 'y'
    elif zone == 'C':
        color = 'c'
    elif zone == 'M':
        color = 'darkmagenta'
    elif zone == 'O':
        color = 'ivory'
    elif zone == 'P':
        color = 'forestgreen'

    patches = [mplPolygon(np.array(shape), True)]
    pc = PatchCollection(patches)
    pc.set_alpha(.7)
    pc.set_facecolor(color)
    pc.set_zorder(2)
    pc.set_linewidth(.1)
    pc.set_edgecolor('k')
    ax.add_collection(pc)

#####################################################################
#				Add Roads shapefile to map 							#
#####################################################################
map.readshapefile('RoadCenterlines/RoadCenterlines_2',
                  'RoadCenterlines_2',
                  linewidth=.2)

#####################################################################
#						TRY CODE HERE 								#
#####################################################################
Пример #22
0
def plot_radial_histogram(angles,
                          counts,
                          all_angles=None,
                          include_labels=False,
                          offset=180.0,
                          direction=-1,
                          closed=False,
                          color=STIM_COLOR):
    if all_angles is None:
        if len(angles) < 2:
            all_angles = np.linspace(0, 315, 8)
        else:
            all_angles = angles

    dth = (all_angles[1] - all_angles[0]) * 0.5

    if len(counts) == 0:
        max_count = 1
    else:
        max_count = max(counts)

    wedges = []
    for count, angle in zip(counts, angles):
        angle = angle*direction + offset
        wedge = mpatches.Wedge((0,0), count, angle-dth, angle+dth)
        wedges.append(wedge)

    wedge_coll = PatchCollection(wedges)
    wedge_coll.set_facecolor(color)
    wedge_coll.set_zorder(2)

    angles_rad = (all_angles*direction + offset)*np.pi/180.0

    if closed:
        border_coll = cplots.radial_circles([max_count])
    else:
        border_coll = cplots.radial_arcs([max_count], 
                                         min(angles_rad), 
                                         max(angles_rad))
    border_coll.set_facecolor((0,0,0,0))
    border_coll.set_zorder(1)

    line_coll = cplots.angle_lines(angles_rad, 0, max_count)
    line_coll.set_edgecolor((0,0,0,1))
    line_coll.set_linestyle(":")
    line_coll.set_zorder(1)

    ax = plt.gca()
    ax.add_collection(wedge_coll)
    ax.add_collection(border_coll)
    ax.add_collection(line_coll)

    if include_labels:
        cplots.add_angle_labels(ax, angles_rad, all_angles.astype(int), max_count, (0,0,0,1), offset=max_count*0.1)
        ax.set(xlim=(-max_count*1.2, max_count*1.2),
               ylim=(-max_count*1.2, max_count*1.2),
               aspect=1.0)
    else:
        ax.set(xlim=(-max_count*1.05, max_count*1.05),
               ylim=(-max_count*1.05, max_count*1.05),
               aspect=1.0)
Пример #23
0
def plot(network,
         margin=0.05,
         ax=None,
         basemap=True,
         bus_colors='b',
         line_colors='g',
         bus_sizes=10,
         line_widths=2,
         title="",
         line_cmap=None,
         bus_cmap=None,
         boundaries=None,
         geometry=False,
         branch_components=['Line', 'Link'],
         jitter=None):
    """
    Plot the network buses and lines using matplotlib and Basemap.

    Parameters
    ----------
    margin : float
        Margin at the sides as proportion of distance between max/min x,y
    ax : matplotlib ax, defaults to plt.gca()
        Axis to which to plot the network
    basemap : bool, default True
        Switch to use Basemap
    bus_colors : dict/pandas.Series
        Colors for the buses, defaults to "b"
    bus_sizes : dict/pandas.Series
        Sizes of bus points, defaults to 10
    line_colors : dict/pandas.Series
        Colors for the lines, defaults to "g" for Lines and "cyan" for
        Links. Colors for branches other than Lines can be
        specified using a pandas Series with a MultiIndex.
    line_widths : dict/pandas.Series
        Widths of lines, defaults to 2. Widths for branches other
        than Lines can be specified using a pandas Series with a
        MultiIndex.
    title : string
        Graph title
    line_cmap : plt.cm.ColorMap/str|dict
        If line_colors are floats, this color map will assign the colors.
        Use a dict to specify colormaps for more than one branch type.
    bus_cmap : plt.cm.ColorMap/str
        If bus_colors are floats, this color map will assign the colors
    boundaries : list of four floats
        Boundaries of the plot in format [x1,x2,y1,y2]
    branch_components : list of str
        Branch components to be plotted, defaults to Line and Link.
    jitter : None|float
        Amount of random noise to add to bus positions to distinguish
        overlapping buses

    Returns
    -------
    bus_collection, branch_collection1, ... : tuple of Collections
        Collections for buses and branches.
    """

    defaults_for_branches = {
        'Link': dict(color="cyan", width=2),
        'Line': dict(color="b", width=2),
        'Transformer': dict(color='green', width=2)
    }

    if not plt_present:
        logger.error("Matplotlib is not present, so plotting won't work.")
        return

    if ax is None:
        ax = plt.gca()

    def compute_bbox_with_margins(margin, x, y):
        #set margins
        pos = np.asarray((x, y))
        minxy, maxxy = pos.min(axis=1), pos.max(axis=1)
        xy1 = minxy - margin * (maxxy - minxy)
        xy2 = maxxy + margin * (maxxy - minxy)
        return tuple(xy1), tuple(xy2)

    x = network.buses["x"]
    y = network.buses["y"]

    if jitter is not None:
        x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
        y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))

    if basemap and basemap_present:
        if boundaries is None:
            (x1, y1), (x2, y2) = compute_bbox_with_margins(margin, x, y)
        else:
            x1, x2, y1, y2 = boundaries
        bmap = Basemap(resolution='l',
                       epsg=network.srid,
                       llcrnrlat=y1,
                       urcrnrlat=y2,
                       llcrnrlon=x1,
                       urcrnrlon=x2,
                       ax=ax)
        bmap.drawcountries()
        bmap.drawcoastlines()

        x, y = bmap(x.values, y.values)
        x = pd.Series(x, network.buses.index)
        y = pd.Series(y, network.buses.index)

    if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index,
                                                       pd.MultiIndex):
        # We are drawing pies to show all the different shares
        assert len(bus_sizes.index.levels[0].difference(network.buses.index)) == 0, \
            "The first MultiIndex level of bus_sizes must contain buses"
        assert isinstance(bus_colors, dict) and set(bus_colors).issuperset(bus_sizes.index.levels[1]), \
            "bus_colors must be a dictionary defining a color for each element " \
            "in the second MultiIndex level of bus_sizes"

        bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)

        patches = []
        for b_i in bus_sizes.index.levels[0]:
            s = bus_sizes.loc[b_i]
            radius = s.sum()**0.5
            ratios = s / s.sum()

            start = 0.25
            for i, ratio in ratios.iteritems():
                patches.append(
                    Wedge((x.at[b_i], y.at[b_i]),
                          radius,
                          360 * start,
                          360 * (start + ratio),
                          facecolor=bus_colors[i]))
                start += ratio
        bus_collection = PatchCollection(patches, match_original=True)
        ax.add_collection(bus_collection)
    else:
        c = pd.Series(bus_colors, index=network.buses.index)
        if c.dtype == np.dtype('O'):
            c.fillna("b", inplace=True)
            c = list(c.values)
        s = pd.Series(bus_sizes, index=network.buses.index,
                      dtype="float").fillna(10)
        bus_collection = ax.scatter(x, y, c=c, s=s, cmap=bus_cmap)

    def as_branch_series(ser):
        if isinstance(ser, dict) and set(ser).issubset(branch_components):
            return pd.Series(ser)
        elif isinstance(ser, pd.Series):
            if isinstance(ser.index, pd.MultiIndex):
                return ser
            index = ser.index
            ser = ser.values
        else:
            index = network.lines.index
        return pd.Series(ser,
                         index=pd.MultiIndex(levels=(["Line"], index),
                                             labels=(np.zeros(len(index)),
                                                     np.arange(len(index)))))

    line_colors = as_branch_series(line_colors)
    line_widths = as_branch_series(line_widths)
    if not isinstance(line_cmap, dict):
        line_cmap = {'Line': line_cmap}

    branch_collections = []
    for c in network.iterate_components(branch_components):
        l_defaults = defaults_for_branches[c.name]
        l_widths = line_widths.get(c.name, l_defaults['width'])
        l_nums = None
        l_colors = line_colors.get(c.name, l_defaults['color'])

        if isinstance(l_colors, pd.Series):
            if issubclass(l_colors.dtype.type, np.number):
                l_nums = l_colors
                l_colors = None
            else:
                l_colors.fillna(l_defaults['color'], inplace=True)

        if not geometry:
            segments = (np.asarray(
                ((c.df.bus0.map(x), c.df.bus0.map(y)),
                 (c.df.bus1.map(x), c.df.bus1.map(y)))).transpose(2, 0, 1))
        else:
            from shapely.wkt import loads
            from shapely.geometry import LineString
            linestrings = c.df.geometry.map(loads)
            assert all(isinstance(ls, LineString) for ls in linestrings), \
                "The WKT-encoded geometry in the 'geometry' column must be composed of LineStrings"
            segments = np.asarray(list(linestrings.map(np.asarray)))
            if basemap and basemap_present:
                segments = np.transpose(
                    bmap(*np.transpose(segments, (2, 0, 1))), (1, 2, 0))

        l_collection = LineCollection(segments,
                                      linewidths=l_widths,
                                      antialiaseds=(1, ),
                                      colors=l_colors,
                                      transOffset=ax.transData)

        if l_nums is not None:
            l_collection.set_array(np.asarray(l_nums))
            l_collection.set_cmap(line_cmap.get(c.name, None))
            l_collection.autoscale()

        ax.add_collection(l_collection)
        l_collection.set_zorder(1)

        branch_collections.append(l_collection)

    bus_collection.set_zorder(2)

    ax.update_datalim(compute_bbox_with_margins(margin, x, y))
    ax.autoscale_view()

    ax.set_title(title)

    return (bus_collection, ) + tuple(branch_collections)
Пример #24
0
def construct_eigenvalue_DOS_plot_haldane(xy,
                                          fig,
                                          dos_ax,
                                          eig_ax,
                                          eigval,
                                          eigvect,
                                          en,
                                          Ni,
                                          Nk,
                                          marker_num=0,
                                          color_scheme='default',
                                          sub_lattice=-1,
                                          PVx=[],
                                          PVy=[],
                                          black_t0lines=False,
                                          mark_t0=True,
                                          title='auto',
                                          normalization=1.,
                                          alpha=0.6,
                                          lw=1,
                                          zorder=10):
    """puts together lattice and DOS plots and draws normal mode magitudes as circles on top
    
    Parameters
    ----------
    xy: array 2nx3
        Equilibrium position of the gyroscopes
    fig :
        figure with lattice and DOS drawn
    dos_ax:
        axis for the DOS plot
    eig_ax
        axis for the eigenvalue plot
    eigval : array of dimension 2nx1
        Eigenvalues of matrix for system
    eigvect : array of dimension 2nx2n
        Eigenvectors of matrix for system.
        Eigvect is stored as NModes x NP array, like: mode0_psi0, mode0_psi1, ... / mode1_psi0, ...
    en: int
        Number of the eigenvalue you are plotting
    
    Returns
    ----------
    fig :
        completed figure for normal mode
    
    [scat_fg, p, f_mark] :
        things to be cleared before next normal mode is drawn
    """

    ppu = leplt.get_points_per_unit()
    s = leplt.absolute_sizer()

    # re_eigvals = sum(abs(real(eigval)))
    # im_eigvals = sum(abs(imag(eigval)))

    ev = eigval[en]
    ev1 = ev

    # Show where current eigenvalue is in DOS plot (red line ticking current eigval)
    if dos_ax is not None:
        (f_mark, ) = dos_ax.plot([ev.real, ev.real], P.ylim(), '-r')
        plt.sca(dos_ax)

    NP = len(xy)

    im1 = np.imag(ev)
    re1 = np.real(ev)
    P.sca(eig_ax)

    if title == 'auto':
        eig_ax.set_title('Mode %d: $\omega=( %0.6f + %0.6f i)$' %
                         (en, re1, im1))
    elif title is not None and title not in ['', 'none']:
        eig_ax.set_title(title)

    # Preallocate ellipsoid plot vars
    shap = eigvect.shape
    angles_arr = np.zeros(NP)
    major_Ax = np.zeros(NP)

    patch = []
    polygon = []
    colors = np.zeros(NP + 2)
    # x_mag = np.zeros(NP)
    # y_mag = np.zeros(NP)

    x0s = np.zeros(NP)
    y0s = np.zeros(NP)

    mag1 = eigvect[en]

    # Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like:
    # x0, y0, x1, y1, ... xNP, yNP.
    mag1x = np.array([mag1[i] for i in range(NP)])
    mag1y = np.array([mag1[i] for i in range(NP)])

    # Pick a series of times to draw out the ellipsoid
    time_arr = np.arange(81) * 2 * pi / (abs(ev1) * 80)
    exp1 = np.exp(1j * ev1 * time_arr)

    # Normalization for the ellipsoids
    lim_mag1 = max(
        np.array([
            np.sqrt(2 * abs(exp1 * mag1x[i])**2) for i in range(len(mag1x))
        ]).flatten())
    mag1x /= lim_mag1
    mag1y /= lim_mag1
    mag1x *= normalization
    mag1y *= normalization

    cw = []
    ccw = []
    lines_1 = []
    for i in range(NP):
        unit = mag1x[i]
        x_disps = 0.5 * (exp1 * unit).real
        y_disps = 0.5 * (exp1 * unit).imag

        x_vals = xy[i, 0] + x_disps
        y_vals = xy[i, 1] + y_disps

        # x_mag[i] = max(x_vals-xy[i,0]).real
        # y_mag[i] = max(y_vals-xy[i,1]).real

        poly_points = array([x_vals, y_vals]).T
        polygon = Polygon(poly_points, True)

        # x0 is the marker_num^th element of x_disps
        x0 = x_disps[marker_num]
        y0 = y_disps[marker_num]

        x0s[i] = x_vals[marker_num]
        y0s[i] = y_vals[marker_num]

        # These are the black lines protruding from pivot point to current position
        lines_1.append([[xy[i, 0], x_vals[marker_num]],
                        [xy[i, 1], y_vals[marker_num]]])

        mag = sqrt(x0**2 + y0**2)
        if mag > 0:
            anglez = np.arccos(x0 / mag)
        else:
            anglez = 0

        if y0 < 0:
            anglez = 2 * np.pi - anglez

        # testangle = arctan2(y0,x0)
        # print ' x0 - x_disps[0] =', x0-x_disps[marker_num]

        angles_arr[i] = anglez

        # print 'polygon = ', poly_points
        patch.append(polygon)

        # Do Fast Fourier Transform (FFT)
        # ff = abs(fft.fft(x_disps + 1j*y_disps))**2
        # ff_freq = fft.fftfreq(len(x_vals), 1)
        # mm_f = ff_freq[ff == max(ff)][0]

        if color_scheme == 'default':
            colors[i] = anglez
        else:
            if sub_lattice[i] == 0:
                colors[i] = 0
            else:
                colors[i] = pi
                # if mm_f > 0:
                #   colors[i] = 0
                # else:
                #   colors[i] = pi

    colors[NP] = 0
    colors[NP + 1] = 2 * pi

    plt.yticks([])
    plt.xticks([])
    # this is the part that puts a dot a t=0 point
    if mark_t0:
        scat_fg = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='k')
        scat_fg2 = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='r')
        scat_fg = [scat_fg, scat_fg2]
    else:
        scat_fg = []

    NP = len(xy)
    try:
        NN = shape(Ni)[1]
    except IndexError:
        NN = 0

    z = np.zeros(NP)

    Rnorm = np.array([x0s, y0s, z]).T

    # Bond Stretches
    inc = 0
    stretches = zeros(3 * len(xy))
    if PVx == [] and PVy == []:
        '''There are no periodic boundaries supplied'''
        for i in range(len(xy)):
            if NN > 0:
                for j, k in zip(Ni[i], Nk[i]):
                    if i < j and abs(k) > 0:
                        n1 = float(linalg.norm(Rnorm[i] - Rnorm[j]))
                        n2 = linalg.norm(xy[i] - xy[j])
                        stretches[inc] = (n1 - n2)
                        inc += 1
    else:
        '''There are periodic boundaries supplied'''
        # get boundary particle indices
        KLabs = np.zeros_like(Nk, dtype='int')
        KLabs[Nk > 0] = 1
        boundary = extract_boundary_from_NL(xy, Ni, KLabs)
        for i in range(len(xy)):
            if NN > 0:
                for j, k in zip(Ni[i], Nk[i]):
                    # if i in boundary and j in boundary:
                    #     col = np.where( Ni[i] == j)[0][0]
                    #     print 'col = ', col
                    #     n1 = float( np.linalg.norm(Rnorm[i]-Rnorm[j]) )
                    #     n2 = np.linalg.norm(R[i] - R[j] )
                    #     stretches[inc] = (n1 - n2)
                    #     inc += 1
                    #
                    #     #test[inc] = [R[i], np.array([R[j,0]+PVx[i,col], R[j,1] + PVy[i,col], 0])]
                    # else:
                    if i < j and abs(k) > 0:
                        n1 = float(np.linalg.norm(Rnorm[i] - Rnorm[j]))
                        n2 = np.linalg.norm(xy[i] - xy[j])
                        stretches[inc] = (n1 - n2)
                        inc += 1

    stretch = np.array(stretches[0:inc])

    # For particles with neighbors, get list of bonds to draw by stretches
    test = list(np.zeros([inc, 1]))
    inc = 0
    xy = np.array([x0s, y0s, z]).T
    if PVx == [] and PVy == []:
        '''There are no periodic boundaries supplied'''
        for i in range(len(xy)):
            if NN > 0:
                for j, k in zip(Ni[i], Nk[i]):
                    if i < j and abs(k) > 0:
                        test[inc] = [xy[(i, j), 0], xy[(i, j), 1]]
                        inc += 1
    else:
        '''There are periodic boundaries supplied'''
        # get boundary particle indices
        KLabs = np.zeros_like(Nk, dtype='int')
        KLabs[Nk > 0] = 1
        boundary = extract_boundary_from_NL(xy, Ni, KLabs)
        for i in range(len(xy)):
            if NN > 0:
                for j, k in zip(Ni[i], Nk[i]):
                    # if i in boundary and j in boundary:
                    #     col = np.where( Ni[i] == j)[0][0]
                    #     print 'i,j = (', i,j, ')'
                    #     print 'PVx[i,col] = ', PVx[i,col]
                    #     print 'PVy[i,col] = ', PVy[i,col]
                    #     test[inc] = [xy[i], np.array([xy[j,0]+PVx[i,col], xy[j,1] + PVy[i,col], 0])]
                    #     #plt.plot([ xy[i,0], xy[j,0]+PVx[i,col]], [xy[i,1], xy[j,1] + PVy[i,col] ], 'k-')
                    #     #plt.plot(xy[:,0], xy[:,1],'b.')
                    #     #plt.show()
                    #     print 'test = ', test
                    #     inc += 1
                    # else:
                    if i < j and abs(k) > 0:
                        test[inc] = [xy[(i, j), 0], xy[(i, j), 1]]
                        inc += 1

    lines = [zip(x, y) for x, y in test]

    # angles[-1] = 0
    # angles[-2] = 2*pi
    lines_st = LineCollection(lines,
                              array=stretch,
                              cmap='seismic',
                              linewidth=8)
    lines_st.set_clim([-1. * 0.25, 1 * 0.25])
    lines_st.set_zorder(2)

    if black_t0lines:
        lines_12 = [zip(x, y) for x, y in lines_1]
        lines_12_st = LineCollection(lines_12, linewidth=0.8)
        lines_12_st.set_color('k')
        eig_ax.add_collection(lines_12_st)
    else:
        lines_12_st = []

    p = PatchCollection(patch, cmap='hsv', lw=lw, alpha=alpha, zorder=zorder)
    p.set_array(array(colors))
    p.set_clim([0, 2 * pi])
    p.set_zorder(1)

    # eig_ax.add_collection(lines_st)
    eig_ax.add_collection(p)

    eig_ax.set_aspect('equal')
    s = leplt.absolute_sizer()

    # erased ev/(2*pi) here npm 2016
    cw_ccw = [cw, ccw, ev]
    # print cw_ccw[1]

    return fig, [scat_fg, p, f_mark, lines_12_st], cw_ccw
Пример #25
0
def plot_hextensor(tensor,
                   image_range=(0, None),
                   channel_range=(0, None),
                   cmap=mymap,
                   norm=None,
                   linewidth=1,
                   edgecolors='k',
                   zorder=None,
                   figname="figure",
                   mask=[]):
    r"""Plot the hexagonal representation of a 4D tensor according to the 
        addressing sheme used by HexagDLy.

        Args:
        tensor:         torch tensor or numpy array containing the hexagonal data points
        image_range:    tuple of ints, range defining the images to be plotted
        channel_range:  tuple of ints, range defining the channels to be plotted
        cmap:           colourmap
        figname:        str, name of figure
        mask:           list of ints that depict the pixels to skip in plots 
                        counting top to bottom left to right from the top left pixel  

    """
    try:
        tensor = tensor.data.numpy()
    except:
        pass
    if norm is None:
        norm = Normalize(tensor.min(), tensor.max())
    if isinstance(edgecolors, np.ndarray) or isinstance(edgecolors, list):
        edgecolors[(edgecolors == '') | (edgecolors == '1')] = 'k'

    inshape = np.shape(tensor[image_range[0]:image_range[1],
                              channel_range[0]:channel_range[1]])
    inexamples = inshape[0]
    inchannels = inshape[1]
    if inexamples != 1 and inchannels != 1:
        print(
            "Choose one image and n channels or one channel an n images to display!"
        )
        sys.exit()
    nimages = max(inexamples, inchannels)
    hexagons = [[] for i in range(nimages)]
    intensities = [[] for i in range(nimages)]
    fig = plt.figure(figname, (5, 5))
    fig.clear()
    nrows = int(np.ceil(np.sqrt(nimages)))
    gs = gridspec.GridSpec(nrows, nrows)
    gs.update(wspace=0, hspace=0)

    for i in range(nimages):
        if inexamples >= inchannels:
            a = i
            b = 0
        else:
            a = 0
            b = i
        npixel = 0
        for x in range(
                np.shape(tensor[image_range[0] + a, channel_range[0] + b])[1]):
            for y in range(
                    np.shape(tensor[image_range[0] + a,
                                    channel_range[0] + b])[0]):
                if npixel not in mask:
                    intensity = tensor[image_range[0] + a,
                                       channel_range[0] + b, y, x]
                    hexagon = RegularPolygon(
                        (x * np.sqrt(3) / 2, -(y + np.mod(x, 2) * 0.5)),
                        6,
                        0.577349,
                        orientation=np.pi / 6,
                    )
                    intensities[i].append(intensity)
                    hexagons[i].append(hexagon)
                npixel += 1
        ax = fig.add_subplot(gs[i])
        ax.set_xlim([
            -1,
            np.shape(tensor[image_range[0] + a, channel_range[0] + b])[1]
        ])
        #embed()
        ax.set_ylim([
            -1.15 *
            np.shape(tensor[image_range[0] + a, channel_range[0] + b])[0] - 1,
            1,
        ])
        ax.set_axis_off()

        p = PatchCollection(np.array(hexagons[i]),
                            cmap=cmap,
                            norm=norm,
                            alpha=0.9,
                            edgecolors="k",
                            linewidth=linewidth)

        p.set_array(np.array(np.array(intensities[i])))
        p.set_linewidth(linewidth)
        p.set_cmap(cmap)
        #p.set_norm(norm)
        p.set_edgecolors(edgecolors)
        p.set_zorder(zorder)
        ax.add_collection(p)
        ax.set_aspect("equal")
        plt.subplots_adjust(top=0.95, bottom=0.05)
    plt.tight_layout()
Пример #26
0
def plot_tsd(ax, df, segs, args, lane=None, ghost_edges=None, ghost_bounds=None):
    """Plot the time-space diagram.

    Take the pre-processed segments and other meta-data, then plot all the line segments.

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        figure axes that will be plotted on
    df : pd.DataFrame
        data used for axes bounds and speed coloring
    segs : list of list of lists
        line segments to be plotted, where each segment is a list of two [x,y] pairs
    args : dict
        parsed arguments
    lane : int, optional
        lane number to be shown in plot title
    ghost_edges : list or set of str
        ghost edge names to be greyed out, default None
    ghost_bounds : tuple
        lower and upper bounds of domain, excluding ghost edges, default None

    Returns
    -------
    None
    """
    norm = plt.Normalize(args.min_speed, args.max_speed)

    xmin, xmax = df['time_step'].min(), df['time_step'].max()
    xbuffer = (xmax - xmin) * 0.025  # 2.5% of range
    ymin, ymax = df['distance'].min(), df['distance'].max()
    ybuffer = (ymax - ymin) * 0.025  # 2.5% of range

    ax.set_xlim(xmin - xbuffer, xmax + xbuffer)
    ax.set_ylim(ymin - ybuffer, ymax + ybuffer)

    lc = LineCollection(segs, cmap=my_cmap, norm=norm)
    lc.set_array(df['speed'].values)
    lc.set_linewidth(1)
    ax.add_collection(lc)
    ax.autoscale()

    rects = []
    if ghost_edges:
        y_domain_min = df[~df['edge_id'].isin(ghost_edges)]['distance'].min()
        y_domain_max = df[~df['edge_id'].isin(ghost_edges)]['distance'].max()
        rects.append(Rectangle((xmin, y_domain_min), args.start - xmin, y_domain_max - y_domain_min))
        rects.append(Rectangle((xmin, ymin), xmax - xmin, y_domain_min - ymin))
        rects.append(Rectangle((xmin, y_domain_max), xmax - xmin, ymax - y_domain_max))
    elif ghost_bounds:
        rects.append(Rectangle((xmin, ghost_bounds[0]), args.start - xmin, ghost_bounds[1] - ghost_bounds[0]))
        rects.append(Rectangle((xmin, ymin), xmax - xmin, ghost_bounds[0] - ymin))
        rects.append(Rectangle((xmin, ghost_bounds[1]), xmax - xmin, ymax - ghost_bounds[1]))
    else:
        rects.append(Rectangle((xmin, ymin), args.start - xmin, ymax - ymin))

    if rects:
        pc = PatchCollection(rects, facecolor='grey', alpha=0.5, edgecolor=None)
        pc.set_zorder(20)
        ax.add_collection(pc)

    if lane:
        ax.set_title('Time-Space Diagram: Lane {}'.format(lane), fontsize=25)
    else:
        ax.set_title('Time-Space Diagram', fontsize=25)
    ax.set_ylabel('Position (m)', fontsize=20)
    ax.set_xlabel('Time (s)', fontsize=20)
    plt.xticks(fontsize=18)
    plt.yticks(fontsize=18)

    cbar = plt.colorbar(lc, ax=ax, norm=norm)
    cbar.set_label('Velocity (m/s)', fontsize=20)
    cbar.ax.tick_params(labelsize=18)
def construct_haldane_eigvect_DOS_plot(xy,
                                       fig,
                                       DOS_ax,
                                       eig_ax,
                                       eigval,
                                       eigvect,
                                       en,
                                       NL,
                                       KL,
                                       marker_num=0,
                                       color_scheme='default',
                                       sub_lattice=-1,
                                       normalization=None):
    """puts together lattice and DOS plots and draws normal mode ellipsoids on top

    Parameters
    ----------
    xy: array 2N x 3
        Equilibrium position of the gyroscopes
    fig :
        figure with lattice and DOS drawn
    DOS_ax:
        axis for the DOS plot
    eig_ax
        axis for the eigenvalue plot
    eigval : array of dimension 2nx1
        Eigenvalues of matrix for system
    eigvect : array of dimension 2nx2n
        Eigenvectors of matrix for system.
        Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like:
        x0, y0, x1, y1, ... xNP, yNP.
    en: int
        Number of the eigenvalue you are plotting

    Returns
    ----------
    fig :
        completed figure for normal mode

    [scat_fg, p, f_mark] :
        things to be cleared before next normal mode is drawn
        """
    s = leplt.absolute_sizer()

    plt.sca(DOS_ax)

    ev = eigval[en]
    ev1 = ev

    # Show where current eigenvalue is in DOS plot
    (f_mark, ) = plt.plot([ev, ev], plt.ylim(), '-r')

    NP = len(xy)

    im1 = np.imag(ev)
    re1 = np.real(ev)
    plt.sca(eig_ax)
    plt.title('Mode %d; $\Omega=( %0.6f + %0.6f i)$' % (en, re1, im1))

    # Preallocate ellipsoid plot vars
    angles_arr = np.zeros(NP)

    patch = []
    colors = np.zeros(NP + 2)

    x0s = np.zeros(NP)
    y0s = np.zeros(NP)

    mag1 = eigvect[en]
    if normalization is None:
        mag1 /= np.max(np.abs(mag1))
    else:
        mag1 *= normalization * float(len(xy))

    # Pick a series of times to draw out the ellipsoid
    time_arr = np.arange(81.0) * 2. * np.pi / float(abs(ev1) * 80)
    exp1 = np.exp(1j * ev1 * time_arr)
    cw = []
    ccw = []
    lines_1 = []
    for i in range(NP):
        x_disps = 0.5 * (exp1 * mag1[i]).real
        y_disps = 0.5 * (exp1 * mag1[i]).imag
        x_vals = xy[i, 0] + x_disps
        y_vals = xy[i, 1] + y_disps

        poly_points = np.array([x_vals, y_vals]).T
        polygon = Polygon(poly_points, True)

        # x0 is the marker_num^th element of x_disps
        x0 = x_disps[marker_num]
        y0 = y_disps[marker_num]

        x0s[i] = x_vals[marker_num]
        y0s[i] = y_vals[marker_num]

        # These are the black lines protruding from pivot point to current position
        lines_1.append([[xy[i, 0], x_vals[marker_num]],
                        [xy[i, 1], y_vals[marker_num]]])

        mag = np.sqrt(x0**2 + y0**2)
        if mag > 0:
            anglez = np.arccos(x0 / mag)
        else:
            anglez = 0

        if y0 < 0:
            anglez = 2 * np.pi - anglez

        angles_arr[i] = anglez
        patch.append(polygon)

        if color_scheme == 'default':
            colors[i] = anglez
        else:
            if sub_lattice[i] == 0:
                colors[i] = 0
            else:
                colors[i] = np.pi
            ccw.append(i)

    colors[NP] = 0
    colors[NP + 1] = 2 * np.pi

    plt.yticks([])
    plt.xticks([])
    # this is the part that puts a dot a t=0 point
    scat_fg = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='DodgerBlue')
    scat_fg2 = eig_ax.scatter(x0s[ccw], y0s[ccw], s=s(.02), c='Red', zorder=3)

    NP = len(xy)
    try:
        NN = np.shape(NL)[1]
    except IndexError:
        NN = 0

    z = np.zeros(NP)

    Rnorm = np.array([x0s, y0s, z]).T

    # Bond Stretches
    inc = 0
    stretches = np.zeros(4 * len(xy))
    for i in range(len(xy)):
        if NN > 0:
            for j, k in zip(NL[i], KL[i]):
                if i < j and abs(k) > 0:
                    n1 = float(linalg.norm(Rnorm[i] - Rnorm[j]))
                    n2 = linalg.norm(xy[i] - xy[j])
                    stretches[inc] = (n1 - n2)
                    inc += 1

    # For particles with neighbors, get list of bonds to draw by stretches
    test = list(np.zeros([inc, 1]))
    inc = 0
    xy = np.array([x0s, y0s, z]).T
    for i in range(len(xy)):
        if NN > 0:
            for j, k in zip(NL[i], KL[i]):
                if i < j and abs(k) > 0:
                    test[inc] = [xy[(i, j), 0], xy[(i, j), 1]]
                    inc += 1

    stretch = np.array(stretches[0:inc])

    # lines connect sites (bonds), while lines_12 draw the black lines from the pinning to location sites
    lines = [zip(x, y) for x, y in test]
    lines_12 = [zip(x, y) for x, y in lines_1]

    lines_st = LineCollection(lines,
                              array=stretch,
                              cmap='seismic',
                              linewidth=8)
    lines_st.set_clim([-1. * 0.25, 1 * 0.25])
    lines_st.set_zorder(2)

    lines_12_st = LineCollection(lines_12, linewidth=0.8)
    lines_12_st.set_color('k')

    p = PatchCollection(patch, cmap='hsv', alpha=0.6)

    p.set_array(np.array(colors))
    p.set_clim([0, 2 * np.pi])
    p.set_zorder(1)

    # eig_ax.add_collection(lines_st)
    eig_ax.add_collection(lines_12_st)
    eig_ax.add_collection(p)
    eig_ax.set_aspect('equal')

    # erased ev/(2*pi) here npm 2016
    cw_ccw = [cw, ccw, ev]
    # print cw_ccw[1]

    return fig, [scat_fg, scat_fg2, p, f_mark, lines_12_st], cw_ccw
Пример #28
0
def create_bus_collection(net, buses=None, size=5, marker="o", patch_type="circle", colors=None,
                          z=None, cmap=None, norm=None, infofunc=None, picker=False,
                          bus_geodata=None, cbar_title="Bus Voltage [pu]", **kwargs):
    """
    Creates a matplotlib patch collection of pandapower buses.

    Input:
        **net** (pandapowerNet) - The pandapower network

    OPTIONAL:
        **buses** (list, None) - The buses for which the collections are created.
            If None, all buses in the network are considered.

        **size** (int, 5) - patch size

        **marker** (str, "o") - patch marker

        **patch_type** (str, "circle") - patch type, can be

                - "circle" for a circle
                - "rect" for a rectangle
                - "poly<n>" for a polygon with n edges

        **infofunc** (function, None) - infofunction for the patch element

        **colors** (list, None) - list of colors for every element

        **z** (array, None) - array of bus voltage magnitudes for colormap. Used in case of given
            cmap. If None net.res_bus.vm_pu is used.

        **cmap** (ListedColormap, None) - colormap for the patch colors

        **norm** (matplotlib norm object, None) - matplotlib norm object

        **picker** (bool, False) - picker argument passed to the patch collection

        **bus_geodata** (DataFrame, None) - coordinates to use for plotting
            If None, net["bus_geodata"] is used

        **cbar_title** (str, "Bus Voltage [pu]") - colormap bar title in case of given cmap

        **kwargs - key word arguments are passed to the patch function

    OUTPUT:
        **pc** - patch collection
    """
    buses = net.bus.index.tolist() if buses is None else list(buses)
    if len(buses) == 0:
        return None
    if bus_geodata is None:
        bus_geodata = net["bus_geodata"]

    coords = zip(bus_geodata.loc[buses, "x"].values, bus_geodata.loc[buses, "y"].values)

    infos = []

    # RegularPolygon has no param width/height, everything else might use defaults
    if not patch_type.startswith("poly"):
        if 'height' not in kwargs and 'width' not in kwargs:
            kwargs['height'] = kwargs['width'] = 2 * size

        if patch_type == "rect":
            kwargs['height'] *= 2
            kwargs['width'] *= 2

    def figmaker(x, y, i):
        if colors is not None:
            kwargs["color"] = colors[i]
        if patch_type == 'ellipse' or patch_type == 'circle':  # circles are just ellipses
            angle = kwargs['angle'] if 'angle' in kwargs else 0
            fig = Ellipse((x, y), angle=angle, **kwargs)
        elif patch_type == "rect":
            fig = Rectangle([x - kwargs['width'] / 2, y - kwargs['height'] / 2], **kwargs)
        elif patch_type.startswith("poly"):
            edges = int(patch_type[4:])
            fig = RegularPolygon([x, y], numVertices=edges, radius=size, **kwargs)
        else:
            logger.error("Wrong patchtype. Please choose a correct patch type.")
        if infofunc:
            infos.append(infofunc(buses[i]))
        return fig

    patches = [figmaker(x, y, i)
               for i, (x, y) in enumerate(coords)
               if x != np.nan]
    pc = PatchCollection(patches, match_original=True, picker=picker)
    pc.bus_indices = np.array(buses)
    if cmap is not None:
        pc.set_cmap(cmap)
        pc.set_norm(norm)
        if z is None and net is not None:
            z = net.res_bus.vm_pu.loc[buses]
        else:
            logger.warning("z is None and no net is provided")
        pc.set_array(np.array(z))
        pc.has_colormap = True
        pc.cbar_title = cbar_title

    pc.patch_type = patch_type
    pc.size = size
    if 'orientation' in kwargs:
        pc.orientation = kwargs['orientation']
    if "zorder" in kwargs:
        pc.set_zorder(kwargs["zorder"])
    pc.info = infos

    return pc
Пример #29
0
def plot_radial_histogram(angles,
                          counts,
                          all_angles=None,
                          include_labels=False,
                          offset=180.0,
                          direction=-1,
                          closed=False,
                          color=STIM_COLOR):
    if all_angles is None:
        if len(angles) < 2:
            all_angles = np.linspace(0, 315, 8)
        else:
            all_angles = angles

    dth = (all_angles[1] - all_angles[0]) * 0.5

    if len(counts) == 0:
        max_count = 1
    else:
        max_count = max(counts)

    wedges = []
    for count, angle in zip(counts, angles):
        angle = angle*direction + offset
        wedge = mpatches.Wedge((0,0), count, angle-dth, angle+dth)
        wedges.append(wedge)

    wedge_coll = PatchCollection(wedges)
    wedge_coll.set_facecolor(color)
    wedge_coll.set_zorder(2)

    angles_rad = (all_angles*direction + offset)*np.pi/180.0

    if closed:
        border_coll = cplots.radial_circles([max_count])
    else:
        border_coll = cplots.radial_arcs([max_count], 
                                         min(angles_rad), 
                                         max(angles_rad))
    border_coll.set_facecolor((0,0,0,0))
    border_coll.set_zorder(1)

    line_coll = cplots.angle_lines(angles_rad, 0, max_count)
    line_coll.set_edgecolor((0,0,0,1))
    line_coll.set_linestyle(":")
    line_coll.set_zorder(1)

    ax = plt.gca()
    ax.add_collection(wedge_coll)
    ax.add_collection(border_coll)
    ax.add_collection(line_coll)

    if include_labels:
        cplots.add_angle_labels(ax, angles_rad, all_angles.astype(int), max_count, (0,0,0,1), offset=max_count*0.1)
        ax.set(xlim=(-max_count*1.2, max_count*1.2),
               ylim=(-max_count*1.2, max_count*1.2),
               aspect=1.0)
    else:
        ax.set(xlim=(-max_count*1.05, max_count*1.05),
               ylim=(-max_count*1.05, max_count*1.05),
               aspect=1.0)
Пример #30
0
def plot(network,
         margin=0.05,
         ax=None,
         geomap=True,
         projection=None,
         bus_colors='b',
         line_colors='g',
         bus_sizes=10,
         line_widths=2,
         title="",
         line_cmap=None,
         bus_cmap=None,
         boundaries=None,
         geometry=False,
         branch_components=['Line', 'Link'],
         jitter=None,
         basemap=None):
    """
    Plot the network buses and lines using matplotlib and Basemap.

    Parameters
    ----------
    margin : float
        Margin at the sides as proportion of distance between max/min x,y
    ax : matplotlib ax, defaults to plt.gca()
        Axis to which to plot the network
    geomap: bool/str, default True
        Switch to use Basemap or Cartopy (depends on what is installed).
        If string is passed, it will be used as a resolution argument.
        For Basemap users 'c' (crude), 'l' (low), 'i' (intermediate),
        'h' (high), 'f' (full) are valid resolutions options.
        For Cartopy users '10m', '50m', '110m' are valid resolutions options.
    projection: cartopy.crs.Projection, defaults to None
        Define the projection of your geomap, only valid if cartopy is
        installed. If None (default) is passed the projection for cartropy
        is set to cartopy.crs.PlateCarree
    bus_colors : dict/pandas.Series
        Colors for the buses, defaults to "b"
    bus_sizes : dict/pandas.Series
        Sizes of bus points, defaults to 10
    line_colors : dict/pandas.Series
        Colors for the lines, defaults to "g" for Lines and "cyan" for
        Links. Colors for branches other than Lines can be
        specified using a pandas Series with a MultiIndex.
    line_widths : dict/pandas.Series
        Widths of lines, defaults to 2. Widths for branches other
        than Lines can be specified using a pandas Series with a
        MultiIndex.
    title : string
        Graph title
    line_cmap : plt.cm.ColorMap/str|dict
        If line_colors are floats, this color map will assign the colors.
        Use a dict to specify colormaps for more than one branch type.
    bus_cmap : plt.cm.ColorMap/str
        If bus_colors are floats, this color map will assign the colors
    boundaries : list of four floats
        Boundaries of the plot in format [x1,x2,y1,y2]
    branch_components : list of str
        Branch components to be plotted, defaults to Line and Link.
    jitter : None|float
        Amount of random noise to add to bus positions to distinguish
        overlapping buses

    Returns
    -------
    bus_collection, branch_collection1, ... : tuple of Collections
        Collections for buses and branches.
    """

    defaults_for_branches = {
        'Link': dict(color="cyan", width=2),
        'Line': dict(color="b", width=2),
        'Transformer': dict(color='green', width=2)
    }

    if not plt_present:
        logger.error("Matplotlib is not present, so plotting won't work.")
        return

    if basemap is not None:
        logger.warning("argument `basemap` is deprecated, "
                       "use `geomap` instead.")
        geomap = basemap

    if cartopy_present and geomap:
        if projection is None:
            projection = get_projection_from_crs(network.srid)

        if ax is None:
            ax = plt.gca(projection=projection)
        else:
            assert isinstance(ax, cartopy.mpl.geoaxes.GeoAxesSubplot), (
                'The passed axis is not a GeoAxesSubplot. You can '
                'create one with: \nimport cartopy.crs as ccrs \n'
                'fig, ax = plt.subplots('
                'subplot_kw={"projection":ccrs.PlateCarree()})')
    elif ax is None:
        ax = plt.gca()

    x, y = network.buses["x"], network.buses["y"]

    if jitter is not None:
        x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
        y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))

    if geomap:
        transform = draw_map(network, x, y, ax, boundaries, margin, geomap)
    else:
        transform = ax.transData

    if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index,
                                                       pd.MultiIndex):
        # We are drawing pies to show all the different shares
        assert len(bus_sizes.index.levels[0].difference(network.buses.index)) == 0, \
            "The first MultiIndex level of bus_sizes must contain buses"
        assert (isinstance(bus_colors, dict) and
                set(bus_colors).issuperset(bus_sizes.index.levels[1])), \
            "bus_colors must be a dictionary defining a color for each element " \
            "in the second MultiIndex level of bus_sizes"

        bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)\
                        * projected_area_factor(ax, network.srid)**2

        patches = []
        for b_i in bus_sizes.index.levels[0]:
            s = bus_sizes.loc[b_i]
            radius = s.sum()**0.5
            if radius == 0.0:
                ratios = s
            else:
                ratios = s / s.sum()

            start = 0.25
            for i, ratio in ratios.iteritems():
                patches.append(
                    Wedge((x.at[b_i], y.at[b_i]),
                          radius,
                          360 * start,
                          360 * (start + ratio),
                          facecolor=bus_colors[i]))
                start += ratio
        bus_collection = PatchCollection(patches,
                                         match_original=True,
                                         transform=transform)
        ax.add_collection(bus_collection)
    else:
        c = pd.Series(bus_colors, index=network.buses.index)
        s = pd.Series(bus_sizes, index=network.buses.index,
                      dtype="float").fillna(10)
        bus_collection = ax.scatter(x,
                                    y,
                                    c=c,
                                    s=s,
                                    cmap=bus_cmap,
                                    edgecolor='face',
                                    transform=transform)

    def as_branch_series(ser):
        if isinstance(ser, dict) and set(ser).issubset(branch_components):
            return pd.Series(ser)
        elif isinstance(ser, pd.Series):
            if isinstance(ser.index, pd.MultiIndex):
                return ser
            index = ser.index
            ser = ser.values
        else:
            index = network.lines.index
        return pd.Series(ser,
                         index=pd.MultiIndex(levels=(["Line"], index),
                                             codes=(np.zeros(len(index)),
                                                    np.arange(len(index)))))

    line_colors = as_branch_series(line_colors)
    line_widths = as_branch_series(line_widths)
    if not isinstance(line_cmap, dict):
        line_cmap = {'Line': line_cmap}

    branch_collections = []
    for c in network.iterate_components(branch_components):
        l_defaults = defaults_for_branches[c.name]
        l_widths = line_widths.get(c.name, l_defaults['width'])
        l_nums = None
        l_colors = line_colors.get(c.name, l_defaults['color'])

        if isinstance(l_colors, pd.Series):
            if issubclass(l_colors.dtype.type, np.number):
                l_nums = l_colors
                l_colors = None
            else:
                l_colors.fillna(l_defaults['color'], inplace=True)

        if not geometry:
            segments = (np.asarray(
                ((c.df.bus0.map(x), c.df.bus0.map(y)),
                 (c.df.bus1.map(x), c.df.bus1.map(y)))).transpose(2, 0, 1))
        else:
            from shapely.wkt import loads
            from shapely.geometry import LineString
            linestrings = c.df.geometry.map(loads)
            assert all(isinstance(ls, LineString) for ls in linestrings), (
                "The WKT-encoded geometry in the 'geometry' column must be "
                "composed of LineStrings")
            segments = np.asarray(list(linestrings.map(np.asarray)))

        l_collection = LineCollection(segments,
                                      linewidths=l_widths,
                                      antialiaseds=(1, ),
                                      colors=l_colors,
                                      transOffset=ax.transData,
                                      transform=transform)

        if l_nums is not None:
            l_collection.set_array(np.asarray(l_nums))
            l_collection.set_cmap(line_cmap.get(c.name, None))
            l_collection.autoscale()

        ax.add_collection(l_collection)
        l_collection.set_zorder(1)

        branch_collections.append(l_collection)

    bus_collection.set_zorder(2)

    ax.update_datalim(compute_bbox_with_margins(margin, x, y))
    ax.autoscale_view()

    if geomap:
        if cartopy_present:
            ax.outline_patch.set_visible(False)
        ax.axis('off')

    ax.set_title(title)

    return (bus_collection, ) + tuple(branch_collections)
Пример #31
0
def draw_network(network,
                 nsize="total-degree",
                 ncolor="group",
                 nshape="o",
                 nborder_color="k",
                 nborder_width=0.5,
                 esize=1.,
                 ecolor="k",
                 ealpha=0.5,
                 max_nsize=5.,
                 max_esize=2.,
                 curved_edges=False,
                 threshold=0.5,
                 decimate=None,
                 spatial=True,
                 restrict_sources=None,
                 restrict_targets=None,
                 show_environment=True,
                 fast=False,
                 size=(600, 600),
                 xlims=None,
                 ylims=None,
                 dpi=75,
                 axis=None,
                 show=False,
                 **kwargs):
    '''
    Draw a given graph/network.

    Parameters
    ----------
    network : :class:`~nngt.Graph` or subclass
        The graph/network to plot.
    nsize : float, array of float or string, optional (default: "total-degree")
        Size of the nodes as a percentage of the canvas length. Otherwise, it
        can be a string that correlates the size to a node attribute among
        "in/out/total-degree", or "betweenness".
    ncolor : float, array of floats or string, optional (default: 0.5)
        Color of the nodes; if a float in [0, 1], position of the color in the
        current palette, otherwise a string that correlates the color to a node
        attribute among "in/out/total-degree", "betweenness" or "group".
    nshape : char or array of chars, optional (default: "o")
        Shape of the nodes (see `Matplotlib markers <http://matplotlib.org/api/
        markers_api.html?highlight=marker#module-matplotlib.markers>`_).
    nborder_color : char, float or array, optional (default: "k")
        Color of the node's border using predefined `Matplotlib colors
        <http://matplotlib.org/api/colors_api.html?highlight=color
        #module-matplotlib.colors>`_).
        or floats in [0, 1] defining the position in the palette.
    nborder_width : float or array of floats, optional (default: 0.5)
        Width of the border in percent of canvas size.
    esize : float, str, or array of floats, optional (default: 0.5)
        Width of the edges in percent of canvas length. Available string values
        are "betweenness" and "weight".
    ecolor : str, char, float or array, optional (default: "k")
        Edge color. If ecolor="group", edges color will depend on the source
        and target groups, i.e. only edges from and toward same groups will
        have the same color.
    max_esize : float, optional (default: 5.)
        If a custom property is entered as `esize`, this normalizes the edge
        width between 0. and `max_esize`.
    decimate : int, optional (default: keep all connections)
        Plot only one connection every `decimate`.
    spatial : bool, optional (default: True)
        If True, use the neurons' positions to draw them.
    restrict_sources : str or list, optional (default: all)
        Only draw edges starting from a restricted set of source nodes.
    restrict_targets : str or list, optional (default: all)
        Only draw edges ending on a restricted set of target nodes.
    show_environment : bool, optional (default: True)
        Plot the environment if the graph is spatial.
    fast : bool, optional (default: False)
        Use a faster algorithm to plot the edges. This method leads to less
        pretty plots and zooming on the graph will make the edges start or
        ending in places that will differ more or less strongly from the actual
        node positions.
    size : tuple of ints, optional (default: (600,600))
        (width, height) tuple for the canvas size (in px).
    dpi : int, optional (default: 75)
        Resolution (dot per inch).
    show : bool, optional (default: True)
        Display the plot immediately.
    '''
    import matplotlib.pyplot as plt
    size_inches = (size[0] / float(dpi), size[1] / float(dpi))
    if axis is None:
        fig = plt.figure(facecolor='white', figsize=size_inches, dpi=dpi)
        axis = fig.add_subplot(111, frameon=0, aspect=1)
    axis.set_axis_off()
    pos, layout = None, None
    # restrict sources and targets
    if nonstring_container(restrict_sources):
        if isinstance(restrict_sources[0], str):
            assert network.is_network(), \
                "`restrict_sources` canbe string only for Network."
            sources = []
            for name in restrict_sources:
                sources.extend(network.population[name].ids)
            restrict_sources = sources
    elif isinstance(restrict_sources, str):
        assert network.is_network(), \
            "`restrict_sources` canbe string only for Network."
        restrict_sources = network.population[restrict_sources].ids
    if nonstring_container(restrict_targets):
        if isinstance(restrict_targets[0], str):
            assert network.is_network(), \
                "`restrict_targets` canbe string only for Network."
            targets = []
            for name in restrict_targets:
                targets.extend(network.population[name].ids)
            restrict_targets = targets
    elif isinstance(restrict_targets, str):
        assert network.is_network(), \
            "`restrict_sources` canbe string only for Network."
        restrict_targets = network.population[restrict_targets].ids
    # get nodes and edges
    n = network.node_nb()
    adj_mat = network.adjacency_matrix(weights=None)
    if restrict_sources is not None:
        adj_mat = adj_mat[restrict_sources, :]
    if restrict_targets is not None:
        adj_mat = adj_mat[:, restrict_targets]
    e = adj_mat.nnz
    # compute properties
    decimate = 1 if decimate is None else decimate
    if isinstance(nsize, str):
        if e:
            nsize = _node_size(network, nsize)
            nsize *= max_nsize
        else:
            nsize = np.ones(n, dtype=float)
    elif isinstance(nsize, float):
        nsize = np.repeat(nsize, n)
    nsize *= 0.01 * size[0]
    if isinstance(esize, str) and e:
        esize = _edge_size(network, esize)
        esize *= max_esize
        esize[esize < threshold] = 0.
    #~ elif isinstance(esize, float):
    #~ esize = np.repeat(esize, e)
    esize *= 0.005 * size[0]  # border on each side (so 0.5 %)
    ncolor = _node_color(network, ncolor)
    c = ncolor
    if not nonstring_container(nborder_color):
        nborder_color = np.repeat(nborder_color, n)
    # check edge color
    group_based = False
    if isinstance(ecolor, float):
        ecolor = np.repeat(ecolor, e)
    elif ecolor == "groups" and network.is_network():
        group_based = True
        c = np.linspace(0, 1, len(network.population))
        ecolor = {}
        for i, src in enumerate(network.population):
            idx1 = network.population[src].ids[0]
            for j, tgt in enumerate(network.population):
                idx2 = network.population[tgt].ids[0]
                if src == tgt:
                    ecolor[(src, tgt)] = ncolor[idx1]
                else:
                    ecolor[(src, tgt)] = \
                        np.abs(0.8*ncolor[idx1] - 0.2*ncolor[idx2])
    # draw
    pos = np.zeros((n, 2))
    if spatial and network.is_spatial():
        if show_environment:
            nngt.geometry.plot.plot_shape(network.shape, axis=axis, show=False)
        pos = network.get_positions()
    else:
        pos[:, 0] = size[0] * (np.random.uniform(size=n) - 0.5)
        pos[:, 1] = size[1] * (np.random.uniform(size=n) - 0.5)
    # make nodes
    nodes = []
    if network.is_network():
        for group in network.population.values():
            idx = group.ids
            if nonstring_container(ncolor):
                c = palette(ncolor[idx[0]])
            for i in idx:
                nodes.append(
                    Circle(pos[i], 0.5 * nsize[i], fc=c, ec=nborder_color[i]))
    else:
        if not isinstance(c, str):
            c = palette(ncolor)
        for i in range(n):
            nodes.append(
                Circle(pos[i], 0.5 * nsize[i], fc=c, ec=nborder_color[i]))
    nodes = PatchCollection(nodes, match_original=True)
    nodes.set_zorder(2)
    axis.add_collection(nodes)
    _set_ax_lim(axis, pos[:, 0], pos[:, 1], xlims, ylims)
    # use quiver to draw the edges
    if e:
        adj_mat = network.adjacency_matrix(weights=None)
        avg_size = np.average(nsize)
        arr_style = ArrowStyle.Simple(head_length=0.15 * avg_size,
                                      head_width=0.1 * avg_size,
                                      tail_width=0.05 * avg_size)
        arrows = []
        if group_based:
            for src_name, src_group in network.population.items():
                for tgt_name, tgt_group in network.population.items():
                    s_ids = src_group.ids
                    if restrict_sources is not None:
                        s_ids = list(set(restrict_sources).intersection(s_ids))
                    t_ids = tgt_group.ids
                    if restrict_targets is not None:
                        t_ids = list(set(restrict_targets).intersection(t_ids))
                    if t_ids and s_ids:
                        s_min, s_max = np.min(s_ids), np.max(s_ids) + 1
                        t_min, t_max = np.min(t_ids), np.max(t_ids) + 1
                        edges = np.array(adj_mat[s_min:s_max,
                                                 t_min:t_max].nonzero(),
                                         dtype=int)
                        edges[0, :] += s_min
                        edges[1, :] += t_min
                        if nonstring_container(esize):
                            keep = (esize > 0)
                            edges = edges[:, keep]
                            esize = esize[keep]
                        if decimate > 1:
                            edges = edges[:, ::decimate]
                            if nonstring_container(esize):
                                esize = esize[::decimate]
                        # plot
                        ec = palette(ecolor[(src_name, tgt_name)])
                        if fast:
                            dl = 0.5 * np.max(nsize)
                            arrow_x = pos[edges[1], 0] - pos[edges[0], 0]
                            arrow_x -= np.sign(arrow_x) * dl
                            arrow_y = pos[edges[1], 1] - pos[edges[0], 1]
                            arrow_x -= np.sign(arrow_y) * dl
                            axis.quiver(pos[edges[0], 0],
                                        pos[edges[0], 1],
                                        arrow_x,
                                        arrow_y,
                                        scale_units='xy',
                                        angles='xy',
                                        scale=1,
                                        alpha=0.5,
                                        width=1.5e-3,
                                        linewidths=0.5 * esize,
                                        edgecolors=ec,
                                        zorder=1)
                        else:
                            for s, t in zip(edges[0], edges[1]):
                                xs, ys = pos[s, 0], pos[s, 1]
                                xt, yt = pos[t, 0], pos[t, 1]
                                dl = 0.5 * nsize[t]
                                dx = xt - xs
                                dx -= np.sign(dx) * dl
                                dy = yt - ys
                                dy -= np.sign(dy) * dl

                                if curved_edges:
                                    arrow = FancyArrowPatch(
                                        posA=(xs, ys),
                                        posB=(xt, yt),
                                        arrowstyle=arr_style,
                                        connectionstyle='arc3,rad=0.1',
                                        alpha=ealpha,
                                        fc=ec,
                                        lw=0.5)
                                    axis.add_patch(arrow)
                                else:
                                    arrows.append(
                                        FancyArrow(xs,
                                                   ys,
                                                   dx,
                                                   dy,
                                                   width=0.3 * avg_size,
                                                   head_length=0.7 * avg_size,
                                                   head_width=0.7 * avg_size,
                                                   length_includes_head=True,
                                                   alpha=ealpha,
                                                   fc=ec,
                                                   lw=0.5))
        else:
            edges = np.array(adj_mat.nonzero(), dtype=int)
            s_min, s_max, t_min, t_max = 0, n, 0, n
            if restrict_sources is not None:
                s_min = np.min(restrict_sources)
                s_max = np.min(np.max(restrict_sources) + 1, n)
            if restrict_targets is not None:
                t_min = np.min(restrict_targets)
                t_max = np.min(np.max(restrict_targets) + 1, n)
            edges = np.array(adj_mat[s_min:s_max, t_min:t_max].nonzero(),
                             dtype=int)
            edges[0, :] += s_min
            edges[1, :] += t_min
            # keep only large edges
            if nonstring_container(esize):
                keep = (esize > 0)
                edges = edges[:, keep]
                if nonstring_container(ecolor):
                    ecolor = ecolor[keep]
                esize = esize[keep]
            if decimate > 1:
                edges = edges[:, ::decimate]
                if nonstring_container(esize):
                    esize = esize[::decimate]
                if nonstring_container(ecolor):
                    ecolor = ecolor[::decimate]
            if isinstance(ecolor, str):
                ecolor = [ecolor for i in range(0, e, decimate)]

            if fast:
                dl = 0.5 * np.max(nsize)
                arrow_x = pos[edges[1], 0] - pos[edges[0], 0]
                arrow_x -= np.sign(arrow_x) * dl
                arrow_y = pos[edges[1], 1] - pos[edges[0], 1]
                arrow_x -= np.sign(arrow_y) * dl
                axis.quiver(pos[edges[0], 0],
                            pos[edges[0], 1],
                            arrow_x,
                            arrow_y,
                            scale_units='xy',
                            angles='xy',
                            scale=1,
                            alpha=0.5,
                            width=1.5e-3,
                            linewidths=0.5 * esize,
                            edgecolors=ecolor,
                            zorder=1)
            else:
                for i, (s, t) in enumerate(zip(edges[0], edges[1])):
                    xs, ys = pos[s, 0], pos[s, 1]
                    xt, yt = pos[t, 0], pos[t, 1]

                    if curved_edges:
                        arrow = FancyArrowPatch(posA=(xs, ys),
                                                posB=(xt, yt),
                                                arrowstyle=arr_style,
                                                connectionstyle='arc3,rad=0.1',
                                                alpha=ealpha,
                                                fc=ecolor[i],
                                                lw=0.5)
                        axis.add_patch(arrow)
                    else:
                        dl = 0.5 * nsize[t]
                        dx = xt - xs
                        dx -= np.sign(dx) * dl
                        dy = yt - ys
                        dy -= np.sign(dy) * dl
                        arrows.append(
                            FancyArrow(xs,
                                       ys,
                                       dx,
                                       dy,
                                       width=0.3 * avg_size,
                                       head_length=0.7 * avg_size,
                                       head_width=0.7 * avg_size,
                                       length_includes_head=True,
                                       alpha=ealpha,
                                       fc=ecolor[i],
                                       lw=0.5))

        if not fast:
            arrows = PatchCollection(arrows, match_original=True)
            arrows.set_zorder(1)
            axis.add_collection(arrows)

    if kwargs.get('tight', True):
        plt.tight_layout()
        plt.subplots_adjust(hspace=0.,
                            wspace=0.,
                            left=0.,
                            right=1.,
                            top=1.,
                            bottom=0.)
    if show:
        plt.show()
Пример #32
0
def create_bus_symbol_collection(coords,
                                 buses=None,
                                 size=5,
                                 marker="o",
                                 patch_type="circle",
                                 colors=None,
                                 z=None,
                                 cmap=None,
                                 norm=None,
                                 infofunc=None,
                                 picker=False,
                                 net=None,
                                 **kwargs):
    infos = []

    if 'height' in kwargs and 'width' in kwargs:
        height, width = kwargs['height'], kwargs['width']
    else:
        height, width = size, size

    def figmaker(x, y, i):
        if patch_type == "circle":
            if colors:
                fig = Circle((x, y), size, color=colors[i], **kwargs)
            else:
                fig = Circle((x, y), size, **kwargs)
        elif patch_type == 'ellipse':
            angle = kwargs['angle'] if 'angle' in kwargs else 0
            if colors:
                fig = Ellipse((x, y),
                              width=width,
                              height=height,
                              color=colors[i],
                              **kwargs)
            else:
                fig = Ellipse((x, y),
                              width=width,
                              height=height,
                              angle=angle,
                              **kwargs)
        elif patch_type == "rect":
            if colors:
                fig = Rectangle([x - width, y - height],
                                2 * width,
                                2 * height,
                                color=colors[i],
                                **kwargs)
            else:
                fig = Rectangle([x - width, y - height], 2 * width, 2 * height,
                                **kwargs)
        elif patch_type.startswith("poly"):
            edges = int(patch_type[4:])
            if colors:
                fig = RegularPolygon([x, y],
                                     numVertices=edges,
                                     radius=size,
                                     color=colors[i],
                                     **kwargs)
            else:
                fig = RegularPolygon([x, y],
                                     numVertices=edges,
                                     radius=size,
                                     **kwargs)
        else:
            logger.error(
                "Wrong patchtype. Please choose a correct patch type.")
        if infofunc:
            infos.append(infofunc(buses[i]))
        return fig

    patches = [
        figmaker(x, y, i) for i, (x, y) in enumerate(coords) if x != np.nan
    ]
    pc = PatchCollection(patches, match_original=True, picker=picker)
    pc.bus_indices = np.array(buses)
    if cmap:
        pc.set_cmap(cmap)
        pc.set_norm(norm)
        if z is None and net:
            z = net.res_bus.vm_pu.loc[buses]
        else:
            logger.warning("z is None and no net is provided")
        pc.set_array(np.array(z))
        pc.has_colormap = True
        pc.cbar_title = "Bus Voltage [pu]"

    pc.patch_type = patch_type
    pc.size = size
    if 'orientation' in kwargs:
        pc.orientation = kwargs['orientation']
    if "zorder" in kwargs:
        pc.set_zorder(kwargs["zorder"])
    pc.info = infos
    return pc
Пример #33
0
def plot(n,
         margin=0.05,
         ax=None,
         geomap=True,
         projection=None,
         bus_colors='b',
         bus_alpha=1,
         line_colors={
             'Line': 'g',
             'Link': 'cyan'
         },
         bus_sizes=1e-2,
         line_widths={
             'Line': 2,
             'Link': 2
         },
         flow=None,
         layouter=None,
         title="",
         line_cmap=None,
         bus_cmap=None,
         boundaries=None,
         geometry=False,
         branch_components=['Line', 'Link'],
         jitter=None,
         color_geomap=None):
    """
    Plot the network buses and lines using matplotlib and cartopy.

    Parameters
    ----------
    margin : float
        Margin at the sides as proportion of distance between max/min x,y
    ax : matplotlib ax, defaults to plt.gca()
        Axis to which to plot the network
    geomap: bool/str, default True
        Switch to use Cartopy and draw geographical features.
        If string is passed, it will be used as a resolution argument,
        valid options are '10m', '50m' and '110m'.
    projection: cartopy.crs.Projection, defaults to None
        Define the projection of your geomap, only valid if cartopy is
        installed. If None (default) is passed the projection for cartopy
        is set to cartopy.crs.PlateCarree
    bus_colors : dict/pandas.Series
        Colors for the buses, defaults to "b". If bus_sizes is a pandas.Series
        with a Multiindex, bus_colors defaults to the n.carriers['color']
        column.
    bus_sizes : dict/pandas.Series
        Sizes of bus points, defaults to 1e-2. If a multiindexed Series is passed,
        the function will draw pies for each bus (first index level) with
        segments of different color (second index level). Such a Series is ob-
        tained by e.g. n.generators.groupby(['bus', 'carrier']).p_nom.sum()
    bus_alpha : float
        Adds alpha channel to buses, defaults to 1.
    line_colors : dict/pandas.Series
        Colors for the lines, defaults to "g" for Lines and "cyan" for
        Links. Colors for branches other than Lines can be
        specified using a pandas Series with a MultiIndex.
    line_widths : dict/pandas.Series
        Widths of lines, defaults to 2. Widths for branches other
        than Lines can be specified using a pandas Series with a
        MultiIndex.
    flow : snapshot/pandas.Series/function/string
        Flow to be displayed in the plot, defaults to None. If an element of
        n.snapshots is given, the flow at this timestamp will be
        displayed. If an aggregation function is given, is will be applied
        to the total network flow via pandas.DataFrame.agg (accepts also
        function names). Otherwise flows can be specified by passing a pandas
        Series with MultiIndex including all necessary branch components.
        Use the line_widths argument to additionally adjust the size of the
        flow arrows.
    layouter : networkx.drawing.layout function, default None
        Layouting function from `networkx <https://networkx.github.io/>`_ which
        overrules coordinates given in ``n.buses[['x','y']]``. See
        `list <https://networkx.github.io/documentation/stable/reference/drawing.html#module-networkx.drawing.layout>`_
        of available options.
    title : string
        Graph title
    line_cmap : plt.cm.ColorMap/str|dict
        If line_colors are floats, this color map will assign the colors.
        Use a dict to specify colormaps for more than one branch type.
    bus_cmap : plt.cm.ColorMap/str
        If bus_colors are floats, this color map will assign the colors
    boundaries : list of four floats
        Boundaries of the plot in format [x1,x2,y1,y2]
    branch_components : list of str
        Branch components to be plotted, defaults to Line and Link.
    jitter : None|float
        Amount of random noise to add to bus positions to distinguish
        overlapping buses
    color_geomap : dict or bool
        Specify colors to paint land and sea areas in.
        If True, it defaults to `{'ocean': 'lightblue', 'land': 'whitesmoke'}`.
        If no dictionary is provided, colors are white.

    Returns
    -------
    bus_collection, branch_collection1, ... : tuple of Collections
        Collections for buses and branches.
    """
    defaults_for_branches = pd.Series({
        'Link':
        dict(color="cyan", width=2),
        'Line':
        dict(color="b", width=2),
        'Transformer':
        dict(color='green', width=2)
    }).rename_axis('component')

    x, y = _get_coordinates(n, layouter=layouter)

    if geomap:
        if not cartopy_present:
            logger.warning(
                "Cartopy needs to be installed to use `geomap=True`.")
            geomap = False

        if projection is None:
            projection = get_projection_from_crs(n.srid)

        if ax is None:
            ax = plt.gca(projection=projection)
        else:
            assert isinstance(ax, cartopy.mpl.geoaxes.GeoAxesSubplot), (
                'The passed axis is not a GeoAxesSubplot. You can '
                'create one with: \nimport cartopy.crs as ccrs \n'
                'fig, ax = plt.subplots('
                'subplot_kw={"projection":ccrs.PlateCarree()})')
        transform = draw_map_cartopy(n, x, y, ax, boundaries, margin, geomap,
                                     color_geomap)
        x, y, z = ax.projection.transform_points(transform, x.values,
                                                 y.values).T
        x, y = pd.Series(x, n.buses.index), pd.Series(y, n.buses.index)
    elif ax is None:
        ax = plt.gca()

    if jitter is not None:
        x = x + np.random.uniform(low=-jitter, high=jitter, size=len(x))
        y = y + np.random.uniform(low=-jitter, high=jitter, size=len(y))

    if isinstance(bus_sizes, pd.Series) and isinstance(bus_sizes.index,
                                                       pd.MultiIndex):
        # We are drawing pies to show all the different shares
        assert len(bus_sizes.index.levels[0].difference(n.buses.index)) == 0, \
            "The first MultiIndex level of bus_sizes must contain buses"
        if isinstance(bus_colors, dict):
            bus_colors = pd.Series(bus_colors)
        # case bus_colors isn't a series or dict: look in n.carriers for existent colors
        if not isinstance(bus_colors, pd.Series):
            bus_colors = n.carriers.color.dropna()
        assert bus_sizes.index.levels[1].isin(bus_colors.index).all(), (
            "Colors not defined for all elements in the second MultiIndex "
            "level of bus_sizes, please make sure that all the elements are "
            "included in bus_colors or in n.carriers.color")

        bus_sizes = bus_sizes.sort_index(level=0, sort_remaining=False)
        if geomap:
            bus_sizes *= projected_area_factor(ax, n.srid)**2

        patches = []
        for b_i in bus_sizes.index.levels[0]:
            s = bus_sizes.loc[b_i]
            radius = s.sum()**0.5
            if radius == 0.0:
                ratios = s
            else:
                ratios = s / s.sum()

            start = 0.25
            for i, ratio in ratios.iteritems():
                patches.append(
                    Wedge((x.at[b_i], y.at[b_i]),
                          radius,
                          360 * start,
                          360 * (start + ratio),
                          facecolor=bus_colors[i],
                          alpha=bus_alpha))
                start += ratio
        bus_collection = PatchCollection(patches, match_original=True)
        ax.add_collection(bus_collection)
    else:
        c = pd.Series(bus_colors, index=n.buses.index)
        s = pd.Series(bus_sizes, index=n.buses.index, dtype="float")
        if geomap:
            s *= projected_area_factor(ax, n.srid)**2

        if bus_cmap is not None and c.dtype is np.dtype('float'):
            if isinstance(bus_cmap, str):
                bus_cmap = plt.cm.get_cmap(bus_cmap)
            norm = plt.Normalize(vmin=c.min(), vmax=c.max())
            c = c.apply(lambda cval: bus_cmap(norm(cval)))

        patches = []
        for b_i in s.index:
            radius = s.at[b_i]**0.5
            patches.append(
                Circle((x.at[b_i], y.at[b_i]),
                       radius,
                       facecolor=c.at[b_i],
                       alpha=bus_alpha))
        bus_collection = PatchCollection(patches, match_original=True)
        ax.add_collection(bus_collection)

    def as_branch_series(ser):
        # ensure that this function always return a multiindexed series
        if isinstance(ser, dict) and set(ser).issubset(branch_components):
            return pd.concat(
                {
                    c.name: pd.Series(s, index=c.df.index)
                    for c, s in zip(n.iterate_components(ser.keys()),
                                    ser.values())
                },
                names=['component', 'name'])
        elif isinstance(ser, pd.Series) and isinstance(ser.index,
                                                       pd.MultiIndex):
            return ser.rename_axis(index=['component', 'name'])
        else:
            ser = pd.Series(ser, n.lines.index)
            return pd.concat([ser],
                             axis=0,
                             keys=['Line'],
                             names=['component', 'name']).fillna(0)

    line_colors = as_branch_series(line_colors)
    line_widths = as_branch_series(line_widths)

    if not isinstance(line_cmap, dict):
        line_cmap = {'Line': line_cmap}

    branch_collections = []

    if flow is not None:
        flow = (_flow_ds_from_arg(
            flow, n, branch_components).pipe(as_branch_series).div(
                sum(
                    len(t.df)
                    for t in n.iterate_components(branch_components)) + 100))
        flow = flow.mul(line_widths[flow.index], fill_value=1)
        # update the line width, allows to set line widths separately from flows
        line_widths.update((5 * flow.abs()).pipe(np.sqrt))
        arrows = directed_flow(n,
                               flow,
                               x=x,
                               y=y,
                               ax=ax,
                               geomap=geomap,
                               branch_colors=line_colors,
                               branch_comps=branch_components,
                               cmap=line_cmap['Line'])
        branch_collections.append(arrows)

    for c in n.iterate_components(branch_components):
        l_defaults = defaults_for_branches[c.name]
        l_widths = line_widths.get(c.name, l_defaults['width'])
        l_nums = None
        l_colors = line_colors.get(c.name, l_defaults['color'])

        if isinstance(l_colors, pd.Series):
            if issubclass(l_colors.dtype.type, np.number):
                l_nums = l_colors
                l_colors = None
            else:
                l_colors.fillna(l_defaults['color'], inplace=True)

        if not geometry:
            segments = (np.asarray(
                ((c.df.bus0.map(x), c.df.bus0.map(y)),
                 (c.df.bus1.map(x), c.df.bus1.map(y)))).transpose(2, 0, 1))
        else:
            from shapely.wkt import loads
            from shapely.geometry import LineString
            linestrings = c.df.geometry[lambda ds: ds != ''].map(loads)
            assert all(isinstance(ls, LineString) for ls in linestrings), (
                "The WKT-encoded geometry in the 'geometry' column must be "
                "composed of LineStrings")
            segments = np.asarray(list(linestrings.map(np.asarray)))

        l_collection = LineCollection(segments,
                                      linewidths=l_widths,
                                      antialiaseds=(1, ),
                                      colors=l_colors,
                                      transOffset=ax.transData)

        if l_nums is not None:
            l_collection.set_array(np.asarray(l_nums))
            l_collection.set_cmap(line_cmap.get(c.name, None))
            l_collection.autoscale()

        ax.add_collection(l_collection)
        l_collection.set_zorder(3)

        branch_collections.append(l_collection)

    bus_collection.set_zorder(4)

    ax.update_datalim(compute_bbox_with_margins(margin, x, y))
    ax.autoscale_view()

    if geomap:
        ax.outline_patch.set_visible(False)
        ax.axis('off')
    else:
        ax.set_aspect('equal')

    ax.set_title(title)

    return (bus_collection, ) + tuple(branch_collections)