-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathrun_baseline.py
More file actions
93 lines (69 loc) · 3.24 KB
/
Copy pathrun_baseline.py
File metadata and controls
93 lines (69 loc) · 3.24 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
"""Run the baseline scenario across several seeds and save the collected metrics.
Baseline = ``moderate_lockin`` with the model defaults (no a4/treatment override).
Each seed is run in its own process (the model is single-threaded pure Python),
then we persist two CSVs that the plotting scripts read:
* ``outputs/baseline/timeseries.csv`` -- long form, one row per (seed, tick) with
every per-tick DataCollector model variable. Feeds the mean +/- band plots.
* ``outputs/baseline/summary.csv`` -- one row per seed with the run-level
lifetime metrics from ``summarize_run`` (switch rates, lock-in, price, ...).
Usage:
uv run python scripts/run_baseline.py [--scenario moderate_lockin]
[--seeds 3] [--steps 1000] [--ncores 0] [--out outputs/baseline]
"""
from __future__ import annotations
import argparse
import sys
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
import pandas as pd
from spatial_market_lockin import ModelConfig, SpatialMarketModel
from spatial_market_lockin.metrics import summarize_run
from spatial_market_lockin.scenarios import scenario_overrides
def _run_one(task: tuple[str, int, int]) -> tuple[pd.DataFrame, dict]:
"""Run one seed and return (per-tick time series, lifetime summary).
Pure function of (scenario, seed, steps): safe in a process-pool worker.
"""
scenario, seed, steps = task
model = SpatialMarketModel(
ModelConfig(seed=seed, steps=steps, **scenario_overrides(scenario))
)
model.run()
ts = model.datacollector.get_model_vars_dataframe().reset_index(drop=True)
ts.insert(0, "tick", range(len(ts)))
ts.insert(0, "seed", seed)
summary = {"seed": seed, **summarize_run(model)}
return ts, summary
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--scenario", default="moderate_lockin")
parser.add_argument(
"--seeds", type=int, default=3, help="number of seeds (1..N) to average over"
)
parser.add_argument("--steps", type=int, default=1000)
parser.add_argument(
"--ncores", type=int, default=0, help="pool size; 0 = one worker per seed"
)
parser.add_argument("--out", type=Path, default=Path("outputs/baseline"))
args = parser.parse_args()
seeds = list(range(1, args.seeds + 1))
tasks = [(args.scenario, seed, args.steps) for seed in seeds]
ncores = args.ncores or len(seeds)
print(
f"Baseline run: scenario={args.scenario}, seeds={seeds}, "
f"steps={args.steps}, workers={ncores}",
flush=True,
)
with ProcessPoolExecutor(max_workers=ncores) as pool:
results = list(pool.map(_run_one, tasks))
ts_all = pd.concat([ts for ts, _ in results], ignore_index=True)
summary_all = pd.DataFrame([s for _, s in results])
args.out.mkdir(parents=True, exist_ok=True)
ts_path = args.out / "timeseries.csv"
summary_path = args.out / "summary.csv"
ts_all.to_csv(ts_path, index=False)
summary_all.to_csv(summary_path, index=False)
print(f"\nSaved {len(ts_all)} time-series rows -> {ts_path}", flush=True)
print(f"Saved {len(summary_all)} summary rows -> {summary_path}", flush=True)
if __name__ == "__main__":
main()