PolicyEngine / PolicyEngine/microcosm

Record per-target calibration estimates during optimization

Open
#908 0 comments 0 reactions 0 assignees View on GitHub

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-calibrate functionality.

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-calibrate code 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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.