set_legend method

make @ ba4eae32

-Subproject commit c603498c3d12fb386cfee0a8976517891dca17d0
+Subproject commit ba4eae32e9ba1ed07090b70476a09a8a30781684

tfields/__about__.py

@@ -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"

tfields/plotting/mpl.py

@@ -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):