1"""One simulation run: build population, loop over ticks, record series, write artefacts."""2from __future__ import annotations34import json5import time6from collections import deque7from dataclasses import dataclass, field8from pathlib import Path910import numpy as np11import pandas as pd1213from .. import __version__14from ..agents.base import AgentSpec, BookView, Observation, OwnOrderView15from ..agents.classical import CLASSICAL_POLICIES16from ..assets.fundamentals import FundamentalConfig, Fundamentals, private_signal, public_observe17from ..exchange.accounts import Account18from ..exchange.engine import AssetSpec, Exchange, ExchangeConfig19from ..exchange.orders import Action, OrderType, Side20from ..signals.news import NewsConfig, NewsFeed21from ..utils.hashing import git_revision, stable_json22from ..utils.rng import agent_rng, make_streams23from .config import RunConfig24from .population import build_population252627@dataclass28class RunResult:29 config: RunConfig30 ticks: pd.DataFrame31 trades: pd.DataFrame32 agents: pd.DataFrame33 llm: pd.DataFrame34 events: pd.DataFrame | None35 manifest: dict36 specs: list[AgentSpec] = field(default_factory=list)3738 def save(self, out_dir: str | Path) -> Path:39 out = Path(out_dir)40 out.mkdir(parents=True, exist_ok=True)41 self.config.to_yaml(out / "config.yaml")42 self.ticks.to_parquet(out / "ticks.parquet", index=False)43 self.trades.to_parquet(out / "trades.parquet", index=False)44 self.agents.to_parquet(out / "agents.parquet", index=False)45 self.llm.to_parquet(out / "llm_decisions.parquet", index=False)46 if self.events is not None:47 self.events.to_parquet(out / "events.parquet", index=False)48 with open(out / "manifest.json", "w") as f:49 json.dump(self.manifest, f, indent=2, default=str)50 return out515253class Simulation:54 def __init__(self, cfg: RunConfig, llm_policy=None, progress: bool = False) -> None:55 self.cfg = cfg56 self.streams = make_streams(cfg.seed)57 self.progress = progress58 K = cfg.assets.n_assets59 fc = cfg.fundamentals60 self.fund = Fundamentals(FundamentalConfig(61 n_assets=K, div_interval=fc.div_interval, r_period=fc.r_period, mu=tuple(fc.mu), mu_shift=fc.mu_shift,62 p_switch=fc.p_switch, rho=fc.rho, sigma=fc.sigma, jump_prob=fc.jump_prob, jump_sigma=fc.jump_sigma,63 factor_loading=tuple(fc.factor_loading), x0=fc.x0), self.streams["fundamentals"])64 self.news = NewsFeed(NewsConfig(rate=cfg.news.rate, noise_sd=cfg.news.noise_sd, misleading_share=cfg.news.misleading_share,65 misleading_bias=cfg.news.misleading_bias), self.streams["news"], K)66 ex = cfg.exchange67 self.exchange = Exchange(68 [AssetSpec(asset=a, name=f"A{a}", tick_cents=cfg.assets.tick_cents, price_band_pct=cfg.assets.price_band_pct,69 max_order_qty=cfg.assets.max_order_qty) for a in range(K)],70 ExchangeConfig(maker_fee_bps=ex.maker_fee_bps, taker_fee_bps=ex.taker_fee_bps, min_latency=ex.min_latency,71 halt_threshold=ex.halt_threshold, halt_window=ex.halt_window, halt_duration=ex.halt_duration,72 halt_cancel_resting=ex.halt_cancel_resting, allow_short=ex.allow_short, max_short_lots=ex.max_short_lots,73 margin_check=ex.margin_check, liquidation_fraction=ex.liquidation_fraction, order_ttl=ex.order_ttl,74 position_limit=ex.position_limit, call_auction_interval=ex.call_auction_interval),75 self.streams["tiebreak"], {a: cfg.assets.initial_price_cents // cfg.assets.tick_cents for a in range(K)})76 self.specs = build_population(cfg, self.streams)77 self.policies: dict[int, object] = {}78 self.rngs: dict[int, np.random.Generator] = {}79 for s in self.specs:80 acct = Account(agent_id=s.agent_id, cash=s.cash, positions=dict(s.positions), leverage_cap=s.leverage_cap,81 short_allowed=s.short_allowed, maintenance_margin=cfg.population.maintenance_margin)82 self.exchange.register(acct, s.latency)83 self.rngs[s.agent_id] = agent_rng(cfg.seed, s.agent_id)84 if s.policy_name != "llm":85 self.policies[s.agent_id] = CLASSICAL_POLICIES[s.policy_name](**s.policy_kwargs)86 self.llm_specs = [s for s in self.specs if s.policy_name == "llm"]87 self.llm_policy = llm_policy88 if self.llm_specs and self.llm_policy is None:89 raise ValueError("population has LLM slots but no llm_policy was supplied")90 self.memory: dict[int, deque[str]] = {s.agent_id: deque(maxlen=max(cfg.population.llm.memory, 1)) for s in self.llm_specs}91 self.last_equity: dict[int, float] = {}92 self.pnl_hist: dict[int, list[float]] = {s.agent_id: [] for s in self.llm_specs}93 self.channel: dict[int, deque[str]] = {}94 self.mids = np.zeros((K, cfg.n_ticks + 1))95 self.rows: list[dict] = []96 self.shock_log: list[dict] = []97 self.mm_withdrawn = False98 self.n_llm_orders = 099 self.n_llm_volume = 0100101 # ---- observation -----------------------------------------------------------------------102 def _observe(self, spec: AgentSpec, t: int, trades_window: dict[int, list]) -> Observation:103 ex, cfg = self.exchange, self.cfg104 acct = ex.ledger[spec.agent_id]105 books = {}106 for a, b in ex.books.items():107 bids, asks = b.depth(5)108 books[a] = BookView(best_bid=b.best_bid(), best_ask=b.best_ask(), bids=bids, asks=asks, mid=ex.mid_ticks(a),109 last_price=ex.state[a].last_price, halted=ex.is_halted(a))110 lb = cfg.output.lookback111 hist = {a: self.mids[a, max(0, t - lb):t + 1].copy() if t > 0 else np.array([float(ex.mid_ticks(a))]) for a in ex.books}112 sig, prec = {}, {}113 if spec.signal_precision > 0:114 for a in ex.books:115 sig[a] = private_signal(self.fund, self.rngs[spec.agent_id], spec.signal_precision, a)116 prec[a] = spec.signal_precision117 open_orders = [OwnOrderView(o.order_id, a, int(o.side), o.qty, o.remaining, o.price, t - o.submit_time)118 for a, b in ex.books.items() for o in b.resting_for(spec.agent_id)]119 for oid in ex.pending_by_agent.get(spec.agent_id, ()): # in-flight orders are known to their owner120 o = ex.orders[oid]121 open_orders.append(OwnOrderView(o.order_id, o.asset, int(o.side), o.qty, o.remaining, o.price, t - o.submit_time))122 marks = ex.marks()123 pv = self.fund.public_value()124 return Observation(125 t=t, agent_id=spec.agent_id, books=books, mid_history=hist, recent_trades=trades_window,126 last_dividend={a: float(self.fund.last_dividend[a]) for a in ex.books},127 public_value={a: float(pv[a]) for a in ex.books}, x_hat_public={a: float(self.fund.state.x_hat[a]) for a in ex.books},128 news=self.news.recent(t, cfg.news.window),129 private_signals=sig, signal_precision=prec, cash=acct.cash, free_cash=acct.free_cash,130 positions=dict(acct.positions), equity=acct.equity(marks), open_orders=open_orders,131 leverage_cap=acct.leverage_cap, short_allowed=acct.short_allowed,132 max_order_qty=min(cfg.assets.max_order_qty, cfg.population.llm.max_order_qty if spec.policy_name == "llm" else cfg.assets.max_order_qty),133 tick_cents={a: cfg.assets.tick_cents for a in ex.books}, div_interval=cfg.fundamentals.div_interval,134 ticks_to_dividend=cfg.fundamentals.div_interval - (t % cfg.fundamentals.div_interval), r_period=cfg.fundamentals.r_period,135 memory=list(self.memory.get(spec.agent_id, [])) if cfg.population.llm.memory > 0 else [],136 channel=list(self.channel.get(spec.cohort, [])) if (cfg.population.llm.channel and spec.cohort >= 0) else [],137 realized_pnl_history=self.pnl_hist.get(spec.agent_id, []))138139 # ---- shocks ----------------------------------------------------------------------------140 def _apply_shocks(self, t: int) -> None:141 for s in self.cfg.shocks:142 if s.time != t:143 continue144 ex = self.exchange145 rec = {"t": t, "shock": s.kind, "size": s.size, "asset": s.asset}146 if s.kind == "fundamental_jump":147 a = s.asset or 0148 self.fund.state.x[a] += s.size # latent shock, revealed only through dividends/news149 elif s.kind == "rate":150 self.fund.cfg.r_period = max(1e-4, self.fund.cfg.r_period + s.size)151 elif s.kind == "funding":152 for acct in ex.ledger.accounts.values():153 if acct.leverage_cap > 1.0:154 acct.leverage_cap = max(1.0, acct.leverage_cap * (1.0 - s.size))155 elif s.kind == "liquidity":156 for pol in self.policies.values():157 if getattr(pol, "name", "") == "market_maker":158 pol.withdrawn = True159 self.mm_withdrawn = True160 elif s.kind == "liquidity_end":161 for pol in self.policies.values():162 if getattr(pol, "name", "") == "market_maker":163 pol.withdrawn = False164 self.mm_withdrawn = False165 elif s.kind == "news_burst":166 a = s.asset or 0167 from ..signals.news import NewsItem168 item = NewsItem(time=t, asset=a, value=s.size, credible=True, noise_sd=0.5,169 text=f"Breaking: next dividend expected {abs(s.size):.1f} cents {'above' if s.size > 0 else 'below'} its long-run mean.")170 self.news.items.append(item)171 public_observe(self.fund, a, item.value, item.noise_sd)172 elif s.kind == "placebo":173 pass174 self.shock_log.append(rec)175 self.exchange._log("shock", **rec)176177 # ---- main loop ---------------------------------------------------------------------------178 def run(self) -> RunResult:179 cfg, ex = self.cfg, self.exchange180 t_start = time.perf_counter()181 K = cfg.assets.n_assets182 act_rng = self.streams["activation"]183 trades_seen = 0184 pending_liq: list[tuple[int, Action]] = []185 for a in range(K):186 self.mids[a, 0] = ex.mid_ticks(a)187 for t in range(cfg.n_ticks):188 ex.time = float(t)189 if t > 0 and t % cfg.fundamentals.div_interval == 0:190 d = self.fund.step()191 for a in range(K):192 ex.pay_dividends(a, int(round(d[a])))193 self._apply_shocks(t) # after the dividend: a fundamental jump stays latent until news or the next dividend194 for item in self.news.step(t, self.fund.state.x):195 public_observe(self.fund, item.asset, item.value, item.noise_sd)196 # liquidation orders from the previous tick's margin check197 for aid, act in pending_liq:198 ex.submit(aid, act, float(t))199 pending_liq = []200 # who acts this tick201 u = act_rng.random(len(self.specs))202 active_classical = [s for s, ui in zip(self.specs, u) if s.policy_name != "llm" and ui < s.activity]203 active_llm = [s for s in self.llm_specs if (t - s.decision_offset) % s.decision_interval == 0 and t >= cfg.warmup // 2]204 if cfg.population.llm.decision_cap is not None:205 active_llm = [s for s in active_llm if self._decisions_made(s.agent_id) < cfg.population.llm.decision_cap]206 tw = {a: [(tr.time, tr.price, tr.qty, int(tr.aggressor)) for tr in ex.trades[-30:] if tr.asset == a] for a in range(K)}207 order = list(active_classical)208 self.streams["shuffle"].shuffle(order) # type: ignore[arg-type] # separate stream: activations stay matched across arms209 for s in order:210 obs = self._observe(s, t, tw)211 action = self.policies[s.agent_id].decide(obs, self.rngs[s.agent_id])212 if not action.is_hold:213 ex.submit(s.agent_id, action, float(t))214 if active_llm:215 items = []216 for s in active_llm:217 obs = self._observe(s, t, tw)218 eq = obs.equity219 if s.agent_id in self.last_equity and self.last_equity[s.agent_id] > 0:220 self.pnl_hist[s.agent_id].append(100.0 * (eq / self.last_equity[s.agent_id] - 1.0))221 self.last_equity[s.agent_id] = eq222 items.append((obs, s.provider, s.model, s.prompt_id))223 results = self.llm_policy.decide_batch(items)224 for s, (obs, _, _, _), (action, rec) in zip(active_llm, items, results):225 if not action.is_hold:226 ex.submit(s.agent_id, action, float(t))227 self.n_llm_orders += len(action.orders)228 self._remember(s, obs, rec)229 ex.process_until(float(t + 1))230 liq = ex.margin_check()231 pending_liq = list(liq) # type: ignore[arg-type]232 ex.end_of_tick(t)233 for a in range(K):234 self.mids[a, t + 1] = ex.mid_ticks(a)235 self._record_tick(t, trades_seen)236 trades_seen = len(ex.trades)237 if self.progress and t % 100 == 0:238 print(f" t={t} trades={len(ex.trades)} mid={ex.mid_ticks(0):.1f} V={self.fund.full_value()[0]:.1f}", flush=True)239 return self._finish(time.perf_counter() - t_start)240241 def _decisions_made(self, agent_id: int) -> int:242 return sum(1 for r in self.llm_policy.records if r.agent_id == agent_id)243244 def _remember(self, s: AgentSpec, obs: Observation, rec) -> None:245 cfg = self.cfg.population.llm246 if cfg.memory > 0:247 msg = f"tick {obs.t}: {rec.action}"248 if rec.action == "place":249 msg += f" {rec.side} {rec.quantity} {rec.order_type}" + (f" @ {rec.limit_price:.2f}" if rec.order_type == "limit" else "")250 if not rec.valid:251 msg += f" (REJECTED: {rec.reason})"252 msg += f"; mid was {obs.books[0].mid * self.cfg.assets.tick_cents / 100:.2f}, equity {obs.equity / 100:.2f}"253 self.memory[s.agent_id].append(msg)254 if cfg.channel and s.cohort >= 0 and rec.rationale:255 ch = self.channel.setdefault(s.cohort, deque(maxlen=cfg.cohort_size))256 ch.append(f"trader {s.agent_id} (tick {obs.t}, {rec.action}{' ' + rec.side if rec.side else ''}): {rec.rationale}")257258 def _record_tick(self, t: int, trades_seen: int) -> None:259 ex = self.exchange260 new_trades = ex.trades[trades_seen:]261 fv, pv = self.fund.full_value(), self.fund.public_value()262 marks = ex.marks()263 for a, b in ex.books.items():264 bids, asks = b.depth(5)265 tr = [x for x in new_trades if x.asset == a]266 vol = sum(x.qty for x in tr)267 llm_ids = {s.agent_id for s in self.llm_specs}268 llm_vol = sum(x.qty for x in tr if x.buyer_id in llm_ids or x.seller_id in llm_ids)269 tc = self.cfg.assets.tick_cents270 self.rows.append({271 "t": t, "asset": a, "mid": ex.mid_ticks(a) * tc, "best_bid": (b.best_bid() or np.nan) * tc,272 "best_ask": (b.best_ask() or np.nan) * tc, "spread": (ex.spread_ticks(a) if ex.spread_ticks(a) is not None else np.nan) * tc,273 "depth_bid5": sum(q for _, q in bids), "depth_ask5": sum(q for _, q in asks),274 "last_price": (ex.state[a].last_price or np.nan) * tc, "volume": vol, "llm_volume": llm_vol,275 "signed_volume": sum(int(x.aggressor) * x.qty for x in tr), "n_trades": len(tr),276 "full_value": float(fv[a]), "public_value": float(pv[a]), "x": float(self.fund.state.x[a]),277 "regime": int(self.fund.state.regime[a]), "halted": ex.is_halted(a), "n_halts": ex.state[a].n_halts,278 "eff_spread": float(np.mean([2 * abs(x.price * tc - self._mid_before(x)) for x in tr])) if tr else np.nan,279 "n_resting": len(b),280 "total_equity": sum(acc.equity(marks) for acc in ex.ledger.accounts.values()),281 })282283 def _mid_before(self, trade) -> float:284 # mid at the start of the trade's tick (proxy for the pre-trade quote midpoint)285 a, t = trade.asset, int(trade.time)286 return self.mids[a, t] * self.cfg.assets.tick_cents287288 def _finish(self, wall: float) -> RunResult:289 cfg, ex = self.cfg, self.exchange290 ticks = pd.DataFrame(self.rows)291 tc = cfg.assets.tick_cents292 trades = pd.DataFrame([{"trade_id": x.trade_id, "time": x.time, "asset": x.asset, "price": x.price * tc, "qty": x.qty,293 "buyer": x.buyer_id, "seller": x.seller_id, "aggressor": int(x.aggressor),294 "maker_fee": x.maker_fee, "taker_fee": x.taker_fee} for x in ex.trades])295 marks = ex.marks()296 rows = []297 for s in self.specs:298 acc = ex.ledger[s.agent_id]299 w0 = acc.initial_cash + sum(q * cfg.assets.initial_price_cents for q in acc.initial_positions.values())300 w1 = acc.equity(marks)301 rows.append({"agent_id": s.agent_id, "policy": s.policy_name, "replaced": s.policy_kwargs.get("replaced", ""),302 "provider": s.provider, "model": s.model, "prompt_id": s.prompt_id, "cohort": s.cohort,303 "stratum": s.stratum, "latency": s.latency, "precision": s.signal_precision, "leverage_cap": s.leverage_cap,304 "initial_wealth": w0, "final_wealth": w1, "pnl_pct": 100.0 * (w1 / w0 - 1.0) if w0 else np.nan,305 "cash": acc.cash, "position": acc.positions.get(0, 0), "fees": acc.realized_fees,306 "dividends": acc.dividends_received, "n_orders": acc.n_orders, "n_fills": acc.n_fills,307 "n_rejects": acc.n_rejects, "n_cancels": acc.n_cancels, "traded_qty": acc.traded_qty, "frozen": acc.frozen})308 agents = pd.DataFrame(rows)309 if self.llm_policy is not None and self.llm_policy.records:310 from dataclasses import asdict311 llm = pd.DataFrame([asdict(r) for r in self.llm_policy.records])312 usage = self.llm_policy.budget.usage.as_dict()313 else:314 llm = pd.DataFrame(columns=["t", "agent_id", "provider", "model", "valid"])315 usage = {}316 events = None317 if cfg.output.events:318 events = pd.DataFrame([{"seq": e.seq, "time": e.time, "kind": e.kind, "payload": json.dumps(e.payload)} for e in ex.events])319 n_llm = len(self.llm_specs)320 n_orders = int(agents["n_orders"].sum())321 tot_vol = int(trades["qty"].sum()) if len(trades) else 0322 llm_vol = int(ticks["llm_volume"].sum()) if len(ticks) else 0323 manifest = {324 "name": cfg.name, "seed": cfg.seed, "config_hash": cfg.config_hash(), "code_version": __version__,325 "git_revision": git_revision(Path(__file__).resolve().parents[3]), "n_ticks": cfg.n_ticks, "n_traders": cfg.population.n_traders,326 "llm_share_traders": n_llm / cfg.population.n_traders, "llm_share_orders": self.n_llm_orders / max(n_orders, 1),327 "llm_share_volume": llm_vol / max(tot_vol, 1), "n_llm_agents": n_llm, "n_llm_decisions": int(len(llm)),328 "llm_valid_rate": float(llm["valid"].mean()) if len(llm) else np.nan, "llm_usage": usage,329 "providers": sorted({s.provider for s in self.llm_specs}), "models": sorted({s.model for s in self.llm_specs}),330 "prompt_version": cfg.population.llm.prompt_version, "price_assumptions": cfg.population.llm.price_assumptions,331 "treatment": cfg.treatment, "shocks": self.shock_log, "n_trades": len(trades), "wall_s": wall,332 "final_conservation_ok": True, "n_halts": int(sum(st.n_halts for st in ex.state.values())),333 "n_bankrupt": int(agents["frozen"].sum()), "streams": "SeedSequence(seed).spawn: " + ",".join(self.streams),334 }335 ex.ledger.check_conservation()336 return RunResult(config=cfg, ticks=ticks, trades=trades, agents=agents, llm=llm, events=events, manifest=manifest, specs=self.specs)337338339def make_llm_policy(cfg: RunConfig, cache_path: str | None = None, mock_kwargs: dict | None = None, frozen_only: bool = False):340 from ..llm.policy import LLMPolicy, LLMPolicyConfig341 l = cfg.population.llm342 return LLMPolicy(LLMPolicyConfig(343 models=l.models, prompt_version=l.prompt_version, reasoning_effort=l.reasoning_effort, temperature=l.temperature,344 max_output_tokens=l.max_output_tokens, concurrency=l.concurrency, timeout_s=l.timeout_s, max_retries=l.max_retries,345 usd_cap=l.usd_cap, token_cap=l.token_cap, cache_path=cache_path, use_cache=l.cache, frozen_only=frozen_only,346 adaptive_prompt=l.adaptive_prompt, mu_cents=cfg.fundamentals.mu[0], taker_bps=cfg.exchange.taker_fee_bps,347 order_ttl=cfg.exchange.order_ttl, decision_interval=l.decision_interval, price_assumptions=l.price_assumptions, library_path=l.frozen_library,348 mock_kwargs={"seed": cfg.seed, **(mock_kwargs or {})}))349350351def run_config(cfg: RunConfig, out_dir: str | Path | None = None, cache_path: str | None = None, progress: bool = False,352 mock_kwargs: dict | None = None) -> RunResult:353 llm_policy = make_llm_policy(cfg, cache_path, mock_kwargs) if cfg.population.llm_share > 0 else None354 res = Simulation(cfg, llm_policy=llm_policy, progress=progress).run()355 if out_dir is not None:356 res.save(out_dir)357 return res358