Add versioned computation trace export (policyengine.trace.v1)

changelog.d/versioned-trace-export.added.md

@@ -0,0 +1 @@
+Added `FullTracer.to_trace()` and `Simulation.to_trace()` for versioned computation trace export (`policyengine.trace.v1`).

policyengine_core/simulations/simulation.py

@@ -528,6 +528,52 @@ def trace(self, trace: SimpleTracer) -> None:
         else:
             self.tracer = SimpleTracer()
 
+    def to_trace(
+        self,
+        format: str = "policyengine.trace.v1",
+        *,
+        model: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Export a versioned computation trace for this simulation.
+
+        Requires ``trace=True`` (or assigning ``simulation.trace = True``)
+        so a :class:`~policyengine_core.tracers.FullTracer` is installed.
+
+        :param format: Trace document format identifier.
+        :param model: Optional country-package metadata override. When omitted,
+            a best-effort package name/version is taken from the tax-benefit
+            system module when available.
+        :returns: JSON-compatible trace document (see
+            :meth:`~policyengine_core.tracers.FullTracer.to_trace`).
+        """
+        if not isinstance(self.tracer, FullTracer):
+            raise ValueError(
+                "Computation traces require FullTracer. "
+                "Create the simulation with trace=True, or set simulation.trace = True "
+                "before calculating, then call to_trace()."
+            )
+        if model is None:
+            model = self._model_metadata_for_trace()
+        return self.tracer.to_trace(format=format, model=model)
+
+    def _model_metadata_for_trace(self) -> Optional[Dict[str, Any]]:
+        tax_benefit_system = getattr(self, "tax_benefit_system", None)
+        if tax_benefit_system is None:
+            return None
+        module = type(tax_benefit_system).__module__
+        package_name = module.split(".", 1)[0] if module else None
+        if not package_name or package_name == "policyengine_core":
+            return None
+        metadata: Dict[str, Any] = {"package": package_name.replace("_", "-")}
+        try:
+            import importlib.metadata
+
+            distribution_name = package_name.replace("_", "-")
+            metadata["version"] = importlib.metadata.version(distribution_name)
+        except Exception:
+            pass
+        return metadata
+
     def link_to_entities_instances(self) -> None:
         for _key, entity_instance in self.populations.items():
             entity_instance.simulation = self

policyengine_core/tracers/full_tracer.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import re
 import time
 import typing
-from typing import Dict, Iterator, List, Optional, Union
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 from .. import tracers
 
@@ -14,6 +15,25 @@
     Stack = List[Dict[str, Union[str, Period]]]
 
 
+TRACE_FORMAT_V1 = "policyengine.trace.v1"
+# Matches FlatTrace.key(): ``name<period, (branch)>``
+_TRACE_KEY_RE = re.compile(
+    r"^(?P<name>.+)<(?P<period>[^<>]+), \((?P<branch>[^()]*)\)>$"
+)
+
+
+def parse_trace_key(key: str) -> Dict[str, str]:
+    """Parse a serialized flat-trace node key into name/period/branch."""
+    match = _TRACE_KEY_RE.match(key)
+    if not match:
+        raise ValueError(f"Invalid PolicyEngine trace key: {key!r}")
+    return {
+        "name": match.group("name"),
+        "period": match.group("period").strip(),
+        "branch": match.group("branch").strip(),
+    }
+
+
 class FullTracer:
     _simple_tracer: tracers.SimpleTracer
     _trees: list
@@ -172,3 +192,130 @@ def _browse_node(node):
 
         for node in self._trees:
             yield from _browse_node(node)
+
+    def to_trace(
+        self,
+        format: str = TRACE_FORMAT_V1,
+        *,
+        model: Optional[Dict[str, Any]] = None,
+    ) -> Dict[str, Any]:
+        """Export a versioned, JSON-compatible computation trace.
+
+        Node identity fields (``variable`` / ``period`` / ``branch``) are taken
+        from structured :class:`~policyengine_core.tracers.TraceNode` attributes
+        via :meth:`browse_trace`, not by regex-parsing serialized flat-trace
+        keys. :func:`parse_trace_key` remains available for callers that only
+        have a flat key string (for example parameter maps without a node).
+
+        **v1 intentionally omits** entity/count on ``calculation`` roots and
+        source/version on ``parameters``. Those were listed as optional
+        "minimum useful fields" in the provenance issue and can land in a later
+        format version without breaking ``policyengine.trace.v1`` consumers.
+
+        This document is *per-computation* provenance. It is complementary to
+        PolicyEngine's TRACE Transparent Research Object (TRO) surface in
+        ``policyengine`` (JSON-LD under ``https://policyengine.org/trace/0.1#``),
+        which pins release/composition artifacts. A future TRO may embed or
+        reference a ``policyengine.trace.v1`` payload as its runtime trace; the
+        format id and TRO namespace are intentionally distinct.
+
+        Values are full serialized vectors (household-audit friendly). Optional
+        summarization for population-scale traces is deferred.
+
+        :param format: Trace document format identifier. Currently only
+            ``policyengine.trace.v1`` is supported.
+        :param model: Optional country-package / model metadata
+            (for example ``{"package": "policyengine-us", "version": "…"}``).
+        :returns: A plain ``dict`` ready for ``json.dumps``.
+        :raises ValueError: If ``format`` is not a supported version string.
+        """
+        if format != TRACE_FORMAT_V1:
+            raise ValueError(
+                f"Unsupported trace format {format!r}. "
+                f"Supported formats: {TRACE_FORMAT_V1!r}"
+            )
+
+        engine = self._engine_metadata()
+        flat_trace = self.flat_trace
+        nodes: List[Dict[str, Any]] = []
+        parameters: Dict[str, Any] = {}
+        seen_node_ids: set[str] = set()
+
+        for tree_node in self.browse_trace():
+            node_id = flat_trace.key(tree_node)
+            # Cache reads can revisit the same calculation; keep the first
+            # structured record (same non-overwriting rule as FlatTrace).
+            if node_id in seen_node_ids:
+                continue
+            seen_node_ids.add(node_id)
+
+            parameter_map = {
+                flat_trace.key(parameter): flat_trace.serialize(parameter.value)
+                for parameter in tree_node.parameters
+            }
+            node: Dict[str, Any] = {
+                "id": node_id,
+                "variable": tree_node.name,
+                "period": str(tree_node.period),
+                "branch": tree_node.branch_name,
+                "dependencies": [
+                    flat_trace.key(child) for child in tree_node.children
+                ],
+                "parameters": parameter_map,
+                "value": flat_trace.serialize(tree_node.value),
+                "calculation_time": tree_node.calculation_time(),
+                "formula_time": tree_node.formula_time(),
+            }
+            nodes.append(node)
+
+            for parameter in tree_node.parameters:
+                parameter_id = flat_trace.key(parameter)
+                if parameter_id in parameters:
+                    continue
+                parameters[parameter_id] = {
+                    "id": parameter_id,
+                    "name": parameter.name,
+                    "instant": str(parameter.period),
+                    "branch": parameter.branch_name,
+                    "value": flat_trace.serialize(parameter.value),
+                }
+
+        document: Dict[str, Any] = {
+            "format": format,
+            "engine": engine,
+            "calculation": {
+                "roots": [
+                    {
+                        "id": flat_trace.key(tree),
+                        "variable": tree.name,
+                        "period": str(tree.period),
+                        "branch": tree.branch_name,
+                    }
+                    for tree in self._trees
+                ],
+            },
+            "nodes": nodes,
+            "parameters": parameters,
+        }
+        if model is not None:
+            document["model"] = model
+        return document
+
+    @staticmethod
+    def _engine_metadata() -> Dict[str, Any]:
+        try:
+            from policyengine_core.build_metadata import get_runtime_metadata
+
+            metadata = get_runtime_metadata()
+            engine: Dict[str, Any] = {
+                "package": metadata.get("name", "policyengine-core"),
+                "version": metadata.get("version", "unknown"),
+            }
+            if metadata.get("git_sha"):
+                engine["git_sha"] = metadata["git_sha"]
+            return engine
+        except Exception:
+            return {
+                "package": "policyengine-core",
+                "version": "unknown",
+            }

tests/core/test_tracers.py

@@ -17,6 +17,7 @@
     TraceNode,
     TracingParameterNodeAtInstant,
 )
+from policyengine_core.tracers.full_tracer import parse_trace_key
 
 from .parameters_fancy_indexing.test_fancy_indexing import parameters
 
@@ -529,3 +530,150 @@ def test_browse_trace():
 
     browsed_nodes = [node.name for node in tracer.browse_trace()]
     assert browsed_nodes == ["B", "C", "D", "E", "F"]
+
+
+def test_to_trace_versioned_export():
+    tracer = FullTracer()
+    tracer.record_calculation_start("income_tax", 2017)
+    tracer.record_calculation_start("salary", 2017)
+    tracer.record_calculation_result(np.asarray([2000]))
+    tracer.record_calculation_end()
+    tracer.record_parameter_access("taxes.rate", "2017-01-01", "default", 0.15)
+    tracer.record_calculation_result(np.asarray([300]))
+    tracer.record_calculation_end()
+
+    trace = tracer.to_trace()
+
+    assert trace["format"] == "policyengine.trace.v1"
+    assert trace["engine"]["package"] == "policyengine-core"
+    assert "version" in trace["engine"]
+    assert trace["calculation"]["roots"] == [
+        {
+            "id": "income_tax<2017, (default)>",
+            "variable": "income_tax",
+            "period": "2017",
+            "branch": "default",
+        }
+    ]
+
+    nodes_by_id = {node["id"]: node for node in trace["nodes"]}
+    root = nodes_by_id["income_tax<2017, (default)>"]
+    assert root["variable"] == "income_tax"
+    assert root["period"] == "2017"
+    assert root["branch"] == "default"
+    assert root["dependencies"] == ["salary<2017, (default)>"]
+    assert root["value"] == [300]
+    assert "taxes.rate<2017-01-01, (default)>" in root["parameters"]
+    assert root["parameters"]["taxes.rate<2017-01-01, (default)>"] == approx(0.15)
+    assert "calculation_time" in root
+    assert "formula_time" in root
+
+    assert nodes_by_id["salary<2017, (default)>"]["value"] == [2000]
+    assert "taxes.rate<2017-01-01, (default)>" in trace["parameters"]
+    assert (
+        trace["parameters"]["taxes.rate<2017-01-01, (default)>"]["name"] == "taxes.rate"
+    )
+
+
+def test_to_trace_rejects_unknown_format():
+    tracer = FullTracer()
+    tracer.record_calculation_start("salary", 2017)
+    tracer.record_calculation_end()
+
+    with raises(ValueError, match="Unsupported trace format"):
+        tracer.to_trace(format="policyengine.trace.v0")
+
+
+def test_to_trace_includes_optional_model_metadata():
+    tracer = FullTracer()
+    tracer.record_calculation_start("salary", 2017)
+    tracer.record_calculation_end()
+
+    trace = tracer.to_trace(model={"package": "policyengine-us", "version": "0.0.0"})
+    assert trace["model"] == {"package": "policyengine-us", "version": "0.0.0"}
+
+
+def test_simulation_to_trace_requires_full_tracer():
+    simulation = StubSimulation()
+    simulation.tracer = SimpleTracer()
+
+    with raises(ValueError, match="FullTracer"):
+        simulation.to_trace()
+
+
+def test_simulation_to_trace_exports_full_tracer_document():
+    simulation = StubSimulation()
+    simulation.tracer = FullTracer()
+    simulation.tax_benefit_system = None
+    simulation.tracer.record_calculation_start("salary", 2017)
+    simulation.tracer.record_calculation_result(np.asarray([1000]))
+    simulation.tracer.record_calculation_end()
+
+    trace = simulation.to_trace(model={"package": "policyengine-us", "version": "1.0.0"})
+
+    assert trace["format"] == "policyengine.trace.v1"
+    assert trace["model"] == {"package": "policyengine-us", "version": "1.0.0"}
+    assert trace["calculation"]["roots"][0]["variable"] == "salary"
+    assert trace["nodes"][0]["value"] == [1000]
+
+
+def test_simulation_to_trace_infers_model_metadata_from_tax_benefit_system():
+    class FakeCountrySystem:
+        pass
+
+    FakeCountrySystem.__module__ = "policyengine_us.system"
+
+    simulation = StubSimulation()
+    simulation.tracer = FullTracer()
+    simulation.tax_benefit_system = FakeCountrySystem()
+    simulation.tracer.record_calculation_start("salary", 2017)
+    simulation.tracer.record_calculation_end()
+
+    trace = simulation.to_trace()
+
+    assert trace["model"]["package"] == "policyengine-us"
+
+
+def test_simulation_to_trace_skips_core_package_as_model_metadata():
+    class CoreLikeSystem:
+        pass
+
+    CoreLikeSystem.__module__ = "policyengine_core.country_template"
+
+    simulation = StubSimulation()
+    simulation.tracer = FullTracer()
+    simulation.tax_benefit_system = CoreLikeSystem()
+    simulation.tracer.record_calculation_start("salary", 2017)
+    simulation.tracer.record_calculation_end()
+
+    trace = simulation.to_trace()
+    assert "model" not in trace
+
+
+def test_to_trace_reads_structured_fields_from_trace_nodes():
+    """Nodes/parameters come from TraceNode attributes, not key regex parsing."""
+    tracer = FullTracer()
+    tracer.record_calculation_start("income_tax", 2017, branch_name="reform")
+    tracer.record_parameter_access("taxes.rate", "2017-01-01", "reform", 0.2)
+    tracer.record_calculation_result(np.asarray([100]))
+    tracer.record_calculation_end()
+
+    trace = tracer.to_trace()
+    root = trace["nodes"][0]
+    assert root["variable"] == "income_tax"
+    assert root["period"] == "2017"
+    assert root["branch"] == "reform"
+
+    parameter = trace["parameters"]["taxes.rate<2017-01-01, (reform)>"]
+    assert parameter["name"] == "taxes.rate"
+    assert parameter["instant"] == "2017-01-01"
+    assert parameter["branch"] == "reform"
+    assert parameter["value"] == approx(0.2)
+
+
+def test_parse_trace_key_still_available_for_flat_keys():
+    parsed = parse_trace_key("salary<2017, (default)>")
+    assert parsed == {"name": "salary", "period": "2017", "branch": "default"}
+
+    with raises(ValueError, match="Invalid PolicyEngine trace key"):
+        parse_trace_key("not-a-valid-trace-key")