"""Running a linking scenario: validate, transform, (optionally) simulate.
Two entry points:
``apply_scenario(system, data, scenario)``
The core. Turns external-model shocks into transformed EUROMOD input —
the counterfactual and, when the methodology restructures rows, a matching
baseline built on the same rows. Executes nothing.
``run_scenario(system, scenario, input_path=...)``
Convenience: loads the dataset, applies the scenario and runs both halves
through ``System.run()``.
Applications that need caching, retries or their own response envelope build on
`apply_scenario` and run the frames themselves; that is why nothing here knows
about caches or run ids.
The API exposes *shocks*, not methodologies: the transformation methodology is
resolved from the shock table's channels/metrics and reported back. An optional
``methodology`` pin exists only for exact reproduction or disambiguation.
Constants semantics:
- scenario ``constants`` are context, applied to BOTH runs, so the delta
isolates the shock;
- shock records with ``channel: constant`` are part of the shock and apply to
the COUNTERFACTUAL only (metric = constant name, period = constant group,
op must be ``set``). A constants-only shock table needs no methodology.
Errors are exceptions (`ScenarioError` carrying `.problems`, `RunError`), not
error dicts — callers that need a machine-readable envelope build it themselves.
"""
import hashlib
import json
import logging
from euromod_linking import dimensions, ingest as ingest_mod, registry, shock_table
from euromod_linking.resources import load_schema
from euromod_linking.runner import RunError, execute, frames_identical, resolve_dataset
logger = logging.getLogger(__name__)
[docs]
class ScenarioError(ValueError):
"""A scenario cannot be applied. ``.problems`` lists every problem found,
not just the first, so one call surfaces everything to fix."""
def __init__(self, problems: list[str]):
super().__init__(problems[0] if problems else "invalid scenario")
self.problems = problems
# --- the scenario document ----------------------------------------------------
def _validate_structure(scenario: dict) -> list[str]:
"""JSON-Schema problems with the scenario document ([] = ok)."""
import jsonschema
validator = jsonschema.Draft202012Validator(load_schema("scenario.v1.schema.json"))
problems = []
for err in sorted(validator.iter_errors(scenario), key=lambda e: list(e.path)):
loc = "/".join(map(str, err.path)) or "<root>"
problems.append(f"{loc}: {err.message}")
return problems
def _validate_params(scenario: dict, spec) -> list[str]:
"""Validate scenario['params'] against the method's own params_schema.
additionalProperties: false in that schema is what blocks methodology knobs."""
import jsonschema
problems = []
validator = jsonschema.Draft202012Validator(spec.params_schema)
for err in sorted(validator.iter_errors(scenario.get("params") or {}), key=lambda e: list(e.path)):
loc = "/".join(map(str, err.path)) or "params"
problems.append(f"params/{loc}: {err.message}")
return problems
def _resolve_shocks(scenario: dict) -> tuple:
"""Resolve scenario['shocks'] to (DataFrame, shock_table_id, summary, warnings).
Two forms: {inline: [records]} or {file, mapping} (ingested here). The id is
derived from the canonical records, so re-running either form on the same
shocks yields the same id without anything being stored. Raises ScenarioError.
"""
shocks = scenario.get("shocks") or {}
warnings: list[str] = []
if "inline" in shocks:
try:
df = shock_table.normalize(shocks["inline"])
except shock_table.ShockTableError as e:
raise ScenarioError(e.problems)
summary = shock_table.describe(df, origin="inline")
return df, summary["shock_table_id"], summary, warnings
if "file" in shocks:
try:
spec = ingest_mod.load_mapping(shocks["mapping"])
records, warnings = ingest_mod.ingest(
shocks["file"], spec, country=scenario.get("country_code"))
df = shock_table.normalize(records)
except (ingest_mod.IngestError, shock_table.ShockTableError) as e:
raise ScenarioError(getattr(e, "problems", None) or [str(e)])
summary = shock_table.describe(df, origin="ingest", file=str(shocks["file"]),
mapping=str(shocks["mapping"]),
dimensions=spec.get("dimensions") or {})
return df, summary["shock_table_id"], summary, warnings
raise ScenarioError(["shocks must be one of {inline}, {file, mapping}"])
def _scenario_constants(scenario: dict) -> dict:
"""Scenario constants as the {(name, group): value} form the run takes."""
out = {}
for e in scenario.get("constants") or []:
out[(str(e["name"]).strip(), str(e.get("group") or "").strip())] = str(e["value"])
return out
[docs]
def fingerprint(scenario: dict, shock_table_id: str, methodology: str,
code_fingerprint: str = "") -> str:
"""Canonical scenario hash for a result-cache key. Shocks are represented by
the table's content id, so inline vs stored vs re-ingested identical shocks
hit the same cache entry. Excludes the free-text 'name' label."""
payload = {
"scenario_version": 1,
"methodology": methodology,
# Methodology *code*, not just its name: editing a method must not
# serve results the previous implementation produced.
**({"methodology_code": code_fingerprint} if code_fingerprint else {}),
"cc": str(scenario.get("country_code", "")).upper(),
"system": scenario.get("system_name", ""),
"dataset": scenario.get("dataset_name") or "",
"shocks": shock_table_id,
"constants": sorted((f"{n}|{g}" if g else n, v)
for (n, g), v in _scenario_constants(scenario).items()),
"params": scenario.get("params") or {},
"addons": sorted([list(map(str, a)) for a in (scenario.get("addons") or [])]),
"extensions": sorted([[str(e[0]), bool(e[1])] for e in (scenario.get("extensions") or [])]),
}
blob = json.dumps(payload, sort_keys=True, default=str)
return hashlib.sha256(blob.encode("utf-8")).hexdigest()
# --- running it ---------------------------------------------------------------
[docs]
class NoEffectError(RuntimeError):
"""The methodology transformed the input but the engine produced identical
output — the run is invalid.
Emphatically *not* "the reform has no impact": it means the model never
acted on the transformation (typically a required add-on or switch is not
available for this country/system).
"""
# --- pieces of the plan -------------------------------------------------------
def _split_constant_shocks(shocks):
"""Separate ``channel=constant`` records into {(name, group): value}.
Returns (population_shocks, constant_overrides). Raises ScenarioError when a
constant shock is relative — only ``set`` is meaningful for a parameter.
"""
problems, overrides = [], {}
for t in shocks[shocks["channel"] == "constant"].itertuples():
if t.op != "set":
problems.append(f"Constant shock {t.metric!r}: op must be 'set' (got {t.op!r}; "
"relative constant changes are not supported)")
continue
overrides[(t.metric, str(t.period))] = str(t.value)
if problems:
raise ScenarioError(problems)
return shocks[shocks["channel"] != "constant"], overrides
def _resolve_methodology(scenario: dict, population_shocks):
"""The pinned methodology, else dispatch from the shock channels.
None when the scenario is constants-only."""
pin = scenario.get("methodology")
if pin:
return registry.resolve(pin)
if population_shocks.empty:
return None
return registry.resolve_for_channels(set(population_shocks["channel"].unique()),
set(population_shocks["metric"].unique()))
def _check_contract(spec, population_shocks) -> list[str]:
"""Shocks vs the methodology's declared contract (a safety net for pinned
references — dispatch already guarantees channels and metrics match)."""
problems = []
for ch in sorted(population_shocks["channel"].unique()):
if ch not in spec.channels_consumed:
problems.append(f"Methodology {spec.name} does not consume channel '{ch}' "
f"(consumes: {sorted(spec.channels_consumed)})")
for m in sorted(population_shocks["metric"].unique()):
if spec.metrics_consumed and m not in spec.metrics_consumed:
problems.append(f"Methodology {spec.name} does not consume metric '{m}' "
f"(consumes: {sorted(spec.metrics_consumed)})")
return sorted(set(problems))
[docs]
def run_arguments(scenario: dict, spec) -> tuple[list, list]:
"""Add-ons and extension switches the run needs: the methodology's own
requirements (with ``{cc}`` resolved) plus anything the scenario adds."""
cc = str(scenario["country_code"]).upper()
addon_entries, switch_entries = (spec.addon_requirements or ((), ())) if spec else ((), ())
addons = [[str(x).format(cc=cc) for x in a] if isinstance(a, (list, tuple))
else str(a).format(cc=cc) for a in addon_entries]
switches = [[str(s[0]).format(cc=cc), bool(s[1])] for s in switch_entries]
addons += list(scenario.get("addons") or [])
switches += [[str(e[0]), bool(e[1])] for e in (scenario.get("extensions") or [])]
return addons, switches
def _compatibility(system, spec):
"""The methodology's requirements checked against the model, or None.
None means the check could not run at all — an unreadable model, or the
caller setting ``EUROMOD_SKIP_COMPAT_CHECK``. That is deliberately not a
failure: this check exists to turn a late `NoEffectError` into an early
`ScenarioError`, and it must never become a new way for a working scenario
to be refused."""
from euromod_linking import compat
if compat._skip_requested():
return None
try:
return compat.check_compatibility(system, spec)
except Exception:
logger.debug("compatibility check unavailable for %s", spec.name, exc_info=True)
return None
[docs]
def cell_population(data, population_shocks) -> dict:
"""Weighted people matched by each distinct shock cell.
Surfacing this before a run is what turns a mistyped cell (``deh=34``) into
a visible zero rather than a shock that silently hits nobody.
"""
import pandas as pd
counts = {}
for group in sorted(g for g in population_shocks["group"].unique() if g):
mask = pd.Series(True, index=data.index)
for k, v in dimensions.parse_group(group).items():
series = dimensions.derive_region(data) if k == "region" else data.get(k)
if series is None:
mask = None
break
mask &= (series.str.startswith(v) if k == "region"
else dimensions.matches(series, dimensions.parse_value_spec(v)))
if mask is None:
continue
counts[group] = {"n_rows": int(mask.sum()),
"n_weighted_all_ages": round(float(data.loc[mask, "dwt"].sum()), 1)}
return counts
# --- the core transform -------------------------------------------------------
[docs]
def check_scenario(scenario: dict) -> dict:
"""Validate everything that does not need the model or the data.
Structure, shock resolution, methodology dispatch, declared params and the
methodology's contract. Separated so a caller can reject a malformed
scenario without paying to load a model. Raises ScenarioError with
`.problems`; returns the resolved pieces for `apply_scenario`.
"""
problems = _validate_structure(scenario)
if problems:
raise ScenarioError(problems)
shocks, shock_id, summary, warnings = _resolve_shocks(scenario)
population_shocks, shock_constants = _split_constant_shocks(shocks)
spec = _resolve_methodology(scenario, population_shocks) # MethodLookupError propagates
context_constants = _scenario_constants(scenario)
clash = sorted(set(context_constants) & set(shock_constants))
if clash:
raise ScenarioError([
f"Constants appear both as scenario context and as shocks: {clash}. Context "
"constants apply to both runs; shock constants only to the counterfactual."])
if spec is None:
if scenario.get("params"):
raise ScenarioError(["params require a methodology; a constants-only "
"scenario takes none"])
else:
problems = _validate_params(scenario, spec)
problems += _check_contract(spec, population_shocks)
if problems:
raise ScenarioError(problems)
addons, extensions = run_arguments(scenario, spec)
return {
"spec": spec,
"methodology": spec.name if spec else "constants-only",
"population_shocks": population_shocks,
"shock_constants": shock_constants,
"context_constants": context_constants,
"constants": {**context_constants, **shock_constants},
"shock_table_id": shock_id, "shocks": summary,
"addons": addons, "extensions": extensions,
"warnings": list(warnings),
"counterfactual": None, "baseline": None, "diagnostics": {},
"compatibility": None, # filled by apply_scenario, which has the system
}
[docs]
def apply_scenario(system, data, scenario: dict, *, dataset_name: str | None = None,
validate_only: bool = False) -> dict:
"""Transform input microdata according to a scenario. Executes nothing.
Returns a dict with::
counterfactual the transformed input
baseline the matching baseline (same rows) when the methodology
restructures rows, else None — run the untouched `data`
diagnostics the methodology's account of what it did
methodology the resolved reference, e.g. "lma_labour_alignment"
constants {(name, group): value} for the counterfactual run
context_constants ... applied to BOTH runs
addons, extensions what the runs must activate
compatibility what the methodology needs from the model and whether
this model has it (a CompatibilityReport), or None when
the model could not be inspected
Raises ScenarioError (with .problems) when the scenario cannot be applied.
"""
from euromod_linking.methods.base import MethodContext, MethodError
from euromod_linking.session import adopt
# Read helpers below (income-list expansion, the country's full-time week)
# resolve the model from the process session. Take it from the System the
# caller already handed us, so building the model yourself is enough.
adopt(system)
result = check_scenario(scenario)
spec = result.pop("spec")
population_shocks = result.pop("population_shocks")
shock_constants = result["shock_constants"]
cc = str(scenario["country_code"]).upper()
system_name = scenario["system_name"]
if spec is None:
# Constants-only: the input is unchanged; the two runs differ in overrides.
result["counterfactual"] = data
return result
problems = []
columns = list(data.columns)
missing = [c for c in spec.dataset_requirements if c not in columns]
if missing:
problems.append(f"Dataset lacks required columns {missing}")
method = spec.factory()
problems += method.check_dataset(columns, population_shocks)
# The model side of the contract. A missing add-on or extension switch is
# dropped by the engine without failing, so without this the scenario runs
# both simulations and only then raises NoEffectError. Checking here means
# validate_only catches it too, before anything expensive happens.
result["compatibility"] = _compatibility(system, spec)
if result["compatibility"] is not None:
problems += list(result["compatibility"].problems)
if problems:
raise ScenarioError(problems)
ctx = MethodContext(country_code=cc, system_name=system_name,
dataset_name=dataset_name,
extensions=scenario.get("extensions"))
if validate_only:
# Everything above is validation; the transform below is the expensive
# part (it runs the whole alignment), so stop here — but not before
# asking the methodology what it would target. That preview is what
# makes a mis-sized shock visible while it is still cheap to fix.
diagnostics = {"cell_population": cell_population(data, population_shocks)}
preview = getattr(method, "preview", None)
if callable(preview):
try:
diagnostics.update(preview(data, population_shocks,
scenario.get("params") or {}, ctx))
except MethodError as e:
raise ScenarioError([f"[{spec.name}] {e}"])
result["diagnostics"] = diagnostics
result["warnings"] = _merge_warnings(result["warnings"], diagnostics.get("warnings"))
return result
try:
applied = method.apply(data, population_shocks, scenario.get("params") or {}, ctx)
except MethodError as e:
raise ScenarioError([f"[{spec.name}] {e}"])
diagnostics = dict(applied.diagnostics or {})
diagnostics["cell_population"] = cell_population(data, population_shocks)
if shock_constants:
diagnostics["constant_shocks_applied"] = {
f"{n}|{g}" if g else n: v for (n, g), v in sorted(shock_constants.items())}
result["counterfactual"] = applied.data
result["baseline"] = applied.baseline
result["diagnostics"] = diagnostics
# The methodology's warnings are the ones about the shock itself ("moves
# 0.035 people", "consumes 60% of the pool"). Kept in diagnostics for
# existing readers, but surfaced at top level too — a caller scanning
# `warnings` must not miss them.
result["warnings"] = _merge_warnings(result["warnings"], diagnostics.get("warnings"))
return result
def _merge_warnings(*lists) -> list:
"""Concatenate warning lists, dropping duplicates, preserving order."""
out, seen = [], set()
for items in lists:
for w in items or []:
if w not in seen:
seen.add(w)
out.append(w)
return out
# --- convenience: apply and run -----------------------------------------------
[docs]
def run_scenario(system, scenario: dict, *, input_path: str,
paired_baseline: bool = False) -> dict:
"""Apply a scenario and run both halves through ``System.run()``.
``paired_baseline`` runs the baseline on the counterfactual's own rows, so
the two outputs are observation-paired — needed by any baseline-vs-reform
comparison that works row by row, such as one applying a fixed poverty line
or baseline-defined decile groups.
Returns the apply_scenario dict with ``baseline_output`` and
``counterfactual_output`` frames added. Raises NoEffectError when the
methodology changed people but the outputs are identical.
"""
import pandas as pd
cc = str(scenario["country_code"]).upper()
system_name = scenario["system_name"]
data_file, dataset_used, _, _ = resolve_dataset(
system, cc, scenario.get("dataset_name"), input_path)
data = pd.read_csv(data_file, sep="\t")
plan = apply_scenario(system, data, scenario, dataset_name=dataset_used)
plan["dataset_used"] = dataset_used
baseline_input = plan["baseline"] if (paired_baseline and plan["baseline"] is not None) else data
run_args = dict(country_code=cc, input_path=input_path, dataset_name=dataset_used,
addons=plan["addons"] or None, extensions=plan["extensions"] or None)
plan["baseline_output"] = execute(system, baseline_input,
constants=_as_list(plan["context_constants"]), **run_args)
plan["counterfactual_output"] = execute(system, plan["counterfactual"],
constants=_as_list(plan["constants"]), **run_args)
# Did the transform change the input at all? Ask the data, not the method:
# reading a method's diagnostics would mean knowing each method's private
# keys here, and a shape this did not recognise would silently answer "no"
# and disarm the guard.
if (not frames_identical(baseline_input, plan["counterfactual"])
and frames_identical(plan["baseline_output"], plan["counterfactual_output"])):
raise NoEffectError(
f"[{plan['methodology']}] transformed the input but the simulation output is "
"identical to the baseline: the engine did not act on it. Check that the "
f"required add-ons {plan['addons']} and switches {plan['extensions']} exist "
"for this country/system.")
return plan
def _as_list(constants: dict) -> list | None:
return [{"name": n, "group": g, "value": v}
for (n, g), v in sorted(constants.items())] or None