Repository: pygfx/pygfx (python)

Issue:
Interpretation of colormaps
I was personally surprised with the results of using a very simple colormap `[(0, 0, 0), (1, 1, 1)]` with the image visuals.


![image](https://github.com/pygfx/pygfx/assets/90008/26af706a-6d60-4fd9-b637-f1983be19547)

What I found was that the colormap was sampled as if there were only "two" pixels placed at the position `0.25`, and `0.75`.

This sampling makes sense maybe for images, but it didn't jive with my expectation of colormaps.

In my own shader, I am shifting things and rescaling them so as to sample the full "range" of the colormap. 

1D example
```
           let sizef = f32(textureDimensions(t_colormap));
           let sample_coord = texcoord * (sizef - 1) / sizef + 0.5 / sizef;
           let value = textureSample(t_colormap, s_colormap, sample_coord);
```

For example, sampling a matplotlib colormap at a value of 0.1 will give you something much closer to 0.1 than the equivalent 2 value colormap in pygfx (pygfx will clamp it to 0).

<details><smmary>Matplotlib code</summary>

```python
import numpy as np
from matplotlib.colors import LinearSegmentedColormap

colors = [(0, 0, 0), (1, 1, 1)]

custom_cmap = LinearSegmentedColormap.from_list('custom_cmap', colors)

print(custom_cmap(0.1))
```

</details>

<details> <summary>PyGFX code</summary>

```python
import numpy as np
from wgpu.gui.auto import WgpuCanvas, run
import pygfx as gfx

canvas = WgpuCanvas(size=(512 * 2 + 100, 400), max_fps=999)
renderer = gfx.renderers.WgpuRenderer(canvas, show_fps=True)
scene = gfx.Scene()
camera = gfx.OrthographicCamera(512 * 2 + 50 + 100, 550)
camera.local.y = 256
camera.local.scale_y = -1
camera.local.x = (512 * 2 + 50)/ 2

cmap = gfx.Texture(
    np.asarray([(0, 0, 0), (1, 1, 1)], dtype='float32'),
    dim=1
)

img = np.arange(256, dtype='float32').reshape(16, 16)
geometry = gfx.Geometry(grid=gfx.Texture(img, dim=2))

image0 = gfx.Image(
    geometry,
    gfx.ImageBasicMaterial(clim=(0, 255))
)
image0.local.scale = 512 / 16, 512/16, 1

scene.add(image0)

image1 = gfx.Image(
    geometry,
    gfx.ImageBasicMaterial(clim=(0, 255), map=cmap)
)
image1.local.scale = 512 / 16, 512/16, 1
image1.local.x = 50 + 512

scene.add(image1)

def animate():
    renderer.render(scene, camera)

if __name__ == "__main__":
    canvas.request_draw(animate)
    run()
```

</details>

Maybe there is an easier way to do this by changing how the GPU interpolates a texture?

Thank you for the discussion.

Pull request:
Support for custom colormaps, and lazy loading
Basically what would have been the correct way to close #655 😉 

* [x] Function `create_colormap()` creates a colormap from a list of color values. Adjusted from an implementation in visvis.
* [x] More colormaps (also from visvis, I think they're mostly from Matlab originally).
* [x] Lazy loading of colormaps using module-level `getattr()` so that these textures are only created when needed.  

Patch:
diff --git a/pygfx/utils/cm.py b/pygfx/utils/cm.py
index 835aafa..b69bce9 100644
--- a/pygfx/utils/cm.py
+++ b/pygfx/utils/cm.py
@@ -1,19 +1,32 @@
 """Predefined Colormaps.
 
-The following colormaps are currently available:
-
 .. currentmodule:: pygfx.utils.cm
 
 .. autosummary::
 
+    create_colormap
+    gray
     viridis
     plasma
     inferno
     magma
     cividis
+    cool
+    hot
+    bone
+    copper
+    pink
+    spring
+    summer
+    autumn
+    winter
+    jet
 
 """
 
+import numpy as np
+
+from ..resources import Texture, TextureMap
 from ._cmdata_mpl import (
     _viridis_data,
     _plasma_data,
@@ -22,14 +35,110 @@ from ._cmdata_mpl import (
     _cividis_data,
 )
 
-from ..resources import Texture, TextureMap
+
+_data = {
+    # Simple maps
+    "gray": [(0, 0, 0, 0), (1, 1, 1, 1)],
+    "cool": [(0, 1, 1), (1, 0, 1)],
+    "hot": [(0, 0, 0), (1, 0, 0), (1, 1, 0), (1, 1, 1)],
+    "bone": [(0, 0, 0), (0.333, 0.333, 0.444), (0.666, 0.777, 0.777), (1, 1, 1)],
+    "copper": [(0, 0, 0), (1, 0.7, 0.5)],
+    "pink": [(0.1, 0, 0), (0.75, 0.5, 0.5), (0.9, 0.9, 0.7), (1, 1, 1)],
+    "spring": [(1, 0, 1), (1, 1, 0)],
+    "summer": [(0, 0.5, 0.4), (1, 1, 0.4)],
+    "autumn": [(1, 0, 0), (1, 1, 0)],
+    "winter": [(0, 0, 1), (0, 1, 0.5)],
+    "jet": [
+        (0, 0, 0.5),
+        (0, 0, 1),
+        (0, 0.5, 1),
+        (0, 1, 1),
+        (0.5, 1, 0.5),
+        (1, 1, 0),
+        (1, 0.5, 0),
+        (1, 0, 0),
+        (0.5, 0, 0),
+    ],
+    # Larger maps
+    "viridis": _viridis_data,
+    "plasma": _plasma_data,
+    "inferno": _inferno_data,
+    "magma": _magma_data,
+    "cividis": _cividis_data,
+}
+
+__all__ = sorted([*_data.keys(), "create_colormap"])
+
+
+def create_colormap(input, n=256):
+    """Create a colormap from given data.
+
+    Returns a TextureMap with (at least) n colors.
+
+    If a list or array-like object is given, the values are interpolated to create the map.
+    E.g. ``create_colormap([(0, 0, 0), (1, 1, 1)])`` creates the gray colormap.
+
+    If a dict is given, it must have at least one field of "r", "g", "b", "a".
+    Each value is a sequence of 2 values (i.e. an Nx2 array-like), where the
+    first value specifies the index (0..1), and the second the value for that color.
+    """
+
+    if isinstance(input, dict):
+        # Init data, alpha 1
+        colormap_data = np.zeros((n, 4), "f4")
+        colormap_data[:, 3] = 1.0
+        # For each channel ...
+        for channel_index in range(4):
+            channe_name = "rgba"[channel_index]
+            try:
+                values = input[channe_name]
+            except KeyError:
+                continue  # this color not in dict
+            # Get value list and check
+            values = np.asarray(values)
+            if not (values.ndim == 2 and values.shape[1] == 2):
+                raise ValueError(
+                    "The values in the dict to convert to a colormap must be Nx2."
+                )
+            # Init interpolation
+            data = np.zeros((len(values),), dtype="f4")
+            x = np.linspace(0.0, 1.0, n)
+            xp = np.zeros((len(values),), dtype="f4")
+            # Insert values
+            for i, el in enumerate(values):
+                xp[i] = el[0]
+                data[i] = el[1]
+            # Interpolate
+            colormap_data[:, channel_index] = np.interp(x, xp, data)
+
+    else:
+        # Input is list or array-like
+        data = np.array(input, dtype="f4")
+        if not (data.ndim == 2 and data.shape[1] in [3, 4]):
+            raise ValueError("Colormap entries must have 3 or 4 elements.")
+        if data.shape[1] == 3:
+            data = np.hstack([data, np.ones((data.shape[0], 1), "f4")])
+
+        # Apply interpolation (if required)
+        if data.shape[0] == n:
+            colormap_data = data
+        else:
+            x = np.linspace(0.0, 1.0, n)
+            xp = np.linspace(0.0, 1.0, data.shape[0])
+            colormap_data = np.zeros((n, 4), "f4")
+            for i in range(4):
+                colormap_data[:, i] = np.interp(x, xp, data[:, i])
+
+    return TextureMap(Texture(colormap_data, dim=1), filter="nearest", wrap="clamp")
 
 
-__all__ = ["cividis", "inferno", "magma", "plasma", "viridis"]
+def __dir__():
+    return __all__
 
 
-viridis = TextureMap(Texture(_viridis_data, dim=1), filter="nearest", wrap="clamp")
-plasma = TextureMap(Texture(_plasma_data, dim=1), filter="nearest", wrap="clamp")
-inferno = TextureMap(Texture(_inferno_data, dim=1), filter="nearest", wrap="clamp")
-magma = TextureMap(Texture(_magma_data, dim=1), filter="nearest", wrap="clamp")
-cividis = TextureMap(Texture(_cividis_data, dim=1), filter="nearest", wrap="clamp")
+def __getattr__(name):
+    if name not in _data:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    colormap = create_colormap(_data[name])
+    globals()[name] = colormap  # prevent creating it again
+    return colormap


Test patch:
diff --git a/tests/utils/test_colormaps.py b/tests/utils/test_colormaps.py
new file mode 100644
index 0000000..da45160
--- /dev/null
+++ b/tests/utils/test_colormaps.py
@@ -0,0 +1,88 @@
+import numpy as np
+import pygfx as gfx
+from pygfx.utils import cm
+import pytest
+
+
+def test_create_colormap_basic():
+    tm1 = cm.create_colormap([(0, 0, 0), (1, 0.2, 0.3)])
+    tm2 = cm.create_colormap([(0, 0, 0, 1), (1, 0.2, 0.3, 1)])
+
+    assert isinstance(tm1, gfx.TextureMap)
+    assert isinstance(tm2, gfx.TextureMap)
+    assert np.all(tm1.texture.data == tm2.texture.data)
+
+    data = tm1.texture.data
+    assert data.shape == (256, 4)
+
+    assert tuple(data[0]) == (0, 0, 0, 1)
+    assert tuple(data[-1]) == (1, 0.2, 0.3, 1)
+
+    # In the middle
+    assert np.allclose(data[127], (0.5, 0.1, 0.15, 1), atol=0.002)
+    assert np.allclose(data[128], (0.5, 0.1, 0.15, 1), atol=0.002)
+
+    # But it does increase
+    assert np.all(data[128, :3] > data[127, :3])
+
+    # Fails
+
+    with pytest.raises(ValueError):
+        cm.create_colormap([(0, 0), (1, 1)])
+
+    with pytest.raises(ValueError):
+        cm.create_colormap([(0, 0, 0, 0, 0), (1, 1, 1, 1, 1)])
+
+
+def test_create_colormap_size():
+    # Create large colormap
+    tm1 = cm.create_colormap([(0, 0, 0), (1, 0.2, 0.3)], 1024)
+    assert isinstance(tm1, gfx.TextureMap)
+    assert tm1.texture.data.shape == (1024, 4)
+
+    # Use large data with n=256, reduces
+    tm2 = cm.create_colormap(tm1.texture.data, 256)
+    assert isinstance(tm2, gfx.TextureMap)
+    assert tm2.texture.data.shape == (256, 4)
+
+    # Use same size
+    tm3 = cm.create_colormap(tm2.texture.data, 256)
+    assert isinstance(tm3, gfx.TextureMap)
+    assert tm3.texture.data.shape == (256, 4)
+
+    # Check that copy is made
+    assert tm3.texture.data is not tm2.texture.data
+
+
+def test_create_colormap_dict():
+    tm = cm.create_colormap(
+        {
+            "g": [(0, 0), (1, 1)],
+            "b": [(0.5, 0), (1, 1)],
+        },
+        101,
+    )
+
+    assert isinstance(tm, gfx.TextureMap)
+    data = tm.texture.data
+    assert data.shape == (101, 4)
+
+    assert tuple(data[0]) == (0, 0, 0, 1)
+    assert tuple(data[25]) == (0, 0.25, 0, 1)
+    assert tuple(data[50]) == (0, 0.5, 0, 1)
+    assert tuple(data[75]) == (0, 0.75, 0.5, 1)
+    assert tuple(data[100]) == (0, 1, 1, 1)
+
+
+def test_colormap_reuse():
+    # This fails if we don't implement the __getattr__ in cm.py correctly
+    tm1 = cm.viridis
+    tm2 = cm.viridis
+    assert tm1 is tm2
+
+
+if __name__ == "__main__":
+    test_create_colormap_basic()
+    test_create_colormap_size()
+    test_create_colormap_dict()
+    test_colormap_reuse()