pint
pint copied to clipboard
Is there a way to wrap class methods based on a class member's unit?
Consider the following not-so-useful class:
from scipy.stats import norm
def Norm:
def __init__(self, mean, sigma):
self._func = norm(mean, sigma)
def pdf(self, x):
return self._func.pdf(x)
If I want to make this class unit-aware for a fixed unit, that's reasonable straightforward:
import pint
ureg = pint.get_application_registry()
class NormEnergy:
@ureg.wraps(None, (None, 'keV', 'keV'))
def __init__(self, mean, sigma):
self._func = norm(mean, sigma)
@ureg.wraps(None, (None, 'keV'))
def pdf(self, x):
return self._func.pdf(x)
@ureg.wraps('keV', (None, 'dimensionless'))
def ppf(self, q):
return self._func.ppf(q)
But what if I want this class to be unit-aware for an arbitrary but consistent unit? E.g., I can write the constructor with
@ureg.wraps(None, (None, '=A', '=A'))
def __init__(self, mean, sigma):
...
But the other methods then don't have access to the units that the constructor's wrapper identified. Is there any way to write something like
@ureg.wraps(None, (None, self.mean.units))
def pdf(self, x):
...
? The only way I can see to do this right now would be to use a factory function, which feels like overkill here.