Initial project version
This commit is contained in:
sebastiankloth
355
pygrace/pygrace.py
Normal file
355
pygrace/pygrace.py
Normal file
@ -0,0 +1,355 @@
|
||||
import numpy as np
|
||||
from matplotlib.colors import ColorConverter
|
||||
from matplotlib.backends import backend_agg
|
||||
import matplotlib.pyplot as plt
|
||||
import textwrap
|
||||
from itertools import zip_longest
|
||||
from collections import defaultdict
|
||||
|
||||
from . import tex2grace
|
||||
from .tex2grace import latex_to_xmgrace
|
||||
|
||||
|
||||
def update_labels(labels, axis=None):
|
||||
if axis is None:
|
||||
axis = plt.gca()
|
||||
for line, lb in zip_longest(axis.lines, labels, fillvalue=''):
|
||||
line.set_label(lb)
|
||||
|
||||
|
||||
def sanitize_strings(dic):
|
||||
"""Do some sanitization for strings."""
|
||||
for k in dic:
|
||||
if isinstance(dic[k], str):
|
||||
dic[k].replace('{', '{{')
|
||||
dic[k].replace('}', '}}')
|
||||
|
||||
|
||||
def get_world_coords(artist):
|
||||
"""Get the world coordinates of an artist."""
|
||||
trans = artist.axes.transData.inverted()
|
||||
return trans.transform(artist.get_window_extent())
|
||||
|
||||
|
||||
class AgrText:
|
||||
fmt = """string on
|
||||
string loctype view
|
||||
string {position}
|
||||
string char size {size:.2f}
|
||||
string color {color}
|
||||
string def "{value}"
|
||||
"""
|
||||
|
||||
def get_position(self):
|
||||
trans = self.agr_axis.agr_figure.figure.transFigure.inverted()
|
||||
pos = trans.transform(self.text.get_window_extent())[0]
|
||||
pos = (pos + self.agr_figure.offset) * self.agr_figure.pagescale
|
||||
self.position = '{:.5f}, {:.5f}'.format(*pos)
|
||||
|
||||
def __init__(self, text, agr_axis):
|
||||
self.text = text
|
||||
self.agr_axis = agr_axis
|
||||
self.agr_figure = agr_axis.agr_figure
|
||||
|
||||
self.value = latex_to_xmgrace(text.get_text())
|
||||
self.size = text.get_fontsize() * self.agr_figure.fontscale
|
||||
self.color = AgrLine.color_index(text.get_color())
|
||||
self.get_position()
|
||||
sanitize_strings(self.__dict__)
|
||||
|
||||
def __str__(self):
|
||||
return self.fmt.format(**self.__dict__)
|
||||
|
||||
|
||||
class AgrLine:
|
||||
fmt = """hidden {hidden}
|
||||
type {type}
|
||||
legend "{label}"
|
||||
line linestyle {linestyle}
|
||||
line linewidth {linewidth}
|
||||
line color {color}
|
||||
symbol {marker}
|
||||
symbol color {markeredgecolor}
|
||||
symbol fill color {markerfacecolor}
|
||||
symbol fill pattern {markerfill}
|
||||
symbol linewidth {linewidth}
|
||||
"""
|
||||
width_scale = 2
|
||||
cc = ColorConverter()
|
||||
linestyles = {'None': 0, '-': 1, ':': 2, '--': 3, '-.': 5}
|
||||
markers = defaultdict(
|
||||
lambda: 1,
|
||||
{'': 0, 'None': 0, 'o': 1, 's': 2, 'd': 3, '^': 4, '<': 5, 'v': 6, '>': 7, '+': 8, 'x': 9, '*': 10}
|
||||
)
|
||||
fillstyles = ('none', 'full', 'right', 'left', 'bottom', 'top')
|
||||
colors = ['white', 'black']
|
||||
|
||||
def color_index(col):
|
||||
if col not in AgrLine.colors:
|
||||
AgrLine.colors.append(col)
|
||||
return AgrLine.colors.index(col)
|
||||
|
||||
@property
|
||||
def data(self):
|
||||
o = '@type xy\n'
|
||||
for x, y in self.line.get_xydata():
|
||||
o += '{} {}\n'.format(x, y)
|
||||
o += '&'
|
||||
return o
|
||||
|
||||
def get_label(self):
|
||||
lbl = self.line.get_label()
|
||||
self.label = latex_to_xmgrace(lbl) if not lbl.startswith('_line') else ''
|
||||
|
||||
def get_linestyle(self):
|
||||
self.linestyle = self.linestyles[self.line.get_linestyle()]
|
||||
|
||||
def get_linewidth(self):
|
||||
self.linewidth = self.line.get_linewidth() * self.width_scale
|
||||
|
||||
def get_color(self):
|
||||
self.color = AgrLine.color_index(self.line.get_color())
|
||||
|
||||
def get_marker(self):
|
||||
mk = self.line.get_marker()
|
||||
self.marker = self.markers[mk] if mk in self.markers else 1
|
||||
mfc = self.line.get_markerfacecolor()
|
||||
self.markerfacecolor = AgrLine.color_index(mfc)
|
||||
mec = self.line.get_markeredgecolor()
|
||||
self.markeredgecolor = AgrLine.color_index(mec)
|
||||
self.markeredgewidth = self.line.get_markeredgewidth() * self.width_scale
|
||||
self.markerfill = self.fillstyles.index(self.line.get_fillstyle())
|
||||
|
||||
def __init__(self, line, agr_axis):
|
||||
self.agr_axis = agr_axis
|
||||
self.line = line
|
||||
self.hidden = 'false'
|
||||
self.type = 'xy'
|
||||
|
||||
# run all get_ methods
|
||||
for d in dir(self):
|
||||
if d.startswith('get_'):
|
||||
getattr(self, d)()
|
||||
|
||||
sanitize_strings(self.__dict__)
|
||||
|
||||
def __str__(self):
|
||||
return self.fmt.format(**self.__dict__)
|
||||
|
||||
|
||||
class AgrAxis:
|
||||
fmt = """world {world}
|
||||
view {view}
|
||||
title {title}
|
||||
yaxes scale {yscale}
|
||||
yaxis tick major {yticks}
|
||||
xaxis label "{xlabel}"
|
||||
xaxis label place {xlabelpos}
|
||||
xaxis label char size {labelsize}
|
||||
xaxis tick major {xticks}
|
||||
xaxis ticklabel {xticklabel}
|
||||
xaxis ticklabel place {xticklabelpos}
|
||||
xaxis ticklabel char size {ticklabelsize}
|
||||
xaxes scale {xscale}
|
||||
yaxis label "{ylabel}"
|
||||
yaxis label place {ylabelpos}
|
||||
yaxis label char size {labelsize}
|
||||
yaxis ticklabel {yticklabel}
|
||||
yaxis ticklabel place {yticklabelpos}
|
||||
yaxis ticklabel char size {ticklabelsize}
|
||||
legend {legend}
|
||||
legend loctype world
|
||||
legend {legend_pos}
|
||||
legend char size {legend_fontsize}
|
||||
"""
|
||||
|
||||
def get_world(self):
|
||||
xmin, xmax = self.axis.get_xlim()
|
||||
ymin, ymax = self.axis.get_ylim()
|
||||
self.world = '{:.3}, {:.3}, {:.3}, {:.3}'.format(xmin, ymin, xmax, ymax)
|
||||
box = self.axis.get_position()
|
||||
fx, fy = self.axis.figure.get_size_inches()
|
||||
sx, sy = self.agr_figure.pagescale
|
||||
offx, offy = self.agr_figure.offset
|
||||
self.view = '{:.3}, {:.3}, {:.3}, {:.3}'.format(
|
||||
(box.xmin + offx)*sx, (box.ymin + offy)*sy,
|
||||
(box.xmax + offx)*sx, (box.ymax + offy)*sy
|
||||
)
|
||||
|
||||
def get_title(self):
|
||||
self.title = latex_to_xmgrace(self.axis.get_title())
|
||||
|
||||
def get_xyaxis(self):
|
||||
self.xlabel = latex_to_xmgrace(self.axis.get_xlabel())
|
||||
xpos = self.axis.xaxis.get_label_position()
|
||||
self.xlabelpos = 'normal' if xpos == 'bottom' else 'opposite'
|
||||
|
||||
self.ylabel = latex_to_xmgrace(self.axis.get_ylabel())
|
||||
ypos = self.axis.yaxis.get_label_position()
|
||||
self.ylabelpos = 'normal' if ypos == 'left' else 'opposite'
|
||||
|
||||
xsc = self.axis.get_xscale()
|
||||
self.xscale = 'Logarithmic' if 'log' in xsc else 'Normal'
|
||||
xticks = self.axis.get_xticks()
|
||||
if xsc == 'log':
|
||||
self.xticks = (xticks[1:] / xticks[:-1]).mean()
|
||||
else:
|
||||
self.xticks = (xticks[1:] - xticks[:-1]).mean()
|
||||
|
||||
self.xticklabel = 'on' if any([x.get_visible() for x in self.axis.get_xticklabels()]) else 'off'
|
||||
xpos = self.axis.xaxis.get_ticks_position()
|
||||
if xpos == 'both':
|
||||
self.xticklabelpos = 'both'
|
||||
elif xpos == 'top':
|
||||
self.xticklabelpos = 'opposite'
|
||||
else:
|
||||
self.xticklabelpos = 'normal'
|
||||
|
||||
xtlpos = set([x.get_position()[1] for x in self.axis.get_xticklabels() if x.get_visible()])
|
||||
if len(xtlpos) == 0:
|
||||
self.xticklabel = 'off'
|
||||
self.xticklabelpos = 'normal'
|
||||
elif len(xtlpos) == 1:
|
||||
self.xticklabel = 'on'
|
||||
self.xticklabelpos = 'opposite' if 1 in xtlpos else 'normal'
|
||||
else:
|
||||
self.xticklabel = 'on'
|
||||
self.xticklabelpos = 'both'
|
||||
|
||||
ytlpos = set([x.get_position()[0] for x in self.axis.get_yticklabels() if x.get_visible()])
|
||||
if len(ytlpos) == 0:
|
||||
self.yticklabel = 'off'
|
||||
self.yticklabelpos = 'normal'
|
||||
elif len(ytlpos) == 1:
|
||||
self.yticklabel = 'on'
|
||||
self.yticklabelpos = 'opposite' if 1 in ytlpos else 'normal'
|
||||
else:
|
||||
self.yticklabel = 'on'
|
||||
self.yticklabelpos = 'both'
|
||||
|
||||
ysc = self.axis.get_yscale()
|
||||
self.yscale = 'Logarithmic' if 'log' in ysc else 'Normal'
|
||||
yticks = self.axis.get_yticks()
|
||||
if ysc == 'log':
|
||||
self.yticks = (yticks[1:] / yticks[:-1]).mean()
|
||||
else:
|
||||
self.yticks = (yticks[1:] - yticks[:-1]).mean()
|
||||
|
||||
self.labelsize = self.axis.xaxis.get_label().get_fontsize() * self.agr_figure.fontscale
|
||||
fs = self.axis.xaxis.get_ticklabels()[0].get_fontsize()
|
||||
self.ticklabelsize = fs * self.agr_figure.fontscale
|
||||
|
||||
def get_legend(self):
|
||||
leg = self.axis.get_legend()
|
||||
if leg is None:
|
||||
self.legend = 'off'
|
||||
self.legend_pos = '0, 0'
|
||||
self.legend_fontsize = 0
|
||||
else:
|
||||
self.legend = 'on'
|
||||
for lbl, line in zip(leg.get_texts(), leg.get_lines()):
|
||||
pass
|
||||
pos = get_world_coords(leg)
|
||||
self.legend_pos = '{:.3f}, {:.3f}'.format(*pos.diagonal())
|
||||
self.legend_fontsize = leg.get_texts()[0].get_fontsize() * self.agr_figure.fontscale
|
||||
|
||||
def __init__(self, axis, agr_fig):
|
||||
self.agr_figure = agr_fig
|
||||
self.axis = axis
|
||||
|
||||
# run all get_ methods
|
||||
for d in dir(self):
|
||||
if d.startswith('get_'):
|
||||
getattr(self, d)()
|
||||
|
||||
sanitize_strings(self.__dict__)
|
||||
self.lines = {'s{}'.format(i): AgrLine(l, self) for i, l in enumerate(axis.lines)}
|
||||
self.texts = [AgrText(t, self) for t in self.axis.texts]
|
||||
|
||||
def __str__(self):
|
||||
o = self.fmt.format(**self.__dict__)
|
||||
for k, l in self.lines.items():
|
||||
o += textwrap.indent(str(l), prefix=k + ' ')
|
||||
for txt in self.texts:
|
||||
o += 'with string\n'
|
||||
o += textwrap.indent(str(txt), prefix=' ')
|
||||
return o
|
||||
|
||||
|
||||
class AgrFigure:
|
||||
dpi = 96
|
||||
fontscale = 0.5
|
||||
fmt = """@version 50125
|
||||
@page size {page}
|
||||
"""
|
||||
data_fmt = "@target {target}\n{data}\n"""
|
||||
|
||||
def collect_data(self):
|
||||
d = {}
|
||||
for ia, ax in self.axes.items():
|
||||
for il, ln in ax.lines.items():
|
||||
d['{}.{}'.format(ia.upper(), il.upper())] = ln.data
|
||||
return d
|
||||
|
||||
def get_figprops(self):
|
||||
fx, fy = self.figure.get_size_inches()
|
||||
scx, scy = (1 + self.offset)
|
||||
# scy = (1 + self.offset_vertical)
|
||||
self.page = '{}, {}'.format(int(scx * self.dpi * fx), int(scy * self.dpi * fy))
|
||||
self.fontscale = AgrFigure.fontscale / min(fx * scx, fy * scy)
|
||||
|
||||
self.pagescale = np.array([fx, fy]) / min(fx * scx, fy * scy)
|
||||
|
||||
def __init__(self, figure, offset_horizontal=0, offset_vertical=0, convert_latex=True):
|
||||
tex2grace.do_latex_conversion = convert_latex
|
||||
self.figure = figure
|
||||
self.offset = np.array([offset_horizontal, offset_vertical])
|
||||
# make sure to draw the figure...
|
||||
canv = backend_agg.FigureCanvasAgg(figure)
|
||||
canv.draw()
|
||||
|
||||
# run all get_ methods
|
||||
for d in dir(self):
|
||||
if d.startswith('get_'):
|
||||
getattr(self, d)()
|
||||
sanitize_strings(self.__dict__)
|
||||
|
||||
self.axes = {'g{}'.format(i): AgrAxis(ax, self) for i, ax in enumerate(self.figure.axes)}
|
||||
|
||||
def __str__(self):
|
||||
o = self.fmt.format(**self.__dict__)
|
||||
|
||||
for i, col in enumerate(AgrLine.colors):
|
||||
# in matplotlib-1.5 to_rgb can not handle 'none', this was fixed in 2.0
|
||||
rgb = [int(x * 255) for x in AgrLine.cc.to_rgba(col)[:3]]
|
||||
o += '@map color {i} to ({rgb[0]}, {rgb[1]}, {rgb[2]}), "{col}"\n'.format(i=i, rgb=rgb, col=col)
|
||||
|
||||
for k, ax in self.axes.items():
|
||||
o += textwrap.indent("on\nhidden false\ntype XY\nstacked false\n", prefix='@{} '.format(k))
|
||||
o += '@with {}\n'.format(k)
|
||||
o += textwrap.indent(str(ax), prefix='@ ')
|
||||
|
||||
for k, d in self.collect_data().items():
|
||||
o += self.data_fmt.format(target=k, data=d)
|
||||
|
||||
return o
|
||||
|
||||
|
||||
def saveagr(fname, figure=None, offset_x=0, offset_y=0, convert_latex=True):
|
||||
"""
|
||||
Save figure as xmgrace plot.
|
||||
|
||||
If no figure is provided, this will save the current figure.
|
||||
|
||||
Args:
|
||||
fname: Filename of the agr plot
|
||||
figure (opt.): Matplotlib figure to save, if None gcf() is used.
|
||||
offset_x, offset_y (opt.): Add an offest in x or y direction to the xmgrace plot.
|
||||
convert_latex (opt.): If latex strings will be converted to xmgrace.
|
||||
"""
|
||||
if figure is None:
|
||||
figure = plt.gcf()
|
||||
with open(fname, 'w') as f:
|
||||
af = AgrFigure(figure, offset_horizontal=offset_x, offset_vertical=offset_y,
|
||||
convert_latex=convert_latex)
|
||||
f.write(str(af))
|
||||
Reference in New Issue
Block a user