Skip to content

Commit

Permalink
Merge pull request #7010 from mhvk/quantity-speedups
Browse files Browse the repository at this point in the history 
Speed up basic quantity operations.
  • Loading branch information
@mhvk
mhvk committed Dec 21, 2017
2 parents c61f761 + f379050 commit 4efa2c6
Show file tree
Hide file tree
Showing 2 changed files with 91 additions and 54 deletions.
40 changes: 31 additions & 9 deletions astropy/units/quantity.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
# AstroPy
from .core import (Unit, dimensionless_unscaled, get_current_unit_registry,
UnitBase, UnitsError, UnitTypeError)
from .utils import is_effectively_unity
from .format.latex import Latex
from ..utils.compat import NUMPY_LT_1_13, NUMPY_LT_1_14
from ..utils.compat.misc import override__dir__
Expand Down Expand Up @@ -391,6 +392,10 @@ def __new__(cls, value, unit=None, dtype=None, copy=True, order=None,
return value.to(unit)

def __array_finalize__(self, obj):
# If we're a new object or viewing an ndarray, nothing has to be done.
if obj is None or obj.__class__ is np.ndarray:
return

# If our unit is not set and obj has a valid one, use it.
if self._unit is None:
unit = getattr(obj, '_unit', None)
Expand All @@ -399,9 +404,8 @@ def __array_finalize__(self, obj):

# Copy info if the original had `info` defined. Because of the way the
# DataInfo works, `'info' in obj.__dict__` is False until the
# `info` attribute is accessed or set. Note that `obj` can be an
# ndarray which doesn't have a `__dict__`.
if 'info' in getattr(obj, '__dict__', ()):
# `info` attribute is accessed or set.
if 'info' in obj.__dict__:
self.info = obj.info

def __array_prepare__(self, obj, context=None):
Expand Down Expand Up @@ -627,9 +631,9 @@ def __array_ufunc__(self, function, method, *inputs, **kwargs):
kwargs['out'] = (out_array,) if function.nout == 1 else out_array

# Same for inputs, but here also convert if necessary.
arrays = tuple((converter(input_.value) if converter else
getattr(input_, 'value', input_))
for input_, converter in zip(inputs, converters))
arrays = [(converter(input_.value) if converter else
getattr(input_, 'value', input_))
for input_, converter in zip(inputs, converters)]

# Call our superclass's __array_ufunc__
result = super().__array_ufunc__(function, method, *arrays, **kwargs)
Expand Down Expand Up @@ -734,6 +738,9 @@ def _new_view(self, obj=None, unit=None):
unit = self.unit
quantity_subclass = self.__class__
else:
# In principle, could gain time by testing unit is self.unit
# as well, and then quantity_subclass = self.__class__, but
# Constant relies on going through `__quantity_subclass__`.
unit = Unit(unit)
quantity_subclass = getattr(unit, '_quantity_class', Quantity)
if isinstance(self, quantity_subclass):
Expand All @@ -746,6 +753,8 @@ def _new_view(self, obj=None, unit=None):
# convert python and numpy scalars, which cannot be viewed as arrays
# and thus not as Quantity either, to zero-dimensional arrays.
# (These are turned back into scalar in `.value`)
# Note that for an ndarray input, the np.array call takes only double
# ``obj.__class is np.ndarray``. So, not worth special-casing.
if obj is None:
obj = self.view(np.ndarray)
else:
Expand Down Expand Up @@ -865,11 +874,24 @@ def to_value(self, unit=None, equivalencies=[]):
--------
to : Get a new instance in a different unit.
"""
value = self.view(np.ndarray)
if unit is not None:
if unit is None or unit is self.unit:
value = self.view(np.ndarray)
else:
unit = Unit(unit)
if unit != self.unit:
# We want a view if the unit does not change. One could check
# with "==", but that calculates the scale that we need anyway.
# TODO: would be better for `unit.to` to have an in-place flag.
try:
scale = self.unit._to(unit)
except Exception:
# Short-cut failed; try default (maybe equivalencies help).
value = self._to_value(unit, equivalencies)
else:
value = self.view(np.ndarray)
if not is_effectively_unity(scale):
# not in-place!
value = value * scale

return value if self.shape else value.item()

value = property(to_value,
Expand Down
105 changes: 60 additions & 45 deletions astropy/units/quantity_helper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -307,40 +307,52 @@ def helper_two_arg_dimensionless(f, unit1, unit2):
UFUNC_HELPERS[np.logaddexp2] = helper_two_arg_dimensionless


def get_converters_and_unit(f, *units):

def get_converters_and_unit(f, unit1, unit2):
converters = [None, None]
# no units for any input -- e.g., np.add(a1, a2, out=q)
if all(unit is None for unit in units):
return converters, dimensionless_unscaled
# By default, we try adjusting unit2 to unit1, so that the result will
# be unit1 as well. But if there is no second unit, we have to try
# adjusting unit1 (to dimensionless, see below).
if unit2 is None:
if unit1 is None:
# No units for any input -- e.g., np.add(a1, a2, out=q)
return converters, dimensionless_unscaled
changeable = 0
# swap units.
unit2 = unit1
unit1 = None
elif unit2 is unit1:
# ensure identical units is fast ("==" is slow, so avoid that).
return converters, unit1
else:
changeable = 1

fixed, changeable = (1, 0) if units[1] is None else (0, 1)
if units[fixed] is None:
# Try to get a converter from unit2 to unit1.
if unit1 is None:
try:
converters[changeable] = get_converter(units[changeable],
converters[changeable] = get_converter(unit2,
dimensionless_unscaled)
except UnitsError:
# special case: would be OK if unitless number is zero, inf, nan
converters[fixed] = False
return converters, units[changeable]
converters[1-changeable] = False
return converters, unit2
else:
return converters, dimensionless_unscaled

else:
try:
converters[changeable] = get_converter(units[changeable],
units[fixed])
converters[changeable] = get_converter(unit2, unit1)
except UnitsError:
raise UnitConversionError(
"Can only apply '{0}' function to quantities "
"with compatible dimensions"
.format(f.__name__))

return converters, units[fixed]
return converters, unit1


def helper_twoarg_invariant(f, unit1, unit2):
return get_converters_and_unit(f, unit1, unit2)
# This used to be a separate function that just called get_converters_and_unit.
# Using it directly saves a few us; keeping the clearer name.
helper_twoarg_invariant = get_converters_and_unit


UFUNC_HELPERS[np.add] = helper_twoarg_invariant
Expand Down Expand Up @@ -441,27 +453,30 @@ def converters_and_unit(function, method, *args):
UnitTypeError : when the conversion to the required (or consistent) units
is not possible.
"""
# Check whether we even support this ufunc
if function in UNSUPPORTED_UFUNCS:
raise TypeError("Cannot use function '{0}' with quantities"
.format(function.__name__))

# Check whether we support this ufunc, by getting the helper function
# (defined above) which returns a list of function(s) that convert the
# input(s) to the unit required for the ufunc, as well as the unit the
# result will have (a tuple of units if there are multiple outputs).
try:
ufunc_helper = UFUNC_HELPERS[function]
except KeyError:
if function in UNSUPPORTED_UFUNCS:
raise TypeError("Cannot use function '{0}' with quantities"
.format(function.__name__))
else:
raise TypeError("Unknown ufunc {0}. Please raise issue on "
"https://github.com/astropy/astropy"
.format(function.__name__))

if method == '__call__' or (method == 'outer' and function.nin == 2):
# Find out the units of the arguments passed to the ufunc; usually,
# at least one is a quantity, but for two-argument ufuncs, the second
# could also be a Numpy array, etc. These are given unit=None.
units = [getattr(arg, 'unit', None) for arg in args]

# If the ufunc is supported, then we call a helper function (defined
# above) which returns a list of function(s) that converts the input(s)
# to the unit required for the ufunc, as well as the unit the output
# will have (this is a tuple of units if there are multiple outputs).
if function in UFUNC_HELPERS:
converters, result_unit = UFUNC_HELPERS[function](function, *units)
else:
raise TypeError("Unknown ufunc {0}. Please raise issue on "
"https://github.com/astropy/astropy"
.format(function.__name__))
# Determine possible conversion functions, and the result unit.
converters, result_unit = ufunc_helper(function, *units)

if any(converter is False for converter in converters):
# for two-argument ufuncs with a quantity and a non-quantity,
Expand Down Expand Up @@ -516,30 +531,29 @@ def converters_and_unit(function, method, *args):
result_unit = dimensionless_unscaled

else: # methods for which the unit should stay the same
if method == 'at':
unit = getattr(args[0], 'unit', None)
units = [unit]
if function.nin == 2:
units.append(getattr(args[2], 'unit', None))
nin = function.nin
unit = getattr(args[0], 'unit', None)
if method == 'at' and nin <= 2:
if nin == 1:
units = [unit]
else:
units = [unit, getattr(args[2], 'unit', None)]

converters, result_unit = UFUNC_HELPERS[function](function, *units)
converters, result_unit = ufunc_helper(function, *units)

# ensure there is no 'converter' for indices (2nd argument)
converters.insert(1, None)

elif (method in ('reduce', 'accumulate', 'reduceat') and
function.nin == 2):
unit = getattr(args[0], 'unit', None)
converters, result_unit = UFUNC_HELPERS[function](function,
unit, unit)
elif method in {'reduce', 'accumulate', 'reduceat'} and nin == 2:
converters, result_unit = ufunc_helper(function, unit, unit)
converters = converters[:1]
if method == 'reduceat':
# add 'scale' for indices (2nd argument)
converters += [None]

else:
if method in ('reduce', 'accumulate', 'reduceat',
'outer') and function.nin != 2:
if method in {'reduce', 'accumulate',
'reduceat', 'outer'} and nin != 2:
raise ValueError("{0} only supported for binary functions"
.format(method))

Expand All @@ -554,9 +568,10 @@ def converters_and_unit(function, method, *args):
"Quantity instance as the result is not a "
"Quantity.".format(function.__name__, method))

if converters[0] is not None or (unit is not None and
(not result_unit.is_equivalent(unit) or
result_unit.to(unit) != 1.)):
if (converters[0] is not None or
(unit is not None and unit is not result_unit and
(not result_unit.is_equivalent(unit) or
result_unit.to(unit) != 1.))):
raise UnitsError("Cannot use '{1}' method on ufunc {0} with a "
"Quantity instance as it would change the unit."
.format(function.__name__, method))
Expand Down

0 comments on commit 4efa2c6

Please sign in to comment.