Source code for mmf.plotting.matplotlib.figures

"""Helper functions for drawing figures for inclusion in LaTeX
documents using matplotlib.

BUGS:
 Ensure that latex figures do not have unicode!

"""

import os
import unicodedata

import pylab

_latex_template = r"""
\begin{figure}[htb]
  \begin{centering}
    %(psfrags)s
    \includegraphics[width=%(width)s\columnwidth]{%(filename)s}
    \caption{\label{fig:%(name)s}
    } 
  \end{centering}
\end{figure}
"""
_usetex = False

class Figures(object):
    """Container of a bunch of figures."""
    def __init__(self):
        self.figs = {}
        self.LaTeX = {}

    def setup_figure(self,name=None,
                     width=1.0,height=1.0,
                     margin_factors=None,
                     col_width_pt=246.0):
        """
        Use this to setup the figure properties.
        :Parameters:
            width : Fraction of columnwidth for figure
            height : Fraction of natural height (in terms of golden mean)
            col_width_pt : Get this from LaTeX using \showthe\columnwidth
        """
        fig = self.figs.get(name, pylab.figure())

        # Custom fields added to fig.
        fig.name = name
        fig.frags = {}          # \psfrag pairs.
        fig.width = width       # Width

        if fig.name is None:
            fig.name = "fig_%i"%(fig.number,)

        inches_per_pt = 1.0/72.27               # Convert pt to inch
        golden_mean = (pylab.sqrt(5)-1.0)/2.0   # Aesthetic ratio
        col_width = col_width_pt*inches_per_pt  # width in inches
        fig_height = col_width*golden_mean      # height in inches
        fig_width = width*col_width
        fig_height = height*fig_height
        fig_size =  [fig_width,fig_height]

        # Now we need to allocate space for the labels.  These should be
        # absolute because the text size does not change.
        m_factors = {'top': 0.5,
                     'left': 2.8,
                     'bot': 2.0,
                     'right': 0.5}

        if margin_factors is not None:
            m_factors.update(margin_factors)

        space_top = m_factors['top']*12.0*inches_per_pt
        space_left = m_factors['left']*12.0*inches_per_pt
        space_bottom = m_factors['bot']*12.0*inches_per_pt
        space_right = m_factors['right']*12.0*inches_per_pt

        # Compute axes size:
        axes_left = space_left/fig_width
        axes_bottom = space_bottom/fig_height
        axes_width = 1.0-(space_left+space_right)/fig_width
        axes_height = 1.0-(space_bottom+space_top)/fig_height

        axes_size = [axes_left,axes_bottom,
                     axes_width,axes_height]

        params = {'backend': 'ps',
                  'axes.labelsize': 8,
                  'text.fontsize': 10,
                  'legend.fontsize': 8,
                  'xtick.labelsize': 8,
                  'ytick.labelsize': 8,
                  'text.usetex': _usetex,
                  'figure.figsize': fig_size}
        pylab.rcParams.update(params)
        fig.set_size_inches(fig_size)

        pylab.figure(fig.number)
        pylab.clf()
        pylab.axes(axes_size)
        self.figs[fig.name] = fig

    def savefig(self, ext="eps"):
        psfrags = []

        fig = pylab.gcf()
        name = fig.name
        filename = ".".join([name,ext])
        for ax in fig.get_axes():
            for t in ax.get_xticklabels():
                text = t.get_text()
                fig.frags[text] = r"\footnotesize{%s}"%(text,)

            for t in ax.get_yticklabels():
                text = t.get_text()
                fig.frags[text] = r"\footnotesize{%s}"%(text,)
                
        for frag in fig.frags:
            psfrags.append(r"\psfrag{%s}{%s}"%(
                    frag, fig.frags[frag]))
        psfrags = "\n    ".join(psfrags)
        self.LaTeX[name] = _latex_template%\
            dict(name=name,
                 filename=filename,
                 width=fig.width,
                 psfrags=psfrags)
        
        pylab.savefig(filename)
        try:
            print self.LaTeX[name]
        except UnicodeEncodeError:
            print unicodedata.normalize('NFKD',
                                        self.LaTeX[name]).\
                                        encode('ascii','ignore')
        
    def text(self, text, latex=None):
        """Use this to format text for labels etc."""
        fig = pylab.gcf()
        if latex is None:
            latex = text
        fig.frags[text] = latex
        return text