"""Standardized model export format (`ModelPackage`).
This is the versioned interchange format for a compiled `prophys`
model. It carries no dependency on any downstream consumer, and 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
derived again rather than only read,
- a `unit_description` per attribute 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.
Format 1.2 adds `surrogates`, one `SurrogateExport` per learned surrogate
of the model, carrying the native program or tree spec, weights, support
region, head posterior, training card and a parity checksum that the
standalone runtime library reproduces without Python.
"""
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 written by this prophys (the ``format_version`` field of
#: every exported package).
FORMAT_VERSION = "1.2"
#: 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", "1.2"})
[docs]
@dataclasses.dataclass
class AttributeExport:
"""The exported marginal distribution of one model attribute.
Which of `grid`, `edges` and `samples` is set depends on
`representation`. `probs` is set for ``"quantized"`` and
``"histogram"``.
"""
name: str
"""Attribute name."""
unit: str
"""Unit string of the attribute (e.g. ``"kW"``)."""
domain: str
"""Free text domain description of the attribute, as declared in the model."""
representation: str # "quantized" | "histogram" | "sampler"
"""``"quantized"``, ``"histogram"`` or ``"sampler"``."""
support: tuple[float, float]
"""``(min, max)`` of the Monte Carlo draws the export was built from."""
# "quantized": grid (list[float]) + probs (list[float])
grid: list[float] | None = None
"""``"quantized"``: the support points (bin centres)."""
probs: list[float] | None = None
"""``"quantized"`` and ``"histogram"``: probability of each point or bin,
summing to one."""
# "histogram": edges (n+1,) + probs (n,)
edges: list[float] | None = None
"""``"histogram"``: the ``n + 1`` bin edges, equally spaced."""
# "sampler": raw NumPy samples, JAX-free and picklable
samples: list[float] | None = None
"""``"sampler"``: the raw Monte Carlo draws as plain floats."""
mean: float | None = None
"""Monte Carlo mean of the draws."""
# -- 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:
"""Versioned export of a compiled model's distributions as JSON, with provenance.
Built by :func:`export_model` (or ``CompiledModel.export()``), written
with :meth:`save` and read back with :meth:`load`. It carries each
attribute's marginal distribution, the parameter values, a
:class:`~prophys.export.RunManifest` recording what produced it, optional
diagnostic reports, and learned surrogates in a form the standalone
runtime library evaluates without Python.
Example:
>>> package = export_model(compiled, params=fit.raw_params, representation="histogram")
>>> package.save("model.json")
>>> ModelPackage.load("model.json").attribute("energy").mean
"""
format_version: str
"""Format version of the package, e.g. ``"1.2"``."""
model_name: str
"""Readable name of the source model."""
params: dict[str, Any]
"""Parameter values the package was exported at, in constrained space."""
attributes: list[AttributeExport]
"""One exported distribution per attribute."""
correlation: dict[str, Any] | None = None
"""Reserved for a dependence summary between attributes, not filled by
:func:`export_model`."""
# -- 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`."""
# -- format 1.2 --------------------------------------------------------
surrogates: list | None = None
"""Learned surrogates of the model as `prophys.export.surrogate` exports."""
[docs]
def surrogate(self, name: str):
"""Return the exported surrogate called `name`.
Raises:
KeyError: If the package has no such surrogate.
"""
for s in self.surrogates or []:
if s.name == name:
return s
raise KeyError(f"No surrogate {name!r} in this package")
[docs]
def to_dict(self) -> dict[str, Any]:
"""Return the package as a dict ready for JSON (the content of the saved file)."""
surr = self.surrogates
self.surrogates = None
data = dataclasses.asdict(self)
self.surrogates = surr
if self.manifest is not None:
data["manifest"] = self.manifest.to_dict()
if surr is not None:
data["surrogates"] = [s.to_dict() for s in surr]
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`), or when the package has no
manifest because it was written in format 1.0.
"""
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)
[docs]
def attribute(self, name: str) -> AttributeExport:
"""Return the exported distribution of the attribute called `name`.
Raises:
KeyError: If the package has no such attribute.
"""
for attr in self.attributes:
if attr.name == name:
return attr
raise KeyError(f"No attribute {name!r} in this package; it has {[a.name for a in self.attributes]}")
[docs]
def to_json(self, indent: int = 2) -> str:
"""Serialize the package to a JSON string."""
return json.dumps(self.to_dict(), indent=indent)
[docs]
def save(self, path: str) -> None:
"""Write the package to a JSON file at `path`."""
with open(path, "w") as f:
f.write(self.to_json())
[docs]
@classmethod
def load(cls, path: str) -> "ModelPackage":
"""Read a package from a JSON file written by :meth:`save`.
Files of format 1.0, 1.1 and 1.2 load, and fields this version does
not know are ignored.
Raises:
ValueError: If the file's format version is not supported.
"""
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
if data.get("surrogates"):
from .surrogate import SurrogateExport
data["surrogates"] = [SurrogateExport.from_dict(x) for x in data["surrogates"]]
known_pkg = {f.name for f in dataclasses.fields(cls)}
return cls(**{k: v for k, v in data.items() if k in known_pkg})
[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:
"""Export a compiled model's attribute distributions as a :class:`ModelPackage`.
Each attribute is sampled `n_samples` times and its marginal
distribution serialized according to `representation`:
* ``"quantized"``: PMF on the ``2**n_qubits`` centres of equally spaced
bins over the sample range (for example for consumers in the style of
quantum amplitude estimation).
* ``"histogram"``: histogram with ``2**n_qubits`` equally wide bins.
* ``"sampler"``: the raw Monte Carlo samples as plain floats, for
consumers that just want data to resample from.
Every attribute also records its mean with the Monte Carlo standard
error of the same draw. Learned surrogates of the model are exported
alongside, except those with `Physics` layers, which are skipped.
Args:
compiled: The compiled model.
attribute_names: Attributes to export. Defaults to all.
representation: ``"quantized"``, ``"histogram"`` or ``"sampler"``.
n_qubits: Resolution: ``2**n_qubits`` points or bins.
n_samples: Monte Carlo draws per attribute.
seed: PRNG seed of the draws.
params: Unconstrained parameter values to export at. Pass a fit's
`raw_params` so the distributions and the manifest record the
fitted model rather than its declared initial values.
observations: Observations of a fit, recorded in the manifest as
counts.
diagnostics: Name to diagnostic report (e.g. the output of
:func:`~prophys.engine.report_card`), serialized in the manifest
and the package.
notes: Free text recorded in the manifest.
Returns:
The :class:`ModelPackage`.
Raises:
ValueError: If `representation` is unknown.
"""
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,
)
from ..nn.graph import find_surrogates
from .surrogate import export_surrogate
surrogates = []
for sg in find_surrogates(compiled.model).values():
if all(st.kind == "native" for st in getattr(sg, "stages", [])) and len(getattr(sg, "stages", [None])) <= 1:
surrogates.append(export_surrogate(sg, compiled._full_raw_params(params)))
return ModelPackage(
format_version=FORMAT_VERSION,
model_name=repr(compiled.model),
params=exported_params,
attributes=attr_exports,
manifest=manifest,
surrogates=surrogates or None,
diagnostics=manifest.diagnostics or None,
)
