"""Standardized model export format (`ModelPackage`).
This is the versioned interchange format for a compiled `prophys`
model. It carries no dependency on any downstream consumer;
consumers implement their own import logic against this documented shape.
Format 1.1 adds two things to 1.0:
- a :class:`~prophys.export.manifest.RunManifest`, recording which prophys
version, model structure, parameter values, seed and Monte-Carlo settings
produced the exported distributions, so a result can be traced and
re-derived rather than only read;
- per-attribute `unit_description` and optional `diagnostics`, so a reader
learns what the numbers are measured in and how well the model that
produced them was doing.
A 1.0 file loads unchanged: the new fields are optional and default to
absent.
"""
from __future__ import annotations
import dataclasses
import json
from typing import Any
import jax
import numpy as np
from .. import _ffi, units
from .manifest import (
ReproducibilityError,
RunManifest,
build_manifest,
manifest_from_dict,
)
from .manifest import verify as verify_manifest
FORMAT_VERSION = "1.1"
#: Format versions this module can load. A file written by a newer prophys
#: is refused with its version named, rather than being read with fields
#: silently missing.
SUPPORTED_VERSIONS = frozenset({"1.0", "1.1"})
[docs]
@dataclasses.dataclass
class AttributeExport:
name: str
unit: str
domain: str
representation: str # "quantized" | "histogram" | "sampler"
support: tuple[float, float]
# "quantized": grid (list[float]) + probs (list[float])
grid: list[float] | None = None
probs: list[float] | None = None
# "histogram": edges (n+1,) + probs (n,)
edges: list[float] | None = None
# "sampler": raw NumPy samples, JAX-free and picklable
samples: list[float] | None = None
mean: float | None = None
# -- format 1.1 --------------------------------------------------------
unit_description: str | None = None
"""The attribute's unit expanded into its dimension and SI scale (see
`prophys.units.describe`), so a consumer can check compatibility without
reimplementing this package's unit parser."""
std_error: float | None = None
"""Monte-Carlo standard error of `mean`, from the same draw the
distribution was built from. `None` when the export did not compute
it."""
n_samples: int | None = None
"""Draws the exported representation was built from."""
[docs]
@dataclasses.dataclass
class ModelPackage:
format_version: str
model_name: str
params: dict[str, Any]
attributes: list[AttributeExport]
correlation: dict[str, Any] | None = None
# -- format 1.1 --------------------------------------------------------
manifest: RunManifest | None = None
"""Provenance for this export: versions, structure hash, seed, licensing.
`None` only for a package loaded from a 1.0 file."""
diagnostics: dict[str, Any] | None = None
"""Diagnostic reports attached to the export, keyed by name — typically
the output of :func:`prophys.engine.scoring.report_card`."""
def to_dict(self) -> dict[str, Any]:
data = dataclasses.asdict(self)
if self.manifest is not None:
data["manifest"] = self.manifest.to_dict()
return data
[docs]
def verify(self, compiled, *, check_params: bool = True) -> None:
"""Check that `compiled` is the model this package came from.
Raises :class:`~prophys.export.manifest.ReproducibilityError` when
the structure hash or the parameters differ; see
:func:`prophys.export.manifest.verify`.
"""
if self.manifest is None:
raise ReproducibilityError(
"This package carries no manifest (it was written in format "
"1.0), so there is nothing to verify against."
)
verify_manifest(self.manifest, compiled, check_params=check_params)
def attribute(self, name: str) -> AttributeExport:
for attr in self.attributes:
if attr.name == name:
return attr
raise KeyError(
f"No attribute {name!r} in this package; it has "
f"{[a.name for a in self.attributes]}"
)
def to_json(self, indent: int = 2) -> str:
return json.dumps(self.to_dict(), indent=indent)
def save(self, path: str) -> None:
with open(path, "w") as f:
f.write(self.to_json())
@classmethod
def load(cls, path: str) -> "ModelPackage":
with open(path) as f:
data = json.load(f)
version = data.get("format_version")
if version not in SUPPORTED_VERSIONS:
raise ValueError(
f"{path} is format version {version!r}; this prophys "
f"({FORMAT_VERSION}) reads {sorted(SUPPORTED_VERSIONS)}."
)
# Unknown per-attribute fields are dropped rather than raising, so a
# file written by a later prophys within a supported version still
# loads with the fields this one understands.
known = {f.name for f in dataclasses.fields(AttributeExport)}
attrs = [
AttributeExport(**{k: v for k, v in a.items() if k in known})
for a in data["attributes"]
]
data = dict(data)
if data.get("manifest"):
data["manifest"] = manifest_from_dict(data["manifest"])
data["attributes"] = attrs
return cls(**data)
[docs]
def export_model(
compiled,
attribute_names: list[str] | None = None,
representation: str = "quantized",
n_qubits: int = 8,
n_samples: int = 4096,
seed: int = 0,
params: Any = None,
observations: Any = None,
diagnostics: Any = None,
notes: str = "",
) -> ModelPackage:
"""Build a `ModelPackage` from a `CompiledModel`.
`representation` controls how each attribute's marginal distribution is
serialized:
- ``"quantized"``: PMF on ``2**n_qubits`` support points (e.g. for
quantum-amplitude-estimation-style consumers).
- ``"histogram"``: variable-width-free histogram with ``2**n_qubits``
equal-width bins.
- ``"sampler"``: raw Monte-Carlo samples, JAX-free (plain floats), for
consumers that just want data to resample from.
`params` is the unconstrained parameter dict to export at — pass a fit's
`raw_params` so the exported distributions and the manifest both record
the fitted model rather than its declared initial values.
`observations` and `diagnostics` are recorded in the manifest as
provenance; `notes` is free text carried alongside.
"""
names = attribute_names or [a.name for a in compiled.model.attributes]
key = jax.random.PRNGKey(seed)
attr_exports = []
for name in names:
key, subkey = jax.random.split(key)
attr = compiled.model.attribute(name)
samples = np.asarray(compiled.sample(name, subkey, (n_samples,), params))
lo, hi = float(np.min(samples)), float(np.max(samples))
mean = float(np.mean(samples))
# The exported mean is a Monte-Carlo estimate like any other, so it
# is exported with the error of the same draw rather than as a bare
# number a consumer would read as exact.
error = float(np.asarray(_ffi.stats_mean_error(samples.ravel()))[1])
common = dict(
unit_description=units.describe(attr.unit),
std_error=error,
n_samples=int(n_samples),
)
if representation == "sampler":
attr_exports.append(
AttributeExport(
name=name, unit=attr.unit, domain=attr.domain,
representation="sampler", support=(lo, hi),
samples=samples.tolist(), mean=mean, **common,
)
)
continue
n_bins = 2**n_qubits
counts, edges = np.histogram(samples, bins=n_bins, range=(lo, hi))
probs = (counts / max(counts.sum(), 1)).tolist()
if representation == "histogram":
attr_exports.append(
AttributeExport(
name=name, unit=attr.unit, domain=attr.domain,
representation="histogram", support=(lo, hi),
edges=edges.tolist(), probs=probs, mean=mean, **common,
)
)
elif representation == "quantized":
grid = ((edges[:-1] + edges[1:]) / 2.0).tolist()
attr_exports.append(
AttributeExport(
name=name, unit=attr.unit, domain=attr.domain,
representation="quantized", support=(lo, hi),
grid=grid, probs=probs, mean=mean, **common,
)
)
else:
raise ValueError(f"Unknown representation: {representation!r}")
exported_params = {
k: np.asarray(v).tolist()
for k, v in compiled.constrained_params(params).items()
}
manifest = build_manifest(
compiled,
params=params,
observations=observations,
diagnostics=diagnostics,
notes=notes,
)
return ModelPackage(
format_version=FORMAT_VERSION,
model_name=repr(compiled.model),
params=exported_params,
attributes=attr_exports,
manifest=manifest,
diagnostics=manifest.diagnostics or None,
)