path: root/gnuplot.py

#!/usr/bin/env python
import defapp
from subprocess import Popen, PIPE
from optparse import make_option as o

FORMATS = ['png', 'eps', 'pdf', 'show', 'ps']
STYLES  = ['points', 'lines', 'linespoints']

TERMINALS = {
    'png' : 'png',
    'eps' : 'postscript eps',
    'ps'  : 'postscript',
    'pdf' : 'pdf',
    }

class CommandBuffer(object):
    def __init__(self):
        self.cmds = []

    def __call__(self, cmd):
        self.cmds += [cmd]

    def __str__(self):
        return '\n'.join([str(x) for x in self.cmds])


class Label(object):
    def __init__(self, x, y, txt):
        self.x   = x
        self.y   = y
        self.txt = txt
        # support more label stuff optionally...

    def gnuplot_cmd(self):
        return '"%s" at %f,%f center' % (self.txt, self.x, self.y)

    def __str__(self):
        return 'set label %s' % self.gnuplot_cmd()


class FileGraph(object):
    def __init__(self, fname, xcol=1, ycol=2, error=None,
                 title=None, style=None, smooth=None):
        self.fname = fname
        self.xcol  = xcol
        self.ycol  = ycol
        self.error = error
        self.title = title
        self.style = style # replace with linetype, lines/points etc.
        self.smooth = smooth

    def __str__(self):
        return self.gnuplot_cmd()

    def gnuplot_cmd(self, default_style=None):
        if self.error:
            using_txt = "%s:%s:%s" % (self.xcol, self.ycol, self.error)
            style     = self.style if self.style else "errorbars"
        else:
            using_txt = "%s:%s" % (self.xcol, self.ycol)
            style     = self.style if self.style else default_style
        style_txt = " with %s" % style if style else ""
        title_txt = ' title "%s"' % self.title if self.title else ""
        smooth_txt = ' smooth %s' % self.smooth if self.smooth else ""
        return "'%s' using %s%s%s%s" % \
            (self.fname, using_txt, title_txt, style_txt, smooth_txt)


class LiteralGraph(object):
    def __init__(self, literal_expr, title=None, style=None):
        self.expr  = literal_expr
        self.title = title
        self.style = style

    def __str__(self):
        return self.gnuplot_cmd()

    def gnuplot_cmd(self, default_style=None):
        style     = self.style if self.style else default_style
        style_txt = " with %s" % style if style else ""
        title_txt = " title '%s'" % self.title if self.title else ""
        return "%s %s%s" % \
            (self.expr, title_txt, style_txt)


label = Label

def curve(fname=None, literal=None, **kargs):
    if fname:
        return FileGraph(fname, **kargs)
    elif literal:
        return LiteralGraph(literal, **kargs)


class Plot(object):
    def __init__(self):
        self.format  = None
        self.output  = None
        self.size    = None

        self.font      = None
        self.font_size = None
        self.enhanced_text = None

        self.monochrome   = None
        self.dashed_lines = None
        self.rounded_caps = None

        self.xrange = None
        self.xticks = None
        self.xlog   = None
        self.xlabel = None

        self.yrange = None
        self.yticks = None
        self.ylog   = None
        self.ylabel = None

        self.key    = None
        self.title  = None

        self.labels = []
        self.curves = []

        self.default_style = None # for plotted curves
        self.line_styles = []

    def gnuplot_commands(self, cmd_buf=None):
        if cmd_buf:
            g = cmd_buf
        else:
            g = CommandBuffer()

        def isset(x):
            return not (x is None)

        if isset(self.format):
            term = [TERMINALS[self.format]]
        else:
            term = 'x11'

        ps_like = self.format in ['eps', 'pdf', 'ps']

        if isset(self.monochrome) and ps_like:
            term += ['monochrome' if self.monochrome else 'color']

        if isset(self.enhanced_text):
            term += ['enhanced' if self.enhanced_text else 'noenhanced']

        if isset(self.dashed_lines) and ps_like:
            term += ['dashed' if self.dashed_lines else 'solid']
        if isset(self.font) and ps_like:
            term += ['font "%s%s"' %
                     (self.font,
                      (",%s" % self.font_size) if isset(self.font_size)
                      else "")
                     ]

        if isset(self.rounded_caps):
            term += ['rounded' if self.rounded_caps else 'butt']

        if isset(self.size):
            term += ['size %s,%s' % (self.size)]

        if isset(self.font_size) and self.format == 'png':
            term += [self.font_size]

        g("set terminal %s" % " ".join(term))
        if isset(self.output):
            g("set out '/dev/null'")

        if self.xrange:
            g("set xrange [%s:%s]" % self.xrange)
        if self.xticks:
            g("set xtics %s, %s" % self.xticks)
        if self.xlabel:
            g("set xlabel '%s'" % self.xlabel)

        if self.yrange:
            g("set yrange [%s:%s]" % self.yrange)
        if self.yticks:
            g("set ytics %s, %s" % self.yticks)
        if self.ylabel:
            g("set ylabel '%s'" % self.ylabel)

        if self.key:
            g('set key %s' % self.key)
        if self.title:
            g('set title "%s"' % self.title)

        logscale = ""
        if self.xlog:
            logscale += "x"
        if self.ylog:
            logscale += "y"
        if logscale:
            g("set logscale %s" % logscale)

        for ls in self.line_styles:
            g("set style line %d %s" % (ls[0], ls[1]))

        plots = [c.gnuplot_cmd(self.default_style) for c in self.curves]
        if plots:
            g("plot " + ", ".join(plots))

        for l in self.labels:
            g("set label " + l.gnuplot_cmd())

        if isset(self.output):
            g('set out "%s"' % self.output)

        g("replot")

        if isset(self.output):
            g("set out") # close file

        return g


    def gnuplot_exec(self):
        pipe2gnuplot(self.gnuplot_commands())

    def gnuplot_save(self, fname):
        f = open(fname, 'w')
        f.write(str(self.gnuplot_commands()))
        f.close()

def gnuplot_cmd(graphs, title=None, ylabel=None, xlabel=None,
                format='show', term_opts=None,
                style='linespoints', xrange=None,
                yrange=None,
                xticks=None, yticks=None,
                labels=[],
                key='below',
                logscale=None,
                fname=None):
    g = CommandBuffer()
    if format == 'png':
        terminal = 'png'
        if term_opts is None:
            term_opts = 'size 1024,768 large'
    elif format == 'eps':
        terminal = 'postscript eps'
        if term_opts is None:
            term_opts = 'color blacktext solid linewidth 1.0'
    elif format == 'ps':
        terminal = 'postscript'
        if term_opts is None:
            term_opts = 'color enhanced'
    elif format == 'pdf':
        terminal = 'pdf'
        if term_opts is None:
            term_opts = 'color enhanced'
    if format != 'show':
        g('set terminal %s %s' % (terminal, term_opts))
        g("set out '/dev/null'")
    if xlabel:
        g("set xlabel '%s'" % xlabel)
    if ylabel:
        g("set ylabel '%s'" % ylabel)
    if title:
        g("set title '%s'" % title)
    if xrange:
        g("set xrange [%s:%s]" % xrange)
    if yrange:
        g("set yrange [%s:%s]" % yrange)
    if xticks:
        if isinstance(xticks, str):
            g('set xtics {0}'.format(xticks))
        else:
            g("set xtics %s, %s" % xticks)
    if yticks:
        g("set ytics %s, %s" % yticks)
    if logscale:
        if type(logscale) == tuple:
            for x in logscale:
                g("set logscale %s" % x)
        else:
            g("set logscale %s" % logscale)
    g('set key %s' % key)
    plot = []
    for gr in graphs:
        if type(gr) == str:
            # literal plot command
            plot.append(str(gr))
        elif type(gr) == FileGraph:
            # formatter object
            plot.append(gr.gnuplot_cmd(style))
        elif len(gr) == 4:
            par = (gr[0], gr[1], gr[2], gr[3], style)
            plot += ["'%s' using %s:%s title '%s' with %s" % par]
        elif len(gr) == 6:
            par = (gr[0], gr[1], gr[2], gr[3], gr[4], style, gr[5])
            plot += ["'%s' using %s:%s:%s:%s with %s title '%s'" % par]
        elif len(gr) == 5:
            plot += ["'%s' using %s:%s title '%s' with %s" % gr]
        elif len(gr) == 3:
            plot += ["%s title '%s' with %s" % gr]
    if plot:
        g('plot ' + ', '.join(plot))
    for (x, y, txt) in labels:
        g('set label "%s" at %f,%f center' % (txt, x, y))
    if format != 'show' and fname:
        g("set out '%s.%s'" % (fname, format))
    if plot:
        g('replot')
    if format != 'show' and fname:
        g('set out')
    return g

def pipe2gnuplot(cmds):
    proc = Popen(['gnuplot'], stdin=PIPE)
    proc.stdin.write(str(cmds))
    proc.stdin.close()
    proc.wait()

def gnuplot(*args, **kargs):
        cmd = gnuplot_cmd(*args, **kargs)
        pipe2gnuplot(cmd)

def eps2pdf(file):
    Popen(['ps2pdf', '-dEPSCrop', '%s.eps' % file]).wait()

options = [
    o('-f', '--format', action='store', dest='format', type='choice',
      choices=FORMATS, help='output format'),
    o('-o', '--output', action='store', dest='out', help='Output file name.'),
    o('-s', '--style', action='store', dest='style', type='choice',
      choices=STYLES, help='line style'),
    o('-t', '--title', action='store', dest='title'),
    o(None, '--xlabel', action='store', dest='xlabel'),
    o(None, '--ylabel', action='store', dest='ylabel'),
    o(None, '--xrange', action='store', dest='xrange', nargs=2, type='float'),
    o(None, '--yrange', action='store', dest='yrange', nargs=2, type='float'),
    o(None, '--xticks', action='store', dest='xticks', nargs=2, type='float'),
    o(None, '--yticks', action='store', dest='yticks', nargs=2, type='float'),
    o(None, '--term-opts', action='store', dest='term_opts'),
    ]

defaults = {
    'out'    : 'graph',
    'format' : 'show',
    'style'  : 'linespoints',
    'title'  : None,
    'xlabel' : None,
    'ylabel' : None,
    'xrange' : None,
    'yrange' : None,
    'xticks' : None,
    'yticks' : None,
    'term_opts' : None,
    }

class GnuPlotter(defapp.App):
    def __init__(self):
        defapp.App.__init__(self, options, defaults, no_std_opts=True)

    def make_cmd(self, graphs):
        return gnuplot_cmd(graphs,
                           title=self.options.title,
                           format=self.options.format,
                           style=self.options.style,
                           xlabel=self.options.xlabel,
                           ylabel=self.options.ylabel,
                           xrange=self.options.xrange,
                           yrange=self.options.yrange,
                           xticks=self.options.xticks,
                           yticks=self.options.yticks,
                           fname=self.options.out,
                           term_opts=self.options.term_opts)

    def get_cmd(self, args):
        graphs = []
        while len(args) >= 4:
            g =  args[0:4]
            graphs += [g]
            args = args[4:]
        if args:
            self.err("Warning: Ignoring trailing args. Args:", *args)
        return self.make_cmd(graphs)

    def default(self, _):
        cmd = self.get_cmd(list(self.args))
        pipe2gnuplot(cmd)

    def do_cmd(self, _):
        cmd = self.get_cmd(self.args[1:])
        self.out(cmd)

    def do_linedemo(self, _):
        graphs = ["%d title 'ls %d'" % (x, x) #, x) #  with linespoints ls %d
                  for x in xrange(1,10)]
        graphs += ["%d title 'lt %d' with lines lt %d" % (10 + x, x, x)
                   for x in xrange(1,10)]
        self.options.yrange = (0, 20)
        self.options.xrange = (0, 10)
        self.options.xticks = (0, 10)
        self.options.yticks = (0, 1)
        self.options.title = "%s Default Line Styles" % self.options.format.upper()
        self.options.ylabel = "Styles"
        self.options.term_opts = "" # avoid default

        cmd = self.make_cmd(graphs)
        pipe2gnuplot(cmd)


if __name__ == "__main__":
    GnuPlotter().launch()