Add set_data for ContourSet by larsoner · Pull Request #32271 · matplotlib/matplotlib · GitHub
Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions doc/release/next_whats_new/contourset_set_data.rst
29 changes: 29 additions & 0 deletions galleries/examples/animation/simple_anim.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,35 @@ def animate(i):

plt.show()

# %%
# Contours are updated the same way, with `.ContourSet.set_data`. Recontouring
# the existing artist is faster than removing the contour set and making a new
# one, and it keeps the contours in the same place in the draw order, which
# matters when blitting. The levels are not recomputed, so the colors mean
# the same thing in every frame.

fig, ax = plt.subplots()

X, Y = np.meshgrid(np.linspace(-3, 3, 100), np.linspace(-3, 3, 100))


def f(t):
return np.sin(X + t) * np.cos(Y - t)


cs = ax.contour(X, Y, f(0), levels=np.linspace(-0.9, 0.9, 7))


def animate_contour(i):
cs.set_data(X, Y, f(i / 25)) # update the data.
return cs,


ani_contour = animation.FuncAnimation(
fig, animate_contour, interval=20, blit=True, save_count=50)

plt.show()

# %%
#
# .. tags::
Expand Down
20 changes: 20 additions & 0 deletions lib/matplotlib/cbook.py
Original file line number Diff line number Diff line change
Expand Up @@ -2175,6 +2175,26 @@ def _setattr_cm(obj, **kwargs):
setattr(obj, attr, orig)


# TODO: Could be used in Line2D.set_data and other setters that mutate an artist
# more than once, and so can leave it half-updated when they raise.
@contextlib.contextmanager
def _safe_state_update(obj):
"""
Context manager to make an in-place update of *obj* all-or-nothing.

Yields a snapshot of ``obj.__dict__``, which is restored if the body raises, so
that an update mutating *obj* incrementally -- and able to fail partway through
-- leaves it as it was rather than half-updated.
"""
state = obj.__dict__.copy()
try:
yield state
except Exception:
obj.__dict__.clear()
obj.__dict__.update(state)
raise


class _OrderedSet(collections.abc.MutableSet):
def __init__(self):
self._od = collections.OrderedDict()
Expand Down
3 changes: 3 additions & 0 deletions lib/matplotlib/cbook.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -158,6 +158,9 @@ def _array_perimeter(arr: np.ndarray) -> np.ndarray: ...
def _unfold(arr: np.ndarray, axis: int, size: int, step: int) -> np.ndarray: ...
def _array_patch_perimeters(x: np.ndarray, rstride: int, cstride: int) -> np.ndarray: ...
def _setattr_cm(obj: Any, **kwargs) -> contextlib.AbstractContextManager[None]: ...
def _safe_state_update(
obj: Any,
) -> contextlib.AbstractContextManager[dict[str, Any]]: ...

class _OrderedSet(collections.abc.MutableSet):
def __init__(self) -> None: ...
Expand Down
74 changes: 71 additions & 3 deletions lib/matplotlib/contour.py
Original file line number Diff line number Diff line change
Expand Up @@ -891,6 +891,67 @@ def legend_elements(self, variable_name='x', str_format=str):

return artists, labels

def set_data(self, *args, **kwargs):
"""
Set new data and recompute the contours.

Call signatures::

set_data(Z)
set_data(X, Y, Z)

This reuses the existing artist, which is faster than removing the
contour set and creating a new one, and keeps its styling and its
place in the draw tree.

.. versionadded:: 3.12

Parameters
----------
*args
The new data, interpreted as by the function that created this
contour set, i.e. `~.Axes.contour`, `~.Axes.contourf`,
`~.Axes.tricontour` or `~.Axes.tricontourf`.

**kwargs
Only the keywords of that function that control *how the contours are
computed* are accepted, and they keep their current values if not
given: *corner_mask*, *algorithm* and the *xunits*/*yunits* unit
keywords for `~.Axes.contour`, or *triangles* and *mask* for
`~.Axes.tricontour`. Keywords that control how the contours *look*
(*levels*, *colors*, *cmap*, *linewidths*, ...) raise `TypeError`; set
those with the corresponding `.Collection` setters, or create a new
contour set.

Notes
-----
The *levels* are not recomputed; the new data is contoured at the
levels already in use, so that the colors and the colorbar stay
valid. Create a new contour set if you need different levels.

Existing contour labels are not moved. The Axes limits are not
rescaled, as for other artists' data setters.
"""
# _process_args() both validates and mutates: by the time it can fail it has
# already rebound the levels, zmin/zmax and the contour generator, so a
# rejected call would otherwise leave the contour set half-updated.
with cbook._safe_state_update(self) as state:
kwargs = self._process_args(*args, **kwargs)
if kwargs:
raise TypeError(
f"set_data() got unexpected keyword arguments {[*kwargs]}")
if not np.array_equal(state['levels'], self.levels):
# Only reachable through the ContourSet(ax, levels, allsegs)
# signature. Keeping the levels fixed is what lets us skip
# reprocessing the layers and the colors here, so enforce it rather
# than silently mismatching them.
raise ValueError("set_data() cannot change the contour levels")
if self._paths is state['_paths']: # Not set by _process_args.
self._paths = self._make_paths_from_contour_generator()
self.sticky_edges.x[:] = [self._mins[0], self._maxs[0]]
self.sticky_edges.y[:] = [self._mins[1], self._maxs[1]]
self.stale = True

def _process_args(self, *args, **kwargs):
"""
Process *args* and *kwargs*; override in derived classes.
Expand Down Expand Up @@ -937,8 +998,6 @@ def _process_args(self, *args, **kwargs):

def _make_paths_from_contour_generator(self):
"""Compute ``paths`` using C extension."""
if self._paths is not None:
return self._paths
cg = self._contour_generator
empty_path = Path(np.empty((0, 2)))
vertices_and_codes = (
Expand Down Expand Up @@ -1315,6 +1374,11 @@ class QuadContourSet(ContourSet):
%(contour_set_attributes)s
"""

# Set on the first _process_args; a later call from set_data keeps these rather
# than falling back to the rcParams.
_algorithm = None
_corner_mask = None

def _process_args(self, *args, corner_mask=None, algorithm=None, **kwargs):
"""
Process args and kwargs.
Expand All @@ -1332,10 +1396,13 @@ def _process_args(self, *args, corner_mask=None, algorithm=None, **kwargs):
else:
import contourpy

algorithm = mpl._val_or_rc(algorithm, 'contour.algorithm')
algorithm = mpl._val_or_rc(
algorithm or self._algorithm, 'contour.algorithm')
mpl.rcParams.validate["contour.algorithm"](algorithm)
self._algorithm = algorithm

if corner_mask is None:
corner_mask = self._corner_mask
if corner_mask is None:
if self._algorithm == "mpl2005":
# mpl2005 does not support corner_mask=True so if not
Expand Down Expand Up @@ -1530,6 +1597,7 @@ def _initialize_x_y(self, z):
Returns
-------
`~.contour.QuadContourSet`
Use `.ContourSet.set_data` to recontour new data with the same artist.

Other Parameters
----------------
Expand Down
1 change: 1 addition & 0 deletions lib/matplotlib/contour.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,7 @@ class ContourSet(ContourLabeler, Collection):
def legend_elements(
self, variable_name: str = ..., str_format: Callable[[float], str] = ...
) -> tuple[list[Artist], list[str]]: ...
def set_data(self, *args, **kwargs) -> None: ...
def find_nearest_contour(
self, x: float, y: float, indices: Iterable[int] | None = ..., pixel: bool = ...
) -> tuple[int, int, int, float, float, float]: ...
Expand Down
101 changes: 100 additions & 1 deletion lib/matplotlib/tests/test_contour.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@

import contourpy
import numpy as np
from numpy.testing import assert_array_almost_equal, assert_array_almost_equal_nulp
from numpy.testing import (
assert_array_almost_equal, assert_array_almost_equal_nulp,
assert_array_equal)
import matplotlib as mpl
from matplotlib import pyplot as plt, rc_context, ticker
from matplotlib.colors import LogNorm, same_color
Expand Down Expand Up @@ -896,3 +898,100 @@ def test_clabel_manual_subset():
cs = ax.contour([[1, 2], [3, 4]], levels=[1.5, 2.5, 3.5])
# Attempt to label only one specific level manually
ax.clabel(cs, levels=[2.5], manual=[(0.5, 0.5)])


@pytest.mark.parametrize("levels", [None, [-0.5, 0, 0.5]])
@pytest.mark.parametrize("plotter", ["contour", "contourf"])
@check_figures_equal()
def test_contour_set_data(fig_test, fig_ref, plotter, levels):
x, y = np.meshgrid(np.linspace(-3, 3, 20), np.linspace(-3, 3, 20))
z1 = np.sin(x) * np.cos(y)
z2 = np.cos(x) * np.sin(y)

cs = getattr(fig_test.subplots(), plotter)(x, y, z1, levels=levels)
orig_levels = cs.levels.copy()
cs.set_data(x, y, z2)
# The levels are kept, also where they were autoscaled from the old data.
assert_array_equal(cs.levels, orig_levels)
getattr(fig_ref.subplots(), plotter)(x, y, z2, levels=orig_levels)


@check_figures_equal()
def test_contour_set_data_z_only(fig_test, fig_ref):
z1 = np.arange(25).reshape((5, 5)) % 7
z2 = np.arange(25).reshape((5, 5)) % 5
fig_test.subplots().contour(z1, levels=[1, 2, 3]).set_data(z2)
fig_ref.subplots().contour(z2, levels=[1, 2, 3])


def test_contour_set_data_updates_artist_state():
x, y = np.meshgrid(np.linspace(0, 1, 5), np.linspace(0, 1, 5))
fig = plt.figure()
cs = fig.add_subplot().contour(x, y, x + y, levels=[0.5, 1, 1.5])
fig.canvas.draw()
assert not cs.stale
cs.set_data(2 * x - 1, 2 * y - 1, x + y)
assert cs.stale
assert cs.sticky_edges.x == [-1, 1]
assert cs.sticky_edges.y == [-1, 1]


def test_contour_set_data_changing_levels():
# The (levels, allsegs) signature is the only way set_data could change the
# levels, which would leave the colors and the colorbar mismatched.
segs = [[np.array([[0.0, 0.0], [1.0, 1.0]])]]
cs = mpl.contour.ContourSet(plt.figure().add_subplot(), [0.5], segs)
cs.set_data([0.5], segs) # Same levels, new segments: fine.
with pytest.raises(ValueError, match="cannot change the contour levels"):
cs.set_data([0.25, 0.75], segs * 2)
assert_array_equal(cs.levels, [0.5]) # Rolled back, not half-applied.


def test_contour_set_data_keeps_algorithm_and_corner_mask():
z = np.array([[1.0, 2.0], [3.0, 4.0]])
cs = plt.figure().add_subplot().contour(
z, algorithm='mpl2005', corner_mask=False)
cs.set_data(z * 2)
assert cs._algorithm == 'mpl2005'
assert cs._corner_mask is False
# They describe how the contours are computed, so they can also be re-set.
cs.set_data(z * 2, algorithm='serial', corner_mask=True)
assert cs._algorithm == 'serial'
assert cs._corner_mask is True


def _assert_set_data_rejects(exc, match, args=(), **kwargs):
# A rejected set_data() must raise and leave the contour set exactly as it was,
# rather than with the levels or the contour generator already swapped out.
z = np.arange(9).reshape((3, 3))
fig = plt.figure()
cs = fig.add_subplot().contour(z)
before = (cs.levels.copy(), [p.vertices.copy() for p in cs.get_paths()],
cs.zmin, cs.zmax, cs._contour_generator)
with pytest.raises(exc, match=match):
cs.set_data(*(args or (z,)), **kwargs)
assert_array_equal(cs.levels, before[0])
for path, vertices in zip(cs.get_paths(), before[1], strict=True):
assert_array_equal(path.vertices, vertices)
assert (cs.zmin, cs.zmax, cs._contour_generator) == before[2:]
fig.canvas.draw() # Must not raise.


@pytest.mark.parametrize("kwargs", [
{"levels": [1, 2]}, {"colors": "red"}, {"cmap": "plasma"},
{"linewidths": 2}, {"linestyles": "dashed"}, {"extend": "both"},
{"alpha": 0.5}, {"hatches": ["/"]}, {"zorder": 5},
])
def test_contour_set_data_rejects_style_kwargs(kwargs):
# Only the keywords that feed the contour generator are accepted; anything
# affecting the appearance would leave the levels and the colors mismatched.
_assert_set_data_rejects(TypeError, "unexpected keyword arguments", **kwargs)


@pytest.mark.parametrize("args", [
(np.arange(5), np.arange(5), np.empty((3, 4))), # mismatched shapes
(np.arange(7),), # z is not 2D
(np.empty((3, 3)),) * 5, # too many arguments
])
def test_contour_set_data_rejects_bad_data(args):
_assert_set_data_rejects((TypeError, ValueError), None, args)
28 changes: 28 additions & 0 deletions lib/matplotlib/tests/test_triangulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -1323,6 +1323,34 @@ def test_tricontourset_reuse():
assert tcs3._contour_generator == tcs1._contour_generator


@pytest.mark.parametrize("plotter", ["tricontour", "tricontourf"])
@check_figures_equal()
def test_tricontour_triangles_kwarg(fig_test, fig_ref, plotter):
# Passing triangles by keyword used to fall through to Collection.set().
x = [0.0, 1.0, 2.0, 0.0, 1.0, 0.0]
y = [0.0, 0.0, 0.0, 1.0, 1.0, 2.0]
z = [0.0, 1.0, 2.0, 1.0, 2.0, 3.0]
triangles = [[0, 1, 3], [1, 4, 3], [1, 2, 4], [3, 4, 5]]
levels = [0.5, 1.5, 2.5]
getattr(fig_test.subplots(), plotter)(x, y, z, triangles=triangles,
levels=levels)
getattr(fig_ref.subplots(), plotter)(x, y, triangles, z, levels=levels)


@check_figures_equal()
def test_tricontour_set_data(fig_test, fig_ref):
x = [0.0, 0.5, 1.0, 0.0, 0.5, 0.0]
y = [0.0, 0.0, 0.0, 0.5, 0.5, 1.0]
z1 = [0.0, 1.0, 2.0, 1.0, 2.0, 3.0]
z2 = [3.0, 2.0, 1.0, 2.0, 1.0, 0.0]
levels = [0.5, 1.5, 2.5]
triangles = [[0, 1, 3], [1, 4, 3], [1, 2, 4], [3, 4, 5]]
cs = fig_test.subplots().tricontour(x, y, z1, triangles=triangles,
levels=levels)
cs.set_data(x, y, z2, triangles=triangles)
fig_ref.subplots().tricontour(x, y, z2, triangles=triangles, levels=levels)


@check_figures_equal()
def test_triplot_with_ls(fig_test, fig_ref):
x = [0, 2, 1]
Expand Down
6 changes: 4 additions & 2 deletions lib/matplotlib/tri/_tricontour.py
Loading