colors="random" for line, line collection and scatter (#175)

kushalkolar · web-flow · commit 48f6f1466e2b · 2023-04-18T22:11:40.000-04:00
diff --git a/fastplotlib/graphics/features/_colors.py b/fastplotlib/graphics/features/_colors.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from ._base import GraphicFeature, GraphicFeatureIndexable, cleanup_slice, FeatureEvent
-from ...utils import make_colors, get_cmap_texture
+from ...utils import make_colors, get_cmap_texture, make_pygfx_colors
 from pygfx import Color
 
 
@@ -80,10 +80,15 @@ def __init__(self, parent, colors, n_colors: int, alpha: float = 1.0, collection
                     f"Valid iterable color arguments must be a `tuple` or `list` representing RGBA values or "
                     f"an iterable of `str` with the same length as the number of datapoints."
                 )
+        elif isinstance(colors, str):
+            if colors == "random":
+                data = np.random.rand(n_colors, 4)
+                data[:, -1] = alpha
+            else:
+                data = make_pygfx_colors(colors, n_colors)
         else:
             # assume it's a single color, use pygfx.Color to parse it
-            c = Color(colors)
-            data = np.repeat(np.array([c]), n_colors, axis=0)
+            data = make_pygfx_colors(colors, n_colors)
 
         if alpha != 1.0:
             data[:, -1] = alpha
diff --git a/fastplotlib/graphics/line_collection.py b/fastplotlib/graphics/line_collection.py
@@ -27,6 +27,7 @@ def __init__(
             z_position: Union[List[float], float] = None,
             thickness: Union[float, List[float]] = 2.0,
             colors: Union[List[np.ndarray], np.ndarray] = "w",
+            alpha: float = 1.0,
             cmap: Union[List[str], str] = None,
             name: str = None,
             *args,
@@ -131,33 +132,44 @@ def __init__(
                                  "with the same length as the data")
         else:
             if isinstance(colors, np.ndarray):
+                # single color for all lines in the collection as RGBA
                 if colors.shape == (4,):
                     single_color = True
 
+                # colors specified for each line as array of shape [n_lines, RGBA]
                 elif colors.shape == (len(data), 4):
                     single_color = False
 
                 else:
                     raise ValueError(
-                        "numpy array colors argument must be of shape (4,) or (len(data), 4)"
+                        f"numpy array colors argument must be of shape (4,) or (n_lines, 4)."
+                        f"You have pass the following shape: {colors.shape}"
                     )
 
             elif isinstance(colors, str):
-                single_color = True
-                colors = pygfx.Color(colors)
+                if colors == "random":
+                    colors = np.random.rand(len(data), 4)
+                    colors[:, -1] = alpha
+                    single_color = False
+                else:
+                    # parse string color
+                    single_color = True
+                    colors = pygfx.Color(colors)
 
             elif isinstance(colors, (tuple, list)):
                 if len(colors) == 4:
+                    # single color specified as (R, G, B, A) tuple or list
                     if all([isinstance(c, (float, int)) for c in colors]):
                         single_color = True
 
                 elif len(colors) == len(data):
+                    # colors passed as list/tuple of colors, such as list of string
                     single_color = False
 
                 else:
                     raise ValueError(
                         "tuple or list colors argument must be a single color represented as [R, G, B, A], "
-                        "or must be a str of tuple/list with the same length as the data"
+                        "or must be a tuple/list of colors represented by a string with the same length as the data"
                     )
 
         self._set_world_object(pygfx.Group())
diff --git a/fastplotlib/utils/functions.py b/fastplotlib/utils/functions.py
@@ -1,5 +1,5 @@
 import numpy as np
-from pygfx import Texture
+from pygfx import Texture, Color
 from collections import OrderedDict
 from typing import *
 from pathlib import Path
@@ -147,3 +147,10 @@ def quick_min_max(data: np.ndarray) -> Tuple[float, float]:
         data = data[tuple(sl)]
 
     return float(np.nanmin(data)), float(np.nanmax(data))
+
+
+def make_pygfx_colors(colors, n_colors):
+    """parse and make colors array using pyfx.Color"""
+    c = Color(colors)
+    colors_array = np.repeat(np.array([c]), n_colors, axis=0)
+    return colors_array
