import itertools
import cmap as cmap_lib
import numpy as np
from .line import LineGraphic
from .scatter import ScatterGraphic
from .image import ImageGraphic
from ._collection_base import GraphicCollection, cmap_across_graphics
from .selectors import (
LinearSelector,
LinearRegionSelector,
RectangleSelector,
PolygonSelector,
)
from ..utils import calculate_figure_shape
class PositionsCollection(GraphicCollection):
"""A collection of positions-based graphics (lines, scatters); adds selectors spanning all graphics."""
# appended as a Notes section to the generated constructor stub (scripts/generate_graphics_stubs.py),
# since the child graphic's docstring describes `cmap`/`cmap_transform` for a single graphic only
_stub_constructor_notes = """\
Notes
-----
``cmap`` and ``cmap_transform`` apply across the collection. A single ``cmap`` gives each graphic
one color spread across the colormap, selected by a 1D ``cmap_transform`` (one value per graphic).
An iterable of colormaps gives each graphic its own colormap along its datapoints, with a
per-graphic ``cmap_transform``."""
def __init__(
self, data, *, cmap=None, cmap_transform=None, cmap_range=None, **kwargs
):
super().__init__(data, **kwargs)
self._set_cmap(cmap, cmap_transform, cmap_range)
def _set_cmap(self, cmap, cmap_transform=None, cmap_range=None):
"""
A single cmap (str or ``cmap_lib.Colormap``) gives each graphic a uniform color.
An iterable of cmaps gives each graphic its own colormap.
"""
if hasattr(cmap, "__next__"):
# an iterator (itertools.repeat/cycle, a generator, ...): one cmap per graphic,
# materialized so re-applying it (e.g. each frame in NDPositions) stays stable
cmap = list(itertools.islice(cmap, len(self)))
self._cmap = cmap
self._cmap_transform = cmap_transform
self._cmap_range = cmap_range
if cmap is None:
if cmap_transform is not None:
raise ValueError("must pass `cmap` if passing `cmap_transform`")
return
single_cmap = isinstance(cmap, (str, cmap_lib.Colormap))
# a single cmap needs a 1D transform (across graphics), an iterable needs a 2D transform (per-graphic)
if cmap_transform is not None and single_cmap == (
np.ndim(cmap_transform[0]) >= 1
):
raise ValueError(
"`cmap` and `cmap_transform` must match: a single `cmap` uses a 1D transform, "
"an iterable of cmaps uses a 2D transform"
)
if single_cmap:
self.colors[:] = cmap_across_graphics(
cmap, len(self), cmap_transform, cmap_range
)
return
if len(cmap) != len(self):
raise IndexError(
f"len(cmap) must equal the number of graphics, got {len(cmap)} cmaps for "
f"{len(self)} graphics"
)
if cmap_transform is not None and len(cmap_transform) != len(self):
raise IndexError(
f"len(cmap_transform) must equal the number of graphics, got "
f"{len(cmap_transform)} `cmap_transform` values for {len(self)} graphics"
)
if np.ndim(cmap_range) == 2 and len(cmap_range) != len(self):
raise IndexError(
f"len(cmap_range) must equal the number of graphics, got {len(cmap_range)} "
f"`cmap_range` values for {len(self)} graphics"
)
transforms = (
cmap_transform if cmap_transform is not None else itertools.repeat(None)
)
ranges = (
cmap_range if np.ndim(cmap_range) == 2 else itertools.repeat(cmap_range)
)
for graphic, one_cmap, transform, rng in zip(
self.graphics, cmap, transforms, ranges
):
graphic.cmap = one_cmap
if transform is not None:
graphic.cmap_transform = transform
if rng is not None:
graphic.cmap_range = rng
@property
def cmap(self) -> str | list | None:
"""
get or set the colormap(s) across the collection
A single colormap gives each graphic one color, spread across the colormap. An iterable of
colormaps gives each graphic its own colormap along its datapoints.
"""
return self._cmap
@cmap.setter
def cmap(self, value):
self._set_cmap(value, self._cmap_transform, self._cmap_range)
@property
def cmap_transform(self) -> np.ndarray | None:
"""
get or set the cmap_transform across the collection
With a single ``cmap`` the transform is 1D, one value per graphic, selecting each graphic's
color. With an iterable of colormaps the transform is per-graphic, coloring each graphic
along its datapoints.
"""
return self._cmap_transform
@cmap_transform.setter
def cmap_transform(self, value):
self._set_cmap(self._cmap, value, self._cmap_range)
@property
def cmap_range(self) -> tuple[float, float] | None:
"""get or set the cmap_range of the graphics in the collection"""
return self._cmap_range
@cmap_range.setter
def cmap_range(self, value):
self._set_cmap(self._cmap, self._cmap_transform, value)
def add_linear_selector(
self, selection: float = None, padding: float = 0.0, axis: str = "x", **kwargs
) -> LinearSelector:
"""
Add a :class:`.LinearSelector`.
Parameters
----------
selection: float, optional
initial position of the selector along ``axis``, computed from the data if not given
padding: float, default 0.0
extra padding along the orthogonal axis to make the selector easier to grab
axis: str, default "x"
axis the selector moves along
**kwargs
passed to :class:`.LinearSelector`
Returns
-------
LinearSelector
"""
bounds_init, limits, size, center = self._get_linear_selector_init_args(
axis, padding
)
if selection is None:
selection = bounds_init[0]
selector = LinearSelector(
selection=selection, limits=limits, axis=axis, parent=self, **kwargs
)
self._plot_area.add_graphic(selector, center=False)
return selector
def add_linear_region_selector(
self,
selection: tuple[float, float] = None,
padding: float = 0.0,
axis: str = "x",
**kwargs,
) -> LinearRegionSelector:
"""
Add a :class:`.LinearRegionSelector`.
Parameters
----------
selection: (float, float), optional
initial bounds of the region along ``axis``, computed from the data if not given
padding: float, default 0.0
extra padding along the orthogonal axis to make the selector easier to grab
axis: str, default "x"
axis the selector spans
**kwargs
passed to :class:`.LinearRegionSelector`
Returns
-------
LinearRegionSelector
"""
bounds_init, limits, size, center = self._get_linear_selector_init_args(
axis, padding
)
if selection is None:
selection = bounds_init
selector = LinearRegionSelector(
selection=selection,
limits=limits,
size=size,
center=center,
axis=axis,
parent=self,
**kwargs,
)
self._plot_area.add_graphic(selector, center=False)
return selector
def add_rectangle_selector(
self, selection: tuple[float, float, float, float] = None, **kwargs
) -> RectangleSelector:
"""
Add a :class:`.RectangleSelector`.
Parameters
----------
selection: (float, float, float, float), optional
initial (xmin, xmax, ymin, ymax), computed from the data if not given
**kwargs
passed to :class:`.RectangleSelector`
Returns
-------
RectangleSelector
"""
bbox = self.world_object.get_world_bounding_box()
xdata = np.concatenate(self.data[:, :, 0])
xmin, xmax = np.nanmin(xdata), np.nanmax(xdata)
# y from the world bounding box so that the graphics' offsets, e.g. a stack's, are included
ymin, ymax = bbox[0, 1], bbox[1, 1]
yspan = ymax - ymin
if selection is None:
# the first quarter along x, the full y extent
selection = (xmin, xmin + (xmax - xmin) / 4, ymin, ymax)
limits = (xmin, xmax, ymin - yspan / 2, ymax + yspan / 2)
selector = RectangleSelector(
selection=selection, limits=limits, parent=self, **kwargs
)
self._plot_area.add_graphic(selector, center=False)
return selector
def add_polygon_selector(
self, selection: list[tuple[float, float]] = None, **kwargs
) -> PolygonSelector:
"""
Add a :class:`.PolygonSelector`.
Parameters
----------
selection: list of (float, float), optional
initial polygon points; if not given, you draw the polygon by clicking
**kwargs
passed to :class:`.PolygonSelector`
Returns
-------
PolygonSelector
"""
bbox = self.world_object.get_world_bounding_box()
xdata = np.concatenate(self.data[:, :, 0])
xmin, xmax = np.nanmin(xdata), np.nanmax(xdata)
# y from the world bounding box so that the graphics' offsets, e.g. a stack's, are included
ymin, ymax = bbox[0, 1], bbox[1, 1]
yspan = ymax - ymin
limits = (xmin, xmax, ymin - yspan / 2, ymax + yspan / 2)
selector = PolygonSelector(selection, limits, parent=self, **kwargs)
self._plot_area.add_graphic(selector, center=False)
return selector
def _get_linear_selector_init_args(self, axis: str, padding: float):
bbox = self.world_object.get_world_bounding_box()
axis_index = {"x": 0, "y": 1}[axis]
orthogonal_index = 1 - axis_index
data = np.concatenate(self.data[:, :, axis_index])
vmin, vmax = np.nanmin(data), np.nanmax(data)
# the first quarter along `axis`
bounds = (vmin, vmin + (vmax - vmin) / 4)
limits = (vmin, vmax)
# size and center on the orthogonal axis, from the world bounding box so that the
# graphics' offsets, e.g. a stack's, are included
size = np.ptp(bbox[:, orthogonal_index]) * 1.5 + padding
center = bbox[:, orthogonal_index].mean()
return bounds, limits, size, center
[docs]
class LineCollection(PositionsCollection):
_child_type = LineGraphic
[docs]
class ScatterCollection(PositionsCollection):
_child_type = ScatterGraphic
[docs]
class ImageCollection(GraphicCollection):
_child_type = ImageGraphic
[docs]
class ImageGrid(ImageCollection):
def __init__(
self,
data,
*,
shape: tuple[int, int] = None,
separation: tuple[float, float] = (0.0, 0.0),
offsets: np.ndarray = None,
**kwargs,
):
"""
Lay out a collection of images in a grid.
If ``offsets`` is given it is used directly as the per-image position. Otherwise the images
are placed row-major into a grid of ``shape`` (rows, columns), each cell sized to the
largest image so the rows and columns line up, with ``separation`` world-space gaps between
them. By default there is separation space between the images.
Parameters
----------
data: list of array-like
one image per grid cell
shape: (int, int), optional
grid (n_rows, n_cols); defaults to a roughly square grid that fits all the images
separation: (float, float), default (0.0, 0.0)
world-space (row, column) gaps between the images
offsets: array-like, optional
explicit (x, y, z) offset per image; when given, ``shape`` and ``separation`` are ignored
**kwargs
passed to :class:`.ImageCollection`, e.g. ``cmap``, ``vmin``, ``vmax``
"""
super().__init__(data, **kwargs)
n = len(self)
if offsets is None:
if shape is None:
shape = calculate_figure_shape(n) # roughly square (rows, cols)
if np.prod(shape) < n:
raise ValueError(
f"grid shape {shape} has fewer cells than the {n} images"
)
rows, cols = np.divmod(np.arange(n), shape[1])
# cell size = the largest image, via the data accessor, so rows and columns line up
sizes = np.array(
[image.shape[:2] for image in self.data[:]]
) # (rows, cols) per image
cell_height, cell_width = sizes.max(axis=0)
row_sep, col_sep = separation
offsets = np.zeros((n, 3))
offsets[:, 0] = cols * (cell_width + col_sep) # x, left to right
offsets[:, 1] = -rows * (cell_height + row_sep) # y, top row first
self.offsets[:] = offsets
class GraphicStack:
"""
Mixin that stacks a collection's graphics along the axes in ``separation_axis``. Each graphic is
offset by its index times the data max plus the ``separation`` gap, so the graphics are evenly
spaced and do not overlap; pass per-graphic ``steps`` to space them individually. Set
``separation`` or ``separation_axis`` to (re)stack, e.g. after changing the data.
"""
def __init__(
self,
data,
*,
separation: tuple[float, float, float] = (0.0, 0.0, 0.0),
separation_axis: str = "y",
steps: np.ndarray = None,
**kwargs,
):
"""
Create a stack of graphics.
Parameters
----------
data: list of array-like
one entry per graphic; its length is the number of graphics in the stack
separation: (float, float, float), default (0.0, 0.0, 0.0)
(x, y, z) gap between successive graphics, added to the step along the corresponding
stacking axis
separation_axis: str, default "y"
axes to stack along, any combination of "x", "y", "z", e.g. "y", "xy", "xyz"
steps: [n_graphics, 3] array-like, optional
per-graphic step along each (x, y, z) axis, i.e. the max each graphic reaches. When
``None`` (default) a single max over all the data sets one uniform step. When given, each
graphic is offset by the cumulative step of the graphics before it, plus ``separation``.
**kwargs
passed to the collection, e.g. ``colors``, ``thickness``, ``sizes``
"""
super().__init__(data, **kwargs)
self._separation = self._check_separation(separation)
self._steps = self._check_steps(steps)
self.separation_axis = separation_axis # (re)stacks
def _check_separation(self, separation) -> np.ndarray:
separation = np.asarray(separation, dtype=float)
if separation.shape != (3,):
raise ValueError(
f"separation must be an (x, y, z) iterable, got shape {separation.shape}"
)
return separation
def _check_steps(self, steps) -> np.ndarray | None:
if steps is None:
return None
steps = np.asarray(steps, dtype=float)
if steps.shape != (len(self), 3):
raise ValueError(
f"steps must be a [n_graphics, 3] array, got shape {steps.shape} for "
f"{len(self)} graphics"
)
return steps
@property
def separation(self) -> np.ndarray:
"""get or set the (x, y, z) gap added to the step along the stacking axes"""
return self._separation
@separation.setter
def separation(self, value: tuple[float, float, float]):
self._separation = self._check_separation(value)
self._restack()
@property
def steps(self) -> np.ndarray | None:
"""get or set the per-graphic (x, y, z) steps used to space the stack, ``None`` to auto-determine"""
return self._steps
@steps.setter
def steps(self, value: np.ndarray | None):
self._steps = self._check_steps(value)
self._restack()
@property
def separation_axis(self) -> str:
"""get or set the axes to stack along, e.g. "y", "xy", "xyz\" """
return self._separation_axis
@separation_axis.setter
def separation_axis(self, value: str):
if not set(value).issubset("xyz"):
raise ValueError(
f"separation_axis must be a combination of 'x', 'y', 'z', got {value!r}"
)
self._separation_axis = value
self._restack()
def _restack(self):
axes = [{"x": 0, "y": 1, "z": 2}[axis] for axis in self._separation_axis]
offsets = np.zeros((len(self), 3))
if self._steps is None:
# one max over all the data gives the step to stack by along each stacking axis,
# reduce each graphic first so the whole dataset is never concatenated
step = np.max([view.max(axis=0) for view in self.data[:, :, axes]], axis=0)
offsets[:, axes] = np.arange(len(self))[:, np.newaxis] * (
step + self._separation[axes]
)
else:
# per-graphic steps: offset each graphic past the previous ones by their cumulative step
offsets[1:, axes] = np.cumsum(
self._steps[:-1, axes] + self._separation[axes], axis=0
)
self.offsets[:] = offsets
[docs]
class LineStack(GraphicStack, LineCollection):
pass
[docs]
class ScatterStack(GraphicStack, ScatterCollection):
pass
