Exemple #1
0
class Visualize:
    def __init__(self, v, t, e, fig, win, axesLimit=[-3,3.5,-2,2]):
        self.e = e.copy()
        self.p = [Polygon(v[ti]) for ti in t]
        self.p = PatchCollection(self.p, edgecolors='none')
        self.l = LineCollection(v[e[:,:2]])

        win = win or fig.canvas.manager.window
        if fig is None: fig = gcf()
        fig.clf()
        ax = fig.add_axes([0.02,0.02,.98,.98])
        ax.axis('scaled')
        ax.axis(axesLimit)
        ax.set_autoscale_on(False)
        self.axis, self.fig, self.win = ax, fig, win

        ax.add_collection(self.p)
        ax.add_collection(self.l)
        # ax.add_collection(self.l1)
        # ax.add_collection(self.l2)

    def update(self, title, phi):
        norm = Normalize(phi.min(), phi.max())
        self.p.set_norm(norm)
        self.l.set_norm(norm)
        self.p.set_array(phi)
        self.l.set_array(phi[self.e[:,2:]].mean(1))
        if not self.__dict__.has_key('colorbar'):
            self.colorbar = self.fig.colorbar(self.p)
        self.win.set_title(title)
        #self.fig.canvas.set_window_title(title)
        self.fig.canvas.draw()
Exemple #2
0
class Visualize:
    def __init__(self, v, t, e, fig, win, axesLimit=[-3, 3.5, -2, 2]):
        self.e = e.copy()
        self.p = [Polygon(v[ti]) for ti in t]
        self.p = PatchCollection(self.p, edgecolors='none')
        self.l = LineCollection(v[e[:, :2]])

        win = win or fig.canvas.manager.window
        if fig is None: fig = gcf()
        fig.clf()
        ax = fig.add_axes([0.02, 0.02, .98, .98])
        ax.axis('scaled')
        ax.axis(axesLimit)
        ax.set_autoscale_on(False)
        self.axis, self.fig, self.win = ax, fig, win

        ax.add_collection(self.p)
        ax.add_collection(self.l)
        # ax.add_collection(self.l1)
        # ax.add_collection(self.l2)

    def update(self, title, phi):
        norm = Normalize(phi.min(), phi.max())
        self.p.set_norm(norm)
        self.l.set_norm(norm)
        self.p.set_array(phi)
        self.l.set_array(phi[self.e[:, 2:]].mean(1))
        if not self.__dict__.has_key('colorbar'):
            self.colorbar = self.fig.colorbar(self.p)
        self.win.set_title(title)
        #self.fig.canvas.set_window_title(title)
        self.fig.canvas.draw()
Exemple #3
0
  def render(self,ax,alpha=.7,usePyLeaflet=False,
      minValueThreshold=-float('Inf'),logScale=True,colormapname='BlueRed'):
    if not usePyLeaflet or colormapname=='nothingRed':
      alpha=1.0

    patches = []
    values = []
    colorvalues = []
    for d in self.countPerGeohash.keys():
      try:
        if (self.countPerGeohash[d]>minValueThreshold):
          bbox = geohash.bbox(d)
          rect = mpatches.Rectangle([ bbox['w'],
                        bbox['s']],
                        bbox['e'] - bbox['w'],
                        bbox['n'] - bbox['s'],
                        ec='none', lw=.1, fc='red', alpha=alpha)
          patches.append(rect)
          # values.append(self.countPerGeohash[d] \
          #   if self.countPerGeohash[d]<3 else self.countPerGeohash[d]+10)
          # colorvalues.append(self.countPerGeohash[d] \
          #   if self.countPerGeohash[d]<3 else self.countPerGeohash[d]+10)
          values.append(self.countPerGeohash[d])
          colorvalues.append(self.countPerGeohash[d])
      except KeyError:
        print("'"+d +"' is not a valid geohash.")

    try:
      maxval = max(values)
      minval = min(values)
    except ValueError:
      print('heatmap appears to be empty...')
      maxval = 1
      minval = 0

    p = PatchCollection(patches,cmap=plt.get_cmap(colormapname),alpha=alpha)
  #   if usePyLeaflet:
    if (len(values)<100):
      p.set_edgecolors(np.array(['black' for x in values]))
    else:
      p.set_edgecolors(np.array(['#333333' if x<=2 \
        else ('#666666' if x<=10 else 'black') for x in values]))
  #   else:
  #     p.set_edgecolors(np.array(['white' for x in values]))
    p.set_array(np.array(colorvalues))
    if logScale:
      p.set_norm(colors.LogNorm(vmin=.01, vmax=maxval))
    else:
      p.set_norm(colors.Normalize(vmin=0, vmax=maxval))
    ax.add_collection(p)
    ax.set_xlim(self.bbox['w'], self.bbox['e'])
    ax.set_ylim(self.bbox['s'], self.bbox['n'])
    divider = make_axes_locatable(ax)
    cbar = plt.colorbar(p)
    cbar.set_clim(vmin=max(0,minval),vmax=maxval)
    cbar.update_normal(p)
    return
Exemple #4
0
def draw1DColumn(ax,
                 x,
                 val,
                 thk,
                 width=30,
                 ztopo=0,
                 cmin=1,
                 cmax=1000,
                 cmap=None,
                 name=None,
                 textoffset=0.0):
    """Draw a 1D column (e.g., from a 1D inversion) on a given ax.

    Examples
    --------
    >>> import numpy as np
    >>> import matplotlib.pyplot as plt
    >>> from pygimli.mplviewer import draw1DColumn
    >>> thk = [1, 2, 3, 4]
    >>> val = thk
    >>> fig, ax = plt.subplots()
    >>> draw1DColumn(ax, 0.5, val, thk, width=0.1, cmin=1, cmax=4, name="VES")
    <matplotlib.collections.PatchCollection object at ...>
    >>> ax.set_ylim(-np.sum(thk), 0)
    (-10, 0)
    """
    z = -np.hstack((0., np.cumsum(thk), np.sum(thk) * 1.5)) + ztopo
    recs = []
    for i in range(len(val)):
        recs.append(Rectangle((x - width / 2., z[i]), width, z[i + 1] - z[i]))

    pp = PatchCollection(recs)
    col = ax.add_collection(pp)

    pp.set_edgecolor(None)
    pp.set_linewidths(0.0)

    if cmap is not None:
        if isinstance(cmap, str):
            pp.set_cmap(pg.mplviewer.cmapFromName(cmap))
        else:
            pp.set_cmap(cmap)

    pp.set_norm(colors.LogNorm(cmin, cmax))
    pp.set_array(np.array(val))
    pp.set_clim(cmin, cmax)
    if name:
        ax.text(x + textoffset, ztopo, name, ha='center', va='bottom')

    updateAxes_(ax)

    return col
Exemple #5
0
def plot_polygons(geoms,
                  ax,
                  values=None,
                  colormap='Set1',
                  facecolor=None,
                  edgecolor=None,
                  alpha=1.0,
                  linewidth=1.0,
                  **kwargs):
    """Makes a MatPlotLib PatchCollection out of Polygon and/or MultiPolygon geometries 
     Thanks to http://stackoverflow.com/a/33753927 and David Sullivan"""

    # init list to store
    patches = []
    newvals = []

    for polynum in range(len(geoms)):  # for  polygon # i
        poly = geoms.iloc[polynum]  # find data.geometry[i]
        if type(
                poly
        ) != shapely.geometry.polygon.Polygon:  # if that is not a shapely Polygon object
            for currpoly in poly.geoms:  # then for data.geometry[i].geoms
                a = np.asarray(
                    currpoly.exterior
                )  # make a an array of those exterior values and
                patches.append(Polygon(a))  # append ato patches
                if values is not None:  # if values, add value to newvals
                    newvals.append(values.iloc[polynum])
        else:
            a = np.asarray(poly.exterior)
            patches.append(Polygon(a))
            if values is not None:
                newvals.append(values.iloc[polynum])

    patches = PatchCollection(patches,
                              facecolor=facecolor,
                              linewidth=linewidth,
                              edgecolor=edgecolor,
                              alpha=alpha,
                              **kwargs)
    if values is not None:
        patches.set_array(np.asarray(newvals))
        patches.set_cmap(colormap)
        norm = matplotlib.colors.Normalize()
        norm.autoscale(newvals)
        patches.set_norm(norm)
    ax.add_collection(patches, autolim=True)
    ax.autoscale_view()
    return patches
Exemple #6
0
def draw1DColumn(ax, x, val, thk, width=30, ztopo=0, cmin=1, cmax=1000,
                 cmap=None, name=None, textoffset=0.0):
    """Draw a 1D column (e.g., from a 1D inversion) on a given ax.

    Examples
    --------
    >>> import numpy as np
    >>> import matplotlib.pyplot as plt
    >>> from pygimli.mplviewer import draw1DColumn
    >>> thk = [1, 2, 3, 4]
    >>> val = thk
    >>> fig, ax = plt.subplots()
    >>> draw1DColumn(ax, 0.5, val, thk, width=0.1, cmin=1, cmax=4, name="VES")
    <matplotlib.collections.PatchCollection object at ...>
    >>> ax.set_ylim(-np.sum(thk), 0)
    (-10, 0)
    """
    z = -np.hstack((0., np.cumsum(thk), np.sum(thk) * 1.5)) + ztopo
    recs = []
    for i in range(len(val)):
        recs.append(Rectangle((x - width / 2., z[i]), width, z[i + 1] - z[i]))

    pp = PatchCollection(recs)
    col = ax.add_collection(pp)

    pp.set_edgecolor(None)
    pp.set_linewidths(0.0)

    if cmap is not None:
        if isinstance(cmap, str):
            pp.set_cmap(pg.mplviewer.cmapFromName(cmap))
        else:
            pp.set_cmap(cmap)

    pp.set_norm(colors.LogNorm(cmin, cmax))
    pp.set_array(np.array(val))
    pp.set_clim(cmin, cmax)
    if name:
        ax.text(x+textoffset, ztopo, name, ha='center', va='bottom')

    updateAxes_(ax)

    return col
Exemple #7
0
def draw_patches(list_of_polygons,
                 intensity_list=None,
                 cm='inferno',
                 empty_face=False,
                 norm=None,
                 alpha=None):
    patches = [Polygon(polygon, lw=0.1) for polygon in list_of_polygons]
    pc = PatchCollection(patches)
    if not empty_face:
        if intensity_list is None:
            raise ValueError('values not provided for Voronoi plot.')
        pc.set_array(np.array(intensity_list))
        pc.set_edgecolor('face')
        pc.set_cmap(cm)
        pc.set_norm(norm)
    else:
        pc.set_edgecolor([0.3, 0.3, 0.5])
        pc.set_facecolor([0, 0, 0, 0])
    return pc
Exemple #8
0
def draw_map(map: GeographicalMap, mode='random', data=None):
    fig, ax = plt.subplots(figsize=(18, 10))
    ax.set_xlim(map.xmin, map.xmax)
    ax.set_ylim(map.ymin, map.ymax)

    patches = PatchCollection(
        [matplotlib.patches.Polygon(p, fill=False) for p in map.polygons])

    to_line = lambda e: map.centroids[(e.v1, e.v2), :]
    land_edges = [to_line(e) for e in map.edges.values() if e.type == 'land']
    air_edges = [to_line(e) for e in map.edges.values() if e.type == 'air']
    land_edges = LineCollection(land_edges, linewidths=0.3, colors='red')
    air_edges = LineCollection(air_edges, linewidths=0.1, colors='green')

    ax.add_collection(patches)

    if mode == 'random':
        patches.set_array(np.random.randint(0, 20, size=map.N))
        patches.set_cmap(matplotlib.cm.jet)
    elif mode == 'population':
        patches.set_array(map.population)
        patches.set_cmap(matplotlib.cm.jet)
        patches.set_norm(matplotlib.colors.LogNorm())
        plt.colorbar(patches, ax=ax)
    elif mode == 'graph':
        patches.set_color('black')
        patches.set_facecolor('white')
        patches.set_linewidth(0.1)
        ax.scatter(map.centroids[:, 0], map.centroids[:, 1], s=5)

        # Plot edges.
        ax.add_collection(land_edges)
        ax.add_collection(air_edges)
    elif mode == 'data':
        patches.set_array(data)
        patches.set_cmap(matplotlib.cm.jet)
        plt.colorbar(patches, ax=ax)

    plt.show()
Exemple #9
0
def patchValMap(vals, xvec=None, yvec=None, ax=None, cMin=None, cMax=None,
                logScale=None, label=None, dx=1, dy=None, **kwargs):
    """ plot previously generated (generateVecMatrix) y map (category)

    Parameters
    ----------
    A : iterable
        to show
    xvec : dict {i:num}
        dict (must match A.shape[0])
    ymap : iterable
        vector for x axis (must match A.shape[0])
    ax : mpl.axis
        axis to plot, if not given a new figure is created
    cMin/cMax : float
        minimum/maximum color values
    logScale : bool
        logarithmic colour scale [min(A)>0]
    label : string
        colorbar label
    """
    if cMin is None:
        cMin = np.min(vals)
    if cMax is None:
        cMax = np.max(vals)
    if logScale is None:
        logScale = (cMin > 0.0)
    if logScale:
        norm = LogNorm(vmin=cMin, vmax=cMax)
    else:
        norm = Normalize(vmin=cMin, vmax=cMax)
    if 'ax' is None:
        fig, ax = plt.subplots()

    recs = []
    if dy is None:  # map y values to unique
        ymap = {xy: ii for ii, xy in enumerate(np.unique(yvec))}
        for i in range(len(vals)):
            recs.append(Rectangle((xvec[i]-dx/2, ymap[yvec[i]]-0.5), dx, 1))
    else:
        for i in range(len(vals)):
            recs.append(Rectangle((xvec[i]-dx/2, yvec[i]-dy/2), dx, dy))

    pp = PatchCollection(recs)
    col = ax.add_collection(pp)
    pp.set_edgecolor(None)
    pp.set_linewidths(0.0)
    if 'cmap' in kwargs:
        pp.set_cmap(kwargs.pop('cmap'))
    pp.set_norm(norm)
    pp.set_array(np.array(vals))
    pp.set_clim(cMin, cMax)
    ax.set_xlim(min(xvec)-dx/2, max(xvec)+dx/2)
    ax.set_ylim(len(ymap)-0.5, -0.5)

    updateAxes_(ax)
    cbar = None
    if kwargs.pop('colorBar', True):
        cbar = pg.mplviewer.createColorbar(col, cMin=cMin, cMax=cMax, nLevs=5,
                                           label=label)
    return ax, cbar, ymap
Exemple #10
0
def patchMatrix(A, xmap=None, ymap=None, ax=None, cMin=None, cMax=None,
                logScale=None, label=None, dx=1, **kwargs):
    """ plot previously generated (generateVecMatrix) matrix

    Parameters
    ----------
    A : numpy.array2d
        matrix to show
    xmap : dict {i:num}
        dict (must match A.shape[0])
    ymap : iterable
        vector for x axis (must match A.shape[0])
    ax : mpl.axis
        axis to plot, if not given a new figure is created
    cMin/cMax : float
        minimum/maximum color values
    logScale : bool
        logarithmic colour scale [min(A)>0]
    label : string
        colorbar label
    """
    mat = np.ma.masked_where(A == 0.0, A, False)
    if cMin is None:
        cMin = np.min(mat)
    if cMax is None:
        cMax = np.max(mat)
    if logScale is None:
        logScale = (cMin > 0.0)
    if logScale:
        norm = LogNorm(vmin=cMin, vmax=cMax)
    else:
        norm = Normalize(vmin=cMin, vmax=cMax)
    if 'ax' is None:
        fig, ax = plt.subplots()

    iy, ix = np.nonzero(A)  # != 0)
    recs = []
    vals = []
    for i in range(len(ix)):
        recs.append(Rectangle((ix[i]-dx/2, iy[i]-0.5), dx, 1))
        vals.append(A[iy[i], ix[i]])

    pp = PatchCollection(recs)
    col = ax.add_collection(pp)
    pp.set_edgecolor(None)
    pp.set_linewidths(0.0)
    if 'cmap' in kwargs:
        pp.set_cmap(kwargs.pop('cmap'))
    pp.set_norm(norm)
    pp.set_array(np.array(vals))
    pp.set_clim(cMin, cMax)
    xval = [k for k in xmap.keys()]
    ax.set_xlim(min(xval)-dx/2, max(xval)+dx/2)
    ax.set_ylim(len(ymap)+0.5, -0.5)

    updateAxes_(ax)
    cbar = None
    if kwargs.pop('colorBar', True):
        cbar = pg.mplviewer.createColorbar(col, cMin=cMin, cMax=cMax, nLevs=5,
                                           label=label)
    return ax, cbar
# and finally we use the dipping constraint matrix
fop.setConstraints(Cdip)
res = inv.run(vals, error, lam=15)
print(('Cg-9/2m: ' + '{:.1f} ' * 6).format(*fop(res), inv.chi2()))
pg.show(mesh, res, ax=ax[1, 1], **kw)
ax[1, 1].set_title("I=[10/2], dip=25")
np.testing.assert_array_less(inv.chi2(), 1.2)

# plot the position of the priors
patches = [CirclePolygon(po, 0.2) for po in pos]
for ai in ax.flat:
    p = PatchCollection(patches, cmap=kw['cMap'])
    p.set_facecolor(None)
    p.set_array(np.array(vals))
    p.set_norm(LogNorm(kw['cMin'], kw['cMax']))
    ai.add_collection(p)

# %%
# Note that all four regularization operators fit the data equivalently but
# the images (i.e. how the gaps between the data points are filled) are quite
# different. This is something we should have in mind using regularization.

# %%
# Generating geostatistical media
# -------------------------------
# For generating geostatistical media, one can use the function
# generateGeostatisticalModel. It computes a correlation matrix and multiplies
# it with a pseudo-random (randn) series. The arguments are the same as for the
# correlation or constraint matrices.
model = pg.utils.generateGeostatisticalModel(mesh, I=[20, 4])
Exemple #12
0
def plot_stack(npixels, hp, cr, snr, filename=None, scale=None):
    lon, lat = hp.healpix_to_lonlat(npixels)
    boundaries = hp.boundaries_lonlat(npixels, 1)

    patches = []
    for blon, blat in zip(*boundaries):
        patches.append(
            Polygon(np.array([blon.value, blat.value]).T, closed=True))

    if not scale:
        vmin_cr, vmax_cr = cr.flatten().min(), cr.flatten().max()
        vmin_snr, vmax_snr = snr.flatten().min(), snr.flatten().max()
        scale = [vmin_cr, vmax_cr, vmin_snr, vmax_snr]
    else:
        vmin_cr, vmax_cr = scale[0], scale[1]
        vmin_snr, vmax_snr = scale[2], scale[3]

    norm_cr = Normalize(vmin=vmin_cr, vmax=vmax_cr)
    norm_snr = Normalize(vmin=vmin_snr, vmax=vmax_snr)

    fig, axs = plt.subplots(2, 3, constrained_layout=False, figsize=(5.5, 4))
    for i, eband in enumerate(["6", "7", "8"]):
        # Count-rate "images"
        pcm_cr = axs[0, i].scatter(lon,
                                   lat,
                                   c=cr[:, i],
                                   s=1,
                                   vmin=vmin_cr,
                                   vmax=vmax_cr)

        p = PatchCollection(patches, alpha=1)
        p.set_array(cr[:, i])
        p.set_norm(norm_cr)
        axs[0, i].add_collection(p)

        axs[0, i].set_title(f"Energy band {eband}")
        axs[0, i].set_xticks([])
        axs[0, i].set_yticks([])

        # signal-to-noise ratio "images"
        pcm_snr = axs[1, i].scatter(lon,
                                    lat,
                                    c=snr[:, i],
                                    s=1,
                                    vmin=vmin_snr,
                                    vmax=vmax_snr)

        p = PatchCollection(patches, alpha=1)
        p.set_array(snr[:, i])
        p.set_norm(norm_snr)
        axs[1, i].add_collection(p)

        axs[1, i].set_xticks([])
        axs[1, i].set_yticks([])

        if i == 0:
            axs[0, i].set_ylabel("Stack net CR (median)")
            axs[1, i].set_ylabel("Stack SNR (median)")

    plt.tight_layout()

    fig.colorbar(pcm_cr, ax=axs[0, :], shrink=0.6, location='bottom', pad=0.02)
    fig.colorbar(pcm_snr,
                 ax=axs[1, :],
                 shrink=0.6,
                 location='bottom',
                 pad=0.02)

    if filename:
        fig.savefig(filename, bbox_inches='tight', pad_inches=0)
        plt.close()
    else:
        plt.show()

    return scale
Exemple #13
0
def patchValMap(vals,
                xvec=None,
                yvec=None,
                ax=None,
                cMin=None,
                cMax=None,
                logScale=None,
                label=None,
                dx=1,
                dy=None,
                **kwargs):
    """Plot previously generated (generateVecMatrix) y map (category).

    Parameters
    ----------
    vals : iterable
        Data values to show.
    xvec : dict {i:num}
        dict (must match vals.shape[0])
    ymap : iterable
        vector for x axis (must match vals.shape[0])
    ax : mpl.axis
        axis to plot, if not given a new figure is created
    cMin/cMax : float
        minimum/maximum color values
    logScale : bool
        logarithmic colour scale [min(vals)>0]
    label : string
        colorbar label
    ** kwargs:
        * circular : bool
            Plot in polar coordinates.
    """
    if cMin is None:
        cMin = np.min(vals)
    if cMax is None:
        cMax = np.max(vals)

    if logScale is None:
        logScale = (cMin > 0.0)

    norm = None
    if logScale and cMin > 0:
        norm = LogNorm(vmin=cMin, vmax=cMax)
    else:
        norm = Normalize(vmin=cMin, vmax=cMax)

    if ax is None:
        ax = plt.subplots()[1]

    recs = []

    circular = kwargs.pop('circular', False)
    if circular:
        recs = [None] * len(xvec)
        if dy is None:  # map y values to unique
            ymap = {xy: ii for ii, xy in enumerate(np.unique(yvec))}

            xyMap = {}
            for i, y in enumerate(yvec):
                if y not in xyMap:
                    xyMap[y] = []
                xyMap[y].append(i)

            maxR = max(ymap.values())  # what's that for? not used
            dR = 1 / (len(ymap.values()) + 1)

            dOff = np.pi / 2  # what's that for? not used

            for y, xIds in xyMap.items():
                r = 1. - dR * (ymap[y] + 1)
                # ax.plot(r * np.cos(xvec[xIds]),
                #         r * np.sin(xvec[xIds]), 'o')

                # print(y, ymap[y])
                for i in xIds:
                    phi = xvec[i]
                    x = r * np.cos(phi)  # what's that for? not used
                    y = r * np.sin(phi)

                    dPhi = (xvec[1] - xvec[0])

                    recs[i] = Wedge((0., 0.),
                                    r + dR / 1.5,
                                    (phi - dPhi) * 360 / (2 * np.pi),
                                    (phi + dPhi) * 360 / (2 * np.pi),
                                    width=dR,
                                    zorder=1 + r)
                    # if i < 5:
                    #     ax.text(x, y, str(i))
                    # pg.wait()
        else:
            raise ("Implementme")
    else:
        if dy is None:  # map y values to unique
            ymap = {xy: ii for ii, xy in enumerate(np.unique(yvec))}
            for i in range(len(vals)):
                recs.append(
                    Rectangle((xvec[i] - dx / 2, ymap[yvec[i]] - 0.5), dx, 1))
        else:
            for i in range(len(vals)):
                recs.append(
                    Rectangle((xvec[i] - dx / 2, yvec[i] - dy / 2), dx, dy))
        ax.set_xlim(min(xvec) - dx / 2, max(xvec) + dx / 2)
        ax.set_ylim(len(ymap) - 0.5, -0.5)

    pp = PatchCollection(recs)
    # ax.clear()
    col = ax.add_collection(pp)
    pp.set_edgecolor(None)
    pp.set_linewidths(0.0)

    if circular:
        pp.set_edgecolor('black')
        pp.set_linewidths(0.1)

    cmap = pg.mplviewer.cmapFromName(**kwargs)
    if kwargs.pop('markOutside', False):
        cmap.set_bad('grey')
        cmap.set_under('darkgrey')
        cmap.set_over('lightgrey')
        cmap.set_bad('black')
    pp.set_cmap(cmap)

    pp.set_norm(norm)
    pp.set_array(vals)
    pp.set_clim(cMin, cMax)

    updateAxes_(ax)
    cbar = kwargs.pop('colorBar', True)
    ori = kwargs.pop('orientation', 'horizontal')
    if cbar in ['horizontal', 'vertical']:
        ori = cbar
        cbar = True

    if cbar is True:  # not for cbar=1, which is really confusing!
        cbar = pg.mplviewer.createColorBar(col,
                                           cMin=cMin,
                                           cMax=cMax,
                                           nLevs=5,
                                           label=label,
                                           orientation=ori)
    elif cbar is not False:
        # .. cbar is an already existing cbar .. so we update its values
        pg.mplviewer.updateColorBar(cbar,
                                    cMin=cMin,
                                    cMax=cMax,
                                    nLevs=5,
                                    label=label)

    updateAxes_(ax)
    return ax, cbar, ymap
Exemple #14
0
def create_bus_symbol_collection(coords,
                                 buses=None,
                                 size=5,
                                 marker="o",
                                 patch_type="circle",
                                 colors=None,
                                 z=None,
                                 cmap=None,
                                 norm=None,
                                 infofunc=None,
                                 picker=False,
                                 net=None,
                                 **kwargs):
    infos = []

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

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

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

    pc.patch_type = patch_type
    pc.size = size
    if 'orientation' in kwargs:
        pc.orientation = kwargs['orientation']
    if "zorder" in kwargs:
        pc.set_zorder(kwargs["zorder"])
    pc.info = infos
    return pc
Exemple #15
0
def showStitchedModels(mods, axes=None, cmin=None, cmax=None, **kwargs):
    """
        Show several 1d block models as (stitched) section.
    """
    x = kwargs.pop('x', np.arange(len(mods)))
    topo = kwargs.pop('topo', x*0)

    nlay = int(np.floor((len(mods[0]) - 1) / 2.)) + 1
    if cmin is None or cmax is None:
        cmin = 1e9
        cmax = 1e-9
        for model in mods:
            res = np.asarray(model)[nlay - 1:nlay * 2 - 1]
            cmin = min(cmin, min(res))
            cmax = max(cmax, max(res))

    if kwargs.pop('sameSize', True):  # all having the same width
        dx = np.ones_like(x)*np.median(np.diff(x))
    else:
        dx = np.diff(x) * 1.05
        dx = np.hstack((dx, dx[-1]))

    x1 = x - dx / 2
    if axes is None:
        fig, ax = plt.subplots()
    else:
        ax = axes
        fig = ax.figure

#    ax.plot(x, x * 0., 'k.')
    zm = kwargs.pop('zm', None)
    maxz = 0.
    if zm is not None:
        maxz = zm
    recs = []
    RES = []
    for i, mod in enumerate(mods):
        mod1 = np.asarray(mod)
        res = mod1[nlay - 1:]
        RES.extend(res)

        thk = mod1[:nlay - 1]
        thk = np.hstack((thk, thk[-1]))
        z = np.hstack((0., np.cumsum(thk)))
        if zm is not None:
            thk[-1] = zm - z[-2]
            z[-1] = zm
        else:
            maxz = max(maxz, z[-1])

        for j in range(len(thk)):
            recs.append(Rectangle((x1[i], topo[i]-z[j]), dx[i], -thk[j]))

    pp = PatchCollection(recs, edgecolors=kwargs.pop('edgecolors', 'none'))
    pp.set_edgecolor(kwargs.pop('edgecolors', 'none'))
    pp.set_linewidths(0.0)
    ax.add_collection(pp)
    if 'cmap' in kwargs:
        pp.set_cmap(kwargs['cmap'])

    print(cmin, cmax)
    norm = LogNorm(cmin, cmax)
    pp.set_norm(norm)
    pp.set_array(np.array(RES))
#    pp.set_clim(cmin, cmax)
    ax.set_ylim((-maxz, max(topo)))
    ax.set_xlim((x1[0], x1[-1] + dx[-1]))

    cbar = None
    if kwargs.pop('colorBar', True):
        cbar = plt.colorbar(pp, ax=ax, norm=norm, orientation='horizontal',
                            aspect=60)  # , ticks=[1, 3, 10, 30, 100, 300])
        if 'ticks' in kwargs:
            cbar.set_ticks(kwargs['ticks'])
#        cbar.autoscale_None()
    if axes is None:  # newly created fig+ax
        return fig, ax
    else:  # already given, better give back color bar
        return cbar
Exemple #16
0
class InstrumentView:
    def __init__(self,
                 scipp_obj=None,
                 bins=None,
                 masks=None,
                 cmap=None,
                 log=None,
                 vmin=None,
                 vmax=None,
                 aspect=None,
                 size=1,
                 projection=None,
                 nan_color=None,
                 filename=None,
                 continuous_update=None,
                 dim=None):

        self.fig2d = None
        self.fig3d = None
        self.scatter2d = None
        self.scatter3d = None
        self.outline = None
        self.size = size
        self.aspect = aspect
        self.do_update = None
        self.figurewidget = widgets.Output()
        self.figure2d = False
        self.figure3d = False
        self.image = None
        self.nan_color = nan_color
        self.log = log
        self.current_projection = None
        self.dim = dim

        self.data_arrays = {}
        tp = type(scipp_obj)
        if tp is sc.Dataset or tp is sc.DatasetProxy:
            for key in sorted(scipp_obj.keys()):
                var = scipp_obj[key]
                if self.dim in var.dims:
                    self.data_arrays[key] = var
        elif tp is sc.DataArray or tp is sc.DataProxy:
            self.data_arrays[scipp_obj.name] = scipp_obj
        else:
            raise RuntimeError("Unknown input type: {}. Allowed inputs "
                               "are a Dataset or a DataArray (and their "
                               "respective proxies).".format(tp))

        self.globs = {"cmap": cmap, "log": log, "vmin": vmin, "vmax": vmax}
        self.params = {}
        self.hist_data_array = {}
        self.scalar_map = {}
        self.minmax = {}

        # Find the min/max time-of-flight limits and store them
        self.minmax["tof"] = [np.Inf, np.NINF, 1]
        for key, data_array in self.data_arrays.items():
            bins_here = bins
            if data_array.sparse_dim is not None and bins_here is None:
                bins_here = True
            if bins_here is not None:
                dim = None if data_array.sparse_dim is not None else self.dim
                spdim = None if data_array.sparse_dim is None else self.dim
                var = make_bins(data_array=data_array,
                                sparse_dim=spdim,
                                dim=dim,
                                bins=bins_here,
                                padding=(data_array.sparse_dim is not None))
            else:
                var = data_array.coords[self.dim]
            self.minmax["tof"][0] = min(self.minmax["tof"][0], var.values[0])
            self.minmax["tof"][1] = max(self.minmax["tof"][1], var.values[-1])
            self.minmax["tof"][2] = var.shape[0]

        # Rebin all DataArrays to common Tof axis
        self.rebin_data(np.linspace(*self.minmax["tof"]))

        # Create dropdown menu to select the DataArray
        keys = list(self.hist_data_array.keys())
        self.dropdown = widgets.Dropdown(options=keys,
                                         description="Select entry:",
                                         layout={"width": "initial"})
        self.dropdown.observe(self.change_data_array, names="value")

        # Store current active data entry (DataArray)
        self.key = keys[0]

        # Create a Tof slider and its label
        self.tof_dim_indx = self.hist_data_array[self.key].dims.index(self.dim)
        self.slider = widgets.IntSlider(
            value=0,
            min=0,
            step=1,
            description=str(self.dim).replace("Dim.", ""),
            max=self.hist_data_array[self.key].shape[self.tof_dim_indx] - 1,
            continuous_update=continuous_update,
            readout=False)
        self.slider.observe(self.update_colors, names="value")
        self.label = widgets.Label()

        # Add text boxes to change number of bins/bin size
        self.nbins = widgets.Text(value=str(
            self.hist_data_array[self.key].shape[self.tof_dim_indx]),
                                  description="Number of bins:",
                                  style={"description_width": "initial"})
        self.nbins.on_submit(self.update_nbins)

        tof_values = self.hist_data_array[self.key].coords[self.dim].values
        self.bin_size = widgets.Text(value=str(tof_values[1] - tof_values[0]),
                                     description="Bin size:")
        self.bin_size.on_submit(self.update_bin_size)

        projections = [
            "3D", "Cylindrical X", "Cylindrical Y", "Cylindrical Z",
            "Spherical X", "Spherical Y", "Spherical Z"
        ]

        # Create toggle buttons to change projection
        self.buttons = {}
        for p in projections:
            self.buttons[p] = widgets.Button(
                description=p,
                disabled=False,
                button_style=("info" if (p == projection) else ""))
            self.buttons[p].on_click(self.change_projection)
        items = [self.buttons["3D"]]
        for x in "XYZ":
            items.append(self.buttons["Cylindrical {}".format(x)])
            items.append(self.buttons["Spherical {}".format(x)])
            if x != "Z":
                items.append(widgets.Label())

        self.togglebuttons = widgets.GridBox(
            items,
            layout=widgets.Layout(grid_template_columns="repeat(3, 150px)"))

        # Place widgets in boxes
        self.vbox = widgets.VBox([
            widgets.HBox([self.dropdown, self.slider, self.label]),
            widgets.HBox([self.nbins, self.bin_size]), self.togglebuttons
        ])
        self.box = widgets.VBox([self.figurewidget, self.vbox])
        self.box.layout.align_items = "center"

        # Protect against uninstalled ipyvolume
        if ipv is None and projection == "3D":
            print("Warning: 3D projection requires ipyvolume to be "
                  "installed. Use conda/pip install ipyvolume. Reverting to "
                  "2D projection.")
            self.buttons[projections[1]].button_style = "info"
            self.buttons["3D"].button_style = ""
            self.buttons["3D"].disabled = True

        # Render the plot here instead of at the top level because to capture
        # the matplotlib output (if a 2D projection is requested to begin with,
        # the output widget needs to be displayed first, before any mpl figure
        # is displayed.
        render_plot(widgets=self.box, filename=filename, ipv=ipv)

        # Get detector positions
        self.det_pos = np.array(
            sn.position(self.hist_data_array[self.key]).values)
        # Find extents of the detectors
        for i, x in enumerate("xyz"):
            self.minmax[x] = [
                np.amin(self.det_pos[:, i]),
                np.amax(self.det_pos[:, i])
            ]

        # Update the figure
        self.change_projection(self.buttons[projection])

        # Create members object
        self.members = {
            "widgets": {
                "sliders": self.slider,
                "buttons": self.buttons,
                "text": {
                    "nbins": self.nbins,
                    "bin_size": self.bin_size
                },
                "dropdown": self.dropdown
            },
            "fig2d": self.fig2d,
            "fig3d": self.fig3d,
            "scatter2d": self.scatter2d,
            "scatter3d": self.scatter3d,
            "outline": self.outline
        }

        return

    def rebin_data(self, bins):
        """
        Rebin the original data to Tof given some bins. This is executed both
        on first instrument display and when either the number of bins or the
        bin width is changed.
        """
        for key, data_array in self.data_arrays.items():

            # Histogram the data in the Tof dimension
            if data_array.sparse_dim is not None:
                self.hist_data_array[key] = histogram_sparse_data(
                    data_array, data_array.sparse_dim, bins)
            else:
                self.hist_data_array[key] = sc.rebin(
                    data_array, self.dim,
                    make_bins(data_array=data_array,
                              dim=self.dim,
                              bins=bins,
                              padding=False))

            # Parse input parameters for colorbar
            self.params[key] = parse_params(
                globs=self.globs, array=self.hist_data_array[key].values)
            cmap = cm.get_cmap(self.params[key]["cmap"])
            cmap.set_bad(color=self.nan_color)
            self.scalar_map[key] = cm.ScalarMappable(
                cmap=cmap, norm=self.params[key]["norm"])
        return

    def update_colors(self, change):
        self.do_update(change)
        self.label.value = name_with_unit(
            var=self.hist_data_array[self.key].coords[self.dim],
            name=value_to_string(self.hist_data_array[self.key].coords[
                self.dim].values[change["new"]]))
        return

    def change_projection(self, owner):

        if owner.description == self.current_projection:
            owner.button_style = "info"
            return
        if self.current_projection is not None:
            self.buttons[self.current_projection].button_style = ""

        # Temporarily disable automatic plotting in notebook
        if plt.isinteractive():
            plt.ioff()
            re_enable_interactive = True
        else:
            re_enable_interactive = False

        update_children = False

        if owner.description == "3D":
            self.projection_3d()
            self.do_update = self.update_colors_3d
        else:
            if self.current_projection == "3D" or \
               self.current_projection is None:
                update_children = True
            self.projection_2d(owner.description, update_children)
            self.do_update = self.update_colors_2d

        self.update_colors({"new": self.slider.value})

        self.current_projection = owner.description
        self.buttons[owner.description].button_style = "info"

        # Re-enable automatic plotting in notebook
        if re_enable_interactive:
            plt.ion()

        return

    def projection_3d(self):
        # Initialise Figure
        if not self.figure3d:
            self.fig3d = ipv.figure(width=config.plot.width,
                                    height=config.plot.height,
                                    animation=0,
                                    lighting=False)
            max_size = 0.0
            dx = {"x": 0, "y": 0, "z": 0}
            for ax in dx.keys():
                dx[ax] = np.ediff1d(self.minmax[ax])
            max_size = np.amax(list(dx.values()))
            # Make plot outline if aspect ratio is to be conserved
            if self.aspect == "equal":
                arrays = dict()
                for ax, s in dx.items():
                    diff = max_size - s
                    arrays[ax] = [
                        self.minmax[ax][0] - 0.5 * diff,
                        self.minmax[ax][1] + 0.5 * diff
                    ]

                outl_x, outl_y, outl_z = np.meshgrid(arrays["x"],
                                                     arrays["y"],
                                                     arrays["z"],
                                                     indexing="ij")
                self.outline = ipv.plot_wireframe(outl_x,
                                                  outl_y,
                                                  outl_z,
                                                  color="black")
            # Try to guess marker size
            perc_size = 100.0 * self.size / max_size
            self.scatter3d = ipv.scatter(x=self.det_pos[:, 0],
                                         y=self.det_pos[:, 1],
                                         z=self.det_pos[:, 2],
                                         marker="square_2d",
                                         size=perc_size)
            self.figure3d = True

        self.figurewidget.clear_output()
        self.box.children = tuple([self.figurewidget, ipv.gcc(), self.vbox])
        return

    def update_colors_3d(self, change):
        arr = self.hist_data_array[self.key][self.dim, change["new"]].values
        if self.log:
            arr = np.ma.masked_where(arr <= 0, arr)
        self.scatter3d.color = self.scalar_map[self.key].to_rgba(arr)
        return

    def projection_2d(self, projection, update_children):
        # Initialise figure if we switched from 3D view, if not re-use current
        # figure.
        if update_children:
            self.box.children = tuple([self.figurewidget, self.vbox])
        if not self.figure2d:
            self.fig2d, self.ax = plt.subplots(
                1,
                1,
                figsize=(config.plot.width / config.plot.dpi,
                         config.plot.height / config.plot.dpi))

        if update_children:
            with self.figurewidget:
                disp.display(self.fig2d)

        # Compute cylindrical or spherical projections
        permutations = {"X": [0, 2, 1], "Y": [1, 0, 2], "Z": [2, 1, 0]}
        axis = projection[-1]

        theta = np.arctan2(self.det_pos[:, permutations[axis][2]],
                           self.det_pos[:, permutations[axis][1]])
        if projection.startswith("Cylindrical"):
            z_or_phi = self.det_pos[:, permutations[axis][0]]
        elif projection.startswith("Spherical"):
            z_or_phi = np.arcsin(
                self.det_pos[:, permutations[axis][0]] /
                np.sqrt(self.det_pos[:, 0]**2 + self.det_pos[:, 1]**2 +
                        self.det_pos[:, 2]**2))

        # Create the scatter
        if not self.figure2d:
            patches = []
            for x, y in zip(theta, z_or_phi):
                patches.append(
                    Rectangle((x - 0.5 * self.size, y - 0.5 * self.size),
                              self.size, self.size))

            self.scatter2d = PatchCollection(
                patches,
                cmap=self.params[self.key]["cmap"],
                norm=self.params[self.key]["norm"],
                array=self.hist_data_array[self.key][self.dim,
                                                     self.slider.value].values)
            self.ax.add_collection(self.scatter2d)
            self.save_xy = np.array([theta, z_or_phi]).T
            if self.params[self.key]["cbar"]:
                self.cbar = plt.colorbar(self.scatter2d, ax=self.ax)
                self.cbar.ax.set_ylabel(
                    name_with_unit(var=self.hist_data_array[self.key],
                                   name=""))
                self.cbar.ax.yaxis.set_label_coords(-1.1, 0.5)
            self.figure2d = True
        else:
            self.scatter2d.set_offset_position("data")
            self.scatter2d.set_offsets(
                np.array([theta, z_or_phi]).T - self.save_xy)
        self.ax.set_xlim(
            [np.amin(theta) - self.size,
             np.amax(theta) + self.size])
        self.ax.set_ylim(
            [np.amin(z_or_phi) - self.size,
             np.amax(z_or_phi) + self.size])
        return

    def update_colors_2d(self, change):
        self.scatter2d.set_array(
            self.hist_data_array[self.key][self.dim, change["new"]].values)
        self.fig2d.canvas.draw_idle()
        return

    def update_nbins(self, owner):
        try:
            nbins = int(owner.value)
        except ValueError:
            print("Warning: could not convert value: {} to an "
                  "integer.".format(owner.value))
            return
        # self.rebin_data(nbins, from_nbins_text)
        self.rebin_data(
            np.linspace(self.minmax["tof"][0], self.minmax["tof"][1],
                        nbins + 1))
        x = self.hist_data_array[self.key].coords[self.dim].values
        self.bin_size.value = str(x[1] - x[0])
        self.update_slider()
        return

    def update_bin_size(self, owner):
        try:
            binw = float(owner.value)
        except ValueError:
            print("Warning: could not convert value: {} to a "
                  "float.".format(owner.value))
            return
        self.rebin_data(
            np.arange(self.minmax["tof"][0], self.minmax["tof"][1], binw))
        self.nbins.value = str(
            self.hist_data_array[self.key].shape[self.tof_dim_indx])
        self.update_slider()
        return

    def update_slider(self):
        """
        Try to replace the slider around the same position
        """

        # Compute percentage position
        perc_pos = self.slider.value / self.slider.max
        # Compute new position
        nbins = int(self.nbins.value)
        new_pos = int(perc_pos * nbins)
        # Either place new upper boundary first, or change slider value first
        if new_pos > self.slider.max:
            self.slider.max = nbins
            self.slider.value = new_pos
        else:
            self.slider.value = new_pos
            self.slider.max = nbins
        return

    def change_data_array(self, change):
        self.key = change["new"]
        if self.scatter2d is not None:
            # Apparently, you have to set norm, clim on PatchCollection and
            # clim on the colorbar to get this working. Only setting norm
            # seems to work only on the first change.
            self.scatter2d.set_norm(self.params[self.key]["norm"])
            self.scatter2d.set_clim(vmin=self.params[self.key]["vmin"],
                                    vmax=self.params[self.key]["vmax"])
            self.cbar.set_clim(vmin=self.params[self.key]["vmin"],
                               vmax=self.params[self.key]["vmax"])
        self.update_colors({"new": self.slider.value})
def create_bus_collection(net, buses=None, size=5, marker="o", patch_type="circle", colors=None,
                          z=None, cmap=None, norm=None, infofunc=None, picker=False,
                          bus_geodata=None, cbar_title="Bus Voltage [pu]", **kwargs):
    """
    Creates a matplotlib patch collection of pandapower buses.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    infos = []

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

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

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

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

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

    return pc
Exemple #18
0
    def hexPlot(self,
                what='tallies',
                fixed=None,
                ax=None,
                cmap=None,
                logColor=False,
                xlabel=None,
                ylabel=None,
                logx=False,
                logy=False,
                loglog=False,
                title=None,
                normalizer=None,
                cbarLabel=None,
                borderpad=2.5,
                **kwargs):
        """
        Create and return a hexagonal mesh plot.

        Parameters
        ----------
        what: {'tallies', 'errors', 'scores'}
            Quantity to plot
        fixed: None or dict
            Dictionary of slicing arguments to pass to :meth:`slice`
        {ax}
        {cmap}
        {logColor}
        {xlabel}
        {ylabel}
        {logx}
        {logy}
        {loglog}
        {title}
        borderpad: int or float
            Percentage of total plot to apply as a border. A value of
            zero means that the extreme edges of the hexagons will touch
            the x and y axis.
        {kwargs} :class:`matplotlib.patches.RegularPolygon`

        Raises
        ------
        AttributeError
            If :attr:`pitch` and :attr:`hexType` are not set.
        """
        borderpad = max(0, float(borderpad))
        if fixed and ('xcoord' in fixed or 'ycoord' in fixed):
            raise KeyError("Refusing to restrict along one of the hexagonal "
                           "dimensions {x/y}coord")
        for attr in {'pitch', 'hexType'}:
            if getattr(self, attr) is None:
                raise AttributeError("{} is not set.".format(attr))

        for key in {'color', 'fc', 'facecolor', 'orientation'}:
            checkClearKwargs(key, 'hexPlot', **kwargs)
        ec = kwargs.get('ec', None) or kwargs.get('edgecolor', None)
        if ec is None:
            ec = 'k'
        kwargs['ec'] = kwargs['edgecolor'] = ec
        if 'figure' in kwargs and kwargs['figure'] is not None:
            fig = kwargs['figure']
            if not isinstance(fig, Figure):
                raise TypeError(
                    "Expected 'figure' to be of type Figure, is {}".format(
                        type(fig)))
            if len(fig.axes) != 1 and not ax:
                raise TypeError("Don't know where to place the figure since"
                                "'figure' argument has multiple axes.")
            if ax and fig.axes and ax not in fig.axes:
                raise IndexError("Passed argument for 'figure' and 'ax', "
                                 "but ax is not attached to figure.")
            ax = ax or (fig.axes[0] if fig.axes else axes())
        alpha = kwargs.get('alpha', None)

        ny = len(self.indexes['ycoord'])
        nx = len(self.indexes['xcoord'])
        data = self.slice(fixed, what)
        if data.shape != (ny, nx):
            raise IndexError("Constrained data does not agree with hexagonal "
                             "grid structure. Coordinate grid: {}. "
                             "Constrained shape: {}".format((ny, nx),
                                                            data.shape))
        nItems = ny * nx
        patches = empty(nItems, dtype=object)
        values = empty(nItems)
        coords = self.grids['COORD']

        ax = ax or axes()
        pos = 0
        xmax, ymax = [
            -inf,
        ] * 2
        xmin, ymin = [
            inf,
        ] * 2
        radius = self.pitch / sqrt(3)

        for xy, val in zip(coords, data.flat):
            values[pos] = val
            h = RegularPolygon(xy, 6, radius, self.__hexRot, **kwargs)
            verts = h.get_verts()
            vmins = verts.min(0)
            vmaxs = verts.max(0)
            xmax = max(xmax, vmaxs[0])
            xmin = min(xmin, vmins[0])
            ymax = max(ymax, vmaxs[1])
            ymin = min(ymin, vmins[1])
            patches[pos] = h
            pos += 1
        normalizer = normalizerFactory(values, normalizer, logColor,
                                       coords[:, 0], coords[:, 1])
        pc = PatchCollection(patches, cmap=cmap, alpha=alpha)
        pc.set_array(values)
        pc.set_norm(normalizer)
        ax.add_collection(pc)

        addColorbar(ax, pc, None, cbarLabel)

        formatPlot(
            ax,
            loglog=loglog,
            logx=logx,
            logy=logy,
            xlabel=xlabel or "X [cm]",
            ylabel=ylabel or "Y [cm]",
            title=title,
        )
        setAx_xlims(ax, xmin, xmax, pad=borderpad)
        setAx_ylims(ax, ymin, ymax, pad=borderpad)

        return ax
Exemple #19
0
def plot_polygon_collection(ax, geoms, colors_or_values, plot_values=False,
                            vmin=None, vmax=None, cmap=None,
                            edgecolor='black', alpha=0.5, linewidth=1.0,label="Area",norm=None,  **kwargs):
    """
    Plots a collection of Polygon and MultiPolygon geometries to `ax`

    Parameters
    ----------

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

    geoms : a sequence of `N` Polygons and/or MultiPolygons (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
    """


    
    components, component_colors_or_values = _flatten_multi_geoms(
        geoms, colors_or_values)

    # PatchCollection does not accept some kwargs.
    collection = PatchCollection([PolygonPatch(poly) for poly in components],
                                 linewidth=linewidth, edgecolor=edgecolor,
                                 alpha=alpha,label=label, **kwargs)

    if plot_values:
        collection.set_array(np.array(component_colors_or_values))
        collection.set_cmap(cmap)
        collection.set_norm(norm)
        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:
            collection.set_edgecolor(edgecolor)

    ax.add_collection(collection, autolim=True)
    ax.autoscale_view()
    return collection
def colorflood(gridobject, arr, **kwargs):
    """
    Method for color flooding an array for a Modflow grid.
    """

    #imports
    import matplotlib.cm as cm
    import matplotlib.colors as colors
    from matplotlib.patches import Polygon
    from matplotlib.collections import PatchCollection

    #set defaults from kwargs
    ax = matplotlib.pyplot.gca()
    arrexclude = []
    arrmin = arr.min()
    arrmax = arr.max()
    alpha = 1.0
    cb = False
    if kwargs.has_key('ax'): ax = kwargs['ax']
    if kwargs.has_key('arrexclude'): arrexclude = kwargs['arrexclude']
    if kwargs.has_key('arrmin'): arrmin = kwargs['arrmin']
    if kwargs.has_key('arrmax'): arrmax = kwargs['arrmax']
    if kwargs.has_key('alpha'): alpha = kwargs['alpha']
    if kwargs.has_key('cb'): cb = kwargs['cb'] #plot colorbar?

    norm = colors.normalize(arrmin, arrmax)
    patches = []
    for nodenumber in range(gridobject.nodes):
        nodeobj = gridobject.get_nodeobj(nodenumber)
        (k, i, j) = gridobject.get_indices(nodenumber)
        x, y = nodeobj.position[0], nodeobj.position[1]
        dx, dy = nodeobj.dxdydz[0], nodeobj.dxdydz[1]
        xmin = x - dx / 2.
        xmax = x + dx / 2.
        ymin = y - dy / 2.
        ymax = y + dy / 2.
        vertices = [ (xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
        if arr[k, i, j] in arrexclude:
            #cannot get visible and alpha to work, so this is the hack.
            vertices = [(0,0), (0,0), (0,0), (0,0) ]
        poly = Polygon(vertices, linewidth=0., fc = None, 
                       ec='none')
        patches.append(poly)

    #set up the patchcollection for faster coloring
    p = PatchCollection(patches, match_original=True)
    p.set_array(arr.reshape(-1))
    p.set_norm(norm)
    p.set_edgecolor(None)
    p.set_alpha(alpha)
    ax.add_collection(p)
    
    #create the colorbar
    if cb:
        pc = []
        color = cm.jet(norm(arrmin))
        poly = Polygon( ((0,0), (1,0), (1,1), (0,1)), linewidth=0., 
            fc = color, ec='none', alpha=0.4)
        pc.append(poly)
        color = cm.jet(norm(arrmax))
        poly = Polygon( ((0,0), (1,0), (1,1), (0,1)), linewidth=0., 
            fc = color, ec='none', alpha=0.4)
        pc.append(poly)
        p = PatchCollection(pc, match_original=True)
        p.set_array(numpy.array([arrmin, arrmax]))
        matplotlib.pyplot.colorbar(p, shrink=0.5)

    return    
Exemple #21
0
def showStitchedModels(models,
                       ax=None,
                       x=None,
                       cMin=None,
                       cMax=None,
                       thk=None,
                       logScale=True,
                       title=None,
                       zMin=0,
                       zMax=0,
                       zLog=False,
                       **kwargs):
    """Show several 1d block models as (stitched) section.

    Parameters
    ----------
    model : iterable of iterable (np.ndarray or list of np.array)
        1D models (consisting of thicknesses and values) to plot
    ax : matplotlib axes [None - create new]
        axes object to plot in
    x : iterable
        positions of individual models
    cMin/cMax : float [None - autodetection from range]
        minimum and maximum colorscale range
    logScale : bool [True]
        use logarithmic color scaling
    zMin/zMax : float [0 - automatic]
        range for z (y axis) limits
    zLog : bool
        use logarithmic z (y axis) instead of linear
    topo : iterable
        vector of elevation for shifting
    thk : iterable
        vector of layer thicknesses for all models
    Returns
    -------
    ax : matplotlib axes [None - create new]
        axes object to plot in
    """
    if x is None:
        x = np.arange(len(models))

    topo = kwargs.pop('topo', np.zeros_like(x))

    fig = None
    if ax is None:
        fig, ax = plt.subplots()

    dxmed2 = np.median(np.diff(x)) / 2.
    patches = []
    zMinLimit = 9e99
    zMaxLimit = 0

    if thk is not None:
        nlay = len(models[0])
    else:
        nlay = int(np.floor((len(models[0]) + 1) / 2.))

    vals = np.zeros((len(models), nlay))
    for i, imod in enumerate(models):
        if thk is not None:  # take only resistivity from model
            vals[i, :] = imod
            thki = thk
        else:  # extract thickness from model vector
            if isinstance(imod, pg.Vector):
                vals[i, :] = imod[nlay - 1:2 * nlay - 1]
                thki = np.asarray(imod[:nlay - 1])
            else:
                vals[i, :] = imod[nlay - 1:2 * nlay - 1]
                thki = imod[:nlay - 1]

        if zMax > 0:
            z = np.hstack((0., np.cumsum(thki), zMax))
        else:
            thki = np.hstack((thki, thki[-1] * 3))
            z = np.hstack((0., np.cumsum(thki)))

        z = topo[i] - z
        zMinLimit = min(zMinLimit, z[-1])
        zMaxLimit = max(zMaxLimit, z[0])

        for j in range(nlay):
            rect = Rectangle((x[i] - dxmed2, z[j]), dxmed2 * 2,
                             z[j + 1] - z[j])
            patches.append(rect)

    p = PatchCollection(patches)  # , cmap=cmap, linewidths=0)
    if cMin is not None:
        p.set_clim(cMin, cMax)

    setMappableData(p, vals.ravel(), logScale=logScale)
    ax.add_collection(p)

    if logScale:
        norm = colors.LogNorm(cMin, cMax)
        p.set_norm(norm)

    if 'cMap' in kwargs:
        p.set_cmap(kwargs['cMap'])


#    ax.set_ylim((zMaxLimit, zMin))
    ax.set_ylim((zMinLimit, zMaxLimit))

    if zLog:
        ax.set_yscale("log", nonposy='clip')

    ax.set_xlim((min(x) - dxmed2, max(x) + dxmed2))

    if title is not None:
        ax.set_title(title)

    if kwargs.pop('colorBar', True):
        cb = pg.viewer.mpl.createColorBar(p, cMin=cMin, cMax=cMax, nLevs=5)
        #    cb = plt.colorbar(p, orientation='horizontal',aspect=50,pad=0.1)
        if 'cticks' in kwargs:
            xt = np.unique(np.clip(kwargs['cticks'], cMin, cMax))
            cb.set_ticks(xt)
            cb.set_ticklabels([str(xti) for xti in xt])
        if 'label' in kwargs:
            cb.set_label(kwargs['label'])

    plt.draw()
    return ax  # maybe return cb as well?
Exemple #22
0
def patchValMap(vals, xvec=None, yvec=None, ax=None, cMin=None, cMax=None,
                logScale=None, label=None, dx=1, dy=None, **kwargs):
    """Plot previously generated (generateVecMatrix) y map (category).

    Parameters
    ----------
    vals : iterable
        to show
    xvec : dict {i:num}
        dict (must match vals.shape[0])
    ymap : iterable
        vector for x axis (must match vals.shape[0])
    ax : mpl.axis
        axis to plot, if not given a new figure is created
    cMin/cMax : float
        minimum/maximum color values
    logScale : bool
        logarithmic colour scale [min(vals)>0]
    label : string
        colorbar label
    """
    if cMin is None:
        cMin = np.min(vals)
    if cMax is None:
        cMax = np.max(vals)

    if logScale is None:
        logScale = (cMin > 0.0)

    norm = None
    if logScale and cMin > 0:
        norm = LogNorm(vmin=cMin, vmax=cMax)
    else:
        norm = Normalize(vmin=cMin, vmax=cMax)

    if 'ax' is None:
        ax = plt.subplots()[1]

    recs = []
    if dy is None:  # map y values to unique
        ymap = {xy: ii for ii, xy in enumerate(np.unique(yvec))}
        for i in range(len(vals)):
            recs.append(Rectangle((xvec[i] - dx / 2, ymap[yvec[i]] - 0.5),
                                  dx, 1))
    else:
        for i in range(len(vals)):
            recs.append(Rectangle((xvec[i] - dx / 2, yvec[i] - dy / 2),
                                  dx, dy))

    pp = PatchCollection(recs)
    # ax.clear()
    col = ax.add_collection(pp)
    pp.set_edgecolor(None)
    pp.set_linewidths(0.0)

    cmap = pg.mplviewer.cmapFromName(**kwargs)
    cmap.set_bad('grey')
    if kwargs.pop('markOutside', True):
        cmap.set_under('darkgrey')
        cmap.set_bad('lightgrey')
    pp.set_cmap(cmap)

    pp.set_norm(norm)
    pp.set_array(np.array(vals))
    pp.set_clim(cMin, cMax)
    ax.set_xlim(min(xvec) - dx / 2, max(xvec) + dx / 2)
    ax.set_ylim(len(ymap) - 0.5, -0.5)

    updateAxes_(ax)
    cbar = kwargs.pop('colorBar', True)

    if cbar is True:  # not for cbar=1, which is really confusing!
        cbar = pg.mplviewer.createColorBar(col, cMin=cMin, cMax=cMax,
                                           nLevs=5, label=label)
    elif cbar is not False:  # what the hell is this?
        pg.mplviewer.updateColorBar(cbar, cMin=cMin, cMax=cMax,
                                    nLevs=5, label=label)

    return ax, cbar, ymap
Exemple #23
0
def create_bus_collection(net, buses=None, size=5, marker="o", patch_type="circle", colors=None,
                          cmap=None, norm=None, infofunc=None, **kwargs):
    """
    Creates a matplotlib patch collection of pandapower buses.
    
    Input:

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

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

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

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

    pc.patch_type = patch_type
    pc.size = size
    if "zorder" in kwargs:
        pc.set_zorder(kwargs["zorder"])
    pc.info = infos
    return pc
Exemple #24
0
def showStitchedModels_Redundant(mods,
                                 ax=None,
                                 cmin=None,
                                 cmax=None,
                                 **kwargs):
    """Show several 1d block models as (stitched) section."""
    x = kwargs.pop('x', np.arange(len(mods)))
    topo = kwargs.pop('topo', x * 0)

    nlay = int(np.floor((len(mods[0]) - 1) / 2.)) + 1
    if cmin is None or cmax is None:
        cmin = 1e9
        cmax = 1e-9
        for model in mods:
            res = np.asarray(model)[nlay - 1:nlay * 2 - 1]
            cmin = min(cmin, min(res))
            cmax = max(cmax, max(res))

    if kwargs.pop('sameSize', True):  # all having the same width
        dx = np.ones_like(x) * np.median(np.diff(x))
    else:
        dx = np.diff(x) * 1.05
        dx = np.hstack((dx, dx[-1]))

    x1 = x - dx / 2
    if ax is None:
        fig, ax = plt.subplots()
    else:
        ax = ax
        fig = ax.figure

#    ax.plot(x, x * 0., 'k.')
    zm = kwargs.pop('zm', None)
    maxz = 0.
    if zm is not None:
        maxz = zm
    recs = []
    RES = []
    for i, mod in enumerate(mods):
        mod1 = np.asarray(mod)
        res = mod1[nlay - 1:]
        RES.extend(res)

        thk = mod1[:nlay - 1]
        thk = np.hstack((thk, thk[-1]))
        z = np.hstack((0., np.cumsum(thk)))
        if zm is not None:
            thk[-1] = zm - z[-2]
            z[-1] = zm
        else:
            maxz = max(maxz, z[-1])

        for j, _ in enumerate(thk):
            recs.append(Rectangle((x1[i], topo[i] - z[j]), dx[i], -thk[j]))

    pp = PatchCollection(recs, edgecolors=kwargs.pop('edgecolors', 'none'))
    pp.set_edgecolor(kwargs.pop('edgecolors', 'none'))
    pp.set_linewidths(0.0)
    ax.add_collection(pp)
    if 'cmap' in kwargs:
        pp.set_cmap(kwargs['cmap'])

    print(cmin, cmax)
    norm = colors.LogNorm(cmin, cmax)
    pp.set_norm(norm)
    pp.set_array(np.array(RES))
    #    pp.set_clim(cmin, cmax)
    ax.set_ylim((-maxz, max(topo)))
    ax.set_xlim((x1[0], x1[-1] + dx[-1]))

    cbar = None
    if kwargs.pop('colorBar', True):
        cbar = plt.colorbar(pp,
                            ax=ax,
                            norm=norm,
                            orientation='horizontal',
                            aspect=60)  # , ticks=[1, 3, 10, 30, 100, 300])
        if 'ticks' in kwargs:
            cbar.set_ticks(kwargs['ticks'])


#        cbar.autoscale_None()
    if ax is None:  # newly created fig+ax
        return fig, ax
    else:  # already given, better give back color bar
        return cbar
Exemple #25
0
def patchValMap(vals, xvec=None, yvec=None, ax=None, cMin=None, cMax=None,
                logScale=None, label=None, dx=1, dy=None, **kwargs):
    """Plot previously generated (generateVecMatrix) y map (category).

    Parameters
    ----------
    vals : iterable
        Data values to show.
    xvec : dict {i:num}
        dict (must match vals.shape[0])
    ymap : iterable
        vector for x axis (must match vals.shape[0])
    ax : mpl.axis
        axis to plot, if not given a new figure is created
    cMin/cMax : float
        minimum/maximum color values
    logScale : bool
        logarithmic colour scale [min(vals)>0]
    label : string
        colorbar label
    ** kwargs:
        * circular : bool
            Plot in polar coordinates.
    """
    if cMin is None:
        cMin = np.min(vals)
    if cMax is None:
        cMax = np.max(vals)

    if logScale is None:
        logScale = (cMin > 0.0)

    norm = None
    if logScale and cMin > 0:
        norm = LogNorm(vmin=cMin, vmax=cMax)
    else:
        norm = Normalize(vmin=cMin, vmax=cMax)

    if ax is None:
        ax = plt.subplots()[1]

    recs = []

    circular = kwargs.pop('circular', False)
    if circular:
        recs = [None] * len(xvec)
        if dy is None:  # map y values to unique
            ymap = {xy: ii for ii, xy in enumerate(np.unique(yvec))}

            xyMap = {}
            for i, y in enumerate(yvec):
                if y not in xyMap:
                    xyMap[y] = []
                xyMap[y].append(i)

            # maxR = max(ymap.values())  # what's that for? not used
            dR = 1 / (len(ymap.values())+1)
            # dOff = np.pi / 2  # what's that for? not used

            for y, xIds in xyMap.items():
                r = 1. - dR*(ymap[y]+1)
                # ax.plot(r * np.cos(xvec[xIds]),
                #         r * np.sin(xvec[xIds]), 'o')

                # print(y, ymap[y])
                for i in xIds:
                    phi = xvec[i]
                    # x = r * np.cos(phi)  # what's that for? not used
                    y = r * np.sin(phi)

                    dPhi = (xvec[1] - xvec[0])

                    recs[i] = Wedge((0., 0.), r + dR/1.5,
                                    (phi - dPhi)*360/(2*np.pi),
                                    (phi + dPhi)*360/(2*np.pi),
                                    width=dR,
                                    zorder=1+r)
                    # if i < 5:
                    #     ax.text(x, y, str(i))
                    # pg.wait()
        else:
            raise("Implementme")
    else:
        if dy is None:  # map y values to unique
            ymap = {xy: ii for ii, xy in enumerate(np.unique(yvec))}
            for i in range(len(vals)):
                recs.append(Rectangle((xvec[i] - dx / 2, ymap[yvec[i]] - 0.5),
                                      dx, 1))
        else:
            for i in range(len(vals)):
                recs.append(Rectangle((xvec[i] - dx / 2, yvec[i] - dy / 2),
                                      dx, dy))
        ax.set_xlim(min(xvec) - dx / 2, max(xvec) + dx / 2)
        ax.set_ylim(len(ymap) - 0.5, -0.5)

    pp = PatchCollection(recs)
    # ax.clear()
    col = ax.add_collection(pp)
    pp.set_edgecolor(None)
    pp.set_linewidths(0.0)

    if circular:
        pp.set_edgecolor('black')
        pp.set_linewidths(0.1)

    cmap = pg.mplviewer.cmapFromName(**kwargs)
    if kwargs.pop('markOutside', False):
        cmap.set_bad('grey')
        cmap.set_under('darkgrey')
        cmap.set_over('lightgrey')
        cmap.set_bad('black')
    pp.set_cmap(cmap)

    pp.set_norm(norm)
    pp.set_array(vals)
    pp.set_clim(cMin, cMax)

    updateAxes_(ax)
    cbar = kwargs.pop('colorBar', True)
    ori = kwargs.pop('orientation', 'horizontal')
    if cbar in ['horizontal', 'vertical']:
        ori = cbar
        cbar = True

    if cbar is True:  # not for cbar=1, which is really confusing!
        cbar = pg.mplviewer.createColorBar(col, cMin=cMin, cMax=cMax,
                                           nLevs=5, label=label,
                                           orientation=ori)
    elif cbar is not False:
        # .. cbar is an already existing cbar .. so we update its values
        pg.mplviewer.updateColorBar(cbar, cMin=cMin, cMax=cMax,
                                    nLevs=5, label=label)

    updateAxes_(ax)
    return ax, cbar, ymap
Exemple #26
0
def patchMatrix(mat,
                xmap=None,
                ymap=None,
                ax=None,
                cMin=None,
                cMax=None,
                logScale=None,
                label=None,
                dx=1,
                **kwargs):
    """Plot previously generated (generateVecMatrix) matrix.

    Parameters
    ----------
    mat : numpy.array2d
        matrix to show
    xmap : dict {i:num}
        dict (must match A.shape[0])
    ymap : iterable
        vector for x axis (must match A.shape[0])
    ax : mpl.axis
        axis to plot, if not given a new figure is created
    cMin/cMax : float
        minimum/maximum color values
    logScale : bool
        logarithmic colour scale [min(A)>0]
    label : string
        colorbar label
    dx : float
        width of the matrix elements (by default 1)
    """
    mat = np.ma.masked_where(mat == 0.0, mat, False)
    if cMin is None:
        cMin = np.min(mat)
    if cMax is None:
        cMax = np.max(mat)
    if logScale is None:
        logScale = (cMin > 0.0)
    if logScale:
        norm = LogNorm(vmin=cMin, vmax=cMax)
    else:
        norm = Normalize(vmin=cMin, vmax=cMax)

    if 'ax' is None:
        ax = plt.subplots()[1]

    iy, ix = np.nonzero(mat)  # != 0)
    recs = []
    vals = []
    for i, _ in enumerate(ix):
        recs.append(Rectangle((ix[i] - dx / 2, iy[i] - 0.5), dx, 1))
        vals.append(mat[iy[i], ix[i]])

    pp = PatchCollection(recs)
    col = ax.add_collection(pp)
    pp.set_edgecolor(None)
    pp.set_linewidths(0.0)
    if 'cmap' in kwargs:
        pp.set_cmap(kwargs.pop('cmap'))
    if 'cMap' in kwargs:
        pp.set_cmap(kwargs.pop('cMap'))
    pp.set_norm(norm)
    pp.set_array(np.array(vals))
    pp.set_clim(cMin, cMax)
    xval = [k for k in xmap.keys()]
    ax.set_xlim(min(xval) - dx / 2, max(xval) + dx / 2)
    ax.set_ylim(len(ymap) + 0.5, -0.5)

    updateAxes_(ax)
    cbar = None
    if kwargs.pop('colorBar', True):
        ori = kwargs.pop('orientation', 'horizontal')
        cbar = pg.mplviewer.createColorBar(col,
                                           cMin=cMin,
                                           cMax=cMax,
                                           nLevs=5,
                                           label=label,
                                           orientation=ori)
    return ax, cbar
Exemple #27
0
def plot_filled_pixels(
    map, save_filename=None, title='', ra_range=[], dec_range=[],
    colorbar_range=[], log=False, colorbar_label='Surface Brightness (Jy/sr)',
    ra_cut=None,
    overplot_points=False, point_ras=None, point_decs=None, point_values=None,
    overplot_points_vmin=-np.pi, overplot_points_vmax=np.pi,
    overplot_points_colormap='seismic',
    overplot_mwa_beam_contours=False,
    mwa_beam_center_ras=[0, 60], mwa_beam_center_decs=[-27, -27],
    overplot_hera_band=False,
    overplot_bright_sources=False,
    overplot_sgp=False,
    big=False,
    galactic_coord_contours=False
):

    if map.coords == '':
        print('WARNING: No map coordinate scheme supplied.')
        print('Assuming equatorial coordinates.')
        coords = 'equatorial'
    else:
        coords = map.coords

    # Set branch cut location
    if ra_cut is None:
        if len(ra_range) == 2:
            ra_cut = (ra_range[1]+ra_range[0])/2.+12.
        else:
            ra_cut = 18.

    if overplot_points:
        point_ras[np.where(point_ras > ra_cut)] -= 24.
        point_ras[np.where(point_ras < ra_cut-24.)] += 24.

    map.get_ra_dec(ra_cut=ra_cut*15.)

    # Limit pixel calculation when axis ranges are set
    if len(ra_range) == 2 or len(dec_range) == 2:
        if len(ra_range) != 2:
            ra_range = np.array([np.min(map.ra_arr), np.max(map.ra_arr)])/15.
        if len(dec_range) != 2:
            dec_range = [np.min(map.dec_arr), np.max(map.dec_arr)]
        use_indices = np.arange(len(map.signal_arr))[
            (map.ra_arr/15.>ra_range[0]) & (map.ra_arr/15.<ra_range[1])
            & (map.dec_arr>dec_range[0]) & (map.dec_arr<dec_range[1])
        ]
        use_map = healpix_utils.HealpixMap(
            map.signal_arr[use_indices], map.pix_arr[use_indices], map.nside,
            nest=map.nest, coords=coords
        )
    else:
        use_map = map

    use_map.get_pixel_corners(ra_cut=ra_cut*15.)
    patches = []
    for ind in range(len(use_map.signal_arr)):
        polygon = Polygon(np.stack([
            use_map.pix_corner_ras_arr[ind]/15.,
            use_map.pix_corner_decs_arr[ind]
        ], axis=1))
        patches.append(polygon)
    colors = use_map.signal_arr

    # Establish axis ranges
    if len(ra_range) != 2:
        ra_range = np.array([
            np.amin(use_map.pix_corner_ras_arr),
            np.amax(use_map.pix_corner_ras_arr)
        ])/15.
    if len(dec_range) != 2:
        dec_range = np.array([
            np.amin(use_map.pix_corner_decs_arr),
            np.amax(use_map.pix_corner_decs_arr)
        ])

    cm_use = 'Greys_r'
    #cm_use = 'viridis'
    collection = PatchCollection(patches, cmap=cm_use, lw=0.05)
    collection.set_array(np.array(colors))  # set the data colors
    collection.set_edgecolor('face')  # make the face and edge colors match
    if log:  # set the color bar to a log scale
        collection.set_norm(LogNorm())
    if len(colorbar_range) == 2:  # set the colorbar min and max
        collection.set_clim(vmin=colorbar_range[0], vmax=colorbar_range[1])
    else:
        signal_mean = np.mean(colors)
        signal_std = np.std(colors)
        collection.set_clim(
            vmin=max([min(colors), signal_mean-5*signal_std]),
            vmax=min([max(colors), signal_mean+5*signal_std])
        )

    plt.rcParams.update({'font.size': 9})
    plt.rcParams['axes.facecolor']='gray'
    if big:
        fig, ax = plt.subplots(figsize=(10, 4), dpi=600)
    else:
        fig, ax = plt.subplots(figsize=(6, 0.6*4), dpi=600)
    ax.add_collection(collection)  # plot data

    plt.xlabel('RA (hours)')
    plt.ylabel('Dec (degrees)')
    #plt.axis('equal')
    ax.set_aspect(1./15.)
    #ax.set_facecolor('gray')  # make plot background gray
    if galactic_coord_contours:
        npoints_ra = 200
        npoints_dec = 100
        coord_ra_vals = np.linspace(ra_range[0], ra_range[1], num=npoints_ra)
        coord_dec_vals = np.linspace(dec_range[0], dec_range[1], num=npoints_dec)
        phi_eq = np.radians(coord_ra_vals*15.)
        for phi_ind, phi_val in enumerate(phi_eq):
            if phi_val < 0:
                phi_eq[phi_ind] += 2*np.pi
        theta_eq = np.radians(90. - coord_dec_vals)
        rot = hp.rotator.Rotator(coord=['C', 'G'])
        theta_eq_grid, phi_eq_grid = np.meshgrid(theta_eq, phi_eq)
        theta_gal, phi_gal = rot(theta_eq_grid.flatten(), phi_eq_grid.flatten())
        theta_gal = 90. - np.degrees(theta_gal.reshape((npoints_ra, npoints_dec)))
        phi_gal = np.degrees(phi_gal.reshape((npoints_ra, npoints_dec))) + 180.
        theta_cont = plt.contour(
            coord_ra_vals, coord_dec_vals,
            theta_gal.T, levels=np.arange(-90, 90, 15),
            colors='white', linestyles=['solid'], linewidths=0.2
        )
        plt.clabel(theta_cont, inline=True, fontsize=5, fmt="%.0f$^\circ$")
        # Plot nonzero contours to aviod branch cut
        phi_gal_nonzero = np.copy(phi_gal)
        phi_gal_nonzero[np.where(phi_gal_nonzero < 20.)] = np.nan
        phi_gal_nonzero[np.where(phi_gal_nonzero > 340.)] = np.nan
        phi_cont_nonzero = plt.contour(
            coord_ra_vals, coord_dec_vals,
            phi_gal_nonzero.T, levels=np.arange(45, 360, 45),
            colors='white', linestyles=['solid'], linewidths=0.2
        )
        plt.clabel(phi_cont_nonzero, inline=True, fontsize=5, fmt="%.0f$^\circ$")
        # Plot zero contour
        phi_gal_zero = np.copy(phi_gal)
        phi_gal_zero[np.where(phi_gal_zero > 180.)] = phi_gal_zero[np.where(phi_gal_zero > 180.)] - 360.
        phi_gal_zero[np.where(phi_gal_zero > 20.)] = np.nan
        phi_gal_zero[np.where(phi_gal_zero < -20.)] = np.nan
        phi_cont_zero = plt.contour(
            coord_ra_vals, coord_dec_vals,
            phi_gal_zero.T, levels=[0],
            colors='white', linestyles=['solid'], linewidths=0.2
        )
        plt.clabel(phi_cont_zero, inline=True, fontsize=5, fmt="%.0f$^\circ$")
    if overplot_points:
        cm_overplot_points = plt.cm.get_cmap(overplot_points_colormap)
        plt.scatter(
            point_ras, point_decs, c=point_values,
            vmin=overplot_points_vmin, vmax=overplot_points_vmax, s=30,
            cmap=cm_overplot_points, edgecolor='black', linewidth=.5
        )
    if overplot_mwa_beam_contours:
        beam_ras, beam_decs, beam_val = get_mwa_beam()
        for beam_ind, use_beam_center_ra in enumerate(mwa_beam_center_ras):
            use_beam_center_dec = mwa_beam_center_decs[beam_ind]
            plt.contour(
                (beam_ras+use_beam_center_ra)/15., beam_decs+use_beam_center_dec,
                beam_val, levels=[.5],
                colors='cyan', linestyles=['solid'], linewidths=0.7
            )
    if overplot_hera_band:
        hera_band_center = -30.
        hera_band_width = 11.
        plt.plot(
            ra_range, np.full(2, hera_band_center+hera_band_width/2),
            '--', color='cyan', linewidth=0.7
        )
        plt.plot(
            ra_range, np.full(2, hera_band_center-hera_band_width/2),
            '--', color='cyan', linewidth=0.7
        )
    if overplot_bright_sources:
        source_names = [
            'Pictor A', 'Fornax A'
        ]
        named_source_ras = np.array([
            79.9572, 50.6738
        ])/15.
        named_source_decs = np.array([
            -45.7788, -37.2083
        ])
        plt.plot(
            named_source_ras, named_source_decs, 'x',
            color='yellow', markersize=3
        )
        for source_ind, name in enumerate(source_names):
            plt.annotate(
                name, (named_source_ras[source_ind]-2/15., named_source_decs[source_ind]),
                fontsize=8., color='black',
                path_effects=[patheffects.withStroke(linewidth=0.5, foreground="white")]
            )
    if overplot_sgp:
        sgp_ra = 0.857222
        sgp_dec = -27.1283
        plt.plot(
            [sgp_ra], [sgp_dec], '+', color='red', markersize=6
        )
        plt.annotate(
            'SGP', (sgp_ra, sgp_dec-8), fontsize=8,
            horizontalalignment='center', color='black',
            path_effects=[patheffects.withStroke(linewidth=0.5, foreground="white")]
        )

    plt.axis([ra_range[1], ra_range[0], dec_range[0], dec_range[1]])
    plt.title(title)
    cbar = fig.colorbar(collection, ax=ax, extend='both')  # add colorbar
    # label colorbar
    cbar.ax.set_ylabel(colorbar_label, rotation=270, labelpad=15)
    if save_filename is not None:
        print('Saving plot to {}'.format(save_filename))
        plt.savefig(save_filename, format='png', dpi=600)
        plt.close()
    else:
        plt.show()
Exemple #28
0
def GridStrain(pos, disp, k, par, plotpar, plotst):
    '''
	GridStrain computes the infinitesimal strain of a network
	of stations with displacements in x (east) and y (north).
	Strain in z is assumed to be zero (plane strain)

	USE: cent,eps,ome,pstrain,rotc = GridStrain(pos,disp,k,par,plotpar,plotst)

	pos = nstations x 2 matrix with x (east) and y (north)
		positions of stations in meters
	disp = nstations x 2 matrix with x (east) and y (north)
		displacements of stations in meters
	k = Type of computation: Delaunay (k = 0), nearest
		neighbor (k = 1), or distance weighted (k = 2)
	par = Parameters for nearest neighbor or distance
		weighted computation. If Delaunay (k = 0), enter
		a scalar corresponding to the minimum internal
		angle of a triangle valid for computation.
		If nearest neighbor (k = 1), input a 1 x 3 vector
		with grid spacing, number of nearest neighbors,
		and maximum distance to neighbors. If distance
		weighted (k = 2), input a 1 x 2 vector with grid
		spacing and distance weighting factor alpha
	plotpar = Parameter to color the cells: Max elongation
		(plotpar = 0), minimum elongation
		(plotpar = 1), rotation (plotpar = 2),
		or dilatation (plotpar = 3)
	plotst = A flag to plot the stations (1) or not (0)
	cent = ncells x 2 matrix with x and y positions of cells
		centroids
	eps = 3 x 3 x ncells array	with strain tensors of
		the cells
	ome = 3 x 3 x ncells array with rotation tensors of
		the cells
	pstrain = 3 x 3 x ncells array with magnitude and
		orientation of principal strains of
		the cells
	rotc = ncells x 3 matrix with rotation components
		of cells

	NOTE: Input/Output angles are in radians. Output
		azimuths are given with respect to North
		pos, disp, grid spacing, max. distance to
		neighbors, and alpha should be in meters
	
	GridStrain uses functions lscov and InfStrain

	Python function translated from the Matlab function
	GridStrain in Allmendinger et al. (2012)
	'''
    pi = np.pi
    # If Delaunay
    if k == 0:
        # Indexes of triangles vertices
        # Use function Delaunay
        tri = Delaunay(pos)
        inds = tri.simplices
        # Number of cells
        ncells = np.size(inds, 0)
        # Number of stations per cell = 3
        nstat = 3
        # Centers of cells
        cent = np.zeros((ncells, 2))
        for i in range(0, ncells):
            # Triangle vertices
            v1x = pos[inds[i, 0], 0]
            v2x = pos[inds[i, 1], 0]
            v3x = pos[inds[i, 2], 0]
            v1y = pos[inds[i, 0], 1]
            v2y = pos[inds[i, 1], 1]
            v3y = pos[inds[i, 2], 1]
            # Center of cell
            cent[i, 0] = (v1x + v2x + v3x) / 3.0
            cent[i, 1] = (v1y + v2y + v3y) / 3.0
            # Triangle internal angles
            s1 = np.sqrt((v3x - v2x)**2 + (v3y - v2y)**2)
            s2 = np.sqrt((v1x - v3x)**2 + (v1y - v3y)**2)
            s3 = np.sqrt((v2x - v1x)**2 + (v2y - v1y)**2)
            a1 = np.arccos((v2x-v1x)*(v3x-v1x)/(s3*s2)+\
               (v2y-v1y)*(v3y-v1y)/(s3*s2))
            a2 = np.arccos((v3x-v2x)*(v1x-v2x)/(s1*s3)+\
               (v3y-v2y)*(v1y-v2y)/(s1*s3))
            a3 = np.arccos((v2x-v3x)*(v1x-v3x)/(s1*s2)+\
               (v2y-v3y)*(v1y-v3y)/(s1*s2))
            # If any of the internal angles is less than
            # specified minimum, invalidate triangle
            if a1 < par or a2 < par or a3 < par:
                inds[i, :] = np.zeros(3)
    # If nearest neighbor or distance weighted
    else:
        # Construct grid
        xmin = min(pos[:, 0])
        xmax = max(pos[:, 0])
        ymin = min(pos[:, 1])
        ymax = max(pos[:, 1])
        cellsx = int(np.ceil((xmax - xmin) / par[0]))
        cellsy = int(np.ceil((ymax - ymin) / par[0]))
        xgrid = np.arange(xmin, (xmin + (cellsx + 1) * par[0]), par[0])
        ygrid = np.arange(ymin, (ymin + (cellsy + 1) * par[0]), par[0])
        XX, YY = np.meshgrid(xgrid, ygrid)
        # Number of cells
        ncells = cellsx * cellsy
        # Number of stations per cell (nstat) and
        # other parameters
        # If nearest neighbor
        if k == 1:
            nstat = par[1]  # max neighbors
            sqmd = par[2]**2  # max squared distance
        # If distance weighted
        elif k == 2:
            nstat = np.size(pos, 0)  # all stations
            dalpha = 2.0 * par[1] * par[1]  # 2*alpha*alpha
        # Cells' centers
        cent = np.zeros((ncells, 2))
        count = 0
        for i in range(0, cellsy):
            for j in range(0, cellsx):
                cent[count, 0] = (XX[i, j] + XX[i, j + 1]) / 2.0
                cent[count, 1] = (YY[i, j] + YY[i + 1, j]) / 2.0
                count += 1
        # Initialize stations indexes for cells to -1
        inds = np.ones((ncells, nstat), dtype=int) * -1
        # Initialize weight matrix for distance weighted
        wv = np.zeros((ncells, nstat * 2))
        # For all cells set stations indexes
        for i in range(0, ncells):
            # Initialize sq distances to -1.0
            sds = np.ones(nstat) * -1.0
            # For all stations
            for j in range(0, np.size(pos, 0)):
                # Sq distance from cell center to station
                dx = cent[i, 0] - pos[j, 0]
                dy = cent[i, 1] - pos[j, 1]
                sd = dx**2 + dy**2
                # If nearest neighbor
                if k == 1:
                    # If within the max sq distance
                    if sd <= sqmd:
                        minsd = min(sds)
                        mini = np.argmin(sds)
                        # If less than max neighbors
                        if minsd == -1.0:
                            sds[mini] = sd
                            inds[i, mini] = j
                        # If max neighbors
                        else:
                            # If sq distance is less
                            # than neighbors max sq distance
                            maxsd = max(sds)
                            maxi = np.argmax(sds)
                            if sd < maxsd:
                                sds[maxi] = sd
                                inds[i, maxi] = j
                # If distance weighted
                elif k == 2:
                    # All stations indexes
                    inds[i, :] = np.arange(nstat)
                    # Eq. 8.27: Weight factor
                    weight = np.exp(-sd / dalpha)
                    wv[i, j * 2] = weight
                    wv[i, j * 2 + 1] = weight

    # Initialize arrays
    y = np.zeros(nstat * 2)
    M = np.zeros((nstat * 2, 6))
    e = np.zeros((3, 3))
    eps = np.zeros((3, 3, ncells))
    ome = np.zeros((3, 3, ncells))
    pstrain = np.zeros((3, 3, ncells))
    rotc = np.zeros((ncells, 3))

    # For each cell
    for i in range(0, ncells):
        # If required minimum number of stations
        if min(inds[i, :]) >= 0:
            # Eq. 8.24: Displacements column vector y
            # and design matrix M. X1 = North, X2 = East
            for j in range(0, nstat):
                ic = inds[i, j]
                y[j * 2] = disp[ic, 1]
                y[j * 2 + 1] = disp[ic, 0]
                M[j * 2, :] = [1., 0., pos[ic, 1], pos[ic, 0], 0., 0.]
                M[j * 2 + 1, :] = [0., 1., 0., 0., pos[ic, 1], pos[ic, 0]]
            # Eqs. 8.25-8.26: Find x using function lscov
            # If Delaunay or nearest neighbor
            if k == 0 or k == 1:
                x = lscov(M, y)
            # If distance weighted
            elif k == 2:
                x = lscov(M, y, wv[i, :])
            # Displacement gradient tensor
            for j in range(0, 2):
                e[j, 0] = x[j * 2 + 2]
                e[j, 1] = x[j * 2 + 3]
            # Compute strain
            eps[:,:,i],ome[:,:,i],pstrain[:,:,i],\
             rotc[i,:],_ = InfStrain(e)

    # Variable to plot
    # If maximum principal strain
    if plotpar == 0:
        vp = pstrain[0, 0, :]
        lcb = "emax"
    # If minimum principal strain
    elif plotpar == 1:
        vp = pstrain[2, 0, :]
        lcb = "emin"
    # If rotation:
    # For plane strain, rotation = rotc(3)
    elif plotpar == 2:
        vp = rotc[:, 2] * 180 / pi
        lcb = "Rotation (deg)"
    # If dilatation
    elif plotpar == 3:
        vp = pstrain[0, 0, :] + pstrain[1, 0, :] + pstrain[2, 0, :]
        lcb = "dilatation"

    # Make a figure
    fig, ax = plt.subplots()
    fig.set_size_inches(15.0, 7.5)

    # Patches and colors for cells
    patches = []
    colors = []

    # Fill cells patches and colors
    # If Delaunay
    if k == 0:
        for i in range(0, ncells):
            # If minimum number of stations
            if min(inds[i, :]) >= 0:
                xpyp = [[pos[inds[i,0],0],pos[inds[i,0],1]],\
                  [pos[inds[i,1],0],pos[inds[i,1],1]],\
                  [pos[inds[i,2],0],pos[inds[i,2],1]]]
                # length in km
                xpyp = np.divide(xpyp, 1e3)
                polygon = Polygon(xpyp, True)
                patches.append(polygon)
                colors.append(vp[i])
    # If nearest neighbor or distance weighted
    if k == 1 or k == 2:
        count = 0
        for i in range(0, cellsy):
            for j in range(0, cellsx):
                # If minimum number of stations
                if min(inds[count, :]) >= 0:
                    xpyp = [[XX[i,j],YY[i,j]],[XX[i,j+1],YY[i,j+1]],\
                     [XX[i+1,j+1],YY[i+1,j+1]],[XX[i+1,j],YY[i+1,j]]]
                    # length in km
                    xpyp = np.divide(xpyp, 1e3)
                    polygon = Polygon(xpyp, True)
                    patches.append(polygon)
                    colors.append(vp[count])
                count += 1

    # Collect cells patches
    pcoll = PatchCollection(patches)
    # Cells colors
    pcoll.set_array(np.array(colors))
    # Color map is blue to red
    pcoll.set_cmap('bwr')
    # Positive values are red, negative are
    # blue and zero is white
    vmin = min(vp)
    vmax = max(vp)
    norm = mcolors.TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax)
    pcoll.set_norm(norm)

    # Draw cells
    ax.add_collection(pcoll)

    # Plot stations
    if plotst == 1:
        plt.plot(pos[:, 0] * 1e-3, pos[:, 1] * 1e-3, 'k.', markersize=2)

    # Axes
    plt.axis('equal')
    plt.xlabel('x (km)')
    plt.ylabel('y (km)')

    # Color bar with nice ticks
    intv = (vmax - vmin) * 0.25
    ticks = [vmin, vmin + intv, vmin + 2 * intv, vmin + 3 * intv, vmax]
    lticks = ['{:.2e}'.format(ticks[0]),\
      '{:.2e}'.format(ticks[1]),'{:.2e}'.format(ticks[2]),\
      '{:.2e}'.format(ticks[3]),'{:.2e}'.format(ticks[4])]
    cbar = fig.colorbar(pcoll, label=lcb, ticks=ticks)
    cbar.ax.set_yticklabels(lticks)

    # Show plot
    plt.show()

    return cent, eps, ome, pstrain, rotc
Exemple #29
0
cmap = matplotlib.cm.get_cmap('Spectral')

fig = plt.figure()
ax = fig.add_subplot(111)

ax.set_xlim(-100, 100)
ax.set_ylim(-100, 100)

norm = matplotlib.colors.LogNorm(vmin=min(values), vmax=3000)

center = (0, 0)
patches = [Circle(center, r) for r in reversed(radius)]

p = PatchCollection(patches)
p.set_norm(norm)
p.set_linewidth(0)
p.set_array(np.array(values[::-1]))
ax.add_collection(p)
ax.set_xlabel('Distance in [AU]')
ax.set_ylabel('Distance in [AU]')
cbar = fig.colorbar(p, ax=ax)
cbar.set_label('Surface density [kg/m^2]')
''' 
fig, ax = plt.subplots()
 
ax.set_xlim(-100, 100)
ax.set_ylim(-100, 100)
ax.set_xlabel('Radius in [AU]')
ax.set_ylabel('Radius in [AU]')