flextool 4.0.0__py3-none-any.whl

flextool/scenario_comparison/dispatch_mappings.py

@@ -0,0 +1,205 @@
+"""Load and combine dispatch mapping parquet files across scenarios.
+
+Exports:
+    load_dispatch_mappings  — load all mapping files from one scenario folder
+    combine_dispatch_mappings — combine across scenarios → DispatchMappings
+    get_group_node_multiindex — build (group, node) MultiIndex for a scenario
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pandas as pd
+
+from flextool.lean_parquet import read_lean_parquet
+from flextool.scenario_comparison.data_models import DispatchMappings
+
+
+def load_dispatch_mappings(parquet_dir: Path) -> dict[str, pd.DataFrame | None]:
+    """
+    Load all dispatch-related mapping files from a single parquet directory.
+
+    Parameters:
+    -----------
+    parquet_dir : Path
+        Directory containing the parquet files for one scenario
+
+    Returns:
+    --------
+    dict : Dictionary containing all dispatch mapping dataframes (raw, unfiltered)
+    """
+    parquet_dir = Path(parquet_dir)
+    mappings: dict[str, pd.DataFrame | None] = {}
+
+    # File mapping: key -> (filename, required)
+    # node_inflow is excluded — it's already in combined_dfs as 'node_inflow__dt'
+    # ``dispatch_groups`` lists the groups that have on-disk dispatch tables.
+    # The model has two separate sets — nodeGroupIndicators (summary
+    # indicators) and nodeGroupDispatch (dispatch table output) — and a
+    # project may flag either independently. We use nodeGroupDispatch
+    # because that's the set that gates the actual time-series files
+    # the dispatch view consumes.
+    file_mapping = {
+        'dispatch_groups': ('nodeGroupDispatch.parquet', True),
+        'group_node': ('group_node.parquet', True),
+        'group_process_node': ('group_process_node.parquet', False),
+        # ProcessGroup aggregation files
+        'processGroup_Unit_to_group': ('nodeGroupDispatch__processGroup_Unit_to_group.parquet', False),
+        'processGroup_Group_to_unit': ('nodeGroupDispatch__processGroup_Group_to_unit.parquet', False),
+        'processGroup_Connection': ('nodeGroupDispatch__processGroup_Connection.parquet', False),
+        # ProcessGroup member files
+        'processGroup_unit_to_node_members': ('nodeGroupDispatch__processGroup__process__unit__to_node.parquet', False),
+        'processGroup_node_to_unit_members': ('nodeGroupDispatch__processGroup__process__node__to_unit.parquet', False),
+        'processGroup_connection_to_node_members': ('nodeGroupDispatch__processGroup__process__connection__to_node.parquet', False),
+        'processGroup_node_to_connection_members': ('nodeGroupDispatch__processGroup__process__node__to_connection.parquet', False),
+        # Not_in_aggregate files
+        'not_in_aggregate_unit_to_node': ('nodeGroupDispatch__process__unit__to_node_Not_in_aggregate.parquet', False),
+        'not_in_aggregate_node_to_unit': ('nodeGroupDispatch__process__node__to_unit_Not_in_aggregate.parquet', False),
+        'not_in_aggregate_connection_to_node': ('nodeGroupDispatch__process__connection__to_node_Not_in_aggregate.parquet', False),
+        'not_in_aggregate_node_to_connection': ('nodeGroupDispatch__process__node__to_connection_Not_in_aggregate.parquet', False),
+        'not_in_aggregate_connection': ('nodeGroupDispatch__connection_Not_in_aggregate.parquet', False),
+        # Process fully inside (for internal losses)
+        'process_fully_inside': ('nodeGroupDispatch__process_fully_inside.parquet', False),
+    }
+
+    for key, (filename, required) in file_mapping.items():
+        filepath = parquet_dir / filename
+        if filepath.exists():
+            mappings[key] = read_lean_parquet(filepath)
+        elif required:
+            print(f"Warning: Required file {filepath} not found")
+            mappings[key] = None
+        else:
+            mappings[key] = None
+
+    # nodeGroupDispatch.parquet stores its single membership column under the
+    # set name 'nodeGroupDispatch'; downstream consumers expect 'group'
+    # (matching group_node / group_process_node).
+    dg_df = mappings.get('dispatch_groups')
+    if dg_df is not None and 'nodeGroupDispatch' in dg_df.columns:
+        mappings['dispatch_groups'] = dg_df.rename(columns={'nodeGroupDispatch': 'group'})
+
+    return mappings
+
+
+def combine_dispatch_mappings(
+    scenario_folders: dict[str, str],
+    parquet_subdir: str,
+) -> DispatchMappings:
+    """
+    Load and combine dispatch mapping files across all scenarios.
+
+    Each mapping key's DataFrames are concatenated and indexed by 'scenario'.
+    Use ``DispatchMappings.get_for_scenario(field, scenario)`` for per-scenario views.
+
+    Parameters:
+    -----------
+    scenario_folders : dict
+        Mapping of scenario names to folder paths
+    parquet_subdir : str
+        Subdirectory within each folder containing parquet files
+
+    Returns:
+    --------
+    DispatchMappings : Combined mapping dataclass with 'scenario' in the index
+    """
+    # Collect per-key lists: {key: [df, df, ...]}
+    collected: dict[str, list[pd.DataFrame]] = {}
+
+    for scen_name, folder_path in scenario_folders.items():
+        pq_dir = Path(folder_path) / parquet_subdir / scen_name
+        mappings = load_dispatch_mappings(pq_dir)
+        for key, df in mappings.items():
+            if df is None or df.empty:
+                continue
+            # Re-tag the scenario column to the folder identity (the loop
+            # key) so two folders holding the same model scenario stay
+            # distinct.  No-op when the baked-in tag already equals the
+            # folder.  Mirrors combine_parquet_files; see
+            # specs/dispatch_scenario_identity_retag.md.
+            if 'scenario' in df.columns:
+                df = df.assign(scenario=scen_name)
+            if key not in collected:
+                collected[key] = []
+            collected[key].append(df)
+
+    # Concatenate and set scenario as index
+    combined: dict[str, pd.DataFrame] = {}
+    for key, dfs in collected.items():
+        combined_df = pd.concat(dfs, ignore_index=True)
+        if 'scenario' in combined_df.columns:
+            combined_df = combined_df.set_index('scenario')
+        combined[key] = combined_df
+
+    return DispatchMappings(**combined)
+
+
+def data_scenario_tags(mappings: DispatchMappings) -> list[str]:
+    """Return the scenario identities present in the combined dispatch data.
+
+    ``combine_dispatch_mappings`` re-tags each scenario to its folder
+    identity at load time (see
+    specs/dispatch_scenario_identity_retag.md), so the values returned here
+    are the folder names — already the canonical dispatch slice key.  This
+    reads them back from the data (e.g. so a caller can enumerate exactly
+    the folders that carry dispatch groups, skipping ones that don't).
+
+    Identities are returned first-seen-ordered (mirroring the folder order
+    that ``combine_dispatch_mappings`` concatenated), de-duplicated.
+    Returns ``[]`` when the data carries no scenario index/column.
+    """
+    df = mappings.dispatch_groups
+    if df is None or df.empty:
+        return []
+    if df.index.name == 'scenario':
+        return list(dict.fromkeys(df.index.tolist()))
+    if 'scenario' in df.columns:
+        return list(dict.fromkeys(df['scenario'].tolist()))
+    return []
+
+
+def resolve_data_scenario_tag(mappings: DispatchMappings, fallback: str) -> str:
+    """Return the single scenario identity present in the dispatch data.
+
+    Convenience wrapper over :func:`data_scenario_tags` for the
+    single-scenario case (post re-tag this is the folder identity).
+    Returns *fallback* when the identity is absent or not unambiguously
+    single-valued (e.g. combined multi-scenario mappings — use
+    :func:`data_scenario_tags` there instead).
+    """
+    tags = data_scenario_tags(mappings)
+    if len(tags) == 1:
+        return tags[0]
+    return fallback
+
+
+def get_group_node_multiindex(
+    group_node_df: pd.DataFrame | None,
+    scenario: str | None = None,
+) -> pd.MultiIndex | None:
+    """
+    Create a MultiIndex from group_node dataframe for a specific scenario.
+
+    Parameters:
+    -----------
+    group_node_df : pd.DataFrame
+        DataFrame with columns [scenario, group, node]
+    scenario : str, optional
+        Scenario name to filter by. If None, uses all rows.
+
+    Returns:
+    --------
+    pd.MultiIndex : MultiIndex with levels [group, node]
+    """
+    if group_node_df is None:
+        return None
+    if scenario is None:
+        df = group_node_df
+    elif 'scenario' in group_node_df.columns:
+        df = group_node_df[group_node_df['scenario'] == scenario]
+    elif group_node_df.index.name == 'scenario':
+        df = group_node_df.xs(scenario)
+    else:
+        df = group_node_df
+    return pd.MultiIndex.from_frame(df[['group', 'node']])