"""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.
"""
from __future__ import annotations
import dataclasses
import json
from typing import Any
import jax
import numpy as np
FORMAT_VERSION = "1.0"
[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
[docs]
@dataclasses.dataclass
class ModelPackage:
format_version: str
model_name: str
params: dict[str, Any]
attributes: list[AttributeExport]
correlation: dict[str, Any] | None = None
def to_dict(self) -> dict[str, Any]:
return dataclasses.asdict(self)
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)
attrs = [AttributeExport(**a) for a in data["attributes"]]
data = dict(data)
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,
) -> 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.
"""
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,)))
lo, hi = float(np.min(samples)), float(np.max(samples))
mean = float(np.mean(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,
)
)
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,
)
)
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,
)
)
else:
raise ValueError(f"Unknown representation: {representation!r}")
params = {k: np.asarray(v).tolist() for k, v in compiled.constrained_params().items()}
return ModelPackage(
format_version=FORMAT_VERSION,
model_name=repr(compiled.model),
params=params,
attributes=attr_exports,
)