Skip to content

Fix unit support with plot and pint #4803

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Sep 14, 2015
Merged

Fix unit support with plot and pint #4803

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0925e38
Add test for unit conversion.
dopplershift Sep 14, 2015
1d43e39
Remove some direct calls to atleast_1d.
dopplershift Jul 28, 2015
9b65233
Add test_units to whitelist.
dopplershift Sep 14, 2015
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
1 change: 1 addition & 0 deletions lib/matplotlib/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1465,6 +1465,7 @@ def tk_window_focus():
'matplotlib.tests.test_tightlayout',
'matplotlib.tests.test_transforms',
'matplotlib.tests.test_triangulation',
'matplotlib.tests.test_units',
'matplotlib.tests.test_widgets',
'matplotlib.tests.test_cycles',
'matplotlib.tests.test_labeled_data_unpacking',
Expand Down
13 changes: 8 additions & 5 deletions lib/matplotlib/axes/_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@
import matplotlib

from matplotlib import cbook
from matplotlib.cbook import _string_to_bool, iterable, index_of, get_label
from matplotlib.cbook import (_check_1d, _string_to_bool, iterable,
index_of, get_label)
from matplotlib import docstring
import matplotlib.colors as mcolors
import matplotlib.lines as mlines
Expand Down Expand Up @@ -214,8 +215,10 @@ def _xy_from_xy(self, x, y):
if by:
y = self.axes.convert_yunits(y)

x = np.atleast_1d(x) # like asanyarray, but converts scalar to array
y = np.atleast_1d(y)
# like asanyarray, but converts scalar to array, and doesn't change
# existing compatible sequences
x = _check_1d(x)
y = _check_1d(y)
if x.shape[0] != y.shape[0]:
raise ValueError("x and y must have same first dimension")
if x.ndim > 2 or y.ndim > 2:
Expand Down Expand Up @@ -353,8 +356,8 @@ def _plot_args(self, tup, kwargs):
kwargs['label'] = get_label(tup[-1], None)

if len(tup) == 2:
x = np.atleast_1d(tup[0])
y = np.atleast_1d(tup[-1])
x = _check_1d(tup[0])
y = _check_1d(tup[-1])
else:
x, y = index_of(tup[-1])

Expand Down
15 changes: 15 additions & 0 deletions lib/matplotlib/cbook.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2192,6 +2192,21 @@ def is_math_text(s):
return even_dollars


def _check_1d(x):
'''
Converts a sequence of less than 1 dimension, to an array of 1
dimension; leaves everything else untouched.
'''
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about instead using something like the following?

def asarray_units_okay(x):
    if type(x) in units.registry:
        return x
    else:
        return np.asarray(x)

def atleast_1d_units_okay(x);
    x = asarray_units_okay(x)
    if x.ndim == 0:
        x = x[np.newaxis]
    return x

I am nervous about preemptively checking all allowed operations.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This breaks anything using units that relies upon having atleast_1d called; this includes the "simple" unit framework in the examples.

Otherwise, the decomposition seems fine.

if not hasattr(x, 'shape') or len(x.shape) < 1:
return np.atleast_1d(x)
else:
try:
x[:, None]
return x
except (IndexError, TypeError):
return np.atleast_1d(x)


def _reshape_2D(X):
"""
Converts a non-empty list or an ndarray of two or fewer dimensions
Expand Down
57 changes: 57 additions & 0 deletions lib/matplotlib/tests/test_units.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import matplotlib.pyplot as plt
import matplotlib.units as munits
import numpy as np

try:
# mock in python 3.3+
from unittest.mock import MagicMock
except ImportError:
from mock import MagicMock


# Tests that the conversion machinery works properly for classes that
# work as a facade over numpy arrays (like pint)
def test_numpy_facade():
# Basic class that wraps numpy array and has units
class Quantity(object):
def __init__(self, data, units):
self.magnitude = data
self.units = units

def to(self, new_units):
return Quantity(self.magnitude, new_units)

def __getattr__(self, attr):
return getattr(self.magnitude, attr)

def __getitem__(self, item):
return self.magnitude[item]

# 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)
else:
return Quantity(value, axis.get_units()).to(unit).magnitude

qc.convert = MagicMock(side_effect=convert)
qc.axisinfo = MagicMock(return_value=None)
qc.default_units = MagicMock(side_effect=lambda x, a: x.units)

# Register the class
munits.registry[Quantity] = qc

# Simple test
t = Quantity(np.linspace(0, 10), 'sec')
d = Quantity(30 * np.linspace(0, 10), 'm/s')

fig, ax = plt.subplots(1, 1)
l, = plt.plot(t, d)
ax.yaxis.set_units('inch')

assert qc.convert.called
assert qc.axisinfo.called
assert qc.default_units.called