diff --git a/src/plopp/backends/pythreejs/scatter3d.py b/src/plopp/backends/pythreejs/scatter3d.py index 20050b91..72bdbd11 100644 --- a/src/plopp/backends/pythreejs/scatter3d.py +++ b/src/plopp/backends/pythreejs/scatter3d.py @@ -2,7 +2,7 @@ # Copyright (c) 2023 Scipp contributors (https://github.com/scipp) import uuid -from typing import Literal +from typing import Literal, cast import numpy as np import pythreejs as p3 @@ -75,8 +75,12 @@ def __init__( self._x = x self._y = y self._z = z - self._color = color - self._artist_number = artist_number + self._color = None + if colormapper is None: + self._color = np.array( + to_rgb(f'C{artist_number}' if color is None else color), + dtype='float32', + ) # TODO: remove pixel_size in the next release self._size = size if pixel_size is None else pixel_size @@ -124,18 +128,8 @@ def _make_positions(self) -> np.ndarray: def _make_colors(self) -> np.ndarray: if self._colormapper is not None: return self._colormapper.rgba(self._data)[..., :3].astype('float32') - else: - return np.broadcast_to( - np.array( - to_rgb( - f'C{self._artist_number}' - if self._color is None - else self._color - ), - dtype='float32', - ), - (self._data.coords[self._x].shape[0], 3), - ) + color = cast(np.ndarray, self._color) + return np.broadcast_to(color, (self._data.coords[self._x].shape[0], 3)) def _make_geometry( self, positions: np.ndarray, colors: np.ndarray @@ -209,9 +203,25 @@ def color(self) -> np.ndarray: return self.geometry.attributes['color'].array @color.setter - def color(self, val: np.ndarray): + def color( + self, + val: str + | tuple[float, float, float] + | tuple[float, float, float, float] + | np.ndarray, + ) -> None: + if isinstance(val, (str, tuple)) or val.ndim == 1: + color = val.tolist() if isinstance(val, np.ndarray) else val + if self._colormapper is None: + self._color = np.array(to_rgb(color), dtype='float32') + val = self._make_colors() self.geometry.attributes['color'].array = val + @property + def single_color(self) -> np.ndarray | None: + """The single RGB color, or ``None`` if a colormapper is in use.""" + return self._color + @property def geometry(self) -> p3.BufferGeometry: """ diff --git a/src/plopp/widgets/clip3d.py b/src/plopp/widgets/clip3d.py index f9ab6921..a1168645 100644 --- a/src/plopp/widgets/clip3d.py +++ b/src/plopp/widgets/clip3d.py @@ -503,6 +503,11 @@ def update_controls(self): self._nodes.clear() self.update_state() + if self._view.colormapper is None: + for source_id, node in self._nodes.items(): + color = self._view.artists[source_id].single_color + if color is not None: + self._view.artists[node.id].color = color def _set_opacity(self, change: dict[str, Any]): """ diff --git a/tests/backends/pythreejs/pythreejs_scatter3d_test.py b/tests/backends/pythreejs/pythreejs_scatter3d_test.py index 257c4975..cf7864fd 100644 --- a/tests/backends/pythreejs/pythreejs_scatter3d_test.py +++ b/tests/backends/pythreejs/pythreejs_scatter3d_test.py @@ -4,6 +4,8 @@ import numpy as np import pytest import scipp as sc +from matplotlib import cycler, rc_context +from matplotlib.colors import to_rgb from plopp.backends.pythreejs.canvas import Canvas from plopp.backends.pythreejs.scatter3d import Scatter3d @@ -30,6 +32,66 @@ def test_update(): assert sc.identical(scat._data, da * 2.5) +def test_update_with_different_number_of_points_preserves_default_color(): + with rc_context({'axes.prop_cycle': cycler(color=['#102030'])}): + scat = Scatter3d(canvas=Canvas(), data=scatter(20), x='x', y='y', z='z') + color = scat.color[0].copy() + + with rc_context({'axes.prop_cycle': cycler(color=['#708090'])}): + scat.update(scatter(30)) + + np.testing.assert_array_equal(scat.color, np.broadcast_to(color, (30, 3))) + + +@pytest.mark.parametrize( + 'color', + [ + 'red', + (0.1, 0.2, 0.3), + (0.1, 0.2, 0.3, 0.4), + np.array([0.1, 0.2, 0.3]), + np.array([0.1, 0.2, 0.3, 0.4]), + ], +) +def test_set_single_color(color): + scat = Scatter3d(canvas=Canvas(), data=scatter(20), x='x', y='y', z='z') + + scat.color = color + + expected = np.array(to_rgb(color), dtype='float32') + np.testing.assert_array_equal(scat.single_color, expected) + np.testing.assert_array_equal(scat.color, np.broadcast_to(expected, (20, 3))) + + +def test_set_single_color_is_preserved_when_number_of_points_changes(): + scat = Scatter3d(canvas=Canvas(), data=scatter(20), x='x', y='y', z='z') + scat.color = 'red' + + scat.update(scatter(30)) + + expected = np.array(to_rgb('red'), dtype='float32') + np.testing.assert_array_equal(scat.single_color, expected) + np.testing.assert_array_equal(scat.color, np.broadcast_to(expected, (30, 3))) + + +def test_single_color_is_none_with_colormapper(): + canvas = Canvas() + scat = Scatter3d( + canvas=canvas, + data=scatter(20), + x='x', + y='y', + z='z', + colormapper=ColorMapper(canvas=canvas), + ) + colors = scat.color.copy() + + scat.color = 'red' + + assert scat.single_color is None + np.testing.assert_array_equal(scat.color, colors) + + def test_bounding_box(): da = scatter() pix = 0.5 diff --git a/tests/widgets/clip3d_test.py b/tests/widgets/clip3d_test.py index ab5d0b18..c81f2032 100644 --- a/tests/widgets/clip3d_test.py +++ b/tests/widgets/clip3d_test.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: BSD-3-Clause # Copyright (c) 2023 Scipp contributors (https://github.com/scipp) +import numpy as np import pytest import scipp as sc +from matplotlib import cycler, rc_context from plopp import Node from plopp.data.testing import data_array, scatter @@ -9,6 +11,40 @@ from plopp.widgets import ClippingManager +def test_cut_colors_match_original_artists(): + a = scatter(seed=1) + b = scatter(seed=2) + b.coords['x'] += sc.scalar(60, unit='m') + nodes = [Node(a), Node(b)] + with rc_context({'axes.prop_cycle': cycler(color=['#102030', '#405060'])}): + fig = scatter3dfigure(*nodes, x='x', y='y', z='z', cbar=False) + clip = ClippingManager(fig) + + with rc_context({'axes.prop_cycle': cycler(color=['#708090', '#a0b0c0'])}): + clip.add_y_cut.click() + + for source in nodes: + cut = clip._nodes[source.id] + source_artist = fig.artists[source.id] + cut_artist = fig.artists[cut.id] + assert np.array_equal(cut_artist.single_color, source_artist.single_color) + assert np.array_equal(cut_artist.color[0], source_artist.color[0]) + + npoints = [fig.artists[node.id].data.shape[0] for node in clip._nodes.values()] + ycut = clip.cuts[-1] + ycut.slider.value = [ycut.slider.min, ycut.slider.value[1]] + clip.update_state() + assert [ + fig.artists[node.id].data.shape[0] for node in clip._nodes.values() + ] != npoints + for source in nodes: + cut = clip._nodes[source.id] + source_artist = fig.artists[source.id] + cut_artist = fig.artists[cut.id] + assert np.array_equal(cut_artist.single_color, source_artist.single_color) + assert np.array_equal(cut_artist.color[0], source_artist.color[0]) + + @pytest.mark.parametrize('multiple_nodes', [False, True]) def test_add_remove_cuts(multiple_nodes): a = scatter()