Skip to content

Commit 5871e22

Browse files
committed
Merge pull request #3586 from mhvk/quantity-scale-to-scaler
Use converters in Quantity ufuncs instead of just scaling (closes #2496)
2 parents e47b028 + 44a166a commit 5871e22

5 files changed

Lines changed: 106 additions & 78 deletions

File tree

‎CHANGES.rst‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,9 @@ Bug Fixes
295295

296296
- ``astropy.units``
297297

298+
- Ensure equivalencies that do more than just scale a ``Quantity`` are
299+
properly handled also in ``ufunc`` evaluations. [#2496, #3586]
300+
298301
- ``astropy.utils``
299302

300303
- ``astropy.vo``

‎astropy/units/core.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -786,7 +786,7 @@ def _is_equivalent(self, other, equivalencies=[]):
786786

787787
return False
788788

789-
def _apply_equivalences(self, unit, other, equivalencies):
789+
def _apply_equivalencies(self, unit, other, equivalencies):
790790
"""
791791
Internal function (used from `_get_converter`) to apply
792792
equivalence pairs.
@@ -846,7 +846,7 @@ def _get_converter(self, other, equivalencies=[]):
846846
try:
847847
scale = self._to(other)
848848
except UnitsError:
849-
return self._apply_equivalences(
849+
return self._apply_equivalencies(
850850
self, other, self._normalize_equivalencies(equivalencies))
851851
return lambda val: scale * _condition_arg(val)
852852

‎astropy/units/quantity.py‎

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -292,23 +292,23 @@ def __array_prepare__(self, obj, context=None):
292292
# should be multiplied before being passed to the ufunc, as well as
293293
# the unit the output from the ufunc will have.
294294
if function in UFUNC_HELPERS:
295-
scales, result_unit = UFUNC_HELPERS[function](function, *units)
295+
converters, result_unit = UFUNC_HELPERS[function](function, *units)
296296
else:
297297
raise TypeError("Unknown ufunc {0}. Please raise issue on "
298298
"https://github.com/astropy/astropy"
299299
.format(function.__name__))
300300

301-
if any(scale == 0. for scale in scales):
301+
if any(converter is False for converter in converters):
302302
# for two-argument ufuncs with a quantity and a non-quantity,
303303
# the quantity normally needs to be dimensionless, *except*
304304
# if the non-quantity can have arbitrary unit, i.e., when it
305305
# is all zero, infinity or NaN. In that case, the non-quantity
306306
# can just have the unit of the quantity
307307
# (this allows, e.g., `q > 0.` independent of unit)
308-
maybe_arbitrary_arg = args[scales.index(0.)]
308+
maybe_arbitrary_arg = args[converters.index(False)]
309309
try:
310310
if _can_have_arbitrary_unit(maybe_arbitrary_arg):
311-
scales = [1., 1.]
311+
converters = [None, None]
312312
else:
313313
raise UnitsError("Can only apply '{0}' function to "
314314
"dimensionless quantities when other "
@@ -363,8 +363,10 @@ def __array_prepare__(self, obj, context=None):
363363
# decomposed, which involves being scaled by a float, but since
364364
# the array is an integer the output then gets converted to an int
365365
# and truncated.
366-
if(any(not np.can_cast(arg, obj.dtype) for arg in args) or
367-
np.any(np.array(scales, dtype=obj.dtype) != np.array(scales))):
366+
result_dtype = np.result_type(*(args + tuple(
367+
(float if converter and converter(1.) % 1. != 0. else int)
368+
for converter in converters)))
369+
if not np.can_cast(result_dtype, obj.dtype):
368370
raise TypeError("Arguments cannot be cast safely to inplace "
369371
"output with dtype={0}".format(self.dtype))
370372

@@ -382,7 +384,7 @@ def __array_prepare__(self, obj, context=None):
382384
# the issue is that we can't actually scale the inputs since that
383385
# would be changing the objects passed to the ufunc, which would not
384386
# be expected by the user.
385-
if any(scale != 1. for scale in scales):
387+
if any(converters):
386388

387389
# If self is both output and input (which happens for in-place
388390
# operations), input will get overwritten with junk. To avoid
@@ -409,7 +411,7 @@ def __array_prepare__(self, obj, context=None):
409411
result._contiguous = self.copy()
410412

411413
# ensure we remember the scales we need
412-
result._scales = scales
414+
result._converters = converters
413415

414416
# unit output will get (setting _unit could prematurely change input
415417
# if obj is self, which happens for in-place operations; see above)
@@ -432,10 +434,10 @@ def __array_wrap__(self, obj, context=None):
432434

433435
# We now need to re-calculate quantities for which the input
434436
# needed to be scaled.
435-
if hasattr(obj, '_scales'):
437+
if hasattr(obj, '_converters'):
436438

437-
scales = obj._scales
438-
del obj._scales
439+
converters = obj._converters
440+
del obj._converters
439441

440442
# For in-place operations, input will get overwritten with
441443
# junk. To avoid that, we hid it in a new object in
@@ -456,19 +458,18 @@ def __array_wrap__(self, obj, context=None):
456458

457459
# Set the inputs, rescaling as necessary
458460
inputs = []
459-
for arg, scale in zip(args, scales):
460-
if scale != 1.:
461-
inputs.append(arg.value * scale)
462-
else: # for scale==1, input is not necessarily a Quantity
461+
for arg, converter in zip(args, converters):
462+
if converter:
463+
inputs.append(converter(arg.value))
464+
else: # with no conversion, input can be non-Quantity.
463465
inputs.append(getattr(arg, 'value', arg))
464466

465467
# For output arrays that require scaling, we can reuse the
466468
# output array to perform the scaling in place, as long as the
467469
# array is not integral. Here, we set the obj_array to `None`
468470
# when it can not be used to store the scaled result.
469-
if(result_unit is not None and
470-
any(not np.can_cast(scaled_arg, obj_array.dtype)
471-
for scaled_arg in inputs)):
471+
if not (result_unit is None or
472+
np.can_cast(np.result_type(*inputs), obj_array.dtype)):
472473
obj_array = None
473474

474475
# Re-compute the output using the ufunc

0 commit comments

Comments
 (0)