problem.py

"""Shared definitions for the sensitivity analysis (OFAT + Sobol).

The SA sweeps 5 parameters on the ``moderate_lockin`` baseline and measures the
lock-in response family (conditioned + unconditioned) plus the price response.
Everything here is imported by both the Sobol and OFAT run scripts so the
parameter space, integer handling, fixed baseline, and response list are defined
in exactly one place.

Method follows the professor's notebook (SALib problem dict + Saltelli sampling
+ ``sobol.analyze``); only the Mesa-2 ``BatchRunner`` is replaced by a direct
run of our Mesa-3 model (``run_config``), since BatchRunner was removed in Mesa 3.
"""

from __future__ import annotations

import gc
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

from spatial_market_lockin import ModelConfig, SpatialMarketModel
from spatial_market_lockin.metrics import summarize_run
from spatial_market_lockin.scenarios import scenario_overrides

# --- the SALib problem (5 swept parameters) --------------------------------
PROBLEM: dict = {
    "num_vars": 5,
    "names": [
        "a4",  # lock-in premium weight
        "loyalty_threshold_k",  # consecutive purchases to count as loyal (int)
        "loss_aversion",  # lambda
        "belief_update_weight",  # alpha (EMA step)
        "prior_price_mean",  # cold-start belief (dominant switching lever)
    ],
    "bounds": [
        [0.0, 2.0],
        [1, 10],
        [1.0, 4.0],
        [0.05, 1.0],
        [1.2, 2.2],
    ],
}

# Parameters that must be cast to int before constructing ModelConfig. Saltelli
# samples floats; we round to the nearest valid integer (as the notebook does
# for wolf_gain_from_food).
INTEGER_PARAMS = {"loyalty_threshold_k"}

# Response columns analysed for sensitivity. All come from one summarize_run, so
# adding responses costs nothing at run time. Lock-in family first (the primary
# metrics, conditioned + unconditioned), then the price response.
RESPONSES = [
    "revealed_suboptimal_rate",  # headline welfare lock-in
    "revealed_lock_in",  # 1-seller fraction
    "streak_suboptimal_rate",  # explored-then-trapped
    "switch_rate",  # unconditioned, per-life
    "switch_rate_min2_purchases",  # conditioned: >= 2 purchases
    "switch_rate_min3_purchases",  # conditioned: >= 3 purchases
    "switching_propensity",  # conditioned, lifespan-neutral
    "avg_transaction_price",  # market price level
    "price_cv",  # cross-seller price dispersion
]

# Fixed baseline: the moderate_lockin substrate (a2=0.5, a1=0.3, num_sellers=20,
# visual_range=12, psi_tick=0.25, prior=1.7, ...). Swept parameters override
# their baseline value per sample; prior_price_mean is swept, so its 1.7 here is
# only used as the OFAT hold-out value for the other parameters' sweeps.
BASELINE: dict = scenario_overrides("moderate_lockin")

# Hold-out values for OFAT: when sweeping one parameter, the other four sit here.
OFAT_BASELINE: dict = {
    "a4": 1.0,
    "loyalty_threshold_k": 3,
    "loss_aversion": 2.25,
    "belief_update_weight": 0.3,
    "prior_price_mean": 1.7,
}


def params_from_values(values) -> dict:
    """Map a sample row (sequence aligned with PROBLEM['names']) to a config
    override dict, casting integer-valued parameters."""

    out: dict = {}
    for name, value in zip(PROBLEM["names"], values):
        out[name] = int(round(value)) if name in INTEGER_PARAMS else float(value)
    return out


def run_config(param_overrides: dict, seed: int, steps: int) -> dict:
    """Run one model and return its run summary, tagged with the inputs.

    Pure function of (param_overrides, seed, steps): safe to call from a process
    pool worker. Builds on the fixed BASELINE; swept params override it.
    """

    overrides = dict(BASELINE)
    overrides.update(param_overrides)
    config = ModelConfig(seed=seed, steps=steps, sa_mode=True, **overrides)
    model = None
    try:
        model = SpatialMarketModel(config)
        model.run()
        summary = summarize_run(model)
    finally:
        del model
        gc.collect()

    # Keep only what the analysis needs (inputs + responses), so partial CSVs
    # stay small and picklable.
    row = {name: param_overrides[name] for name in PROBLEM["names"]}
    row["seed"] = seed
    for key in RESPONSES:
        row[key] = summary.get(key, float("nan"))
    return row