fixed buts w/ facet_wrap

Author: Greg

ggplot/geoms/facet_grid.py

@@ -49,8 +49,11 @@ def __radd__(self, gg):
 
         combos = []
         for x_i in sorted(x.unique()):
-            for y_i in sorted(y.unique()):
-                combos.append((x_i, y_i))
+            if y is not None:
+                for y_i in sorted(y.unique()):
+                    combos.append((x_i, y_i))
+            else:
+                combos.append((x_i, 1))
         gg.facet_pairs = combos
 
         return deepcopy(gg)

ggplot/geoms/facet_wrap.py

@@ -2,20 +2,21 @@
 import math
 
 class facet_wrap(object):
-    def __init__(self, x=None, ncol=None, nrow=None, scales="free"):
+    def __init__(self, x=None, y=None, ncol=None, nrow=None, scales="free"):
         self.x = x
+        self.y = y
         self.ncol = ncol
         self.nrow = nrow
         self.scales = scales
 
     def __radd__(self, gg):
 
         x = gg.data.get(self.x)
-        if x is None:
-            n_wide = 1
-        else:
-            n_wide = x.nunique()
+        y = gg.data.get(self.y)
+        if x is not None:
             gg.n_dim_x = x.nunique()
+        if y is not None:
+            gg.n_dim_x *= y.nunique()
 
         n_wide = self.ncol
         n_high = self.nrow
@@ -34,6 +35,8 @@ def __radd__(self, gg):
         gg.n_wide, gg.n_high = int(n_wide), int(n_high)
         facets = []
         facets.append(self.x)
+        if self.y:
+            facets.append(self.y)
         gg.facets = facets
         gg.facet_type = "wrap"
         gg.facet_scales = self.scales

ggplot/ggplot.py

@@ -115,7 +115,6 @@ def __init__(self, aesthetics, data):
         self.colormap = plt.cm.Blues
 
     def __repr__(self):
-        # TODO: Currently plotting the wrong number of facets.
         if self.facet_type=="grid":
             fig, axs = plt.subplots(self.n_high, self.n_wide,
                     sharex=True, sharey=True)
@@ -128,13 +127,10 @@ def __repr__(self):
             for extra_plot in axs.flatten()[-extra_subplots:]:
                 extra_plot.axis('off')
 
-            # TODO: This isn't working
-            plots = [None for i in range(self.n_dim_x)]
-            for i in range(self.n_dim_x):
-                idx = (i % self.n_high) * self.n_wide + (i % self.n_wide)
-                plots[idx] = (i % self.n_wide, i % self.n_high)
-
-            plots = [plot for plot in plots if plot is not None]
+            plots = []
+            for x in range(self.n_wide):
+                for y in range(self.n_high):
+                    plots.append((x, y))
             plots = sorted(plots, key=lambda x: x[1] + x[0] * self.n_high + 1)
 
         else:
@@ -146,7 +142,7 @@ def __repr__(self):
         # dimensions of the plot remain the same
         if self.facets:
             cntr = 0
-            if len(self.facets)==2:
+            if len(self.facets)==2 and self.facet_type!="wrap":
                 # store the extreme x and y coordinates of each pair of axes			
                 axis_extremes = np.zeros(shape=(self.n_high * self.n_wide, 4))
                 xlab_offset = .15
@@ -183,10 +179,11 @@ def __repr__(self):
                         )
                     plt.subplot(self.n_wide, self.n_high, pos)
 
-                # Handle the different scale types here (free|free_y|free_x|None) and 
-                # also make sure that only the left column gets y scales and the bottom
-                # row gets x scales
-                scale_facet(self.n_wide, self.n_high, self.facet_pairs, self.facet_scales)
+                # Handle the different scale types here 
+                # (free|free_y|free_x|None) and also make sure that only the
+                # left column gets y scales and the bottom row gets x scales
+                scale_facet_grid(self.n_wide, self.n_high, 
+                        self.facet_pairs, self.facet_scales)
 
             else:
                 for facet, frame in self.data.groupby(self.facets):
@@ -213,9 +210,13 @@ def __repr__(self):
                                 for callback in callbacks:
                                     fn = getattr(axs[cntr], callback['function'])
                                     fn(*callback['args'])
-                    #TODO: selective titles
-                    plt.title(facet)
+                    title = facet
+                    if isinstance(facet, tuple):
+                        title = ", ".join(facet)
+                    plt.table(cellText=[[title]], loc='top',
+                            cellLoc='center', cellColours=[['lightgrey']])
                     cntr += 1
+                scale_facet_wrap(self.n_wide, self.n_high, [], self.facet_scales)
         else:
             for layer in self._get_layers(self.data):
                 for geom in self.geoms:

ggplot/scales/__init__.py

@@ -1,7 +1,8 @@
 from .scale_colour_gradient import scale_colour_gradient
 from .scale_colour_gradient import scale_colour_gradient as scale_color_gradient
 from .scale_colour_gradient import scale_colour_gradient as scale_colour_gradient2
-from .scale_facet import scale_facet
+from .scale_facet import scale_facet_grid
+from .scale_facet import scale_facet_wrap
 from .scale_x_continuous import scale_x_continuous
 from .scale_x_date import scale_x_date
 from .scale_y_continuous import scale_y_continuous

ggplot/scales/scale_facet.py

@@ -4,8 +4,15 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from .utils import calc_axis_breaks
+import sys
 
-def scale_facet(xdim, ydim, facet_pairs, scaletype):
+
+def scale_facet_wrap(xdim, ydim, positions, scaletype):
+    if scaletype is not None:
+        sys.stderr.write("facet_wrap scales not yet implemented!\n")
+
+
+def scale_facet_grid(xdim, ydim, facet_pairs, scaletype):
     # everyone gets the same scales
     if scaletype is None:
         min_x, max_x = 999999999, -999999999