iframe-proxy | Sunbelt Computer Software

895 lines (689 loc) · 21.6 KB
"""This file contains code for use with "Think Stats",
by Allen B. Downey, available from greenteapress.com
Copyright 2014 Allen B. Downey
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html
import math
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import warnings
class _Brewer(object):
    """Encapsulates a nice sequence of colors.
    Shades of blue that look good in color and can be distinguished
    in grayscale (up to a point).
    Borrowed from http://colorbrewer2.org/
    color_iter = None
    colors = [
        "#f7fbff",
        "#deebf7",
        "#c6dbef",
        "#9ecae1",
        "#6baed6",
        "#4292c6",
        "#2171b5",
        "#08519c",
        "#08306b",
    ][::-1]
    # lists that indicate which colors to use depending on how many are used
    which_colors = [
        [],
        [1],
        [1, 3],
        [0, 2, 4],
        [0, 2, 4, 6],
        [0, 2, 3, 5, 6],
        [0, 2, 3, 4, 5, 6],
        [0, 1, 2, 3, 4, 5, 6],
        [0, 1, 2, 3, 4, 5, 6, 7],
        [0, 1, 2, 3, 4, 5, 6, 7, 8],
    current_figure = None
    @classmethod
    def Colors(cls):
        """Returns the list of colors."""
        return cls.colors
    @classmethod
    def ColorGenerator(cls, num):
        """Returns an iterator of color strings.
        n: how many colors will be used
        """
        for i in cls.which_colors[num]:
            yield cls.colors[i]
        raise StopIteration("Ran out of colors in _Brewer.")
    @classmethod
    def InitIter(cls, num):
        """Initializes the color iterator with the given number of colors."""
        cls.color_iter = cls.ColorGenerator(num)
        fig = plt.gcf()
        cls.current_figure = fig
    @classmethod
    def ClearIter(cls):
        """Sets the color iterator to None."""
        cls.color_iter = None
        cls.current_figure = None
    @classmethod
    def GetIter(cls, num):
        """Gets the color iterator."""
        fig = plt.gcf()
        if fig != cls.current_figure:
            cls.InitIter(num)
            cls.current_figure = fig
        if cls.color_iter is None:
            cls.InitIter(num)
        return cls.color_iter
def _UnderrideColor(options):
    """If color is not in the options, chooses a color."""
    if "color" in options:
        return options
    # get the current color iterator; if there is none, init one
    color_iter = _Brewer.GetIter(5)
        options["color"] = next(color_iter)
    except StopIteration:
        # if you run out of colors, initialize the color iterator
        # and try again
        warnings.warn("Ran out of colors.  Starting over.")
        _Brewer.ClearIter()
        _UnderrideColor(options)
    return options
def PrePlot(num=None, rows=None, cols=None):
    """Takes hints about what's coming.
    num: number of lines that will be plotted
    rows: number of rows of subplots
    cols: number of columns of subplots
    if num:
        _Brewer.InitIter(num)
    if rows is None and cols is None:
        return
    if rows is not None and cols is None:
        cols = 1
    if cols is not None and rows is None:
        rows = 1
    # resize the image, depending on the number of rows and cols
    size_map = {
        (1, 1): (8, 6),
        (1, 2): (12, 6),
        (1, 3): (12, 6),
        (1, 4): (12, 5),
        (1, 5): (12, 4),
        (2, 2): (10, 10),
        (2, 3): (16, 10),
        (3, 1): (8, 10),
        (4, 1): (8, 12),
    if (rows, cols) in size_map:
        fig = plt.gcf()
        fig.set_size_inches(*size_map[rows, cols])
    # create the first subplot
    if rows > 1 or cols > 1:
        ax = plt.subplot(rows, cols, 1)
        global SUBPLOT_ROWS, SUBPLOT_COLS
        SUBPLOT_ROWS = rows
        SUBPLOT_COLS = cols
        ax = plt.gca()
    return ax
def SubPlot(plot_number, rows=None, cols=None, **options):
    """Configures the number of subplots and changes the current plot.
    rows: int
    cols: int
    plot_number: int
    options: passed to subplot
    rows = rows or SUBPLOT_ROWS
    cols = cols or SUBPLOT_COLS
    return plt.subplot(rows, cols, plot_number, **options)
def _Underride(d, **options):
    """Add key-value pairs to d only if key is not in d.
    If d is None, create a new dictionary.
    d: dictionary
    options: keyword args to add to d
    if d is None:
        d = {}
    for key, val in options.items():
        d.setdefault(key, val)
    return d
    """Clears the figure and any hints that have been set."""
    global LOC
    LOC = None
    _Brewer.ClearIter()
    plt.clf()
    fig = plt.gcf()
    fig.set_size_inches(8, 6)
def Figure(**options):
    """Sets options for the current figure."""
    _Underride(options, figsize=(6, 8))
    plt.figure(**options)
def Plot(obj, ys=None, style="", **options):
    """Plots a line.
      obj: sequence of x values, or Series, or anything with Render()
      ys: sequence of y values
      style: style string passed along to plt.plot
      options: keyword args passed to plt.plot
    options = _UnderrideColor(options)
    label = getattr(obj, "label", "_nolegend_")
    options = _Underride(options, linewidth=3, alpha=0.7, label=label)
    xs = obj
    if ys is None:
        if hasattr(obj, "Render"):
            xs, ys = obj.Render()
        if isinstance(obj, pd.Series):
            ys = obj.values
            xs = obj.index
    if ys is None:
        plt.plot(xs, style, **options)
        plt.plot(xs, ys, style, **options)
def Vlines(xs, y1, y2, **options):
    """Plots a set of vertical lines.
      xs: sequence of x values
      y1: sequence of y values
      y2: sequence of y values
      options: keyword args passed to plt.vlines
    options = _UnderrideColor(options)
    options = _Underride(options, linewidth=1, alpha=0.5)
    plt.vlines(xs, y1, y2, **options)
def Hlines(ys, x1, x2, **options):
    """Plots a set of horizontal lines.
      ys: sequence of y values
      x1: sequence of x values
      x2: sequence of x values
      options: keyword args passed to plt.vlines
    options = _UnderrideColor(options)
    options = _Underride(options, linewidth=1, alpha=0.5)
    plt.hlines(ys, x1, x2, **options)
def axvline(x, **options):
    """Plots a vertical line.
      x: x location
      options: keyword args passed to plt.axvline
    options = _UnderrideColor(options)
    options = _Underride(options, linewidth=1, alpha=0.5)
    plt.axvline(x, **options)
def axhline(y, **options):
    """Plots a horizontal line.
      y: y location
      options: keyword args passed to plt.axhline
    options = _UnderrideColor(options)
    options = _Underride(options, linewidth=1, alpha=0.5)
    plt.axhline(y, **options)
def tight_layout(**options):
    """Adjust subplots to minimize padding and margins."""
    options = _Underride(
        options, wspace=0.1, hspace=0.1, left=0, right=1, bottom=0, top=1
    plt.tight_layout()
    plt.subplots_adjust(**options)
def FillBetween(xs, y1, y2=None, where=None, **options):
    """Fills the space between two lines.
      xs: sequence of x values
      y1: sequence of y values
      y2: sequence of y values
      where: sequence of boolean
      options: keyword args passed to plt.fill_between
    options = _UnderrideColor(options)
    options = _Underride(options, linewidth=0, alpha=0.5)
    plt.fill_between(xs, y1, y2, where, **options)
def Bar(xs, ys, **options):
    """Plots a line.
      xs: sequence of x values
      ys: sequence of y values
      options: keyword args passed to plt.bar
    options = _UnderrideColor(options)
    options = _Underride(options, linewidth=0, alpha=0.6)
    plt.bar(xs, ys, **options)
def Scatter(xs, ys=None, **options):
    """Makes a scatter plot.
    xs: x values
    ys: y values
    options: options passed to plt.scatter
    options = _Underride(options, color="blue", alpha=0.2, s=30, edgecolors="none")
    if ys is None and isinstance(xs, pd.Series):
        ys = xs.values
        xs = xs.index
    plt.scatter(xs, ys, **options)
def HexBin(xs, ys, **options):
    """Makes a scatter plot.
    xs: x values
    ys: y values
    options: options passed to plt.scatter
    options = _Underride(options, cmap=matplotlib.cm.Blues)
    plt.hexbin(xs, ys, **options)
def Pdf(pdf, **options):
    """Plots a Pdf, Pmf, or Hist as a line.
      pdf: Pdf, Pmf, or Hist object
      options: keyword args passed to plt.plot
    low, high = options.pop("low", None), options.pop("high", None)
    n = options.pop("n", 101)
    xs, ps = pdf.Render(low=low, high=high, n=n)
    options = _Underride(options, label=pdf.label)
    Plot(xs, ps, **options)
def Pdfs(pdfs, **options):
    """Plots a sequence of PDFs.
    Options are passed along for all PDFs.  If you want different
    options for each pdf, make multiple calls to Pdf.
      pdfs: sequence of PDF objects
      options: keyword args passed to plt.plot
    for pdf in pdfs:
        Pdf(pdf, **options)
def Hist(hist, **options):
    """Plots a Pmf or Hist with a bar plot.
    The default width of the bars is based on the minimum difference
    between values in the Hist.  If that's too small, you can override
    it by providing a width keyword argument, in the same units
    as the values.
      hist: Hist or Pmf object
      options: keyword args passed to plt.bar
    # find the minimum distance between adjacent values
    xs, ys = hist.Render()
    # see if the values support arithmetic
        xs[0] - xs[0]
    except TypeError:
        # if not, replace values with numbers
        labels = [str(x) for x in xs]
        xs = np.arange(len(xs))
        plt.xticks(xs + 0.5, labels)
    if "width" not in options:
        try:
            options["width"] = 0.9 * np.diff(xs).min()
        except TypeError:
            warnings.warn(
                "Hist: Can't compute bar width automatically."
                "Check for non-numeric types in Hist."
                "Or try providing width option."
    options = _Underride(options, label=hist.label)
    options = _Underride(options, align="center")
    if options["align"] == "left":
        options["align"] = "edge"
    elif options["align"] == "right":
        options["align"] = "edge"
        options["width"] *= -1
    Bar(xs, ys, **options)
def Hists(hists, **options):
    """Plots two histograms as interleaved bar plots.
    Options are passed along for all PMFs.  If you want different
    options for each pmf, make multiple calls to Pmf.
      hists: list of two Hist or Pmf objects
      options: keyword args passed to plt.plot
    for hist in hists:
        Hist(hist, **options)
def Pmf(pmf, **options):
    """Plots a Pmf or Hist as a line.
      pmf: Hist or Pmf object
      options: keyword args passed to plt.plot
    xs, ys = pmf.Render()
    width = options.pop("width", None)
    if width is None:
        try:
            width = np.diff(xs).min()
        except TypeError:
            warnings.warn(
                "Pmf: Can't compute bar width automatically."
                "Check for non-numeric types in Pmf."
                "Or try providing width option."
    points = []
    lastx = np.nan
    lasty = 0
    for x, y in zip(xs, ys):
        if (x - lastx) > 1e-5:
            points.append((lastx, 0))
            points.append((x, 0))
        points.append((x, lasty))
        points.append((x, y))
        points.append((x + width, y))
        lastx = x + width
        lasty = y
    points.append((lastx, 0))
    pxs, pys = zip(*points)
    align = options.pop("align", "center")
    if align == "center":
        pxs = np.array(pxs) - width / 2.0
    if align == "right":
        pxs = np.array(pxs) - width
    options = _Underride(options, label=pmf.label)
    Plot(pxs, pys, **options)
def Pmfs(pmfs, **options):
    """Plots a sequence of PMFs.
    Options are passed along for all PMFs.  If you want different
    options for each pmf, make multiple calls to Pmf.
      pmfs: sequence of PMF objects
      options: keyword args passed to plt.plot
    for pmf in pmfs:
        Pmf(pmf, **options)
def Diff(t):
    """Compute the differences between adjacent elements in a sequence.
        t: sequence of number
    Returns:
        sequence of differences (length one less than t)
    diffs = [t[i + 1] - t[i] for i in range(len(t) - 1)]
    return diffs
def Cdf(cdf, complement=False, transform=None, **options):
    """Plots a CDF as a line.
      cdf: Cdf object
      complement: boolean, whether to plot the complementary CDF
      transform: string, one of 'exponential', 'pareto', 'weibull', 'gumbel'
      options: keyword args passed to plt.plot
    Returns:
      dictionary with the scale options that should be passed to
      Config, Show or Save.
    xs, ps = cdf.Render()
    xs = np.asarray(xs)
    ps = np.asarray(ps)
    scale = dict(xscale="linear", yscale="linear")
    for s in ["xscale", "yscale"]:
        if s in options:
            scale[s] = options.pop(s)
    if transform == "exponential":
        complement = True
        scale["yscale"] = "log"
    if transform == "pareto":
        complement = True
        scale["yscale"] = "log"
        scale["xscale"] = "log"
    if complement:
        ps = [1.0 - p for p in ps]
    if transform == "weibull":
        xs = np.delete(xs, -1)
        ps = np.delete(ps, -1)
        ps = [-math.log(1.0 - p) for p in ps]
        scale["xscale"] = "log"
        scale["yscale"] = "log"
    if transform == "gumbel":
        xs = np.delete(xs, 0)
        ps = np.delete(ps, 0)
        ps = [-math.log(p) for p in ps]
        scale["yscale"] = "log"
    options = _Underride(options, label=cdf.label)
    Plot(xs, ps, **options)
    return scale
def Cdfs(cdfs, complement=False, transform=None, **options):
    """Plots a sequence of CDFs.
    cdfs: sequence of CDF objects
    complement: boolean, whether to plot the complementary CDF
    transform: string, one of 'exponential', 'pareto', 'weibull', 'gumbel'
    options: keyword args passed to plt.plot
    for cdf in cdfs:
        Cdf(cdf, complement, transform, **options)
def Contour(obj, pcolor=False, contour=True, imshow=False, **options):
    """Makes a contour plot.
    d: map from (x, y) to z, or object that provides GetDict
    pcolor: boolean, whether to make a pseudocolor plot
    contour: boolean, whether to make a contour plot
    imshow: boolean, whether to use plt.imshow
    options: keyword args passed to plt.pcolor and/or plt.contour
        d = obj.GetDict()
    except AttributeError:
        d = obj
    _Underride(options, linewidth=3, cmap=matplotlib.cm.Blues)
    xs, ys = zip(*d.keys())
    xs = sorted(set(xs))
    ys = sorted(set(ys))
    X, Y = np.meshgrid(xs, ys)
    lam = lambda x, y: d.get((x, y), 0)
    func = np.vectorize(lam)
    Z = func(X, Y)
    x_formatter = matplotlib.ticker.ScalarFormatter(useOffset=False)
    axes = plt.gca()
    axes.xaxis.set_major_formatter(x_formatter)
    if pcolor:
        plt.pcolormesh(X, Y, Z, **options)
    if contour:
        cs = plt.contour(X, Y, Z, **options)
        plt.clabel(cs, inline=1, fontsize=10)
    if imshow:
        extent = xs[0], xs[-1], ys[0], ys[-1]
        plt.imshow(Z, extent=extent, **options)
def Pcolor(xs, ys, zs, pcolor=True, contour=False, **options):
    """Makes a pseudocolor plot.
    pcolor: boolean, whether to make a pseudocolor plot
    contour: boolean, whether to make a contour plot
    options: keyword args passed to plt.pcolor and/or plt.contour
    _Underride(options, linewidth=3, cmap=matplotlib.cm.Blues)
    X, Y = np.meshgrid(xs, ys)
    x_formatter = matplotlib.ticker.ScalarFormatter(useOffset=False)
    axes = plt.gca()
    axes.xaxis.set_major_formatter(x_formatter)
    if pcolor:
        plt.pcolormesh(X, Y, Z, **options)
    if contour:
        cs = plt.contour(X, Y, Z, **options)
        plt.clabel(cs, inline=1, fontsize=10)
def Text(x, y, s, **options):
    """Puts text in a figure.
    x: number
    y: number
    s: string
    options: keyword args passed to plt.text
    options = _Underride(
        options, fontsize=16, verticalalignment="top", horizontalalignment="left"
    plt.text(x, y, s, **options)
LEGEND = True
def Config(**options):
    """Configures the plot.
    Pulls options out of the option dictionary and passes them to
    the corresponding plt functions.
    names = [
        "title",
        "xlabel",
        "ylabel",
        "xscale",
        "yscale",
        "xticks",
        "yticks",
        "axis",
        "xlim",
        "ylim",
    for name in names:
        if name in options:
            getattr(plt, name)(options[name])
    global LEGEND
    LEGEND = options.get("legend", LEGEND)
    # see if there are any elements with labels;
    # if not, don't draw a legend
    ax = plt.gca()
    handles, labels = ax.get_legend_handles_labels()
    if LEGEND and len(labels) > 0:
        global LOC
        LOC = options.get("loc", LOC)
        frameon = options.get("frameon", True)
        try:
            plt.legend(loc=LOC, frameon=frameon)
        except UserWarning:
            pass
    # x and y ticklabels can be made invisible
    val = options.get("xticklabels", None)
    if val is not None:
        if val == "invisible":
            ax = plt.gca()
            labels = ax.get_xticklabels()
            plt.setp(labels, visible=False)
    val = options.get("yticklabels", None)
    if val is not None:
        if val == "invisible":
            ax = plt.gca()
            labels = ax.get_yticklabels()
            plt.setp(labels, visible=False)
def set_font_size(title_size=16, label_size=16, ticklabel_size=14, legend_size=14):
    """Set font sizes for the title, labels, ticklabels, and legend."""
    def set_text_size(texts, size):
        for text in texts:
            text.set_size(size)
    ax = plt.gca()
    # TODO: Make this function more robust if any of these elements
    # is missing.
    # title
    ax.title.set_size(title_size)
    # x axis
    ax.xaxis.label.set_size(label_size)
    set_text_size(ax.xaxis.get_ticklabels(), ticklabel_size)
    # y axis
    ax.yaxis.label.set_size(label_size)
    set_text_size(ax.yaxis.get_ticklabels(), ticklabel_size)
    # legend
    legend = ax.get_legend()
    if legend is not None:
        set_text_size(legend.texts, legend_size)
def bigger_text():
    sizes = dict(title_size=16, label_size=16, ticklabel_size=14, legend_size=14)
    set_font_size(**sizes)
def Show(**options):
    """Shows the plot.
    For options, see Config.
    options: keyword args used to invoke various plt functions
    clf = options.pop("clf", True)
    Config(**options)
    plt.show()
    if clf:
        Clf()
def Plotly(**options):
    """Shows the plot.
    For options, see Config.
    options: keyword args used to invoke various plt functions
    clf = options.pop("clf", True)
    Config(**options)
    import plotly.plotly as plotly
    url = plotly.plot_mpl(plt.gcf())
    if clf:
        Clf()
    return url
def Save(root=None, formats=None, **options):
    """Saves the plot in the given formats and clears the figure.
    For options, see Config.
    Note: With a capital S, this is the original save, maintained for
    compatibility.  New code should use save(), which works better
    with my newer code, especially in Jupyter notebooks.
      root: string filename root
      formats: list of string formats
      options: keyword args used to invoke various plt functions
    clf = options.pop("clf", True)
    save_options = {}
    for option in ["bbox_inches", "pad_inches"]:
        if option in options:
            save_options[option] = options.pop(option)
    # TODO: falling Config inside Save was probably a mistake, but removing
    # it will require some work
    Config(**options)
    if formats is None:
        formats = ["pdf", "png"]
        formats.remove("plotly")
        Plotly(clf=False)
    except ValueError:
        pass
    if root:
        for fmt in formats:
            SaveFormat(root, fmt, **save_options)
    if clf:
        Clf()
def save(root, formats=None, **options):
    """Saves the plot in the given formats and clears the figure.
    For options, see plt.savefig.
      root: string filename root
      formats: list of string formats
      options: keyword args passed to plt.savefig
    if formats is None:
        formats = ["pdf", "png"]
        formats.remove("plotly")
        Plotly(clf=False)
    except ValueError:
        pass
    for fmt in formats:
        SaveFormat(root, fmt, **options)
def SaveFormat(root, fmt="eps", **options):
    """Writes the current figure to a file in the given format.
      root: string filename root
      fmt: string format
    _Underride(options, dpi=300)
    filename = "%s.%s" % (root, fmt)
    print("Writing", filename)
    plt.savefig(filename, format=fmt, **options)
# provide aliases for calling functions with lower-case names
preplot = PrePlot
subplot = SubPlot
figure = Figure
plot = Plot
vlines = Vlines
hlines = Hlines
fill_between = FillBetween
text = Text
scatter = Scatter
pmfs = Pmfs
hist = Hist
hists = Hists
diff = Diff
cdfs = Cdfs
contour = Contour
pcolor = Pcolor
config = Config
show = Show
def main():
    color_iter = _Brewer.ColorGenerator(7)
    for color in color_iter:
        print(color)
if __name__ == "__main__":
Sunbelt Computer Software

PL/B Language Development and Support

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Sunbelt Computer Software

PL/B Language Development and Support

FilesExpand file tree

thinkplot.py

Latest commit

History

thinkplot.py

File metadata and controls
