[AstroPy] Obscure question on subclassing units.Quantity

[E. Madison Bray]

On Sun, Sep 27, 2020 at 6:04 AM Livesey, Nathaniel J (US 3290)
<nathaniel.j.livesey at jpl.nasa.gov> wrote:
>
> Dear all,
>
>    I’ve subclassed astropy.Quantity (following, I believe, the various rules discussed for that in the astropy and numpy documentation).  Specifically, I’m implementing “dual algebra" (for computing values and Jacobians simultaneously, see the Wikipedia article on automatic differentiation if you’re curious).
>
>    It’s mostly working pretty well.  However, I’m having trouble getting it to invoke my __array_ufunc__ under certain circumstances, and I think it’s defaulting to the units.Quantity method.  Specifically, if “x” is an instance of my new “dual” class, then, saying:
>
> y = x * units.m
>
> (where units is astropy.units), my __array_ufunc__ is not invoked, thus I never get to my class’ multiply or rmultiply method. The consequence is that the Jacobian information attached to x doesn’t get handled correctly before being transferred to y.
>
>    If instead I do
>
> y = x * (1.0*units.m)
> or
> y = units.m * x
>
> then my __array_ufunc__ does get invoked, and everything works fine.  However, while this is arguably a workaround it’s unsatisfactory, and highly undesirable to require any “users” of my new class (mostly me I admit) to remember to do it that way. The consequences of forgetting are such that it’s hard to track down the place where the mistake was made.
> (I’ve not yet gotten round to overloading <<, but I still want the above to work.)
>
> Does anyone have a suggestion of how I can get my __array_ufunc__ to be invoked in this case?

Hello,

Taking a step back a bit, it would be good to confirm that
__array_ufunc__ is invoked when multiplying a Quantity by a Unit even
with the Quantity base class.  In fact it's not, so you can take some
relief in knowing the issue is not particular to your code, nor
subclassing Quantities in general.  To confirm, I just ran a bit of
code in PDB and set a breakpoint on Quantity.__array_ufunc__:

>>> from astropy import units as u
>>> import pdb
>>> q = u.Quantity([1, 2, 3])
>>> pdb.run("q * u.m")
> <string>(1)<module>()
(Pdb) s
--Call--
> astropy/units/quantity.py(934)__mul__()
-> def __mul__(self, other):
(Pdb) b Quantity.__array_ufunc__
Breakpoint 1 at astropy/units/quantity.py:434
(Pdb) c
>>> pdb.run("u.m * q")
> <string>(1)<module>()
(Pdb) s
--Call--
> astropy/units/core.py(683)__mul__()
-> def __mul__(self, m):
(Pdb) from .quantity import Quantity
(Pdb) b Quantity.__array_ufunc__
Breakpoint 2 at astropy/units/quantity.py:434
(Pdb) c
> astropy/units/quantity.py(457)__array_ufunc__()
-> converters, unit = converters_and_unit(function, method, *inputs)
(Pdb) where
  /usr/lib/python3.6/bdb.py(434)run()
-> exec(cmd, globals, locals)
  <string>(1)<module>()
  astropy/units/core.py(697)__mul__()
-> return Quantity(1, self) * m
  astropy/units/quantity.py(943)__mul__()
-> return super().__mul__(other)
> astropy/units/quantity.py(457)__array_ufunc__()
-> converters, unit = converters_and_unit(function, method, *inputs)


TL;DR, there happens to be a code path [1] in UnitBase.__mul__ (which
is called in the `u.m * q` case) where the LHS is converted to a
Quantity of magnitude 1 when multiplying by anything else that isn't a
unit, so Quantity.__array_ufunc__ is called in this case.  In the case
of `q * u.m` it just goes straight through Quantity.__mul__ [2] which
in effect just constructs a new Quantity (or appropriate subclass
thereof) and assigns it the unit from the RHS.  In this case there is
no need to call the __array_ufunc__ method.

This asymmetry could be considered an oversight, though a quantity
expert should comment on that.

In the meantime, I found a quick workaround for this *specific* case
by overriding __mul__ in a Quantity subclass:

>>> class Subquantity(u.Quantity):
...     def __array_ufunc__(self, function, method, *inputs, **kwargs):
...         print('my ufunc', function, method, inputs, kwargs)
...         return super().__array_ufunc__(function, method, *inputs, **kwargs)
...     def __mul__(self, other):
...         if isinstance(other, (u.UnitBase, str)):
...             return self * u.Quantity(1., other)
...         return super().__mul__(other)
...
>>> q = Subquantity([1, 2, 3])
>>> q * u.m
my ufunc <ufunc 'multiply'> __call__ (<Subquantity [1., 2., 3.]>,
<Quantity 1. m>) {}
<Subquantity [1., 2., 3.] m>

Hope that helps,

Madison


[1] https://github.com/astropy/astropy/blob/2f9deb3677e5ec0e1336c48030ad44d4ba20f1d4/astropy/units/core.py#L683
[2] https://github.com/astropy/astropy/blob/2f9deb3677e5ec0e1336c48030ad44d4ba20f1d4/astropy/units/quantity.py#L945