diff --git a/src/plopp/plotting/inspector.py b/src/plopp/plotting/inspector.py index 89104787..311cb195 100644 --- a/src/plopp/plotting/inspector.py +++ b/src/plopp/plotting/inspector.py @@ -3,6 +3,7 @@ from typing import Literal +import numpy as np import scipp as sc from ..core import Node @@ -34,6 +35,35 @@ def _slice_xy(da: sc.DataArray, xy: dict[str, dict[str, int]]) -> sc.DataArray: return da[y['dim'], y['value']][x['dim'], x['value']] +def _coord_to_centers(da: sc.DataArray, dim: str) -> sc.Variable: + coord = da.coords[dim] + if da.coords.is_edges(dim, dim=dim): + return sc.midpoints(coord, dim=dim) + return coord + + +def _mask_outside_polygon( + da: sc.DataArray, poly: dict[str, dict[str, sc.Variable]] +) -> sc.DataArray: + from matplotlib.path import Path + + xdim = poly['x']['dim'] + ydim = poly['y']['dim'] + x = _coord_to_centers(da, xdim) + y = _coord_to_centers(da, ydim) + vx = poly['x']['value'].to(unit=x.unit).values + vy = poly['y']['value'].to(unit=y.unit).values + verts = np.column_stack([vx, vy]) + path = Path(verts) + xx = sc.broadcast(x, sizes={**x.sizes, **y.sizes}) + yy = sc.broadcast(y, sizes={**x.sizes, **y.sizes}) + points = np.column_stack([xx.values.ravel(), yy.values.ravel()]) + inside = sc.array( + dims=yy.dims, values=path.contains_points(points).reshape(yy.shape) + ) + return da.assign_masks(__inside_polygon=~inside).sum({*x.dims, *y.dims}) + + def inspector( obj: Plottable, dim: str | None = None, @@ -51,6 +81,7 @@ def inspector( logc: bool | None = None, mask_cmap: str = 'gray', mask_color: str = 'black', + mode: Literal['point', 'polygon'] = 'point', nan_color: str | None = None, norm: Literal['linear', 'log'] | None = None, operation: Literal['sum', 'mean', 'min', 'max'] = 'sum', @@ -70,16 +101,23 @@ def inspector( Inspector takes in a three-dimensional input and applies a reduction operation (``'sum'`` by default) along one of the dimensions specified by ``dim``. It displays the result as a two-dimensional image. - In addition, an 'inspection' tool is available in the toolbar which allows to place - markers on the image which perform slicing at that position to retain only the third - dimension and displays the resulting one-dimensional slice on the right hand side - figure. - - Controls: - - Click to make new point - - Drag existing point to move it + In addition, an 'inspection' tool is available in the toolbar. In ``mode='point'`` + it allows placing point markers on the image to slice at that position, retaining + only the third dimension and displaying the resulting one-dimensional slice in the + right-hand side figure. In ``mode='polygon'`` it allows drawing a polygon to compute + the total intensity inside the polygon as a function of the third dimension. + + Controls (point mode): + - Left-click to make new points + - Left-click and hold on point to move point - Middle-click to delete point + Controls (polygon mode): + - Left-click to make new polygons + - Left-click and hold on polygon vertex to move vertex + - Right-click and hold to drag/move the entire polygon + - Middle-click to delete polygon + Notes ----- @@ -123,6 +161,9 @@ def inspector( Colormap to use for masks. mask_color: Color of masks (overrides ``mask_cmap``). + mode: + Select ``'point'`` for point inspection or ``'polygon'`` for polygon selection + with total intensity inside the polygon plotted as a function of ``dim``. nan_color: Color to use for NaN values. norm: @@ -207,17 +248,28 @@ def inspector( ylabel=ylabel, **kwargs, ) + from ..widgets import Box, PointsTool, PolygonTool - from ..widgets import Box, PointsTool + if mode == 'point': + tool = PointsTool( + figure=f2d, + input_node=bin_edges_node, + func=_slice_xy, + destination=f1d, + tooltip="Activate inspector tool", + ) + elif mode == 'polygon': + tool = PolygonTool( + figure=f2d, + input_node=bin_edges_node, + func=_mask_outside_polygon, + destination=f1d, + tooltip="Activate polygon inspector tool", + ) + else: + raise ValueError(f'Mode "{mode}" is unknown.') - pts = PointsTool( - figure=f2d, - input_node=bin_edges_node, - func=_slice_xy, - destination=f1d, - tooltip="Activate inspector tool", - ) - f2d.toolbar['inspect'] = pts + f2d.toolbar['inspect'] = tool out = [f2d, f1d] if orientation == 'horizontal': out = [out] diff --git a/src/plopp/widgets/__init__.pyi b/src/plopp/widgets/__init__.pyi index 225f8a65..b7516d4f 100644 --- a/src/plopp/widgets/__init__.pyi +++ b/src/plopp/widgets/__init__.pyi @@ -4,7 +4,7 @@ from .box import Box, HBar, VBar from .checkboxes import Checkboxes from .clip3d import Clip3dTool, ClippingPlanes -from .drawing import DrawingTool, PointsTool +from .drawing import DrawingTool, PointsTool, PolygonTool from .linesave import LineSaveTool from .slice import RangeSliceWidget, SliceWidget, slice_dims from .toolbar import Toolbar, make_toolbar_canvas2d, make_toolbar_canvas3d @@ -21,6 +21,7 @@ __all__ = [ "HBar", "LineSaveTool", "PointsTool", + "PolygonTool", "RangeSliceWidget", "SliceWidget", "ToggleTool", diff --git a/src/plopp/widgets/drawing.py b/src/plopp/widgets/drawing.py index a65b9b45..8bcc307d 100644 --- a/src/plopp/widgets/drawing.py +++ b/src/plopp/widgets/drawing.py @@ -189,3 +189,73 @@ def _make_points(**kwargs): **kwargs: Additional arguments are forwarded to the ``ToggleTool`` constructor. """ + + +def _get_polygon_info(artist, figure): + """ + Convert the raw polygon vertices to a dict containing the dimensions of + each axis, and arrays with units. + """ + xs, ys = artist.x, artist.y + return { + 'x': { + 'dim': figure.canvas.dims['x'], + 'value': sc.array( + dims=['vertex'], + values=xs, + unit=figure.canvas.units['x'], + ), + }, + 'y': { + 'dim': figure.canvas.dims['y'], + 'value': sc.array( + dims=['vertex'], + values=ys, + unit=figure.canvas.units['y'], + ), + }, + } + + +def _make_polygons(**kwargs): + """ + Intermediate function needed for giving to `partial` to avoid making mpltoolbox a + hard dependency. + """ + from mpltoolbox import Polygons + + return Polygons(**kwargs) + + +PolygonTool = partial( + DrawingTool, + tool=partial(_make_polygons, mec='w'), + get_artist_info=_get_polygon_info, + icon='draw-polygon', +) +""" +Tool to draw polygon selections on a figure. + +Controls: + - Left-click to make new polygons + - Left-click and hold on polygon vertex to move vertex + - Right-click and hold to drag/move the entire polygon + - Middle-click to delete polygon + +Parameters +---------- +figure: + The figure where the tool will draw the polygon. +input_node: + The node that provides the raw data which is shown in ``figure``. +func: + The function to be used to make a node whose parents will be the ``input_node`` + and a node yielding the current state of the tool (current position, size). +destination: + Where the output from the ``func`` node will be then sent on. This can either + be a figure, or another graph node. +value: + Activate the tool upon creation if ``True``. +**kwargs: + Additional arguments are forwarded to the ``ToggleTool`` constructor. +""" diff --git a/tests/plotting/inspector_test.py b/tests/plotting/inspector_test.py index 5b776acd..0ac62550 100644 --- a/tests/plotting/inspector_test.py +++ b/tests/plotting/inspector_test.py @@ -63,6 +63,51 @@ def test_line_creation(): assert len(fig1d.artists) == 2 +def _finalize_polygon(tool, points): + """ + Click the polygon vertices and force creation without needing a close gesture. + """ + for x, y in points: + tool.click(x, y) + # Finalize to trigger the on_create callback without needing to close the polygon. + tool._finalize_owner() + + +def _coord_margin(values, index): + """ + Return a small margin around a coordinate to build a polygon around a bin center. + """ + if index == 0: + return (values[1] - values[0]) / 4.0 + if index == len(values) - 1: + return (values[-1] - values[-2]) / 4.0 + return ( + min(values[index] - values[index - 1], values[index + 1] - values[index]) / 4.0 + ) + + +def _polygon_vertices(da, xdim, ydim, x_index, y_indices): + """ + Build a rectangle polygon around the given x index and y index range. + """ + xvalues = da.coords[xdim].values + yvalues = da.coords[ydim].values + y_start = y_indices[0] + y_end = y_indices[-1] + x0 = xvalues[x_index] + y0 = yvalues[y_start] + y1 = yvalues[y_end] + x_margin = _coord_margin(xvalues, x_index) + y_margin_start = _coord_margin(yvalues, y_start) + y_margin_end = _coord_margin(yvalues, y_end) + return [ + (x0 - x_margin, y0 - y_margin_start), + (x0 + x_margin, y0 - y_margin_start), + (x0 + x_margin, y1 + y_margin_end), + (x0 - x_margin, y1 + y_margin_end), + ] + + @pytest.mark.usefixtures('_use_ipympl') def test_line_removal(): da = pp.data.data3d() @@ -83,6 +128,173 @@ def test_line_removal(): assert len(fig1d.artists) == 0 +def _polygon_case( + *, + dims, + values, + coords, + x_index, + y_indices, + keep_dim=None, +): + """ + Create a polygon test case with input data and the selection indices to compare. + """ + return ( + sc.DataArray( + data=sc.array(dims=dims, values=values), + coords={ + name: sc.array(dims=[name], values=vals) + for name, vals in coords.items() + }, + ), + x_index, + y_indices, + keep_dim, + ) + + +@pytest.mark.usefixtures('_use_ipympl') +@pytest.mark.parametrize( + ("da", "x_index", "y_indices", "keep_dim"), + [ + _polygon_case( + dims=["yy", "xx", "zz"], + values=[ + [[0, 1], [2, 3], [4, 5]], + [[6, 7], [8, 9], [10, 11]], + [[12, 13], [14, 15], [16, 17]], + ], + coords={ + "yy": [0.0, 10.0, 20.0], + "xx": [0.0, 5.0, 10.0], + "zz": [1.0, 2.0], + }, + x_index=1, + y_indices=[1], + keep_dim=None, + ), + _polygon_case( + dims=["row", "depth", "col"], + values=[ + [[0, 1, 2], [3, 4, 5], [6, 7, 8]], + [[9, 10, 11], [12, 13, 14], [15, 16, 17]], + [[18, 19, 20], [21, 22, 23], [24, 25, 26]], + [[27, 28, 29], [30, 31, 32], [33, 34, 35]], + ], + coords={ + "row": [0.0, 100.0, 200.0, 300.0], + "col": [0.0, 10.0, 20.0], + "depth": [0.0, 1.0, 2.0], + }, + x_index=1, + y_indices=[1, 2], + keep_dim="row", + ), + ], +) +def test_polygon_mode_data_values(da, x_index, y_indices, keep_dim): + """ + Polygon selection should produce the expected 1D curve over the keep dimension. + """ + if keep_dim is None: + keep_dim = da.dims[-1] + ip = pp.inspector(da, mode='polygon', dim=keep_dim) + fig2d = ip[0][0] + fig1d = ip[0][1] + fig2d.toolbar['inspect'].value = True + tool = fig2d.toolbar['inspect']._tool + xdim = fig2d.canvas.dims['x'] + ydim = fig2d.canvas.dims['y'] + keep_dim = keep_dim + points = _polygon_vertices(da, xdim, ydim, x_index=x_index, y_indices=y_indices) + _finalize_polygon(tool, points) + assert fig1d.canvas.dims == {'x': keep_dim} + assert len(fig1d.artists) == 1 + line = next(iter(fig1d.artists.values())) + if len(y_indices) == 1: + expected = da[ydim, y_indices[0]][xdim, x_index] + else: + expected = da[ydim, y_indices][xdim, x_index].sum(ydim) + expected = expected.drop_coords( + [name for name in expected.coords if name != keep_dim] + ) + assert sc.identical(line._data, expected) + + +@pytest.mark.usefixtures('_use_ipympl') +def test_polygon_mode_respects_masks_on_keep_dim(): + """ + Masks on the keep dimension should propagate to the 1D output. + """ + da = sc.DataArray( + data=sc.array( + dims=["yy", "xx", "zz"], + values=[ + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]], + ], + ), + coords={ + "yy": sc.array(dims=["yy"], values=[0.0, 1.0]), + "xx": sc.array(dims=["xx"], values=[0.0, 1.0]), + "zz": sc.array(dims=["zz"], values=[0.0, 1.0, 2.0]), + }, + masks={ + "masked": sc.array(dims=["zz"], values=[False, True, False]), + }, + ) + ip = pp.inspector(da, mode="polygon") + fig2d = ip[0][0] + fig1d = ip[0][1] + fig2d.toolbar["inspect"].value = True + tool = fig2d.toolbar["inspect"]._tool + xdim = fig2d.canvas.dims["x"] + ydim = fig2d.canvas.dims["y"] + points = _polygon_vertices(da, xdim, ydim, x_index=0, y_indices=[0, 1]) + _finalize_polygon(tool, points) + line = next(iter(fig1d.artists.values())) + expected = da[ydim, [0, 1]][xdim, 0].sum(ydim) + expected = expected.drop_coords( + [name for name in expected.coords if name != da.dims[-1]] + ) + assert sc.identical(line._data, expected) + assert "masked" in line._data.masks + + +@pytest.mark.usefixtures('_use_ipympl') +def test_polygon_mode_preserves_keep_dim_binedges(): + """ + Bin-edge coordinates on the keep dimension should be preserved in the 1D output. + """ + da = sc.DataArray( + data=sc.array( + dims=["yy", "xx", "zz"], + values=[ + [[0.0, 1.0], [2.0, 3.0]], + [[4.0, 5.0], [6.0, 7.0]], + ], + ), + coords={ + "yy": sc.array(dims=["yy"], values=[0.0, 1.0]), + "xx": sc.array(dims=["xx"], values=[0.0, 1.0]), + "zz": sc.array(dims=["zz"], values=[0.0, 1.0, 2.0]), + }, + ) + assert da.coords.is_edges("zz", dim="zz") + ip = pp.inspector(da, mode="polygon") + fig2d = ip[0][0] + fig1d = ip[0][1] + fig2d.toolbar["inspect"].value = True + tool = fig2d.toolbar["inspect"]._tool + xdim = fig2d.canvas.dims["x"] + ydim = fig2d.canvas.dims["y"] + points = _polygon_vertices(da, xdim, ydim, x_index=1, y_indices=[0, 1]) + _finalize_polygon(tool, points) + line = next(iter(fig1d.artists.values())) + assert line._data.coords.is_edges("zz", dim="zz") + + @pytest.mark.usefixtures('_use_ipympl') def test_operation(): da = pp.data.data3d()