8000 Merge pull request #1471 from astrofrog/contour-2d-shape-mismatch · matplotlib/matplotlib@8de33d0 · GitHub
[go: up one dir, main page]
Skip to content

Commit 8de33d0

Browse files
dmcdougalldmcdougall
committed
Merge pull request #1471 from astrofrog/contour-2d-shape-mismatch
Improved checking logic of _check_xyz in contour.py
2 parents 17e1cf8 + 75336c4 commit 8de33d0

File tree

3 files changed

+171
-9
lines changed

3 files changed

+171
-9
lines changed

‎lib/matplotlib/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1077,6 +1077,7 @@ def tk_window_focus():
10771077
'matplotlib.tests.test_cbook',
10781078
'matplotlib.tests.test_colorbar',
10791079
'matplotlib.tests.test_colors',
1080+
'matplotlib.tests.test_contour',
10801081
'matplotlib.tests.test_dates',
10811082
'matplotlib.tests.test_delaunay',
10821083
'matplotlib.tests.test_figure',

‎lib/matplotlib/contour.py

Lines changed: 31 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1424,20 +1424,42 @@ def _check_xyz(self, args, kwargs):
14241424
x = np.asarray(x, dtype=np.float64)
14251425
y = np.asarray(y, dtype=np.float64)
14261426
z = ma.asarray(args[2], dtype=np.float64)
1427+
14271428
if z.ndim != 2:
14281429
raise TypeError("Input z must be a 2D array.")
14291430
else:
14301431
Ny, Nx = z.shape
1431-
if x.shape == z.shape and y.shape == z.shape:
1432-
return x, y, z
1433-
if x.ndim != 1 or y.ndim != 1:
1432+
1433+
if x.ndim != y.ndim:
1434+
raise TypeError("Number of dimensions of x and y should match.")
1435+
1436+
if x.ndim == 1:
1437+
1438+
nx, = x.shape
1439+
ny, = y.shape
1440+
1441+
if nx != Nx:
1442+
raise TypeError("Length of x must be number of columns in z.")
1443+
1444+
if ny != Ny:
1445+
raise TypeError("Length of y must be number of rows in z.")
1446+
1447+
x, y = np.meshgrid(x, y)
1448+
1449+
elif x.ndim == 2:
1450+
1451+
if x.shape != z.shape:
1452+
raise TypeError("Shape of x does not match that of z: found "
1453+
"{0} instead of {1}.".format(x.shape, z.shape))
1454+
1455+
if y.shape != z.shape:
1456+
raise TypeError("Shape of y does not match that of z: found "
1457+
"{0} instead of {1}.".format(y.shape, z.shape))
1458+
1459+
else:
1460+
14341461
raise TypeError("Inputs x and y must be 1D or 2D.")
1435-
nx, = x.shape
1436-
ny, = y.shape
1437-
if nx != Nx or ny != Ny:
1438-
raise TypeError("Length of x must be number of columns in z,\n" +
1439-
"and length of y must be number of rows.")
1440-
x, y = np.meshgrid(x, y)
1462+
14411463
return x, y, z
14421464

14431465
def _initialize_x_y(self, z):

‎lib/matplotlib/tests/test_contour.py

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
import numpy as np
2+
3+
from matplotlib.testing.decorators import cleanup
4+
from matplotlib import pyplot as plt
5+
6+
7+
@cleanup
8+
def test_contour_shape_1d_valid():
9+
10+
x = np.arange(10)
11+
y = np.arange(9)
12+
z = np.random.random((9, 10))
13+
14+
fig = plt.figure()
15+
ax = fig.add_subplot(111)
16+
ax.contour(x, y, z)
17+
18+
19+
@cleanup
20+
def test_contour_shape_2d_valid():
21+
22+
x = np.arange(10)
23+
y = np.arange(9)
24+
xg, yg = np.meshgrid(x, y)
25+
z = np.random.random((9, 10))
26+
27+
fig = plt.figure()
28+
ax = fig.add_subplot(111)
29+
ax.contour(xg, yg, z)
30+
31+
32+
@cleanup
33+
def test_contour_shape_mismatch_1():
34+
35+
x = np.arange(9)
36+
y = np.arange(9)
37+
z = np.random.random((9, 10))
38+
39+
fig = plt.figure()
40+
ax = fig.add_subplot(111)
41+
42+
try:
43+
ax.contour(x, y, z)
44+
except TypeError as exc:
45+
assert exc.args[0] == 'Length of x must be number of columns in z.'
46+
47+
48+
@cleanup
49+
def test_contour_shape_mismatch_2():
50+
51+
x = np.arange(10)
52+
y = np.arange(10)
53+
z = np.random.random((9, 10))
54+
55+
fig = plt.figure()
56+
ax = fig.add_subplot(111)
57+
58+
try:
59+
ax.contour(x, y, z)
60+
except TypeError as exc:
61+
assert exc.args[0] == 'Length of y must be number of rows in z.'
62+
63+
64+
@cleanup
65+
def test_contour_shape_mismatch_3():
66+
67+
x = np.arange(10)
68+
y = np.arange(10)
69+
xg, yg = np.meshgrid(x, y)
70+
z = np.random.random((9, 10))
71+
72+
fig = plt.figure()
73+
ax = fig.add_subplot(111)
74+
75+
try:
76+
ax.contour(xg, y, z)
77+
except TypeError as exc:
78+
assert exc.args[0] == 'Number of dimensions of x and y should match.'
79+
80+
try:
81+
ax.contour(x, yg, z)
82+
except TypeError as exc:
83+
assert exc.args[0] == 'Number of dimensions of x and y should match.'
84+
85+
86+
@cleanup
87+
def test_contour_shape_mismatch_4():
88+
89+
g = np.random.random((9, 10))
90+
b = np.random.random((9, 9))
91+
z = np.random.random((9, 10))
92+
93+
fig = plt.figure()
94+
ax = fig.add_subplot(111)
95+
96+
try:
97+
ax.contour(b, g, z)
98+
except TypeError as exc:
99+
print exc.args[0]
100+
assert exc.args[0] == 'Shape of x does not match that of z: ' + \
101+
'found (9, 9) instead of (9, 10).'
102+
103+
try:
104+
ax.contour(g, b, z)
105+
except TypeError as exc:
106+
assert exc.args[0] == 'Shape of y does not match that of z: ' + \
107+
'found (9, 9) instead of (9, 10).'
108+
109+
110+
@cleanup
111+
def test_contour_shape_invalid_1():
112+
113+
x = np.random.random((3, 3, 3))
114+
y = np.random.random((3, 3, 3))
115+
z = np.random.random((9, 10))
116+
117+
fig = plt.figure()
118+
ax = fig.add_subplot(111)
119+
120+
try:
121+
ax.contour(x, y, z)
122+
except TypeError as exc:
123+
assert exc.args[0] == 'Inputs x and y must be 1D or 2D.'
124+
125+
126+
@cleanup
127+
def test_contour_shape_invalid_2():
128+
129+
x = np.random.random((3, 3, 3))
130+
y = np.random.random((3, 3, 3))
131+
z = np.random.random((3, 3, 3))
132+
133+
fig = plt.figure()
134+
ax = fig.add_subplot(111)
135+
136+
try:
137+
ax.contour(x, y, z)
138+
except TypeError as exc:
139+
assert exc.args[0] == 'Input z must be a 2D array.'

0 commit comments

Comments
 (0)
0