python-3.x - astropy：如何检查参数是否具有正确的单位

Question

在 Python 中astropy，我如何检查函数的参数不仅具有正确的单位，而且完全有单位？

我很熟悉is_equivalent()，所以要检查它M有质量单位，我可以说

assert M.unit.is_equivalent(u.g)

True如果，说，它返回M = 1e12 * u.Msun。但是如果M根本没有任何单位——例如，如果我不小心通过了M = 1e12——我会得到

AttributeError: 'Quantity' object has no '_normalize_equivalencies' member

一种方法是首先乘以Meg ，如果它确实(u.m/u.m)有一个单位，则它等于Unit(dimensionless)并且不会改变，同时它将一个类似的数字转换为然后可以检查而不会给出错误的数字。M1e12<Quantity 1.e+12>

不过，这似乎不是正确的方法。有没有更好的办法？

标签： python-3.xastropy

Answer 1

如评论中所述，如果您的目标是专门检查函数参数的单位，则可以使用quantity_input专门用于此目的的实用程序装饰器（并将根据需要引发ValueErroror TypeErroror or UnitsError，这通常是无效函数的更合适的例外论点比裸露的AssertionError）。

我特别喜欢使用函数注释的 Python 3 特定语法。这里有些例子：

>>> from astropy import units as u
>>> @u.quantity_input
... def my_function(M: u.g):
...     print(M)

正确单位：

>>> my_function(1e12*u.Msun)                                                
1000000000000.0 solMass
>>> my_function(1e12*u.g)                                                   
1000000000000.0 g

参数不正确的例子：

>>> my_function(1e12)
Traceback (most recent call last)
...
TypeError: Argument 'M' to function 'my_function' has no 'unit' attribute. You may want to pass in an astropy Quantity instead.

>>> my_function(1e12*u.dimensionless_unscaled)
Traceback (most recent call last)
...
UnitsError: Argument 'M' to function 'my_function' must be in units convertible to 'g'.

如果您有许多函数M作为参数必须具有质量单位，您还可以为其创建简写，例如：

>>> M_mass = u.quantity_input(M=u.g)
>>> @M_mass
... def my_function(M):
...     print(M)

相同的方法可用于不同单元中的多个参数的函数。