🌐 AI搜索 & 代理 主页
Skip to content

Commit 0f557d5

Browse files
trygvradtrygvrad
committed
expose multivariate plotting functionality to top level functions imshow, pcolor, pcolormesh, and Collection
1 parent d5c43e0 commit 0f557d5

17 files changed

+406
-79
lines changed

‎lib/matplotlib/axes/_axes.py‎

Lines changed: 39 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6060,6 +6060,7 @@ def imshow(self, X, cmap=None, norm=None, *, aspect=None,
60606060
- (M, N): an image with scalar data. The values are mapped to
60616061
colors using normalization and a colormap. See parameters *norm*,
60626062
*cmap*, *vmin*, *vmax*.
6063+
- (v, M, N): if coupled with a cmap that supports v scalars
60636064
- (M, N, 3): an image with RGB values (0-1 float or 0-255 int).
60646065
- (M, N, 4): an image with RGBA values (0-1 float or 0-255 int),
60656066
i.e. including transparency.
@@ -6069,15 +6070,16 @@ def imshow(self, X, cmap=None, norm=None, *, aspect=None,
60696070
60706071
Out-of-range RGB(A) values are clipped.
60716072
6072-
%(cmap_doc)s
6073+
6074+
%(multi_cmap_doc)s
60736075
60746076
This parameter is ignored if *X* is RGB(A).
60756077
6076-
%(norm_doc)s
6078+
%(multi_norm_doc)s
60776079
60786080
This parameter is ignored if *X* is RGB(A).
60796081
6080-
%(vmin_vmax_doc)s
6082+
%(multi_vmin_vmax_doc)s
60816083
60826084
This parameter is ignored if *X* is RGB(A).
60836085
@@ -6156,6 +6158,9 @@ def imshow(self, X, cmap=None, norm=None, *, aspect=None,
61566158
See :doc:`/gallery/images_contours_and_fields/image_antialiasing` for
61576159
a discussion of image antialiasing.
61586160
6161+
Only 'data' is available when using `~matplotlib.colors.BivarColormap`
6162+
or `~matplotlib.colors.MultivarColormap`
6163+
61596164
alpha : float or array-like, optional
61606165
The alpha blending value, between 0 (transparent) and 1 (opaque).
61616166
If *alpha* is an array, the alpha blending values are applied pixel
@@ -6261,6 +6266,7 @@ def imshow(self, X, cmap=None, norm=None, *, aspect=None,
62616266
if aspect is not None:
62626267
self.set_aspect(aspect)
62636268

6269+
X = mcolorizer._ensure_multivariate_data(X, im.norm.n_components)
62646270
im.set_data(X)
62656271
im.set_alpha(alpha)
62666272
if im.get_clip_path() is None:
@@ -6416,9 +6422,10 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
64166422
64176423
Parameters
64186424
----------
6419-
C : 2D array-like
6425+
C : 2D or 3D array-like
64206426
The color-mapped values. Color-mapping is controlled by *cmap*,
6421-
*norm*, *vmin*, and *vmax*.
6427+
*norm*, *vmin*, and *vmax*. 3D arrays are supported only if the
6428+
cmap supports v channels, where v is the size along the first axis.
64226429
64236430
X, Y : array-like, optional
64246431
The coordinates of the corners of quadrilaterals of a pcolormesh::
@@ -6465,11 +6472,11 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
64656472
See :doc:`/gallery/images_contours_and_fields/pcolormesh_grids`
64666473
for more description.
64676474
6468-
%(cmap_doc)s
6475+
%(multi_cmap_doc)s
64696476
6470-
%(norm_doc)s
6477+
%(multi_norm_doc)s
64716478
6472-
%(vmin_vmax_doc)s
6479+
%(multi_vmin_vmax_doc)s
64736480
64746481
%(colorizer_doc)s
64756482
@@ -6544,8 +6551,17 @@ def pcolor(self, *args, shading=None, alpha=None, norm=None, cmap=None,
65446551
if shading is None:
65456552
shading = mpl.rcParams['pcolor.shading']
65466553
shading = shading.lower()
6547-
X, Y, C, shading = self._pcolorargs('pcolor', *args, shading=shading,
6548-
kwargs=kwargs)
6554+
6555+
if colorizer is None:
6556+
cmap = mcolorizer._ensure_cmap(cmap, accept_multivariate=True)
6557+
C = mcolorizer._ensure_multivariate_data(args[-1], cmap.n_variates)
6558+
else:
6559+
C = mcolorizer._ensure_multivariate_data(args[-1],
6560+
colorizer.cmap.n_variates)
6561+
6562+
X, Y, C, shading = self._pcolorargs('pcolor', *args[:-1], C,
6563+
shading=shading, kwargs=kwargs)
6564+
65496565
linewidths = (0.25,)
65506566
if 'linewidth' in kwargs:
65516567
kwargs['linewidths'] = kwargs.pop('linewidth')
@@ -6620,6 +6636,7 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
66206636
- (M, N) or M*N: a mesh with scalar data. The values are mapped to
66216637
colors using normalization and a colormap. See parameters *norm*,
66226638
*cmap*, *vmin*, *vmax*.
6639+
- (v, M, N): if coupled with a cmap that supports v scalars
66236640
- (M, N, 3): an image with RGB values (0-1 float or 0-255 int).
66246641
- (M, N, 4): an image with RGBA values (0-1 float or 0-255 int),
66256642
i.e. including transparency.
@@ -6656,11 +6673,11 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
66566673
expanded as needed into the appropriate 2D arrays, making a
66576674
rectangular grid.
66586675
6659-
%(cmap_doc)s
6676+
%(multi_cmap_doc)s
66606677
6661-
%(norm_doc)s
6678+
%(multi_norm_doc)s
66626679
6663-
%(vmin_vmax_doc)s
6680+
%(multi_vmin_vmax_doc)s
66646681
66656682
%(colorizer_doc)s
66666683
@@ -6783,7 +6800,15 @@ def pcolormesh(self, *args, alpha=None, norm=None, cmap=None, vmin=None,
67836800
shading = mpl._val_or_rc(shading, 'pcolor.shading').lower()
67846801
kwargs.setdefault('edgecolors', 'none')
67856802

6786-
X, Y, C, shading = self._pcolorargs('pcolormesh', *args,
6803+
if colorizer is None:
6804+
cmap = mcolorizer._ensure_cmap(cmap, accept_multivariate=True)
6805+
C = mcolorizer._ensure_multivariate_data(args[-1], cmap.n_variates)
6806+
else:
6807+
C = mcolorizer._ensure_multivariate_data(args[-1],
6808+
colorizer.cmap.n_variates)
6809+
6810+
6811+
X, Y, C, shading = self._pcolorargs('pcolormesh', *args[:-1], C,
67876812
shading=shading, kwargs=kwargs)
67886813
coords = np.stack([X, Y], axis=-1)
67896814

‎lib/matplotlib/axes/_axes.pyi‎

Lines changed: 19 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,13 @@ from matplotlib.collections import (
1212
QuadMesh,
1313
)
1414
from matplotlib.colorizer import Colorizer
15-
from matplotlib.colors import Colormap, Normalize
15+
from matplotlib.colors import (
16+
Colormap,
17+
BivarColormap,
18+
MultivarColormap,
19+
Norm,
20+
Normalize,
21+
)
1622
from matplotlib.container import BarContainer, ErrorbarContainer, StemContainer
1723
from matplotlib.contour import ContourSet, QuadContourSet
1824
from matplotlib.image import AxesImage, PcolorImage
@@ -490,14 +496,14 @@ class Axes(_AxesBase):
490496
def imshow(
491497
self,
492498
X: ArrayLike | PIL.Image.Image,
493-
cmap: str | Colormap | None = ...,
494-
norm: str | Normalize | None = ...,
499+
cmap: str | Colormap | BivarColormap | MultivarColormap | None = ...,
500+
norm: str | Norm | None = ...,
495501
*,
496502
aspect: Literal["equal", "auto"] | float | None = ...,
497503
interpolation: str | None = ...,
498504
alpha: float | ArrayLike | None = ...,
499-
vmin: float | None = ...,
500-
vmax: float | None = ...,
505+
vmin: float | tuple[float] | None = ...,
506+
vmax: float | tuple[float] | None = ...,
501507
colorizer: Colorizer | None = ...,
502508
origin: Literal["upper", "lower"] | None = ...,
503509
extent: tuple[float, float, float, float] | None = ...,
@@ -514,10 +520,10 @@ class Axes(_AxesBase):
514520
*args: ArrayLike,
515521
shading: Literal["flat", "nearest", "auto"] | None = ...,
516522
alpha: float | None = ...,
517-
norm: str | Normalize | None = ...,
518-
cmap: str | Colormap | None = ...,
519-
vmin: float | None = ...,
520-
vmax: float | None = ...,
523+
norm: str | Norm | None = ...,
524+
cmap: str | Colormap | BivarColormap | MultivarColormap | None = ...,
525+
vmin: float | tuple[float] | None = ...,
526+
vmax: float | tuple[float] | None = ...,
521527
colorizer: Colorizer | None = ...,
522528
data=...,
523529
**kwargs
@@ -526,10 +532,10 @@ class Axes(_AxesBase):
526532
self,
527533
*args: ArrayLike,
528534
alpha: float | None = ...,
529-
norm: str | Normalize | None = ...,
530-
cmap: str | Colormap | None = ...,
531-
vmin: float | None = ...,
532-
vmax: float | None = ...,
535+
norm: str | Norm | None = ...,
536+
cmap: str | Colormap | BivarColormap | MultivarColormap | None = ...,
537+
vmin: float | tuple[float] | None = ...,
538+
vmax: float | tuple[float] | None = ...,
533539
colorizer: Colorizer | None = ...,
534540
shading: Literal["flat", "nearest", "gouraud", "auto"] | None = ...,
535541
antialiased: bool = ...,

‎lib/matplotlib/collections.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2352,6 +2352,8 @@ def set_array(self, A):
23522352
h, w = height, width
23532353
ok_shapes = [(h, w, 3), (h, w, 4), (h, w), (h * w,)]
23542354
if A is not None:
2355+
if hasattr(self, 'norm'):
2356+
A = mcolorizer._ensure_multivariate_data(A, self.norm.n_components)
23552357
shape = np.shape(A)
23562358
if shape not in ok_shapes:
23572359
raise ValueError(
@@ -2609,7 +2611,7 @@ def _get_unmasked_polys(self):
26092611
mask = (mask[0:-1, 0:-1] | mask[1:, 1:] | mask[0:-1, 1:] | mask[1:, 0:-1])
26102612
arr = self.get_array()
26112613
if arr is not None:
2612-
arr = np.ma.getmaskarray(arr)
2614+
arr = self._getmaskarray(arr)
26132615
if arr.ndim == 3:
26142616
# RGB(A) case
26152617
mask |= np.any(arr, axis=-1)

‎lib/matplotlib/collections.pyi‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,12 @@ from numpy.typing import ArrayLike, NDArray
77
from . import colorizer, transforms
88
from .backend_bases import MouseEvent
99
from .artist import Artist
10-
from .colors import Normalize, Colormap
10+
from .colors import (
11+
Colormap,
12+
BivarColormap,
13+
MultivarColormap,
14+
Norm,
15+
)
1116
from .lines import Line2D
1217
from .path import Path
1318
from .patches import Patch
@@ -29,8 +34,8 @@ class Collection(colorizer.ColorizingArtist):
2934
antialiaseds: bool | Sequence[bool] | None = ...,
3035
offsets: tuple[float, float] | Sequence[tuple[float, float]] | None = ...,
3136
offset_transform: transforms.Transform | None = ...,
32-
norm: Normalize | None = ...,
33-
cmap: Colormap | None = ...,
37+
norm: Norm | None = ...,
38+
cmap: Colormap | BivarColormap | MultivarColormap | None = ...,
3439
colorizer: colorizer.Colorizer | None = ...,
3540
pickradius: float = ...,
3641
hatch: str | None = ...,

‎lib/matplotlib/colorizer.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -600,6 +600,18 @@ def get_array(self):
600600
"""
601601
return self._A
602602

603+
def _getmaskarray(self, A):
604+
"""
605+
Similar to np.ma.getmaskarray but also handles the case where
606+
the data has multiple fields.
607+
608+
The return array always has the same shape as the input, and dtype bool
609+
"""
610+
mask = np.ma.getmaskarray(A)
611+
if isinstance(self.norm, colors.MultiNorm):
612+
mask = np.any(mask.view('bool').reshape((*A.shape, -1)), axis=-1)
613+
return mask
614+
603615
def changed(self):
604616
"""
605617
Call this whenever the mappable is changed to notify all the

‎lib/matplotlib/colorizer.pyi‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,13 @@ class Colorizer:
99
callbacks: cbook.CallbackRegistry
1010
def __init__(
1111
self,
12-
cmap: str | colors.Colormap | None = ...,
12+
cmap: str | colors.Colormap | colors.BivarColormap | colors.MultivarColormap | None = ...,
1313
norm: str | colors.Norm | None = ...,
1414
) -> None: ...
1515
@property
1616
def norm(self) -> colors.Norm: ...
1717
@norm.setter
18-
def norm(self, norm: colors.Norm | str | None) -> None: ...
18+
def norm(self, norm: colors.Norm | str | tuple[str] | None) -> None: ...
1919
def to_rgba(
2020
self,
2121
x: np.ndarray,
@@ -26,28 +26,28 @@ class Colorizer:
2626
def autoscale(self, A: ArrayLike) -> None: ...
2727
def autoscale_None(self, A: ArrayLike) -> None: ...
2828
@property
29-
def cmap(self) -> colors.Colormap: ...
29+
def cmap(self) -> colors.Colormap | colors.BivarColormap | colors.MultivarColormap: ...
3030
@cmap.setter
31-
def cmap(self, cmap: colors.Colormap | str | None) -> None: ...
31+
def cmap(self, cmap: colors.Colormap | colors.BivarColormap | colors.MultivarColormap | str | None) -> None: ...
3232
def get_clim(self) -> tuple[float, float]: ...
3333
def set_clim(self, vmin: float | tuple[float, float] | None = ..., vmax: float | None = ...) -> None: ...
3434
def changed(self) -> None: ...
3535
@property
36-
def vmin(self) -> float | None: ...
36+
def vmin(self) -> float | tuple[float] | None: ...
3737
@vmin.setter
38-
def vmin(self, value: float | None) -> None: ...
38+
def vmin(self, value: float | tuple[float] | None) -> None: ...
3939
@property
40-
def vmax(self) -> float | None: ...
40+
def vmax(self) -> float | tuple[float] | None: ...
4141
@vmax.setter
42-
def vmax(self, value: float | None) -> None: ...
42+
def vmax(self, value: float | tuple[float] | None) -> None: ...
4343
@property
44-
def clip(self) -> bool: ...
44+
def clip(self) -> bool | tuple[bool, ...]: ...
4545
@clip.setter
46-
def clip(self, value: bool) -> None: ...
46+
def clip(self, value: ArrayLike | bool) -> None: ...
4747

4848

4949
class _ColorizerInterface:
50-
cmap: colors.Colormap
50+
cmap: colors.Colormap | colors.BivarColormap | colors.MultivarColormap
5151
colorbar: colorbar.Colorbar | None
5252
callbacks: cbook.CallbackRegistry
5353
def to_rgba(
@@ -60,8 +60,8 @@ class _ColorizerInterface:
6060
def get_clim(self) -> tuple[float, float]: ...
6161
def set_clim(self, vmin: float | tuple[float, float] | None = ..., vmax: float | None = ...) -> None: ...
6262
def get_alpha(self) -> float | None: ...
63-
def get_cmap(self) -> colors.Colormap: ...
64-
def set_cmap(self, cmap: str | colors.Colormap) -> None: ...
63+
def get_cmap(self) -> colors.Colormap | colors.BivarColormap | colors.MultivarColormap: ...
64+
def set_cmap(self, cmap: str | colors.Colormap | colors.BivarColormap | colors.MultivarColormap) -> None: ...
6565
@property
6666
def norm(self) -> colors.Norm: ...
6767
@norm.setter
@@ -75,7 +75,7 @@ class _ScalarMappable(_ColorizerInterface):
7575
def __init__(
7676
self,
7777
norm: colors.Norm | None = ...,
78-
cmap: str | colors.Colormap | None = ...,
78+
cmap: str | colors.Colormap | colors.BivarColormap | colors.MultivarColormap | None = ...,
7979
*,
8080
colorizer: Colorizer | None = ...,
8181
**kwargs

0 commit comments

Comments
 (0)