add method to choose link colors

petercorke · petercorke · commit f963bc5aab6b · 2020-11-10T06:57:51.000+10:00
useful for STL models
diff --git a/roboticstoolbox/robot/Robot.py b/roboticstoolbox/robot/Robot.py
@@ -11,6 +11,8 @@
 # from roboticstoolbox.tools import xacro
 from pathlib import PurePath, PurePosixPath
 from scipy.optimize import minimize, Bounds, LinearConstraint
+from matplotlib import colors
+from matplotlib import cm
 from roboticstoolbox.tools.null import null
 from ansitable import ANSITable, Column
 
@@ -267,6 +269,52 @@ def dyntable(self):
             table.row(link.name, *link._dyn2list())
         return str(table)
 
+
+
+    def linkcolormap(self, linkcolors="viridis"):
+        """
+        Create a colormap for robot joints
+
+        :param linkcolors: list of colors or colormap, defaults to "viridis"
+        :type linkcolors: list or str, optional
+        :return: color map
+        :rtype: matplotlib.colors.ListedColormap
+
+        - ``cm = robot.linkcolormap()`` is an n-element colormap that gives a
+          unique color for every link.  The RGBA colors for link ``j`` are
+          ``cm(j)``.
+        - ``cm = robot.linkcolormap(cmap)`` as above but ``cmap`` is the name
+          of a valid matplotlib colormap.  The default, example above, is the
+          ``viridis`` colormap.
+        - ``cm = robot.linkcolormap(list of colors)`` as above but a 
+          colormap is created from a list of n color names given as strings,
+          tuples or hexstrings.
+        
+        .. runblock:: pycon
+
+            >>> import roboticstoolbox as rtb
+            >>> robot = rtb.models.DH.Puma560()
+            >>> cm = robot.linkcolormap("inferno")
+            >>> print(cm(range(6))) # cm(i) is 3rd color in colormap
+            >>> cm = robot.linkcolormap(['red', 'g', (0,0.5,0), '#0f8040', 'yellow', 'cyan'])
+            >>> print(cm(range(6)))
+
+        .. note::
+            
+            - Colormaps have 4-elements: red, green, blue, alpha (RGBA)
+            - Names of supported colors and colormaps are defined in the matplotlib
+              documentation.
+
+                - `Specifying colors <https://matplotlib.org/3.1.0/tutorials/colors/colors.html#sphx-glr-tutorials-colors-colors-py>`_
+                - `Colormaps <https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html#sphx-glr-tutorials-colors-colormaps-py>`_
+        """
+
+        if isinstance(linkcolors, list) and len(linkcolors) == self.n:
+            # provided a list of color names
+            return colors.ListedColormap(linkcolors)
+        else:
+            # assume it is a colormap name
+            return cm.get_cmap(linkcolors, 6)
 # --------------------------------------------------------------------- #
     @property
     def name(self):
