8000 Merge pull request #2583 from mhvk/quantity-numpy-ufunc · astropy/astropy@5b18ebc · GitHub
[go: up one dir, main page]

Skip to content

Commit 5b18ebc

Browse files
authored
Merge pull request #2583 from mhvk/quantity-numpy-ufunc
Quantity.__array_ufunc__
2 parents 326ac86 + cfc3852 commit 5b18ebc

File tree

12 files changed

+750
-241
lines changed

12 files changed

+750
-241
lines changed

CHANGES.rst

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,11 @@ New Features
187187
- ``Quantity`` now supports the ``@`` operator for matrix multiplication that
188188
was introduced in Python 3.5, for all supported versions of numpy. [#6144]
189189

190+
- ``Quantity`` supports the new ``__array_ufunc__`` protocol introduced in
191+
numpy 1.13. As a result, operations that involve unit conversion will be
192+
sped up considerably (by up to a factor of two for costly operations such
193+
as trigonometric ones). [#2583]
194+
190195
- ``astropy.utils``
191196

192197
- Added a new ``dataurl_mirror`` configuration item in ``astropy.utils.data``

astropy/coordinates/angles.py

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -444,6 +444,18 @@ def _repr_latex_(self):
444444
formatter={'str_kind': lambda x: x})
445445

446446

447+
def _no_angle_subclass(obj):
448+
"""Return any Angle subclass objects as an Angle objects.
449+
450+
This is used to ensure that Latitute and Longitude change to Angle
451+
objects when they are used in calculations (such as lon/2.)
452+
"""
453+
if isinstance(obj, tuple):
454+
return tuple(_no_angle_subclass(_obj) for _obj in obj)
455+
456+
return obj.view(Angle) if isinstance(obj, Angle) else obj
457+
458+
447459
class Latitude(Angle):
448460
"""
449461
Latitude-like angle(s) which must be in the range -90 to +90 deg.
@@ -525,11 +537,11 @@ def __setitem__(self, item, value):
525537
# Any calculation should drop to Angle
526538
def __array_wrap__(self, obj, context=None):
527539
obj = super(Angle, self).__array_wrap__(obj, context=context)
540+
return _no_angle_subclass(obj)
528541

529-
if isinstance(obj, Angle):
530-
return obj.view(Angle)
531-
532-
return obj
542+
def __array_ufunc__(self, *args, **kwargs):
543+
results = super(Latitude, self).__array_ufunc__(*args, **kwargs)
544+
return _no_angle_subclass(results)
533545

534546

535547
class Longitude(Angle):
@@ -644,8 +656,8 @@ def __array_finalize__(self, obj):
644656
# Any calculation should drop to Angle
645657
def __array_wrap__(self, obj, context=None):
646658
obj = super(Angle, self).__array_wrap__(obj, context=context)
659+
return _no_angle_subclass(obj)
647660

648-
if isinstance(obj, Angle):
649-
return obj.view(Angle)
650-
651-
return obj
661+
def __array_ufunc__(self, *args, **kwargs):
662+
results = super(Longitude, self).__array_ufunc__(*args, **kwargs)
663+
return _no_angle_subclass(results)

astropy/modeling/functional_models.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2185,19 +2185,21 @@ def evaluate(cls, x, y, amplitude, x_0, y_0, radius):
21852185
r = np.sqrt((x - x_0) ** 2 + (y - y_0) ** 2) / (radius / cls._rz)
21862186

21872187
if isinstance(r, Quantity):
2188-
r = r.decompose().value
2188+
# scipy function cannot handle Quantity, so turn into array.
2189+
r = r.to_value(u.dimensionless_unscaled)
21892190

21902191
# Since r can be zero, we have to take care to treat that case
21912192
# separately so as not to raise a numpy warning
21922193
z = np.ones(r.shape)
21932194
rt = np.pi * r[r > 0]
21942195
z[r > 0] = (2.0 * cls._j1(rt) / rt) ** 2
2195-
z *= amplitude
21962196

21972197
if isinstance(amplitude, Quantity):
2198-
return Quantity(z, unit=amplitude.unit, copy=False)
2199-
else:
2200-
return z
2198+
# make z quantity too, otherwise in-place multiplication fails.
2199+
z = Quantity(z, u.dimensionless_unscaled, copy=False)
2200+
2201+
z *= amplitude
2202+
return z
22012203

22022204
@property
22032205
def input_units(self):

0 commit comments

Comments
 (0)
0