Skip to content

Commit 1c6fca3

Browse files
mhvkmhvk
committed
Add support for the clip ufunc (new in numpy 1.17).
1 parent 2d99bed commit 1c6fca3

4 files changed

Lines changed: 142 additions & 18 deletions

File tree

‎astropy/units/quantity_helper/converters.py‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -166,31 +166,34 @@ def converters_and_unit(function, method, *args):
166166
converters, result_unit = ufunc_helper(function, *units)
167167

168168
if any(converter is False for converter in converters):
169-
# for two-argument ufuncs with a quantity and a non-quantity,
169+
# for multi-argument ufuncs with a quantity and a non-quantity,
170170
# the quantity normally needs to be dimensionless, *except*
171171
# if the non-quantity can have arbitrary unit, i.e., when it
172172
# is all zero, infinity or NaN. In that case, the non-quantity
173173
# can just have the unit of the quantity
174174
# (this allows, e.g., `q > 0.` independent of unit)
175-
maybe_arbitrary_arg = args[converters.index(False)]
176175
try:
177-
if can_have_arbitrary_unit(maybe_arbitrary_arg):
178-
converters = [None, None]
179-
else:
180-
raise UnitConversionError(
181-
"Can only apply '{0}' function to "
182-
"dimensionless quantities when other "
183-
"argument is not a quantity (unless the "
184-
"latter is all zero/infinity/nan)"
185-
.format(function.__name__))
176+
# Don't fold this loop in the test above: this rare case
177+
# should not make the common case slower.
178+
for i, converter in enumerate(converters):
179+
if converter is not False:
180+
continue
181+
if can_have_arbitrary_unit(args[i]):
182+
converters[i] = None
183+
else:
184+
raise UnitConversionError(
185+
"Can only apply '{0}' function to "
186+
"dimensionless quantities when other "
187+
"argument is not a quantity (unless the "
188+
"latter is all zero/infinity/nan)"
189+
.format(function.__name__))
186190
except TypeError:
187191
# _can_have_arbitrary_unit failed: arg could not be compared
188192
# with zero or checked to be finite. Then, ufunc will fail too.
189193
raise TypeError("Unsupported operand type(s) for ufunc {0}: "
190-
"'{1}' and '{2}'"
191-
.format(function.__name__,
192-
args[0].__class__.__name__,
193-
args[1].__class__.__name__))
194+
"'{1}'".format(function.__name__,
195+
','.join([arg.__class__.__name__
196+
for arg in args])))
194197

195198
# In the case of np.power and np.float_power, the unit itself needs to
196199
# be modified by an amount that depends on one of the input values,

‎astropy/units/quantity_helper/helpers.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,39 @@ def helper_divmod(f, unit1, unit2):
296296
return converters, (dimensionless_unscaled, result_unit)
297297

298298

299+
def helper_clip(f, unit1, unit2, unit3):
300+
# Treat the array being clipped as primary.
301+
converters = [None]
302+
if unit1 is None:
303+
result_unit = dimensionless_unscaled
304+
try:
305+
converters += [(None if unit is None else
306+
get_converter(unit, dimensionless_unscaled))
307+
for unit in (unit2, unit3)]
308+
except UnitsError:
309+
raise UnitConversionError(
310+
"Can only apply '{0}' function to quantities with "
311+
"compatible dimensions".format(f.__name__))
312+
313+
else:
314+
result_unit = unit1
315+
for unit in unit2, unit3:
316+
try:
317+
converter = get_converter(_d(unit), result_unit)
318+
except UnitsError:
319+
if unit is None:
320+
# special case: OK if unitless number is zero, inf, nan
321+
converters.append(False)
322+
else:
323+
raise UnitConversionError(
324+
"Can only apply '{0}' function to quantities with "
325+
"compatible dimensions".format(f.__name__))
326+
else:
327+
converters.append(converter)
328+
329+
return converters, result_unit
330+
331+
299332
# list of ufuncs:
300333
# http://docs.scipy.org/doc/numpy/reference/ufuncs.html#available-ufuncs
301334

@@ -407,3 +440,6 @@ def helper_divmod(f, unit1, unit2):
407440
UFUNC_HELPERS[np.heaviside] = helper_heaviside
408441
UFUNC_HELPERS[np.float_power] = helper_power
409442
UFUNC_HELPERS[np.divmod] = helper_divmod
443+
# Check for clip ufunc; note that np.clip is a wrapper function, not the ufunc.
444+
if isinstance(getattr(np.core.umath, 'clip', None), np.ufunc):
445+
UFUNC_HELPERS[np.core.umath.clip] = helper_clip

‎astropy/units/tests/test_quantity.py‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -496,11 +496,10 @@ def test_incompatible_units(self):
496496

497497
def test_non_number_type(self):
498498
q1 = u.Quantity(11.412, unit=u.meter)
499-
type_err_msg = ("Unsupported operand type(s) for ufunc add: "
500-
"'Quantity' and 'dict'")
501499
with pytest.raises(TypeError) as exc:
502500
q1 + {'a': 1}
503-
assert exc.value.args[0] == type_err_msg
501+
assert exc.value.args[0].startswith(
502+
"Unsupported operand type(s) for ufunc add:")
504503

505504
with pytest.raises(TypeError):
506505
q1 + u.meter

‎astropy/units/tests/test_quantity_ufuncs.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -868,6 +868,92 @@ def test_ufunc_inplace_non_standard_dtype(self):
868868
a4 += u.Quantity(10, u.mm, dtype=np.int64)
869869

870870

871+
@pytest.mark.skipif(not hasattr(np.core.umath, 'clip'),
872+
reason='no clip ufunc available')
873+
class TestClip:
874+
"""Test the clip ufunc.
875+
876+
In numpy, this is hidden behind a function that does not backwards
877+
compatibility checks. We explicitly test the ufunc here.
878+
"""
879+
def setup(self):
880+
self.clip = np.core.umath.clip
881+
882+
def test_clip_simple(self):
883+
q = np.arange(-1., 10.) * u.m
884+
q_min = 125 * u.cm
885+
q_max = 0.0055 * u.km
886+
result = self.clip(q, q_min, q_max)
887+
assert result.unit == q.unit
888+
expected = self.clip(q.value, q_min.to_value(q.unit),
889+
q_max.to_value(q.unit)) * q.unit
890+
assert np.all(result == expected)
891+
892+
def test_clip_unitless_parts(self):
893+
q = np.arange(-1., 10.) * u.m
894+
qlim = 0.0055 * u.km
895+
# one-sided
896+
result1 = self.clip(q, -np.inf, qlim)
897+
expected1 = self.clip(q.value, -np.inf, qlim.to_value(q.unit)) * q.unit
898+
assert np.all(result1 == expected1)
899+
result2 = self.clip(q, qlim, np.inf)
900+
expected2 = self.clip(q.value, qlim.to_value(q.unit), np.inf) * q.unit
901+
assert np.all(result2 == expected2)
902+
# Zero
903+
result3 = self.clip(q, np.zeros(q.shape), qlim)
904+
expected3 = self.clip(q.value, 0, qlim.to_value(q.unit)) * q.unit
905+
assert np.all(result3 == expected3)
906+
# Two unitless parts, array-shaped.
907+
result4 = self.clip(q, np.zeros(q.shape), np.full(q.shape, np.inf))
908+
expected4 = self.clip(q.value, 0, np.inf) * q.unit
909+
assert np.all(result4 == expected4)
910+
911+
def test_clip_dimensionless(self):
912+
q = np.arange(-1., 10.) * u.dimensionless_unscaled
913+
result = self.clip(q, 200 * u.percent, 5.)
914+
expected = self.clip(q, 2., 5.)
915+
assert result.unit == u.dimensionless_unscaled
916+
assert np.all(result == expected)
917+
918+
def test_clip_ndarray(self):
919+
a = np.arange(-1., 10.)
920+
result = self.clip(a, 200 * u.percent, 5. * u.dimensionless_unscaled)
921+
assert isinstance(result, u.Quantity)
922+
expected = self.clip(a, 2., 5.) * u.dimensionless_unscaled
923+
assert np.all(result == expected)
924+
925+
def test_clip_quantity_inplace(self):
926+
q = np.arange(-1., 10.) * u.m
927+
q_min = 125 * u.cm
928+
q_max = 0.0055 * u.km
929+
expected = self.clip(q.value, q_min.to_value(q.unit),
930+
q_max.to_value(q.unit)) * q.unit
931+
result = self.clip(q, q_min, q_max, out=q)
932+
assert result is q
933+
assert np.all(result == expected)
934+
935+
def test_clip_ndarray_dimensionless_output(self):
936+
a = np.arange(-1., 10.)
937+
q = np.zeros_like(a) * u.m
938+
expected = self.clip(a, 2., 5.) * u.dimensionless_unscaled
939+
result = self.clip(a, 200 * u.percent, 5. * u.dimensionless_unscaled,
940+
out=q)
941+
assert result is q
942+
assert result.unit == u.dimensionless_unscaled
943+
assert np.all(result == expected)
944+
945+
def test_clip_errors(self):
946+
q = np.arange(-1., 10.) * u.m
947+
with pytest.raises(u.UnitsError):
948+
self.clip(q, 0, 1*u.s)
949+
with pytest.raises(u.UnitsError):
950+
self.clip(q.value, 0, 1*u.s)
951+
with pytest.raises(u.UnitsError):
952+
self.clip(q, -1, 0.)
953+
with pytest.raises(u.UnitsError):
954+
self.clip(q, 0., 1.)
955+
956+
871957
class TestUfuncAt:
872958
"""Test that 'at' method for ufuncs (calculates in-place at given indices)
873959

0 commit comments

Comments
 (0)