1"""Figures for Models B and C and the LLM laboratory (Model A figures are produced by scripts/10_model_a.py).23Reads results/model_b/*.json, results/model_c/{summary,stress_summary}.csv, results/llm_lab/cells.csv.4Writes results/figures/figB_*.pdf, figC_*.pdf, figL_*.pdf and the architecture/timing diagrams.5"""6from __future__ import annotations78import json9import sys10from pathlib import Path1112import numpy as np13import pandas as pd1415sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'src'))16import matplotlib.pyplot as plt # noqa: E40217from matplotlib.patches import FancyArrowPatch, FancyBboxPatch # noqa: E4021819from age.analysis import plotstyle as ps # noqa: E4022021ROOT = Path(__file__).resolve().parents[1]22RB, RC, RL = ROOT / 'results' / 'model_b', ROOT / 'results' / 'model_c', ROOT / 'results' / 'llm_lab'23FIG = ROOT / 'results' / 'figures'24FIG.mkdir(parents=True, exist_ok=True)25ps.use()262728# ------------------------------------------------------------------ architecture and timing diagrams29def box(ax, xy, w, h, text, fc='#f0efec', ec=ps.AXIS):30 ax.add_patch(FancyBboxPatch(xy, w, h, boxstyle='round,pad=0.02,rounding_size=0.02', fc=fc, ec=ec, lw=0.8))31 ax.text(xy[0] + w / 2, xy[1] + h / 2, text, ha='center', va='center', fontsize=7.5)323334def arrow(ax, a, b, text='', color=ps.TEXT2):35 ax.add_patch(FancyArrowPatch(a, b, arrowstyle='-|>', mutation_scale=9, color=color, lw=0.8))36 if text:37 ax.text((a[0] + b[0]) / 2, (a[1] + b[1]) / 2 + 0.02, text, fontsize=6.5, ha='center', color=color)383940fig, ax = plt.subplots(figsize=(8.0, 5.4))41ax.set_xlim(0, 1); ax.set_ylim(0, 1.02); ax.axis('off')42box(ax, (0.02, 0.66), 0.28, 0.24, 'Households $i\\in[0,1]$\nown capital $a_i$, compute shares $\\omega_i$\nchoose $c,\\ l,\\ a\',\\ d$', fc='#cde2fb')43box(ax, (0.36, 0.66), 0.28, 0.24, 'Foundation models $m=1..M$\npolicies $\\pi_m$ (no balance sheet)\nsignal $s_m=\\theta+\\eta_m$', fc='#fde3d7')44box(ax, (0.70, 0.66), 0.28, 0.24, 'Compute owners\nstock $K_Q$, rental $p_Q$\n(broad or concentrated)', fc='#d4f2e5')45box(ax, (0.02, 0.10), 0.28, 0.24, 'Firm (CRS)\ntasks $x\\leq\\iota$: compute\n$x>\\iota$: labour', fc='#fdf0cc')46box(ax, (0.36, 0.10), 0.28, 0.24, 'Compute market\n$Q^P+\\bar q D=K_Q$', fc='#f0efec')47box(ax, (0.70, 0.10), 0.28, 0.24, 'Energy sector\n$p_E$ per unit', fc='#f0efec')48def arr(a, b, text, dx=0, dy=0.03, color=ps.TEXT2):49 ax.add_patch(FancyArrowPatch(a, b, arrowstyle='-|>', mutation_scale=9, color=color, lw=0.8))50 ax.text((a[0] + b[0]) / 2 + dx, (a[1] + b[1]) / 2 + dy, text, fontsize=6.5, ha='center', color=color)51arr((0.30, 0.80), (0.36, 0.80), 'delegation $d_i$, fee $\\phi$', dy=0.025)52arr((0.64, 0.80), (0.70, 0.80), 'inference $\\bar q D$', dy=0.025)53arr((0.12, 0.66), (0.12, 0.34), 'labour $l$,\ncapital $a\'$', dx=-0.07, dy=0)54arr((0.20, 0.34), (0.20, 0.66), 'wage $w$,\nrental $R^K$', dx=0.07, dy=0)55arr((0.30, 0.22), (0.36, 0.22), 'production\ncompute $Q^P$', dy=0.04)56arr((0.84, 0.66), (0.60, 0.34), '$K_Q$ supply', dx=0.07, dy=0.0)57ax.add_patch(FancyArrowPatch((0.84, 0.90), (0.16, 0.90), arrowstyle='-|>', mutation_scale=9, color=ps.CAT[1], lw=0.8, connectionstyle='arc3,rad=0.25'))58ax.text(0.50, 0.965, 'rents $p_Q K_Q$ distributed by ownership shares $\\omega_i$', fontsize=6.5, ha='center', color=ps.CAT[1])59arr((0.64, 0.22), (0.70, 0.22), 'energy $\\bar e D$', dy=0.025)60ax.set_title('Architecture and ownership flows (Model A): every arrow is a priced flow with a payer and a payee', fontsize=8.5)61ps.save(fig, FIG / 'fig_architecture.pdf')6263fig, ax = plt.subplots(figsize=(7.4, 1.9))64ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.axis('off')65steps = ['$z_t$ realised;\nspot markets clear\n$(w,R^K,p_Q)$', 'delegation $d_{it}$\nchosen', 'signals $s_{it}$\nor $s_{mt}$ arrive',66 'savings / allocation\nby human or $\\pi_m$;\nfees paid', '$K_{t+1}=\\int a_{i,t+1}$;\nmodel noise enters', '$\\theta_t=\\varepsilon_{t+1}$\nrevealed at $t+1$']67for k, s in enumerate(steps):68 box(ax, (0.01 + k * 0.165, 0.25), 0.15, 0.55, s, fc='#f0efec' if k % 2 else '#cde2fb')69 if k < len(steps) - 1:70 arrow(ax, (0.16 + k * 0.165, 0.52), (0.175 + k * 0.165, 0.52))71ax.set_title('Within-period timing (Model A)', fontsize=8.5)72ps.save(fig, FIG / 'fig_timing.pdf')7374# ------------------------------------------------------------------ Model B75scen = {}76for f in sorted(RB.glob('*.json')):77 scen[f.stem] = json.load(open(f))78if scen:79 order = [s for s in ['no_ai', 'production_only', 'delegation_only', 'full_broad', 'full_concentrated', 'scarce_broad',80 'scarce_concentrated', 'full_broad_tight_credit', 'full_broad_cheap_fee', 'full_broad_no_mistakes'] if s in scen]81 labels = {'no_ai': 'no AI', 'production_only': 'production only', 'delegation_only': 'delegation only', 'full_broad': 'full, broad',82 'full_concentrated': 'full, concentrated', 'scarce_broad': 'scarce, broad', 'scarce_concentrated': 'scarce, concentrated',83 'full_broad_tight_credit': 'full, no borrowing', 'full_broad_cheap_fee': 'full, cheap fee', 'full_broad_no_mistakes': 'full, no mistakes'}84 # Lorenz curves85 fig, axs = plt.subplots(1, 2, figsize=(7.4, 3.3))86 for k, s in enumerate(['no_ai', 'delegation_only', 'full_broad', 'full_concentrated', 'scarce_concentrated']):87 if s in scen:88 axs[0].plot(scen[s]['lorenz_x'], scen[s]['lorenz_y'], color=ps.CAT[k], label=f"{labels[s]} (Gini {scen[s]['gini']:.2f})")89 axs[0].plot([0, 1], [0, 1], color=ps.AXIS, lw=0.6)90 axs[0].set(xlabel='population share', ylabel='wealth share', title='(a) Lorenz curves')91 axs[0].legend(fontsize=7)92 for k, s in enumerate(['delegation_only', 'full_broad', 'full_concentrated', 'scarce_broad']):93 if s in scen:94 axs[1].plot(np.arange(1, 11), scen[s]['deleg_by_decile'], marker='o', color=ps.CAT[k], label=labels[s])95 axs[1].set(xlabel='wealth decile', ylabel='delegation share within decile', title='(b) Who delegates (H1)')96 axs[1].legend(fontsize=7)97 ps.save(fig, FIG / 'figB_lorenz_delegation.pdf')98 # shares and top shares99 fig, axs = plt.subplots(1, 2, figsize=(7.4, 3.3))100 y = np.arange(len(order))101 L = [scen[s]['shares']['labor'] for s in order]102 Kc = [scen[s]['shares']['capital'] for s in order]103 Q = [scen[s]['shares']['compute'] for s in order]104 axs[0].barh(y, L, color=ps.CAT[0], label='labour', height=0.6)105 axs[0].barh(y, Kc, left=L, color=ps.CAT[1], label='capital', height=0.6)106 axs[0].barh(y, Q, left=np.array(L) + np.array(Kc), color=ps.CAT[2], label='compute', height=0.6)107 axs[0].set_yticks(y); axs[0].set_yticklabels([labels[s] for s in order], fontsize=7)108 axs[0].set(xlabel='income share', title='(a) Factor income shares'); axs[0].legend(fontsize=7, loc='lower right'); axs[0].invert_yaxis()109 axs[1].plot([scen[s]['top10'] for s in order], y, 'o', color=ps.CAT[0], label='top 10%')110 axs[1].plot([scen[s]['top1'] for s in order], y, 's', color=ps.CAT[1], label='top 1%')111 axs[1].plot([scen[s]['gini'] for s in order], y, '^', color=ps.CAT[2], label='Gini')112 axs[1].set_yticks(y); axs[1].set_yticklabels([]); axs[1].invert_yaxis()113 axs[1].set(xlabel='share / index', title='(b) Wealth concentration (H3)'); axs[1].legend(fontsize=7)114 ps.save(fig, FIG / 'figB_shares_concentration.pdf')115 # welfare by percentile relative to no_ai: consumption-equivalent from value functions (CRRA)116 if 'no_ai' in scen:117 gamma = scen['no_ai']['params'].get('gamma', 2.0)118 fig, ax = plt.subplots(figsize=(5.4, 3.2))119 pcts = [10, 25, 50, 75, 90, 99]120 for k, s in enumerate(['delegation_only', 'full_broad', 'full_concentrated', 'scarce_concentrated']):121 if s not in scen:122 continue123 ce = []124 for q in pcts:125 v1, v0 = scen[s][f'V_p{q}'], scen['no_ai'][f'V_p{q}']126 # V = sum beta^t (c^(1-g)-1)/(1-g): CE = ((v1*(1-g)*(1-beta)+1)/(v0*(1-g)*(1-beta)+1))^(1/(1-g)) - 1127 b = scen[s]['params'].get('beta', 0.96)128 ce.append(100 * (((v1 * (1 - gamma) * (1 - b) + 1) / (v0 * (1 - gamma) * (1 - b) + 1)) ** (1 / (1 - gamma)) - 1))129 ax.plot(pcts, ce, marker='o', color=ps.CAT[k], label=labels[s])130 ax.axhline(0, color=ps.AXIS, lw=0.6)131 ax.set(xlabel='wealth percentile', ylabel='consumption-equivalent gain vs no AI (%)', title='Welfare by wealth percentile (Model B)')132 ax.legend(fontsize=7)133 ps.save(fig, FIG / 'figB_welfare_percentile.pdf')134135# ------------------------------------------------------------------ Model C136if (RC / 'summary.csv').exists():137 d = pd.read_csv(RC / 'summary.csv')138 fig, axs = plt.subplots(1, 3, figsize=(9.6, 3.1))139 for k, design in enumerate(['walrasian', 'cda', 'bilateral']):140 ax = axs[k]141 for j, M in enumerate([1, 2, 4, 8]):142 g = d[(d.design == design) & (d.leverage_cap == 0) & ((d.M == M) | (d.D == 0))].sort_values('D')143 ax.plot(g.D, g.of_corr, color=ps.CAT[j], label=f'M = {M}')144 ax.fill_between(g.D, g.of_corr_lo, g.of_corr_hi, color=ps.CAT[j], alpha=0.15, lw=0)145 ax.set(xlabel='delegation share $D$', title=f'{design}', ylim=(0, 1))146 if k == 0:147 ax.set_ylabel('order-flow correlation | public info')148 ax.legend(fontsize=7)149 fig.suptitle('Common-model herding (H2): order-flow correlation rises with $D$ and falls with $M$', fontsize=9)150 ps.save(fig, FIG / 'figC_herding.pdf')151 fig, axs = plt.subplots(1, 2, figsize=(7.4, 3.2))152 for j, M in enumerate([1, 2, 4, 8]):153 g = d[(d.design == 'cda') & (d.leverage_cap == 0) & ((d.M == M) | (d.D == 0))].sort_values('D')154 axs[0].plot(g.D, g.vol * 100, color=ps.CAT[j], label=f'M = {M}')155 axs[0].fill_between(g.D, g.vol_lo * 100, g.vol_hi * 100, color=ps.CAT[j], alpha=0.15, lw=0)156 axs[1].plot(g.D, g.n_trades, color=ps.CAT[j], label=f'M = {M}')157 axs[0].set(xlabel='delegation share $D$', ylabel='return volatility (% per period)', title='(a) Volatility (CDA)')158 axs[1].set(xlabel='delegation share $D$', ylabel='trades per period', title='(b) Trading volume: no-trade with one model')159 axs[0].legend(fontsize=7)160 ps.save(fig, FIG / 'figC_volatility_volume.pdf')161 fig, axs = plt.subplots(1, 3, figsize=(9.6, 3.0))162 for k, (key, lab) in enumerate([('of_corr_deleg', 'order-flow corr. among delegators'), ('vol', 'return volatility'), ('defaults', 'defaults per run')]):163 for j, design in enumerate(['walrasian', 'cda', 'bilateral']):164 g = d[(d.design == design) & (d.leverage_cap == 0.5) & (d.M == 1) & (d.D > 0)].sort_values('D')165 axs[k].plot(g.D, g[key], color=ps.CAT[j], label=design)166 axs[k].fill_between(g.D, g[key + '_lo'], g[key + '_hi'], color=ps.CAT[j], alpha=0.15, lw=0)167 axs[k].set(xlabel='delegation share $D$', title=lab)168 axs[0].legend(fontsize=7)169 fig.suptitle('Leveraged agents (borrowing cap 0.5 × equity, margin 50%), M = 1', fontsize=9)170 ps.save(fig, FIG / 'figC_designs.pdf')171if (RC / 'stress_summary.csv').exists():172 s = pd.read_csv(RC / 'stress_summary.csv')173 s = s[~s.control.astype(bool) & (s.persist == 1)]174 fig, axs = plt.subplots(1, 2, figsize=(7.6, 3.3))175 combos = [(0.5, 0.0, 'unlevered'), (0.5, 0.95, 'levered, margin 50%'), (0.2, 0.95, 'levered, margin 20%')]176 designs = [('walrasian', 1), ('walrasian', 4), ('cda', 1), ('cda', 4)]177 w = 0.26178 for j, (margin, pre, lab) in enumerate(combos):179 vals, lo, hi, fs = [], [], [], []180 for design, M in designs:181 g = s[(s.design == design) & (s.M == M) & (s.margin == margin) & (s.preload == pre)]182 vals.append(100 * g.shock_return.iloc[0]); lo.append(100 * g.shock_return_lo.iloc[0]); hi.append(100 * g.shock_return_hi.iloc[0])183 fs.append(g.shock_forced.iloc[0])184 x = np.arange(len(designs)) + (j - 1) * w185 axs[0].bar(x, vals, w, color=ps.CAT[j], label=lab, yerr=[np.array(vals) - np.array(lo), np.array(hi) - np.array(vals)], error_kw=dict(lw=0.6))186 axs[1].bar(x, fs, w, color=ps.CAT[j], label=lab)187 for ax in axs:188 ax.set_xticks(np.arange(len(designs))); ax.set_xticklabels([f'{d}\nM={M}' for d, M in designs], fontsize=7)189 axs[0].axhline(0, color=ps.AXIS, lw=0.6)190 axs[0].set(ylabel='cumulative return over the shock window (%)', title='(a) Price response to a false shared signal')191 axs[1].set(ylabel='forced sales (shares)', title='(b) Forced sales (H4)')192 axs[0].legend(fontsize=7)193 ps.save(fig, FIG / 'figC_stress.pdf')194195# ------------------------------------------------------------------ LLM lab196if (RL / 'cells.csv').exists():197 c = pd.read_csv(RL / 'cells.csv')198 g = c.groupby(['policy', 'D', 'M', 'temperature']).agg(of=('of_corr', 'mean'), of_sd=('of_corr', 'std'), raw=('raw_of_corr', 'mean'),199 vol=('vol', 'mean'), trades=('n_trades', 'mean'), invalid=('invalid', 'mean'),200 clipped=('clipped', 'mean')).reset_index()201 pols = [p for p in ['bayesian', 'surrogate', 'llm_shared', 'llm_diverse', 'llm_shared_varied_prompts', 'llm_shared_memory'] if p in set(g.policy)]202 names = {'bayesian': 'Bayesian control', 'surrogate': 'surrogate (rule)', 'llm_shared': 'one LLM (Qwen 7B)', 'llm_diverse': 'three LLMs',203 'llm_shared_varied_prompts': 'one LLM, varied prompts', 'llm_shared_memory': 'one LLM, memory'}204 fig, axs = plt.subplots(1, 2, figsize=(7.6, 3.4))205 w = 0.25206 for j, (D, M) in enumerate([(0.0, 1), (1.0, 1), (1.0, 3)]):207 vals = [g[(g.policy == p) & (g.D == D) & (g.M == M) & (g.temperature == 0.0)].of.mean() for p in pols]208 errs = [g[(g.policy == p) & (g.D == D) & (g.M == M) & (g.temperature == 0.0)].of_sd.mean() for p in pols]209 axs[0].bar(np.arange(len(pols)) + (j - 1) * w, vals, w, yerr=errs, color=ps.CAT[j], label=['private signals', 'shared model signal', '3 model signals'][j], error_kw=dict(lw=0.6))210 axs[0].set_xticks(np.arange(len(pols))); axs[0].set_xticklabels([names[p] for p in pols], rotation=25, ha='right', fontsize=7)211 axs[0].set(ylabel='order-flow correlation | public info', title='(a) Herding by policy and signal structure')212 axs[0].legend(fontsize=7)213 for j, (D, M) in enumerate([(0.0, 1), (1.0, 1), (1.0, 3)]):214 vals = [g[(g.policy == p) & (g.D == D) & (g.M == M) & (g.temperature == 0.0)].trades.mean() for p in pols]215 axs[1].bar(np.arange(len(pols)) + (j - 1) * w, vals, w, color=ps.CAT[j])216 axs[1].set_xticks(np.arange(len(pols))); axs[1].set_xticklabels([names[p] for p in pols], rotation=25, ha='right', fontsize=7)217 axs[1].set(ylabel='trades per period', title='(b) Trading volume')218 ps.save(fig, FIG / 'figL_llm_lab.pdf')219 # inference cost frontier: runtime vs validity per policy220 fig, ax = plt.subplots(figsize=(5.2, 3.2))221 r = c.groupby('policy').agg(rt=('runtime_s', 'mean'), inv=('invalid', 'mean'), cl=('clipped', 'mean')).reset_index()222 for k, row in r.iterrows():223 ax.scatter(row.rt, row.inv, color=ps.CAT[k % 8], s=30)224 ax.annotate(names.get(row.policy, row.policy), (row.rt, row.inv), xytext=(4, 3), textcoords='offset points', fontsize=7)225 ax.set(xlabel='wall-clock seconds per run (local 4-bit inference)', ylabel='invalid actions per run', title='Inference cost frontier', xscale='symlog')226 ps.save(fig, FIG / 'figL_cost_frontier.pdf')227print('figures written to', FIG)228