"""Streaming corpus statistics for first- and second-order primitive design.""" from __future__ import annotations import json from collections.abc import Iterable, Sequence from dataclasses import asdict, dataclass from pathlib import Path import numpy as np from .boundary import source_differences from .configuration import ActionCodec2Config from .trajectory import TrajectoryLike, prepare_trajectory __all__ = [ "DEFAULT_QUANTILES", "AxisStatistics", "CorpusStatistics", "analyze_trajectories", ] DEFAULT_QUANTILES = (0.001, 0.01, 0.5, 0.99, 0.999) class _StreamingAxis: def __init__(self, sample_size: int, rng: np.random.Generator) -> None: self.sample_size = sample_size self.rng = rng self.count = 0 self.minimum = np.inf self.maximum = -np.inf self.mean = 0.0 self.m2 = 0.0 self.sample = np.empty(0, dtype=np.float32) self.priority = np.empty(0, dtype=np.float64) def update(self, raw: np.ndarray) -> None: values = np.asarray(raw, dtype=np.float32).reshape(-1) if not len(values): return batch_count = len(values) batch_mean = float(values.mean(dtype=np.float64)) batch_m2 = float(np.square(values.astype(np.float64) - batch_mean).sum()) total = self.count + batch_count delta = batch_mean - self.mean self.mean += delta * batch_count / total self.m2 += batch_m2 + delta * delta * self.count * batch_count / total self.count = total self.minimum = min(self.minimum, float(values.min())) self.maximum = max(self.maximum, float(values.max())) priority = self.rng.random(batch_count) candidates = np.concatenate((self.sample, values)) candidate_priority = np.concatenate((self.priority, priority)) if len(candidates) > self.sample_size: keep = np.argpartition(candidate_priority, self.sample_size - 1)[: self.sample_size] candidates = candidates[keep] candidate_priority = candidate_priority[keep] self.sample = candidates self.priority = candidate_priority def finish(self, quantiles: Sequence[float]) -> AxisStatistics: if self.count == 0: return AxisStatistics(0, None, None, None, None, {}, {}) signed = np.quantile(self.sample, quantiles) absolute = np.quantile(np.abs(self.sample), quantiles) return AxisStatistics( samples=self.count, minimum=self.minimum, maximum=self.maximum, mean=self.mean, std=float(np.sqrt(self.m2 / self.count)), quantiles={f"{q:g}": float(value) for q, value in zip(quantiles, signed, strict=True)}, absolute_quantiles={ f"{q:g}": float(value) for q, value in zip(quantiles, absolute, strict=True) }, ) @dataclass(frozen=True) class AxisStatistics: """Statistics for one global codebook coordinate at one signal order.""" samples: int minimum: float | None maximum: float | None mean: float | None std: float | None quantiles: dict[str, float] absolute_quantiles: dict[str, float] @dataclass(frozen=True) class CorpusStatistics: """Per-axis input, first-difference and optional second-difference measurements.""" dimension: int episodes: int frames: int active_values: int binary_dims: list[int] trajectory_arities: dict[str, int] origin_max_abs: list[float | None] order_0: list[AxisStatistics] order_1: list[AxisStatistics] quantiles: list[float] quantile_sample_size: int order_2: list[AxisStatistics] | None = None def to_dict(self) -> dict[str, object]: payload = asdict(self) payload["schema"] = "actioncodec2_first_order_statistics/1" if self.order_2 is None: payload.pop("order_2") else: payload["schema"] = "actioncodec2_differential_statistics/2" return payload def save(self, path: str | Path) -> None: destination = Path(path) destination.parent.mkdir(parents=True, exist_ok=True) destination.write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8") def _validate_options( dimension: int, binary_dims: Sequence[int], quantiles: Sequence[float], sample_size: int, ) -> tuple[set[int], tuple[float, ...]]: if dimension < 1: raise ValueError("dimension must be positive") binary = {int(value) for value in binary_dims} if any(value < 0 or value >= dimension for value in binary): raise ValueError("binary_dims contains a dimension outside the codebook") checked_quantiles = tuple(float(value) for value in quantiles) if not checked_quantiles or any(not 0.0 <= value <= 1.0 for value in checked_quantiles): raise ValueError("quantiles must contain values in [0, 1]") if sample_size < 1: raise ValueError("quantile_sample_size must be positive") return binary, checked_quantiles def analyze_trajectories( trajectories: Iterable[TrajectoryLike], *, dimension: int, binary_dims: Sequence[int] = (), quantiles: Sequence[float] = DEFAULT_QUANTILES, quantile_sample_size: int = 1_000_000, seed: int = 0, config: ActionCodec2Config | None = None, ) -> CorpusStatistics: 'Measure variable-width canonical absolute trajectories in one pass.' binary, checked_quantiles = _validate_options( dimension, binary_dims, quantiles, quantile_sample_size ) rng = np.random.default_rng(seed) order_0 = [_StreamingAxis(quantile_sample_size, rng) for _ in range(dimension)] order_1 = [_StreamingAxis(quantile_sample_size, rng) for _ in range(dimension)] order_2 = ( None if config is None else [_StreamingAxis(quantile_sample_size, rng) for _ in range(dimension)] ) if config is not None and config.physical.input_dim != dimension: raise ValueError("statistics dimension must match config.physical.input_dim") origin = np.full(dimension, -np.inf, dtype=np.float64) arities: dict[int, int] = {} episodes = frames = active_values = 0 for index, raw in enumerate(trajectories): values, active_dims = prepare_trajectory( raw, dimension=dimension, index=index, dtype=np.float32 ) episodes += 1 frames += len(values) active_values += values.size arities[len(active_dims)] = arities.get(len(active_dims), 0) + 1 for local, global_dim in enumerate(active_dims): axis = values[:, local] order_0[global_dim].update(axis) origin[global_dim] = max(origin[global_dim], abs(float(axis[0]))) if global_dim not in binary and config is None: delta = np.diff(axis, prepend=np.float32(0.0)) order_1[global_dim].update(delta) if config is not None: for dims, first, second in source_differences(raw, config).values(): for local, global_dim in enumerate(dims): if global_dim not in binary: order_1[global_dim].update(first[:, local]) order_2[global_dim].update(second[:, local]) if episodes == 0: raise ValueError("statistics require at least one complete trajectory") return CorpusStatistics( dimension=dimension, episodes=episodes, frames=frames, active_values=active_values, binary_dims=sorted(binary), trajectory_arities={str(key): value for key, value in sorted(arities.items())}, origin_max_abs=[None if value == -np.inf else float(value) for value in origin], order_0=[axis.finish(checked_quantiles) for axis in order_0], order_1=[axis.finish(checked_quantiles) for axis in order_1], quantiles=list(checked_quantiles), quantile_sample_size=quantile_sample_size, order_2=None if order_2 is None else [axis.finish(checked_quantiles) for axis in order_2], )