Esempio n. 1
0
def draw_edges(tree, steiner_pos, diams=None, fig=None, ax=None, lim=None,
               colorbar=True):
    '''Draw edges with given positions.'''
    if fig is None:
        fig, ax = new_fig()
    if lim is None:
        lim = diam_min, diam_max

    pos = merge_pos(tree, steiner_pos)
    nodes = tree.get_nodes()
    arcs = tree.get_arcs()
    x = np.array([pos[n][0] for n in nodes])
    y = np.array([pos[n][1] for n in nodes])

    segments = [(pos[u], pos[v]) for (u,v) in arcs]

    if diams is None:
        lines = LineCollection(segments, colors='k', zorder=1)
    else:
        diams = np.array([diams[a] for a in arcs])
        lw = 7*diams + 1
        lines = LineCollection(segments, linewidths=lw, zorder=1)
        # set colors
        lines.set_array(diams)
        lines.set_cmap(_diam_cmap)
        lines.set_clim(*lim)
        if colorbar:
            plt.colorbar(lines, orientation='horizontal')
    ax.add_collection(lines)
Esempio n. 2
0
    def __plot_all(self, spectrum):
        total = len(spectrum)
        count = 0.0
        for timeStamp in spectrum:
            if self.settings.fadeScans:
                alpha = (total - count) / total
            else:
                alpha = 1

            data = spectrum[timeStamp].items()
            peakF, peakL = self.extent.get_peak_fl()

            segments, levels = self.__create_segments(data)
            if segments is not None:
                lc = LineCollection(segments)
                lc.set_array(numpy.array(levels))
                lc.set_norm(self.__get_norm(self.settings.autoL, self.extent))
                lc.set_cmap(self.colourMap)
                lc.set_linewidth(self.lineWidth)
                lc.set_gid('plot')
                lc.set_alpha(alpha)
                self.axes.add_collection(lc)
                count += 1

        return peakF, peakL
Esempio n. 3
0
    def __plot_all(self):
        total = len(self.data)
        count = 0.0
        for timeStamp in self.data:
            if len(self.data[timeStamp]) < 2:
                self.parent.threadPlot = None
                return None, None

            if self.fade:
                alpha = (total - count) / total
            else:
                alpha = 1

            data = self.data[timeStamp].items()
            peakF, peakL = self.extent.get_peak_fl()

            segments, levels = self.__create_segments(data)
            lc = LineCollection(segments)
            lc.set_array(numpy.array(levels))
            lc.set_norm(self.__get_norm(self.autoL, self.extent))
            lc.set_cmap(self.colourMap)
            lc.set_linewidth(self.lineWidth)
            lc.set_gid('plot')
            lc.set_alpha(alpha)
            self.axes.add_collection(lc)
            count += 1

        return peakF, peakL
Esempio n. 4
0
def plot_linestring_collection(ax, geoms, values=None, color=None,
                               cmap=None, vmin=None, vmax=None, **kwargs):
    """
    Plots a collection of LineString and MultiLineString geometries to `ax`

    Parameters
    ----------

    ax : matplotlib.axes.Axes
        where shapes will be plotted

    geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be
            mixed)

    values : a sequence of `N` values, optional
        Values will be mapped to colors using vmin/vmax/cmap. They should
        have 1:1 correspondence with the geometries (not their components).

    color : single color or sequence of `N` colors
        Cannot be used together with `values`.

    Returns
    -------

    collection : matplotlib.collections.Collection that was plotted

    """
    from matplotlib.collections import LineCollection

    geoms, values = _flatten_multi_geoms(geoms, values)
    if None in values:
        values = None

    # LineCollection does not accept some kwargs.
    if 'markersize' in kwargs:
        del kwargs['markersize']

    # color=None gives black instead of default color cycle
    if color is not None:
        kwargs['color'] = color

    segments = [np.array(linestring)[:, :2] for linestring in geoms]
    collection = LineCollection(segments, **kwargs)

    if values is not None:
        collection.set_array(np.asarray(values))
        collection.set_cmap(cmap)
        collection.set_clim(vmin, vmax)

    ax.add_collection(collection, autolim=True)
    ax.autoscale_view()
    return collection
Esempio n. 5
0
    def __plot_single(self, points):
        data = points.items()
        peakF, peakL = max(data, key=lambda item: item[1])

        segments, levels = self.__create_segments(data)
        lc = LineCollection(segments)
        lc.set_array(numpy.array(levels))
        lc.set_norm(self.__get_norm(self.autoL, self.extent))
        lc.set_cmap(self.colourMap)
        lc.set_linewidth(self.lineWidth)
        lc.set_gid('plot')
        self.axes.add_collection(lc)

        return peakF, peakL
Esempio n. 6
0
    def plot(self, time, beam=None, 
        maxground=2000, maxalt=500, step=1,
        showrefract=False, nr_cmap='jet_r', nr_lim=[0.8, 1.], 
        raycolor='0.3', title=False, zorder=2, alpha=1, 
        fig=None, rect=111, ax=None, aax=None):
        """Plot ray paths
        
        Parameters
        ----------
        time : datetime.datetime
            time of rays
        beam: Optional[ ]
            beam number
        maxground : Optional[int]
            maximum ground range [km]
        maxalt : Optional[int]
            highest altitude limit [km]
        step : Optional[int]
            step between each plotted ray (in number of ray steps)
        showrefract : Optional[bool]
            show refractive index along ray paths (supersedes raycolor)
        nr_cmap : Optional[str]
            color map name for refractive index coloring
        nr_lim : Optional[list, float]
            refractive index plotting limits
        raycolor : Optional[float]
            color of ray paths
        title : Optional[bool]
            Show default title
        zorder : Optional[int]

        alpha : Optional[int]

        fig : Optional[pylab.figure]
            object (default to gcf)
        rect : Optional[int]
            subplot spcification
        ax : Optional[ ]
            Existing main axes
        aax : Optional[ ]
            Existing auxialary axes

        Returns
        -------
        ax : matplotlib.axes
            object containing formatting
        aax : matplotlib.axes
            object containing data
        cbax : matplotlib.axes
            object containing colorbar

        Example
        -------
            # Show ray paths with colored refractive index along path
            import datetime as dt
            from davitpy.models import raydarn
            sTime = dt.datetime(2012, 11, 18, 5)
            rto = raydarn.RtRun(sTime, rCode='bks', beam=12, title=True)
            rto.readRays() # read rays into memory
            ax, aax, cbax = rto.rays.plot(sTime, step=10, showrefract=True, nr_lim=[.85,1])
            ax.grid()

        written by Sebastien, 2013-04

        """
        import datetime as dt
        from davitpy.utils import plotUtils
        from matplotlib.collections import LineCollection
        import matplotlib.pyplot as plt
        import numpy as np
        from types import MethodType

        # Set up axes
        if not ax and not aax:
            ax, aax = plotUtils.curvedEarthAxes(fig=fig, rect=rect, 
                maxground=maxground, maxalt=maxalt)
        else:
            ax = ax
            aax = aax
            if hasattr(ax, 'time'):
                time = ax.time
            if hasattr(ax, 'beam'):
                beam = ax.beam

        # make sure that the required time and beam are present
        # Allow a 60 second difference between the requested time and the time
        # available.
        keys    = np.array(self.paths.keys())
        diffs   = np.abs(keys-time)
        if diffs.min() < dt.timedelta(minutes=1):
            time = keys[diffs.argmin()]

        assert (time in self.paths.keys()), logging.error('Unkown time %s' % time)
        if beam:
            assert (beam in self.paths[time].keys()), logging.error('Unkown beam %s' % beam)
        else:
            beam = self.paths[time].keys()[0]
        
        for ir, (el, rays) in enumerate( sorted(self.paths[time][beam].items()) ):
            if not ir % step:
                if not showrefract:
                    aax.plot(rays['th'], rays['r']*1e-3, c=raycolor, 
                        zorder=zorder, alpha=alpha)
                else:
                    points = np.array([rays['th'], rays['r']*1e-3]).T.reshape(-1, 1, 2)
                    segments = np.concatenate([points[:-1], points[1:]], axis=1)
                    lcol = LineCollection( segments, zorder=zorder, alpha=alpha)
                    _ = lcol.set_cmap( nr_cmap )
                    _ = lcol.set_norm( plt.Normalize(*nr_lim) )
                    _ = lcol.set_array( rays['nr'] )
                    _ = aax.add_collection( lcol )

        # Plot title with date ut time and local time
        if title:
            stitle = _getTitle(time, beam, self.header, self.name)
            ax.set_title( stitle )

        # Add a colorbar when plotting refractive index
        if showrefract:
            cbax = plotUtils.addColorbar(lcol, ax)
            _ = cbax.set_ylabel("refractive index")
        else: cbax = None

        # Declare a new method to show range markers
        # This method is only available after rays have been plotted
        # This ensures that the markers match the plotted rays
        def showRange(self, markers=None, 
            color='.8', s=2, zorder=3, 
            **kwargs):
            """Plot ray paths
            
            Parameters
            ----------
            markers : Optional[ ]
                range markers. Defaults to every 250 km
            color : Optional[float]

            s : Optional[int]

            zorder : Optional[int]

            **kwargs :

            Returns
            -------
            coll :
                a collection of range markers

            Notes
            -----
            Parameters other than markers are borrowed from matplotlib.pyplot.scatter

            Example
            -------
                # Add range markers to an existing ray plot
                ax, aax, cbax = rto.rays.plot(sTime, step=10)
                rto.rays.showRange()

            written by Sebastien, 2013-04

            """

            if not markers:
                markers = np.arange(0, 5000, 250)
            
            x, y = [], []
            for el, rays in self.paths[time][beam].items():
                for rm in markers:
                    inds = (rays['gran']*1e-3 >= rm)
                    if inds.any():
                        x.append( rays['th'][inds][0] )
                        y.append( rays['r'][inds][0]*1e-3 )
            coll = aax.scatter(x, y, 
                color=color, s=s, zorder=zorder, **kwargs)

            return coll
        # End of new method

        # Assign new method
        self.showRange = MethodType(showRange, self)

        ax.beam = beam
        return ax, aax, cbax
Esempio n. 7
0
def plot(n, margin=None, ax=None, geomap=True, projection=None,
         bus_colors='cadetblue', bus_alpha=1, bus_sizes=2e-2, bus_cmap=None,
         line_colors='rosybrown', link_colors='darkseagreen',
         transformer_colors='orange',
         line_widths=1.5, link_widths=1.5, transformer_widths=1.5,
         line_cmap=None, link_cmap=None, transformer_cmap=None,
         flow=None, branch_components=None, layouter=None, title="",
         boundaries=None, geometry=False, 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 "cadetblue". If bus_sizes is a
        pandas.Series with a Multiindex, bus_colors defaults to the
        n.carriers['color'] column.
    bus_alpha : float
        Adds alpha channel to buses, defaults to 1.
    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_cmap : plt.cm.ColorMap/str
        If bus_colors are floats, this color map will assign the colors
    line_colors : str/pandas.Series
        Colors for the lines, defaults to 'rosybrown'.
    link_colors : str/pandas.Series
        Colors for the links, defaults to 'darkseagreen'.
    transfomer_colors : str/pandas.Series
        Colors for the transfomer, defaults to 'orange'.
    line_widths : dict/pandas.Series
        Widths of lines, defaults to 1.5
    link_widths : dict/pandas.Series
        Widths of links, defaults to 1.5
    transformer_widths : dict/pandas.Series
        Widths of transformer, defaults to 1.5
    line_cmap : plt.cm.ColorMap/str|dict
        If line_colors are floats, this color map will assign the colors.
    link_cmap : plt.cm.ColorMap/str|dict
        If link_colors are floats, this color map will assign the colors.
    transformer_cmap : plt.cm.ColorMap/str|dict
        If transformer_colors are floats, this color map will assign the colors.
    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
    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.
    """
    x, y = _get_coordinates(n, layouter=layouter)
    if boundaries is None and margin:
        boundaries = sum(zip(*compute_bbox_with_margins(margin, x, y)), ())

    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, 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)
        if boundaries:
            ax.set_extent(boundaries, crs=transform)
    elif ax is None:
        ax = plt.gca()
    if not geomap and boundaries:
        ax.axis(boundaries)

    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(title)

    # Plot buses:

    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 = 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 = 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, zorder=5)
        ax.add_collection(bus_collection)

    # Plot branches:
    if isinstance(line_widths, pd.Series):
        if isinstance(line_widths.index, pd.MultiIndex):
            raise TypeError("Index of argument 'line_widths' is a Multiindex, "
                            "this is not support since pypsa v0.17. "
                            "Set differing widths with arguments 'line_widths', "
                            "'link_widths' and 'transformer_widths'.")
    if isinstance(line_colors, pd.Series):
        if isinstance(line_colors.index, pd.MultiIndex):
            raise TypeError("Index of argument 'line_colors' is a Multiindex, "
                            "this is not support since pypsa v0.17. "
                            "Set differing colors with arguments 'line_colors', "
                            "'link_colors' and 'transformer_colors'.")

    if branch_components is None:
        branch_components = n.branch_components

    branch_colors = {'Line': line_colors, 'Link': link_colors,
                     'Transformer': transformer_colors}
    branch_widths = {'Line': line_widths, 'Link': link_widths,
                     'Transformer': transformer_widths}
    branch_cmap = {'Line': line_cmap, 'Link': link_cmap,
                   'Transformer': transformer_cmap}

    branch_collections = []
    arrow_collections = []

    if flow is not None:
        rough_scale = sum(len(n.df(c)) for c in branch_components) + 100
        flow = _flow_ds_from_arg(flow, n, branch_components) / rough_scale

    for c in n.iterate_components(branch_components):
        b_widths = as_branch_series(branch_widths[c.name], 'width', c.name, n)
        b_colors = as_branch_series(branch_colors[c.name], 'color', c.name, n)
        b_nums = None
        b_cmap = branch_cmap[c.name]
        b_flow = flow.get(c.name, None) if flow is not None else None

        if issubclass(b_colors.dtype.type, np.number):
            b_nums = b_colors
            b_colors = None

        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)))

        if b_flow is not None:
            coords = pd.DataFrame({'x1': c.df.bus0.map(x), 'y1': c.df.bus0.map(y),
                                   'x2': c.df.bus1.map(x), 'y2': c.df.bus1.map(y)})
            b_flow = b_flow.mul(b_widths[b_flow.index], fill_value=0)
            # update the line width, allows to set line widths separately from flows
            b_widths.update((5 * b_flow.abs()).pipe(np.sqrt))
            area_factor = projected_area_factor(ax, n.srid)
            f_collection = directed_flow(coords, b_flow, b_colors, area_factor,
                                         b_cmap)
            if b_nums is not None:
                f_collection.set_array(np.asarray(b_nums))
                f_collection.set_cmap(b_cmap)
                f_collection.autoscale()
            arrow_collections.append(f_collection)
            ax.add_collection(f_collection)


        b_collection = LineCollection(segments, linewidths=b_widths,
                                      antialiaseds=(1,), colors=b_colors,
                                      transOffset=ax.transData)

        if b_nums is not None:
            b_collection.set_array(np.asarray(b_nums))
            b_collection.set_cmap(b_cmap)
            b_collection.autoscale()

        ax.add_collection(b_collection)
        b_collection.set_zorder(3)
        branch_collections.append(b_collection)

    if boundaries is None:
        ax.autoscale()

    return (bus_collection,) + tuple(branch_collections) + tuple(arrow_collections)
Esempio n. 8
0
for p in ps:
    gm = cmpy.machines.GoldenMean(bias=p)
    data = gm.cmech_quantities(['HX0L', 'hmu'], maxL)
    lines1.append(zip(range(maxL + 1), data[0]))
    data[1].mask[0] = False
    lines2.append(zip(range(maxL + 1), data[1, :]))

fig = plt.figure()
fig.set_figheight(5)
fig.set_figwidth(10)
fig.subplots_adjust(left=0.09, right=0.86)
ax1 = fig.add_subplot(1, 2, 1)

lc1 = LineCollection(lines1)
lc1.set_array(ps)
lc1.set_cmap(plt.cm.jet)

ax1.add_collection(lc1)

ax1.set_xlabel("$L$ [symbols]")
ax1.set_ylabel("$H[X_0^L]$ [bits]")
ax1.axis('auto')
ax1.set_xlim((0, maxL))

ax2 = fig.add_subplot(1, 2, 2)

lc2 = LineCollection(lines2)
lc2.set_array(ps)
lc2.set_cmap(plt.cm.jet)

ax2.add_collection(lc2)
Esempio n. 9
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,
                        **kwds):
    """Draw the edges of the graph G

    This draws only the edges of the graph G.

    pos is a dictionary keyed by vertex with a two-tuple
    of x-y positions as the value.
    See networkx.layout for functions that compute node positions.

    edgelist is an optional list of the edges in G to be drawn.
    If provided, only the edges in edgelist will be drawn. 

    edgecolor can be a list of matplotlib color letters such as 'k' or
    'b' that lists the color of each edge; the list must be ordered in
    the same way as the edge list. Alternatively, this list can contain
    numbers and those number are mapped to a color scale using the color
    map edge_cmap.  Finally, it can also be a list of (r,g,b) or (r,g,b,a)
    tuples, in which case these will be used directly to color the edges.  If
    the latter mode is used, you should not provide a value for alpha, as it
    would be applied globally to all lines.
    
    For directed graphs, "arrows" (actually just thicker stubs) are drawn
    at the head end.  Arrows can be turned off with keyword arrows=False.

    See draw_networkx for the list of other optional parameters.

    """
    try:
        import matplotlib
        import matplotlib.pylab as pylab
        import matplotlib.cbook as cb
        from matplotlib.colors import colorConverter, Colormap
        from matplotlib.collections import LineCollection
        import numpy
    except ImportError:
        raise ImportError, "Matplotlib required for draw()"
    except RuntimeError:
        pass  # unable to open display

    if ax is None:
        ax = pylab.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 = LineCollection(
        edge_pos,
        colors=edge_colors,
        linewidths=lw,
        antialiaseds=(1, ),
        linestyle=style,
        transOffset=ax.transData,
    )

    edge_collection.set_zorder(1)  # edges go behind nodes
    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)

    # need 0.87.7 or greater for edge colormaps
    mpl_version = matplotlib.__version__
    if mpl_version.endswith('.svn'):
        mpl_version = matplotlib.__version__[0:-4]
    elif mpl_version.endswith('svn'):
        mpl_version = matplotlib.__version__[0:-3]
    elif mpl_version.endswith('pre'):
        mpl_version = matplotlib.__version__[0:-3]
    if map(int, mpl_version.split('.')) >= [0, 87, 7]:
        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()
            pylab.axes(ax)
            pylab.sci(edge_collection)


#    else:
#        sys.stderr.write(\
#            """matplotlib version >= 0.87.7 required for colormapped edges.
#        (version %s detected)."""%matplotlib.__version__)
#        raise UserWarning(\
#            """matplotlib version >= 0.87.7 required for colormapped edges.
#        (version %s detected)."""%matplotlib.__version__)

    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
        a_pos = []
        p = 1.0 - 0.25  # make head segment 25 percent of edge length
        for src, dst in edge_pos:
            x1, y1 = src
            x2, y2 = dst
            dx = x2 - x1  # x offset
            dy = y2 - y1  # y offset
            d = numpy.sqrt(float(dx**2 + dy**2))  # length of edge
            if d == 0:  # source and target at same position
                continue
            if dx == 0:  # vertical edge
                xa = x2
                ya = dy * p + y1
            if dy == 0:  # horizontal edge
                ya = y2
                xa = dx * p + x1
            else:
                theta = numpy.arctan2(dy, dx)
                xa = p * d * numpy.cos(theta) + x1
                ya = p * d * numpy.sin(theta) + y1

            a_pos.append(((xa, ya), (x2, y2)))

        arrow_collection = LineCollection(
            a_pos,
            colors=arrow_colors,
            linewidths=[4 * ww for ww in lw],
            antialiaseds=(1, ),
            transOffset=ax.transData,
        )

        arrow_collection.set_zorder(1)  # edges go behind nodes
        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 edge_collection
Esempio n. 10
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.
       If not specified a spring layout positioning will be computed.
       See networkx.layout for functions that compute node positions.

    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

    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.pylab as pylab
        import matplotlib.cbook as cb
        from matplotlib.colors import colorConverter, Colormap
        from matplotlib.collections import LineCollection
        import numpy
    except ImportError:
        raise ImportError("Matplotlib required for draw()")
    except RuntimeError:
        print("Matplotlib unable to open display")
        raise

    if ax is None:
        ax = pylab.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 = 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()
        pylab.sci(edge_collection)

    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
        a_pos = []
        p = 1.0 - 0.25  # make head segment 25 percent of edge length
        for src, dst in edge_pos:
            x1, y1 = src
            x2, y2 = dst
            dx = x2 - x1  # x offset
            dy = y2 - y1  # y offset
            d = numpy.sqrt(float(dx**2 + dy**2))  # length of edge
            if d == 0:  # source and target at same position
                continue
            if dx == 0:  # vertical edge
                xa = x2
                ya = dy * p + y1
            if dy == 0:  # horizontal edge
                ya = y2
                xa = dx * p + x1
            else:
                theta = numpy.arctan2(dy, dx)
                xa = p * d * numpy.cos(theta) + x1
                ya = p * d * numpy.sin(theta) + y1

            a_pos.append(((xa, ya), (x2, y2)))

        arrow_collection = LineCollection(
            a_pos,
            colors=arrow_colors,
            linewidths=[4 * ww for ww in lw],
            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 edge_collection
Esempio n. 11
0
def create_line_collection(net,
                           lines=None,
                           line_geodata=None,
                           use_bus_geodata=False,
                           infofunc=None,
                           cmap=None,
                           norm=None,
                           picker=False,
                           z=None,
                           cbar_title="Line Loading [%]",
                           clim=None,
                           **kwargs):
    """
    Creates a matplotlib line collection of pandapower lines.

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

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

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

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

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

    """
    lines = net.line.index.tolist() if lines is None else list(lines)
    if len(lines) == 0:
        return None
    if line_geodata is None:
        line_geodata = net["line_geodata"]
    if len(lines) == 0:
        return None

    if use_bus_geodata:
        data = [
            ([(net.bus_geodata.at[a, "x"], net.bus_geodata.at[a, "y"]),
              (net.bus_geodata.at[b, "x"], net.bus_geodata.at[b, "y"])],
             infofunc(line) if infofunc else [])
            for line, (a,
                       b) in net.line.loc[lines,
                                          ["from_bus", "to_bus"]].iterrows()
            if a in net.bus_geodata.index.values
            and b in net.bus_geodata.index.values
        ]
    else:
        data = [(line_geodata.loc[line, "coords"],
                 infofunc(line) if infofunc else []) for line in lines
                if line in line_geodata.index.values]

    if len(data) == 0:
        return None

    data, info = list(zip(*data))

    # This would be done anyways by matplotlib - doing it explicitly makes it a) clear and
    # b) prevents unexpected behavior when observing colors being "none"
    lc = LineCollection(data, picker=picker, **kwargs)
    lc.line_indices = np.array(lines)
    if cmap:
        if z is None:
            z = net.res_line.loading_percent.loc[lines]
        lc.set_cmap(cmap)
        lc.set_norm(norm)
        if clim is not None:
            lc.set_clim(clim)
        lc.set_array(np.array(z))
        lc.has_colormap = True
        lc.cbar_title = cbar_title
    lc.info = info
    return lc
Esempio n. 12
0
def plot_dft_angles(snap_times,
                    angles,
                    amps,
                    sim_time,
                    logx=False,
                    logy=False,
                    plot_wheel=False,
                    lw=1.5,
                    cmap='hsv',
                    wheel_axis=None,
                    label_size=10):
    """
  Plot 2D-DFT amplitudes at a specific frequency as a function of time and angle
  """

    color_steps = 2056

    if logy: pl.gca().set_yscale('log')
    if logx: pl.gca().set_xscale('log')

    ax = pl.gca()
    custom_axes()
    ax.set_xlim(snap_times[1] if pl.gca().get_xscale() == 'log' else 0,
                sim_time)

    det_lines = LineCollection(
        [list(zip(snap_times[1:], amp[1:])) for amp in amps],
        linewidths=lw,
        linestyles='-')

    ax.set_ylim((np.amin(amps), np.amax(amps)))

    det_lines.set_array(angles)
    det_lines.set_clim(0, 180)
    det_lines.set_cmap(cmap)

    ax.add_collection(det_lines)
    pl.sci(det_lines)

    pos = ax.get_position()
    wheel_pos = [
        pos.x0 + pos.width / 20., pos.y0 + pos.height * 0.7, pos.width / 4.0,
        pos.height / 4.0
    ]

    if plot_wheel is True:
        if wheel_axis is None:
            wheel_axis = pl.gcf().add_axes(wheel_pos, projection='polar')
        else:
            wheel_axis.projection = 'polar'

        wheel_axis._direction = 2 * np.pi

        norm = pl.colors.Normalize(0.0, 180.)

        cb = pl.colorbar.ColorbarBase(wheel_axis,
                                      cmap=cm.get_cmap(cmap, color_steps),
                                      norm=norm,
                                      orientation='horizontal',
                                      ticks=[0, 30, 60, 90, 120, 150])
        cb.ax.tick_params(labelsize=label_size)
        cb.outline.set_visible(False)

    return ax
Esempio n. 13
0
def draw_edges(G,
               pos,
               ax,
               edgelist=None,
               width=1.0,
               width_adjuster=50,
               edge_color='k',
               style='solid',
               alpha=None,
               edge_cmap=None,
               edge_vmin=None,
               edge_vmax=None,
               traversal_weight=1.0,
               edge_delengthify=0.15,
               arrows=True,
               label=None,
               zorder=1,
               **kwds):
    """
    Code cleaned-up version of networkx.draw_networkx_edges

    New args:

    width_adjuster - the line width is generated from the weight if present, use this adjuster to thicken the lines (multiply)
    """
    if edgelist is None:
        edgelist = G.edges()

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

    # set edge positions
    edge_pos = [(pos[e[0]], pos[e[1]]) for e in edgelist]
    new_ep = []
    for e in edge_pos:
        x, y = e[0]
        dx, dy = e[1]

        # Get edge length
        elx = (dx - x) * edge_delengthify
        ely = (dy - y) * edge_delengthify

        x += elx
        y += ely
        dx -= elx
        dy -= ely

        new_ep.append(((x, y), (dx, dy)))
    edge_pos = numpy.asarray(new_ep)

    if numpy.iterable(width):
        lw = width

    else:
        #print [G.get_edge_data(n[0], n[1])['weight'] for n in edgelist]
        # see if I can find an edge attribute:
        if 'weight' in G.get_edge_data(edgelist[0][0],
                                       edgelist[0][1]):  # Test an edge
            lw = [
                0.5 +
                ((G.get_edge_data(n[0], n[1])['weight'] - traversal_weight) *
                 width_adjuster) for n in edgelist
            ]
        else:
            lw = (width, )
    if not is_string_like(edge_color) and numpy.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(
                [numpy.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 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 = LineCollection(edge_pos,
                                     colors=edge_colors,
                                     linewidths=lw,
                                     antialiaseds=(1, ),
                                     linestyle=style,
                                     transOffset=ax.transData,
                                     zorder=zorder)

    edge_collection.set_label(label)
    ax.add_collection(edge_collection)

    #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 None and edge_vmax is None:
            edge_collection.autoscale()

        else:
            edge_collection.set_clim(edge_vmin, edge_vmax)
    # 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()
    '''
    return (edge_collection)
Esempio n. 14
0
def create_line_collection(net, lines=None, line_geodata=None, bus_geodata=None,
                           use_bus_geodata=False, infofunc=None,
                           cmap=None, norm=None, picker=False, z=None,
                           cbar_title="Line Loading [%]", clim=None, **kwargs):
    """
    Creates a matplotlib line collection of pandapower lines.

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

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

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

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

        **use_bus_geodata** (bool, False) - Defines whether bus or line geodata are used.

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

        **cmap** - colormap for the patch colors

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

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

        **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.

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

        **clim** (tuple of floats, None) - setting the norm limits for image scaling

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

    OUTPUT:
        **lc** - line collection
    """
    if use_bus_geodata:
        linetab = net.line if lines is None else net.line.loc[lines]
    lines = net.line.index.tolist() if lines is None else list(lines)
    if len(lines) == 0:
        return None
    if line_geodata is None:
        line_geodata = net["line_geodata"]
    if bus_geodata is None:
        bus_geodata = net["bus_geodata"]
    if len(lines) == 0:
        return None

    lines_with_geo = []
    if use_bus_geodata:
        data = []
        buses_with_geodata = bus_geodata.index.values
        bg_dict = bus_geodata.to_dict() #transforming to dict to make lookup faster
        for line, fb, tb in zip(linetab.index, linetab.from_bus.values, linetab.to_bus.values):
            if fb in buses_with_geodata and tb in buses_with_geodata:
                lines_with_geo.append(line)
                data.append(([(bg_dict["x"][fb], bg_dict["y"][fb]),
                              (bg_dict["x"][tb], bg_dict["y"][tb])],
                             infofunc(line) if infofunc else[]))
        lines_without_geo = set(lines)-set(lines_with_geo)
        if lines_without_geo:
            logger.warning("Could not plot lines %s. Bus geodata is missing for those lines!"
                           % lines_without_geo)
    else:
        data = []
        for line in lines:
            if line in line_geodata.index.values:
                lines_with_geo.append(line)
                data.append((line_geodata.loc[line, "coords"], infofunc(line) if infofunc else []))

        lines_without_geo = set(lines)-set(lines_with_geo)
        if len(lines_without_geo) > 0:
            logger.warning("Could not plot lines %s. Line geodata is missing for those lines!"
                           % lines_without_geo)

    if len(data) == 0:
        return None

    data, info = list(zip(*data))

    # This would be done anyways by matplotlib - doing it explicitly makes it a) clear and
    # b) prevents unexpected behavior when observing colors being "none"
    lc = LineCollection(data, picker=picker, **kwargs)
    lc.line_indices = np.array(lines_with_geo)
    if cmap is not None:
        if z is None:
            z = net.res_line.loading_percent.loc[lines_with_geo]
        lc.set_cmap(cmap)
        lc.set_norm(norm)
        if clim is not None:
            lc.set_clim(clim)
        lc.set_array(np.array(z))
        lc.has_colormap = True
        lc.cbar_title = cbar_title
    lc.info = info

    return lc
Esempio n. 15
0
def create_line_collection(net,
                           lines=None,
                           line_geodata=None,
                           bus_geodata=None,
                           use_bus_geodata=False,
                           infofunc=None,
                           cmap=None,
                           norm=None,
                           picker=False,
                           z=None,
                           cbar_title="Line Loading [%]",
                           clim=None,
                           **kwargs):
    """
    Creates a matplotlib line collection of pandapower lines.

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

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

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

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

        **use_bus_geodata** (bool, False) - Defines whether bus or line geodata are used.

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

        **cmap** - colormap for the patch colors

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

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

        **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.

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

        **clim** (tuple of floats, None) - setting the norm limits for image scaling

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

    OUTPUT:
        **lc** - line collection
    """
    lines = net.line.index.tolist() if lines is None else list(lines)
    if len(lines) == 0:
        return None
    if line_geodata is None:
        line_geodata = net["line_geodata"]
    if bus_geodata is None:
        bus_geodata = net["bus_geodata"]
    if len(lines) == 0:
        return None

    if use_bus_geodata:
        data = [
            ([(bus_geodata.at[a, "x"], bus_geodata.at[a, "y"]),
              (bus_geodata.at[b, "x"], bus_geodata.at[b, "y"])],
             infofunc(line) if infofunc else [])
            for line, (a,
                       b) in net.line.loc[lines,
                                          ["from_bus", "to_bus"]].iterrows()
            if a in bus_geodata.index.values and b in bus_geodata.index.values
        ]
    else:
        data = [(line_geodata.loc[line, "coords"],
                 infofunc(line) if infofunc else []) for line in lines
                if line in line_geodata.index.values]

    if len(data) == 0:
        return None

    data, info = list(zip(*data))

    # This would be done anyways by matplotlib - doing it explicitly makes it a) clear and
    # b) prevents unexpected behavior when observing colors being "none"
    lc = LineCollection(data, picker=picker, **kwargs)
    lc.line_indices = np.array(lines)
    if cmap:
        if z is None:
            z = net.res_line.loading_percent.loc[lines]
        lc.set_cmap(cmap)
        lc.set_norm(norm)
        if clim is not None:
            lc.set_clim(clim)
        lc.set_array(np.array(z))
        lc.has_colormap = True
        lc.cbar_title = cbar_title
    lc.info = info

    return lc
Esempio n. 16
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)
Esempio n. 17
0
    def plot(self, time, beam=None, 
        maxground=2000, maxalt=500, step=1,
        showrefract=False, nr_cmap='jet_r', nr_lim=[0.8, 1.], 
        raycolor='0.3', title=False, zorder=2, alpha=1, 
        fig=None, rect=111, ax=None, aax=None):
        """Plot ray paths
        
        **Args**: 
            * **time** (datetime.datetime): time of rays
            * [**beam**]: beam number
            * [**maxground**]: maximum ground range [km]
            * [**maxalt**]: highest altitude limit [km]
            * [**step**]: step between each plotted ray (in number of ray steps)
            * [**showrefract**]: show refractive index along ray paths (supersedes raycolor)
            * [**nr_cmap**]: color map name for refractive index coloring
            * [**nr_lim**]: refractive index plotting limits
            * [**raycolor**]: color of ray paths
            * [**rect**]: subplot spcification
            * [**fig**]: A pylab.figure object (default to gcf)
            * [**title**]: Show default title
            * [**ax**]: Existing main axes
            * [**aax**]: Existing auxialary axes
        **Returns**:
            * **ax**: matplotlib.axes object containing formatting
            * **aax**: matplotlib.axes object containing data
            * **cbax**: matplotlib.axes object containing colorbar
        **Example**:
            ::

                # Show ray paths with colored refractive index along path
                import datetime as dt
                from models import raydarn
                sTime = dt.datetime(2012, 11, 18, 5)
                rto = raydarn.RtRun(sTime, rCode='bks', beam=12, title=True)
                rto.readRays() # read rays into memory
                ax, aax, cbax = rto.rays.plot(sTime, step=10, showrefract=True, nr_lim=[.85,1])
                ax.grid()
                
        written by Sebastien, 2013-04
        """
        from utils import plotUtils
        from matplotlib.collections import LineCollection
        import matplotlib.pyplot as plt
        import numpy as np
        from types import MethodType

        # Set up axes
        if not ax and not aax:
            ax, aax = plotUtils.curvedEarthAxes(fig=fig, rect=rect, 
                maxground=maxground, maxalt=maxalt)
        else:
            ax = ax
            aax = aax
            if hasattr(ax, 'time'):
                time = ax.time
            if hasattr(ax, 'beam'):
                beam = ax.beam

        # make sure that the required time and beam are present
        assert (time in self.paths.keys()), 'Unkown time %s' % time
        if beam:
            assert (beam in self.paths[time].keys()), 'Unkown beam %s' % beam
        else:
            beam = self.paths[time].keys()[0]
        
        for ir, (el, rays) in enumerate( sorted(self.paths[time][beam].items()) ):
            if not ir % step:
                if not showrefract:
                    aax.plot(rays['th'], rays['r']*1e-3, c=raycolor, 
                        zorder=zorder, alpha=alpha)
                else:
                    points = np.array([rays['th'], rays['r']*1e-3]).T.reshape(-1, 1, 2)
                    segments = np.concatenate([points[:-1], points[1:]], axis=1)
                    lcol = LineCollection( segments, zorder=zorder, alpha=alpha)
                    _ = lcol.set_cmap( nr_cmap )
                    _ = lcol.set_norm( plt.Normalize(*nr_lim) )
                    _ = lcol.set_array( rays['nr'] )
                    _ = aax.add_collection( lcol )

        # Plot title with date ut time and local time
        if title:
            stitle = _getTitle(time, beam, self.header, self.name)
            ax.set_title( stitle )

        # Add a colorbar when plotting refractive index
        if showrefract:
            cbax = plotUtils.addColorbar(lcol, ax)
            _ = cbax.set_ylabel("refractive index")
        else: cbax = None

        # Declare a new method to show range markers
        # This method is only available after rays have been plotted
        # This ensures that the markers match the plotted rays
        def showRange(self, markers=None, 
            color='.8', s=2, zorder=3, 
            **kwargs):
            """Plot ray paths
            
            **Args**: 
                * [**markers**]: range markers. Defaults to every 250 km
                * All other keywords are borrowed from :func:`matplotlib.pyplot.scatter`
            **Returns**:
                * **coll**: a collection of range markers
            **Example**:
                ::

                    # Add range markers to an existing ray plot
                    ax, aax, cbax = rto.rays.plot(sTime, step=10)
                    rto.rays.showRange()
                    
            written by Sebastien, 2013-04
            """

            if not markers:
                markers = np.arange(0, 5000, 250)
            
            x, y = [], []
            for el, rays in self.paths[time][beam].items():
                for rm in markers:
                    inds = (rays['gran']*1e-3 >= rm)
                    if inds.any():
                        x.append( rays['th'][inds][0] )
                        y.append( rays['r'][inds][0]*1e-3 )
            coll = aax.scatter(x, y, 
                color=color, s=s, zorder=zorder, **kwargs)

            return coll
        # End of new method

        # Assign new method
        self.showRange = MethodType(showRange, self)

        ax.beam = beam
        return ax, aax, cbax
Esempio n. 18
0
ax.axvline(0, lw=0.5, color='lightgray')
ax.axhline(1, lw=0.5, color='lightgray')

fpr, tpr, t = roc_curve(df['label'],
                        df['probabilities'],
                        drop_intermediate=False)

nth = 5000
fpr, tpr, t = fpr[::nth], tpr[::nth], t[::nth]

points = np.array([fpr, tpr]).T.reshape((-1, 1, 2))
segments = np.concatenate([points[:-1], points[1:]], axis=1)
col = LineCollection(segments, zorder=3)
col.set_array(t)
col.set_clim(0, 1)
col.set_cmap('inferno')
ax.add_collection(col)

rocs = np.array([
    roc_auc_score(fold['label'], fold['probabilities'])
    for _, fold in df.groupby('cv_fold')
])

ax.set_title(f'ROC-AUC: {rocs.mean():.4f} ± {rocs.std():.4f}')

ax.set_ylim(0, 1.025)
ax.set_xlim(-0.025, 1)
ax.set_aspect(1)
ax.set_xlabel('False Positive Rate')
ax.set_ylabel('True Positive Rate')
Esempio n. 19
0
def draw_networkx_edges(G, pos,
                        edgelist=None,
                        width=1.0,
                        edge_color='k',
                        style='solid',
                        alpha=1.0,
                        arrowstyle='-|>',
                        arrowsize=10,
                        edge_cmap=None,
                        edge_vmin=None,
                        edge_vmax=None,
                        ax=None,
                        arrows=True,
                        label=None,
                        node_size=300,
                        nodelist=None,
                        node_shape="o",
                        **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, or array of floats
       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.
       Note: Arrows will be the same color as edges.

    arrowstyle : str, optional (default='-|>')
       For directed graphs, choose the style of the arrow heads.
       See :py:class: `matplotlib.patches.ArrowStyle` for more
       options.

    arrowsize : int, optional (default=10)
       For directed graphs, choose the size of the arrow head head's length and
       width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
       `mutation_scale` for more info.

    label : [None| string]
       Label for legend

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

    list of matplotlib.patches.FancyArrowPatch
        `FancyArrowPatch` instances of the directed edges

    Depending whether the drawing includes arrows or not.

    Notes
    -----
    For directed graphs, arrows are drawn at the head end.  Arrows can be
    turned off with keyword arrows=False. Be sure to include `node_size' as a
    keyword argument; arrows are drawn considering the size of nodes.

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

    >>> G = nx.DiGraph()
    >>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
    >>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
    >>> alphas = [0.3, 0.4, 0.5]
    >>> for i, arc in enumerate(arcs):  # change alpha values of arcs
    ...     arc.set_alpha(alphas[i])

    Also see the NetworkX drawing examples at
    https://networkx.github.io/documentation/latest/auto_examples/index.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.cbook as cb
        from matplotlib.colors import colorConverter, Colormap, Normalize
        from matplotlib.collections import LineCollection
        from matplotlib.patches import FancyArrowPatch
        import numpy as np
    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 = list(G.edges())

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

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

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

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

    if not is_string_like(edge_color) \
            and cb.iterable(edge_color) \
            and len(edge_color) == len(edge_pos):
        if np.alltrue([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 np.alltrue([not is_string_like(c) for c in edge_color]):
            # If color specs are given as (rgb) or (rgba) tuples, we're OK
            if np.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 contain color names or numbers')
    else:
        if is_string_like(edge_color) or len(edge_color) == 1:
            edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
        else:
            msg = 'edge_color must be a color or list of one color per edge'
            raise ValueError(msg)

    if (not G.is_directed() or 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 by
        # 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(np.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()
        return edge_collection

    arrow_collection = None

    if G.is_directed() and arrows:
        # Note: Waiting for someone to implement arrow to intersection with
        # marker.  Meanwhile, this works well for polygons with more than 4
        # sides and circle.

        def to_marker_edge(marker_size, marker):
            if marker in "s^>v<d":  # `large` markers need extra space
                return np.sqrt(2 * marker_size) / 2
            else:
                return np.sqrt(marker_size) / 2

        # Draw arrows with `matplotlib.patches.FancyarrowPatch`
        arrow_collection = []
        mutation_scale = arrowsize  # scale factor of arrow head
        arrow_colors = edge_colors
        if arrow_colors is None:
            if edge_cmap is not None:
                assert(isinstance(edge_cmap, Colormap))
            else:
                edge_cmap = plt.get_cmap()  # default matplotlib colormap
            if edge_vmin is None:
                edge_vmin = min(edge_color)
            if edge_vmax is None:
                edge_vmax = max(edge_color)
            color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)

        for i, (src, dst) in enumerate(edge_pos):
            x1, y1 = src
            x2, y2 = dst
            arrow_color = None
            line_width = None
            shrink_source = 0  # space from source to tail
            shrink_target = 0  # space from  head to target
            if cb.iterable(node_size):  # many node sizes
                src_node, dst_node = edgelist[i]
                index_node = nodelist.index(dst_node)
                marker_size = node_size[index_node]
                shrink_target = to_marker_edge(marker_size, node_shape)
            else:
                shrink_target = to_marker_edge(node_size, node_shape)
            if arrow_colors is None:
                arrow_color = edge_cmap(color_normal(edge_color[i]))
            elif len(arrow_colors) > 1:
                arrow_color = arrow_colors[i]
            else:
                arrow_color = arrow_colors[0]
            if len(lw) > 1:
                line_width = lw[i]
            else:
                line_width = lw[0]
            arrow = FancyArrowPatch((x1, y1), (x2, y2),
                                    arrowstyle=arrowstyle,
                                    shrinkA=shrink_source,
                                    shrinkB=shrink_target,
                                    mutation_scale=mutation_scale,
                                    color=arrow_color,
                                    linewidth=line_width,
                                    zorder=1)  # arrows go behind nodes

            # There seems to be a bug in matplotlib to make collections of
            # FancyArrowPatch instances. Until fixed, the patches are added
            # individually to the axes instance.
            arrow_collection.append(arrow)
            ax.add_patch(arrow)

    # update view
    minx = np.amin(np.ravel(edge_pos[:, :, 0]))
    maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
    miny = np.amin(np.ravel(edge_pos[:, :, 1]))
    maxy = np.amax(np.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()

    return arrow_collection
Esempio n. 20
0
for p in ps:
	gm = cmpy.machines.BiasedCoin(bias = p)
	data = gm.cmech_quantities(['HX0L', 'hmu'], maxL)
	lines1.append(zip(range(maxL+1),data[0]))
	data[1].mask[0] = False
	lines2.append(zip(range(maxL+1),data[1,:]))

fig = plt.figure()
fig.set_figheight(5)
fig.set_figwidth(10)
fig.subplots_adjust(left = 0.09, right = 0.86)
ax1 = fig.add_subplot(1,2,1)

lc1 = LineCollection(lines1)
lc1.set_array(ps)
lc1.set_cmap(plt.cm.jet)

ax1.add_collection(lc1)

ax1.set_xlabel("$L$ [symbols]")
ax1.set_ylabel("$H[X_0^L]$ [bits]")
ax1.axis('auto')
ax1.set_xlim((0,maxL))

ax2 = fig.add_subplot(1,2,2)

lc2 = LineCollection(lines2)
lc2.set_array(ps)
lc2.set_cmap(plt.cm.jet)

ax2.add_collection(lc2)
Esempio n. 21
0
class SunPlotPy(wx.Frame, Spatial, Grid ):
    """ 
    The main frame of the application
    """
    title = 'sunplot(py)'

    # Plotting options
    autoclim=True
    showedges=False
    bgcolor='k'
    textcolor='w'
    cmap='RdBu'
    particlesize = 1.8
    particlecolor = 'm'

    # other flags
    collectiontype='cells'
    oldcollectiontype='cells'

    # 
    tindex=0 
    depthlevs = [0., 10., 100., 200., 300., 400., 500.,\
        1000.,2000.,3000.,4000.,5000]

    _FillValue=999999
    
    def __init__(self):
        wx.Frame.__init__(self, None, -1, self.title)
        
        self.create_menu()
        self.create_status_bar()
        self.create_main_panel()
        
        #self.draw_figure()

    def create_menu(self):
        self.menubar = wx.MenuBar()
        
        ###
        # File Menu
        ###
        menu_file = wx.Menu()
        # Load a hydro output file
        m_expt = menu_file.Append(-1, "&Open file\tCtrl-O", "Open netcdf file")
        self.Bind(wx.EVT_MENU, self.on_open_file, m_expt)

        # Load a grid file
        m_grid = menu_file.Append(-1, "&Load grid\tCtrl-G", "Load SUNTANS grid from folder")
        self.Bind(wx.EVT_MENU, self.on_load_grid, m_grid)

        # Load a particle file
        m_part = menu_file.Append(-1, "&Load PTM file\tCtrl-Shift-P", "Load a PTM file")
        self.Bind(wx.EVT_MENU, self.on_load_ptm, m_part)

        # Save current scene as an animation
        m_anim = menu_file.Append(-1,"&Save animation of current scene\tCtrl-S","Save animation")
        self.Bind(wx.EVT_MENU, self.on_save_anim, m_anim)

        # Save the current figure
        m_prin = menu_file.Append(-1,"&Print current scene\tCtrl-P","Save figure")
        self.Bind(wx.EVT_MENU, self.on_save_fig, m_prin)



        menu_file.AppendSeparator()
        # Exit
        m_exit = menu_file.Append(-1, "E&xit\tCtrl-X", "Exit")
        self.Bind(wx.EVT_MENU, self.on_exit, m_exit)

        ###
        # Tools menu
        ###
        menu_tools = wx.Menu()
        m_gridstat = menu_tools.Append(-1, "&Plot grid size statistics", "SUNTANS grid size")
        self.Bind(wx.EVT_MENU, self.on_plot_gridstat, m_gridstat)

        m_countcells = menu_tools.Append(-1, "&Count # grid cells", "Grid cell count")
        self.Bind(wx.EVT_MENU, self.on_count_cells, m_countcells)

        m_overlaybathy = menu_tools.Append(-1, "&Overlay depth contours", "Depth overlay")
        self.Bind(wx.EVT_MENU, self.on_overlay_bathy, m_overlaybathy)

        
        ###
        # Help Menu
        ###
        menu_help = wx.Menu()
        m_about = menu_help.Append(-1, "&About\tF1", "About the demo")
        self.Bind(wx.EVT_MENU, self.on_about, m_about)
        
        
        # Add all of the menu bars
        self.menubar.Append(menu_file, "&File")
        self.menubar.Append(menu_tools, "&Tools")
        self.menubar.Append(menu_help, "&Help")
        self.SetMenuBar(self.menubar)

    def create_main_panel(self):
        """ Creates the main panel with all the controls on it:
             * mpl canvas 
             * mpl navigation toolbar
             * Control panel for interaction
        """
        self.panel = wx.Panel(self)
        
        # Create the mpl Figure and FigCanvas objects. 
        # 5x4 inches, 100 dots-per-inch
        #
        self.dpi = 100
        #self.fig = Figure((7.0, 6.0), dpi=self.dpi,facecolor=self.bgcolor)
        self.fig = Figure((7.0, 6.0), dpi=self.dpi)
        self.canvas = FigCanvas(self.panel, -1, self.fig)
        
        
        # Since we have only one plot, we can use add_axes 
        # instead of add_subplot, but then the subplot
        # configuration tool in the navigation toolbar wouldn't
        # work.
        #
        self.axes = self.fig.add_subplot(111)
        #SetAxColor(self.axes,self.textcolor,self.bgcolor)
        
        # Bind the 'pick' event for clicking on one of the bars
        #
        #self.canvas.mpl_connect('pick_event', self.on_pick)
        
        ########
        # Create widgets
        ########
        self.variable_list = wx.ComboBox(
            self.panel, 
            size=(200,-1),
            choices=['Select a variable...'],
            style=wx.CB_READONLY)
        self.variable_list.Bind(wx.EVT_COMBOBOX, self.on_select_variable)
        
        self.time_list = wx.ComboBox(
            self.panel, 
            size=(200,-1),
            choices=['Select a time step...'],
            style=wx.CB_READONLY)
        self.time_list.Bind(wx.EVT_COMBOBOX, self.on_select_time)

        self.depthlayer_list = wx.ComboBox(
            self.panel, 
            size=(200,-1),
            choices=['Select a vertical layer...'],
            style=wx.CB_READONLY)
        self.depthlayer_list.Bind(wx.EVT_COMBOBOX, self.on_select_depth)

        self.show_edge_check = wx.CheckBox(self.panel, -1, 
            "Show Edges",
            style=wx.ALIGN_RIGHT)
        self.show_edge_check.Bind(wx.EVT_CHECKBOX, self.on_show_edges)

        if USECMOCEAN:
            cmaps=[]
            for cmap in cm.cmapnames:
                cmaps.append(cmap)
                cmaps.append(cmap+'_r') # Add all reverse map options
        else:
            # Use matplotlib standard
            cmaps = list(matplotlib.cm.datad.keys())

        cmaps.sort()
        self.colormap_list = wx.ComboBox(
            self.panel, 
            size=(100,-1),
            choices=cmaps,
            style=wx.CB_READONLY)
        self.colormap_list.Bind(wx.EVT_COMBOBOX, self.on_select_cmap)
        self.colormap_label = wx.StaticText(self.panel, -1,"Colormap:")

        self.clim_check = wx.CheckBox(self.panel, -1, 
            "Manual color limits ",
            style=wx.ALIGN_RIGHT)
        self.clim_check.Bind(wx.EVT_CHECKBOX, self.on_clim_check)

        self.climlow = wx.TextCtrl(
            self.panel, 
            size=(100,-1),
            style=wx.TE_PROCESS_ENTER)
        self.climlow.Bind(wx.EVT_TEXT_ENTER, self.on_climlow)
        
        self.climhigh = wx.TextCtrl(
            self.panel, 
            size=(100,-1),
            style=wx.TE_PROCESS_ENTER)
        self.climhigh.Bind(wx.EVT_TEXT_ENTER, self.on_climhigh)
 


        # Labels
        self.variable_label = wx.StaticText(self.panel, -1,"Variable:",size=(200,-1))
        self.time_label = wx.StaticText(self.panel, -1,"Time step:",size=(200,-1))
        self.depth_label = wx.StaticText(self.panel, -1,"Vertical level:",size=(200,-1))


        # Create the navigation toolbar, tied to the canvas
        #
        self.toolbar = NavigationToolbar(self.canvas)
        #self.toolbar.toolitems[8][3]='my_save_fig'

        #def my_save_fig(self,*args):
        #    print 'saving figure'
        #    return "break"

        
        #########
        # Layout with box sizers
        #########
        
        self.vbox = wx.BoxSizer(wx.VERTICAL)
        self.vbox.Add(self.canvas, 1, wx.LEFT | wx.TOP | wx.GROW)
        self.vbox.Add(self.toolbar, 0, wx.EXPAND)

        self.vbox.AddSpacer(10)
        #self.vbox.Add((-1,25))

        flags = wx.ALIGN_LEFT | wx.ALL | wx.ALIGN_CENTER_VERTICAL

        self.hbox0 = wx.BoxSizer(wx.HORIZONTAL)
        self.hbox0.Add(self.show_edge_check, 0, border=10, flag=flags)
        self.hbox0.Add(self.colormap_label, 0, border=10, flag=flags)
        self.hbox0.Add(self.colormap_list, 0, border=10, flag=flags)
        self.hbox0.Add(self.clim_check, 0, border=10, flag=flags)
        self.hbox0.Add(self.climlow, 0, border=10, flag=flags)
        self.hbox0.Add(self.climhigh, 0, border=10, flag=flags)

        self.vbox.AddSpacer(5)
        self.hbox1 = wx.BoxSizer(wx.HORIZONTAL)
        self.hbox1.Add(self.variable_label, 0, border=10, flag=flags)
        self.hbox1.Add(self.time_label, 0, border=10, flag=flags)
        self.hbox1.Add(self.depth_label, 0, border=10, flag=flags)

        self.vbox.AddSpacer(5)
        self.hbox2 = wx.BoxSizer(wx.HORIZONTAL)
        self.hbox2.Add(self.variable_list, 0, border=10, flag=flags)
        self.hbox2.Add(self.time_list, 0, border=10, flag=flags)
        self.hbox2.Add(self.depthlayer_list, 0, border=10, flag=flags)
       
        self.vbox.Add(self.hbox1, 0, flag = wx.ALIGN_LEFT | wx.TOP)
        self.vbox.Add(self.hbox2, 0, flag = wx.ALIGN_LEFT | wx.TOP)
        self.vbox.Add(self.hbox0, 0, flag = wx.ALIGN_LEFT | wx.TOP)
        
        self.panel.SetSizer(self.vbox)
        self.vbox.Fit(self)
    
    ##########
    # Event functions
    ##########

    def create_figure(self):
        """ 
        Creates the figure
        """
        # Find the colorbar limits if unspecified
        if self.autoclim:
            self.clim = [self.data.min(),self.data.max()]
            self.climlow.SetValue('%3.1f'%self.clim[0])
            self.climhigh.SetValue('%3.1f'%self.clim[1])
         
        if 'collection' in self.__dict__:
            #self.collection.remove()
            self.axes.collections.remove(self.collection)
        else:
            # First call - set the axes limits
            self.axes.set_aspect('equal')
            self.axes.set_xlim(self.xlims)
            self.axes.set_ylim(self.ylims)
 

        if self.collectiontype=='cells':
            self.collection = PolyCollection(self.xy,cmap=self.cmap)
            self.collection.set_array(np.array(self.data[:]))
            if not self.showedges:
                self.collection.set_edgecolors(self.collection.to_rgba(np.array((self.data[:])))) 
        elif self.collectiontype=='edges':
            xylines = [self.xp[self.edges],self.yp[self.edges]]
            linesc = [list(zip(xylines[0][ii,:],xylines[1][ii,:])) for ii in range(self.Ne)]
            self.collection = LineCollection(linesc,array=np.array(self.data[:]),cmap=self.cmap)

        self.collection.set_clim(vmin=self.clim[0],vmax=self.clim[1])

        self.axes.add_collection(self.collection)    
        self.title=self.axes.set_title(self.genTitle(),color=self.textcolor)
        self.axes.set_xlabel('Easting [m]')
        self.axes.set_ylabel('Northing [m]')

        # create a colorbar

        if 'cbar' not in self.__dict__:
            self.cbar = self.fig.colorbar(self.collection)
            #SetAxColor(self.cbar.ax.axes,self.textcolor,self.bgcolor)
        else:
            #pass
            print('Updating colorbar...')
            #self.cbar.check_update(self.collection)
            self.cbar.on_mappable_changed(self.collection)

        self.canvas.draw()
   
    def update_figure(self):
        if self.autoclim:
            self.clim = [self.data.min(),self.data.max()]
            self.climlow.SetValue('%3.1f'%self.clim[0])
            self.climhigh.SetValue('%3.1f'%self.clim[1])
        else:
            self.clim = [float(self.climlow.GetValue()),\
                float(self.climhigh.GetValue())]
 
        # check whether it is cell or edge type
        if self.hasDim(self.variable,self.griddims['Ne']):
            self.collectiontype='edges'
        elif self.hasDim(self.variable,self.griddims['Nc']):
            self.collectiontype='cells'

        # Create a new figure if the variable has gone from cell to edge of vice
        # versa
        if not self.collectiontype==self.oldcollectiontype:
            self.create_figure()
            self.oldcollectiontype=self.collectiontype

        self.collection.set_array(np.array(self.data[:]))
        self.collection.set_clim(vmin=self.clim[0],vmax=self.clim[1])

        # Cells only
        if self.collectiontype=='cells':
            if not self.showedges:
                self.collection.set_edgecolors(self.collection.to_rgba(np.array((self.data[:])))) 
            else:
                self.collection.set_edgecolors('k')
                self.collection.set_linewidths(0.2)

        # Update the title
        self.title=self.axes.set_title(self.genTitle(),color=self.textcolor)

        #Update the colorbar
        self.cbar.update_normal(self.collection)

        # redraw the figure
        self.canvas.draw()
    
    def on_pick(self, event):
        # The event received here is of the type
        # matplotlib.backend_bases.PickEvent
        #
        # It carries lots of information, of which we're using
        # only a small amount here.
        # 
        box_points = event.artist.get_bbox().get_points()
        msg = "You've clicked on a bar with coords:\n %s" % box_points
        
        dlg = wx.MessageDialog(
            self, 
            msg, 
            "Click!",
            wx.OK | wx.ICON_INFORMATION)

        dlg.ShowModal() 
        dlg.Destroy()        
    
    def on_select_variable(self, event):
        vname = event.GetString()
        self.flash_status_message("Selecting variable: %s"%vname)
        # update the spatial object and load the data
        self.variable = vname
        self.loadData(variable=self.variable)

        # Check if the variable has a depth coordinate
        depthstr = ['']
        # If so populate the vertical layer box
        if self.hasDim(self.variable,self.griddims['Nk']):
            depthstr = ['%3.1f'%self.z_r[k] for k in range(self.Nkmax)]
            depthstr += ['surface','seabed']

        elif self.hasDim(self.variable,'Nkw'):
            depthstr = ['%3.1f'%self.z_w[k] for k in range(self.Nkmax+1)]

        self.depthlayer_list.SetItems(depthstr)

        # Update the plot
        self.update_figure()



    def on_select_time(self, event):
        self.tindex = event.GetSelection()
        # Update the object time index and reload the data
        if self.plot_type=='hydro':
            if not self.tstep==self.tindex:
                self.tstep=self.tindex
                self.loadData()
                self.flash_status_message("Selecting variable: %s..."%event.GetString())

                # Update the plot
                self.update_figure()
        elif self.plot_type=='particles':
            self.PTM.plot(self.tindex,ax=self.axes,\
                xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
        
            self.canvas.draw()


    def on_select_depth(self, event):
        kindex = event.GetSelection()
        if not self.klayer[0]==kindex:
            # Check if its the seabed or surface value
            if kindex>=self.Nkmax:
                kindex=event.GetString()
            self.klayer = [kindex]
            self.loadData()       
            self.flash_status_message("Selecting depth: %s..."%event.GetString())

            # Update the plot
            self.update_figure()

    def on_open_file(self, event):
        file_choices = "SUNTANS NetCDF (*.nc)|*.nc*|UnTRIM NetCDF (*.nc)|*.nc*|All Files (*.*)|*.*"
        
        dlg = wx.FileDialog(
            self, 
            message="Open SUNTANS file...",
            defaultDir=os.getcwd(),
            defaultFile="",
            wildcard=file_choices,
            style= wx.FD_MULTIPLE)
        
        if dlg.ShowModal() == wx.ID_OK:
            self.plot_type='hydro'

            path = dlg.GetPaths()

            # Initialise the class
            if dlg.GetFilterIndex() == 0 or dlg.GetFilterIndex() > 1: #SUNTANS
                self.flash_status_message("Opening SUNTANS file: %s" % path)
		try:
		    Spatial.__init__(self, path, _FillValue=self._FillValue)
		except:
		    Spatial.__init__(self, path, _FillValue=-999999)
            startvar='dv'
            
            # Populate the drop down menus
            vnames = self.listCoordVars()
            self.variable_list.SetItems(vnames)
            
            # Update the time drop down list
            if 'time' in self.__dict__:
                self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.time]
            else:
                # Assume that it is a harmonic-type file
                self.timestr = self.nc.Constituent_Names.split()

            self.time_list.SetItems(self.timestr)

            # Draw the depth
            if startvar in vnames:
                self.variable=startvar
                self.loadData()
                self.create_figure()

    def on_load_grid(self, event):
        
        dlg = wx.DirDialog(
            self, 
            message="Open SUNTANS grid from folder...",
            defaultPath=os.getcwd(),
            style= wx.DD_DEFAULT_STYLE)
        
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()

            # Initialise the class
            self.flash_status_message("Opening SUNTANS grid from folder: %s" % path)
            Grid.__init__(self,path)

            # Plot the Grid
            if 'collection' in self.__dict__:
                self.axes.collections.remove(self.collection)

            self.axes,self.collection = self.plotmesh(ax=self.axes,edgecolors='y')

            # redraw the figure
            self.canvas.draw()

    def on_load_ptm(self, event):
        file_choices = "PTM NetCDF (*.nc)|*.nc|PTM Binary (*_bin.out)|*_bin.out|All Files (*.*)|*.*"
        
        dlg = wx.FileDialog(
            self, 
            message="Open PTM file...",
            defaultDir=os.getcwd(),
            defaultFile="",
            wildcard=file_choices,
            style= wx.FD_MULTIPLE)
        
        if dlg.ShowModal() == wx.ID_OK:
            self.plot_type = 'particles'
            path = dlg.GetPath()

            # Initialise the class
            if dlg.GetFilterIndex() == 0: #SUNTANS
                self.flash_status_message("Opening PTM netcdf file: %s" % path)
                self.PTM = PtmNC(path)
            elif dlg.GetFilterIndex() == 1: #PTM
                self.flash_status_message("Opening PTM binary file: %s" % path)
                self.PTM = PtmBin(path)

            self.Nt = self.PTM.nt
            
            # Update the time drop down list
            self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.PTM.time]
            self.time_list.SetItems(self.timestr)

            # Plot the first time step
            if 'xlims' in self.__dict__:
                self.PTM.plot(self.PTM.nt-1,ax=self.axes,xlims=self.xlims,\
                ylims=self.ylims,color=self.particlecolor,\
                fontcolor='w',markersize=self.particlesize)
            else:
                self.PTM.plot(self.PTM.nt-1,ax=self.axes,fontcolor='w',\
                    color=self.particlecolor,markersize=self.particlesize)
            # redraw the figure
            self.canvas.draw()

        
    def on_show_edges(self,event):
        sender=event.GetEventObject()
        self.showedges = sender.GetValue()

        # Update the figure
        self.update_figure()

    def on_clim_check(self,event):
        sender=event.GetEventObject()
        if sender.GetValue() == True:
            self.autoclim=False
            self.update_figure()
        else:
            self.autoclim=True
       

    def on_climlow(self,event):
        self.clim[0] = event.GetString()
        #self.update_figure()

    def on_climhigh(self,event):
        self.clim[1] = event.GetString()
        #self.update_figure()

    def on_select_cmap(self,event):
        self.cmap=event.GetString()
        if USECMOCEAN:
            self.collection.set_cmap(getattr(cm,self.cmap))
        else:
            self.collection.set_cmap(self.cmap)

        # Update the figure
        self.update_figure()

    def on_save_fig(self,event):
        """
        Save a figure of the current scene to a file
        """
        file_choices = " (*.png)|*.png| (*.pdf)|*.pdf |(*.jpg)|*.jpg |(*.eps)|*eps "
        filters=['.png','.pdf','.png','.png']

        
        dlg = wx.FileDialog(
            self, 
            message="Save figure to file...",
            defaultDir=os.getcwd(),
            defaultFile="",
            wildcard=file_choices,
            style= wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)

        if dlg.ShowModal() == wx.ID_OK:

            path = dlg.GetPath()
            ext = filters[dlg.GetFilterIndex()]
            if ext in path:
                outfile=path
            else:
                outfile = path+ext

            self.fig.savefig(outfile)

            


    def on_save_anim(self,event):
        """
        Save an animation of the current scene to a file
        """
        file_choices = "Quicktime (*.mov)|*.mov| (*.gif)|*.gif| (*.avi)|*.avi |(*.mp4)|*.mp4 "
        filters=['.mov','.gif','.avi','.mp4']

        
        dlg = wx.FileDialog(
            self, 
            message="Output animation file...",
            defaultDir=os.getcwd(),
            defaultFile="",
            wildcard=file_choices,
            style= wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)

        if dlg.ShowModal() == wx.ID_OK:

            path = dlg.GetPath()
            ext = filters[dlg.GetFilterIndex()]
            if ext in path:
                outfile=path
            else:
                outfile = path+ext
            self.flash_status_message("Saving figure to file: %s" %outfile)
            self.flash_status_message("Saving animation to file: %s" %outfile)

            # Create the animation
            #self.tstep = range(self.Nt) # Use all time steps for animation
            #self.animate(cbar=self.cbar,cmap=self.cmap,\
            #    xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
            def initanim():
                if not self.plot_type=='particles':
                    return (self.title, self.collection)
                else:
                    return (self.PTM.title,self.PTM.p_handle)

            def updateScalar(i):
                if not self.plot_type=='particles':
                    self.tstep=[i]
                    self.loadData()
                    self.update_figure()
                    return (self.title,self.collection)
                elif self.plot_type=='particles':
                    self.PTM.plot(i,ax=self.axes,\
                        xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
                    return (self.PTM.title,self.PTM.p_handle)

            self.anim = animation.FuncAnimation(self.fig, \
                updateScalar, init_func = initanim, frames=self.Nt, interval=50, blit=True)

            if ext=='.gif':
                self.anim.save(outfile,writer='imagemagick',fps=6)
            elif ext=='.mp4':
                print('Saving html5 video...')
                # Ensures html5 compatibility
                self.anim.save(outfile,fps=6,\
                    writer='ffmpeg',\
                    bitrate=3600,extra_args=['-vcodec','libx264'])
                    #writer='mencoder',
                    #bitrate=3600,extra_args=['-ovc','x264']) # mencoder options
            else:
                self.anim.save(outfile,writer='mencoder',fps=6,bitrate=3600)

            # Return the figure back to its status
            del self.anim
            self.tstep=self.tindex
            if not self.plot_type=='particles':
                self.loadData()
                self.update_figure()

            # Bring up a dialog box
            dlg2= wx.MessageDialog(self, 'Animation complete.', "Done", wx.OK)
            dlg2.ShowModal()
            dlg2.Destroy()

    def on_exit(self, event):
        self.Destroy()
        
    def on_about(self, event):
        msg = """ SUNTANS NetCDF visualization tool
        
            *Author: Matt Rayson
            *Institution: Stanford University
            *Created: October 2013
        """
        dlg = wx.MessageDialog(self, msg, "About", wx.OK)
        dlg.ShowModal()
        dlg.Destroy()

    def on_count_cells(self,eveny):
        msg = "Total 3-D grid cells = %d"%(self.count_cells())
        dlg = wx.MessageDialog(self, msg, "No. cells", wx.OK)
        dlg.ShowModal()
        dlg.Destroy()

    def on_overlay_bathy(self,event):
        # Plot depth contours
        print('Plotting contours...')
        self.contourf(z=self.dv, clevs=self.depthlevs,\
            ax=self.axes,\
            filled=False, colors='0.5', linewidths=0.5, zorder=1e6)
        print('Done')
   
    def on_plot_gridstat(self, event):
        """
        Plot the grid size histogram in a new figure
        """
        matplotlib.pyplot.figure()
        self.plothist()
        matplotlib.pyplot.show()


    def create_status_bar(self):
        self.statusbar = self.CreateStatusBar()

    def flash_status_message(self, msg, flash_len_ms=1500):
        self.statusbar.SetStatusText(msg)
        self.timeroff = wx.Timer(self)
        self.Bind(
            wx.EVT_TIMER, 
            self.on_flash_status_off, 
            self.timeroff)
        self.timeroff.Start(flash_len_ms, oneShot=True)
    
    def on_flash_status_off(self, event):
        self.statusbar.SetStatusText('')
Esempio n. 22
0
class SunPlotPy(Spatial, QMainWindow):
    """ 
    The main frame of the application
    """
    title = 'sunplot(py)'

    # Plotting options
    autoclim = True
    showedges = False
    bgcolor = 'k'
    textcolor = 'w'
    cmap = 'RdBu'
    particlesize = 1.8
    particlecolor = 'm'

    # other flags
    collectiontype = 'cells'
    oldcollectiontype = 'cells'

    #
    tindex = 0
    depthlevs = [0., 10., 100., 200., 300., 400., 500.,\
        1000.,2000.,3000.,4000.,5000]

    _FillValue = 999999

    def __init__(self, parent=None):
        #wx.Frame.__init__(self, None, -1, self.title)
        QMainWindow.__init__(self, parent)
        #super(SunPlotPy, self).__init__(parent)

        self.create_menu()
        #self.create_status_bar()
        self.create_main_panel()

        #self.draw_figure()

    def create_menu(self):
        self.file_menu = self.menuBar().addMenu("&File")

        ###
        # File Menu
        ###
        # Load a hydro output file
        load_file_action = self.create_action("&Open file",\
                shortcut="ctrl-o", slot=self.on_open_file,
                tip="open netcdf file")

        load_grid_action = self.create_action("&Open grid",\
                shortcut="ctrl-g", slot=self.on_load_grid,
                tip="open suntans grid")

        save_anim_action = self.create_action("&Save animation",\
                shortcut="ctrl-a", slot=self.on_save_anim,
                tip="animate current scene")

        quit_action = self.create_action("&Exit",\
                shortcut="ctrl-x", slot=self.close,
                tip="Close the application")

        self.add_actions(self.file_menu,
                (load_file_action, load_grid_action,\
                save_anim_action, None, quit_action))

        #    self.Bind(wx.EVT_MENU, self.on_open_file, m_expt)
        ###
        # Tools menu
        ###
        self.tools_menu = self.menuBar().addMenu("&Tools")

        ###
        # File Menu
        ###
        # Load a hydro output file
        load_stat_action = self.create_action("&Plot grid size statistics",\
                slot=self.on_plot_gridstat,
                tip="grid stats")

        self.add_actions(self.tools_menu, (load_stat_action, ))

    #    # Load a grid file
    #    m_grid = menu_file.Append(-1, "&Load grid\tCtrl-G", "Load SUNTANS grid from folder")
    #    self.Bind(wx.EVT_MENU, self.on_load_grid, m_grid)

    #    # Load a particle file
    #    m_part = menu_file.Append(-1, "&Load PTM file\tCtrl-Shift-P", "Load a PTM file")
    #    self.Bind(wx.EVT_MENU, self.on_load_ptm, m_part)

    #    # Save current scene as an animation
    #    m_anim = menu_file.Append(-1,"&Save animation of current scene\tCtrl-S","Save animation")
    #    self.Bind(wx.EVT_MENU, self.on_save_anim, m_anim)

    #    # Save the current figure
    #    m_prin = menu_file.Append(-1,"&Print current scene\tCtrl-P","Save figure")
    #    self.Bind(wx.EVT_MENU, self.on_save_fig, m_prin)

    #    menu_file.AppendSeparator()
    #    # Exit
    #    m_exit = menu_file.Append(-1, "E&xit\tCtrl-X", "Exit")
    #    self.Bind(wx.EVT_MENU, self.on_exit, m_exit)

    #    ###
    #    # Tools menu
    #    ###
    #    menu_tools = wx.Menu()
    #    m_gridstat = menu_tools.Append(-1, "&Plot grid size statistics", "SUNTANS grid size")
    #    self.Bind(wx.EVT_MENU, self.on_plot_gridstat, m_gridstat)

    #    m_countcells = menu_tools.Append(-1, "&Count # grid cells", "Grid cell count")
    #    self.Bind(wx.EVT_MENU, self.on_count_cells, m_countcells)

    #    m_overlaybathy = menu_tools.Append(-1, "&Overlay depth contours", "Depth overlay")
    #    self.Bind(wx.EVT_MENU, self.on_overlay_bathy, m_overlaybathy)

    #
    #    ###
    #    # Help Menu
    #    ###
    #    menu_help = wx.Menu()
    #    m_about = menu_help.Append(-1, "&About\tF1", "About the demo")
    #    self.Bind(wx.EVT_MENU, self.on_about, m_about)
    #
    #
    #    # Add all of the menu bars
    #    self.menubar.Append(menu_file, "&File")
    #    self.menubar.Append(menu_tools, "&Tools")
    #    self.menubar.Append(menu_help, "&Help")
    #    self.SetMenuBar(self.menubar)

    def add_actions(self, target, actions):
        for action in actions:
            if action is None:
                target.addSeparator()
            else:
                target.addAction(action)

    def create_action(
        self,
        text,
        slot=None,
        shortcut=None,
        icon=None,
        tip=None,
        checkable=False,
    ):
        #signal="triggered()"):
        action = QAction(text, self)
        if icon is not None:
            action.setIcon(QIcon(":/%s.png" % icon))
        if shortcut is not None:
            action.setShortcut(shortcut)

        if tip is not None:
            action.setToolTip(tip)
            action.setStatusTip(tip)
        if slot is not None:
            #self.connect(action, SIGNAL(signal), slot)
            # Qt5 style
            action.triggered.connect(slot)
        if checkable:
            action.setCheckable(True)
        return action

    def create_main_panel(self):
        """ Creates the main panel with all the controls on it:
             * mpl canvas 
             * mpl navigation toolbar
             * Control panel for interaction
        """
        self.panel = QWidget()

        # Create the mpl Figure and FigCanvas objects.
        # 5x4 inches, 100 dots-per-inch
        #
        self.dpi = 100
        #self.fig = Figure((7.0, 6.0), dpi=self.dpi,facecolor=self.bgcolor)
        self.fig = Figure((7.0, 6.0), dpi=self.dpi)
        #self.canvas = FigCanvas(self.panel, -1, self.fig)
        self.canvas = FigureCanvas(self.fig)
        self.canvas.setParent(self.panel)

        # Since we have only one plot, we can use add_axes
        # instead of add_subplot, but then the subplot
        # configuration tool in the navigation toolbar wouldn't
        # work.
        #
        self.axes = self.fig.add_subplot(111)
        #SetAxColor(self.axes,self.textcolor,self.bgcolor)

        # Bind the 'pick' event for clicking on one of the bars
        #
        #self.canvas.mpl_connect('pick_event', self.on_pick)

        #########
        ## Create widgets
        #########
        self.variable_list = QComboBox()
        self.variable_list.addItem("Select a variable...")
        self.variable_list.activated[str].connect(self.on_select_variable)

        self.time_list = QComboBox()
        self.time_list.addItem("Select a time...")
        self.time_list.activated[int].connect(self.on_select_time)

        self.depthlayer_list = QComboBox()
        self.depthlayer_list.addItem("Select a vertical layer...")
        self.depthlayer_list.activated[int].connect(self.on_select_depth)

        self.show_edge_check = QCheckBox('Show Edges', self)
        #self.show_edge_check.toggle()
        self.show_edge_check.stateChanged.connect(self.on_show_edges)

        cmaps = list(matplotlib.cm.datad.keys())
        cmaps.sort()
        self.colormap_list = QComboBox()
        self.colormap_list.clear()
        self.colormap_list.addItems(cmaps)
        self.colormap_list.activated[str].connect(self.on_select_cmap)

        self.clim_check = QCheckBox('Manual color limits', self)
        #self.show_edge_check.toggle()
        self.clim_check.stateChanged.connect(self.on_clim_check)

        #self.clim_check = wx.CheckBox(self.panel, -1,
        #    "Manual color limits ",
        #    style=wx.ALIGN_RIGHT)
        #self.clim_check.Bind(wx.EVT_CHECKBOX, self.on_clim_check)

        self.climlow = QLineEdit()
        self.climlow.textChanged[str].connect(self.on_climlow)

        self.climhigh = QLineEdit()
        self.climhigh.textChanged[str].connect(self.on_climhigh)

        ## Labels
        #self.variable_label = wx.StaticText(self.panel, -1,"Variable:",size=(200,-1))
        #self.time_label = wx.StaticText(self.panel, -1,"Time step:",size=(200,-1))
        #self.depth_label = wx.StaticText(self.panel, -1,"Vertical level:",size=(200,-1))

        # Create the navigation toolbar, tied to the canvas
        #
        self.toolbar = NavigationToolbar(self.canvas, self)
        #self.toolbar.toolitems[8][3]='my_save_fig'

        #def my_save_fig(self,*args):
        #     print 'saving figure'
        ##    return "break"

        #########
        # Layout with box sizers
        #########
        hbox = QHBoxLayout()
        for w in [self.variable_list, self.time_list, self.depthlayer_list]:
            hbox.addWidget(w)
            hbox.setAlignment(w, Qt.AlignVCenter)

        hbox1 = QHBoxLayout()
        for w in [
                self.show_edge_check, self.colormap_list, self.clim_check,
                self.climlow, self.climhigh
        ]:
            hbox1.addWidget(w)
            hbox1.setAlignment(w, Qt.AlignVCenter)

        self.vbox = QVBoxLayout()
        self.vbox.addWidget(self.canvas)
        self.vbox.addWidget(self.toolbar)
        self.vbox.addLayout(hbox)
        self.vbox.addLayout(hbox1)

        self.panel.setLayout(self.vbox)
        self.setCentralWidget(self.panel)

#
###########
## Event functions
###########

    def create_figure(self):
        """ 
        Creates the figure
        """
        # Find the colorbar limits if unspecified
        if self.autoclim:
            self.clim = [self.data.min(), self.data.max()]
            self.climlow.setText('%3.1f' % self.clim[0])
            self.climhigh.setText('%3.1f' % self.clim[1])

        if 'collection' in self.__dict__:
            #self.collection.remove()
            self.axes.collections.remove(self.collection)
        else:
            # First call - set the axes limits
            self.axes.set_aspect('equal')
            self.axes.set_xlim(self.xlims)
            self.axes.set_ylim(self.ylims)

        if self.collectiontype == 'cells':
            self.collection = PolyCollection(self.xy, cmap=self.cmap)
            self.collection.set_array(np.array(self.data[:]))
            if not self.showedges:
                self.collection.set_edgecolors(
                    self.collection.to_rgba(np.array((self.data[:]))))
        elif self.collectiontype == 'edges':
            xylines = [self.xp[self.edges], self.yp[self.edges]]
            linesc = [
                list(zip(xylines[0][ii, :], xylines[1][ii, :]))
                for ii in range(self.Ne)
            ]
            self.collection = LineCollection(linesc,
                                             array=np.array(self.data[:]),
                                             cmap=self.cmap)

        self.collection.set_clim(vmin=self.clim[0], vmax=self.clim[1])

        self.axes.add_collection(self.collection)
        self.title = self.axes.set_title(self.genTitle(), color=self.textcolor)
        self.axes.set_xlabel('Easting [m]')
        self.axes.set_ylabel('Northing [m]')

        # create a colorbar

        if 'cbar' not in self.__dict__:
            self.cbar = self.fig.colorbar(self.collection)
            #SetAxColor(self.cbar.ax.axes,self.textcolor,self.bgcolor)
        else:
            #pass
            print('Updating colorbar...')
            #self.cbar.check_update(self.collection)
            self.cbar.on_mappable_changed(self.collection)

        self.canvas.draw()

    def update_figure(self):
        if self.autoclim:
            self.clim = [self.data.min(), self.data.max()]
            self.climlow.setText('%3.1f' % self.clim[0])
            self.climhigh.setText('%3.1f' % self.clim[1])
        else:
            self.clim = [float(self.climlow.text()),\
                float(self.climhigh.text())]

        # check whether it is cell or edge type
        if self.hasDim(self.variable, self.griddims['Ne']):
            self.collectiontype = 'edges'
        elif self.hasDim(self.variable, self.griddims['Nc']):
            self.collectiontype = 'cells'

        # Create a new figure if the variable has gone from cell to edge of vice
        # versa
        if not self.collectiontype == self.oldcollectiontype:
            self.create_figure()
            self.oldcollectiontype = self.collectiontype

        self.collection.set_array(np.array(self.data[:]))
        self.collection.set_clim(vmin=self.clim[0], vmax=self.clim[1])

        # Cells only
        if self.collectiontype == 'cells':
            if not self.showedges:
                self.collection.set_edgecolors(
                    self.collection.to_rgba(np.array((self.data[:]))))
            else:
                self.collection.set_edgecolors('k')
                self.collection.set_linewidths(0.2)

        # Update the title
        self.title = self.axes.set_title(self.genTitle(), color=self.textcolor)

        #Update the colorbar
        self.cbar.update_normal(self.collection)

        # redraw the figure
        self.canvas.draw()

    #def on_pick(self, event):
    #    # The event received here is of the type
    #    # matplotlib.backend_bases.PickEvent
    #    #
    #    # It carries lots of information, of which we're using
    #    # only a small amount here.
    #    #
    #    box_points = event.artist.get_bbox().get_points()
    #    msg = "You've clicked on a bar with coords:\n %s" % box_points
    #
    #    dlg = wx.MessageDialog(
    #        self,
    #        msg,
    #        "Click!",
    #        wx.OK | wx.ICON_INFORMATION)

    #    dlg.ShowModal()
    #    dlg.Destroy()
    #
    def on_select_variable(self, event):
        #vname = event.GetString()
        vname = event
        #self.flash_status_message("Selecting variable: %s"%vname)
        # update the spatial object and load the data
        self.variable = vname
        self.loadData(variable=self.variable)

        # Check if the variable has a depth coordinate
        depthstr = ['']
        # If so populate the vertical layer box
        if self.hasDim(self.variable, self.griddims['Nk']):
            depthstr = ['%3.1f' % self.z_r[k] for k in range(self.Nkmax)]
            depthstr += ['surface', 'seabed']

        elif self.hasDim(self.variable, 'Nkw'):
            depthstr = ['%3.1f' % self.z_w[k] for k in range(self.Nkmax + 1)]

        self.depthlayer_list.clear()
        self.depthlayer_list.addItems(depthstr)

        # Update the plot
        self.update_figure()

    def on_select_time(self, event):

        self.tindex = event  #
        # Update the object time index and reload the data
        #if self.plot_type=='hydro':
        if not self.tstep == self.tindex:
            self.tstep = self.tindex
            self.loadData()
            #self.flash_status_message("Selecting variable: %s..."%event.GetString())

            # Update the plot
            self.update_figure()
        #elif self.plot_type=='particles':
        #    self.PTM.plot(self.tindex,ax=self.axes,\
        #        xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
        #
        #    self.canvas.draw()

    def on_select_depth(self, event):
        print(event)
        kindex = event
        if not self.klayer[0] == kindex:
            # Check if its the seabed or surface value
            if kindex >= self.Nkmax:
                kindex = event.GetString()
            self.klayer = [kindex]
            self.loadData()
            #self.flash_status_message("Selecting depth: %s..."%event.GetString())

            # Update the plot
            self.update_figure()

    def on_open_file(self, event):
        file_choices = "SUNTANS NetCDF (*.nc);;All Files (*.*)"

        dlg = QFileDialog.getOpenFileNames(self, "Open SUNTANS file...", "",
                                           file_choices)

        path = dlg[0]

        if len(path) == 0:
            return

        self.statusBar().showMessage("Opening SUNTANS file: %s" % path)
        try:
            Spatial.__init__(self, path, _FillValue=self._FillValue)
        except:
            Spatial.__init__(self, path, _FillValue=-999999)
        startvar = 'dv'

        self.statusBar().clearMessage()

        # Populate the drop down menus
        vnames = self.listCoordVars()
        self.variable_list.clear()
        self.variable_list.addItems(vnames)
        # Update the time drop down list
        if 'time' in self.__dict__:
            self.timestr = [
                datetime.strftime(tt, '%d-%b-%Y %H:%M:%S') for tt in self.time
            ]
        else:
            # Assume that it is a harmonic-type file
            self.timestr = self.nc.Constituent_Names.split()

        self.time_list.clear()
        self.time_list.addItems(self.timestr)

        # Draw the depth
        if startvar in vnames:
            self.variable = startvar
            self.loadData()
            self.create_figure()

    def on_load_grid(self, event):

        dir_ = QFileDialog.getExistingDirectory(None, 'Select a SUNTANS grid folder:',\
                '~/', QFileDialog.ShowDirsOnly)
        print(dir_)

        if dir_ is not None:
            path = dir_

            # Initialise the class
            #self.flash_status_message("Opening SUNTANS grid from folder: %s" % path)
            Grid.__init__(self, path)

            # Plot the Grid
            if 'collection' in self.__dict__:
                self.axes.collections.remove(self.collection)

            self.axes, self.collection = self.plotmesh(ax=self.axes,
                                                       edgecolors='y')

            # redraw the figure
            self.canvas.draw()

    #def on_load_ptm(self, event):
    #    file_choices = "PTM NetCDF (*.nc)|*.nc|PTM Binary (*_bin.out)|*_bin.out|All Files (*.*)|*.*"
    #
    #    dlg = wx.FileDialog(
    #        self,
    #        message="Open PTM file...",
    #        defaultDir=os.getcwd(),
    #        defaultFile="",
    #        wildcard=file_choices,
    #        style= wx.FD_MULTIPLE)
    #
    #    if dlg.ShowModal() == wx.ID_OK:
    #        self.plot_type = 'particles'
    #        path = dlg.GetPath()

    #        # Initialise the class
    #        if dlg.GetFilterIndex() == 0: #SUNTANS
    #            self.flash_status_message("Opening PTM netcdf file: %s" % path)
    #            self.PTM = PtmNC(path)
    #        elif dlg.GetFilterIndex() == 1: #PTM
    #            self.flash_status_message("Opening PTM binary file: %s" % path)
    #            self.PTM = PtmBin(path)

    #        self.Nt = self.PTM.nt
    #
    #        # Update the time drop down list
    #        self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.PTM.time]
    #        self.time_list.SetItems(self.timestr)

    #        # Plot the first time step
    #        if self.__dict__.has_key('xlims'):
    #            self.PTM.plot(self.PTM.nt-1,ax=self.axes,xlims=self.xlims,\
    #            ylims=self.ylims,color=self.particlecolor,\
    #            fontcolor='w',markersize=self.particlesize)
    #        else:
    #            self.PTM.plot(self.PTM.nt-1,ax=self.axes,fontcolor='w',\
    #                color=self.particlecolor,markersize=self.particlesize)
    #        # redraw the figure
    #        self.canvas.draw()

    #
    def on_show_edges(self, event):
        if event > 0:
            self.showedges = True
        else:
            self.showedges = False

        # Update the figure
        self.update_figure()

    def on_clim_check(self, event):
        if event > 0:
            self.autoclim = False
            self.update_figure()
        else:
            self.autoclim = True

    def on_climlow(self, event):
        try:
            self.clim[0] = float(event)
        except:
            return  # do nothing

    #    self.clim[0] = event.GetString()
    #    #self.update_figure()

    def on_climhigh(self, event):
        try:
            self.clim[1] = float(event)
        except:
            return  # do nothing

    #    print event
    #    self.clim[1] = event.GetString()
    #    #self.update_figure()

    def on_select_cmap(self, event):
        self.cmap = event
        self.collection.set_cmap(self.cmap)

        # Update the figure
        self.update_figure()

    #def on_save_fig(self,event):
    #    """
    #    Save a figure of the current scene to a file
    #    """
    #    file_choices = " (*.png)|*.png| (*.pdf)|*.pdf |(*.jpg)|*.jpg |(*.eps)|*eps "
    #    filters=['.png','.pdf','.png','.png']

    #
    #    dlg = wx.FileDialog(
    #        self,
    #        message="Save figure to file...",
    #        defaultDir=os.getcwd(),
    #        defaultFile="",
    #        wildcard=file_choices,
    #        style= wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)

    #    if dlg.ShowModal() == wx.ID_OK:

    #        path = dlg.GetPath()
    #        ext = filters[dlg.GetFilterIndex()]
    #        if ext in path:
    #            outfile=path
    #        else:
    #            outfile = path+ext

    #        self.fig.savefig(outfile)

    #

    def on_save_anim(self, event):
        """
        Save an animation of the current scene to a file
        """
        #file_choices = "Quicktime (*.mov)|*.mov| (*.gif)|*.gif| (*.avi)|*.avi |(*.mp4)|*.mp4 "
        #filters=['.mov','.gif','.avi','.mp4']
        filters = "Movie formats (*.mp4 *.avi *.gif);;All files (*.*)"


        dir_ = QFileDialog.getSaveFileName(None, 'Save animation to file:',\
                '~/', filters)
        print(dir_)

        #dlg = wx.FileDialog(
        #    self,
        #    message="Output animation file...",
        #    defaultDir=os.getcwd(),
        #    defaultFile="",
        #    wildcard=file_choices,
        #    style= wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)

        if dir_ is not None:

            outfile = dir_[0]
            ext = outfile[-4::]

            # Create the animation
            #self.tstep = range(self.Nt) # Use all time steps for animation
            #self.animate(cbar=self.cbar,cmap=self.cmap,\
            #    xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
            def initanim():
                return (self.title, self.collection)
                #if not self.plot_type=='particles':
                #    return (self.title, self.collection)
                #else:
                #    return (self.PTM.title,self.PTM.p_handle)

            def updateScalar(i):
                self.tstep = [i]
                self.loadData()
                self.update_figure()
                return (self.title, self.collection)

                #if not self.plot_type=='particles':
                #    self.tstep=[i]
                #    self.loadData()
                #    self.update_figure()
                #    return (self.title,self.collection)
                #elif self.plot_type=='particles':
                #    self.PTM.plot(i,ax=self.axes,\
                #        xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
                #    return (self.PTM.title,self.PTM.p_handle)

            self.anim = animation.FuncAnimation(self.fig, \
                updateScalar, init_func = initanim, frames=self.Nt, interval=50, blit=True)

            if ext == '.gif':
                self.anim.save(outfile, writer='imagemagick', fps=6)
            elif ext == '.mp4':
                print('Saving html5 video...')
                # Ensures html5 compatibility
                self.anim.save(outfile,fps=6,\
                    writer='ffmpeg',\
                    bitrate=3600,extra_args=['-vcodec','libx264'])
                #writer='mencoder',
                #bitrate=3600,extra_args=['-ovc','x264']) # mencoder options
            else:
                self.anim.save(outfile, writer='mencoder', fps=6, bitrate=3600)

            # Return the figure back to its status
            del self.anim
            self.tstep = self.tindex
            self.loadData()
            self.update_figure()
            print('Finished saving animation to %s' % outfile)
            print(72 * '#')

            #if not self.plot_type=='particles':
            #    self.loadData()
            #    self.update_figure()

            # Bring up a dialog box
            #dlg2= wx.MessageDialog(self, 'Animation complete.', "Done", wx.OK)
            #dlg2.ShowModal()
            #dlg2.Destroy()

    #def on_exit(self, event):
    #    self.Destroy()
    #
    #def on_about(self, event):
    #    msg = """ SUNTANS NetCDF visualization tool
    #
    #        *Author: Matt Rayson
    #        *Institution: Stanford University
    #        *Created: October 2013
    #    """
    #    dlg = wx.MessageDialog(self, msg, "About", wx.OK)
    #    dlg.ShowModal()
    #    dlg.Destroy()

    #def on_count_cells(self,eveny):
    #    msg = "Total 3-D grid cells = %d"%(self.count_cells())
    #    dlg = wx.MessageDialog(self, msg, "No. cells", wx.OK)
    #    dlg.ShowModal()
    #    dlg.Destroy()

    #def on_overlay_bathy(self,event):
    #    # Plot depth contours
    #    print 'Plotting contours...'
    #    self.contourf(z=self.dv, clevs=self.depthlevs,\
    #        ax=self.axes,\
    #        filled=False, colors='0.5', linewidths=0.5, zorder=1e6)
    #    print 'Done'

    def on_plot_gridstat(self, event):
        """
        Plot the grid size histogram in a new figure
        """
        matplotlib.pyplot.figure()
        self.plothist()
        matplotlib.pyplot.show()
Esempio n. 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_types=['Line', 'Link']):
    """
    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_types : list of str or pypsa.component
        Branch types to be plotted, defaults to Line and Link.

    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)
    }

    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 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)

    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, 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 t in network.iterate_components(branch_types):
        l_defaults = defaults_for_branches[t.name]
        l_widths = line_widths.get(t.name, l_defaults['width'])
        l_nums = None
        if t.name in line_colors:
            l_colors = line_colors[t.name]

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

        if not geometry:
            segments = (np.asarray(
                ((t.df.bus0.map(x), t.df.bus0.map(y)),
                 (t.df.bus1.map(x), t.df.bus1.map(y)))).transpose(2, 0, 1))
        else:
            from shapely.wkt import loads
            from shapely.geometry import LineString
            linestrings = t.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(t.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)
Esempio n. 24
0
def draw_networkx_edges(
    G,
    pos,
    edgelist=None,
    width=1.0,
    edge_color="k",
    style="solid",
    alpha=None,
    arrowstyle="-|>",
    arrowsize=10,
    edge_cmap=None,
    edge_vmin=None,
    edge_vmax=None,
    ax=None,
    arrows=True,
    label=None,
    node_size=300,
    nodelist=None,
    node_shape="o",
    connectionstyle=None,
    min_source_margin=0,
    min_target_margin=0,
):
    """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, or array of floats
       Line width of edges (default=1.0)

    edge_color : color or array of colors (default='k')
       Edge color. Can be a single color or a sequence of colors with the same
       length as edgelist. Color can be string, or rgb (or rgba) tuple of
       floats from 0-1. 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=None)

    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.
       Note: Arrows will be the same color as edges.

    arrowstyle : str, optional (default='-|>')
       For directed graphs, choose the style of the arrow heads.
       See :py:class: `matplotlib.patches.ArrowStyle` for more
       options.

    arrowsize : int, optional (default=10)
       For directed graphs, choose the size of the arrow head head's length and
       width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
       `mutation_scale` for more info.

    connectionstyle : str, optional (default=None)
       Pass the connectionstyle parameter to create curved arc of rounding
       radius rad. For example, connectionstyle='arc3,rad=0.2'.
       See :py:class: `matplotlib.patches.ConnectionStyle` and
       :py:class: `matplotlib.patches.FancyArrowPatch` for more info.

    label : [None| string]
       Label for legend

    min_source_margin : int, optional (default=0)
       The minimum margin (gap) at the begining of the edge at the source.

    min_target_margin : int, optional (default=0)
       The minimum margin (gap) at the end of the edge at the target.

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

    list of matplotlib.patches.FancyArrowPatch
        `FancyArrowPatch` instances of the directed edges

    Depending whether the drawing includes arrows or not.

    Notes
    -----
    For directed graphs, arrows are drawn at the head end.  Arrows can be
    turned off with keyword arrows=False. Be sure to include `node_size` as a
    keyword argument; arrows are drawn considering the size of nodes.

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

    >>> G = nx.DiGraph()
    >>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
    >>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
    >>> alphas = [0.3, 0.4, 0.5]
    >>> for i, arc in enumerate(arcs):  # change alpha values of arcs
    ...     arc.set_alpha(alphas[i])

    Also see the NetworkX drawing examples at
    https://networkx.org/documentation/latest/auto_examples/index.html

    See Also
    --------
    draw()
    draw_networkx()
    draw_networkx_nodes()
    draw_networkx_labels()
    draw_networkx_edge_labels()
    """
    import matplotlib.pyplot as plt
    from matplotlib.colors import colorConverter, Colormap, Normalize
    from matplotlib.collections import LineCollection
    from matplotlib.patches import FancyArrowPatch
    import numpy as np

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

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

    if len(edgelist) == 0:  # no edges!
        if not G.is_directed() or not arrows:
            return LineCollection(None)
        else:
            return []

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

    # FancyArrowPatch handles color=None different from LineCollection
    if edge_color is None:
        edge_color = "k"

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

    # Check if edge_color is an array of floats and map to edge_cmap.
    # This is the only case handled differently from matplotlib
    if (
        np.iterable(edge_color)
        and (len(edge_color) == len(edge_pos))
        and np.alltrue([isinstance(c, Number) for c in edge_color])
    ):
        if edge_cmap is not None:
            assert isinstance(edge_cmap, Colormap)
        else:
            edge_cmap = plt.get_cmap()
        if edge_vmin is None:
            edge_vmin = min(edge_color)
        if edge_vmax is None:
            edge_vmax = max(edge_color)
        color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
        edge_color = [edge_cmap(color_normal(e)) for e in edge_color]

    if not G.is_directed() or not arrows:
        edge_collection = LineCollection(
            edge_pos,
            colors=edge_color,
            linewidths=width,
            antialiaseds=(1,),
            linestyle=style,
            transOffset=ax.transData,
            alpha=alpha,
        )

        edge_collection.set_cmap(edge_cmap)
        edge_collection.set_clim(edge_vmin, edge_vmax)

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

        return edge_collection

    arrow_collection = None

    if G.is_directed() and arrows:
        # Note: Waiting for someone to implement arrow to intersection with
        # marker.  Meanwhile, this works well for polygons with more than 4
        # sides and circle.

        def to_marker_edge(marker_size, marker):
            if marker in "s^>v<d":  # `large` markers need extra space
                return np.sqrt(2 * marker_size) / 2
            else:
                return np.sqrt(marker_size) / 2

        # Draw arrows with `matplotlib.patches.FancyarrowPatch`
        arrow_collection = []
        mutation_scale = arrowsize  # scale factor of arrow head

        # FancyArrowPatch doesn't handle color strings
        arrow_colors = colorConverter.to_rgba_array(edge_color, alpha)
        for i, (src, dst) in enumerate(edge_pos):
            x1, y1 = src
            x2, y2 = dst
            shrink_source = 0  # space from source to tail
            shrink_target = 0  # space from  head to target
            if np.iterable(node_size):  # many node sizes
                source, target = edgelist[i][:2]
                source_node_size = node_size[nodelist.index(source)]
                target_node_size = node_size[nodelist.index(target)]
                shrink_source = to_marker_edge(source_node_size, node_shape)
                shrink_target = to_marker_edge(target_node_size, node_shape)
            else:
                shrink_source = shrink_target = to_marker_edge(node_size, node_shape)

            if shrink_source < min_source_margin:
                shrink_source = min_source_margin

            if shrink_target < min_target_margin:
                shrink_target = min_target_margin

            if len(arrow_colors) == len(edge_pos):
                arrow_color = arrow_colors[i]
            elif len(arrow_colors) == 1:
                arrow_color = arrow_colors[0]
            else:  # Cycle through colors
                arrow_color = arrow_colors[i % len(arrow_colors)]

            if np.iterable(width):
                if len(width) == len(edge_pos):
                    line_width = width[i]
                else:
                    line_width = width[i % len(width)]
            else:
                line_width = width

            arrow = FancyArrowPatch(
                (x1, y1),
                (x2, y2),
                arrowstyle=arrowstyle,
                shrinkA=shrink_source,
                shrinkB=shrink_target,
                mutation_scale=mutation_scale,
                color=arrow_color,
                linewidth=line_width,
                connectionstyle=connectionstyle,
                linestyle=style,
                zorder=1,
            )  # arrows go behind nodes

            # There seems to be a bug in matplotlib to make collections of
            # FancyArrowPatch instances. Until fixed, the patches are added
            # individually to the axes instance.
            arrow_collection.append(arrow)
            ax.add_patch(arrow)

    # update view
    minx = np.amin(np.ravel(edge_pos[:, :, 0]))
    maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
    miny = np.amin(np.ravel(edge_pos[:, :, 1]))
    maxy = np.amax(np.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()

    ax.tick_params(
        axis="both",
        which="both",
        bottom=False,
        left=False,
        labelbottom=False,
        labelleft=False,
    )

    return arrow_collection
Esempio n. 25
0
    lines1.append(zip(range(maxL + 1), data[0]))
    data[1].mask[0] = False
    er = gm.entropy_rate()
    data[1] -= er
    data[1] /= er
    lines2.append(zip(range(maxL + 1), data[1, :]))

fig = plt.figure()
fig.set_figheight(5)
fig.set_figwidth(10)
fig.subplots_adjust(left=0.09, right=0.86)
ax1 = fig.add_subplot(1, 2, 1)

lc1 = LineCollection(lines1)
lc1.set_array(ps)
lc1.set_cmap(cmap)

ax1.add_collection(lc1)

ax1.set_xlabel(r"$L$ [symbols]")
ax1.set_ylabel(r"$1-(\mathbf{E}(L) / \mathbf{E})$")
ax1.axis('auto')
ax1.set_xlim((0, maxL))

ax2 = fig.add_subplot(1, 2, 2)

lc2 = LineCollection(lines2)
lc2.set_array(ps)
lc2.set_cmap(cmap)

ax2.add_collection(lc2)
Esempio n. 26
0
def draw_networkx_edges(
    G,
    pos,
    edgelist=None,
    width=1.0,
    edge_color="k",
    style="solid",
    alpha=None,
    arrowstyle="-|>",
    arrowsize=3,
    edge_cmap=None,
    edge_vmin=None,
    edge_vmax=None,
    ax=None,
    arrows=True,
    label=None,
    node_size=300,
    nodelist=None,
    node_shape="o",
    connectionstyle=None,
    min_source_margin=0,
    min_target_margin=0,
):
    """Draw the edges of the graph G. Adjusted from networkx."""
    try:
        import matplotlib.pyplot as plt
        from matplotlib.colors import colorConverter, Colormap, Normalize
        from matplotlib.collections import LineCollection
        from matplotlib.patches import FancyArrowPatch
        from numbers import Number
    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 = list(G.edges())

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

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

    # FancyArrowPatch handles color=None different from LineCollection
    if edge_color is None:
        edge_color = "k"

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

    # Check if edge_color is an array of floats and map to edge_cmap.
    # This is the only case handled differently from matplotlib
    if (
        np.iterable(edge_color)
        and (len(edge_color) == len(edge_pos))
        and np.alltrue([isinstance(c, Number) for c in edge_color])
    ):
        if edge_cmap is not None:
            assert isinstance(edge_cmap, Colormap)
        else:
            edge_cmap = plt.get_cmap()
        if edge_vmin is None:
            edge_vmin = min(edge_color)
        if edge_vmax is None:
            edge_vmax = max(edge_color)
        color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)
        edge_color = [edge_cmap(color_normal(e)) for e in edge_color]

    if not G.is_directed() or not arrows:
        edge_collection = LineCollection(
            edge_pos,
            colors=edge_color,
            linewidths=width,
            antialiaseds=(1,),
            linestyle=style,
            transOffset=ax.transData,
            alpha=alpha,
        )

        edge_collection.set_cmap(edge_cmap)
        edge_collection.set_clim(edge_vmin, edge_vmax)

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

        return edge_collection

    arrow_collection = None

    if G.is_directed() and arrows:
        # Note: Waiting for someone to implement arrow to intersection with
        # marker.  Meanwhile, this works well for polygons with more than 4
        # sides and circle.

        def to_marker_edge(marker_size, marker):
            if marker in "s^>v<d":  # `large` markers need extra space
                return np.sqrt(2 * marker_size) / 2
            else:
                return np.sqrt(marker_size) / 2

        # Draw arrows with `matplotlib.patches.FancyarrowPatch`
        arrow_collection = []
        mutation_scale = arrowsize  # scale factor of arrow head

        # FancyArrowPatch doesn't handle color strings
        arrow_colors = colorConverter.to_rgba_array(edge_color, alpha)
        for i, (src, dst) in enumerate(edge_pos):
            x1, y1 = src
            x2, y2 = dst
            shrink_source = 0  # space from source to tail
            shrink_target = 0  # space from  head to target
            if np.iterable(node_size):  # many node sizes
                source, target = edgelist[i][:2]
                source_node_size = node_size[nodelist.index(source)]
                target_node_size = node_size[nodelist.index(target)]
                shrink_source = to_marker_edge(source_node_size, node_shape)
                shrink_target = to_marker_edge(target_node_size, node_shape)
            else:
                shrink_source = shrink_target = to_marker_edge(node_size, node_shape)

            if shrink_source < min_source_margin:
                shrink_source = min_source_margin

            if shrink_target < min_target_margin:
                shrink_target = min_target_margin

            if len(arrow_colors) == len(edge_pos):
                arrow_color = arrow_colors[i]
            elif len(arrow_colors) == 1:
                arrow_color = arrow_colors[0]
            else:  # Cycle through colors
                arrow_color = arrow_colors[i % len(arrow_colors)]

            if np.iterable(width):
                if len(width) == len(edge_pos):
                    line_width = width[i]
                else:
                    line_width = width[i % len(width)]
            else:
                line_width = width

            arrow = FancyArrowPatch(
                (x1, y1),
                (x2, y2),
                arrowstyle=arrowstyle,
                shrinkA=shrink_source,
                shrinkB=shrink_target,
                mutation_scale=mutation_scale,
                color=arrow_color,
                linewidth=line_width,
                connectionstyle=connectionstyle,
                linestyle=style,
                zorder=1,
            )  # arrows go behind nodes

            # There seems to be a bug in matplotlib to make collections of
            # FancyArrowPatch instances. Until fixed, the patches are added
            # individually to the axes instance.
            arrow_collection.append(arrow)
            ax.add_patch(arrow)

    # update view
    minx = np.amin(np.ravel(edge_pos[:, :, 0]))
    maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
    miny = np.amin(np.ravel(edge_pos[:, :, 1]))
    maxy = np.amax(np.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()

    ax.tick_params(
        axis="both",
        which="both",
        bottom=False,
        left=False,
        labelbottom=False,
        labelleft=False,
    )

    return arrow_collection
Esempio n. 27
0
    def plot(self,
             time,
             beam=None,
             maxground=2000,
             maxalt=500,
             iscat=True,
             gscat=True,
             title=False,
             weighted=False,
             cmap='hot_r',
             fig=None,
             rect=111,
             ax=None,
             aax=None,
             zorder=4):
        """Plot scatter on ground/altitude profile
        
        **Args**: 
            * **time** (datetime.datetime): time of profile
            * [**beam**]: beam number
            * [**iscat**] (bool): show ionospheric scatter
            * [**gscat**] (bool): show ground scatter
            * [**maxground**]: maximum ground range [km]
            * [**maxalt**]: highest altitude limit [km]
            * [**rect**]: subplot spcification
            * [**fig**]: A pylab.figure object (default to gcf)
            * [**ax**]: Existing main axes
            * [**aax**]: Existing auxialary axes
            * [**title**]: Show default title
            * [**weighted**] (bool): plot ionospheric scatter relative strength (based on background density and range)
            * [**cmap**]: colormap used for weighted ionospheric scatter
        **Returns**:
            * **ax**: matplotlib.axes object containing formatting
            * **aax**: matplotlib.axes object containing data
            * **cbax**: matplotlib.axes object containing colorbar
        **Example**:
            ::

                # Show ionospheric scatter
                import datetime as dt
                from models import raydarn
                sTime = dt.datetime(2012, 11, 18, 5)
                rto = raydarn.RtRun(sTime, rCode='bks', beam=12)
                rto.readRays() # read rays into memory
                ax, aax, cbax = rto.rays.plot(sTime, title=True)
                rto.readScatter() # read scatter into memory
                rto.scatter.plot(sTime, ax=ax, aax=aax)
                ax.grid()
                
        written by Sebastien, 2013-04
        """
        from utils import plotUtils
        from matplotlib.collections import LineCollection
        import matplotlib.pyplot as plt
        import numpy as np

        # Set up axes
        if not ax and not aax:
            ax, aax = plotUtils.curvedEarthAxes(fig=fig,
                                                rect=rect,
                                                maxground=maxground,
                                                maxalt=maxalt)
        else:
            ax = ax
            aax = aax
            if hasattr(ax, 'beam'):
                beam = ax.beam

        # make sure that the required time and beam are present
        assert (time in self.isc.keys()
                or time in self.gsc.keys()), 'Unkown time %s' % time
        if beam:
            assert (beam in self.isc[time].keys()), 'Unkown beam %s' % beam
        else:
            beam = self.isc[time].keys()[0]

        if gscat and time in self.gsc.keys():
            for ir, (el,
                     rays) in enumerate(sorted(self.gsc[time][beam].items())):
                if len(rays['r']) == 0: continue
                _ = aax.scatter(rays['th'],
                                ax.Re * np.ones(rays['th'].shape),
                                color='0',
                                zorder=zorder)

        if iscat and time in self.isc.keys():
            if weighted:
                wmin = np.min([
                    r['w'].min() for r in self.isc[time][beam].values()
                    if r['nstp'] > 0
                ])
                wmax = np.max([
                    r['w'].max() for r in self.isc[time][beam].values()
                    if r['nstp'] > 0
                ])

            for ir, (el,
                     rays) in enumerate(sorted(self.isc[time][beam].items())):
                if rays['nstp'] == 0: continue
                t = rays['th']
                r = rays['r'] * 1e-3
                spts = np.array([t, r]).T.reshape(-1, 1, 2)
                h = rays['h'] * 1e-3
                rel = np.radians(rays['rel'])
                r = np.sqrt(r**2 + h**2 + 2 * r * h * np.sin(rel))
                t = t + np.arcsin(h / r * np.cos(rel))
                epts = np.array([t, r]).T.reshape(-1, 1, 2)
                segments = np.concatenate([spts, epts], axis=1)
                lcol = LineCollection(segments, zorder=zorder)
                if weighted:
                    _ = lcol.set_cmap(cmap)
                    _ = lcol.set_norm(plt.Normalize(0, 1))
                    _ = lcol.set_array((rays['w'] - wmin) / wmax)
                else:
                    _ = lcol.set_color('0')
                _ = aax.add_collection(lcol)

            # Plot title with date ut time and local time
            if title:
                stitle = _getTitle(time, beam, self.header, None)
                ax.set_title(stitle)

            # If weighted, plot ionospheric scatter with colormap
            if weighted:
                # Add a colorbar
                cbax = plotUtils.addColorbar(lcol, ax)
                _ = cbax.set_ylabel("Ionospheric Scatter")
            else:
                cbax = None

        ax.beam = beam
        return ax, aax, cbax
Esempio n. 28
0
def draw_networkx_edges(G, pos,
                        edgelist=None,
                        width=1.0,
                        edge_color='k',
                        style='solid',
                        alpha=1.0,
                        arrowstyle='-|>',
                        arrowsize=10,
                        edge_cmap=None,
                        edge_vmin=None,
                        edge_vmax=None,
                        ax=None,
                        arrows=True,
                        label=None,
                        node_size=300,
                        nodelist=None,
                        node_shape="o",
                        **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, or array of floats
       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.
       Note: Arrows will be the same color as edges.

    arrowstyle : str, optional (default='-|>')
       For directed graphs, choose the style of the arrow heads.
       See :py:class: `matplotlib.patches.ArrowStyle` for more
       options.

    arrowsize : int, optional (default=10)
       For directed graphs, choose the size of the arrow head head's length and
       width. See :py:class: `matplotlib.patches.FancyArrowPatch` for attribute
       `mutation_scale` for more info.

    label : [None| string]
       Label for legend

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

    list of matplotlib.patches.FancyArrowPatch
        `FancyArrowPatch` instances of the directed edges

    Depending whether the drawing includes arrows or not.

    Notes
    -----
    For directed graphs, arrows are drawn at the head end.  Arrows can be
    turned off with keyword arrows=False. Be sure to include `node_size` as a
    keyword argument; arrows are drawn considering the size of nodes.

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

    >>> G = nx.DiGraph()
    >>> G.add_edges_from([(1, 2), (1, 3), (2, 3)])
    >>> arcs = nx.draw_networkx_edges(G, pos=nx.spring_layout(G))
    >>> alphas = [0.3, 0.4, 0.5]
    >>> for i, arc in enumerate(arcs):  # change alpha values of arcs
    ...     arc.set_alpha(alphas[i])

    Also see the NetworkX drawing examples at
    https://networkx.github.io/documentation/latest/auto_examples/index.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.cbook as cb
        from matplotlib.colors import colorConverter, Colormap, Normalize
        from matplotlib.collections import LineCollection
        from matplotlib.patches import FancyArrowPatch
        import numpy as np
    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 = list(G.edges())

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

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

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

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

    if not is_string_like(edge_color) \
            and cb.iterable(edge_color) \
            and len(edge_color) == len(edge_pos):
        if np.alltrue([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 np.alltrue([not is_string_like(c) for c in edge_color]):
            # If color specs are given as (rgb) or (rgba) tuples, we're OK
            if np.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 contain color names or numbers')
    else:
        if is_string_like(edge_color) or len(edge_color) == 1:
            edge_colors = (colorConverter.to_rgba(edge_color, alpha), )
        else:
            msg = 'edge_color must be a color or list of one color per edge'
            raise ValueError(msg)

    if (not G.is_directed() or 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 by
        # 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(np.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()
        return edge_collection

    arrow_collection = None

    if G.is_directed() and arrows:
        # Note: Waiting for someone to implement arrow to intersection with
        # marker.  Meanwhile, this works well for polygons with more than 4
        # sides and circle.

        def to_marker_edge(marker_size, marker):
            if marker in "s^>v<d":  # `large` markers need extra space
                return np.sqrt(2 * marker_size) / 2
            else:
                return np.sqrt(marker_size) / 2

        # Draw arrows with `matplotlib.patches.FancyarrowPatch`
        arrow_collection = []
        mutation_scale = arrowsize  # scale factor of arrow head
        arrow_colors = edge_colors
        if arrow_colors is None:
            if edge_cmap is not None:
                assert(isinstance(edge_cmap, Colormap))
            else:
                edge_cmap = plt.get_cmap()  # default matplotlib colormap
            if edge_vmin is None:
                edge_vmin = min(edge_color)
            if edge_vmax is None:
                edge_vmax = max(edge_color)
            color_normal = Normalize(vmin=edge_vmin, vmax=edge_vmax)

        for i, (src, dst) in enumerate(edge_pos):
            x1, y1 = src
            x2, y2 = dst
            arrow_color = None
            line_width = None
            shrink_source = 0  # space from source to tail
            shrink_target = 0  # space from  head to target
            if cb.iterable(node_size):  # many node sizes
                src_node, dst_node = edgelist[i]
                index_node = nodelist.index(dst_node)
                marker_size = node_size[index_node]
                shrink_target = to_marker_edge(marker_size, node_shape)
            else:
                shrink_target = to_marker_edge(node_size, node_shape)
            if arrow_colors is None:
                arrow_color = edge_cmap(color_normal(edge_color[i]))
            elif len(arrow_colors) > 1:
                arrow_color = arrow_colors[i]
            else:
                arrow_color = arrow_colors[0]
            if len(lw) > 1:
                line_width = lw[i]
            else:
                line_width = lw[0]
            arrow = FancyArrowPatch((x1, y1), (x2, y2),
                                    arrowstyle=arrowstyle,
                                    shrinkA=shrink_source,
                                    shrinkB=shrink_target,
                                    mutation_scale=mutation_scale,
                                    color=arrow_color,
                                    linewidth=line_width,
                                    zorder=1)  # arrows go behind nodes

            # There seems to be a bug in matplotlib to make collections of
            # FancyArrowPatch instances. Until fixed, the patches are added
            # individually to the axes instance.
            arrow_collection.append(arrow)
            ax.add_patch(arrow)

    # update view
    minx = np.amin(np.ravel(edge_pos[:, :, 0]))
    maxx = np.amax(np.ravel(edge_pos[:, :, 0]))
    miny = np.amin(np.ravel(edge_pos[:, :, 1]))
    maxy = np.amax(np.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()

    return arrow_collection
Esempio n. 29
0
    def plot(self,
             time,
             beam=None,
             maxground=2000,
             maxalt=500,
             step=1,
             showrefract=False,
             nr_cmap='jet_r',
             nr_lim=[0.8, 1.],
             raycolor='0.3',
             title=False,
             zorder=2,
             alpha=1,
             fig=None,
             rect=111,
             ax=None,
             aax=None):
        """Plot ray paths
        
        **Args**: 
            * **time** (datetime.datetime): time of rays
            * [**beam**]: beam number
            * [**maxground**]: maximum ground range [km]
            * [**maxalt**]: highest altitude limit [km]
            * [**step**]: step between each plotted ray (in number of ray steps)
            * [**showrefract**]: show refractive index along ray paths (supersedes raycolor)
            * [**nr_cmap**]: color map name for refractive index coloring
            * [**nr_lim**]: refractive index plotting limits
            * [**raycolor**]: color of ray paths
            * [**rect**]: subplot spcification
            * [**fig**]: A pylab.figure object (default to gcf)
            * [**title**]: Show default title
            * [**ax**]: Existing main axes
            * [**aax**]: Existing auxialary axes
        **Returns**:
            * **ax**: matplotlib.axes object containing formatting
            * **aax**: matplotlib.axes object containing data
            * **cbax**: matplotlib.axes object containing colorbar
        **Example**:
            ::

                # Show ray paths with colored refractive index along path
                import datetime as dt
                from models import raydarn
                sTime = dt.datetime(2012, 11, 18, 5)
                rto = raydarn.RtRun(sTime, rCode='bks', beam=12, title=True)
                rto.readRays() # read rays into memory
                ax, aax, cbax = rto.rays.plot(sTime, step=10, showrefract=True, nr_lim=[.85,1])
                ax.grid()
                
        written by Sebastien, 2013-04
        """
        from utils import plotUtils
        from matplotlib.collections import LineCollection
        import matplotlib.pyplot as plt
        import numpy as np
        from types import MethodType

        # Set up axes
        if not ax and not aax:
            ax, aax = plotUtils.curvedEarthAxes(fig=fig,
                                                rect=rect,
                                                maxground=maxground,
                                                maxalt=maxalt)
        else:
            ax = ax
            aax = aax
            if hasattr(ax, 'time'):
                time = ax.time
            if hasattr(ax, 'beam'):
                beam = ax.beam

        # make sure that the required time and beam are present
        assert (time in self.paths.keys()), 'Unkown time %s' % time
        if beam:
            assert (beam in self.paths[time].keys()), 'Unkown beam %s' % beam
        else:
            beam = self.paths[time].keys()[0]

        for ir, (el,
                 rays) in enumerate(sorted(self.paths[time][beam].items())):
            if not ir % step:
                if not showrefract:
                    aax.plot(rays['th'],
                             rays['r'] * 1e-3,
                             c=raycolor,
                             zorder=zorder,
                             alpha=alpha)
                else:
                    points = np.array([rays['th'],
                                       rays['r'] * 1e-3]).T.reshape(-1, 1, 2)
                    segments = np.concatenate([points[:-1], points[1:]],
                                              axis=1)
                    lcol = LineCollection(segments, zorder=zorder, alpha=alpha)
                    _ = lcol.set_cmap(nr_cmap)
                    _ = lcol.set_norm(plt.Normalize(*nr_lim))
                    _ = lcol.set_array(rays['nr'])
                    _ = aax.add_collection(lcol)

        # Plot title with date ut time and local time
        if title:
            stitle = _getTitle(time, beam, self.header, self.name)
            ax.set_title(stitle)

        # Add a colorbar when plotting refractive index
        if showrefract:
            cbax = plotUtils.addColorbar(lcol, ax)
            _ = cbax.set_ylabel("refractive index")
        else:
            cbax = None

        # Declare a new method to show range markers
        # This method is only available after rays have been plotted
        # This ensures that the markers match the plotted rays
        def showRange(self, markers=None, color='.8', s=2, zorder=3, **kwargs):
            """Plot ray paths
            
            **Args**: 
                * [**markers**]: range markers. Defaults to every 250 km
                * All other keywords are borrowed from :func:`matplotlib.pyplot.scatter`
            **Returns**:
                * **coll**: a collection of range markers
            **Example**:
                ::

                    # Add range markers to an existing ray plot
                    ax, aax, cbax = rto.rays.plot(sTime, step=10)
                    rto.rays.showRange()
                    
            written by Sebastien, 2013-04
            """

            if not markers:
                markers = np.arange(0, 5000, 250)

            x, y = [], []
            for el, rays in self.paths[time][beam].items():
                for rm in markers:
                    inds = (rays['gran'] * 1e-3 >= rm)
                    if inds.any():
                        x.append(rays['th'][inds][0])
                        y.append(rays['r'][inds][0] * 1e-3)
            coll = aax.scatter(x, y, color=color, s=s, zorder=zorder, **kwargs)

            return coll

        # End of new method

        # Assign new method
        self.showRange = MethodType(showRange, self)

        ax.beam = beam
        return ax, aax, cbax
Esempio n. 30
0
def draw_networkx_edges(G, pos,
                        edgelist=None,
                        width=1.0,
                        edge_color='k',
                        style='solid',
                        alpha=1.0,
                        edge_cmap=None,
                        edge_vmin=None,
                        edge_vmax=None, 
                        ax=None,
                        arrows=True,
                        **kwds):
    """Draw the edges of the graph G

    This draws only the edges of the graph G.

    pos is a dictionary keyed by vertex with a two-tuple
    of x-y positions as the value.
    See networkx_v099.layout for functions that compute node positions.

    edgelist is an optional list of the edges in G to be drawn.
    If provided, only the edges in edgelist will be drawn. 

    edgecolor can be a list of matplotlib color letters such as 'k' or
    'b' that lists the color of each edge; the list must be ordered in
    the same way as the edge list. Alternatively, this list can contain
    numbers and those number are mapped to a color scale using the color
    map edge_cmap.
    
    For directed graphs, "arrows" (actually just thicker stubs) are drawn
    at the head end.  Arrows can be turned off with keyword arrows=False.

    See draw_networkx_v099 for the list of other optional parameters.

    """
    if ax is None:
        ax=matplotlib.pylab.gca()

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

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

    # set edge positions
    edge_pos=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 matplotlib.numerix.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 matplotlib.numerix.alltrue([not cb.is_string_like(c) 
                                         for c in edge_color]):
            # 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 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 = LineCollection(edge_pos,
                                colors       = edge_colors,
                                linewidths   = lw,
                                antialiaseds = (1,),
                                linestyle    = style,     
                                transOffset = ax.transData,             
                                )
    edge_collection.set_alpha(alpha)

    # need 0.87.7 or greater for edge colormaps
    mpl_version=matplotlib.__version__
    if mpl_version.endswith('svn'):
        mpl_version=matplotlib.__version__[0:-3]
    if mpl_version.endswith('pre'):
        mpl_version=matplotlib.__version__[0:-3]
    if map(int,mpl_version.split('.'))>=[0,87,7]:
        if edge_colors is None:
            if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap))
            edge_collection.set_array(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()
            matplotlib.pylab.sci(edge_collection)

#    else:
#        sys.stderr.write(\
#            """matplotlib version >= 0.87.7 required for colormapped edges.
#        (version %s detected)."""%matplotlib.__version__)
#        raise UserWarning(\
#            """matplotlib version >= 0.87.7 required for colormapped edges.
#        (version %s detected)."""%matplotlib.__version__)

    arrow_collection=None

    if G.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 = ( colorConverter.to_rgba('k', alpha), )
        a_pos=[]
        p=1.0-0.25 # make head segment 25 percent of edge length
        for src,dst in edge_pos:
            x1,y1=src
            x2,y2=dst
            dx=x2-x1 # x offset
            dy=y2-y1 # y offset
            d=sqrt(float(dx**2+dy**2)) # length of edge
            if d==0: # source and target at same position
                continue
            if dx==0: # vertical edge
                xa=x2
                ya=dy*p+y1
            if dy==0: # horizontal edge
                ya=y2
                xa=dx*p+x1
            else:
                theta=arctan2(dy,dx)
                xa=p*d*cos(theta)+x1
                ya=p*d*sin(theta)+y1
                
            a_pos.append(((xa,ya),(x2,y2)))

        arrow_collection = LineCollection(a_pos,
                                colors       = arrow_colors,
                                linewidths   = [4*ww for ww in lw],
                                antialiaseds = (1,),
                                transOffset = ax.transData,             
                                )
        
    # update view        
    minx = amin(ravel(edge_pos[:,:,0]))
    maxx = amax(ravel(edge_pos[:,:,0]))
    miny = amin(ravel(edge_pos[:,:,1]))
    maxy = amax(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()

    edge_collection.set_zorder(1) # edges go behind nodes            
    ax.add_collection(edge_collection)
    if arrow_collection:
        arrow_collection.set_zorder(1) # edges go behind nodes            
        ax.add_collection(arrow_collection)


    return edge_collection
Esempio n. 31
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, 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.
    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 = {
        '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 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)\
                        * 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):
        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[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)
Esempio n. 32
0
    def plot(self, time, beam=None, maxground=2000, maxalt=500, step=1,
        showrefract=False, nr_cmap='jet_r', nr_lim=[0.8, 1.], 
        raycolor='0.4',  
        fig=None, rect=111):
        """Plot ray paths
        
        **Args**: 
            * **time** (datetime.datetime): time of rays
            * [**beam**]: beam number
            * [**maxground**]: maximum ground range [km]
            * [**maxalt**]: highest altitude limit [km]
            * [**step**]: step between each plotted ray (in number of ray steps)
            * [**showrefract**]: show refractive index along ray paths (supersedes raycolor)
            * [**nr_cmap**]: color map name for refractive index coloring
            * [**nr_lim**]: refractive index plotting limits
            * [**raycolor**]: color of ray paths
            * [**rect**]: subplot spcification
            * [**fig**]: A pylab.figure object (default to gcf)
        **Returns**:
            * **ax**: matplotlib.axes object containing formatting
            * **aax**: matplotlib.axes object containing data
            * **cbax**: matplotlib.axes object containing colorbar
        **Example**:
            ::

                # Show ray paths with colored refractive index along path
                import datetime as dt
                from models import raydarn
                sTime = dt.datetime(2012, 11, 18, 5)
                rto = raydarn.rtRun(sTime, rCode='bks', beam=12)
                rto.readRays() # read rays into memory
                ax, aax, cbax = rto.rays.plot(sTime, step=2, showrefract=True, nr_lim=[.85,1])
                ax.grid()
                
        written by Sebastien, 2013-04
        """
        from utils import plotUtils
        from mpl_toolkits.axes_grid1 import make_axes_locatable
        from matplotlib.collections import LineCollection
        import matplotlib.pyplot as plt
        import numpy as np

        ax, aax = plotUtils.curvedEarthAxes(fig=fig, rect=rect, 
            maxground=maxground, maxalt=maxalt)

        # make sure that the required time and beam are present
        assert (time in self.paths.keys()), 'Unkown time %s' % time
        if beam:
            assert (beam in self.paths[time].keys()), 'Unkown beam %s' % beam
        else:
            beam = self.paths[time].keys()[0]
        
        for ir, (el, rays) in enumerate( sorted(self.paths[time][beam].items()) ):
            if not ir % step:
                if not showrefract:
                    aax.plot(rays['th'], rays['r']*1e-3, c=raycolor, zorder=2)
                else:
                    points = np.array([rays['th'], rays['r']*1e-3]).T.reshape(-1, 1, 2)
                    segments = np.concatenate([points[:-1], points[1:]], axis=1)
                    lcol = LineCollection( segments )
                    lcol.set_cmap( nr_cmap )
                    lcol.set_norm( plt.Normalize(*nr_lim) )
                    lcol.set_array( rays['nr'] )
                    aax.add_collection( lcol )
        # Add a colorbar when plotting refractive index
        if showrefract:
            from mpl_toolkits.axes_grid1 import SubplotDivider, LocatableAxes, Size

            fig1 = ax.get_figure()
            divider = SubplotDivider(fig1, *ax.get_geometry(), aspect=True)

            # axes for colorbar
            cbax = LocatableAxes(fig1, divider.get_position())

            h = [Size.AxesX(ax), # main axes
                 Size.Fixed(0.1), # padding
                 Size.Fixed(0.2)] # colorbar
            v = [Size.AxesY(ax)]

            divider.set_horizontal(h)
            divider.set_vertical(v)

            ax.set_axes_locator(divider.new_locator(nx=0, ny=0))
            cbax.set_axes_locator(divider.new_locator(nx=2, ny=0))

            fig1.add_axes(cbax)

            cbax.axis["left"].toggle(all=False)
            cbax.axis["top"].toggle(all=False)
            cbax.axis["bottom"].toggle(all=False)
            cbax.axis["right"].toggle(ticklabels=True, label=True)

            plt.colorbar(lcol, cax=cbax)
            cbax.set_ylabel("refractive index")

        return ax, aax, cbax
Esempio n. 33
0
def plot_linestring_collection(ax,
                               geoms,
                               colors_or_values,
                               plot_values,
                               vmin=None,
                               vmax=None,
                               cmap=None,
                               linewidth=1.0,
                               **kwargs):
    """
    Plots a collection of LineString and MultiLineString geometries to `ax`

    Parameters
    ----------

    ax : matplotlib.axes.Axes
        where shapes will be plotted

    geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be mixed)

    colors_or_values : a sequence of `N` values or RGBA tuples
        It should have 1:1 correspondence with the geometries (not their components).

    plot_values : bool
        If True, `colors_or_values` is interpreted as a list of values, and will
        be mapped to colors using vmin/vmax/cmap (which become required).
        Otherwise `colors_or_values` is interpreted as a list of colors.

    Returns
    -------

    collection : matplotlib.collections.Collection that was plotted
    """

    from matplotlib.collections import LineCollection

    components, component_colors_or_values = _flatten_multi_geoms(
        geoms, colors_or_values)

    # LineCollection does not accept some kwargs.
    if 'markersize' in kwargs:
        del kwargs['markersize']
    segments = [np.array(linestring)[:, :2] for linestring in components]
    collection = LineCollection(segments, linewidth=linewidth, **kwargs)

    if plot_values:
        collection.set_array(np.array(component_colors_or_values))
        collection.set_cmap(cmap)
        collection.set_clim(vmin, vmax)
    else:
        # set_color magically sets the correct combination of facecolor and
        # edgecolor, based on collection type.
        collection.set_color(component_colors_or_values)

        # If the user set facecolor and/or edgecolor explicitly, the previous
        # call to set_color might have overridden it (remember, the 'color' may
        # have come from plot_series, not from the user). The user should be
        # able to override matplotlib's default behavior, by setting them again
        # after set_color.
        if 'facecolor' in kwargs:
            collection.set_facecolor(kwargs['facecolor'])

        if 'edgecolor' in kwargs:
            collection.set_edgecolor(kwargs['edgecolor'])

    ax.add_collection(collection, autolim=True)
    ax.autoscale_view()
    return collection
Esempio n. 34
0
    def plot(self, time, beam=None, maxground=2000, maxalt=500,
        iscat=True, gscat=True, title=False, weighted=False, cmap='hot_r', 
        fig=None, rect=111, ax=None, aax=None, zorder=4):
        """Plot scatter on ground/altitude profile
        
        Parameters
        ----------
        time : datetime.datetime
            time of profile
        beam : Optional[ ]
            beam number
        maxground : Optional[int]
            maximum ground range [km]
        maxalt : Optional[int]
            highest altitude limit [km]
        iscat : Optional[bool]
            show ionospheric scatter
        gscat : Optional[bool]
            show ground scatter
        title : Optional[bool]
            Show default title
        weighted : Optional[bool]
            plot ionospheric scatter relative strength (based on background density and range)
        cmap : Optional[str]
            colormap used for weighted ionospheric scatter
        fig : Optional[pylab.figure]
            object (default to gcf)
        rect : Optional[int]
            subplot spcification
        ax : Optional[ ]
            Existing main axes
        aax : Optional[ ]
            Existing auxialary axes
        zorder : Optional[int]


        Returns
        -------
        ax : matplotlib.axes
            object containing formatting
        aax : matplotlib.axes
            object containing data
        cbax : matplotlib.axes
            object containing colorbar

        Example
        -------
            # Show ionospheric scatter
            import datetime as dt
            from models import raydarn
            sTime = dt.datetime(2012, 11, 18, 5)
            rto = raydarn.RtRun(sTime, rCode='bks', beam=12)
            rto.readRays() # read rays into memory
            ax, aax, cbax = rto.rays.plot(sTime, title=True)
            rto.readScatter() # read scatter into memory
            rto.scatter.plot(sTime, ax=ax, aax=aax)
            ax.grid()

        written by Sebastien, 2013-04

        """
        from davitpy.utils import plotUtils
        from matplotlib.collections import LineCollection
        import matplotlib.pyplot as plt
        import numpy as np

        # Set up axes
        if not ax and not aax:
            ax, aax = plotUtils.curvedEarthAxes(fig=fig, rect=rect, 
                maxground=maxground, maxalt=maxalt)
        else:
            ax = ax
            aax = aax
            if hasattr(ax, 'beam'):
                beam = ax.beam

        # make sure that the required time and beam are present
        assert (time in self.isc.keys() or time in self.gsc.keys()), logging.error('Unkown time %s' % time)
        if beam:
            assert (beam in self.isc[time].keys()), logging.error('Unkown beam %s' % beam)
        else:
            beam = self.isc[time].keys()[0]

        if gscat and time in self.gsc.keys():
            for ir, (el, rays) in enumerate( sorted(self.gsc[time][beam].items()) ):
                if len(rays['r']) == 0: continue
                _ = aax.scatter(rays['th'], ax.Re*np.ones(rays['th'].shape), 
                    color='0', zorder=zorder)

        if iscat and time in self.isc.keys():
            if weighted:
                wmin = np.min( [ r['w'].min() for r in self.isc[time][beam].values() if r['nstp'] > 0] )
                wmax = np.max( [ r['w'].max() for r in self.isc[time][beam].values() if r['nstp'] > 0] )

            for ir, (el, rays) in enumerate( sorted(self.isc[time][beam].items()) ):
                if rays['nstp'] == 0: continue
                t = rays['th']
                r = rays['r']*1e-3
                spts = np.array([t, r]).T.reshape(-1, 1, 2)
                h = rays['h']*1e-3
                rel = np.radians( rays['rel'] )
                r = np.sqrt( r**2 + h**2 + 2*r*h*np.sin( rel ) )
                t = t + np.arcsin( h/r * np.cos( rel ) )
                epts = np.array([t, r]).T.reshape(-1, 1, 2)
                segments = np.concatenate([spts, epts], axis=1)
                lcol = LineCollection( segments, zorder=zorder )
                if weighted:
                    _ = lcol.set_cmap( cmap )
                    _ = lcol.set_norm( plt.Normalize(0, 1) )
                    _ = lcol.set_array( ( rays['w'] - wmin ) / wmax )
                else:
                    _ = lcol.set_color('0')
                _ = aax.add_collection( lcol )

            # Plot title with date ut time and local time
            if title:
                stitle = _getTitle(time, beam, self.header, None)
                ax.set_title( stitle )

            # If weighted, plot ionospheric scatter with colormap
            if weighted:
                # Add a colorbar
                cbax = plotUtils.addColorbar(lcol, ax)
                _ = cbax.set_ylabel("Ionospheric Scatter")
            else: cbax = None

        ax.beam = beam
        return ax, aax, cbax
Esempio n. 35
0
def draw_networkx_edges(G,
                        pos,
                        edgelist=None,
                        width=1.0,
                        edge_color='k',
                        style='solid',
                        alpha=1.0,
                        edge_cmap=None,
                        edge_vmin=None,
                        edge_vmax=None,
                        ax=None,
                        arrows=True,
                        arrowstyle='thick',
                        label=None,
                        **kwds):
    try:
        import matplotlib
        import matplotlib.pyplot as plt
        import matplotlib.cbook as cb
        import matplotlib.patches as patches
        from matplotlib.colors import colorConverter, Colormap
        from matplotlib.collections import LineCollection
        from matplotlib.path import Path
        import numpy
    except ImportError:
        raise ImportError("Matplotlib required for draw()")
    except RuntimeError:
        print("Matplotlib unable to open display")
        raise
    # print "drawing_edges"

    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])
    # for e in edge_pos:
    #   print e
    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 = LineCollection(
        edge_pos,
        colors=edge_colors,
        linewidths=lw,
        antialiaseds=(1, ),
        linestyle=style,
        transOffset=ax.transData,
    )

    # print type(edge_collection)

    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-fix
        # draws arrows at each
        # waiting for someone else to implement arrows that will work
        arrow_colors = edge_colors
        a_pos = []
        p = .1  # make arrows 10% of total length
        angle = 2.7  #angle for arrows
        for src, dst in edge_pos:
            x1, y1 = src
            x2, y2 = dst
            dx = x2 - x1  # x offset
            dy = y2 - y1  # y offset
            d = numpy.sqrt(float(dx**2 + dy**2))  # length of edge
            theta = numpy.arctan2(dy, dx)
            if d == 0:  # source and target at same position
                continue
            if dx == 0:  # vertical edge
                xa = x2
                ya = dy + y1
            if dy == 0:  # horizontal edge
                ya = y2
                xa = dx + x1
            else:
                # xa = p*d*numpy.cos(theta)+x1
                # ya = p*d*numpy.sin(theta)+y1
                #corrects the endpoints to better draw
                x2 -= .04 * numpy.cos(theta)
                y2 -= .04 * numpy.sin(theta)
                lx1 = p * d * numpy.cos(theta + angle) + (x2)
                lx2 = p * d * numpy.cos(theta - angle) + (x2)
                ly1 = p * d * numpy.sin(theta + angle) + (y2)
                ly2 = p * d * numpy.sin(theta - angle) + (y2)

            a_pos.append(((lx1, ly1), (x2, y2)))
            a_pos.append(((lx2, ly2), (x2, y2)))

        arrow_collection = LineCollection(
            a_pos,
            colors=arrow_colors,
            linewidths=[1 * ww for ww in lw],
            antialiaseds=(1, ),
            transOffset=ax.transData,
        )

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

    #drawing self loops

    d = 1
    c = 0.0707
    selfedges = []
    verts = [
        (0.1 * d - 0.1 * d, 0.0),  # P0
        (c * d - 0.1 * d, c * d),  # P0
        (0.0 - 0.1 * d, 0.1 * d),  # P0
        (-c * d - 0.1 * d, c * d),  # P0
        (-0.1 * d - 0.1 * d, 0.0),  # P0
        (-c * d - 0.1 * d, -c * d),  # P0
        (0.0 - 0.1 * d, -0.1 * d),  # P0
        (c * d - 0.1 * d, -c * d),  # P0
        (0.1 * d - 0.1 * d, 0.0)
    ]
    # print verts

    codes = [
        Path.MOVETO,
        Path.LINETO,
        Path.LINETO,
        Path.LINETO,
        Path.LINETO,
        Path.LINETO,
        Path.LINETO,
        Path.LINETO,
        Path.LINETO,
    ]

    for e in edge_pos:
        if (numpy.array_equal(e[0], e[1])):
            nodes = verts[:]
            for i in range(len(nodes)):
                nodes[i] += e[0]
            # print nodes
            path = Path(nodes, codes)
            patch = patches.PathPatch(path,
                                      color=None,
                                      facecolor=None,
                                      edgecolor=edge_colors[0],
                                      fill=False,
                                      lw=4)
            ax.add_patch(patch)

    # 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)
    # print ax
    ax.update_datalim(corners)
    ax.autoscale_view()

    #    if arrow_collection:

    return edge_collection
Esempio n. 36
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_types=['Line', 'TransportLink', 'Link']):
    """
    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
        TransportLinks and 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_types : list of str or pypsa.component
        Branch types to be plotted, defaults to Line, TransportLink and Link.

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

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

    if not plt_present:
        print("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 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)

    c = pd.Series(bus_colors, index=network.buses.index)
    if c.dtype == np.dtype('O'):
        c.fillna("b", inplace=True)
    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, 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 t in network.iterate_components(branch_types):
        l_defaults = defaults_for_branches[t.name]
        l_widths = line_widths.get(t.name, l_defaults['width'])
        l_nums = None
        if t.name in line_colors:
            l_colors = line_colors[t.name]

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

        if not geometry:
            segments = (np.asarray(((t.df.bus0.map(x),
                                     t.df.bus0.map(y)),
                                    (t.df.bus1.map(x),
                                     t.df.bus1.map(y))))
                        .transpose(2, 0, 1))
        else:
            from shapely.wkt import loads
            from shapely.geometry import LineString
            linestrings = t.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(t.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)
Esempio n. 37
0
def draw_networkx_edges(G, pos,
                        edgelist=None,
                        width=1.0,
                        edge_color='k',
                        style='solid',
                        alpha=1.0,
                        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, or array of floats
       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.github.io/documentation/latest/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.cbook as cb
        from matplotlib.colors import colorConverter, Colormap
        from matplotlib.collections import LineCollection
        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 = list(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 = 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
        a_pos = []
        p = 1.0-0.25  # make head segment 25 percent of edge length
        for src, dst in edge_pos:
            x1, y1 = src
            x2, y2 = dst
            dx = x2-x1   # x offset
            dy = y2-y1   # y offset
            d = numpy.sqrt(float(dx**2 + dy**2))  # length of edge
            if d == 0:   # source and target at same position
                continue
            if dx == 0:  # vertical edge
                xa = x2
                ya = dy*p+y1
            if dy == 0:  # horizontal edge
                ya = y2
                xa = dx*p+x1
            else:
                theta = numpy.arctan2(dy, dx)
                xa = p*d*numpy.cos(theta)+x1
                ya = p*d*numpy.sin(theta)+y1

            a_pos.append(((xa, ya), (x2, y2)))

        arrow_collection = LineCollection(a_pos,
                                colors=arrow_colors,
                                linewidths=[4*ww for ww in lw],
                                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 edge_collection
Esempio n. 38
0
def plot_linestring_collection(ax,
                               geoms,
                               values=None,
                               color=None,
                               cmap=None,
                               vmin=None,
                               vmax=None,
                               **kwargs):
    """
    Plots a collection of LineString and MultiLineString geometries to `ax`

    Parameters
    ----------

    ax : matplotlib.axes.Axes
        where shapes will be plotted

    geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be
            mixed)

    values : a sequence of `N` values, optional
        Values will be mapped to colors using vmin/vmax/cmap. They should
        have 1:1 correspondence with the geometries (not their components).

    color : single color or sequence of `N` colors
        Cannot be used together with `values`.

    Returns
    -------

    collection : matplotlib.collections.Collection that was plotted

    """
    from matplotlib.collections import LineCollection

    geoms, multiindex = _flatten_multi_geoms(geoms, range(len(geoms)))
    if values is not None:
        values = np.take(values, multiindex)

    # LineCollection does not accept some kwargs.
    if "markersize" in kwargs:
        del kwargs["markersize"]

    # color=None gives black instead of default color cycle
    if color is not None:
        if pd.api.types.is_list_like(color):
            kwargs["color"] = np.take(color, multiindex)
        else:
            kwargs["color"] = color

    segments = [np.array(linestring)[:, :2] for linestring in geoms]
    collection = LineCollection(segments, **kwargs)

    if values is not None:
        collection.set_array(np.asarray(values))
        collection.set_cmap(cmap)
        if "norm" not in kwargs:
            collection.set_clim(vmin, vmax)

    ax.add_collection(collection, autolim=True)
    ax.autoscale_view()
    return collection
Esempio n. 39
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,
         layouter=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 cartopy/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". If bus_sizes is a pandas.Series
        with a Multiindex, bus_colors defaults to the network.carriers['color']
        column.
    bus_sizes : dict/pandas.Series
        Sizes of bus points, defaults to 10. 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. network.generators.groupby(['bus', 'carrier']).p_nom.sum()
    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.
    layouter : networkx.drawing.layout function, default None
        Layouting function from `networkx <https://networkx.github.io/>`_ which
        overrules coordinates given in ``network.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
    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 = _get_coordinates(network, layouter=layouter)

    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"
        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 = network.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 network.carrier.color")

        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)
Esempio n. 40
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)
Esempio n. 41
0
    def run(self):
        if self.data is None:
            self.parent.threadPlot = None
            return

        peakF = None
        peakL = None
        total = len(self.data)
        if total > 0:
            self.parent.clear_plots()
            lc = None
            if self.average:
                avg = OrderedDict()
                count = len(self.data)

                for timeStamp in self.data:

                    if len(self.data[timeStamp]) < 2:
                        return

                    for x, y in self.data[timeStamp].items():
                        if x in avg:
                            avg[x] = (avg[x] + y) / 2
                        else:
                            avg[x] = y

                data = avg.items()
                peakF, peakL = max(data, key=lambda item: item[1])

                segments, levels = self.create_segments(data)
                lc = LineCollection(segments)
                lc.set_array(numpy.array(levels))
                lc.set_norm(self.get_norm(self.autoL, self.extent))
                lc.set_cmap(self.colourMap)
                lc.set_linewidth(self.lineWidth)
                lc.set_gid('plot')
                self.axes.add_collection(lc)
                self.parent.lc = lc
            else:
                count = 0.0
                for timeStamp in self.data:

                    if len(self.data[timeStamp]) < 2:
                        self.parent.threadPlot = None
                        return

                    if self.fade:
                        alpha = (total - count) / total
                    else:
                        alpha = 1

                    data = self.data[timeStamp].items()
                    peakF, peakL = self.extent.get_peak_fl()

                    segments, levels = self.create_segments(data)
                    lc = LineCollection(segments)
                    lc.set_array(numpy.array(levels))
                    lc.set_norm(self.get_norm(self.autoL, self.extent))
                    lc.set_cmap(self.colourMap)
                    lc.set_linewidth(self.lineWidth)
                    lc.set_gid('plot')
                    lc.set_alpha(alpha)
                    self.axes.add_collection(lc)
                    count += 1

            if self.annotate:
                self.annotate_plot(peakF, peakL)

        if total > 0:
            self.parent.scale_plot()
            self.parent.redraw_plot()

        self.parent.threadPlot = None
Esempio n. 42
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,
                        **kwds):
    """Draw the edges of the graph G

    This draws only the edges of the graph G.

    pos is a dictionary keyed by vertex with a two-tuple
    of x-y positions as the value.
    See networkx.layout for functions that compute node positions.

    edgelist is an optional list of the edges in G to be drawn.
    If provided, only the edges in edgelist will be drawn. 

    edgecolor can be a list of matplotlib color letters such as 'k' or
    'b' that lists the color of each edge; the list must be ordered in
    the same way as the edge list. Alternatively, this list can contain
    numbers and those number are mapped to a color scale using the color
    map edge_cmap.  Finally, it can also be a list of (r,g,b) or (r,g,b,a)
    tuples, in which case these will be used directly to color the edges.  If
    the latter mode is used, you should not provide a value for alpha, as it
    would be applied globally to all lines.
    
    For directed graphs, "arrows" (actually just thicker stubs) are drawn
    at the head end.  Arrows can be turned off with keyword arrows=False.

    See draw_networkx for the list of other optional parameters.

    """
    try:
        import matplotlib
        import matplotlib.pylab as pylab
        import numpy as np
        from matplotlib.colors import colorConverter,Colormap
        from matplotlib.collections import LineCollection
    except ImportError:
        raise ImportError("Matplotlib required for draw()")
    except RuntimeError:
        pass # unable to open display

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

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

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

    # set edge positions
    edge_pos=np.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 np.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 np.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 np.alltrue([cb.iterable(c) and len(c) in (3,4)
                             for c in edge_color]):
                edge_colors = tuple(edge_color)
                alpha=None
            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 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 = LineCollection(edge_pos,
                                     colors       = edge_colors,
                                     linewidths   = lw,
                                     antialiaseds = (1,),
                                     linestyle    = style,     
                                     transOffset = ax.transData,             
                                     )

    # 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)

    # need 0.87.7 or greater for edge colormaps.  No checks done, this will
    # just not work with an older mpl
    if edge_colors is None:
        if edge_cmap is not None: assert(isinstance(edge_cmap, Colormap))
        edge_collection.set_array(np.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()
        pylab.sci(edge_collection)

    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 = ( colorConverter.to_rgba('k', alpha), )
        a_pos=[]
        p=1.0-0.25 # make head segment 25 percent of edge length
        for src,dst in edge_pos:
            x1,y1=src
            x2,y2=dst
            dx=x2-x1 # x offset
            dy=y2-y1 # y offset
            d=np.sqrt(float(dx**2+dy**2)) # length of edge
            if d==0: # source and target at same position
                continue
            if dx==0: # vertical edge
                xa=x2
                ya=dy*p+y1
            if dy==0: # horizontal edge
                ya=y2
                xa=dx*p+x1
            else:
                theta=np.arctan2(dy,dx)
                xa=p*d*np.cos(theta)+x1
                ya=p*d*np.sin(theta)+y1
                
            a_pos.append(((xa,ya),(x2,y2)))

        arrow_collection = LineCollection(a_pos,
                                colors       = arrow_colors,
                                linewidths   = [4*ww for ww in lw],
                                antialiaseds = (1,),
                                transOffset = ax.transData,             
                                )
        
    # update view        
    minx = np.amin(np.ravel(edge_pos[:,:,0]))
    maxx = np.amax(np.ravel(edge_pos[:,:,0]))
    miny = np.amin(np.ravel(edge_pos[:,:,1]))
    maxy = np.amax(np.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()

    edge_collection.set_zorder(1) # edges go behind nodes            
    ax.add_collection(edge_collection)
    if arrow_collection:
        arrow_collection.set_zorder(1) # edges go behind nodes            
        ax.add_collection(arrow_collection)
        
    return edge_collection
Esempio n. 43
0
class SunPlotPy(wx.Frame, Spatial, Grid ):
    """ 
    The main frame of the application
    """
    title = 'sunplot(py)'

    # Plotting options
    autoclim=True
    showedges=False
    bgcolor='k'
    textcolor='w'
    cmap='RdBu'
    particlesize = 1.8
    particlecolor = 'm'

    # other flags
    collectiontype='cells'
    oldcollectiontype='cells'

    # 
    tindex=0 
    depthlevs = [0., 10., 100., 200., 300., 400., 500.,\
        1000.,2000.,3000.,4000.,5000]

    _FillValue=999999
    
    def __init__(self):
        wx.Frame.__init__(self, None, -1, self.title)
        
        self.create_menu()
        self.create_status_bar()
        self.create_main_panel()
        
        #self.draw_figure()

    def create_menu(self):
        self.menubar = wx.MenuBar()
        
        ###
        # File Menu
        ###
        menu_file = wx.Menu()
        # Load a hydro output file
        m_expt = menu_file.Append(-1, "&Open file\tCtrl-O", "Open netcdf file")
        self.Bind(wx.EVT_MENU, self.on_open_file, m_expt)

        # Load a grid file
        m_grid = menu_file.Append(-1, "&Load grid\tCtrl-G", "Load SUNTANS grid from folder")
        self.Bind(wx.EVT_MENU, self.on_load_grid, m_grid)

        # Load a particle file
        m_part = menu_file.Append(-1, "&Load PTM file\tCtrl-Shift-P", "Load a PTM file")
        self.Bind(wx.EVT_MENU, self.on_load_ptm, m_part)

        # Save current scene as an animation
        m_anim = menu_file.Append(-1,"&Save animation of current scene\tCtrl-S","Save animation")
        self.Bind(wx.EVT_MENU, self.on_save_anim, m_anim)

        # Save the current figure
        m_prin = menu_file.Append(-1,"&Print current scene\tCtrl-P","Save figure")
        self.Bind(wx.EVT_MENU, self.on_save_fig, m_prin)



        menu_file.AppendSeparator()
        # Exit
        m_exit = menu_file.Append(-1, "E&xit\tCtrl-X", "Exit")
        self.Bind(wx.EVT_MENU, self.on_exit, m_exit)

        ###
        # Tools menu
        ###
        menu_tools = wx.Menu()
        m_gridstat = menu_tools.Append(-1, "&Plot grid size statistics", "SUNTANS grid size")
        self.Bind(wx.EVT_MENU, self.on_plot_gridstat, m_gridstat)

        m_countcells = menu_tools.Append(-1, "&Count # grid cells", "Grid cell count")
        self.Bind(wx.EVT_MENU, self.on_count_cells, m_countcells)

        m_overlaybathy = menu_tools.Append(-1, "&Overlay depth contours", "Depth overlay")
        self.Bind(wx.EVT_MENU, self.on_overlay_bathy, m_overlaybathy)

        
        ###
        # Help Menu
        ###
        menu_help = wx.Menu()
        m_about = menu_help.Append(-1, "&About\tF1", "About the demo")
        self.Bind(wx.EVT_MENU, self.on_about, m_about)
        
        
        # Add all of the menu bars
        self.menubar.Append(menu_file, "&File")
        self.menubar.Append(menu_tools, "&Tools")
        self.menubar.Append(menu_help, "&Help")
        self.SetMenuBar(self.menubar)

    def create_main_panel(self):
        """ Creates the main panel with all the controls on it:
             * mpl canvas 
             * mpl navigation toolbar
             * Control panel for interaction
        """
        self.panel = wx.Panel(self)
        
        # Create the mpl Figure and FigCanvas objects. 
        # 5x4 inches, 100 dots-per-inch
        #
        self.dpi = 100
        #self.fig = Figure((7.0, 6.0), dpi=self.dpi,facecolor=self.bgcolor)
        self.fig = Figure((7.0, 6.0), dpi=self.dpi)
        self.canvas = FigCanvas(self.panel, -1, self.fig)
        
        
        # Since we have only one plot, we can use add_axes 
        # instead of add_subplot, but then the subplot
        # configuration tool in the navigation toolbar wouldn't
        # work.
        #
        self.axes = self.fig.add_subplot(111)
        #SetAxColor(self.axes,self.textcolor,self.bgcolor)
        
        # Bind the 'pick' event for clicking on one of the bars
        #
        #self.canvas.mpl_connect('pick_event', self.on_pick)
        
        ########
        # Create widgets
        ########
        self.variable_list = wx.ComboBox(
            self.panel, 
            size=(200,-1),
            choices=['Select a variable...'],
            style=wx.CB_READONLY)
        self.variable_list.Bind(wx.EVT_COMBOBOX, self.on_select_variable)
        
        self.time_list = wx.ComboBox(
            self.panel, 
            size=(200,-1),
            choices=['Select a time step...'],
            style=wx.CB_READONLY)
        self.time_list.Bind(wx.EVT_COMBOBOX, self.on_select_time)

        self.depthlayer_list = wx.ComboBox(
            self.panel, 
            size=(200,-1),
            choices=['Select a vertical layer...'],
            style=wx.CB_READONLY)
        self.depthlayer_list.Bind(wx.EVT_COMBOBOX, self.on_select_depth)

        self.show_edge_check = wx.CheckBox(self.panel, -1, 
            "Show Edges",
            style=wx.ALIGN_RIGHT)
        self.show_edge_check.Bind(wx.EVT_CHECKBOX, self.on_show_edges)

        if USECMOCEAN:
            cmaps=[]
            for cmap in cm.cmapnames:
                cmaps.append(cmap)
                cmaps.append(cmap+'_r') # Add all reverse map options
        else:
            # Use matplotlib standard
            cmaps = matplotlib.cm.datad.keys()

        cmaps.sort()
        self.colormap_list = wx.ComboBox(
            self.panel, 
            size=(100,-1),
            choices=cmaps,
            style=wx.CB_READONLY)
        self.colormap_list.Bind(wx.EVT_COMBOBOX, self.on_select_cmap)
        self.colormap_label = wx.StaticText(self.panel, -1,"Colormap:")

        self.clim_check = wx.CheckBox(self.panel, -1, 
            "Manual color limits ",
            style=wx.ALIGN_RIGHT)
        self.clim_check.Bind(wx.EVT_CHECKBOX, self.on_clim_check)

        self.climlow = wx.TextCtrl(
            self.panel, 
            size=(100,-1),
            style=wx.TE_PROCESS_ENTER)
        self.climlow.Bind(wx.EVT_TEXT_ENTER, self.on_climlow)
        
        self.climhigh = wx.TextCtrl(
            self.panel, 
            size=(100,-1),
            style=wx.TE_PROCESS_ENTER)
        self.climhigh.Bind(wx.EVT_TEXT_ENTER, self.on_climhigh)
 


        # Labels
        self.variable_label = wx.StaticText(self.panel, -1,"Variable:",size=(200,-1))
        self.time_label = wx.StaticText(self.panel, -1,"Time step:",size=(200,-1))
        self.depth_label = wx.StaticText(self.panel, -1,"Vertical level:",size=(200,-1))


        # Create the navigation toolbar, tied to the canvas
        #
        self.toolbar = NavigationToolbar(self.canvas)
        #self.toolbar.toolitems[8][3]='my_save_fig'

        #def my_save_fig(self,*args):
        #    print 'saving figure'
        #    return "break"

        
        #########
        # Layout with box sizers
        #########
        
        self.vbox = wx.BoxSizer(wx.VERTICAL)
        self.vbox.Add(self.canvas, 1, wx.LEFT | wx.TOP | wx.GROW)
        self.vbox.Add(self.toolbar, 0, wx.EXPAND)

        self.vbox.AddSpacer(10)
        #self.vbox.Add((-1,25))

        flags = wx.ALIGN_LEFT | wx.ALL | wx.ALIGN_CENTER_VERTICAL

        self.hbox0 = wx.BoxSizer(wx.HORIZONTAL)
        self.hbox0.Add(self.show_edge_check, 0, border=10, flag=flags)
        self.hbox0.Add(self.colormap_label, 0, border=10, flag=flags)
        self.hbox0.Add(self.colormap_list, 0, border=10, flag=flags)
        self.hbox0.Add(self.clim_check, 0, border=10, flag=flags)
        self.hbox0.Add(self.climlow, 0, border=10, flag=flags)
        self.hbox0.Add(self.climhigh, 0, border=10, flag=flags)

        self.vbox.AddSpacer(5)
        self.hbox1 = wx.BoxSizer(wx.HORIZONTAL)
        self.hbox1.Add(self.variable_label, 0, border=10, flag=flags)
        self.hbox1.Add(self.time_label, 0, border=10, flag=flags)
        self.hbox1.Add(self.depth_label, 0, border=10, flag=flags)

        self.vbox.AddSpacer(5)
        self.hbox2 = wx.BoxSizer(wx.HORIZONTAL)
        self.hbox2.Add(self.variable_list, 0, border=10, flag=flags)
        self.hbox2.Add(self.time_list, 0, border=10, flag=flags)
        self.hbox2.Add(self.depthlayer_list, 0, border=10, flag=flags)
       
        self.vbox.Add(self.hbox1, 0, flag = wx.ALIGN_LEFT | wx.TOP)
        self.vbox.Add(self.hbox2, 0, flag = wx.ALIGN_LEFT | wx.TOP)
        self.vbox.Add(self.hbox0, 0, flag = wx.ALIGN_LEFT | wx.TOP)
        
        self.panel.SetSizer(self.vbox)
        self.vbox.Fit(self)
    
    ##########
    # Event functions
    ##########

    def create_figure(self):
        """ 
        Creates the figure
        """
        # Find the colorbar limits if unspecified
        if self.autoclim:
            self.clim = [self.data.min(),self.data.max()]
            self.climlow.SetValue('%3.1f'%self.clim[0])
            self.climhigh.SetValue('%3.1f'%self.clim[1])
         
        if self.__dict__.has_key('collection'):
            #self.collection.remove()
            self.axes.collections.remove(self.collection)
        else:
            # First call - set the axes limits
            self.axes.set_aspect('equal')
            self.axes.set_xlim(self.xlims)
            self.axes.set_ylim(self.ylims)
 

        if self.collectiontype=='cells':
            self.collection = PolyCollection(self.xy,cmap=self.cmap)
            self.collection.set_array(np.array(self.data[:]))
            if not self.showedges:
                self.collection.set_edgecolors(self.collection.to_rgba(np.array((self.data[:])))) 
        elif self.collectiontype=='edges':
            xylines = [self.xp[self.edges],self.yp[self.edges]]
            linesc = [zip(xylines[0][ii,:],xylines[1][ii,:]) for ii in range(self.Ne)]
            self.collection = LineCollection(linesc,array=np.array(self.data[:]),cmap=self.cmap)

        self.collection.set_clim(vmin=self.clim[0],vmax=self.clim[1])

        self.axes.add_collection(self.collection)    
        self.title=self.axes.set_title(self.genTitle(),color=self.textcolor)
        self.axes.set_xlabel('Easting [m]')
        self.axes.set_ylabel('Northing [m]')

        # create a colorbar

        if not self.__dict__.has_key('cbar'):
            self.cbar = self.fig.colorbar(self.collection)
            #SetAxColor(self.cbar.ax.axes,self.textcolor,self.bgcolor)
        else:
            #pass
            print 'Updating colorbar...'
            #self.cbar.check_update(self.collection)
            self.cbar.on_mappable_changed(self.collection)

        self.canvas.draw()
   
    def update_figure(self):
        if self.autoclim:
            self.clim = [self.data.min(),self.data.max()]
            self.climlow.SetValue('%3.1f'%self.clim[0])
            self.climhigh.SetValue('%3.1f'%self.clim[1])
        else:
            self.clim = [float(self.climlow.GetValue()),\
                float(self.climhigh.GetValue())]
 
        # check whether it is cell or edge type
        if self.hasDim(self.variable,self.griddims['Ne']):
            self.collectiontype='edges'
        elif self.hasDim(self.variable,self.griddims['Nc']):
            self.collectiontype='cells'

        # Create a new figure if the variable has gone from cell to edge of vice
        # versa
        if not self.collectiontype==self.oldcollectiontype:
            self.create_figure()
            self.oldcollectiontype=self.collectiontype

        self.collection.set_array(np.array(self.data[:]))
        self.collection.set_clim(vmin=self.clim[0],vmax=self.clim[1])

        # Cells only
        if self.collectiontype=='cells':
            if not self.showedges:
                self.collection.set_edgecolors(self.collection.to_rgba(np.array((self.data[:])))) 
            else:
                self.collection.set_edgecolors('k')
                self.collection.set_linewidths(0.2)

        # Update the title
        self.title=self.axes.set_title(self.genTitle(),color=self.textcolor)

        #Update the colorbar
        self.cbar.update_normal(self.collection)

        # redraw the figure
        self.canvas.draw()
    
    def on_pick(self, event):
        # The event received here is of the type
        # matplotlib.backend_bases.PickEvent
        #
        # It carries lots of information, of which we're using
        # only a small amount here.
        # 
        box_points = event.artist.get_bbox().get_points()
        msg = "You've clicked on a bar with coords:\n %s" % box_points
        
        dlg = wx.MessageDialog(
            self, 
            msg, 
            "Click!",
            wx.OK | wx.ICON_INFORMATION)

        dlg.ShowModal() 
        dlg.Destroy()        
    
    def on_select_variable(self, event):
        vname = event.GetString()
        self.flash_status_message("Selecting variable: %s"%vname)
        # update the spatial object and load the data
        self.variable = vname
        self.loadData(variable=self.variable)

        # Check if the variable has a depth coordinate
        depthstr = ['']
        # If so populate the vertical layer box
        if self.hasDim(self.variable,self.griddims['Nk']):
            depthstr = ['%3.1f'%self.z_r[k] for k in range(self.Nkmax)]
            depthstr += ['surface','seabed']

        elif self.hasDim(self.variable,'Nkw'):
            depthstr = ['%3.1f'%self.z_w[k] for k in range(self.Nkmax+1)]

        self.depthlayer_list.SetItems(depthstr)

        # Update the plot
        self.update_figure()



    def on_select_time(self, event):
        self.tindex = event.GetSelection()
        # Update the object time index and reload the data
        if self.plot_type=='hydro':
            if not self.tstep==self.tindex:
                self.tstep=self.tindex
                self.loadData()
                self.flash_status_message("Selecting variable: %s..."%event.GetString())

                # Update the plot
                self.update_figure()
        elif self.plot_type=='particles':
            self.PTM.plot(self.tindex,ax=self.axes,\
                xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
        
            self.canvas.draw()


    def on_select_depth(self, event):
        kindex = event.GetSelection()
        if not self.klayer[0]==kindex:
            # Check if its the seabed or surface value
            if kindex>=self.Nkmax:
                kindex=event.GetString()
            self.klayer = [kindex]
            self.loadData()       
            self.flash_status_message("Selecting depth: %s..."%event.GetString())

            # Update the plot
            self.update_figure()

    def on_open_file(self, event):
        file_choices = "SUNTANS NetCDF (*.nc)|*.nc*|UnTRIM NetCDF (*.nc)|*.nc*|All Files (*.*)|*.*"
        
        dlg = wx.FileDialog(
            self, 
            message="Open SUNTANS file...",
            defaultDir=os.getcwd(),
            defaultFile="",
            wildcard=file_choices,
            style= wx.FD_MULTIPLE)
        
        if dlg.ShowModal() == wx.ID_OK:
            self.plot_type='hydro'

            path = dlg.GetPaths()

            # Initialise the class
            if dlg.GetFilterIndex() == 0 or dlg.GetFilterIndex() > 1: #SUNTANS
                self.flash_status_message("Opening SUNTANS file: %s" % path)
		try:
		    Spatial.__init__(self, path, _FillValue=self._FillValue)
		except:
		    Spatial.__init__(self, path, _FillValue=-999999)
                startvar='dv'
            if dlg.GetFilterIndex()==1: #UnTRIM
                self.flash_status_message("Opening UnTRIMS file: %s" % path)
                #Spatial.__init__(self,path,gridvars=untrim_gridvars,griddims=untrim_griddims)
                UNTRIMSpatial.__init__(self,path)
                startvar='Mesh2_face_depth'
            
            # Populate the drop down menus
            vnames = self.listCoordVars()
            self.variable_list.SetItems(vnames)
            
            # Update the time drop down list
            if self.__dict__.has_key('time'):
                self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.time]
            else:
                # Assume that it is a harmonic-type file
                self.timestr = self.nc.Constituent_Names.split()

            self.time_list.SetItems(self.timestr)

            # Draw the depth
            if startvar in vnames:
                self.variable=startvar
                self.loadData()
                self.create_figure()

    def on_load_grid(self, event):
        
        dlg = wx.DirDialog(
            self, 
            message="Open SUNTANS grid from folder...",
            defaultPath=os.getcwd(),
            style= wx.DD_DEFAULT_STYLE)
        
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()

            # Initialise the class
            self.flash_status_message("Opening SUNTANS grid from folder: %s" % path)
            Grid.__init__(self,path)

            # Plot the Grid
            if self.__dict__.has_key('collection'):
                self.axes.collections.remove(self.collection)

            self.axes,self.collection = self.plotmesh(ax=self.axes,edgecolors='y')

            # redraw the figure
            self.canvas.draw()

    def on_load_ptm(self, event):
        file_choices = "PTM NetCDF (*.nc)|*.nc|PTM Binary (*_bin.out)|*_bin.out|All Files (*.*)|*.*"
        
        dlg = wx.FileDialog(
            self, 
            message="Open PTM file...",
            defaultDir=os.getcwd(),
            defaultFile="",
            wildcard=file_choices,
            style= wx.FD_MULTIPLE)
        
        if dlg.ShowModal() == wx.ID_OK:
            self.plot_type = 'particles'
            path = dlg.GetPath()

            # Initialise the class
            if dlg.GetFilterIndex() == 0: #SUNTANS
                self.flash_status_message("Opening PTM netcdf file: %s" % path)
                self.PTM = PtmNC(path)
            elif dlg.GetFilterIndex() == 1: #PTM
                self.flash_status_message("Opening PTM binary file: %s" % path)
                self.PTM = PtmBin(path)

            self.Nt = self.PTM.nt
            
            # Update the time drop down list
            self.timestr = [datetime.strftime(tt,'%d-%b-%Y %H:%M:%S') for tt in self.PTM.time]
            self.time_list.SetItems(self.timestr)

            # Plot the first time step
            if self.__dict__.has_key('xlims'):
                self.PTM.plot(self.PTM.nt-1,ax=self.axes,xlims=self.xlims,\
                ylims=self.ylims,color=self.particlecolor,\
                fontcolor='w',markersize=self.particlesize)
            else:
                self.PTM.plot(self.PTM.nt-1,ax=self.axes,fontcolor='w',\
                    color=self.particlecolor,markersize=self.particlesize)
            # redraw the figure
            self.canvas.draw()

        
    def on_show_edges(self,event):
        sender=event.GetEventObject()
        self.showedges = sender.GetValue()

        # Update the figure
        self.update_figure()

    def on_clim_check(self,event):
        sender=event.GetEventObject()
        if sender.GetValue() == True:
            self.autoclim=False
            self.update_figure()
        else:
            self.autoclim=True
       

    def on_climlow(self,event):
        self.clim[0] = event.GetString()
        #self.update_figure()

    def on_climhigh(self,event):
        self.clim[1] = event.GetString()
        #self.update_figure()

    def on_select_cmap(self,event):
        self.cmap=event.GetString()
        if USECMOCEAN:
            self.collection.set_cmap(getattr(cm,self.cmap))
        else:
            self.collection.set_cmap(self.cmap)

        # Update the figure
        self.update_figure()

    def on_save_fig(self,event):
        """
        Save a figure of the current scene to a file
        """
        file_choices = " (*.png)|*.png| (*.pdf)|*.pdf |(*.jpg)|*.jpg |(*.eps)|*eps "
        filters=['.png','.pdf','.png','.png']

        
        dlg = wx.FileDialog(
            self, 
            message="Save figure to file...",
            defaultDir=os.getcwd(),
            defaultFile="",
            wildcard=file_choices,
            style= wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)

        if dlg.ShowModal() == wx.ID_OK:

            path = dlg.GetPath()
            ext = filters[dlg.GetFilterIndex()]
            if ext in path:
                outfile=path
            else:
                outfile = path+ext

            self.fig.savefig(outfile)

            


    def on_save_anim(self,event):
        """
        Save an animation of the current scene to a file
        """
        file_choices = "Quicktime (*.mov)|*.mov| (*.gif)|*.gif| (*.avi)|*.avi |(*.mp4)|*.mp4 "
        filters=['.mov','.gif','.avi','.mp4']

        
        dlg = wx.FileDialog(
            self, 
            message="Output animation file...",
            defaultDir=os.getcwd(),
            defaultFile="",
            wildcard=file_choices,
            style= wx.FD_SAVE | wx.FD_OVERWRITE_PROMPT)

        if dlg.ShowModal() == wx.ID_OK:

            path = dlg.GetPath()
            ext = filters[dlg.GetFilterIndex()]
            if ext in path:
                outfile=path
            else:
                outfile = path+ext
            self.flash_status_message("Saving figure to file: %s" %outfile)
            self.flash_status_message("Saving animation to file: %s" %outfile)

            # Create the animation
            #self.tstep = range(self.Nt) # Use all time steps for animation
            #self.animate(cbar=self.cbar,cmap=self.cmap,\
            #    xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
            def initanim():
                if not self.plot_type=='particles':
                    return (self.title, self.collection)
                else:
                    return (self.PTM.title,self.PTM.p_handle)

            def updateScalar(i):
                if not self.plot_type=='particles':
                    self.tstep=[i]
                    self.loadData()
                    self.update_figure()
                    return (self.title,self.collection)
                elif self.plot_type=='particles':
                    self.PTM.plot(i,ax=self.axes,\
                        xlims=self.axes.get_xlim(),ylims=self.axes.get_ylim())
                    return (self.PTM.title,self.PTM.p_handle)

            self.anim = animation.FuncAnimation(self.fig, \
                updateScalar, init_func = initanim, frames=self.Nt, interval=50, blit=True)

            if ext=='.gif':
                self.anim.save(outfile,writer='imagemagick',fps=6)
            elif ext=='.mp4':
                print 'Saving html5 video...'
                # Ensures html5 compatibility
                self.anim.save(outfile,writer='mencoder',fps=6,\
                    bitrate=3600,extra_args=['-ovc','x264']) # mencoder options
                    #bitrate=3600,extra_args=['-vcodec','libx264'])
            else:
                self.anim.save(outfile,writer='mencoder',fps=6,bitrate=3600)

            # Return the figure back to its status
            del self.anim
            self.tstep=self.tindex
            if not self.plot_type=='particles':
                self.loadData()
                self.update_figure()

            # Bring up a dialog box
            dlg2= wx.MessageDialog(self, 'Animation complete.', "Done", wx.OK)
            dlg2.ShowModal()
            dlg2.Destroy()

    def on_exit(self, event):
        self.Destroy()
        
    def on_about(self, event):
        msg = """ SUNTANS NetCDF visualization tool
        
            *Author: Matt Rayson
            *Institution: Stanford University
            *Created: October 2013
        """
        dlg = wx.MessageDialog(self, msg, "About", wx.OK)
        dlg.ShowModal()
        dlg.Destroy()

    def on_count_cells(self,eveny):
        msg = "Total 3-D grid cells = %d"%(self.count_cells())
        dlg = wx.MessageDialog(self, msg, "No. cells", wx.OK)
        dlg.ShowModal()
        dlg.Destroy()

    def on_overlay_bathy(self,event):
        # Plot depth contours
        print 'Plotting contours...'
        self.contourf(z=self.dv, clevs=self.depthlevs,\
            ax=self.axes,\
            filled=False, colors='0.5', linewidths=0.5, zorder=1e6)
        print 'Done'
   
    def on_plot_gridstat(self, event):
        """
        Plot the grid size histogram in a new figure
        """
        matplotlib.pyplot.figure()
        self.plothist()
        matplotlib.pyplot.show()


    def create_status_bar(self):
        self.statusbar = self.CreateStatusBar()

    def flash_status_message(self, msg, flash_len_ms=1500):
        self.statusbar.SetStatusText(msg)
        self.timeroff = wx.Timer(self)
        self.Bind(
            wx.EVT_TIMER, 
            self.on_flash_status_off, 
            self.timeroff)
        self.timeroff.Start(flash_len_ms, oneShot=True)
    
    def on_flash_status_off(self, event):
        self.statusbar.SetStatusText('')
Esempio n. 44
0
	data = m.cmech_quantities(['El','hmu'], maxL)
	lines1.append(zip(range(maxL+1),data[0]))
	#data[1].mask[0] = False
	#er = m.entropy_rate()
	#data[1] -= er
	#lines2.append(zip(range(maxL+1),data[1,:]))

fig = plt.figure()
#fig.set_figheight(5)
#fig.set_figwidth(10)
#fig.subplots_adjust(left = 0.09, right = 0.86)
ax1 = fig.add_subplot(1,1,1)

lc1 = LineCollection(lines1)
lc1.set_array(ps)
lc1.set_cmap(cmap)

ax1.add_collection(lc1)

ax1.set_xlabel(r"$L$ [symbols]")
ax1.set_ylabel(r"$\mathbf{E}(L)$ [bits]")
ax1.axis('auto')
ax1.set_xlim((0,maxL))

#ax2 = fig.add_subplot(1,2,2)

#lc2 = LineCollection(lines2)
#lc2.set_array(ps)
#lc2.set_cmap(cmap)

#ax2.add_collection(lc2)
Esempio n. 45
0
def plot_linestring_collection(ax, geoms, colors_or_values, plot_values,
                               vmin=None, vmax=None, cmap=None,
                               linewidth=1.0, **kwargs):
    """
    Plots a collection of LineString and MultiLineString geometries to `ax`

    Parameters
    ----------

    ax : matplotlib.axes.Axes
        where shapes will be plotted

    geoms : a sequence of `N` LineStrings and/or MultiLineStrings (can be mixed)

    colors_or_values : a sequence of `N` values or RGBA tuples
        It should have 1:1 correspondence with the geometries (not their components).

    plot_values : bool
        If True, `colors_or_values` is interpreted as a list of values, and will
        be mapped to colors using vmin/vmax/cmap (which become required).
        Otherwise `colors_or_values` is interpreted as a list of colors.

    Returns
    -------

    collection : matplotlib.collections.Collection that was plotted
    """

    from matplotlib.collections import LineCollection

    components, component_colors_or_values = _flatten_multi_geoms(
        geoms, colors_or_values)

    # LineCollection does not accept some kwargs.
    if 'markersize' in kwargs:
        del kwargs['markersize']
    segments = [np.array(linestring)[:, :2] for linestring in components]
    collection = LineCollection(segments,
                                linewidth=linewidth, **kwargs)

    if plot_values:
        collection.set_array(np.array(component_colors_or_values))
        collection.set_cmap(cmap)
        collection.set_clim(vmin, vmax)
    else:
        # set_color magically sets the correct combination of facecolor and
        # edgecolor, based on collection type.
        collection.set_color(component_colors_or_values)

        # If the user set facecolor and/or edgecolor explicitly, the previous
        # call to set_color might have overridden it (remember, the 'color' may
        # have come from plot_series, not from the user). The user should be
        # able to override matplotlib's default behavior, by setting them again
        # after set_color.
        if 'facecolor' in kwargs:
            collection.set_facecolor(kwargs['facecolor'])

        if 'edgecolor' in kwargs:
            collection.set_edgecolor(kwargs['edgecolor'])

    ax.add_collection(collection, autolim=True)
    ax.autoscale_view()
    return collection