Enable again to plot states with RGBA tuple color on the Bloch sphere

Author: PositroniumJS

doc/changes/2678.bugfix

@@ -0,0 +1,2 @@
+Enable again to plot a state as a point on the Bloch sphere with a RGBA tuple.
+Enable to do the same as a vector.

qutip/bloch.py

@@ -4,18 +4,16 @@
 from typing import Literal
 
 import numpy as np
-from numpy import (outer, cos, sin, ones)
-
+from numpy import cos, ones, outer, sin
 from packaging.version import parse as parse_version
 
 from . import Qobj, expect, sigmax, sigmay, sigmaz
 
 try:
     import matplotlib
     import matplotlib.pyplot as plt
-    from mpl_toolkits.mplot3d import Axes3D
     from matplotlib.patches import FancyArrowPatch
-    from mpl_toolkits.mplot3d import proj3d
+    from mpl_toolkits.mplot3d import Axes3D, proj3d
 
     # Define a custom _axes3D function based on the matplotlib version.
     # The auto_add_to_figure keyword is new for matplotlib>=3.4.
@@ -385,8 +383,9 @@ def add_states(self, state: Qobj,
         kind : {'vector', 'point'}
             Type of object to plot.
 
-        colors : array_like
+        colors : str or array_like
             Optional array with colors for the states.
+            The colors can be a string or a RGB or RGBA tuple.
 
         alpha : float, default=1.
             Transparency value for the vectors. Values between 0 and 1.
@@ -404,14 +403,18 @@ def add_states(self, state: Qobj,
             colors = np.asarray(colors)
 
             if colors.ndim == 0:
+                colors = np.repeat(colors, state.shape[0])
+            
+            elif colors.ndim == 1 and np.isdtype(colors.dtype, ("integral", "real floating")):
                 colors = colors[np.newaxis]
-
-            if colors.shape != state.shape:
+                colors = np.repeat(colors, [state.shape[0]], axis=0)
+            
+            if colors.shape[0] != state.shape[0]:
                 raise ValueError("The included colors are not valid. "
-                                 "colors must be equivalent to a 1D array "
-                                 "with the same size as the number of states.")
+                                "colors must have the same size as state.")
+
         else:
-            colors = np.array([None] * state.size)
+            colors = np.array([None] * state.shape[0])
 
         for k, st in enumerate(state):
             vec = _state_to_cartesian_coordinates(st)
@@ -420,6 +423,9 @@ def add_states(self, state: Qobj,
                 self.add_vectors(vec, colors=[colors[k]], alpha=alpha)
             elif kind == 'point':
                 self.add_points(vec, colors=[colors[k]], alpha=alpha)
+            else:
+                raise ValueError("The included kind is not valid. "
+                                 f"It should be vector or point, not {kind}.")
 
     def add_vectors(self, vectors, colors=None, alpha=1.0):
         """Add a list of vectors to Bloch sphere.
@@ -429,8 +435,9 @@ def add_vectors(self, vectors, colors=None, alpha=1.0):
         vectors : array_like
             Array with vectors of unit length or smaller.
 
-        colors : array_like
+        colors : str or array_like
             Optional array with colors for the vectors.
+            The colors can be a string or a RGB or RGBA tuple.
 
         alpha : float, default=1.
             Transparency value for the vectors. Values between 0 and 1.
@@ -448,16 +455,20 @@ def add_vectors(self, vectors, colors=None, alpha=1.0):
                 "index represents the iteration over the vectors and the "
                 "second index represents the position in 3D of vector head.")
 
-        n_vectors = vectors.shape[0]
         if colors is None:
-            colors = np.array([None] * n_vectors)
+            colors = np.array([None] * vectors.shape[0])
         else:
             colors = np.asarray(colors)
 
-        if colors.ndim != 1 or colors.size != n_vectors:
-            raise ValueError("The included colors are not valid. colors must "
-                             "be equivalent to a 1D array with the same "
-                             "size as the number of vectors. ")
+            if colors.ndim == 0:
+                colors = np.repeat(colors, vectors.shape[0])
+
+            if (
+                colors.shape[0] != vectors.shape[0]
+                or colors.ndim == 2 and not np.isdtype(colors.dtype, ("integral", "real floating"))
+            ):
+                raise ValueError("The included colors are not valid. "
+                                 "colors must have the same size as vectors.")
 
         for k, vec in enumerate(vectors):
             self.vectors.append(vec)

qutip/tests/test_bloch.py

@@ -488,14 +488,20 @@ def plot_vector_ref(self, fig, vector_kws):
             dict(vectors=[(1, 0, 1), (1, 1, 0)], alpha=0.5),
         ], id="alpha-multiple-vector-sets"),
         pytest.param(
-            dict(vectors=(0, 0, 1), colors=['y']), id="color-y"),
+            dict(vectors=(0, 0, 1), colors=['#4f52cc']), id="color-hex"),
         pytest.param(
             dict(vectors=[(0, 0, 1), (0, 1, 0)], colors=['y', 'y']),
             id="color-two-y"),
         pytest.param([
             dict(vectors=[(0, 0, 1)], colors=['y']),
             dict(vectors=[(1, 0, 1)], colors=['g']),
         ], id="color-yg"),
+        pytest.param(
+            dict(vectors=[(0, 0, 1), (0, 1, 0)], colors=[(0.4, 0.7, 0.5), (0.1, 0.2, 0.8)]),
+            id="color-RGB"),
+        pytest.param(
+            dict(vectors=[(0, 0, 1), (0, 1, 0)], colors=[(0.4, 0.7, 0.5, 0.9), (0.1, 0.2, 0.8, 0.4)]),
+            id="color-RGBA"),
     ])
     @check_pngs_equal
     def test_vector(self, vector_kws, fig_test, fig_ref):
@@ -536,9 +542,8 @@ def test_vector_errors_color_length(self, vectors, colors):
             b.add_vectors(vectors, colors=colors)
             b.render()
 
-        err_msg = ("The included colors are not valid. colors must "
-                   "be equivalent to a 1D array with the same "
-                   "size as the number of vectors. ")
+        err_msg = ("The included colors are not valid. "
+                   "colors must have the same size as vectors.")
         assert str(err.value) == err_msg
 
     @check_pngs_equal