Units as method attributes

[IvarStefansson]

Variables are tagged with units when created in a model. We have been thinking about providing units for parameters (and possibly derived quantities). So far, we have partially implemented this in the documentation. I've been thinking of ways to improve this which would allow automatic parsing. Here's a suggestion using a decorator:

# Define utility function (in material_constants.py)
def add_units(units):  
    """Adds an si_units attribute to a function."""
    def si_unit_decorator(func):
        setattr(func, "si_units", units)
        return func
    return si_unit_decorator

# Illustration of how to apply using residual aperture as example.
# Build on the existing SolidConstants class to avoid redefinition.
class ImprovedSolidConstants(pp.SolidConstants):
    @add_units("m")  # We'd need to decorate each function like this.
    def residual_aperture(self):
        """Residual aperture [m].

        Returns:
            Residual aperture.

        """
        return self.convert_units(
            self.constants["residual_aperture"], 
            self.residual_aperture.si_units)  # This line could be just "m", as in current version

print(ImprovedSolidConstants.residual_aperture.si_units)  # Returns "m" as specified in decorator
my_constants = ImprovedSolidConstants()
my_constants.set_units(pp.Units())
print(my_constants.residual_aperture())  # Returns default value of 0.1

Caveat: Mypy might not be thrilled.

[Yuriyzabegaev]

I like the idea and can take part in implementing it. What do you think about automatically converting the units in the decorator if needed? That would allow us for getting rid of self.convert_units line, that would duplicate the information about the dimension for the user. However, I don't see how to implement it without having a global variable responsible for units.

[IvarStefansson]

Yes, I was thinking the same. I tried sketching an implementation, but couldn't immediately find an elegant solution. But I'm sure we can think of something if the others agree it's worth pursuing.

[keileg]

Interesting approach. I have the feeling you have been thinking deeper on this than me, @IvarStefansson, so let me know if you want to discuss at some point.

Also, it would be good to have a roadmap for where we want to go with the units in a wider context before we start an implementation.

[IvarStefansson]

Agreed. Yes, let's talk through this soon.

[Yuriyzabegaev]

Here is the minimal working example of how it can be implemented. If this looks good on the conceptual level to @keileg, @IvarStefansson and @IngridKJ, the rest of the work is just to insert it to all the parameters, which is mostly mechanical work.

from porepy import Units, MaterialConstants


def add_units(units: str):
    def decorator(func: callable):
        def inner(self: "PorepyModel", *args, **kwargs):
            result_si = func(self, *args, **kwargs)

            # Access the units from the model.
            # Assuming, this decorator is always applied to instance methods, so the
            # first argument is always "self"
            return self.constants.convert_units(value=result_si, units=units)

        return inner

    return decorator


class PorepyModel:

    def __init__(self, constants: MaterialConstants) -> None:
        self.constants = constants

    @add_units("m^2")
    def permeability(self):
        return 5e-14

    @add_units("kg*m^-3")
    def density(self):
        return 1e-3


units = Units(kg=1e9, m=1e-1)
constants = MaterialConstants({})
constants.set_units(units)
model = PorepyModel(constants)

print(model.permeability())
print(model.density())

[IvarStefansson]

Looks very good to me at first sight.

[Yuriyzabegaev]

Here is a sketch of how to track the derived quantities' units. I believe that's what @IvarStefansson originally wanted.

from porepy.models.units import Units
from porepy.models.material_constants import MaterialConstants
from typing import Callable


DERIVE_UNITS = True


class AdUnit:

    def __init__(self, units: str) -> None:
        self.units: str = units

    def __add__(self, other) -> "AdUnit":
        if isinstance(other, AdUnit):
            assert other.units == self.units
        return self

    def __mul__(self, other) -> "AdUnit":
        if not isinstance(other, AdUnit):
            return self

        return AdUnit(f"{self.units}*{other.units}")

    def __str__(self) -> str:
        return self.units


def add_units(units: str) -> Callable[[Callable], Callable]:
    def decorator(func: Callable) -> Callable:
        def inner(self: "PorepyModel", *args, **kwargs):
            # Currently, this is a global variable.
            # However, this can also be fetched from the model.
            if DERIVE_UNITS:
                return AdUnit(units)

            result_si = func(self, *args, **kwargs)

            # Access the units from the model.
            # Assuming, this decorator is always applied to instance methods, so the
            # first argument is always "self"
            return self.constants.convert_units(value=result_si, units=units)

        return inner

    return decorator


class PorepyModel:

    def __init__(self, constants: MaterialConstants) -> None:
        self.constants = constants

    @add_units("m^2")
    def permeability(self) -> float:
        return 5e-14

    @add_units("kg*m^-3")
    def density(self) -> float:
        return 1e-3

    def derived_quantity(self) -> float:
        return self.permeability() * self.density()


units = Units(kg=1e9, m=1e-1)
constants = MaterialConstants({})
constants.set_units(units)
model = PorepyModel(constants)

print(model.permeability())
print(model.density())
print(model.derived_quantity())

Which gives:

m^2
kg*m^-3
m^2*kg*m^-3

Currently, I used a global variable, however, something more elegant can be done, such as setattr(func, "si_unit", ...) with something like @add_units("derive") or involving a metaclass.