Source code for message_ix_models.model.water.utils

import logging
from collections import defaultdict
from functools import lru_cache
from itertools import product
from typing import Optional, Tuple

import numpy as np
import pandas as pd
import xarray as xr
from sdmx.model.v21 import Code

from message_ix_models import Context
from message_ix_models.model.structure import get_codes
from message_ix_models.util import load_package_data

log = logging.getLogger(__name__)

# Configuration files
    # Information about MESSAGE-water
    ("water", "config"),
    ("water", "set"),
    ("water", "technology"),

[docs]def read_config(context=None): """Read the water model configuration / metadata from file. Numerical values are converted to computation-ready data structures. Returns ------- .Context The current Context, with the loaded configuration. """ context = context or Context.get_instance(0) # if context.nexus_set == 'nexus': if "water set" in context: # Already loaded return context # Load water configuration for parts in METADATA: # Key for storing in the context key = " ".join(parts) # Actual filename parts; ends with YAML _parts = list(parts) _parts[-1] += ".yaml" context[key] = load_package_data(*_parts) return context
[docs]@lru_cache() def map_add_on(rtype=Code): """Map addon & type_addon in ``sets.yaml``.""" dims = ["add_on", "type_addon"] # Retrieve configuration context = read_config() # Assemble group information result = defaultdict(list) for indices in product(*[context["water set"][d]["add"] for d in dims]): # Create a new code by combining two result["code"].append( Code( id="".join( for c in indices), name=", ".join( for c in indices), ) ) # Tuple of the values along each dimension result["index"].append(tuple( for c in indices)) if rtype == "indexers": # Three tuples of members along each dimension indexers = zip(*result["index"]) indexers = { d: xr.DataArray(list(i), dims="consumer_group") for d, i in zip(dims, indexers) } indexers["consumer_group"] = xr.DataArray( [ for c in result["code"]], dims="consumer_group", ) return indexers elif rtype is Code: return sorted(result["code"], key=str) else: raise ValueError(rtype)
def add_commodity_and_level(df, default_level=None): # Add input commodity and level t_info = Context.get_instance()["water set"]["technology"]["add"] c_info = get_codes("commodity") @lru_cache() def t_cl(t): input = t_info[t_info.index(t)].anno["input"] # Commodity must be specified commodity = input["commodity"] # Use the default level for the commodity in the RES (per # commodity.yaml) level = ( input.get("level", "water_supply") or c_info[c_info.index(commodity)].anno.get("level", None) or default_level ) return commodity, level def func(row): row[["commodity", "level"]] = t_cl(row["technology"]) return row return df.apply(func, axis=1)
[docs]def map_yv_ya_lt( periods: Tuple[int, ...], lt: int, ya: Optional[int] = None ) -> pd.DataFrame: """All meaningful combinations of (vintage year, active year) given `periods`. Parameters ---------- labels : pandas.DataFrame Each column (dimension) corresponds to one in `df`. Each row represents one matched set of labels for those dimensions. lt : int, lifetime """ if not ya: ya = periods[0]"First active year set as {ya!r}") if not lt: raise ValueError("Add a fixed lifetime parameter 'lt'") # The following lines are the same as # message_ix.tests.test_feature_vintage_and_active_years._generate_yv_ya # - Create a mesh grid using numpy built-ins # - Take the upper-triangular portion (setting the rest to 0) # - Reshape data = np.triu(np.meshgrid(periods, periods, indexing="ij")).reshape((2, -1)) # Filter only non-zero pairs df = pd.DataFrame( filter(sum, zip(data[0, :], data[1, :])), columns=["year_vtg", "year_act"], dtype=np.int64, ) # Select values using the `ya` and `lt` parameters return df[(ya <= df.year_act) & (df.year_act - df.year_vtg <= lt)]