subplot_share_xy.py

#!/usr/bin/env python

import matplotlib.pyplot as plt
import numpy

class SublotsForReuse(object):
    def __init__(self, *args, **kwds):
        self.fig = plt.figure(*args, **kwds)
        self.ax = []

    def add_axes(self, *args, **kwds):
        a = self.fig.add_axes(*args, **kwds)
        self.ax.append(a)
        return a

    def cla_all(self):
        for a in self.ax:
            a.cla()

    def clf(self):
        self.fig.clf()
        self.ax = []

def normalize_ratio(ratio, default_num):
    if ratio is None:
        ratio = [1] * default_num
    ratio = numpy.array(ratio, dtype=float )
    return ratio / ratio.sum()

class SublotsMultipleTimeSeries(SublotsForReuse):
    def __init__(self, linedict, keytable, xratio=None, yratio=None, steps=None,
                 xsize_per_steps=0.01, ysize_per_ynum=1.5,
                 xsize_blank=1.6, ysize_blank=1.2,
                 *args, **kwds):
        if steps is None:
            steps = sum([len(linedict[k][0]) for k in keytable[0]])
        
        self.linedict = linedict
        self.keytable = keytable
        self.keys = linedict.keys()
        self.k2i = dict([(k,i) for (i,k) in enumerate(self.keys)])
        self.xnum = keytable.shape[1]
        self.ynum = keytable.shape[0]
        ynum  = self.ynum
        self.figsize = ( xsize_per_steps * steps + xsize_blank,
                         ysize_per_ynum  * ynum  + ysize_blank )
        
        self.xratio = normalize_ratio(xratio, self.xnum)
        self.yratio = normalize_ratio(yratio, self.ynum)
        self.prepare_rect()
        self.set_rect()
        
        SublotsForReuse.__init__(self, figsize=self.figsize, *args, **kwds)
        self.set_axes()

    def add_axes(self, i, *args, **kwds):
        self.ax[i] = self.fig.add_axes(*args, **kwds)
        return self.ax[i]

    def set_axes(self):
        self.ax = [None] * self.xnum * self.ynum
        a_below = [None] * self.xnum
        for y in xrange(self.ynum):
            for x in xrange(self.xnum):
                i = self.k2i[self.keytable[y,x]]
                r = self.rect_by_loc(x,y)
                if y == 0:
                    if x == 0:
                        a = self.add_axes(i,r)
                        a_left = a
                    else:
                        a = self.add_axes(i,r, sharey=a_left )
                        a.set_yticklabels([])
                    a_below[x] = a
                else:
                    if x == 0:
                        a = self.add_axes(i,r, sharex=a_below[x] )
                        #a.set_xticks([])
                        a_left = a
                    else:
                        a = self.add_axes(i,r, sharex=a_below[x], sharey=a_left )
                        #a.set_xticks([])
                        #a.set_yticks([])
                

    def plot_all(self):
        for (key, ax) in zip(self.keys, self.ax):
            ax.plot(*self.linedict[key])

    xspace = 0.1
    yspace = 0.1
    def prepare_rect(self):
        self.width      = self.xratio * (1-self.xspace)
        self.height     = self.yratio * (1-self.yspace)
        self.left       = numpy.zeros_like(self.width )
        self.bottom     = numpy.zeros_like(self.height)
        self.left  [0]  = self.xspace * 0.5
        self.bottom[0]  = self.yspace * 0.5
        self.left  [1:] = numpy.cumsum(self.width [:-1]) + self.xspace * 0.5
        self.bottom[1:] = numpy.cumsum(self.height[:-1]) + self.yspace * 0.5

    def find_loc(self,key):
        for x in xrange(self.xnum):
            for y in xrange(self.ynum):
                if key == self.keytable[y,x]:
                    return (x,y)

    def rect_by_loc(self, x,y):
        return [self.left[x], self.bottom[y], self.width[x], self.height[y]]

    def rect_by_key(self, key):
        return self.rect_by_loc(*self.find_loc(key))
    
    def set_rect(self):
        self.rect = []
        for key in self.keys:
            self.rect.append(self.rect_by_key(key))

num = 3
keytable = numpy.array([
    ['a%d' % i for i in xrange(num) ], 
    ['b%d' % i for i in xrange(num) ], 
    ['c%d' % i for i in xrange(num) ], 
    ['d%d' % i for i in xrange(num) ], 
    ])

x = numpy.arange(0.0, 2.0, 0.01)
linedict = dict(
    [(k, [x, numpy.random.ranf(len(x))]) for k in keytable.ravel()]
    )

smts = SublotsMultipleTimeSeries(linedict, keytable, steps=len(x)*num)
smts.plot_all()
smts.fig.show()