augur/core: state-vector simulation refactor — plan + Phases 0–3b scaffolding - #1682
Open
agentydragon wants to merge 43 commits into
Open
augur/core: state-vector simulation refactor — plan + Phases 0–3b scaffolding#1682agentydragon wants to merge 43 commits into
agentydragon wants to merge 43 commits into
Conversation
Multi-phase plan to migrate the engine from the current ~20-matrix scratchpad-style state representation to a clean state-vector + append-only action-log shape. Captures target frame schemas, per-month order of operations, and a ~9-PR migration path.
… Phase 0)
First phase of the state-vector simulation refactor (see
augur/plans/state_vector_simulation_refactor.md). Pre-compute the
per-(rollout, month) inputs that don't depend on policy decisions into
a single long-form polars frame:
- PROPERTY_NET_CASH_FLOW (rental income minus operating expenses)
- PROPERTY_SALE_CASH_FLOW (one-shot at sale month)
- PARTNER_CONTRIBUTION_USED
- PROPERTY_TAX_ACCRUAL, HOA_ACCRUAL, INSURANCE_ACCRUAL,
MAINTENANCE_ACCRUAL
`ScheduledCashflows` is the canonical frame plus a per-kind ndarray
cache; engine reads via `amount_at(kind=..., month_position=...)` for
O(1) per-month lookup. Three of the engine's main-loop reads
(disposition.column / net_property_cash_flow / partner_equity.column)
now go through the frame; the four accrual kinds are present but the
existing `property_cost_obligation_specs` tuple still drives obligation
settlement (rewires in a later phase).
Bench: 1.06s -> 0.98s total (small improvement from the column-major
ndarray cache being slightly faster than the wrapper-column reads).
Tests: scheduled_cashflows_test plus the full augur engine suite
(test_e2e, scenario_engine_test, backend_test, property_sale_test,
annual_tax_test, scenario_set_test) all green.
Phase 1 of the state-vector simulation refactor (see augur/plans/state_vector_simulation_refactor.md). Introduce a single SimulationState object that bundles the per-(rollout, asset/account) state the engine carries across the main month loop. Phase 1 scope: cash (single checking account) + SP500 + crypto holdings. PE state, liabilities, and property state come in Phase 3. Engine now rebuilds `state` at every month boundary from the existing 1D `current_cash` / `remaining_*` locals, and the snapshot block reads the cash / SP500-units / SP500-basis / crypto-quantity / crypto-basis matrix columns FROM `state.cash(...)` / `state.holding(...)` instead of directly from the locals. State is the source for the values it carries; remaining matrix lines (SP500 value, crypto value, sale gains, PE state) still snapshot from locals until later phases bring more into the state object. Bench: 1.04s (within tolerance vs Phase 0's 0.98s). Tests: simulation_state_test plus the existing engine suite all green.
Phase 2 of the state-vector simulation refactor (see
augur/plans/state_vector_simulation_refactor.md). Introduce the
append-only cashflow-log shape that will eventually be the source of
truth for cash matrices:
- `CASHFLOW_LOG_SCHEMA` — long-form polars frame keyed by
(rollout_index, month_index, account_id, cause, amount_delta_usd).
- `CashflowCause` — discriminator enum; today covers scheduled
cash-flow kinds (PROPERTY_NET_CASH_FLOW, PROPERTY_SALE_CASH_FLOW,
PARTNER_CONTRIBUTION_USED). Future phases extend with
OBLIGATION_PAYMENT, ASSET_SALE_PROCEEDS, SPEND_POLICY, ANNUAL_TAX_
SETTLEMENT, etc. emitted from inside settlement and the policy
chain.
- `build_cashflow_log_from_scheduled(...)` — fold a
`ScheduledCashflows` frame's cash-flow kinds into the log shape.
- `derive_cash_matrix(...)` — reconstruct a `(rollouts, months)`
cash matrix from log + per-rollout initial balance via cumulative
sum.
Engine still maintains the `cash` matrix from its 1D `current_cash`
local. Phase 2 establishes the log + derivation primitives; later
phases will (a) emit log rows from settlement / policy chain / etc.,
(b) assert derived cash matches maintained cash, (c) drop the
maintained matrix in favor of the derived view.
Tests: action_log_test covers schema, fold, and derive_cash_matrix.
Existing engine suite unchanged (the new module isn't yet wired into
the engine).
Phase 3a of the state-vector simulation refactor. Add a
private_equity AssetHolding to SimulationState alongside SP500 and
crypto. Engine constructs it at both the initial and per-month
state-construction sites from the existing 1D `remaining_private_
equity_units` / `remaining_private_equity_basis` locals; the snapshot
block now reads `remaining_private_equity_units_by_month[:, month]`
and `remaining_private_equity_basis_by_month[:, month]` from
`state.holding("private_equity")`.
Same scaffold shape as Phase 1 (SP500 + crypto). The
`_PrivateEquityFundingState` helper that backs the post-loop
settlement still owns its own per-month matrices and self-propagation
logic; Phase 3b will fold that into SimulationState too.
Tests: existing engine suite all green. No behavior change.
…ies)
Key finding: property/mortgage per-(rollout, month) state is mostly
precomputed and static during the main month loop —
_amortization_arrays() builds the mortgage matrices once, the live
mask is derived once from disposition.sale_month, and the disposition
/ cash-flow / partner-equity frames are built before the loop and
read but not written. The "property-state matrix soup" is real but
the matrices aren't scratchpad; they're a precomputed schedule.
Phase 3b is therefore architectural placement, not data migration.
SimulationState gains `properties: dict[str, PropertyState]` and
`liabilities: dict[str, LiabilityState]` carrying per-month
`(rollouts,)` views of the precomputed matrices. New PropertyState
and LiabilityState dataclasses with a LiabilityKind enum
(MORTGAGE | TAX_PAYABLE) — TAX_PAYABLE is the seam for Phase 4 when
the post-loop pass collapses into the main month loop.
Engine is single-property today (scalar property_id,
mortgage_id = f"mortgage:{property_id}"); the dicts have at most one
entry. dict[id, ...] keying is forward-compatible without forcing
multi-property changes now. Read sites unchanged; matrix reads
migrate to state.liability(...).principal_usd etc. in Phase 5.
Out of scope for 3b: partner equity as a state field, multi-property
support, TAX_PAYABLE actually populated.
User feedback: partner equity is best modeled as the agent's stake
in shared properties, not as a top-level state frame. That implies
the broader restructure of SimulationState around AgentState rather
than flat cash/holdings/liabilities.
Updated shape:
SimulationState:
month_position
agents: dict[actor_id, AgentState] # per-actor state
properties: dict[property_id, PropertyState] # shared property facts
AgentState:
actor_id
cash_by_account # this agent's accounts
holdings # SP500 / crypto / PE owned
liabilities # debts owed (mortgages with property_id link)
property_stakes # ownership stake in each property (partner equity)
PropertyState: live, value, cumulative_depreciation (shared facts)
PropertyStake: ownership_pct, contribution_used, equity_ledger
LiabilityState: + property_id field linking secured liabilities
Phase 1/3a fields (cash_by_account, holdings at SimulationState
top level) move into AgentState. Single-actor scenarios still go
through state.agents[primary_owner_actor_id]. Owner-plus-partner
scenarios get two agent entries.
Mortgage handling decision: keep on borrowing agent's liabilities
dict with property_id link, not on PropertyState. Mirrors real
world: a personal debt secured by a property.
Add the five derived state frame schemas (cash_balance, asset_holding, liability, property_stake, property_state) — all long-form keyed by (rollout, month, actor?, entity-id), with the agent/account/asset/property dimensions in key columns rather than duplicated into per-entity column names. Adding a partner adds rows (actor_id="partner"), not columns; adding a second property adds rows (property_id="property_2"), not columns. Asset kinds are a discriminator in `asset_kind`, not separate columns. Update the corresponding append-only log schemas to carry `actor_id` (cashflow_log, asset_change_log, liability_log) and add a new property_stake_log section for partner-equity-style stake changes. property_state_log keeps no actor_id — properties are shared real-world objects. Add a "Today's dense matrices are projections" subsection showing how the current ScenarioRunArrays `(rollouts, months)` matrices map to filter+pivot views on the long-form frames, preserving wire compatibility without re-shaping the wire schema. Phase 2 caveat: the cashflow log scaffold in action_log.py shipped without `actor_id`. Add it (and to derive_cash_matrix's group_by) before Phase 5 wires emission from the policy chain, otherwise multi-agent scenarios collapse rows from different actors. Also rewrite the "Single working state object" section to match the agent-centric shape (SimulationState → agents/properties; AgentState → cash/holdings/liabilities/property_stakes). The previous flat description had drifted from the post-Phase 3b-design layout.
CASHFLOW_LOG_SCHEMA gains a non-null `actor_id` column between `month_index` and `account_id`. `build_cashflow_log_from_scheduled` takes `actor_id=` and `account_id=` keyword args attributing every folded row to that (actor, account) pair. `derive_cash_matrix` filters by both `actor_id` AND `account_id`. Sets up Phase 5: when log emission lands from the policy chain / settlement, multi-agent scenarios will produce rows for different actors; without actor_id keying the cum_sum would collapse them.
Restructure SimulationState around AgentState. Previous shape held
cash_by_account and holdings flat at the top of SimulationState;
new shape puts them under state.agents[actor_id].
SimulationState:
month_position
agents: dict[actor_id, AgentState]
properties: dict[property_id, PropertyState]
AgentState:
actor_id
cash_by_account -- moved from SimulationState top level
holdings -- moved from SimulationState top level
liabilities -- new (LiabilityBalance with property_id link)
property_stakes -- new (partner equity + owner stake)
New leaf types in simulation_state.py: PropertyState, PropertyStake,
LiabilityBalance (+ LiabilityKind discriminator). Named LiabilityBalance
to avoid name collision with the wire LiabilityState in accounting.py.
Engine state construction goes through _build_state(month_col=...)
closure used at both the initial-state (-1) and per-month sites.
Owner liabilities include the mortgage (with property_id link);
owner + partner stakes populate from PartnerEquityArrays columns.
Tests: all 18 augur/core tests green.
agentydragon
changed the title
Long-form per-(rollout, month, property_id) frame surfacing the
shared property facts (live, value, cumulative_depreciation) that
SimulationState.property(...) views one month at a time. No actor_id
keying — properties are world-objects shared across all agents.
Built from precomputed matrices: property_live_mask, property_value,
and disposition.column('cumulative_property_depreciation_usd').
Not a log — these are facts derived from scenario inputs +
amortization + sale event, no events accumulate. Lives in
action_log.py alongside the other long-form derived frame schemas.
Tests: schema match + post-sale row reads zero live/value with
cumulative depreciation preserved.
…se 4 foundation)
TaxActor observes per-month taxable events (recapture, capital gains,
rental income, property tax, mortgage interest, mortgage balance),
accumulates them per-year, and emits obligations at:
- Quarterly markers (Apr 15, Jun 15, Sep 15, Jan 15 of year N+1)
based on safe-harbor of prior year's actual tax. Year 0 uses
tax_profile.prior_year_tax_usd when supplied; otherwise no
quarterly obligation (the year-end true-up swallows the full
amount, same effect as the no-penalty case).
- Year-end (offset 11 within the tax year): actual year's tax
minus sum of estimated paid this year.
Year-end tax mirrors annual_sale_tax_allocation's per-year math:
salt + qualified-residence-interest deductions, federal NIIT,
California cap-gain treatment, incremental over the user's payroll-
baseline. Reuses federal_income_tax_due_usd / california_income_tax
_due_usd from annual_tax.py.
This is the actor object; wiring it into the main month loop (and
removing the post-loop _settle_required_cash_obligations sweep for
estimated_tax and annual_tax) is the next step. The actual wiring
needs careful handling of where per-month SP500/PE/property gains
get accumulated (today generic_sp500_sale_gain matrix is overwritten
in the snapshot block from policy-chain locals only — property-cost-
driven sales' gain flow needs to be traced through before observe
can read them at the right time).
Tests: quarterly-emit safe-harbor math, prior-year-fallback, Q4-routes
to-prior-year, year-end actual-tax-from-events, observed events drive
year-end obligation amount.
…ge_log User feedback: today's three parallel sale-action-record lists + per-asset-class gain matrices are all the same shape — capital gains realized at (rollout, month, asset_kind, treatment). Unify into one append-only event log; year-end tax is a group-by on tax_treatment, not a sum of separate matrices. Plan updates: - asset_change_log gains a 'tax_treatment' column (LONG_TERM_CAPITAL, SHORT_TERM_CAPITAL, DEPRECIATION_RECAPTURE_1250, null for purchases). asset_kind extended to include PROPERTY. - Property sales emit two rows (capital gain + 1250 recapture) because they have distinct tax-rate buckets. - Year-end tax computation example: filter by tax_treatment, group_by rollout, sum taxable_gain. federal/CA tax math reads bucket sums. - Documents the bug this fixes: generic_sp500_sale_gain[:, month] is overwritten from policy-chain-only locals at scenario_engine.py:1955, so property-cost-driven SP500 sales' gains are silently dropped from the year's taxable income today. - Phase 4 split into 4a (unify sale-event logs, fix the gain-overwrite bug as a side effect) + 4b (TaxActor wiring on top). Sequencing constraint: TaxActor needs the unified events frame as the single per-month gain source.
…atrix
Phase 4a foundation: define the unified capital-gain event log shape
+ derivation function. Replaces (in a future commit) today's three
parallel sale-action-record lists + per-asset-class gain matrices
that lose property-cost-driven SP500 sales' gains at the
generic_sp500_sale_gain[:, month] = sp500_sale - sp500_basis
overwrite at scenario_engine.py:1955.
Schema (one append-only frame):
ASSET_CHANGE_LOG_SCHEMA:
rollout_index, month_index, actor_id, asset_id, asset_kind,
delta_units, delta_basis_usd, cash_proceeds_usd,
taxable_gain_usd, tax_treatment, cause_kind, cause_id
AssetKindForLog: GENERIC_SP500 | CRYPTO | PRIVATE_EQUITY | PROPERTY.
Property sales will emit two rows (one LONG_TERM_CAPITAL, one
DEPRECIATION_RECAPTURE_1250) because section 1250 recapture has a
distinct 25% rate cap from the appreciation portion's LTCG rate.
TaxTreatment: LONG_TERM_CAPITAL | SHORT_TERM_CAPITAL |
DEPRECIATION_RECAPTURE_1250.
derive_per_month_taxable_gain_matrix(events, *, asset_kind=None,
tax_treatment=None, actor_id=None) groups by (rollout, month) and
sums taxable_gain_usd within optional filters. Year-end tax math
calls this with 'tax_treatment=LONG_TERM_CAPITAL' and '=...RECAPTURE'
to get the bucket sums federal/CA tax functions consume.
Tests: filter combinations (kind/treatment/actor), multi-event
accumulation per (rollout, month) — the key invariant proving the
gain-overwrite bug goes away when both code paths emit into the same
log.
This is infrastructure only — engine still uses the old paths.
Next commit wires every sale site to emit into this log and replaces
the gain-matrix overwrites with derive calls.
Phase 4a — wire the unified capital-gain event log into the engine. Builds asset_change_log from sp500_sale_action_records, crypto_sale_action_records, private_equity_sale_action_records, and disposition (LT capital gain + 1250 recapture rows) after the main loop. Then derives generic_sp500_sale_gain from the log via derive_per_month_taxable_gain_matrix(asset_kind=GENERIC_SP500, tax_treatment=LONG_TERM_CAPITAL). Removes the imperative overwrite at the end-of-month snapshot block. That overwrite only saw the within-month-policy-chain locals — property-cost obligation settlement happens earlier in the same iteration, sells SP500 via _settle_required_cash_obligation_at_month_ position, appends to sp500_sale_action_records, and updates remaining_sp500_* locals but never touches sp500_sale / sp500_basis (which reset to zero at line 1820). The derive-from-records path picks up these sales because they're recorded as action records like every other SP500 sale. Helpers in scenario_engine.py: _build_asset_change_log, _sp500_records_to_asset_change_block, _crypto_records_to_asset_change_block, _pe_records_to_asset_change_block, _property_disposition_to_asset_change_blocks. Property sale emits two rows (LONG_TERM_CAPITAL with full proceeds + recapture with zero proceeds) to prevent double-counting in sum-without-filter consumers. Tests: all 19 augur/core tests green.
…_log too Same migration as the SP500 commit one step back: PE gain matrix is now derived from the unified asset_change_log instead of being imperatively populated by two parallel paths (the per-month snapshot from a within-main-loop accumulator local, plus the post-loop tax settlement's += on pe_state.private_equity_sale_taxable_gain_usd). Removed in this commit: - Snapshot write 'private_equity_sale_taxable_gain[:, month] = private_equity_sale_taxable_gain_month' in the engine's end-of-month snapshot block. - Settlement function's 'pe_state.private_equity_sale_taxable_gain_usd[:, M] += pe_application.taxable_gain_usd' in the obligation-funding PE branch. Added: - derive_per_month_taxable_gain_matrix call after the main loop produces the gain matrix for annual_sale_tax_allocation's consumer (gets main-loop PE sale gains). - A second derive AFTER the post-loop tax settlement, since downstream consumers (the _tax_share_for_sale_action loop at :~2570) read the matrix expecting all PE sales including the post-loop-tax-driven ones. Rebuilds the asset_change_log over the now-final record list and re-derives PE. The private_equity_sale_taxable_gain_month local accumulator is left in place for now (no behavior depends on it after the snapshot removal); a follow-up can delete it. Tests: all 19 augur/core green. PE gain matrix produces the same values as before for every test scenario, because every PE sale code path already appends to private_equity_sale_action_records (which the log derives from). The change is structural: one source of truth (the log) instead of two parallel write paths.
Phase 4b prototype (inline TaxActor in main loop) revealed a fundamental ordering issue: today's quarterly estimated tax for year 0 when no prior_year_tax is supplied uses 90% of year-0's actual total tax. That relies on forward knowledge — the simulation precomputes year-0 total via post-loop annual_sale_tax_allocation, then retroactively places 0.9/4 at each year-0 quarterly marker. A true inline observer can't reach forward in time. Three options documented in the plan: - Two-pass simulation (pass 1 fills TaxActor accumulators, pass 2 emits with known totals; ~2x runtime). - Behavior change: year-0 quarterlies = 0 when no prior_year_tax, full residual settles at year-end (more honest semantically; e2e tests would need to update). - Scenario-level expected_year_zero_tax_usd attribute distinct from prior_year_tax_usd. The wired prototype matched today's bit-for-bit for years N>=1 and non-tax-related test paths; the 3 e2e tests that broke are all year-0-no-prior-year-tax scenarios. Decision on the path forward is the gating item for landing Phase 4b.
Move estimated and annual tax obligation emission + settlement from the post-loop _settle_required_cash_obligations sweep into the main month loop, routed through TaxActor. Pre-loop: construct pe_funding_state + tax_actor + 2 obligation accumulators + pre-inline-tax record-list snapshots. End of each main-loop iteration (before snapshot): observe this month's taxable events (SP500/PE gains from records-delta, property gains+recapture from disposition, rental/property-tax/mortgage from precomputed matrices), then check quarterly + year-end markers and settle inline via _settle_required_cash_obligation_at_month_position. Post-loop: emit accumulator row-blocks for the 2 tax kinds. Deleted: _quarterly_estimated_tax_obligation_due_usd, _year_end_tax_obligation_due_usd, _estimated_payments_credit_per_year_usd, plus the _SAFE_HARBOR_* helpers and constants — all replaced by TaxActor. Behavior change: year 0 with no prior_year_tax_usd no longer emits quarterly obligations (today's 90%-of-current-year-tax behavior required forward knowledge — post-loop annual_sale_tax_allocation result placed retroactively at quarterly markers, incompatible with inline observation). The year-end true-up settles the full year tax. Users wanting year-0 quarterlies set prior_year_tax_usd as a user-settable knob on TaxProfile. TaxActor: - _safe_harbor_high_agi_threshold_usd: filing-status-aware (\$150k single/MFJ/HoH, \$75k MFS). - _safe_harbor_year_total: shared helper for quarterly + annual. - annual_obligation_due credits safe-harbor-total (not just paid), so Q1+Q2+Q3+year-end+Q4 sums to actual_tax exactly. - force_year_end parameter clips year-end to last-in-horizon month when natural Dec marker falls past horizon. Engine builds asset_change_log from pre-inline-tax record-list snapshots for annual_sale_tax_allocation so TOTAL_INCOME_TAX_USD reporting matches what TaxActor sized obligations against. Tests: test_quarterly_estimated_tax_first_year_uses_90pct_of_current_year_tax renamed + rewritten to assert the new no-quarterlies-with-no-prior-tax behavior. All other tax tests pass unchanged. All 19 augur/core tests green.
After Phase 4b shipped (inline TaxActor + unified gain log), the engine still falls short of the target architecture in 13 concrete ways. Documented each as a numbered gap (G1-G13) with code-site pointers, then re-prioritized into 7 waves ordered by value-per-risk: Wave 1 — verify + finish unified gains (G8, G3, G4) Wave 2 — wire the cashflow log (G2) Wave 3 — PE scratchpad + value/sale matrix migration (G5, G12) Wave 4 — actor-ify remaining post-loop obligations (G7) Wave 5 — kill duplication + surface TAX_PAYABLE (G9, G10) Wave 6 — state-as-truth + policies-as-actions (G1, G11) Wave 7 — kill sale-record lists + final cleanup (G6, G13) Original Phase 0-6 estimate kept as historical record with landed/ deferred annotations. Total remaining: ~20-25 days / 8-10 PRs. Wave 6 is the structural shift (user's 'state_t → policy → action → state_t+1' shape); Waves 1-5 sharpen data structures without changing control flow.
Move mortgage / special_assessment / outside_rent settlements from post-loop `_settle_required_cash_obligations(...)` sweeps into the main month loop, immediately after the tax settlement block. Each obligation accumulator is constructed pre-loop with the full per-month amount-due matrix, settled one month at a time via `_settle_required_cash_obligation_at_month_position`, and its obligation + funding-decision row blocks are flushed to the event streams at the end (replacing the deleted sweep calls). Settlement order within an iteration: property-cost → tax → mortgage → special_assessment → outside_rent. This matches the previous post-loop order, preserving cash-priority behavior for cash-strapped rollouts. Today's mortgage/special/rent post-loop sweeps were already single-pass-per-month over the same monthly amounts; collapsing them into the main loop just shifts where the same work runs. Removes the second post-hoc loop the user explicitly asked to kill: the simulation is now a single forward DAG (state vector flows month-by-month forward, vectorized across rollouts) for everything except the partner-contribution settlement (deferred — it has its own per-partner state).
… 1.3) Crypto sales now flow through the same tax pipeline as SP500 / PE / property: TaxActor observes the per-month crypto gain (computed from crypto_sale_action_records the same way SP500 gain is), feeds it into the year's long_term_capital_gain, and sizes quarterly + year-end tax obligations accordingly. annual_sale_tax_allocation gains a generic_crypto_sale_gain_usd input + a crypto_sale_tax_usd output that allocates per-month tax to crypto sales the same way it does for SP500. Wire-up: - _TaxYearAccumulator: + generic_crypto_sale_gain_usd - TaxActor.observe_month: + generic_crypto_sale_gain_usd parameter - TaxActor._compute_year_actual_tax: include crypto in long_term_capital_gain - AnnualSaleTaxAllocation: + crypto_sale_tax_usd output - annual_sale_tax_allocation: + generic_crypto_sale_gain_usd input, include in sale_taxable_income, allocate tax to crypto each year - scenario_engine: compute _crypto_gain_this_month from records, pass to observe_month; derive generic_crypto_sale_gain via derive_per_month_taxable_gain_matrix(asset_kind=CRYPTO, LTCG) post-loop; thread through annual_sale_tax_allocation; _record_crypto_sale_journal_entries now takes tax_usd and writes it into the lot disposition's tax_expense_usd (matching SP500/PE) - ReportMetric: + CRYPTO_SALE_TAX_USD (LEDGER_ENTRY tax/generic_crypto_sale_tax) - Terminal aggregation: + total_crypto_sale_tax_usd - TaxPaymentAllocationDetail schema: + generic_crypto_sale_tax_usd + generic_crypto_taxable_gain_usd Addresses two of the explicit user goals in this refactor: "crypto, stock sale, private equity sale are all taxed" and "all taxing that logically flows into the same legal / tax form treatment happens once in the same code without duplication" — crypto now uses the same annual_sale_tax_allocation + TaxActor pipeline SP500/PE already use.
… use acc.obligation_kind)
…ale_tax_allocation precomputed pass-through) The post-loop annual_sale_tax_allocation today calls federal_income_tax_due_usd / california_income_tax_due_usd to compute each year's federal+CA tax — but TaxActor already computed and stored that exact value inline (at year-end markers, for obligation sizing). Two calls of the same year-tax math per year. Make TaxActor store federal and California year totals separately (year_federal_tax_usd, year_california_tax_usd) alongside the combined year_actual_tax_usd. annual_sale_tax_allocation grows a precomputed_year_tax_usd: dict[year, (federal, california)] parameter; when supplied, it skips the year-loop's bracket-walking and uses the provided totals directly. The engine passes TaxActor.year_federal_tax_usd + .year_california_tax_usd in. Result: the year-tax math runs exactly once per (scenario, year); allocation per-month is the only thing annual_sale_tax_allocation still does. Addresses the duplication called out in augur/plans/state_vector_simulation_refactor.md as G10.
…t-loop sweep (Wave 4.1, G7) Move partner-equity contribution obligations from the post-loop `_settle_partner_contribution_obligations(...)` sweep into the main month loop, immediately after the primary owner's tax / mortgage / special-assessment / outside-rent settlement block. Each agreement gets its own pre-loop `_PartnerSettlementContext` carrying the contributing actor's 1D state vectors (cash, sp500 units/basis, crypto quantity/basis), its `_ObligationFundingSources`, and its `_ObligationFundingAccumulator` keyed on the full-horizon contribution_usd matrix. The main loop's per-iteration block iterates the contexts, settles each non-zero month via `_settle_required_cash_obligation_at_month_position`, and rebinds the 1D state on the context for the next iteration. Post-loop the contexts' accumulators flush their obligation + funding-decision row blocks to the event streams (replacing the deleted sweep call). The `_settle_partner_contribution_obligations` helper is no longer called from the engine; it remains defined for now (removed in a follow-up). This drops the last post-loop `_settle_required_cash_obligations(...)` caller in the main scenario flow. The simulation is now a single forward DAG: state vector flows past → future, vectorized across rollouts, no second post-hoc loop after the main time-forward loop.
… (Wave 4.1 followup)
…p in year tax (Wave 1.1, G8) Pins the fix that landed in commit 996bbb4 (derive generic_sp500_sale_gain from the unified asset_change_log instead of overwriting [:, M] at end-of-month). Pre-fix, property-cost-driven SP500 sales' realized gain was silently dropped from TOTAL_INCOME_TAX_USD because the snapshot ran before the obligation-settlement appended its sale records. This test catches a regression by asserting that a property-cost path (mortgage funding via CheckingFloorSellPublicStockPolicy on a zero-basis SP500 position) produces non-zero year tax matching the sum of GENERIC_SP500_SALE_TAX_USD.
…G8 regression test) The Wave 4 mortgage/special_assessment/outside_rent inline migration put those settlements AFTER `TaxActor.observe_month` + the snapshot extends, which silently dropped property-cost-driven asset sales' gains from both the year-tax accumulator AND the post-loop `annual_sale_tax_allocation`. The newly-added G8 regression test caught this: a scenario where mortgage payments force SP500 sales (zero cost basis → all-gain) produced $0 year tax. Reorder one iteration to: policy steps (proactive sales) → mortgage / special / outside_rent settlement (may sell SP500/crypto/PE) → snapshot pre-tax records + TaxActor.observe_month + tax settlement → partner-contribution settlement The snapshot extends capture all sales appended UP TO the observe_month call (proactive policy + property-cost obligation driven), excluding the tax-driven sales appended during tax settlement (which must not recursively tax themselves).
…e 4.2, G7 completion) This was the post-loop sweep wrapper that iterated months internally and carried the suspicious `[:, M+1:]` forward-write deltas the user flagged as the imperative-scratchpad anti-pattern. Now that mortgage / special_assessment / outside_rent / tax / partner-contribution obligations all settle inline in the main month loop, the wrapper has no callers — and with it gone, the entire forward-slice `[:, month_position + 1 :]` write pattern is gone from the engine's obligation path. The remaining forward-slice writes are in the PE funding state (filed as G5; the PE matrix-backed state is the next target) and the property-sale equity-claim freeze (line ~6055, a one-shot post-sale broadcast, not a per-iteration scratchpad).
…ired_cash_obligations helper
…oop (Wave 3.2 partial, G12 partial)
After the unified asset_change_log derivation + the annual_sale_tax_allocation post-loop assignment, these locals carry no information: - `private_equity_sale_tax_month` / `private_equity_sale_tax` pre-loop zero matrix → reassigned wholesale from `annual_tax.private_equity_sale_tax_usd` post-loop. The `private_equity_sale_tax[:, month] = private_equity_sale_tax_month` per-iteration write was a no-op against the eventual reassignment. - `private_equity_sale_taxable_gain` pre-loop zero matrix → reassigned wholesale from `derive_per_month_taxable_gain_matrix` post-loop. - `generic_sp500_sale_tax` pre-loop zero matrix → reassigned wholesale from `annual_tax.generic_sp500_sale_tax_usd` post-loop. - `acquisition_taxable_gain_month` and `private_equity_sale_taxable_gain_month` locals: accumulated each iteration but never read. - `_PrivateEquityFundingState.private_equity_sale_taxable_gain_usd`: dataclass field stopped being mutated when the settlement chain's imperative `[:, M] += taxable_gain` was deleted (in favor of the log-derived path).
…n (reassigned post-loop from log)
The PE branch inside `_settle_required_cash_obligation_at_month_position` was the last surviving `[:, M:]` forward-write scratchpad — the imperative pattern flagged in the original critique. It was also dead code in practice: - `pe_state.remaining_units_by_month` is matrix-backed; the engine initializes it to zeros (line ~1737) and the only non-snapshot write is the settlement function's own `[:, M:] = ... - units_sold` forward-write. Settlement at iteration M reads matrix[:, M], which is zero (the prior iteration's snapshot wrote [:, M-1] from the engine's PE 1D locals, never [:, M]; and the M-1 settlement's forward-write computed `0 - 0 = 0` for [:, M] because matrix[:, M-1] was itself zero when read). - The fixed-point of this chain is units = 0 for all months, so the PE-via-CheckingFloorSellPublicStockPolicy path never actually sold PE units regardless of scenario configuration. - No test exercises PRIVATE_EQUITY in `sale_asset_preference`. Delete the whole infrastructure rather than refactor a dead path: - PE branch inside settlement (lines 5189-5308) - `_PrivateEquityFundingState` dataclass - `_apply_pe_checking_floor_obligation_funding_policy` helper - `PrivateEquityObligationFundingPolicyApplication` dataclass - `pe_funding_state` engine construction - `pe_state` parameter on settlement function - All `pe_state=pe_funding_state` call-site arguments - `PRIVATE_EQUITY` from the `sale_asset_preference` validation allow-list — `CheckingFloorSellPublicStockPolicy` now rejects PE at scenario validation rather than silently doing nothing in the settlement chain. This removes the last `[:, M+1:]` scratchpad pattern from the engine's obligation path and trims 263 lines from scenario_engine.py. PE positions still participate in the simulation via `PrivateEquitySalePolicy` (the discretionary tender / PublicMarket sale path that runs in the policy-step block) and via `Acquisition`-regime liquidity events; only the obligation-funding- chain PE path is gone, and only because it never worked.
… work The old plan was 1557 lines of accumulated narrative: the original 6-phase migration (mostly landed), a re-prioritization into 7 waves (also mostly landed), a 13-gap enumeration (8 closed), and several "additional commits after the re-prioritization" addenda. Rewrite to 409 lines: - Status block reporting the current state (single forward-only loop, no post-hoc passes, crypto/SP500/PE all taxed, year-tax math once per scenario per year). - "What landed" as bullet groups (state object scaffolding, log infra, TaxActor, inline obligation settlement, dead-code purge, value matrices derived, regression test) rather than per-commit narrative. - Remaining gaps renumbered G1-G9 (collapsing the old G1-G13 minus closed items): state-as-truth (G1+G1b), cashflow log (G2), policies-as-actions (G3), drop sale-record lists (G4), TAX_PAYABLE LiabilityBalance (G5), per-month tax allocation into TaxActor (G6), per-month accumulator matrices (G7), state matrices from logs (G8), engine ≤400 LOC (G9). - Roadmap table with effort + risk + dependency graph so the order of attack is unambiguous. - Target architecture section kept (it's the spec we're aiming at) but condensed; redundant per-frame schemas trimmed (full schemas live in augur/core/action_log.py). - Test strategy + What stays untouched + Out of scope kept, condensed. The big historical narrative (original phase list, re-prioritized wave list, the "additional commits" addenda, the historical-effort-estimate appendix) is dropped — the working branch's git log carries that record already. The plan is now a forward-looking spec, not a journal.
…per month with rollouts inside, polars as a candidate physical realization) Two clarifications after a re-read: 1. SimulationState is one object per month, not per rollout. The rollout dimension lives inside the leaves as a bulk axis; every per-month op is a vectorized bulk operation over that axis. No Python `for rollout in ...` loop anywhere in the engine. Add a "Non-goal: per-rollout objects" section spelling this out and pointing at the failure mode (someone reaching for `list[SimulationState]` indexed by rollout). 2. Don't commit the physical realization to numpy. Polars long-form frames are a candidate alternative — and likely a better fit given that the persistent logs (cashflow_log, asset_change_log) are already polars, so working state in polars eliminates the representation transition at log-emission time. Multi-agent / multi-account / multi-asset scales row-wise in polars where it gets awkward as nested dicts. Numpy stays viable for tight arithmetic kernels (mark-to-market, etc.). Concretely: - Rename "Single working state object" subsection to "Single working state object — vectorized across rollouts". - Document both physical realizations side-by-side (nested-dataclass- of-numpy vs long-form polars) with the same sell-policy example expressed in each. - Note in G1 that the leaf-storage decision is part of that PR — the scaffold today is dataclass-of-numpy, but switching to polars during G1 is on the table (and adds ~1 day to the effort). - Update the step() function docstring + non-goal section to drop numpy-specific phrasing; the vectorization invariant holds in either choice.
…epresentation User picked polars over the nested-dataclass-of-numpy alternative. The plan now treats this as a decided architectural choice, not an open Wave 1 question: - "Single working state object" section rewritten around polars long-form frames. The working-state frames at month M (cash_balance_frame, asset_holding_frame, liability_frame, property_stake_frame, property_state_frame, rollout_status_frame) share the SAME schemas as the persistent append-only logs — the working frame is the cross-section of the persistent frame at month M. - `SimulationState` becomes a typed bundle of those frames, not a nested dict of numpy arrays. Reads are polars filters/joins; writes return a new bundle with updated frames. - Sale-policy example rewritten as a single polars expression that joins cash + asset + market frames, computes `sale_usd` via `pl.when` / `pl.min_horizontal`, and outputs a row-shape ready to append to `asset_change_log`. No nested-dict ↔ row materialization. - Failed rollouts represented as rows in `rollout_status_frame`, not by structural removal. - G1 absorbs the storage migration (4 days for read-site migration to polars frames + 1 day to drop the 1D locals → 5 days total). - step() function docstring and non-goal section updated to drop the "numpy ufunc or polars expression" choice phrasing — it's polars. The vectorization invariant (no Python loop over rollouts) is unchanged; the change is which library carries the bulk dimension.
…rep) Per the consolidated plan, the working state under the state-vector refactor is one `SimulationState` per month with polars long-form frames inside (keyed by `rollout_index` + entity-id columns). G1 migrates the engine's read sites onto these frames; this PR lays the seam. Add to `simulation_state.py`: - Per-kind working-state schemas (CASH_BALANCE_FRAME_SCHEMA, ASSET_HOLDING_FRAME_SCHEMA, LIABILITY_FRAME_SCHEMA, PROPERTY_STAKE_FRAME_SCHEMA, PROPERTY_STATE_FRAME_SCHEMA). Same shape as the persistent log schemas (cashflow_log, asset_change_log, liability_log, property_stake_log, property_state_log) modulo the absent `month_index` column — the working frame is the cross-section at one month boundary. - `SimulationStateFrames` dataclass bundling the frames plus `month_position` and `rollout_count`. `from_nested(...)` builds the bundle from the existing nested-dict `SimulationState`. The builder flattens nested-dict cardinality to rows (multi-agent, multi-account, multi-asset all share schemas — added agents/assets add rows, not columns). - `cash_balance(actor_id, account_id)` returns a rollout-sorted `pl.Series`, the polars analogue of today's `(rollouts,)` numpy vector. G1 will fill out the rest of the accessor surface as read sites migrate. Tests assert per-schema correctness, empty-state behavior, multi-agent flattening, and `liability.property_id` nullability for unsecured debts. The existing nested-dict surface (`SimulationState`, `AgentState`, `AssetHolding`, ...) stays in place; both views coexist during G1 so read sites can migrate incrementally, then the nested-dict shape goes away.
…polars frames root-out The previous nested-dict scaffold (SimulationState → AgentState → AssetHolding / LiabilityBalance / PropertyStake / PropertyState dicts of (rollouts,) numpy vectors) pointed the migration the wrong way: it implied future polars frames would be built FROM the nested-dict, when in fact the engine's 1D locals + property/mortgage/partner_equity matrix columns are the actual root sources. The nested-dict view was a structural detour. Delete the nested-dict types (SimulationState, AgentState, AssetHolding, LiabilityBalance, PropertyStake, PropertyState) and the `SimulationStateFrames.from_nested(...)` classmethod that translated between them. Replace with: - Per-kind entry dataclasses (CashEntry, AssetEntry, LiabilityEntry, PropertyStakeEntry, PropertyStateEntry) holding the (id-columns, (rollouts,)-numpy) leaves directly. - `SimulationStateFrames.build(...)` classmethod that takes lists of entries and assembles the polars long-form frames. - Accessors `cash_balance(actor_id, account_id)`, `asset_units(actor_id, asset_id)`, `asset_basis(actor_id, asset_id)` returning rollout-sorted `(rollouts,)` numpy arrays. scenario_engine.py: `_build_state` → `_build_state_frames` builds the bundle directly from the engine's 1D locals + matrix columns. The end-of-month snapshot block reads through the frame accessors. No nested-dict intermediate. AssetKind and LiabilityKind enums are kept as discriminator value sources for the frame columns. simulation_state_test.py rewritten to construct frames directly via `SimulationStateFrames.build(...)`; tests assert per-schema correctness, empty-bundle behavior, multi-agent row flattening, and liability.property_id nullability for unsecured debts.
Owner
Author
|
Status update: this branch may just get thrown away. Since this PR was opened, the project pivoted to a clean rewrite in What's currently on the branch beyond what the PR description covers:
Reasonable defaults:
CI on this branch is likely red against newer devel; not rebased. https://claude.ai/code/session_01WsVZGoEeh3XbPnpeaPWrCt Generated by Claude Code |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
State-vector simulation refactor: plan + scaffolding phases. Each phase
preserves
ScenarioRunArraysbit-for-bit on existing tests; the enginestill maintains its
(rollouts, months)matrices, with the newabstractions sitting alongside as scaffolding until later phases make
them the source of truth.
augur/plans/state_vector_simulation_refactor.md) — targetarchitecture, frame schemas, per-month order of operations,
agent-centric
SimulationStateshape, five derived long-form stateframe schemas, ~9-PR migration path.
ScheduledCashflowsframe —augur/core/scheduled_cashflows.py. Pre-compute per-(rollout,month) inputs that don't depend on policy decisions into a single
long-form polars frame (property net cash flow, property sale cash
flow, partner contribution used, plus property-cost accruals).
SimulationStatescaffold —augur/core/simulation_state.py. Per-rollout state bundlemaintained in parallel with existing locals; snapshot block reads
cash / SP500 / crypto matrix columns through
state.cash(...)/state.holding(...).augur/core/action_log.py) —CASHFLOW_LOG_SCHEMA,CashflowCausediscriminator,build_cashflow_log_from_scheduled(...)fold, andderive_cash_matrix(...)reconstruction via cumulative sum.actor_idcolumn added toCASHFLOW_LOG_SCHEMAand threaded through the fold + derivation,so Phase 5 emission from the policy chain doesn't collapse rows
from different actors.
SimulationStateto PE holdings.SimulationStaterestructure.cash_by_accountandholdingsmove fromSimulationStatetoplevel into
AgentState. NewPropertyState,PropertyStake,LiabilityBalance(+LiabilityKinddiscriminator); single-property mortgage with
property_idlink; owner + partnerproperty-stake population from
PartnerEquityArrayscolumns.All 18 augur/core tests green. Bench within tolerance.
Test plan
bazelisk test //augur/core:test_e2e //augur/core:scenario_engine_test //augur/core:backend_test //augur/core:property_sale_test //augur/core:annual_tax_test //augur/core:scheduled_cashflows_test //augur/core:simulation_state_test //augur/core:action_log_test //augur/core:scenario_set_testbbr test //augur/...(CI will verify)Out of scope (followup PRs from the plan)
Phase 4 — collapse the post-loop
_settle_required_cash_obligationspass into the main month loop. Bigger surface area than the plan first
implied: requires reworking
annual_sale_tax_allocationto accumulateincrementally rather than computing once at end-of-loop (estimated_tax
and annual_tax obligations need YTD income at quarter / year-end
markers). Even the simpler post-loop kinds (special_assessment,
outside_rent) can't trivially move without changing settlement order
relative to tax obligations, which would affect cash-strapped rollouts.
Needs its own design pass.
Phase 5 — emit cashflow / asset-change / liability / stake log rows
from settlement + policy chain + scheduled-cashflow application; assert
derived matrices match maintained matrices; drop the maintained matrices
in favour of derive-from-logs views. Touches every cash/holding/liability
mutation site in the engine — ~30-50 sites across
scenario_engine.py(~6000 lines). Best done after Phase 4 lands so emission can flow through
a uniform main-loop structure rather than the current dual main-loop +
post-loop shape.
Phase 6 — cleanup: delete the per-month snapshot block, retire 1D
locals that the state object now subsumes, shrink
run_scenario_vectorizedto ≤400 lines. Depends on Phase 5.The scaffolding landing in this PR (frame schemas, agent-centric state
object, cashflow-log primitives) is exactly what those phases need to
build on; the structural rework just needs to land separately so each
phase can be reviewed against its own bench / behavioural tests.