replot/replot.py
Phyks (Lucas Verney) d67d185f86 Suppress warning (and legend) if legend is required but no labels are
given.

Also start to implement a basic one-liner plot function, API still in
WIP.
2016-03-01 15:34:11 +01:00

147 lines
4.7 KiB
Python

"""
* Saner default config.
* Matplotlib API methods have an immediate effect on the figure. We do not want
it, then we write a buffer on top of matplotlib API.
"""
import matplotlib.pyplot as plt
import numpy as np
import seaborn
# Cancel seaborn modifications of matplotlib, we do not want to interfere in
# any ways with matplotlib
seaborn.reset_orig()
# TODO: Remove it, this is interfering with matplotlib
plt.rcParams['figure.figsize'] = (10.0, 8.0) # Larger figures by default
plt.rcParams['text.usetex'] = True # Use LaTeX rendering
class Figure():
def __init__(self,
xlabel="", ylabel="", title="", palette="hls",
legend=None):
# TODO: Constants
self.max_colors = 10
self.default_points_number = 1000
self.default_x_interval = np.linspace(-10, 10,
self.default_points_number)
# Set default values for attributes
self.xlabel = xlabel
self.ylabel = ylabel
self.title = title
self.palette = palette
self.legend = legend
self.plots = []
def __enter__(self):
return self
def __exit__(self, exception_type, exception_value, traceback):
self.show()
def show(self):
"""
Actually render and show the figure.
"""
# Tweak matplotlib to use seaborn
seaborn.set()
# Plot using specified color palette
with seaborn.color_palette(self.palette, self.max_colors):
# Create figure
figure, axes = plt.subplots()
# Add plots
for plot in self.plots:
axes.plot(*(plot[0]), **(plot[1]))
# Set properties
axes.set_xlabel(self.xlabel)
axes.set_ylabel(self.ylabel)
axes.set_title(self.title)
if self.legend is not None and self.legend is not False:
self._legend(axes, location=self.legend)
# Draw figure
figure.show()
# Do not forget to restore matplotlib state, in order not to interfere
# with it.
seaborn.reset_orig()
#def palette(self, palette):
# """
# """
# if isinstance(palette, str):
# self.current_palette = palette
# with seaborn.color_palette(self.current_palette, self.max_colors):
# # TODO
# pass
def plot(self, data, *args, **kwargs):
"""
Plot something on the figure.
>>> plot(np.sin)
>>> plot(np.sin, (-1, 1))
>>> plot(np.sin, [-1, -0.9, …, 1])
>>> plot([1, 2, 3], [4, 5, 6])
"""
if hasattr(data, "__call__"):
# We want to plot a function
self._plot_function(data, *args, **kwargs)
else:
# Else, it is a point series, and we just have to store it for
# later plotting.
self.plots.append(((data,) + args, kwargs))
# Automatically set the legend if label is found
# (only do it if legend is not explicitly suppressed)
if "label" in kwargs and self.legend is None:
self.legend = True
def _plot_function(self, data, *args, **kwargs):
"""
"""
# TODO: Better default interval and so on
if len(args) == 0:
# No interval specified, using default one
x_values = self.default_x_interval
elif isinstance(args[0], (list, np.ndarray)):
# List of points specified
x_values = args[0]
elif isinstance(args[0], tuple):
# Interval specified, generate a list of points
x_values = np.linspace(args[0][0], args[0][1],
self.default_points_number)
else:
# TODO: Error
assert False
y_values = [data(i) for i in x_values]
self.plots.append(((x_values, y_values) + args[1:], kwargs))
def _legend(self, axes, location="best"):
"""
"""
if location is True:
# If there should be a legend, but no location provided, put it at
# best location.
location = "best"
# Create aliases for "upper" / "top" and "lower" / "bottom"
location.replace("top ", "upper ")
location.replace("bottom ", "lower ")
# Avoid warning if no labels were given for plots
nb_labelled_plots = sum(["label" in plt[1] for plt in self.plots])
if nb_labelled_plots > 0:
# Add legend
axes.legend(loc=location)
def plot(data, **kwargs):
"""
"""
figure = Figure(**kwargs)
# TODO: Fix API, support every plot type
for plt in data:
figure.plot(plt)
return figure