forked from matplotlib/matplotlib
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_units.py
More file actions
263 lines (211 loc) · 8.68 KB
/
test_units.py
File metadata and controls
263 lines (211 loc) · 8.68 KB
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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
from datetime import datetime, timezone, timedelta
import platform
from unittest.mock import MagicMock
import matplotlib.pyplot as plt
from matplotlib.testing.decorators import check_figures_equal, image_comparison
import matplotlib.units as munits
from matplotlib.category import UnitData
import numpy as np
import pytest
# Basic class that wraps numpy array and has units
class Quantity:
def __init__(self, data, units):
self.magnitude = data
self.units = units
def to(self, new_units):
factors = {('hours', 'seconds'): 3600, ('minutes', 'hours'): 1 / 60,
('minutes', 'seconds'): 60, ('feet', 'miles'): 1 / 5280.,
('feet', 'inches'): 12, ('miles', 'inches'): 12 * 5280}
if self.units != new_units:
mult = factors[self.units, new_units]
return Quantity(mult * self.magnitude, new_units)
else:
return Quantity(self.magnitude, self.units)
def __getattr__(self, attr):
return getattr(self.magnitude, attr)
def __getitem__(self, item):
if np.iterable(self.magnitude):
return Quantity(self.magnitude[item], self.units)
else:
return Quantity(self.magnitude, self.units)
def __array__(self):
return np.asarray(self.magnitude)
@pytest.fixture
def quantity_converter():
# Create an instance of the conversion interface and
# mock so we can check methods called
qc = munits.ConversionInterface()
def convert(value, unit, axis):
if hasattr(value, 'units'):
return value.to(unit).magnitude
elif np.iterable(value):
try:
return [v.to(unit).magnitude for v in value]
except AttributeError:
return [Quantity(v, axis.get_units()).to(unit).magnitude
for v in value]
else:
return Quantity(value, axis.get_units()).to(unit).magnitude
def default_units(value, axis):
if hasattr(value, 'units'):
return value.units
elif np.iterable(value):
for v in value:
if hasattr(v, 'units'):
return v.units
return None
qc.convert = MagicMock(side_effect=convert)
qc.axisinfo = MagicMock(side_effect=lambda u, a:
munits.AxisInfo(label=u, default_limits=(0, 100)))
qc.default_units = MagicMock(side_effect=default_units)
return qc
# Tests that the conversion machinery works properly for classes that
# work as a facade over numpy arrays (like pint)
@image_comparison(['plot_pint.png'], remove_text=False, style='mpl20',
tol=0 if platform.machine() == 'x86_64' else 0.01)
def test_numpy_facade(quantity_converter):
# use former defaults to match existing baseline image
plt.rcParams['axes.formatter.limits'] = -7, 7
# Register the class
munits.registry[Quantity] = quantity_converter
# Simple test
y = Quantity(np.linspace(0, 30), 'miles')
x = Quantity(np.linspace(0, 5), 'hours')
fig, ax = plt.subplots()
fig.subplots_adjust(left=0.15) # Make space for label
ax.plot(x, y, 'tab:blue')
ax.axhline(Quantity(26400, 'feet'), color='tab:red')
ax.axvline(Quantity(120, 'minutes'), color='tab:green')
ax.yaxis.set_units('inches')
ax.xaxis.set_units('seconds')
assert quantity_converter.convert.called
assert quantity_converter.axisinfo.called
assert quantity_converter.default_units.called
# Tests gh-8908
@image_comparison(['plot_masked_units.png'], remove_text=True, style='mpl20',
tol=0 if platform.machine() == 'x86_64' else 0.01)
def test_plot_masked_units():
data = np.linspace(-5, 5)
data_masked = np.ma.array(data, mask=(data > -2) & (data < 2))
data_masked_units = Quantity(data_masked, 'meters')
fig, ax = plt.subplots()
ax.plot(data_masked_units)
def test_empty_set_limits_with_units(quantity_converter):
# Register the class
munits.registry[Quantity] = quantity_converter
fig, ax = plt.subplots()
ax.set_xlim(Quantity(-1, 'meters'), Quantity(6, 'meters'))
ax.set_ylim(Quantity(-1, 'hours'), Quantity(16, 'hours'))
@image_comparison(['jpl_bar_units.png'],
savefig_kwarg={'dpi': 120}, style='mpl20')
def test_jpl_bar_units():
import matplotlib.testing.jpl_units as units
units.register()
day = units.Duration("ET", 24.0 * 60.0 * 60.0)
x = [0 * units.km, 1 * units.km, 2 * units.km]
w = [1 * day, 2 * day, 3 * day]
b = units.Epoch("ET", dt=datetime(2009, 4, 25))
fig, ax = plt.subplots()
ax.bar(x, w, bottom=b)
ax.set_ylim([b - 1 * day, b + w[-1] + (1.001) * day])
@image_comparison(['jpl_barh_units.png'],
savefig_kwarg={'dpi': 120}, style='mpl20')
def test_jpl_barh_units():
import matplotlib.testing.jpl_units as units
units.register()
day = units.Duration("ET", 24.0 * 60.0 * 60.0)
x = [0 * units.km, 1 * units.km, 2 * units.km]
w = [1 * day, 2 * day, 3 * day]
b = units.Epoch("ET", dt=datetime(2009, 4, 25))
fig, ax = plt.subplots()
ax.barh(x, w, left=b)
ax.set_xlim([b - 1 * day, b + w[-1] + (1.001) * day])
def test_empty_arrays():
# Check that plotting an empty array with a dtype works
plt.scatter(np.array([], dtype='datetime64[ns]'), np.array([]))
def test_scatter_element0_masked():
times = np.arange('2005-02', '2005-03', dtype='datetime64[D]')
y = np.arange(len(times), dtype=float)
y[0] = np.nan
fig, ax = plt.subplots()
ax.scatter(times, y)
fig.canvas.draw()
def test_errorbar_mixed_units():
x = np.arange(10)
y = [datetime(2020, 5, i * 2 + 1) for i in x]
fig, ax = plt.subplots()
ax.errorbar(x, y, timedelta(days=0.5))
fig.canvas.draw()
@check_figures_equal(extensions=["png"])
def test_subclass(fig_test, fig_ref):
class subdate(datetime):
pass
fig_test.subplots().plot(subdate(2000, 1, 1), 0, "o")
fig_ref.subplots().plot(datetime(2000, 1, 1), 0, "o")
def test_shared_axis_quantity(quantity_converter):
munits.registry[Quantity] = quantity_converter
x = Quantity(np.linspace(0, 1, 10), "hours")
y1 = Quantity(np.linspace(1, 2, 10), "feet")
y2 = Quantity(np.linspace(3, 4, 10), "feet")
fig, (ax1, ax2) = plt.subplots(2, 1, sharex='all', sharey='all')
ax1.plot(x, y1)
ax2.plot(x, y2)
assert ax1.xaxis.get_units() == ax2.xaxis.get_units() == "hours"
assert ax2.yaxis.get_units() == ax2.yaxis.get_units() == "feet"
ax1.xaxis.set_units("seconds")
ax2.yaxis.set_units("inches")
assert ax1.xaxis.get_units() == ax2.xaxis.get_units() == "seconds"
assert ax1.yaxis.get_units() == ax2.yaxis.get_units() == "inches"
def test_shared_axis_datetime():
# datetime uses dates.DateConverter
y1 = [datetime(2020, i, 1, tzinfo=timezone.utc) for i in range(1, 13)]
y2 = [datetime(2021, i, 1, tzinfo=timezone.utc) for i in range(1, 13)]
fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
ax1.plot(y1)
ax2.plot(y2)
ax1.yaxis.set_units(timezone(timedelta(hours=5)))
assert ax2.yaxis.units == timezone(timedelta(hours=5))
def test_shared_axis_categorical():
# str uses category.StrCategoryConverter
d1 = {"a": 1, "b": 2}
d2 = {"a": 3, "b": 4}
fig, (ax1, ax2) = plt.subplots(1, 2, sharex=True, sharey=True)
ax1.plot(d1.keys(), d1.values())
ax2.plot(d2.keys(), d2.values())
ax1.xaxis.set_units(UnitData(["c", "d"]))
assert "c" in ax2.xaxis.get_units()._mapping.keys()
def test_empty_default_limits(quantity_converter):
munits.registry[Quantity] = quantity_converter
fig, ax1 = plt.subplots()
ax1.xaxis.update_units(Quantity([10], "miles"))
fig.draw_without_rendering()
assert ax1.get_xlim() == (0, 100)
ax1.yaxis.update_units(Quantity([10], "miles"))
fig.draw_without_rendering()
assert ax1.get_ylim() == (0, 100)
fig, ax = plt.subplots()
ax.axhline(30)
ax.plot(Quantity(np.arange(0, 3), "miles"),
Quantity(np.arange(0, 6, 2), "feet"))
fig.draw_without_rendering()
assert ax.get_xlim() == (0, 2)
assert ax.get_ylim() == (0, 30)
fig, ax = plt.subplots()
ax.axvline(30)
ax.plot(Quantity(np.arange(0, 3), "miles"),
Quantity(np.arange(0, 6, 2), "feet"))
fig.draw_without_rendering()
assert ax.get_xlim() == (0, 30)
assert ax.get_ylim() == (0, 4)
fig, ax = plt.subplots()
ax.xaxis.update_units(Quantity([10], "miles"))
ax.axhline(30)
fig.draw_without_rendering()
assert ax.get_xlim() == (0, 100)
assert ax.get_ylim() == (28.5, 31.5)
fig, ax = plt.subplots()
ax.yaxis.update_units(Quantity([10], "miles"))
ax.axvline(30)
fig.draw_without_rendering()
assert ax.get_ylim() == (0, 100)
assert ax.get_xlim() == (28.5, 31.5)