From d1dc14e1e68dc8126525e2a658b93bd4d61aecff Mon Sep 17 00:00:00 2001
From: kushalkolar <kushalkolar@gmail.com>
Date: Wed, 9 Apr 2025 02:15:55 -0400
Subject: [PATCH] custom cmap for positions graphics

---
 examples/scatter/scatter_cmap_transform.py    | 68 +++++++++++++++++++
 fastplotlib/graphics/_positions_base.py       | 13 +++-
 .../graphics/features/_positions_graphics.py  | 15 ++--
 fastplotlib/utils/functions.py                |  2 +-
 4 files changed, 90 insertions(+), 8 deletions(-)
 create mode 100644 examples/scatter/scatter_cmap_transform.py

diff --git a/examples/scatter/scatter_cmap_transform.py b/examples/scatter/scatter_cmap_transform.py
new file mode 100644
index 000000000..f3c5e79b6
--- /dev/null
+++ b/examples/scatter/scatter_cmap_transform.py
@@ -0,0 +1,68 @@
+"""
+Scatter custom cmap
+===================
+
+Use a cmap_transform to define how to map colors to scatter points from a custom defined cmap.
+This is also valid for line graphics.
+
+This is identical to the scatter.py example but the cmap_transform is sometimes a better way to define the colors.
+It may also be more performant if millions of strings for each point do not have to be parsed into colors.
+"""
+
+# test_example = true
+# sphinx_gallery_pygfx_docs = 'screenshot'
+
+import fastplotlib as fpl
+import numpy as np
+
+figure = fpl.Figure(size=(700, 560))
+
+# create a random distribution of 15,000 xyz coordinates
+n_points = 5_000
+
+# dimensions always have to be [n_points, xyz]
+dims = (n_points, 3)
+
+clouds_offset = 15
+
+# create some random clouds
+normal = np.random.normal(size=dims, scale=5)
+# stack the data into a single array
+cloud = np.vstack(
+    [
+        normal - clouds_offset,
+        normal,
+        normal + clouds_offset,
+    ]
+)
+
+# we have 3 clouds, create a 1D array where the value indicates cloud membership
+# this will be used to map the colors from the cmap onto the corresponding point
+cmap_transform = np.empty(cloud.shape[0])
+
+# first cloud is given a value of 0
+cmap_transform[:n_points] = 0
+
+# second cloud is given a value of 1
+cmap_transform[n_points: 2 * n_points] = 1
+
+# 3rd cloud given a value of 2
+cmap_transform[2 * n_points:] = 2
+
+
+figure[0, 0].add_scatter(
+    data=cloud,
+    sizes=3,
+    cmap=["green", "purple", "blue"],  # custom cmap
+    cmap_transform=cmap_transform,  # each element of the cmap_transform maps to the corresponding datapoint
+    alpha=0.6
+)
+
+figure.show()
+
+
+# NOTE: `if __name__ == "__main__"` is NOT how to use fastplotlib interactively
+# please see our docs for using fastplotlib interactively in ipython and jupyter
+if __name__ == "__main__":
+    print(__doc__)
+    fpl.loop.run()
diff --git a/fastplotlib/graphics/_positions_base.py b/fastplotlib/graphics/_positions_base.py
index 5d98d16d1..654a4387e 100644
--- a/fastplotlib/graphics/_positions_base.py
+++ b/fastplotlib/graphics/_positions_base.py
@@ -74,7 +74,7 @@ def __init__(
         colors: str | np.ndarray | tuple[float] | list[float] | list[str] = "w",
         uniform_color: bool = False,
         alpha: float = 1.0,
-        cmap: str | VertexCmap = None,
+        cmap: str | list[str] | tuple[str] | VertexCmap = None,
         cmap_transform: np.ndarray = None,
         isolated_buffer: bool = True,
         size_space: str = "screen",
@@ -94,7 +94,13 @@ def __init__(
             if uniform_color:
                 raise TypeError("Cannot use cmap if uniform_color=True")
 
-            if isinstance(cmap, str):
+            if isinstance(cmap, (str, list, tuple)):
+                if isinstance(cmap, (list, tuple)):
+                    if not all(isinstance(s, str) for s in cmap):
+                        raise TypeError(
+                            "`cmap` argument must be a <str> cmap name, a list/tuple of <str> "
+                            "defining a custom cmap, or an existing `VertexCmap` instance"
+                        )
                 # make colors from cmap
                 if isinstance(colors, VertexColors):
                     # share buffer with existing colors instance for the cmap
@@ -116,7 +122,8 @@ def __init__(
                 self._colors = cmap._vertex_colors
             else:
                 raise TypeError(
-                    "`cmap` argument must be a <str> cmap name or an existing `VertexCmap` instance"
+                    "`cmap` argument must be a <str> cmap name, a list/tuple of <str> "
+                    "defining a custom cmap, or an existing `VertexCmap` instance"
                 )
         else:
             # no cmap given
diff --git a/fastplotlib/graphics/features/_positions_graphics.py b/fastplotlib/graphics/features/_positions_graphics.py
index 868701079..b9a5f2fbf 100644
--- a/fastplotlib/graphics/features/_positions_graphics.py
+++ b/fastplotlib/graphics/features/_positions_graphics.py
@@ -417,7 +417,7 @@ class VertexCmap(BufferManager):
         },
         {
             "dict key": "value",
-            "type": "str",
+            "type": "str | list[str] | tuple[str]",
             "description": "new cmap to set at given slice",
         },
     ]
@@ -425,7 +425,7 @@ class VertexCmap(BufferManager):
     def __init__(
         self,
         vertex_colors: VertexColors,
-        cmap_name: str | None,
+        cmap_name: str | list[str] | tuple[str] | None,
         transform: np.ndarray | None,
         alpha: float = 1.0,
     ):
@@ -442,10 +442,17 @@ def __init__(
         self._alpha = alpha
 
         if self._cmap_name is not None:
-            if not isinstance(self._cmap_name, str):
+            if not isinstance(self._cmap_name, (str, list, tuple)):
                 raise TypeError(
-                    f"cmap name must be of type <str>, you have passed: {self._cmap_name} of type: {type(self._cmap_name)}"
+                    f"cmap name must be of type <str>, or list/tuple of str to define a custom cmap. "
+                    f"You have passed: {self._cmap_name} of type: {type(self._cmap_name)}"
                 )
+            if isinstance(cmap_name, (list, tuple)):
+                if not all(isinstance(s, str) for s in cmap_name):
+                    raise TypeError(
+                        f"cmap name must be of type <str>, or list/tuple of str to define a custom cmap. "
+                        f"You have passed: {self._cmap_name} of type: {type(self._cmap_name)}"
+                    )
 
             if self._transform is not None:
                 self._transform = np.asarray(self._transform)
diff --git a/fastplotlib/utils/functions.py b/fastplotlib/utils/functions.py
index 6ad365e40..f88ce607c 100644
--- a/fastplotlib/utils/functions.py
+++ b/fastplotlib/utils/functions.py
@@ -341,7 +341,7 @@ def normalize_min_max(a):
 
 def parse_cmap_values(
     n_colors: int,
-    cmap_name: str,
+    cmap_name: str | list[str] | tuple[str],
     transform: np.ndarray | list[int | float] = None,
 ) -> np.ndarray:
     """

<!DOCTYPE html PUBLIC '-//W3C//DTD XHTML 1.0 Transitional//EN' 'http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd'>
<html xmlns='http://www.w3.org/1999/xhtml'>
<head>
<title>pFad - Phonifier reborn</title>
<meta http-equiv='Content-Type' content='text/html; charset=utf-8' />
</head>
<body>
<h1>Pfad - The Proxy pFad of &#169; 2024 Garber Painting. All rights reserved.</h1>


<!-- Disclaimer -->
<p>Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.</p>
<br>
<p>Alternative Proxies:</p><p><a href="http://clevelandohioweatherforecast.com/php-proxy/index.php?q=https://patch-diff.githubusercontent.com/raw/fastplotlib/fastplotlib/pull/790.patch" target="_blank">Alternative Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/index.php?u=https://patch-diff.githubusercontent.com/raw/fastplotlib/fastplotlib/pull/790.patch" target="_blank">pFad Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/v3index.php?u=https://patch-diff.githubusercontent.com/raw/fastplotlib/fastplotlib/pull/790.patch" target="_blank">pFad v3 Proxy</a></p><p><a href="http://clevelandohioweatherforecast.com/pFad/v4index.php?u=https://patch-diff.githubusercontent.com/raw/fastplotlib/fastplotlib/pull/790.patch" target="_blank">pFad v4 Proxy</a></p></body>
</html>