flextool 4.0.0__py3-none-any.whl

flextool/scenario_comparison/plan_union.py

@@ -0,0 +1,244 @@
+"""Comparison-mode helpers built on the merged ``default_plots.yaml``.
+
+Each leaf config in ``default_plots.yaml`` carries optional comparison-mode
+add-ons:
+
+* ``scenario_rule`` — single character that says how the ``scenario`` dim is
+  folded into the comparison view (``g``=grouped/coloured bars, ``l``=lines,
+  ``u``=subplots, ``f``=files, ``s``=stacked, ``e``=expand-axis).
+* ``comparison_overrides`` — optional dict of plot-setting overrides applied
+  on top of the single config in comparison mode.  When it includes
+  ``map_dimensions_for_plots``, that explicit value wins over the value
+  auto-derived from ``single + scenario_rule``.
+
+Two helpers live here:
+
+* :func:`derive_comparison_config` — builds a comparison-mode ``PlotConfig``
+  from a single-mode one by inserting ``s`` + ``scenario_rule`` into the
+  column part of ``index_types``/``rules`` and merging in
+  ``comparison_overrides``.
+* :func:`union_raw_data` — reads each viewer scenario's raw result parquet
+  and unions them with ``scenario`` as the outermost column-MultiIndex
+  level, so the derived config can be applied via the standard
+  ``compute_live_plan`` pipeline.
+"""
+from __future__ import annotations
+
+import dataclasses
+import logging
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+try:
+    from flextool.lean_parquet import read_lean_parquet
+except ImportError:  # pragma: no cover — defensive fallback
+    read_lean_parquet = None  # type: ignore[assignment]
+
+
+logger = logging.getLogger(__name__)
+
+
+def _read_parquet(path: Path) -> pd.DataFrame:
+    """Read a parquet file, preferring the lean reader when available."""
+    if read_lean_parquet is not None:
+        return read_lean_parquet(path)
+    return pd.read_parquet(path)
+
+
+def _cfg_get(cfg: Any, name: str) -> Any:
+    """Read *name* from a config that may be a PlotConfig dataclass or dict."""
+    if hasattr(cfg, name):
+        return getattr(cfg, name)
+    if isinstance(cfg, dict):
+        return cfg.get(name)
+    return None
+
+
+def has_comparison_view(cfg: Any) -> bool:
+    """Return ``True`` if *cfg* is renderable in comparison mode.
+
+    A config has a comparison view when it sets ``scenario_rule`` (the
+    ``scenario`` dim is auto-folded into the rules) **or** when its
+    ``comparison_overrides`` supply ``map_dimensions_for_plots`` (the
+    comparison layout is given explicitly, so ``scenario_rule`` is
+    irrelevant).  Mirrors the gate in :func:`derive_comparison_config`.
+    """
+    if _cfg_get(cfg, "scenario_rule") is not None:
+        return True
+    overrides = _cfg_get(cfg, "comparison_overrides")
+    if isinstance(overrides, dict):
+        md = overrides.get("map_dimensions_for_plots")
+        if isinstance(md, (list, tuple)) and len(md) >= 2:
+            return True
+    return False
+
+
+def derive_comparison_config(single_cfg: Any) -> Any:
+    """Build a comparison-mode PlotConfig from a single-mode one.
+
+    The unioned **raw** comparison frame has shape::
+
+        row = (single's row dims)
+        col = ('scenario', *single's col dims)
+
+    so we prepend ``s`` to ``index_types``'s column part and prepend
+    ``scenario_rule`` to the rules string's column part.  The single-mode
+    rules string is otherwise unchanged, which means
+    ``_apply_dimension_rules`` will reshape the data exactly as it does in
+    single mode, with one extra column dim representing scenario.
+
+    ``comparison_overrides`` (also on the single config) provides per-mode
+    overrides for layout knobs.  Overrides win — including
+    ``map_dimensions_for_plots`` if explicitly given (used by the rare
+    configs whose comparison view wants a different visual treatment than
+    "single + scenario_rule prepended").
+
+    ``scenario_rule`` is normally required to mark a config as
+    comparison-renderable.  It may be omitted, however, when
+    ``comparison_overrides`` supplies ``map_dimensions_for_plots`` — that
+    fully specifies the comparison layout (including the scenario dim), so
+    no auto-derivation is needed and ``scenario_rule`` is irrelevant.
+
+    Raises ``ValueError`` when neither ``scenario_rule`` nor an override
+    ``map_dimensions_for_plots`` is set, or when ``map_dimensions_for_plots``
+    is malformed.
+    """
+    # Read overrides first: a comparison config may give
+    # ``map_dimensions_for_plots`` directly here, which fully defines the
+    # comparison layout and makes ``scenario_rule`` unnecessary.
+    overrides = getattr(single_cfg, "comparison_overrides", None)
+    if overrides is None and isinstance(single_cfg, dict):
+        overrides = single_cfg.get("comparison_overrides")
+    overrides = dict(overrides or {})
+    override_md = overrides.get("map_dimensions_for_plots")
+    has_override_md = isinstance(override_md, (list, tuple)) and len(override_md) >= 2
+
+    scenario_rule = getattr(single_cfg, "scenario_rule", None)
+    if scenario_rule is None and isinstance(single_cfg, dict):
+        scenario_rule = single_cfg.get("scenario_rule")
+    if scenario_rule is None and not has_override_md:
+        raise ValueError(
+            "derive_comparison_config: the single config has no "
+            "scenario_rule — define one (e.g. 'g' for grouped bars, "
+            "'l' for lines, 'u' for subplots), or supply "
+            "map_dimensions_for_plots in comparison_overrides, to enable "
+            "comparison view."
+        )
+
+    if scenario_rule is None:
+        # ``comparison_overrides`` fully specifies the comparison layout via
+        # ``map_dimensions_for_plots``; skip auto-derivation and apply the
+        # override as-is below.
+        new_md = None
+    else:
+        if hasattr(single_cfg, "map_dimensions_for_plots"):
+            md = single_cfg.map_dimensions_for_plots
+        else:
+            md = single_cfg.get("map_dimensions_for_plots")
+        if not isinstance(md, (list, tuple)) or len(md) < 2:
+            raise ValueError(
+                f"derive_comparison_config: map_dimensions_for_plots must be a "
+                f"2-element [index_types, rules] list, got {md!r}"
+            )
+        idx, rules = md[0], md[1]
+        if not isinstance(idx, str) or "_" not in idx:
+            raise ValueError(
+                f"derive_comparison_config: index_types must contain '_' "
+                f"separating row and column parts, got {idx!r}"
+            )
+        if not isinstance(rules, str) or "_" not in rules:
+            raise ValueError(
+                f"derive_comparison_config: rules must contain '_' separating "
+                f"row and column parts, got {rules!r}"
+            )
+        row_idx, col_idx = idx.split("_", 1)
+        row_rules, col_rules = rules.split("_", 1)
+
+        # If the single rules already name 'scenario' in the col part (e.g.
+        # ``[d_s, s_b]`` for a costs-by-scenario chart that's the same shape
+        # in both modes), don't auto-prepend another ``s`` — the existing rule
+        # for scenario stands.  ``scenario_rule`` is required to mark the
+        # config as comparison-renderable, but its value is unused on this
+        # branch (and may be overridden via ``comparison_overrides``).
+        if "s" in col_idx:
+            new_idx, new_rules = idx, rules
+        else:
+            new_idx = f"{row_idx}_s{col_idx}"
+            new_rules = f"{row_rules}_{scenario_rule}{col_rules}"
+        new_md = [new_idx, new_rules]
+
+    if hasattr(single_cfg, "__dataclass_fields__"):
+        valid_fields = set(single_cfg.__dataclass_fields__.keys())
+        replace_kwargs = {k: v for k, v in overrides.items() if k in valid_fields}
+        # Two-step replace: derive first, then apply overrides — so a config
+        # whose comparison rules differ beyond just adding scenario can carry
+        # ``map_dimensions_for_plots`` in ``comparison_overrides`` and have it
+        # take precedence over the derived value.
+        derived = (
+            single_cfg if new_md is None
+            else dataclasses.replace(single_cfg, map_dimensions_for_plots=new_md)
+        )
+        return dataclasses.replace(derived, **replace_kwargs)
+
+    if isinstance(single_cfg, dict):
+        new_cfg = dict(single_cfg)
+        if new_md is not None:
+            new_cfg["map_dimensions_for_plots"] = new_md
+        new_cfg.update(overrides)  # overrides win, including map_dimensions_for_plots
+        return new_cfg
+
+    raise TypeError(
+        f"derive_comparison_config: unsupported config type {type(single_cfg)!r}"
+    )
+
+
+def union_raw_data(
+    project_path: Path,
+    scenarios: list[str],
+    result_key: str,
+) -> pd.DataFrame | None:
+    """Read each scenario's raw result parquet and union them.
+
+    Each per-scenario raw parquet at ``output_parquet/<scenario>/<rk>.parquet``
+    already carries a ``scenario`` column-MultiIndex level (added at write
+    time); we drop that local level and re-concat across scenarios with the
+    **viewer-supplied** scenario name as the new outermost ``scenario``
+    level.  This handles the case where a scenario folder name differs from
+    the embedded scenario name (e.g. sensitivity replicas where the folder
+    is ``scenario_test_6h_2`` but the embedded name is ``scenario_test_6h``).
+
+    Returns ``None`` when no raw parquet exists for any of *scenarios*.
+    Missing files for a subset are silently skipped — the union proceeds
+    with whatever is present.
+    """
+    pieces: list[pd.DataFrame] = []
+    found: list[str] = []
+    for s in scenarios:
+        path = Path(project_path) / "output_parquet" / s / f"{result_key}.parquet"
+        if not path.is_file():
+            continue
+        try:
+            df = _read_parquet(path)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("union_raw_data: failed to read %s: %s", path, exc)
+            continue
+        # Strip the embedded scenario level so the keys=... we set below is
+        # the only authoritative scenario label.  Two cases:
+        #  * MultiIndex columns with a 'scenario' level → droplevel.
+        #  * Single-level Index named 'scenario' (e.g. costs_discounted_p_,
+        #    a category × scenario table) → squeeze to a Series so concat
+        #    below adds 'scenario' as the only column level (otherwise the
+        #    inner unnamed level lingers and breaks rules like ``s_b``).
+        if isinstance(df.columns, pd.MultiIndex) and "scenario" in df.columns.names:
+            df = df.droplevel("scenario", axis=1)
+        elif (not isinstance(df.columns, pd.MultiIndex)
+              and df.columns.name == "scenario"
+              and len(df.columns) == 1):
+            df = df.iloc[:, 0]  # Series with the row index preserved
+        pieces.append(df)
+        found.append(s)
+    if not pieces:
+        return None
+    return pd.concat(pieces, axis=1, keys=found, names=["scenario"])

flextool/scenario_comparison/plot_settings_seed.py

@@ -0,0 +1,205 @@
+"""Structured seeding of discovered colors into ``plot_settings.yaml``.
+
+The GUI color/order picker (Stage 6) is now the canonical editor for the
+project ``plot_settings.yaml``, so the file is treated as a plain STRUCTURED
+data file: load with :mod:`pyyaml`, modify the dict, dump with
+``sort_keys=False`` to preserve key order (= stacking order).  Comments are no
+longer preserved in the PROJECT file — they live only in the bundled template
+(schema reference) and as future UI tooltips/labels.
+
+This module provides the single read/modify/write helper used by both seeds:
+
+* the upfront **input-DB** seed
+  (:mod:`flextool.scenario_comparison.input_entity_colors`), which knows the
+  FULL per-class entity set from the input DB and therefore ADDS new names and
+  PRUNES stale ones; and
+* the **dispatch** seed
+  (:func:`flextool.scenario_comparison.orchestrator._seed_dispatch_colors_into_plot_settings`),
+  which discovers only a SUBSET of entities and therefore ADDS only (never
+  prunes).
+
+Add semantics (always): for each entity class / scenario, names not already
+present (case-insensitive) are appended AFTER the existing entries in that
+subsection.  Existing entry ORDER and existing VALUES are never disturbed (a
+user's hand-picked color or order is preserved).
+
+Prune semantics (opt-in, ``prune_entities``): names whose lowercased form is
+NOT in the supplied full set for that class are removed from ``entities.<class>``.
+Only ``entities.*`` is ever pruned; ``categories.*`` (schema-fixed) and
+``scenarios`` are never pruned.
+
+A pass that adds nothing and prunes nothing makes no write (returns ``False``),
+so a freshly-copied, still-commented project file keeps its comments until the
+first real change.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import yaml
+
+# Entity classes managed in the ``entities`` section (the group class is split
+# into nodeGroup/flowGroup by membership upstream).  Legacy ``group`` is NOT
+# listed — old projects keep an inert ``entities.group`` section, never
+# written or pruned here.
+_ENTITY_CLASSES = ("nodeGroup", "flowGroup", "unit", "connection", "node")
+
+
+def dump_plot_settings(data: dict) -> str:
+    """Serialize a ``plot_settings.yaml`` mapping to text.
+
+    Single source of truth for the dump parameters shared by the seed
+    writer here and the GUI color/order picker
+    (:mod:`flextool.gui.dialogs.plot_settings_picker`).  ``sort_keys=False``
+    preserves KEY ORDER (= stacking order); ``default_flow_style=False``
+    keeps the block style; ``allow_unicode=True`` round-trips unicode
+    entity names verbatim.
+    """
+    return yaml.safe_dump(
+        data,
+        sort_keys=False,
+        default_flow_style=False,
+        allow_unicode=True,
+    )
+
+
+def _load_mapping(path: Path) -> dict:
+    """Load *path* as a YAML mapping (``None`` / empty -> ``{}``).
+
+    Raises :class:`ValueError` on a parse error or a non-mapping top level.
+    """
+    text = path.read_text(encoding="utf-8")
+    try:
+        parsed = yaml.safe_load(text)
+    except yaml.YAMLError as exc:  # pragma: no cover - corrupt file
+        raise ValueError(f"Cannot parse {path}: {exc}") from exc
+    if parsed is None:
+        return {}
+    if not isinstance(parsed, dict):
+        raise ValueError(f"{path} is not a YAML mapping")
+    return parsed
+
+
+def _existing_lc(section: dict) -> set[str]:
+    """Lowercased key set of *section* (for case-insensitive de-dup)."""
+    return {str(k).lower() for k in section}
+
+
+def _add_entries(
+    section: dict,
+    wanted: dict[str, str],
+) -> bool:
+    """Append names from *wanted* not already in *section* (case-insensitive).
+
+    Mutates *section* in place; existing entries (order + value) are left
+    untouched.  Returns ``True`` if anything was added.
+    """
+    existing = _existing_lc(section)
+    added = False
+    for name, color in wanted.items():
+        if str(name).lower() in existing:
+            continue
+        section[name] = color
+        existing.add(str(name).lower())
+        added = True
+    return added
+
+
+def _prune_entries(section: dict, keep_lc: set[str]) -> bool:
+    """Remove keys of *section* whose lowercase is NOT in *keep_lc*.
+
+    Mutates *section* in place.  Returns ``True`` if anything was removed.
+    """
+    stale = [k for k in section if str(k).lower() not in keep_lc]
+    for k in stale:
+        del section[k]
+    return bool(stale)
+
+
+def seed_colors_into_plot_settings(
+    path: Path,
+    entity_colors: dict[str, dict[str, str]],
+    scenario_colors: dict[str, str],
+    prune_entities: dict[str, set[str]] | None = None,
+) -> bool:
+    """Structurally add (and optionally prune) colors in ``plot_settings.yaml``.
+
+    *entity_colors* maps entity class (``nodeGroup`` / ``flowGroup`` /
+    ``unit`` / ``connection`` / ``node``) to an ordered ``{name -> color}``
+    mapping (unknown classes are ignored); *scenario_colors* is
+    the ``{name -> color}`` mapping for the ``scenarios`` section.  Color
+    values are written verbatim (``"#RRGGBB"`` strings, matplotlib color
+    names, or ``{color, neg_color}`` dicts — whatever the caller supplies).
+
+    ADD: names not already present (active key, case-insensitively) in their
+    target (sub)section are appended AFTER the existing entries.  Existing
+    entries — their order and their values — are never overwritten.
+
+    PRUNE (only when *prune_entities* is given): for each class present in
+    *prune_entities*, entity names in ``entities.<class>`` whose lowercased
+    name is NOT in ``prune_entities[class]`` are removed.  This is how the
+    input-DB seed drops entities that no longer exist in the DB.  Only
+    ``entities.*`` is pruned — never ``categories.*`` (schema-fixed) nor
+    ``scenarios``.  When *prune_entities* is ``None`` the call is add-only.
+
+    The file is loaded with :func:`yaml.safe_load`, modified, and re-dumped
+    with ``sort_keys=False`` so KEY ORDER (= stacking order) is preserved and
+    ``"#RRGGBB"`` values round-trip as quoted strings.  ``categories.*``,
+    ``scenarios``, and any unmanaged keys are preserved.
+
+    Returns ``True`` when the file was modified, ``False`` when nothing was
+    added or pruned (the no-op / idempotent case — the file is left untouched
+    on disk so a still-commented fresh copy keeps its comments).
+    """
+    path = Path(path)
+    data = _load_mapping(path)
+
+    modified = False
+
+    # --- Entities: add + (optional) prune -----------------------------------
+    entities = data.get("entities")
+    entities_existed = isinstance(entities, dict)
+    if not entities_existed:
+        entities = {}
+
+    for cls in _ENTITY_CLASSES:
+        wanted = entity_colors.get(cls) or {}
+        do_prune = prune_entities is not None and cls in prune_entities
+        if not wanted and not do_prune:
+            continue
+
+        class_map = entities.get(cls)
+        if not isinstance(class_map, dict):
+            class_map = {}
+
+        changed_here = False
+        if wanted and _add_entries(class_map, wanted):
+            changed_here = True
+        if do_prune and _prune_entries(class_map, prune_entities[cls]):
+            changed_here = True
+
+        if changed_here:
+            # Attach the (possibly newly created) class map to the section.
+            entities[cls] = class_map
+            modified = True
+
+    # Only attach the entities section if we actually changed something and
+    # there is content to write — never inject an empty section on a no-op.
+    if modified and entities and not entities_existed:
+        data["entities"] = entities
+
+    # --- Scenarios: add only -------------------------------------------------
+    if scenario_colors:
+        scenarios = data.get("scenarios")
+        if not isinstance(scenarios, dict):
+            scenarios = {}
+        if _add_entries(scenarios, scenario_colors):
+            data["scenarios"] = scenarios
+            modified = True
+
+    if not modified:
+        return False
+
+    path.write_text(dump_plot_settings(data), encoding="utf-8")
+    return True

flextool/schemas/AXIS_CONTRACT.md

@@ -0,0 +1,71 @@
+# FlexTool axis contract
+
+This directory holds two coupled JSON files that together define the
+canonical FlexTool axis vocabularies for the `pl.Enum` dtype activation
+(see `specs/enum_dtype_refactor_plan.md`):
+
+| File | Purpose |
+|---|---|
+| `flextool_axis_contract.json` | Authoritative declaration of every axis the cascade casts to `pl.Enum`. Each axis row names the entity class / parameter map / synthetic-token list that supplies its vocabulary. |
+| `flextool_axis_contract.schema.json` | JSON Schema (draft 2020-12) validating the structure of `flextool_axis_contract.json` itself. |
+
+The contract is consumed by `flextool/spinedb_backend/_axis_enums.py`
+(landing in Phase 1) at runtime to build the per-axis `pl.Enum` dtypes
+that the Backend and the cascade cast against.
+
+## Coupling with `spinedb_schema.json`
+
+The contract migrates **hand-in-hand** with `spinedb_schema.json`.
+Any schema change that adds, renames, or removes an entity class or a
+parameter definition referenced by an axis row REQUIRES updating both
+files in the same commit:
+
+* If you add a new entity class that the cascade joins on, decide
+  whether it deserves its own axis and add a row to `axes`.
+* If you rename an entity class, update every `source` field that
+  references it.
+* If you remove an entity class, remove or repoint the axis row.
+* If you rename or remove a parameter referenced under
+  `source_type: parameter_keys` or `parameter_value_list`, update the
+  `source.entity_class` / `source.parameter` / `source.parameters` field.
+
+The `_template_master_version_compat` block at the top of
+`flextool_axis_contract.json` pins the schema commit hash the contract
+was last verified against. Bump that block whenever a schema-affecting
+edit lands.
+
+## Self-validation
+
+`tests/spinedb_backend/test_axis_contract.py` validates the contract
+against its own schema and cross-checks that every entity class /
+parameter referenced by an axis row exists in `spinedb_schema.json`.
+Run it after every contract edit:
+
+```bash
+python -m pytest tests/spinedb_backend/test_axis_contract.py -v
+```
+
+## Adding or changing an axis
+
+1. Decide the `source_type`:
+   - `entity_class` — one entity class, single member list.
+   - `entity_class_union` — multiple entity classes unioned (e.g. process = unit ∪ connection).
+   - `parameter_keys` — vocabulary lives in the keys of a parameter Map (or the values of an Array).
+   - `parameter_value_list` — vocabulary is the values of a Spine `parameter_value_list`.
+   - `synthetic` — fixed tokens introduced by cascade derivations (no DB source).
+   - `union` — composite axis combining other axes' vocabularies.
+2. Fill in the `source` field — string / list / object per `source_type`.
+3. Add `column_synonyms` if the cascade uses friendly column names that should map to the same enum (mirrors `_AXIS_SYNONYMS` in `flextool/engine_polars/_axis_enums.py`).
+4. Add `filter` if the vocabulary is intersected with another set (e.g. scenario filter, realized periods).
+5. If introducing synthetic tokens that aren't from the DB, register them in `synthetic_token_allowlist` AND in the axis's `tokens` field (for `source_type: synthetic`) or `tokens_default_extension` (for hybrid axes).
+6. Update the schema (`flextool_axis_contract.schema.json`) only if you need a new `source_type` enum value or a new optional field shape.
+7. Run the self-check tests.
+
+## Outstanding review notes
+
+The `_review_notes` array at the bottom of `flextool_axis_contract.json`
+tracks Phase 0 audit findings that need Phase 1 resolution decisions
+(e.g. column-letter collisions between axes like `c` for both commodity
+and constraint, or `b` for both block and branch). These notes are
+informational; Phase 1's `build_axis_enums` is where the disambiguation
+strategy is actually implemented.