pint icon indicating copy to clipboard operation
pint copied to clipboard

Is there a way to wrap class methods based on a class member's unit?

Open bloer opened this issue 3 years ago • 0 comments

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.

bloer avatar Dec 14 '21 21:12 bloer