openscm_units._unit_registry

openscm_units._unit_registry

Definition of our unit registry

See also docs/source/notebooks/design-principles.py

ScmUnitRegistry

Bases: UnitRegistry

Unit registry class.

Provides some convenience methods to add standard units and contexts.

Source code in src/openscm_units/_unit_registry.py

class ScmUnitRegistry(pint.UnitRegistry):
    """
    Unit registry class.

    Provides some convenience methods to add standard units and contexts.
    """

    _contexts_added = False

    def __init__(
        self,
        *args: Any,
        metric_conversions: pd.DataFrame | None = None,
        **kwargs: Any,
    ):
        """
        Initialise the unit registry

        Parameters
        ----------
        metric_conversions
            :obj:`pd.DataFrame` containing the metric conversions.
            ``metric_conversions`` must have an index named ``"Species"`` that
            contains the different species and columns which contain the
            conversion for different metrics (the name of the metrics is taken
            from the column names).If not supplied, the
            ``globalwarmingpotentials`` package is used.

        *args
            Passed to the ``__init__`` method of the super class

        **kwargs
            Passed to the ``__init__`` method of the super class
        """
        self._metric_conversions = metric_conversions
        # If we didn't call init here, we wouldn't need to rebuild the cache
        # below but that also feels like a bad pattern
        super().__init__(*args, **kwargs)

    def add_standards(self) -> None:
        """
        Add standard units.

        Has to be done separately because of pint's weird initialising.
        """
        self._add_gases(_STANDARD_GASES)

        self._add_gases({x: x for x in MIXTURES})

        self.define("yr = 1 * year")
        self.define("a = 1 * year = annum")
        self.define("h = hour")
        self.define("d = day")
        self.define("degreeC = degC")
        self.define("degreeF = degF")
        self.define("kt = 1000 * t")  # since kt is used for "knot" in the defaults
        self.define(
            "Tt = 1000000000000 * t"
        )  # since Tt is used for "tex" in the defaults

        self.define("ppm = [concentrations]")
        self.define("ppb = ppm / 1000")
        self.define("ppt = ppb / 1000")
        # Have to rebuild cache to get right units for ppm as it is defined in
        # pint
        self._build_cache()

    def enable_contexts(
        self,
        *names_or_contexts: str | pint.facets.context.objects.Context,
        **kwargs: Any,
    ) -> None:
        """
        Overload pint's :func:`enable_contexts`

        This ensures we only add contexts once (the first time they are used)
        to avoid (unnecessary) operations.

        Parameters
        ----------
        names_or_contexts
            Names of contexts or :obj:`pint.registry.UnitRegistry.Context`
            objects to enable

        kwargs
            Passed to :meth:`enable_contexts` of the parent class
        """
        if not self._contexts_added:
            self._add_contexts()
        self._contexts_added = True
        super().enable_contexts(*names_or_contexts, **kwargs)

    def _add_mass_emissions_joint_version(self, symbol: str) -> None:
        """
        Add a unit which is the combination of mass and emissions.

        This allows users to units like e.g. ``"tC"`` rather than requiring a
        space between the mass and the emissions i.e. ``"t C"``

        Parameters
        ----------
        symbol
            The unit to add a joint version for
        """
        self.define(f"g{symbol} = g * {symbol}")
        self.define(f"t{symbol} = t * {symbol}")

    def _add_gases(self, gases: dict[str, str | list[str]]) -> None:
        for symbol, value in gases.items():
            if isinstance(value, str):
                # symbol is base unit
                self.define(f"{symbol} = [{value}]")
                if value != symbol:
                    self.define(f"{value} = {symbol}")
            else:
                # symbol has conversion and aliases
                self.define(f"{symbol} = {value[0]}")
                for alias in value[1:]:
                    self.define(f"{alias} = {symbol}")

            self._add_mass_emissions_joint_version(symbol)

            # Add alias for upper case symbol:
            if symbol.upper() != symbol:
                self.define(f"{symbol.upper()} = {symbol}")
                self._add_mass_emissions_joint_version(symbol.upper())

    def _add_contexts(self) -> None:
        """
        Add contexts
        """
        _ch4_context = pint.Context("CH4_conversions")
        _ch4_context = self._add_transformations_to_context(
            _ch4_context,
            "[methane]",
            self.CH4,
            "[carbon]",
            self.C,
            12 / 16,
        )
        self.add_context(_ch4_context)

        _n2o_context = pint.Context("N2O_conversions")
        _n2o_context = self._add_transformations_to_context(
            _n2o_context,
            "[nitrous_oxide]",
            self.nitrous_oxide,
            "[nitrogen]",
            self.nitrogen,
            14 / 44,
        )
        self.add_context(_n2o_context)

        _nox_context = pint.Context("NOx_conversions")
        _nox_context = self._add_transformations_to_context(
            _nox_context,
            "[nitrogen]",
            self.nitrogen,
            "[NOx]",
            self.NOx,
            (14 + 2 * 16) / 14,
        )
        self.add_context(_nox_context)

        _nh3_context = pint.Context("NH3_conversions")
        _nh3_context = self._add_transformations_to_context(
            _nh3_context,
            "[nitrogen]",
            self.nitrogen,
            "[NH3]",
            self.NH3,
            (14 + 3) / 14,
        )
        self.add_context(_nh3_context)

        self._add_metric_conversions()

    def _add_metric_conversions(self) -> None:
        """
        Add metric conversion contexts
        """
        if self._metric_conversions is None:
            metric_conversions = _load_globalwarmingpotentials_frame()
        else:
            metric_conversions = self._metric_conversions

        self._add_metric_conversions_from_df(metric_conversions)

    def _add_metric_conversions_from_df(self, metric_conversions: pd.DataFrame) -> None:
        # could make this public in future
        for col in metric_conversions:
            metric_conversion: pd.Series[float] = metric_conversions[col]
            transform_context = pint.Context(str(col))
            for label, val in metric_conversion.items():
                transform_context = self._add_gwp_to_context(
                    transform_context, str(label), val
                )

            for mixture in MIXTURES:
                constituents = self.split_gas_mixture(1 * self(mixture))
                try:
                    val = sum(
                        c.magnitude * metric_conversion[str(c.units)]
                        for c in constituents
                    )
                except KeyError:  # gwp not available for all constituents
                    continue
                if math.isnan(val):
                    continue
                transform_context = self._add_gwp_to_context(
                    transform_context, mixture, val
                )

            self.add_context(transform_context)

    def _add_gwp_to_context(
        self,
        transform_context: pint.facets.context.objects.Context,
        label: str,
        val: float,
    ) -> pint.facets.context.objects.Context:
        conv_val = (
            val
            * (self("CO2").to_base_units()).magnitude
            / (self(label).to_base_units()).magnitude
        )
        base_unit = next(
            iter(self._get_dimensionality(self(label).to_base_units()._units).keys())
        )

        base_unit_ureg = self(base_unit.replace("[", "").replace("]", ""))

        return self._add_transformations_to_context(
            transform_context,
            base_unit,
            base_unit_ureg,
            "[carbon]",
            self("carbon"),
            conv_val,
        )

    @staticmethod
    def _add_transformations_to_context(  # noqa: PLR0913
        context: pint.facets.context.objects.Context,
        base_unit: str,
        base_unit_ureg: pint.registry.UnitRegistry.Unit
        | pint.registry.UnitRegistry.Quantity,
        other_unit: str,
        other_unit_ureg: pint.registry.UnitRegistry.Unit
        | pint.registry.UnitRegistry.Quantity,
        conv_val: float,
    ) -> pint.facets.context.objects.Context:
        """
        Add all the transformations between units to a context for the two given units

        Transformations are mass x unit per time, mass x unit etc.
        """

        def _get_transform_func(
            forward: bool,
        ) -> pint.facets.context.objects.Transformation:
            if forward:

                def result_forward(
                    ureg: pint.registry.UnitRegistry,
                    value: pint.registry.UnitRegistry.Quantity,
                    **kwargs: Any,
                ) -> pint.registry.UnitRegistry.Quantity:
                    out: pint.registry.UnitRegistry.Quantity = (
                        value * other_unit_ureg / base_unit_ureg * conv_val
                    )

                    return out

                return result_forward  # type: ignore # cannot make pint behave

            def result_backward(
                ureg: pint.registry.UnitRegistry,
                value: pint.registry.UnitRegistry.Quantity,
                **kwargs: Any,
            ) -> pint.registry.UnitRegistry.Quantity:
                out: pint.registry.UnitRegistry.Quantity = (
                    value * (base_unit_ureg / other_unit_ureg) / conv_val
                )

                return out

            return result_backward  # type: ignore # cannot make pint behave

        formatters = [
            "{}",
            "[mass] * {} / [time]",
            "[mass] * {}",
            "{} / [time]",
        ]
        for fmt_str in formatters:
            context.add_transformation(
                fmt_str.format(base_unit),
                fmt_str.format(other_unit),
                _get_transform_func(forward=True),
            )
            context.add_transformation(
                fmt_str.format(other_unit),
                fmt_str.format(base_unit),
                _get_transform_func(forward=False),
            )

        return context

    def split_gas_mixture(
        self, quantity: pint.Quantity
    ) -> list[pint.registry.UnitRegistry.Quantity]:
        """
        Split a gas mixture into constituent gases.

        Parameters
        ----------
        quantity
            Pint quantity to split

        Returns
        -------
        :
            List of constituent gases
        """
        mixture_dimensions = [
            x for x in quantity.dimensionality.keys() if x[1:-1] in MIXTURES
        ]
        if not mixture_dimensions:
            raise ValueError("Dimensions don't contain a gas mixture.")  # noqa: TRY003
        if len(mixture_dimensions) > 1:
            raise NotImplementedError(
                "More than one gas mixture in dimensions is not supported."
            )
        mixture_dimension = mixture_dimensions[0]
        if quantity.dimensionality[mixture_dimension] != 1:
            raise NotImplementedError(
                "Mixture has dimensionality "
                f"{quantity.dimensionality[mixture_dimension]}"
                " != 1, which is not supported."
            )

        mixture = MIXTURES[mixture_dimension[1:-1]]
        mixture_unit = self(mixture_dimension[1:-1])

        ret = []
        for constituent, (fraction_pct, _, _) in mixture.items():
            constituent_unit = self(constituent)
            ret.append(quantity / mixture_unit * fraction_pct / 100 * constituent_unit)

        return ret

__init__(*args, metric_conversions=None, **kwargs)

Initialise the unit registry

Parameters:

Name	Type	Description	Default
metric_conversions	DataFrame | None	:obj:pd.DataFrame containing the metric conversions. metric_conversions must have an index named "Species" that contains the different species and columns which contain the conversion for different metrics (the name of the metrics is taken from the column names).If not supplied, the globalwarmingpotentials package is used.	None
*args	Any	Passed to the __init__ method of the super class	()
**kwargs	Any	Passed to the __init__ method of the super class	{}

Source code in src/openscm_units/_unit_registry.py

def __init__(
    self,
    *args: Any,
    metric_conversions: pd.DataFrame | None = None,
    **kwargs: Any,
):
    """
    Initialise the unit registry

    Parameters
    ----------
    metric_conversions
        :obj:`pd.DataFrame` containing the metric conversions.
        ``metric_conversions`` must have an index named ``"Species"`` that
        contains the different species and columns which contain the
        conversion for different metrics (the name of the metrics is taken
        from the column names).If not supplied, the
        ``globalwarmingpotentials`` package is used.

    *args
        Passed to the ``__init__`` method of the super class

    **kwargs
        Passed to the ``__init__`` method of the super class
    """
    self._metric_conversions = metric_conversions
    # If we didn't call init here, we wouldn't need to rebuild the cache
    # below but that also feels like a bad pattern
    super().__init__(*args, **kwargs)

add_standards()

Add standard units.

Has to be done separately because of pint's weird initialising.

Source code in src/openscm_units/_unit_registry.py

def add_standards(self) -> None:
    """
    Add standard units.

    Has to be done separately because of pint's weird initialising.
    """
    self._add_gases(_STANDARD_GASES)

    self._add_gases({x: x for x in MIXTURES})

    self.define("yr = 1 * year")
    self.define("a = 1 * year = annum")
    self.define("h = hour")
    self.define("d = day")
    self.define("degreeC = degC")
    self.define("degreeF = degF")
    self.define("kt = 1000 * t")  # since kt is used for "knot" in the defaults
    self.define(
        "Tt = 1000000000000 * t"
    )  # since Tt is used for "tex" in the defaults

    self.define("ppm = [concentrations]")
    self.define("ppb = ppm / 1000")
    self.define("ppt = ppb / 1000")
    # Have to rebuild cache to get right units for ppm as it is defined in
    # pint
    self._build_cache()

enable_contexts(*names_or_contexts, **kwargs)

Overload pint's :func:enable_contexts

This ensures we only add contexts once (the first time they are used) to avoid (unnecessary) operations.

Parameters:

Name	Type	Description	Default
names_or_contexts	str | Context	Names of contexts or :obj:pint.registry.UnitRegistry.Context objects to enable	()
kwargs	Any	Passed to :meth:enable_contexts of the parent class	{}

Source code in src/openscm_units/_unit_registry.py

def enable_contexts(
    self,
    *names_or_contexts: str | pint.facets.context.objects.Context,
    **kwargs: Any,
) -> None:
    """
    Overload pint's :func:`enable_contexts`

    This ensures we only add contexts once (the first time they are used)
    to avoid (unnecessary) operations.

    Parameters
    ----------
    names_or_contexts
        Names of contexts or :obj:`pint.registry.UnitRegistry.Context`
        objects to enable

    kwargs
        Passed to :meth:`enable_contexts` of the parent class
    """
    if not self._contexts_added:
        self._add_contexts()
    self._contexts_added = True
    super().enable_contexts(*names_or_contexts, **kwargs)

split_gas_mixture(quantity)

Split a gas mixture into constituent gases.

Parameters:

Name	Type	Description	Default
quantity	Quantity	Pint quantity to split	required

Returns:

Type	Description
list[Quantity]	List of constituent gases

Source code in src/openscm_units/_unit_registry.py

def split_gas_mixture(
    self, quantity: pint.Quantity
) -> list[pint.registry.UnitRegistry.Quantity]:
    """
    Split a gas mixture into constituent gases.

    Parameters
    ----------
    quantity
        Pint quantity to split

    Returns
    -------
    :
        List of constituent gases
    """
    mixture_dimensions = [
        x for x in quantity.dimensionality.keys() if x[1:-1] in MIXTURES
    ]
    if not mixture_dimensions:
        raise ValueError("Dimensions don't contain a gas mixture.")  # noqa: TRY003
    if len(mixture_dimensions) > 1:
        raise NotImplementedError(
            "More than one gas mixture in dimensions is not supported."
        )
    mixture_dimension = mixture_dimensions[0]
    if quantity.dimensionality[mixture_dimension] != 1:
        raise NotImplementedError(
            "Mixture has dimensionality "
            f"{quantity.dimensionality[mixture_dimension]}"
            " != 1, which is not supported."
        )

    mixture = MIXTURES[mixture_dimension[1:-1]]
    mixture_unit = self(mixture_dimension[1:-1])

    ret = []
    for constituent, (fraction_pct, _, _) in mixture.items():
        constituent_unit = self(constituent)
        ret.append(quantity / mixture_unit * fraction_pct / 100 * constituent_unit)

    return ret