PolicyEngine / PolicyEngine/microcosm
Record per-target calibration estimates during optimization
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 0
- Forks
- 4
- Avg merge
- 1d 3h
- Merged PRs (30d)
- 94
Description
Problem
Microcosm currently reports aggregate optimization progress to staging: epoch, phase, iteration, loss, and sparse-selection search metadata. The solver computes the complete target-estimate vector during each epoch, but the progress callback does not expose or persist it. Microcosm writes complete per-target diagnostics only after optimization finishes.
This prevents downstream tools from showing how individual targets move during a calibration run or constructing an intermediate calibration-fit map while optimization is still running.
Goal
Define and implement a versioned, aggregate-only artifact containing per-target estimate snapshots from an in-progress calibration. The artifact should support:
- displaying the most recent intermediate calibration fit;
- tracking a target across epochs and optimization phases;
- retaining enough context to distinguish the current iterate, the best iterate retained by the optimizer, and the selected result from a sparse-selection search;
- UK and US calibration paths through shared
microcosm-calibratefunctionality.
This issue covers producing and storing the data. Calibration Diagnostics UI work belongs in its own repository.
Open design question: every epoch or bounded checkpoints
We should explicitly evaluate writing every epoch. Every-epoch snapshots would provide complete target trajectories, but copying the estimate tensor to CPU, serializing every target, and uploading the growing history may add material runtime and storage costs.
Before selecting a default cadence, benchmark representative UK and US calibrations and record:
- target count and epoch count;
- serialized bytes per snapshot and for the complete run;
- added wall-clock time and accelerator synchronization time;
- local write time and remote upload time;
- the effect of multi-phase and sparse-selection searches, where epoch numbers restart or multiple candidate solutions are evaluated.
The implementation may support an explicit every-epoch mode even if a bounded cadence becomes the production default. If every-epoch remote publication is impractical, preserve full local history when feasible and publish a bounded history plus the latest snapshot.
Proposed data contract
Each snapshot should contain at least:
- schema name and version;
- run and candidate identifiers;
- timestamp;
- optimization phase, epoch, total epochs, and relevant search iteration identifiers;
- whether the values describe the current iterate, best retained iterate, or selected result;
- a digest identifying the exact ordered list of calibration targets;
- one row per target with stable target identity, target value, current estimate, signed relative error, and any fields required to reproduce the dashboard fit calculation.
The data must remain aggregate-only: no household, person, benefit-unit, or source-record values or identifiers.
The storage design should address concurrent dashboard reads without exposing a partially written JSON document. Possible approaches include an atomically replaced latest-snapshot file plus immutable or append-only history chunks. A single ever-growing JSON file should not become the only option if it makes every update increasingly expensive.
Acceptance criteria
- Shared
microcosm-calibratecode can expose a per-target estimate snapshot without changing optimization results or determinism. - Callers can configure snapshot cadence, including an every-epoch mode if benchmarks show it is operationally feasible.
- Snapshot identity remains unambiguous across optimization phases and sparse-selection search iterations.
- A versioned schema validates target identity, finite numeric values, ordering, and the distinction between current, best, and selected estimates.
- The version 2 staging implementation can persist and upload the latest snapshot and the chosen history representation under its aggregate-data content restrictions.
- Completed runs still publish the canonical
calibration_diagnostics.json; intermediate snapshots do not replace final diagnostics. - Tests cover schema validation, atomic replacement or append behavior, cadence, multi-phase identity, and preservation of optimizer results.
- The PR records representative UK and US size and runtime measurements and documents the selected default cadence.
Context
This follows the UK staging work in #896. The current staging callback records only scalar loss progress, while the solver already has the target-estimate tensor available during each epoch.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the solver's per-epoch target-estimate tensor and the staging progress callback described in #896, then trace the shared microcosm-calibrate path and current calibration_diagnostics.json publication. Define and validate the versioned aggregate snapshot contract, including cadence, phase and search identity, and verify atomic or append-only persistence. Done means tests cover validation, cadence, multi-phase identity, storage behavior, and unchanged optimizer results, with UK and US measurements recorded.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, data
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100