Repository: qutip/qutip

Issue:
Add option to include colours to Bloch vectors instead of just cycling from a list
The Bloch sphere is an amazing tool for visualisations. In a project where I wanted to show the evolution of an initial quantum state to a final target, I wanted to put specific colours to the initial and final bloch vector.

Right now, the function `[qutip.bloch.add_states](https://github.com/qutip/qutip/blob/master/qutip/bloch.py#L353)` or `[qutip.bloch.add_vectors](https://github.com/qutip/qutip/blob/master/qutip/bloch.py#L378)` cycles through a fixed list of colours here - https://github.com/qutip/qutip/blob/master/qutip/bloch.py#L563

It would be nice if the user could also add an optional color argument or list of colours when adding a new vector such as `bloch.add_vectors(vec, color="r")` or  `bloch.add_vectors([vec1, vec2], color=["red", "green"])`

In this case I would have liked to set the two colours in the vector myself. 

![image](https://user-images.githubusercontent.com/6968324/84878607-58340100-b08a-11ea-978a-03085378261f.png)



Discussion:
Hey, I tired solving this issue. I changed the add_vectors method such that it takes in an optional argument called 'color'. I created a dictionary called vector_color_specified which stores the vector-index and color as key-value pair.  Then when plot_vectors is called, the color is taken from the dictionary if mentioned, otherwise the default value is used based on the cyclic color list which is already present. I tested my code and it works as you have mentioned above. So, is my approach right? Any advice would be helpful. This is the first open-source issue that I have taken up.
Hi @rajathshetty20, Thanks for the interest. What you described sounds good. Would you like to make a Pull Request? You can go to the Pull Requests page and create one from a branch of your repository. There will be a template guiding you once you choose a branch.

Resolving patch:
diff --git a/qutip/bloch.py b/qutip/bloch.py
--- a/qutip/bloch.py
+++ b/qutip/bloch.py
@@ -88,7 +88,7 @@ class Bloch():
         Color of sphere wireframe.
     frame_width : int {1}
         Width of wireframe.
-    point_color : list {["b","r","g","#CC6600"]}
+    point_default_color : list {["b","r","g","#CC6600"]}
         List of colors for Bloch sphere point markers to cycle through.
         i.e. By default, points 0 and 4 will both be blue ('b').
     point_marker : list {["o","s","d","^"]}
@@ -103,7 +103,7 @@ class Bloch():
     figsize : list {[7,7]}
         Figure size of Bloch sphere plot.  Best to have both numbers the same;
         otherwise you will have a Bloch sphere that looks like a football.
-    vector_color : list {["g","#CC6600","b","r"]}
+    vector_default_color : list {["g","#CC6600","b","r"]}
         List of vector colors to cycle through.
     vector_width : int {5}
         Width of displayed vectors.
@@ -170,7 +170,9 @@ def __init__(self, fig=None, axes=None, view=None, figsize=None,
 
         # ---vector options---
         # List of colors for Bloch vectors, default = ['b','g','r','y']
-        self.vector_color = ['g', '#CC6600', 'b', 'r']
+        self.vector_default_color = ['g', '#CC6600', 'b', 'r']
+        # List that stores the display colors for each vector
+        self.vector_color = []
         #: Width of Bloch vectors, default = 5
         self.vector_width = 3
         #: Style of Bloch vectors, default = '-|>' (or 'simple')
@@ -180,7 +182,9 @@ def __init__(self, fig=None, axes=None, view=None, figsize=None,
 
         # ---point options---
         # List of colors for Bloch point markers, default = ['b','g','r','y']
-        self.point_color = ['b', 'r', 'g', '#CC6600']
+        self.point_default_color = ['b', 'r', 'g', '#CC6600']
+        # List that stores the display colors for each set of points
+        self.point_color = []
         # Size of point markers, default = 25
         self.point_size = [25, 32, 35, 45]
         # Shape of point markers, default = ['o','^','d','s']
@@ -276,13 +280,13 @@ def __str__(self):
         s += "frame_alpha:     " + str(self.frame_alpha) + "\n"
         s += "frame_color:     " + str(self.frame_color) + "\n"
         s += "frame_width:     " + str(self.frame_width) + "\n"
-        s += "point_color:     " + str(self.point_color) + "\n"
+        s += "point_default_color:" + str(self.point_default_color) + "\n"
         s += "point_marker:    " + str(self.point_marker) + "\n"
         s += "point_size:      " + str(self.point_size) + "\n"
         s += "sphere_alpha:    " + str(self.sphere_alpha) + "\n"
         s += "sphere_color:    " + str(self.sphere_color) + "\n"
         s += "figsize:         " + str(self.figsize) + "\n"
-        s += "vector_color:    " + str(self.vector_color) + "\n"
+        s += "vector_default_color:" + str(self.vector_default_color) + "\n"
         s += "vector_width:    " + str(self.vector_width) + "\n"
         s += "vector_style:    " + str(self.vector_style) + "\n"
         s += "vector_mutation: " + str(self.vector_mutation) + "\n"
@@ -316,8 +320,10 @@ def clear(self):
         self.vectors = []
         self.point_style = []
         self.annotations = []
+        self.vector_color = []
+        self.point_color = []
 
-    def add_points(self, points, meth='s'):
+    def add_points(self, points, meth='s', colors=None):
         """Add a list of data points to bloch sphere.
 
         Parameters
@@ -328,6 +334,9 @@ def add_points(self, points, meth='s'):
         meth : str {'s', 'm', 'l'}
             Type of points to plot, use 'm' for multicolored, 'l' for points
             connected with a line.
+        colors : array_like
+            Optional array with colors for the points.
+            A single color for meth 's', and list of colors for meth 'm'
 
         """
         if not isinstance(points[0], (list, ndarray)):
@@ -341,14 +350,31 @@ def add_points(self, points, meth='s'):
                 pnts = points
             self.points.append(pnts)
             self.point_style.append('s')
+            if colors is not None:
+                colors = [colors]
+            else:
+                k = mod(len(self.points), len(self.point_default_color))
+                colors = [self.point_default_color[k]]
         elif meth == 'l':
             self.points.append(points)
             self.point_style.append('l')
+            if colors is not None:
+                colors = [colors]
+            else:
+                k = mod(len(self.points), len(self.point_default_color))
+                colors = [self.point_default_color[k]]
         else:
             self.points.append(points)
             self.point_style.append('m')
-
-    def add_states(self, state, kind='vector'):
+            if colors is None:
+                colors = []
+                for k, point in enumerate(points[0]):
+                    x = mod(k, len(self.point_default_color))
+                    color = self.point_default_color[x]
+                    colors.append(color)
+        self.point_color.append(colors)
+
+    def add_states(self, state, kind='vector', colors=None):
         """Add a state vector Qobj to Bloch sphere.
 
         Parameters
@@ -358,35 +384,61 @@ def add_states(self, state, kind='vector'):
 
         kind : str {'vector','point'}
             Type of object to plot.
+        colors : array_like
+            Optional array with colors for the states.
 
         """
         if isinstance(state, Qobj):
             state = [state]
+        if not isinstance(colors, (list, ndarray)) and colors is not None:
+            colors = [colors]
 
-        for st in state:
+        for k, st in enumerate(state):
             vec = [expect(sigmax(), st),
                    expect(sigmay(), st),
                    expect(sigmaz(), st)]
 
             if kind == 'vector':
-                self.add_vectors(vec)
+                if colors is not None:
+                    self.add_vectors(vec, colors=colors[k])
+                else:
+                    self.add_vectors(vec)
             elif kind == 'point':
-                self.add_points(vec)
+                if colors is not None:
+                    self.add_points(vec, colors=colors[k])
+                else:
+                    self.add_points(vec)
 
-    def add_vectors(self, vectors):
+    def add_vectors(self, vectors, colors=None):
         """Add a list of vectors to Bloch sphere.
 
         Parameters
         ----------
         vectors : array_like
             Array with vectors of unit length or smaller.
+        colors : array_like
+            Optional array with colors for the vectors.
 
         """
         if isinstance(vectors[0], (list, ndarray)):
-            for vec in vectors:
-                self.vectors.append(vec)
+            if colors is not None:
+                for k, vec in enumerate(vectors):
+                    self.vectors.append(vec)
+                    self.vector_color.append(colors[k])
+            else:
+                for vec in vectors:
+                    self.vectors.append(vec)
+                    k = mod(len(self.vectors), len(self.vector_default_color))
+                    color = self.vector_default_color[k]
+                    self.vector_color.append(color)
         else:
             self.vectors.append(vectors)
+            if colors is not None:
+                self.vector_color.append(colors)
+            else:
+                k = mod(len(self.vectors), len(self.vector_default_color))
+                color = self.vector_default_color[k]
+                self.vector_color.append(color)
 
     def add_annotation(self, state_or_vector, text, **kwargs):
         """Add a text or LaTeX annotation to Bloch sphere,
@@ -558,8 +610,6 @@ def plot_vectors(self):
             ys3d = -self.vectors[k][0] * array([0, 1])
             zs3d = self.vectors[k][2] * array([0, 1])
 
-            color = self.vector_color[mod(k, len(self.vector_color))]
-
             if self.vector_style == '':
                 # simple line style
                 self.axes.plot(xs3d, ys3d, zs3d,
@@ -571,7 +621,7 @@ def plot_vectors(self):
                             mutation_scale=self.vector_mutation,
                             lw=self.vector_width,
                             arrowstyle=self.vector_style,
-                            color=color)
+                            color=self.vector_color[k])
 
                 self.axes.add_artist(a)
 
@@ -591,6 +641,7 @@ def plot_points(self):
             else:
                 indperm = arange(num)
             if self.point_style[k] == 's':
+
                 self.axes.scatter(
                     real(self.points[k][1][indperm]),
                     - real(self.points[k][0][indperm]),
@@ -599,31 +650,30 @@ def plot_points(self):
                     alpha=1,
                     edgecolor=None,
                     zdir='z',
-                    color=self.point_color[mod(k, len(self.point_color))],
+                    color=self.point_color[k][0],
                     marker=self.point_marker[mod(k, len(self.point_marker))])
 
             elif self.point_style[k] == 'm':
-                pnt_colors = array(self.point_color *
-                                   int(ceil(num / float(len(self.point_color)))))
+                pnt_colors = self.point_default_color * int(
+                    ceil(num / float(len(self.point_default_color)))
+                )
+                pnt_colors = array(pnt_colors)
 
-                pnt_colors = pnt_colors[0:num]
-                pnt_colors = list(pnt_colors[indperm])
                 marker = self.point_marker[mod(k, len(self.point_marker))]
                 s = self.point_size[mod(k, len(self.point_size))]
                 self.axes.scatter(real(self.points[k][1][indperm]),
                                   -real(self.points[k][0][indperm]),
                                   real(self.points[k][2][indperm]),
                                   s=s, alpha=1, edgecolor=None,
-                                  zdir='z', color=pnt_colors,
+                                  zdir='z', color=self.point_color[k],
                                   marker=marker)
 
             elif self.point_style[k] == 'l':
-                color = self.point_color[mod(k, len(self.point_color))]
                 self.axes.plot(real(self.points[k][1]),
                                -real(self.points[k][0]),
                                real(self.points[k][2]),
                                alpha=0.75, zdir='z',
-                               color=color)
+                               color=self.point_color[k][0])
 
     def plot_annotations(self):
         # -X and Y data are switched for plotting purposes


Test patch: