"""Path-signature dataset loader (Cao et al.)."""
from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
from typing import Any
import numpy as np
from ..._logging import get_logger
log = get_logger(__name__)
from ..base import BaseDataLoader, register_loader
from ..bundle import DataBundle
def _import_standalone_loader(root: Path):
"""Import standalone_data_loader.py from the dataset root."""
module_path = root / "standalone_data_loader.py"
if not module_path.is_file():
raise FileNotFoundError(
f"path_signature loader not found at {module_path}. "
"Link data/path_signature to readout_dataset_cleaned.",
)
spec = importlib.util.spec_from_file_location(
"arcade_path_signature_standalone_loader",
module_path,
)
if spec is None or spec.loader is None:
raise ImportError(f"Cannot load module spec from {module_path}")
mod = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = mod
spec.loader.exec_module(mod)
return mod
def _traces_to_iq(traces: np.ndarray) -> np.ndarray:
"""Convert complex or real traces to IQ format ``(n, t, 2)``."""
tr = np.asarray(traces)
if tr.ndim == 1:
tr = tr[np.newaxis, :]
if tr.ndim != 2:
raise ValueError(f"Expected 2D traces (n, t), got {tr.shape}")
if np.iscomplexobj(tr):
return np.stack(
[np.real(tr).astype(np.float64), np.imag(tr).astype(np.float64)],
axis=-1,
)
tr = tr.astype(np.float64)
return np.stack([tr, np.zeros_like(tr)], axis=-1)
def _subsample_per_class(
traces: np.ndarray,
labels: np.ndarray,
max_per_class: int,
seed: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Randomly keep up to *max_per_class* shots per label (paper protocol)."""
labels = np.asarray(labels, dtype=np.int64)
rng = np.random.default_rng(seed)
parts_t: list[np.ndarray] = []
parts_l: list[np.ndarray] = []
parts_i: list[np.ndarray] = []
for cls in np.unique(labels):
idx = np.where(labels == cls)[0]
rng.shuffle(idx)
take = idx[:max_per_class]
parts_t.append(traces[take])
parts_l.append(labels[take])
parts_i.append(take)
return (
np.concatenate(parts_t, axis=0),
np.concatenate(parts_l, axis=0),
np.concatenate(parts_i, axis=0),
)
[docs]
@register_loader("path_signature")
class PathSignatureLoader(BaseDataLoader):
"""Load Cao et al. path-signature benchmark datasets."""
supported_formats = ()
def __init__(
self,
dataset_name: str = "RIKEN2_Q56_simultaneous",
time_truncation: float = 2.0,
label_source: str = "expected",
max_per_class: int | None = None,
subsample_seed: int = 42,
):
self.dataset_name = dataset_name
self.time_truncation = time_truncation
self.label_source = label_source
self.max_per_class = max_per_class
self.subsample_seed = subsample_seed
[docs]
def load(self, path: str | Path, **kwargs: Any) -> DataBundle:
cfg = kwargs.get("config")
if isinstance(cfg, dict):
kwargs = {**cfg, **{k: v for k, v in kwargs.items() if k != "config"}}
dataset_name = kwargs.get(
"path_signature_dataset",
kwargs.get("dataset_name", self.dataset_name),
)
time_truncation = float(
kwargs.get(
"path_signature_time_truncation",
kwargs.get("time_truncation", self.time_truncation),
)
)
label_source = kwargs.get(
"path_signature_label_source",
kwargs.get("label_source", self.label_source),
)
max_per_class = kwargs.get(
"path_signature_max_per_class",
kwargs.get("max_per_class", self.max_per_class),
)
subsample_seed = int(kwargs.get("subsample_seed", self.subsample_seed))
root = Path(path)
mod = _import_standalone_loader(root)
mod.PROJECT_ROOT = root
mod.data_path_shuxiang = root / "sampled_dataset/dataset_sampled_Square.h5"
mod.data_path_riken = root / "RIKEN/Washed_Dataset/"
mod.data_path_riken2 = root / "RIKEN2/Washed_Dataset/"
mod.data_path_simo = root / "Simone/Washed_Dataset/"
ts, traces, expected, gmm, post, weights = mod.load_dataset(
dataset_name,
time_truncation=time_truncation,
)
label_map = {
"expected": expected,
"gmm": gmm,
"post_measured": post,
}
if label_source not in label_map:
raise ValueError(f"label_source must be one of {list(label_map)}")
labels = np.asarray(label_map[label_source], dtype=np.int64)
post_measured = np.asarray(post, dtype=np.float64)
iq = _traces_to_iq(traces)
if max_per_class is not None:
iq, labels, _sel = _subsample_per_class(
iq, labels, max_per_class, subsample_seed
)
post_measured = post_measured[_sel]
meta: dict[str, Any] = {
"source": str(root),
"format": "path_signature",
"dataset_name": dataset_name,
"label_source": label_source,
"time_truncation_ms": time_truncation,
"post_measured_labels": post_measured,
}
return DataBundle(traces=iq, labels=labels, time=np.asarray(ts), metadata=meta)
[docs]
@classmethod
def from_config(cls, config: dict[str, Any]) -> "PathSignatureLoader":
return cls(
dataset_name=config.get(
"path_signature_dataset",
config.get("dataset_name", "RIKEN2_Q56_simultaneous"),
),
time_truncation=float(
config.get(
"path_signature_time_truncation",
config.get("time_truncation", 2.0),
)
),
label_source=config.get(
"path_signature_label_source",
config.get("label_source", "expected"),
),
max_per_class=config.get(
"path_signature_max_per_class",
config.get("max_per_class"),
),
subsample_seed=int(config.get("subsample_seed", 42)),
)
def load_path_signature_dataset(
root: str | Path,
dataset_name: str,
*,
time_truncation: float = 2.0,
label_source: str = "expected",
max_per_class: int | None = None,
subsample_seed: int = 42,
) -> DataBundle:
"""Load a dataset via ``standalone_data_loader.load_dataset``."""
root = Path(root)
mod = _import_standalone_loader(root)
mod.PROJECT_ROOT = root # type: ignore[attr-defined]
mod.data_path_shuxiang = root / "sampled_dataset/dataset_sampled_Square.h5" # type: ignore[attr-defined]
mod.data_path_riken = root / "RIKEN/Washed_Dataset/" # type: ignore[attr-defined]
mod.data_path_riken2 = root / "RIKEN2/Washed_Dataset/" # type: ignore[attr-defined]
mod.data_path_simo = root / "Simone/Washed_Dataset/" # type: ignore[attr-defined]
ts, traces, expected, gmm, post, weights = mod.load_dataset(
dataset_name,
time_truncation=time_truncation,
)
label_map = {"expected": expected, "gmm": gmm, "post_measured": post}
if label_source not in label_map:
raise ValueError(f"label_source must be one of {list(label_map)}")
labels = np.asarray(label_map[label_source], dtype=np.int64)
post_measured = np.asarray(post, dtype=np.float64)
iq = _traces_to_iq(traces)
if max_per_class is not None:
iq, labels, _sel = _subsample_per_class(iq, labels, max_per_class, subsample_seed)
post_measured = post_measured[_sel]
log.info(
"Subsampled path_signature %s to %d shots (%d/class max)",
dataset_name,
len(labels),
max_per_class,
)
meta: dict[str, Any] = {
"source": str(root),
"format": "path_signature",
"dataset_name": dataset_name,
"label_source": label_source,
"time_truncation_ms": time_truncation,
"post_measured_labels": post_measured,
}
log.info("Loaded path_signature %s: traces %s, labels %s", dataset_name, iq.shape, labels.shape)
return DataBundle(traces=iq, labels=labels, time=np.asarray(ts), metadata=meta)