Skip to content

Commit 43e7210

Browse files
PierreRaybautCopilot
PierreRaybaut
and
Copilot
committed
perf(image): cache colormap lookup tables by colormap content
Every new image item recomputed its 1024-entry lookup table (and the 256-color table) in Python, twice, even for the colormap used by all the other images. Tables are now cached by colormap content, inversion and alpha settings, so that an edited colormap still gets a new table. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 7b7add1 commit 43e7210

3 files changed

Lines changed: 195 additions & 32 deletions

File tree

‎doc/release_notes/release_2.11.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
* **Plot refresh** — Refreshing a plot no longer recomputes the plot layout when the title and axis styles did not change
1212
* **Item tools and item list** — Updating the status of item-related tools and of the item list panel is no longer quadratic in the number of plot items
13+
* **Image creation** — Creating an image item no longer recomputes the colormap lookup table when the same colormap, inversion and transparency settings were already used: the table (more than 1000 colormap evaluations) is computed once and reused, which saves a few milliseconds per image
1314

1415
## PlotPy Version 2.11.1 ##
1516

‎plotpy/items/image/base.py‎

Lines changed: 111 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import os
99
import os.path as osp
10+
from collections import OrderedDict
1011
from typing import TYPE_CHECKING, Any
1112

1213
import numpy as np
@@ -46,6 +47,7 @@
4647
from plotpy.mathutils.arrayfuncs import get_nan_range
4748
from plotpy.mathutils.colormap import FULLRANGE, get_cmap
4849
from plotpy.styles.image import RawImageParam
50+
from plotpy.widgets.colormap.widget import EditableColormap
4951

5052
if TYPE_CHECKING:
5153
import guidata.io
@@ -57,7 +59,110 @@
5759
from plotpy.interfaces import IItemType
5860
from plotpy.plot import BasePlot
5961
from plotpy.styles.base import ItemParameters
60-
from plotpy.widgets.colormap.widget import EditableColormap
62+
63+
#: Lookup tables already computed by :func:`_get_colormap_lut`, most recent last
64+
_LUT_CACHE: OrderedDict[tuple, tuple[tuple[int, ...], np.ndarray]] = OrderedDict()
65+
_LUT_CACHE_SIZE = 32
66+
67+
68+
def _lut_cache_key(
69+
table: EditableColormap, alpha: float, alpha_function: int
70+
) -> tuple | None:
71+
"""Return the key identifying the lookup table of a colormap, if cacheable.
72+
73+
The key is built from the colormap content (not its identity), so that an
74+
edited colormap gets a new lookup table.
75+
76+
Args:
77+
table: Colormap
78+
alpha: Global alpha value
79+
alpha_function: Alpha function (see :class:`plotpy.constants.LUTAlpha`)
80+
81+
Returns:
82+
Cache key, or None if the lookup table must be computed: colormap
83+
subclasses may compute colors differently, and the colors of translucent
84+
stops depend on the order in which the stops were inserted.
85+
"""
86+
if type(table) is not EditableColormap:
87+
return None
88+
stops = tuple((stop.pos, stop.rgb) for stop in table.colorStops())
89+
if any(stop.a != 255 for stop in table.colorStops()):
90+
return None
91+
return (stops, table.mode(), table.invert, alpha, alpha_function)
92+
93+
94+
def _compute_colormap_lut(
95+
table: EditableColormap, alpha: float, alpha_function: int
96+
) -> tuple[tuple[int, ...], np.ndarray]:
97+
"""Compute the color table and the lookup table of a colormap.
98+
99+
Args:
100+
table: Colormap
101+
alpha: Global alpha value
102+
alpha_function: Alpha function (see :class:`plotpy.constants.LUTAlpha`)
103+
104+
Returns:
105+
Color table (256 colors) and lookup table (``LUT_SIZE`` ARGB values)
106+
"""
107+
colors = tuple(table.colorTable(FULLRANGE))
108+
lut = np.zeros((LUT_SIZE,), np.uint32)
109+
for i in range(LUT_SIZE):
110+
if alpha_function == LUTAlpha.NONE.value:
111+
pix_alpha = 1.0
112+
elif alpha_function == LUTAlpha.CONSTANT.value:
113+
pix_alpha = alpha
114+
else:
115+
x = i / float(LUT_SIZE - 1)
116+
if alpha_function == LUTAlpha.LINEAR.value:
117+
# Linear alpha function
118+
pix_alpha = alpha * x
119+
elif alpha_function == LUTAlpha.SIGMOID.value:
120+
# Sigmoid alpha function
121+
pix_alpha = alpha / (1 + np.exp(-10 * x))
122+
elif alpha_function == LUTAlpha.TANH.value:
123+
# Hyperbolic tangent alpha function
124+
pix_alpha = alpha * np.tanh(5 * x)
125+
elif alpha_function == LUTAlpha.STEP.value:
126+
# Fully transparent lowest value and `alpha` transparent elsewhere
127+
pix_alpha = alpha if x > 0 else 0
128+
else:
129+
raise ValueError(f"Invalid alpha function {alpha_function}")
130+
# pylint: disable=too-many-function-args
131+
alpha_channel = max(min(np.uint32(255 * pix_alpha + 0.5), 255), 0) << 24
132+
lut[i] = (
133+
np.uint32((table.rgb(FULLRANGE, i / LUT_MAX)) & 0xFFFFFF) | alpha_channel
134+
)
135+
return colors, lut
136+
137+
138+
def _get_colormap_lut(
139+
table: EditableColormap, alpha: float, alpha_function: int
140+
) -> tuple[tuple[int, ...], np.ndarray]:
141+
"""Return the color table and the lookup table of a colormap.
142+
143+
Computing them takes more than 1000 calls to the colormap: they are cached,
144+
since every new image item usually shows the same colormap.
145+
146+
Args:
147+
table: Colormap
148+
alpha: Global alpha value
149+
alpha_function: Alpha function (see :class:`plotpy.constants.LUTAlpha`)
150+
151+
Returns:
152+
Color table (256 colors) and lookup table (``LUT_SIZE`` ARGB values),
153+
the latter being read-only
154+
"""
155+
key = _lut_cache_key(table, alpha, alpha_function)
156+
if key is not None and key in _LUT_CACHE:
157+
_LUT_CACHE.move_to_end(key)
158+
return _LUT_CACHE[key]
159+
colors, lut = _compute_colormap_lut(table, alpha, alpha_function)
160+
lut.flags.writeable = False
161+
if key is not None:
162+
_LUT_CACHE[key] = (colors, lut)
163+
if len(_LUT_CACHE) > _LUT_CACHE_SIZE:
164+
_LUT_CACHE.popitem(last=False)
165+
return colors, lut
61166

62167

63168
class BaseImageItem(QwtPlotItem):
@@ -466,37 +571,11 @@ def set_color_map(
466571
if invert is not None:
467572
table.invert = invert
468573
self.cmap_table = table
469-
self.cmap = table.colorTable(FULLRANGE)
470-
cmap_a = self.lut[3]
471-
alpha = self.param.alpha
472-
alpha_function = self.param.alpha_function
473-
for i in range(LUT_SIZE):
474-
if alpha_function == LUTAlpha.NONE.value:
475-
pix_alpha = 1.0
476-
elif alpha_function == LUTAlpha.CONSTANT.value:
477-
pix_alpha = alpha
478-
else:
479-
x = i / float(LUT_SIZE - 1)
480-
if alpha_function == LUTAlpha.LINEAR.value:
481-
# Linear alpha function
482-
pix_alpha = alpha * x
483-
elif alpha_function == LUTAlpha.SIGMOID.value:
484-
# Sigmoid alpha function
485-
pix_alpha = alpha / (1 + np.exp(-10 * x))
486-
elif alpha_function == LUTAlpha.TANH.value:
487-
# Hyperbolic tangent alpha function
488-
pix_alpha = alpha * np.tanh(5 * x)
489-
elif alpha_function == LUTAlpha.STEP.value:
490-
# Fully transparent lowest value and `alpha` transparent elsewhere
491-
pix_alpha = alpha if x > 0 else 0
492-
else:
493-
raise ValueError(f"Invalid alpha function {alpha_function}")
494-
# pylint: disable=too-many-function-args
495-
alpha_channel = max(min(np.uint32(255 * pix_alpha + 0.5), 255), 0) << 24
496-
cmap_a[i] = (
497-
np.uint32((table.rgb(FULLRANGE, i / LUT_MAX)) & 0xFFFFFF)
498-
| alpha_channel
499-
)
574+
colors, lut = _get_colormap_lut(
575+
table, self.param.alpha, self.param.alpha_function
576+
)
577+
self.cmap = list(colors)
578+
self.lut[3][:] = lut
500579
plot = self.plot()
501580
if plot:
502581
plot.update_colormap_axis(self)

‎plotpy/tests/unit/test_colormap_lut_cache.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# -*- coding: utf-8 -*-
2+
#
3+
# Licensed under the terms of the BSD 3-Clause
4+
# (see plotpy/LICENSE for details)
5+
6+
"""Image lookup tables are cached by colormap content: a cached table always
7+
matches the colormap, its inversion and the alpha settings."""
8+
9+
from __future__ import annotations
10+
11+
import copy
12+
13+
import numpy as np
14+
import pytest
15+
from guidata.qthelpers import qt_app_context
16+
from qtpy import QtGui as QG
17+
18+
from plotpy.builder import make
19+
from plotpy.constants import LUT_MAX, LUT_SIZE, LUTAlpha
20+
from plotpy.mathutils.colormap import FULLRANGE, get_cmap
21+
22+
23+
def _expected_lut(table, alpha: float, alpha_function: int) -> np.ndarray:
24+
"""Lookup table computed directly from the colormap."""
25+
x = np.arange(LUT_SIZE) / float(LUT_SIZE - 1)
26+
pix_alpha = {
27+
LUTAlpha.NONE.value: np.ones(LUT_SIZE),
28+
LUTAlpha.CONSTANT.value: np.full(LUT_SIZE, alpha),
29+
LUTAlpha.LINEAR.value: alpha * x,
30+
LUTAlpha.SIGMOID.value: alpha / (1 + np.exp(-10 * x)),
31+
LUTAlpha.TANH.value: alpha * np.tanh(5 * x),
32+
LUTAlpha.STEP.value: np.where(x > 0, alpha, 0.0),
33+
}[alpha_function]
34+
alpha_channel = np.clip((255 * pix_alpha + 0.5).astype(np.uint32), 0, 255) << 24
35+
rgb = np.array(
36+
[table.rgb(FULLRANGE, i / LUT_MAX) & 0xFFFFFF for i in range(LUT_SIZE)],
37+
dtype=np.uint32,
38+
)
39+
return rgb | alpha_channel.astype(np.uint32)
40+
41+
42+
def _image(colormap: str = "jet", **kwargs):
43+
return make.image(np.zeros((4, 4)), colormap=colormap, **kwargs)
44+
45+
46+
@pytest.mark.parametrize("alpha_function", list(LUTAlpha))
47+
def test_lut_matches_colormap(alpha_function: LUTAlpha) -> None:
48+
"""Each image gets the lookup table of its colormap, alpha and inversion."""
49+
with qt_app_context(exec_loop=False):
50+
for invert in (False, True, False):
51+
for name in ("jet", "viridis", "gray"):
52+
item = _image(name, alpha_function=alpha_function, alpha=0.6)
53+
item.set_color_map(name, invert=invert)
54+
expected = _expected_lut(get_cmap(name), 0.6, alpha_function.value)
55+
assert np.array_equal(item.lut[3], expected)
56+
assert item.cmap == get_cmap(name).colorTable(FULLRANGE)
57+
58+
59+
def test_lut_follows_edited_colormap() -> None:
60+
"""Editing a colormap after a first use gives a new lookup table."""
61+
with qt_app_context(exec_loop=False):
62+
table = copy.deepcopy(get_cmap("jet"))
63+
first = _image()
64+
first.set_color_map(table)
65+
before = first.lut[3].copy()
66+
67+
table.addColorStop(0.5, QG.QColor(255, 255, 255))
68+
second = _image()
69+
second.set_color_map(table)
70+
expected = _expected_lut(table, 1.0, LUTAlpha.NONE.value)
71+
assert np.array_equal(second.lut[3], expected)
72+
assert not np.array_equal(second.lut[3], before)
73+
74+
75+
def test_items_do_not_share_lookup_tables() -> None:
76+
"""Changing the lookup table of an item leaves the other items unchanged."""
77+
with qt_app_context(exec_loop=False):
78+
first, second = _image(), _image()
79+
first.lut[3][:] = 0
80+
third = _image()
81+
expected = _expected_lut(get_cmap("jet"), 1.0, LUTAlpha.NONE.value)
82+
assert np.array_equal(second.lut[3], expected)
83+
assert np.array_equal(third.lut[3], expected)

0 commit comments

Comments
 (0)