Commit d5d97e5f authored by Daniel Boeckenhoff's avatar Daniel Boeckenhoff
Browse files

set_legend method

parent 6a696389
Subproject commit c603498c3d12fb386cfee0a8976517891dca17d0
Subproject commit ba4eae32e9ba1ed07090b70476a09a8a30781684
......@@ -17,7 +17,7 @@ __all__ = [
"__classifiers__",
]
__version__ = '0.1.2'
__version__ = '0.1.3'
__title__ = os.path.basename(os.path.abspath('.'))
__summary__ = "numpy + sympy implementation of tensor fields with attached coordinate systems"
__keywords__ = "tensors coordinate system trafo sympy numpy"
......
......@@ -14,6 +14,7 @@ from matplotlib.patches import Circle
import mpl_toolkits.mplot3d as plt3D
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.dates as dates
from matplotlib.patches import Rectangle
from itertools import cycle
from functools import partial
import logging
......@@ -645,14 +646,90 @@ def autoscale_3d(axis, array=None, xLim=None, yLim=None, zLim=None):
def set_legend(axis, artists, **kwargs):
"""
Convenience method to set a legend from multiple artists to an axis.
Args:
**kwargs
table (bool): if True, labels containing ',' will be mapped to table
table_title (str): value of the table entry top left - only active
if table
Examples:
>> import tfields
>> import matplotlib.pyplot as plt
>> fig = plt.figure()
>> ax = fig.add_subplot(111)
>> im1 = ax.plot(range(10), pylab.randn(10), "r--", label=(r"$i = 1$,$j = 1$"))
>> im2 = ax.plot(range(10), pylab.randn(10), "g--", label=(r"$i = 1$,$j = 2$"))
>> im3 = ax.plot(range(10), pylab.randn(10), "b--", label=(r"$i = 1$,$j = 3$"))
>> im4 = ax.plot(range(10), pylab.randn(10), "r.", label=(r"$i = 2$,$j = 1$"))
>> im5 = ax.plot(range(10), pylab.randn(10), "g.", label=(r"$i = 2$,$j = 2$"))
>> im6 = ax.plot(range(10), pylab.randn(10), "b.", label=(r"$i = 2$,$j = 3$"))
>> im7 = ax.plot(range(10), pylab.randn(10), "r^", label=(r"$i = 3$,$j = 1$"))
>> im8 = ax.plot(range(10), pylab.randn(10), "g^", label=(r"$i = 3$,$j = 2$"))
>> im9 = ax.plot(range(10), pylab.randn(10), "b^", label=(r"$i = 3$,$j = 3$"))
>> handles = [im1, im2, im3, im4, im5, im6, im7, im8, im9]
>> tfields.plotting.set_legend(ax, handles, table=True)
>> plt.show()
"""
table = kwargs.pop('table', False)
labels = kwargs.pop('labels', None)
ncol = kwargs.pop('ncol', None)
handles = []
for artist in artists:
if isinstance(artist, list):
handles.append(artist[0])
else:
handles.append(artist)
return axis.legend(handles=handles, **kwargs)
if table and labels is None and ncol is None:
table_title = kwargs.pop('table_title', '')
labels = np.array([h.get_label() for h in handles])
labels = [l.split(',') for l in labels]
captions_i = []
captions_j = []
for l in labels:
if l[0] not in captions_i:
captions_i.append(l[0])
if l[1] not in captions_j:
captions_j.append(l[1])
shape = (len(captions_i),
len(captions_j))
# initialize
shape = np.array(shape)
handles = np.array(handles)
# create blank rectangle
extra = Rectangle((0, 0), 1, 1, fc="w", fill=False, edgecolor='none', linewidth=0)
# Create organized list containing all handles for table. Extra represent empty space
handles_table = np.full(shape + 1, extra)
for handle, label in zip(handles, labels):
i = captions_i.index(label[0])
j = captions_j.index(label[1])
if handles_table[i + 1, j + 1] != extra:
raise ValueError("Duplicate label {label}"
.format(**locals()))
handles_table[i + 1, j + 1] = handle
# Define the label captions
labels_table = np.full(shape + 1, '', dtype='S80')
labels_table[0, 0] = table_title
labels_table[0, 1:] = captions_j
labels_table[1:, 0] = captions_i
labels_table = labels_table.astype(str)
handles = list(handles_table.flat)
labels = list(labels_table.flat)
kwargs['ncol'] = shape[0] + 1
kwargs['handletextpad'] = kwargs.pop('handletextpad', -1.5) # negative numbers move to the right
kwargs['columnspacing'] = kwargs.pop('columnspacing', 1.5)
return axis.legend(handles=handles, labels=labels, **kwargs)
class set_zoomable(object):
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment