import os
import shutil as sh
import numpy as np
def _read_line(filename, line_number):
s = None
with open(filename, 'r') as fs:
for i, line in enumerate(fs.readlines()):
if i == line_number:
s = line
return s
class _GnuplotDeletingFile:
def __init__(self, filename):
self.name = filename
def __del__(self):
os.remove(self.name)
class _GnuplotScriptTemp(_GnuplotDeletingFile):
def __init__(self, gnuplot_cmds):
_GnuplotDeletingFile.__init__(self, '.tmp_gnuplot.gpi')
with open(self.name, 'w') as fs:
fs.write(gnuplot_cmds)
class _GnuplotDataTemp(_GnuplotDeletingFile):
def __init__(self, *args):
_GnuplotDeletingFile.__init__(self, '.tmp_gnuplot_data.dat')
data = np.array(args).T
with open(self.name, 'wb') as fs:
np.savetxt(fs, data, delimiter=',')
class _GnuplotDataZMatrixTemp(_GnuplotDeletingFile):
def __init__(self, z_matrix):
_GnuplotDeletingFile.__init__(self, '.tmp_gnuplot_data_z_matrix.dat')
with open(self.name, 'wb') as fs:
np.savetxt(fs, z_matrix, '%.3f', delimiter=',')
[docs]def gnuplot(script_name, args_dict={}, data=[], silent=True):
'''
Call a Gnuplot script, passing it arguments and
datasets.
Args:
scipt_name(str): The name of the Gnuplot script.
args_dict(dict): A dictionary of parameters to pass
to the script. The `key` is the name of the variable
that the `item` will be passed to the Gnuplot script
with.
data(list): A list of lists containing lists to be plotted.
The lists can be accessed by plotting the variable
`data` in the Gnuplot script. The first list in the
list of lists corresponds to the first column in data,
and so on.
silent (bool): `True` if Gnuplot stdout should be silenced,
`False` if not.
Returns:
str: The Gnuplot command used to call the script.
'''
gnuplot_command = 'gnuplot'
if data:
assert 'data' not in args_dict, \
'Can\'t use \'data\' variable twice.'
data_temp = _GnuplotDataTemp(*data)
args_dict['data'] = data_temp.name
if args_dict:
gnuplot_command += ' -e "'
for arg in args_dict.items():
gnuplot_command += arg[0] + '='
if isinstance(arg[1], str):
gnuplot_command += '\'' + arg[1] + '\''
elif isinstance(arg[1], bool):
if arg[1] is True:
gnuplot_command += '1'
else:
gnuplot_command += '0'
elif hasattr(arg[1], '__iter__'):
gnuplot_command += '\'' + ' '.join([str(v) for v in arg[1]]) + '\''
else:
gnuplot_command += str(arg[1])
gnuplot_command += '; '
gnuplot_command = gnuplot_command[:-1]
gnuplot_command += '"'
gnuplot_command += ' ' + script_name
if silent:
gnuplot_command += ' > /dev/null 2>&1'
os.system(gnuplot_command)
return gnuplot_command
[docs]def gnuplot_2d(x, y, filename, title='', x_label='', y_label=''):
'''
Function to produce a general 2D plot.
Args:
x (list): x points.
y (list): y points.
filename (str): Filename of the output image.
title (str): Title of the plot. Default is '' (no title).
x_label (str): x-axis label.
y_label (str): y-axis label.
'''
_, ext = os.path.splitext(filename)
if ext != '.png':
filename += '.png'
gnuplot_cmds = \
'''
set datafile separator ","
set term pngcairo size 30cm,25cm
set out filename
unset key
set border lw 1.5
set grid lt -1 lc rgb "gray80"
set title title
set xlabel x_label
set ylabel y_label
plot filename_data u 1:2 w lp pt 6 ps 0.5
'''
scr = _GnuplotScriptTemp(gnuplot_cmds)
data = _GnuplotDataTemp(x, y)
args_dict = {
'filename': filename,
'filename_data': data.name,
'title': title,
'x_label': x_label,
'y_label': y_label
}
gnuplot(scr.name, args_dict)
[docs]def gnuplot_3d(x, y, z, filename, title='', x_label='', y_label='', z_label=''):
'''
Function to produce a general 3D plot.
Args:
x (list): x points.
y (list): y points.
z (list): z points.
filename (str): Filename of the output image.
title (str): Title of the plot. Default is '' (no title).
x_label (str): x-axis label.
y_label (str): y-axis label.
z_label (str): z-axis label.
'''
_, ext = os.path.splitext(filename)
if ext != '.png':
filename += '.png'
gnuplot_cmds = \
'''
set datafile separator ","
set term pngcairo size 30cm,25cm
set out filename
unset key
set border lw 1.5
set view map
set title title
set xlabel x_label
set ylabel y_label
set zlabel z_label
splot filename_data u 1:2:3 w pm3d
'''
scr = _GnuplotScriptTemp(gnuplot_cmds)
data = _GnuplotDataTemp(x, y, z)
args_dict = {
'filename': filename,
'filename_data': data.name,
'title': title,
'x_label': x_label,
'y_label': y_label,
'z_label': z_label
}
gnuplot(scr.name, args_dict)
[docs]def gnuplot_3d_matrix(z_matrix, filename, title='', x_label='', y_label=''):
'''
Function to produce a general 3D plot from a 2D matrix.
Args:
z_matrix (list): 2D matrix.
filename (str): Filename of the output image.
title (str): Title of the plot. Default is '' (no title).
x_label (str): x-axis label.
y_label (str): y-axis label.
'''
_, ext = os.path.splitext(filename)
if ext != '.png':
filename += '.png'
gnuplot_cmds = \
'''
set datafile separator ","
set term pngcairo size 30cm,25cm
set out filename
unset key
set border lw 1.5
set view map
set title title
set xlabel x_label
set ylabel y_label
splot filename_data matrix w pm3d
'''
scr = _GnuplotScriptTemp(gnuplot_cmds)
data = _GnuplotDataZMatrixTemp(z_matrix)
args_dict = {
'filename': filename,
'filename_data': data.name,
'title': title,
'x_label': x_label,
'y_label': y_label
}
gnuplot(scr.name, args_dict)
[docs]def trim_pad_image(filename, padding=20):
'''
Trims and pads an image.
Args:
filename(str): The filename of the image to be
acted on.
padding(int): The number of pixels in padding to
add to the image after the image has been
trimmed.
'''
os.system('convert %s -trim -bordercolor white -border %i %s' % \
(filename, padding, filename))