-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathoverlay_correlation.py
More file actions
136 lines (116 loc) · 4.5 KB
/
Copy pathoverlay_correlation.py
File metadata and controls
136 lines (116 loc) · 4.5 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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
"""Overlay key per-tick metrics for one scenario and log their correlations.
Overlays (min-max normalised onto one axis, so different scales are comparable):
lock-in rate, total environmental money, transactions/tick, price dispersion
(CV), and average transaction price. A second panel shows the pairwise Pearson
correlation matrix of the *raw* series, to surface co-movement (e.g. trading
depleting money).
Usage:
uv run python scripts/overlay_correlation.py [--scenario moderate_lockin]
[--steps 1500] [--seed 42] [--burn-in 100]
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from spatial_market_lockin import ModelConfig, SpatialMarketModel
from spatial_market_lockin.scenarios import scenario_overrides
# (DataCollector column, legend label, colour)
SERIES = [
("lock_in_rate", "Instantaneous lock-in rate", "#762a83"),
("total_money", "Total environmental money", "#1a9850"),
("num_transactions", "Transactions / tick", "#2166ac"),
("price_cv", "Price dispersion (CV)", "#f1a340"),
("avg_transaction_price", "Avg transaction price", "#d6604d"),
]
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--scenario", default="moderate_lockin")
parser.add_argument("--steps", type=int, default=1500)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument(
"--burn-in",
type=int,
default=100,
help="ticks dropped before correlating (oscillation transient)",
)
args = parser.parse_args()
model = SpatialMarketModel(
ModelConfig(
seed=args.seed, steps=args.steps, **scenario_overrides(args.scenario)
)
)
model.run()
df = model.datacollector.get_model_vars_dataframe().reset_index(drop=True)
df["tick"] = range(len(df))
cols = [c for c, _, _ in SERIES]
labels = {c: lab for c, lab, _ in SERIES}
colors = {c: col for c, _, col in SERIES}
# Correlate on the post-burn-in window (drop the startup transient).
corr = df.loc[df["tick"] >= args.burn_in, cols].corr()
print(f"\nPearson correlations ({args.scenario}, ticks >= {args.burn_in}):\n")
print(corr.round(2).to_string())
fig, (ax_ts, ax_cor) = plt.subplots(
2, 1, figsize=(13, 11), gridspec_kw={"height_ratios": [2.0, 1.4]}
)
# --- top: min-max normalised overlay ---
for col, label, color in SERIES:
s = df[col].astype(float)
lo, hi = s.min(), s.max()
norm = (s - lo) / (hi - lo) if hi > lo else s * 0
ax_ts.plot(
df["tick"],
norm,
color=color,
linewidth=1.3,
alpha=0.85,
label=f"{label} [{lo:.2g}–{hi:.2g}]",
)
ax_ts.axvspan(0, args.burn_in, color="grey", alpha=0.08)
ax_ts.set_title(
f"{args.scenario}: key metrics (min-max normalised, shared axis) — "
f"{args.steps} ticks, seed={args.seed}",
fontsize=12,
fontweight="bold",
)
ax_ts.set_xlabel("Tick")
ax_ts.set_ylabel("Normalised value [0–1]")
ax_ts.legend(
fontsize=9, loc="upper right", framealpha=0.85, title="metric [raw min–max]"
)
ax_ts.grid(True, alpha=0.4)
# --- bottom: correlation heatmap ---
M = corr.to_numpy()
im = ax_cor.imshow(M, cmap="RdBu_r", vmin=-1, vmax=1)
short = [labels[c].replace(" ", "\n") for c in cols]
ax_cor.set_xticks(range(len(cols)), short, fontsize=8)
ax_cor.set_yticks(range(len(cols)), short, fontsize=8)
for i in range(len(cols)):
for j in range(len(cols)):
ax_cor.text(
j,
i,
f"{M[i, j]:.2f}",
ha="center",
va="center",
color="white" if abs(M[i, j]) > 0.55 else "black",
fontsize=9,
fontweight="bold",
)
ax_cor.set_title(
f"Pearson correlation (ticks ≥ {args.burn_in})", fontsize=11, fontweight="bold"
)
fig.colorbar(im, ax=ax_cor, fraction=0.046, pad=0.04)
fig.tight_layout()
images_dir = Path(__file__).resolve().parent.parent / "images"
images_dir.mkdir(exist_ok=True)
out = images_dir / f"overlay_correlation_{args.scenario}.png"
fig.savefig(out, dpi=130, bbox_inches="tight")
plt.close(fig)
print(f"\nSaved figure to {out}")
if __name__ == "__main__":
main()