@@ -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