Source code for message_ix_models.model.snapshot

"""Prepare base models from snapshot data."""
import logging
from pathlib import Path

import pandas as pd
from message_ix import Scenario
from message_ix.models import MACRO
from tqdm import tqdm

from message_ix_models import Spec
from message_ix_models.util import minimum_version
from message_ix_models.util.pooch import SOURCE, fetch

from .build import apply_spec
from .structure import get_codes

log = logging.getLogger(__name__)

[docs]def unpack(path: Path) -> Path: """Unpack :ref:`ixmp-format Excel file <ixmp:excel-data-format>` at `path`. The file is unpacked into a directory with the same name stem as the file (that is, without the :file:`.xlsx` suffix). In this directory are created: - One :file:`.csv.gz` file for each MESSAGE and/or MACRO parameter. - One file :file:`sets.xlsx` with only the :mod:`ixmp` sets, and no parameter data. If the files exist, they are not updated. To force re-unpacking, delete the files. Returns ------- Path Path to the directory containing the unpacked files. """ assert path.suffix == ".xlsx" base = path.with_suffix("") base.mkdir(exist_ok=True) # Get item name -> ixmp type mapping as a pd.Series xf = pd.ExcelFile(path, engine="openpyxl") name_type = xf.parse("ix_type_mapping") # Copied exactly from def parse_item_sheets(name): """Read data for item *name*, possibly across multiple sheets.""" dfs = [xf.parse(name)] # Collect data from repeated sheets due to max_row limit for x in filter(lambda n: n.startswith(name + "("), xf.sheet_names): dfs.append(xf.parse(x)) # pragma: no cover # Concatenate once and return return pd.concat(dfs, axis=0, ignore_index=True) sets_path = base.joinpath("sets.xlsx") sets_path.unlink(missing_ok=True) with pd.ExcelWriter(sets_path, engine="openpyxl") as ew: for _, (name, ix_type) in tqdm(name_type.iterrows()): item_path = base.joinpath(f"{name}.csv.gz") if item_path.exists(): continue df = parse_item_sheets(name) if ix_type == "set": df.to_excel(ew, sheet_name=name, index=False) else: # pragma: no cover df.to_csv(item_path, index=False) name_type.query("ix_type == 'set'").to_excel(ew, sheet_name="ix_type_mapping") return base
[docs]def read_excel(scenario: Scenario, path: Path) -> None: """Similar to :meth:`.Scenario.read_excel`, but using :func:`unpack`.""" base = unpack(path) scenario.read_excel(base.joinpath("sets.xlsx")) parameters = set(scenario.par_list()) with scenario.transact(f"Read snapshot data from {path}"): for p in base.glob("*.csv.gz"): name =".")[0] if name not in parameters: continue # Variable or equation data: don't read data = pd.read_csv(p) # Correct units if name == "inv_cost": data.replace({"unit": {"USD_2005/t ": "USD_2005/t"}}, inplace=True) scenario.add_par(name, data)
[docs]@minimum_version("message_ix 3.5") def load(scenario: Scenario, snapshot_id: int) -> None: """Fetch and load snapshot with ID `snapshot_id` into `scenario`. See also -------- SNAPSHOTS """ path, *_ = fetch(**SOURCE[f"snapshot-{snapshot_id}"]) # Add units spec = Spec() spec.add.set["unit"] = get_codes("unit/snapshot") apply_spec(scenario, spec) # Initialize MACRO items with scenario.transact("Prepare scenario for snapshot data"): MACRO.initialize(scenario) read_excel(scenario, path)