Bug fix in thinkbayes.py

myplot.py

@@ -13,7 +13,7 @@
 # customize some matplotlib attributes
 #matplotlib.rc('figure', figsize=(4, 3))
 
-matplotlib.rc('font', size=14.0)
+#matplotlib.rc('font', size=14.0)
 #matplotlib.rc('axes', labelsize=22.0, titlesize=22.0)
 #matplotlib.rc('legend', fontsize=20.0)
 
@@ -87,14 +87,29 @@ def GetIter(cls):
         return cls.color_iter
 
 
-def PrePlot(num=None):
+def PrePlot(num=None, rows=1, cols=1):
     """Takes hints about what's coming.
 
     num: number of lines that will be plotted
     """
-    Brewer.InitializeIter(num)
+    if num:
+        Brewer.InitializeIter(num)
+
+    # TODO: get sharey and sharex working.  probably means switching
+    # to subplots instead of subplot.
+    # also, get rid of the gray background.
+
+    if rows > 1 or cols > 1:
+        pyplot.subplots(rows, cols, sharey=True)
+        global SUBPLOT_ROWS, SUBPLOT_COLS
+        SUBPLOT_ROWS = rows
+        SUBPLOT_COLS = cols
 
 
+def SubPlot(plot_number):
+    pyplot.subplot(SUBPLOT_ROWS, SUBPLOT_COLS, plot_number)
+
+
 class InfiniteList(list):
     """A list that returns the same value for all indices."""
     def __init__(self, val):
@@ -353,6 +368,32 @@ def Contour(obj, pcolor=False, contour=True, imshow=False, **options):
         pyplot.imshow(Z, extent=extent, **options)
 
 
+def Pcolor(xs, ys, zs, pcolor=True, contour=False, **options):
+    """Makes a pseudocolor plot.
+    
+    xs:
+    ys:
+    zs:
+    pcolor: boolean, whether to make a pseudocolor plot
+    contour: boolean, whether to make a contour plot
+    options: keyword args passed to pyplot.pcolor and/or pyplot.contour
+    """
+    Underride(options, linewidth=3, cmap=matplotlib.cm.Blues)
+
+    X, Y = np.meshgrid(xs, ys)
+    Z = zs
+
+    x_formatter = matplotlib.ticker.ScalarFormatter(useOffset=False)
+    axes = pyplot.gca()
+    axes.xaxis.set_major_formatter(x_formatter)
+
+    if pcolor:
+        pyplot.pcolormesh(X, Y, Z, **options)
+
+    if contour:
+        cs = pyplot.contour(X, Y, Z, **options)
+        pyplot.clabel(cs, inline=1, fontsize=10)
+
 
 def Config(**options):
     """Configures the plot.
@@ -374,13 +415,13 @@ def Config(**options):
         pyplot.xscale(options['xscale'])
 
     if 'xticks' in options:
-        pyplot.xticks(*options['xticks'])
+        pyplot.xticks(options['xticks'])
 
     if 'yscale' in options:
         pyplot.yscale(options['yscale'])
 
     if 'yticks' in options:
-        pyplot.yticks(*options['yticks'])
+        pyplot.yticks(options['yticks'])
 
     if 'axis' in options:
         pyplot.axis(options['axis'])
@@ -436,6 +477,27 @@ def SaveFormat(root, fmt='eps'):
     pyplot.savefig(filename, format=fmt, dpi=300)
 
 
+# provide aliases for calling functons with lower-case names
+preplot = PrePlot
+subplot = SubPlot
+clf = Clf
+figure = Figure
+plot = Plot
+scatter = Scatter
+pmf = Pmf
+pmfs = Pmfs
+hist = Hist
+hists = Hists
+diff = Diff
+cdf = Cdf
+cdfs = Cdfs
+contour = Contour
+pcolor = Pcolor
+config = Config
+show = Show
+save = Save
+
+
 def main():
     color_iter = Brewer.ColorGenerator(7)
     for color in color_iter:

redline.py

@@ -14,7 +14,7 @@
 import random
 import sys
 
-FORMATS = ['pdf', 'eps', 'png']
+FORMATS = ['pdf', 'eps', 'png', 'jpg']
 
 """
 Notation guide:

thinkbayes.py

@@ -29,6 +29,16 @@
 import scipy.stats
 from scipy.special import erf, erfinv
 
+ROOT2 = math.sqrt(2)
+
+def RandomSeed(x):
+    """Initialize the random and numpy.random generators.
+
+    x: int seed
+    """
+    random.seed(x)
+    numpy.random.seed(x)
+
 
 def Odds(p):
     """Computes odds for a given probability.
@@ -385,13 +395,88 @@ def MakeCdf(self, name=None):
         return MakeCdfFromPmf(self, name=name)
 
     def ProbGreater(self, x):
+        """Probability that a sample from this Pmf exceeds x.
+
+        x: number
+
+        returns: float probability
+        """
         t = [prob for (val, prob) in self.d.iteritems() if val > x]
         return sum(t)
 
     def ProbLess(self, x):
+        """Probability that a sample from this Pmf is less than x.
+
+        x: number
+
+        returns: float probability
+        """
         t = [prob for (val, prob) in self.d.iteritems() if val < x]
         return sum(t)
 
+    def __lt__(self, obj):
+        """Less than.
+
+        obj: number or _DictWrapper
+
+        returns: float probability
+        """
+        if isinstance(obj, _DictWrapper):
+            return PmfProbLess(self, obj)
+        else:
+            return self.ProbLess(obj)
+
+    def __gt__(self, obj):
+        """Greater than.
+
+        obj: number or _DictWrapper
+
+        returns: float probability
+        """
+        if isinstance(obj, _DictWrapper):
+            return PmfProbGreater(self, obj)
+        else:
+            return self.ProbGreater(obj)
+
+    def __ge__(self, obj):
+        """Greater than or equal.
+
+        obj: number or _DictWrapper
+
+        returns: float probability
+        """
+        return 1 - (self < obj)
+
+    def __le__(self, obj):
+        """Less than or equal.
+
+        obj: number or _DictWrapper
+
+        returns: float probability
+        """
+        return 1 - (self > obj)
+
+    def __eq__(self, obj):
+        """Less than.
+
+        obj: number or _DictWrapper
+
+        returns: float probability
+        """
+        if isinstance(obj, _DictWrapper):
+            return PmfProbEqual(self, obj)
+        else:
+            return self.Prob(obj)
+
+    def __ne__(self, obj):
+        """Less than.
+
+        obj: number or _DictWrapper
+
+        returns: float probability
+        """
+        return 1 - (self == obj)
+
     def Normalize(self, fraction=1.0):
         """Normalizes this PMF so the sum of all probs is fraction.
 
@@ -1453,7 +1538,7 @@ def EvalPoissonPmf(k, lam):
 
     returns: float probability
     """
-    # don't use the scipy function.  for lam=0 it returns NaN;
+    # don't use the scipy function (yet).  for lam=0 it returns NaN;
     # should be 0.0
     # return scipy.stats.poisson.pmf(k, lam)
 
@@ -1509,7 +1594,7 @@ def MakeExponentialPmf(lam, high, n=200):
     return pmf
 
 
-def StandardGaussianCdf(x, root2=math.sqrt(2)):
+def StandardGaussianCdf(x):
     """Evaluates the CDF of the standard Gaussian distribution.
     
     See http://en.wikipedia.org/wiki/Normal_distribution
@@ -1521,7 +1606,7 @@ def StandardGaussianCdf(x, root2=math.sqrt(2)):
     Returns:
         float
     """
-    return (erf(x / root2) + 1) / 2
+    return (erf(x / ROOT2) + 1) / 2
 
 
 def GaussianCdf(x, mu=0, sigma=1):
@@ -1555,7 +1640,7 @@ def GaussianCdfInverse(p, mu=0, sigma=1):
     Returns:
         float
     """
-    x = root2 * erfinv(2 * p - 1)
+    x = ROOT2 * erfinv(2 * p - 1)
     return mu + x * sigma
 
 
@@ -1587,6 +1672,14 @@ def Random(self):
         """Generates a random variate from this distribution."""
         return random.betavariate(self.alpha, self.beta)
 
+    def Sample(self, n):
+        """Generates a random sample from this distribution.
+
+        n: int sample size
+        """
+        size = n,
+        return numpy.random.beta(self.alpha, self.beta, size)
+
     def EvalPdf(self, x):
         """Evaluates the PDF at x."""
         return x ** (self.alpha - 1) * (1 - x) ** (self.beta - 1)