Source code for genno.compat.pyam.computations

import logging
from pathlib import Path
from typing import Callable, Collection, Optional, Union

import pyam

import genno.computations

from . import util

log = logging.getLogger(__name__)

__all__ = ["as_pyam", "concat", "write_report"]

[docs]def as_pyam( scenario, quantity, # /, # Requires Python 3.8; uncomment if/when support for Python 3.7 is dropped rename=dict(), collapse: Optional[Callable] = None, replace=dict(), drop: Union[Collection[str], str] = "auto", unit=None, ): """Return a :class:`pyam.IamDataFrame` containing `quantity`. Warnings are logged if the arguments result in additional, unhandled columns in the resulting data frame that are not part of the IAMC spec. Parameters ---------- scenario : Any objects with :attr:`model` and :attr:`scenario` attributes of type :class:`str`, e.g. :class:`ixmp.Scenario`. Raises ------ ValueError If the resulting data frame has duplicate values in the standard IAMC index columns. :class:`pyam.IamDataFrame` cannot handle this data. See also -------- .Computer.convert_pyam """ # - Convert to pd.DataFrame # - Rename one dimension to 'year' or 'time' # - Fill variable, unit, model, and scenario columns # - Replace values # - Apply the collapse callback, if given # - Drop any unwanted columns # - Clean units df = ( quantity.to_series() .rename("value") .reset_index() .assign(, unit=quantity.attrs.get("_unit", ""), # TODO accept these from separate strings model=scenario.model, scenario=scenario.scenario, ) .rename(columns=rename) .pipe(collapse or util.collapse) .replace(replace, regex=True) .pipe(util.drop, columns=drop) .pipe(util.clean_units, unit) ) # Raise exception for non-unique data duplicates = df.duplicated(subset=set(df.columns) - {"value"}) if duplicates.any(): raise ValueError( "Duplicate IAMC indices cannot be converted:\n" + str(df[duplicates].drop(columns=["model", "scenario"])) ) return pyam.IamDataFrame(df)
[docs]def concat(*args, **kwargs): """Concatenate *args*, which must all be :class:`pyam.IamDataFrame`.""" if isinstance(args[0], pyam.IamDataFrame): # pyam.concat() takes an iterable of args return pyam.concat(args, **kwargs) else: # genno.computations.concat() takes a variable number of positional arguments return genno.computations.concat(*args, **kwargs)
[docs]def write_report(obj, path: Union[str, Path]) -> None: """Write obj` to the file at `path`. If `obj` is a :class:`pyam.IamDataFrame` and `path` ends with ".csv" or ".xlsx", use :mod:`pyam` methods to write the file to CSV or Excel format, respectively. Otherwise, equivalent to :func:`genno.computations.write_report`. """ if not isinstance(obj, pyam.IamDataFrame): return genno.computations.write_report(obj, path) path = Path(path) if path.suffix == ".csv": obj.to_csv(path) elif path.suffix == ".xlsx": obj.to_excel(path) else: raise ValueError( f"pyam.IamDataFrame can be written to .csv or .xlsx, not {path.suffix}" )