Source code for message_ix_models.util._convert_units

from typing import Mapping, Optional, Tuple
from warnings import catch_warnings, filterwarnings

import pandas as pd
from iam_units import registry

[docs]def series_of_pint_quantity(*args, **kwargs) -> pd.Series: """Suppress a spurious warning. Creating a :class:`pandas.Series` with a list of :class:`pint.Quantity` triggers a warning “The unit of the quantity is stripped when downcasting to ndarray,” even though the entire object is being stored and the unit is **not** stripped. This function suppresses this warning. """ with catch_warnings(): filterwarnings( "ignore", message="The unit of the quantity is stripped when downcasting to ndarray", module="pandas.core.dtypes.cast", ) return pd.Series(*args, **kwargs)
[docs]def convert_units( s: pd.Series, unit_info: Mapping[str, Tuple[float, str, Optional[str]]], store="magnitude", ) -> pd.Series: """Convert units of `s`, for use with :meth:`~pandas.DataFrame.apply`. ```` is used to retrieve a tuple of (`factor`, `input_unit`, `output_unit`) from `unit_info`. The (:class:`float`) values of `s` are converted to :class:`pint.Quantity` with the `input_unit` and factor; then cast to `output_unit`, if provided. Parameters ---------- s : pandas.Series unit_info : dict (str -> tuple) Mapping from quantity name (matched to ````) to 3-tuples of (`factor`, `input_unit`, `output_unit`). `output_unit` may be :obj:`None`. For example, see :data:`.ikarus.UNITS`. store : "magnitude" or "quantity" If "magnitude", the values of the returned series are the magnitudes of the results, with no output units. If "quantity", the values are scalar :class:`~pint.Quantity` objects. Returns ------- pandas.Series Same shape, index, and values as `s`, with output units. """ if store not in "magnitude quantity": raise ValueError(f"{store = }") # Retrieve the information from `unit_info` factor, unit_in, unit_out = unit_info[] # Default: `unit_out` is the same as `unit_in` unit_out = unit_out or unit_in # - Convert the values to a pint.Quantity(array) with the input units # - Convert to output units # - According to `store`, either extract just the magnitude, or store scalar # pint.Quantity objects. # - Reassemble into a series with index matching `s` result = registry.Quantity(factor * s.values, unit_in).to(unit_out) return series_of_pint_quantity( result.magnitude if store == "magnitude" else result.tolist(), index=s.index, dtype=(float if store == "magnitude" else object),, )