1
+ from matplotlib .cbook import iterable
1
2
import matplotlib .pyplot as plt
2
3
from matplotlib .testing .decorators import image_comparison
3
4
import matplotlib .units as munits
@@ -17,42 +18,65 @@ def __init__(self, data, units):
17
18
self .units = units
18
19
19
20
def to (self , new_units ):
20
- return Quantity (self .magnitude , new_units )
21
+ factors = {('hours' , 'seconds' ): 3600 , ('minutes' , 'hours' ): 1 / 60 ,
22
+ ('minutes' , 'seconds' ): 60 , ('feet' , 'miles' ): 1 / 5280. ,
23
+ ('feet' , 'inches' ): 12 , ('miles' , 'inches' ): 12 * 5280 }
24
+ if self .units != new_units :
25
+ mult = factors [self .units , new_units ]
26
+ return Quantity (mult * self .magnitude , new_units )
27
+ else :
28
+ return Quantity (self .magnitude , self .units )
21
29
22
30
def __getattr__ (self , attr ):
23
31
return getattr (self .magnitude , attr )
24
32
25
33
def __getitem__ (self , item ):
26
- return self .magnitude [item ]
34
+ return Quantity (self .magnitude [item ], self .units )
35
+
36
+ def __array__ (self ):
37
+ return np .asarray (self .magnitude )
27
38
28
39
29
40
# Tests that the conversion machinery works properly for classes that
30
41
# work as a facade over numpy arrays (like pint)
42
+ @image_comparison (baseline_images = ['plot_pint' ],
43
+ extensions = ['png' ], remove_text = False , style = 'mpl20' )
31
44
def test_numpy_facade ():
32
45
# Create an instance of the conversion interface and
33
46
# mock so we can check methods called
34
47
qc = munits .ConversionInterface ()
35
48
36
49
def convert (value , unit , axis ):
37
50
if hasattr (value , 'units' ):
38
- return value .to (unit )
51
+ return value .to (unit ).magnitude
52
+ elif iterable (value ):
53
+ try :
54
+ return [v .to (unit ).magnitude for v in value ]
55
+ except AttributeError :
56
+ return [Quantity (v , axis .get_units ()).to (unit ).magnitude
57
+ for v in value ]
39
58
else :
40
59
return Quantity (value , axis .get_units ()).to (unit ).magnitude
41
60
42
61
qc .convert = MagicMock (side_effect = convert )
43
- qc .axisinfo = MagicMock (return_value = None )
62
+ qc .axisinfo = MagicMock (side_effect = lambda u , a : munits . AxisInfo ( label = u ) )
44
63
qc .default_units = MagicMock (side_effect = lambda x , a : x .units )
45
64
46
65
# Register the class
47
66
munits .registry [Quantity ] = qc
48
67
49
68
# Simple test
50
- t = Quantity (np .linspace (0 , 10 ), 'sec ' )
51
- d = Quantity (30 * np .linspace (0 , 10 ), 'm/s ' )
69
+ y = Quantity (np .linspace (0 , 30 ), 'miles ' )
70
+ x = Quantity (np .linspace (0 , 5 ), 'hours ' )
52
71
53
- fig , ax = plt .subplots (1 , 1 )
54
- l , = plt .plot (t , d )
55
- ax .yaxis .set_units ('inch' )
72
+ fig , ax = plt .subplots ()
73
+ fig .subplots_adjust (left = 0.15 ) # Make space for label
74
+ ax .plot (x , y , 'tab:blue' )
75
+ ax .axhline (Quantity (26400 , 'feet' ), color = 'tab:red' )
76
+ ax .axvline (Quantity (120 , 'minutes' ), color = 'tab:green' )
77
+ ax .yaxis .set_units ('inches' )
78
+ ax .xaxis .set_units ('seconds' )
79
+ ax .autoscale_view ()
56
80
57
81
assert qc .convert .called
58
82
assert qc .axisinfo .called
0 commit comments