"""One simulation run: build population, loop over ticks, record series, write artefacts.""" from __future__ import annotations import json import time from collections import deque from dataclasses import dataclass, field from pathlib import Path import numpy as np import pandas as pd from .. import __version__ from ..agents.base import AgentSpec, BookView, Observation, OwnOrderView from ..agents.classical import CLASSICAL_POLICIES from ..assets.fundamentals import FundamentalConfig, Fundamentals, private_signal, public_observe from ..exchange.accounts import Account from ..exchange.engine import AssetSpec, Exchange, ExchangeConfig from ..exchange.orders import Action, OrderType, Side from ..signals.news import NewsConfig, NewsFeed from ..utils.hashing import git_revision, stable_json from ..utils.rng import agent_rng, make_streams from .config import RunConfig from .population import build_population @dataclass class RunResult: config: RunConfig ticks: pd.DataFrame trades: pd.DataFrame agents: pd.DataFrame llm: pd.DataFrame events: pd.DataFrame | None manifest: dict specs: list[AgentSpec] = field(default_factory=list) def save(self, out_dir: str | Path) -> Path: out = Path(out_dir) out.mkdir(parents=True, exist_ok=True) self.config.to_yaml(out / "config.yaml") self.ticks.to_parquet(out / "ticks.parquet", index=False) self.trades.to_parquet(out / "trades.parquet", index=False) self.agents.to_parquet(out / "agents.parquet", index=False) self.llm.to_parquet(out / "llm_decisions.parquet", index=False) if self.events is not None: self.events.to_parquet(out / "events.parquet", index=False) with open(out / "manifest.json", "w") as f: json.dump(self.manifest, f, indent=2, default=str) return out class Simulation: def __init__(self, cfg: RunConfig, llm_policy=None, progress: bool = False) -> None: self.cfg = cfg self.streams = make_streams(cfg.seed) self.progress = progress K = cfg.assets.n_assets fc = cfg.fundamentals self.fund = Fundamentals(FundamentalConfig( n_assets=K, div_interval=fc.div_interval, r_period=fc.r_period, mu=tuple(fc.mu), mu_shift=fc.mu_shift, p_switch=fc.p_switch, rho=fc.rho, sigma=fc.sigma, jump_prob=fc.jump_prob, jump_sigma=fc.jump_sigma, factor_loading=tuple(fc.factor_loading), x0=fc.x0), self.streams["fundamentals"]) self.news = NewsFeed(NewsConfig(rate=cfg.news.rate, noise_sd=cfg.news.noise_sd, misleading_share=cfg.news.misleading_share, misleading_bias=cfg.news.misleading_bias), self.streams["news"], K) ex = cfg.exchange self.exchange = Exchange( [AssetSpec(asset=a, name=f"A{a}", tick_cents=cfg.assets.tick_cents, price_band_pct=cfg.assets.price_band_pct, max_order_qty=cfg.assets.max_order_qty) for a in range(K)], ExchangeConfig(maker_fee_bps=ex.maker_fee_bps, taker_fee_bps=ex.taker_fee_bps, min_latency=ex.min_latency, halt_threshold=ex.halt_threshold, halt_window=ex.halt_window, halt_duration=ex.halt_duration, halt_cancel_resting=ex.halt_cancel_resting, allow_short=ex.allow_short, max_short_lots=ex.max_short_lots, margin_check=ex.margin_check, liquidation_fraction=ex.liquidation_fraction, order_ttl=ex.order_ttl, position_limit=ex.position_limit, call_auction_interval=ex.call_auction_interval), self.streams["tiebreak"], {a: cfg.assets.initial_price_cents // cfg.assets.tick_cents for a in range(K)}) self.specs = build_population(cfg, self.streams) self.policies: dict[int, object] = {} self.rngs: dict[int, np.random.Generator] = {} for s in self.specs: acct = Account(agent_id=s.agent_id, cash=s.cash, positions=dict(s.positions), leverage_cap=s.leverage_cap, short_allowed=s.short_allowed, maintenance_margin=cfg.population.maintenance_margin) self.exchange.register(acct, s.latency) self.rngs[s.agent_id] = agent_rng(cfg.seed, s.agent_id) if s.policy_name != "llm": self.policies[s.agent_id] = CLASSICAL_POLICIES[s.policy_name](**s.policy_kwargs) self.llm_specs = [s for s in self.specs if s.policy_name == "llm"] self.llm_policy = llm_policy if self.llm_specs and self.llm_policy is None: raise ValueError("population has LLM slots but no llm_policy was supplied") self.memory: dict[int, deque[str]] = {s.agent_id: deque(maxlen=max(cfg.population.llm.memory, 1)) for s in self.llm_specs} self.last_equity: dict[int, float] = {} self.pnl_hist: dict[int, list[float]] = {s.agent_id: [] for s in self.llm_specs} self.channel: dict[int, deque[str]] = {} self.mids = np.zeros((K, cfg.n_ticks + 1)) self.rows: list[dict] = [] self.shock_log: list[dict] = [] self.mm_withdrawn = False self.n_llm_orders = 0 self.n_llm_volume = 0 # ---- observation ----------------------------------------------------------------------- def _observe(self, spec: AgentSpec, t: int, trades_window: dict[int, list]) -> Observation: ex, cfg = self.exchange, self.cfg acct = ex.ledger[spec.agent_id] books = {} for a, b in ex.books.items(): bids, asks = b.depth(5) books[a] = BookView(best_bid=b.best_bid(), best_ask=b.best_ask(), bids=bids, asks=asks, mid=ex.mid_ticks(a), last_price=ex.state[a].last_price, halted=ex.is_halted(a)) lb = cfg.output.lookback 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} sig, prec = {}, {} if spec.signal_precision > 0: for a in ex.books: sig[a] = private_signal(self.fund, self.rngs[spec.agent_id], spec.signal_precision, a) prec[a] = spec.signal_precision open_orders = [OwnOrderView(o.order_id, a, int(o.side), o.qty, o.remaining, o.price, t - o.submit_time) for a, b in ex.books.items() for o in b.resting_for(spec.agent_id)] for oid in ex.pending_by_agent.get(spec.agent_id, ()): # in-flight orders are known to their owner o = ex.orders[oid] open_orders.append(OwnOrderView(o.order_id, o.asset, int(o.side), o.qty, o.remaining, o.price, t - o.submit_time)) marks = ex.marks() pv = self.fund.public_value() return Observation( t=t, agent_id=spec.agent_id, books=books, mid_history=hist, recent_trades=trades_window, last_dividend={a: float(self.fund.last_dividend[a]) for a in ex.books}, 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}, news=self.news.recent(t, cfg.news.window), private_signals=sig, signal_precision=prec, cash=acct.cash, free_cash=acct.free_cash, positions=dict(acct.positions), equity=acct.equity(marks), open_orders=open_orders, leverage_cap=acct.leverage_cap, short_allowed=acct.short_allowed, 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), tick_cents={a: cfg.assets.tick_cents for a in ex.books}, div_interval=cfg.fundamentals.div_interval, ticks_to_dividend=cfg.fundamentals.div_interval - (t % cfg.fundamentals.div_interval), r_period=cfg.fundamentals.r_period, memory=list(self.memory.get(spec.agent_id, [])) if cfg.population.llm.memory > 0 else [], channel=list(self.channel.get(spec.cohort, [])) if (cfg.population.llm.channel and spec.cohort >= 0) else [], realized_pnl_history=self.pnl_hist.get(spec.agent_id, [])) # ---- shocks ---------------------------------------------------------------------------- def _apply_shocks(self, t: int) -> None: for s in self.cfg.shocks: if s.time != t: continue ex = self.exchange rec = {"t": t, "shock": s.kind, "size": s.size, "asset": s.asset} if s.kind == "fundamental_jump": a = s.asset or 0 self.fund.state.x[a] += s.size # latent shock, revealed only through dividends/news elif s.kind == "rate": self.fund.cfg.r_period = max(1e-4, self.fund.cfg.r_period + s.size) elif s.kind == "funding": for acct in ex.ledger.accounts.values(): if acct.leverage_cap > 1.0: acct.leverage_cap = max(1.0, acct.leverage_cap * (1.0 - s.size)) elif s.kind == "liquidity": for pol in self.policies.values(): if getattr(pol, "name", "") == "market_maker": pol.withdrawn = True self.mm_withdrawn = True elif s.kind == "liquidity_end": for pol in self.policies.values(): if getattr(pol, "name", "") == "market_maker": pol.withdrawn = False self.mm_withdrawn = False elif s.kind == "news_burst": a = s.asset or 0 from ..signals.news import NewsItem item = NewsItem(time=t, asset=a, value=s.size, credible=True, noise_sd=0.5, text=f"Breaking: next dividend expected {abs(s.size):.1f} cents {'above' if s.size > 0 else 'below'} its long-run mean.") self.news.items.append(item) public_observe(self.fund, a, item.value, item.noise_sd) elif s.kind == "placebo": pass self.shock_log.append(rec) self.exchange._log("shock", **rec) # ---- main loop --------------------------------------------------------------------------- def run(self) -> RunResult: cfg, ex = self.cfg, self.exchange t_start = time.perf_counter() K = cfg.assets.n_assets act_rng = self.streams["activation"] trades_seen = 0 pending_liq: list[tuple[int, Action]] = [] for a in range(K): self.mids[a, 0] = ex.mid_ticks(a) for t in range(cfg.n_ticks): ex.time = float(t) if t > 0 and t % cfg.fundamentals.div_interval == 0: d = self.fund.step() for a in range(K): ex.pay_dividends(a, int(round(d[a]))) self._apply_shocks(t) # after the dividend: a fundamental jump stays latent until news or the next dividend for item in self.news.step(t, self.fund.state.x): public_observe(self.fund, item.asset, item.value, item.noise_sd) # liquidation orders from the previous tick's margin check for aid, act in pending_liq: ex.submit(aid, act, float(t)) pending_liq = [] # who acts this tick u = act_rng.random(len(self.specs)) active_classical = [s for s, ui in zip(self.specs, u) if s.policy_name != "llm" and ui < s.activity] active_llm = [s for s in self.llm_specs if (t - s.decision_offset) % s.decision_interval == 0 and t >= cfg.warmup // 2] if cfg.population.llm.decision_cap is not None: active_llm = [s for s in active_llm if self._decisions_made(s.agent_id) < cfg.population.llm.decision_cap] 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)} order = list(active_classical) self.streams["shuffle"].shuffle(order) # type: ignore[arg-type] # separate stream: activations stay matched across arms for s in order: obs = self._observe(s, t, tw) action = self.policies[s.agent_id].decide(obs, self.rngs[s.agent_id]) if not action.is_hold: ex.submit(s.agent_id, action, float(t)) if active_llm: items = [] for s in active_llm: obs = self._observe(s, t, tw) eq = obs.equity if s.agent_id in self.last_equity and self.last_equity[s.agent_id] > 0: self.pnl_hist[s.agent_id].append(100.0 * (eq / self.last_equity[s.agent_id] - 1.0)) self.last_equity[s.agent_id] = eq items.append((obs, s.provider, s.model, s.prompt_id)) results = self.llm_policy.decide_batch(items) for s, (obs, _, _, _), (action, rec) in zip(active_llm, items, results): if not action.is_hold: ex.submit(s.agent_id, action, float(t)) self.n_llm_orders += len(action.orders) self._remember(s, obs, rec) ex.process_until(float(t + 1)) liq = ex.margin_check() pending_liq = list(liq) # type: ignore[arg-type] ex.end_of_tick(t) for a in range(K): self.mids[a, t + 1] = ex.mid_ticks(a) self._record_tick(t, trades_seen) trades_seen = len(ex.trades) if self.progress and t % 100 == 0: print(f" t={t} trades={len(ex.trades)} mid={ex.mid_ticks(0):.1f} V={self.fund.full_value()[0]:.1f}", flush=True) return self._finish(time.perf_counter() - t_start) def _decisions_made(self, agent_id: int) -> int: return sum(1 for r in self.llm_policy.records if r.agent_id == agent_id) def _remember(self, s: AgentSpec, obs: Observation, rec) -> None: cfg = self.cfg.population.llm if cfg.memory > 0: msg = f"tick {obs.t}: {rec.action}" if rec.action == "place": msg += f" {rec.side} {rec.quantity} {rec.order_type}" + (f" @ {rec.limit_price:.2f}" if rec.order_type == "limit" else "") if not rec.valid: msg += f" (REJECTED: {rec.reason})" msg += f"; mid was {obs.books[0].mid * self.cfg.assets.tick_cents / 100:.2f}, equity {obs.equity / 100:.2f}" self.memory[s.agent_id].append(msg) if cfg.channel and s.cohort >= 0 and rec.rationale: ch = self.channel.setdefault(s.cohort, deque(maxlen=cfg.cohort_size)) ch.append(f"trader {s.agent_id} (tick {obs.t}, {rec.action}{' ' + rec.side if rec.side else ''}): {rec.rationale}") def _record_tick(self, t: int, trades_seen: int) -> None: ex = self.exchange new_trades = ex.trades[trades_seen:] fv, pv = self.fund.full_value(), self.fund.public_value() marks = ex.marks() for a, b in ex.books.items(): bids, asks = b.depth(5) tr = [x for x in new_trades if x.asset == a] vol = sum(x.qty for x in tr) llm_ids = {s.agent_id for s in self.llm_specs} llm_vol = sum(x.qty for x in tr if x.buyer_id in llm_ids or x.seller_id in llm_ids) tc = self.cfg.assets.tick_cents self.rows.append({ "t": t, "asset": a, "mid": ex.mid_ticks(a) * tc, "best_bid": (b.best_bid() or np.nan) * tc, "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, "depth_bid5": sum(q for _, q in bids), "depth_ask5": sum(q for _, q in asks), "last_price": (ex.state[a].last_price or np.nan) * tc, "volume": vol, "llm_volume": llm_vol, "signed_volume": sum(int(x.aggressor) * x.qty for x in tr), "n_trades": len(tr), "full_value": float(fv[a]), "public_value": float(pv[a]), "x": float(self.fund.state.x[a]), "regime": int(self.fund.state.regime[a]), "halted": ex.is_halted(a), "n_halts": ex.state[a].n_halts, "eff_spread": float(np.mean([2 * abs(x.price * tc - self._mid_before(x)) for x in tr])) if tr else np.nan, "n_resting": len(b), "total_equity": sum(acc.equity(marks) for acc in ex.ledger.accounts.values()), }) def _mid_before(self, trade) -> float: # mid at the start of the trade's tick (proxy for the pre-trade quote midpoint) a, t = trade.asset, int(trade.time) return self.mids[a, t] * self.cfg.assets.tick_cents def _finish(self, wall: float) -> RunResult: cfg, ex = self.cfg, self.exchange ticks = pd.DataFrame(self.rows) tc = cfg.assets.tick_cents trades = pd.DataFrame([{"trade_id": x.trade_id, "time": x.time, "asset": x.asset, "price": x.price * tc, "qty": x.qty, "buyer": x.buyer_id, "seller": x.seller_id, "aggressor": int(x.aggressor), "maker_fee": x.maker_fee, "taker_fee": x.taker_fee} for x in ex.trades]) marks = ex.marks() rows = [] for s in self.specs: acc = ex.ledger[s.agent_id] w0 = acc.initial_cash + sum(q * cfg.assets.initial_price_cents for q in acc.initial_positions.values()) w1 = acc.equity(marks) rows.append({"agent_id": s.agent_id, "policy": s.policy_name, "replaced": s.policy_kwargs.get("replaced", ""), "provider": s.provider, "model": s.model, "prompt_id": s.prompt_id, "cohort": s.cohort, "stratum": s.stratum, "latency": s.latency, "precision": s.signal_precision, "leverage_cap": s.leverage_cap, "initial_wealth": w0, "final_wealth": w1, "pnl_pct": 100.0 * (w1 / w0 - 1.0) if w0 else np.nan, "cash": acc.cash, "position": acc.positions.get(0, 0), "fees": acc.realized_fees, "dividends": acc.dividends_received, "n_orders": acc.n_orders, "n_fills": acc.n_fills, "n_rejects": acc.n_rejects, "n_cancels": acc.n_cancels, "traded_qty": acc.traded_qty, "frozen": acc.frozen}) agents = pd.DataFrame(rows) if self.llm_policy is not None and self.llm_policy.records: from dataclasses import asdict llm = pd.DataFrame([asdict(r) for r in self.llm_policy.records]) usage = self.llm_policy.budget.usage.as_dict() else: llm = pd.DataFrame(columns=["t", "agent_id", "provider", "model", "valid"]) usage = {} events = None if cfg.output.events: events = pd.DataFrame([{"seq": e.seq, "time": e.time, "kind": e.kind, "payload": json.dumps(e.payload)} for e in ex.events]) n_llm = len(self.llm_specs) n_orders = int(agents["n_orders"].sum()) tot_vol = int(trades["qty"].sum()) if len(trades) else 0 llm_vol = int(ticks["llm_volume"].sum()) if len(ticks) else 0 manifest = { "name": cfg.name, "seed": cfg.seed, "config_hash": cfg.config_hash(), "code_version": __version__, "git_revision": git_revision(Path(__file__).resolve().parents[3]), "n_ticks": cfg.n_ticks, "n_traders": cfg.population.n_traders, "llm_share_traders": n_llm / cfg.population.n_traders, "llm_share_orders": self.n_llm_orders / max(n_orders, 1), "llm_share_volume": llm_vol / max(tot_vol, 1), "n_llm_agents": n_llm, "n_llm_decisions": int(len(llm)), "llm_valid_rate": float(llm["valid"].mean()) if len(llm) else np.nan, "llm_usage": usage, "providers": sorted({s.provider for s in self.llm_specs}), "models": sorted({s.model for s in self.llm_specs}), "prompt_version": cfg.population.llm.prompt_version, "price_assumptions": cfg.population.llm.price_assumptions, "treatment": cfg.treatment, "shocks": self.shock_log, "n_trades": len(trades), "wall_s": wall, "final_conservation_ok": True, "n_halts": int(sum(st.n_halts for st in ex.state.values())), "n_bankrupt": int(agents["frozen"].sum()), "streams": "SeedSequence(seed).spawn: " + ",".join(self.streams), } ex.ledger.check_conservation() return RunResult(config=cfg, ticks=ticks, trades=trades, agents=agents, llm=llm, events=events, manifest=manifest, specs=self.specs) def make_llm_policy(cfg: RunConfig, cache_path: str | None = None, mock_kwargs: dict | None = None, frozen_only: bool = False): from ..llm.policy import LLMPolicy, LLMPolicyConfig l = cfg.population.llm return LLMPolicy(LLMPolicyConfig( models=l.models, prompt_version=l.prompt_version, reasoning_effort=l.reasoning_effort, temperature=l.temperature, max_output_tokens=l.max_output_tokens, concurrency=l.concurrency, timeout_s=l.timeout_s, max_retries=l.max_retries, usd_cap=l.usd_cap, token_cap=l.token_cap, cache_path=cache_path, use_cache=l.cache, frozen_only=frozen_only, adaptive_prompt=l.adaptive_prompt, mu_cents=cfg.fundamentals.mu[0], taker_bps=cfg.exchange.taker_fee_bps, order_ttl=cfg.exchange.order_ttl, decision_interval=l.decision_interval, price_assumptions=l.price_assumptions, library_path=l.frozen_library, mock_kwargs={"seed": cfg.seed, **(mock_kwargs or {})})) def run_config(cfg: RunConfig, out_dir: str | Path | None = None, cache_path: str | None = None, progress: bool = False, mock_kwargs: dict | None = None) -> RunResult: llm_policy = make_llm_policy(cfg, cache_path, mock_kwargs) if cfg.population.llm_share > 0 else None res = Simulation(cfg, llm_policy=llm_policy, progress=progress).run() if out_dir is not None: res.save(out_dir) return res