1"""Builds the trader population for a run: randomized assignment of policies to trader slots,2stratified by wealth, information and latency (see docs/identification.md).34Assignment is a function of (seed, config) only, so matched pairs that differ in `llm_share`5share the same slot endowments, latencies and signal precisions: the LLM treatment *replaces* the6policy in a subset of slots, chosen by a dedicated RNG stream, holding everything else fixed."""7from __future__ import annotations89import numpy as np1011from ..agents.base import AgentSpec12from .config import RunConfig131415def _stratum(wealth_mult: float, precision: float, latency: float) -> str:16 w = "rich" if wealth_mult > 1.0 else "poor"17 i = "informed" if precision > 0 else "uninformed"18 l_ = "fast" if latency < 0.3 else "slow"19 return f"{w}/{i}/{l_}"202122def build_population(cfg: RunConfig, seed_streams: dict[str, np.random.Generator]) -> list[AgentSpec]:23 pop = cfg.population24 N = pop.n_traders25 rng_slots = seed_streams["slots"] # endowments, latencies, information26 rng_assign = seed_streams["assign"] # policy assignment to slots27 rng_llm = seed_streams["llm_assign"] # which slots become LLM-controlled2829 # 1. slot characteristics (independent of policies)30 wealth = np.exp(rng_slots.standard_normal(N) * pop.endowment.wealth_dispersion)31 wealth /= wealth.mean()32 latency = rng_slots.uniform(pop.latency.min, pop.latency.max, N)33 informed_draw = rng_slots.random(N)3435 # 2. classical policy for each slot (stratified: sort slots by wealth tercile and deal types evenly)36 names = list(pop.mix)37 counts = np.floor(np.array([pop.mix[n] for n in names]) * N).astype(int)38 while counts.sum() < N:39 counts[int(np.argmax(np.array([pop.mix[n] for n in names]) * N - counts))] += 140 deck = np.concatenate([np.full(c, i) for i, c in enumerate(counts)])41 rng_assign.shuffle(deck)42 order = np.argsort(wealth) # stratify by wealth: deal the shuffled deck down the wealth ranking43 policy_idx = np.empty(N, dtype=int)44 policy_idx[order] = deck4546 # 3. LLM slots: replace a fraction of the (non market-maker) slots, stratified by policy47 n_llm = int(round(pop.llm_share * N))48 eligible = [i for i in range(N) if not (pop.llm_keep_market_makers and names[policy_idx[i]] == "market_maker"49 and pop.llm_replace != "market_maker")]50 if pop.llm_replace != "proportional":51 eligible = [i for i in eligible if names[policy_idx[i]] == pop.llm_replace]52 n_llm = min(n_llm, len(eligible))53 llm_slots: set[int] = set()54 if n_llm > 0:55 # stratified draw: proportional within each policy type, remainder at random56 by_type: dict[int, list[int]] = {}57 for i in eligible:58 by_type.setdefault(int(policy_idx[i]), []).append(i)59 quota = {k: int(np.floor(n_llm * len(v) / len(eligible))) for k, v in by_type.items()}60 for k, v in by_type.items():61 chosen = rng_llm.choice(v, size=min(quota[k], len(v)), replace=False)62 llm_slots.update(int(c) for c in chosen)63 rest = [i for i in eligible if i not in llm_slots]64 extra = n_llm - len(llm_slots)65 if extra > 0 and rest:66 llm_slots.update(int(c) for c in rng_llm.choice(rest, size=min(extra, len(rest)), replace=False))6768 # 4. LLM provider / prompt / cohort assignment69 llm = pop.llm70 providers = list(llm.provider_mix)71 pw = np.array([llm.provider_mix[p] for p in providers], dtype=float)72 pw = pw / pw.sum()73 llm_list = sorted(llm_slots)74 prov_of: dict[int, str] = {}75 if llm_list:76 # deterministic proportional allocation then shuffle77 alloc = np.floor(pw * len(llm_list)).astype(int)78 while alloc.sum() < len(llm_list):79 alloc[int(np.argmax(pw * len(llm_list) - alloc))] += 180 plist = np.concatenate([np.full(c, i) for i, c in enumerate(alloc)])81 rng_llm.shuffle(plist)82 prov_of = {s: providers[int(p)] for s, p in zip(llm_list, plist)}83 prompts = ["value", "momentum", "contrarian", "liquidity", "risk_averse", "macro", "quant", "generalist"]84 slots = max(1, min(llm.stagger_slots, llm.decision_interval))85 offsets = (rng_llm.integers(0, slots, size=max(len(llm_list), 1)) * llm.decision_interval) // slots86 cohort_of = {s: (k // max(llm.cohort_size, 1)) for k, s in enumerate(llm_list)} if llm.channel else {}8788 specs: list[AgentSpec] = []89 for i in range(N):90 pname = names[policy_idx[i]]91 kw = dict(pop.policy_kwargs.get(pname, {}))92 prec = 0.093 if pname == "fundamentalist" and informed_draw[i] < pop.informed_share:94 prec = pop.precision95 lev = pop.leverage_cap if pname == "leveraged" else 1.096 cash = int(pop.endowment.cash_cents * wealth[i])97 shares = int(round(pop.endowment.shares * wealth[i]))98 positions = {a: shares for a in range(cfg.assets.n_assets)}99 if pname == "market_maker":100 cash *= 5101 positions = {a: shares * 5 for a in range(cfg.assets.n_assets)}102 if i in llm_slots:103 prov = prov_of[i]104 k = llm_list.index(i)105 prompt_id = prompts[k % len(prompts)] if llm.prompt_set == "diverse" else "generalist"106 llm_prec = prec if llm.informed_share is None else (llm.precision if informed_draw[i] < llm.informed_share else 0.0)107 specs.append(AgentSpec(108 agent_id=i, policy_name="llm", policy_kwargs={"replaced": pname}, cash=cash, positions=positions,109 latency=float(latency[i] if pop.latency.llm_latency is None else pop.latency.llm_latency),110 activity=1.0, decision_interval=llm.decision_interval,111 decision_offset=int(offsets[k]) if llm.stagger else 0, signal_precision=llm_prec,112 leverage_cap=lev, short_allowed=False, stratum=_stratum(wealth[i], llm_prec, latency[i]),113 provider=prov, model=llm.models.get(prov, ""), prompt_id=prompt_id, cohort=cohort_of.get(i, -1)))114 else:115 specs.append(AgentSpec(116 agent_id=i, policy_name=pname, policy_kwargs=kw, cash=cash, positions=positions,117 latency=float(latency[i]), activity=pop.activity.get(pname, 0.1), signal_precision=prec,118 leverage_cap=lev, short_allowed=False, stratum=_stratum(wealth[i], prec, latency[i])))119 return specs120