| """Dense primitive cells to zero-free sparse frames, and sparse training corpora.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
| from ..configuration import ActionCodec2Config, ConfigError, _integer |
| from ..primitives import PrimitiveStore |
|
|
|
|
| @dataclass(frozen=True) |
| class SparseLayout: |
| """Maps dense physical cells to zero-free sparse atoms.""" |
|
|
| bins: tuple[int, ...] |
| zero: tuple[int | None, ...] |
|
|
| def __post_init__(self) -> None: |
| if not self.bins or len(self.bins) != len(self.zero): |
| raise ValueError("bins and zero must have one entry per slot") |
| if len(self.bins) > 15: |
| raise ValueError("sparse frames support at most 15 dense slots") |
| for count, zero in zip(self.bins, self.zero, strict=True): |
| _integer(count, "bin count", positive=True) |
| if zero is not None: |
| _integer(zero, "zero index") |
| if count < 1 or (zero is not None and not 0 <= zero < count): |
| raise ValueError("every slot needs a bin and an in-range zero index") |
| if zero is not None and count < 2: |
| raise ValueError("a motion slot needs at least one nonzero primitive") |
|
|
| @classmethod |
| def from_config(cls, config: ActionCodec2Config, primitives: PrimitiveStore) -> SparseLayout: |
| """Motion slots (order 1 or 2) keep their exact zero primitive; order 0 is change-coded.""" |
| zero: list[int | None] = [None] * len(primitives.bins) |
| for component in config.physical.components.values(): |
| if component.primitive_order == 0: |
| continue |
| for dimension in component.dims: |
| axis = np.asarray(primitives.values[dimension]) |
| matches = np.flatnonzero(axis == 0) |
| if not len(matches): |
| raise ConfigError(f"dimension {dimension} has no exact zero primitive") |
| zero[dimension] = int(matches[0]) |
| return cls(tuple(int(b) for b in primitives.bins), tuple(zero)) |
|
|
| @property |
| def dims(self) -> int: |
| return len(self.bins) |
|
|
| @property |
| def e_dim(self) -> int: |
| """Sparse slot index of the end-of-frame atom.""" |
| return len(self.bins) |
|
|
| @property |
| def sparse_bins(self) -> tuple[int, ...]: |
| """Motion slots drop their zero primitive; ``E`` has exactly one value.""" |
| return tuple(b - (z is not None) for b, z in zip(self.bins, self.zero, strict=True)) + (1,) |
|
|
| @property |
| def native_zero(self) -> list[int]: |
| return [-1 if z is None else int(z) for z in self.zero] |
|
|
| def sparsify(self, cells: np.ndarray) -> np.ndarray: |
| """Dense chunk ``(T, D)`` to sparse frames ``(T, D + 1)``.""" |
| cells = np.asarray(cells, dtype=np.int64) |
| if cells.ndim != 2 or cells.shape[1] != self.dims or not len(cells): |
| raise ValueError(f"cells must have shape (T, {self.dims}) with T >= 1") |
| if np.any(cells >= np.asarray(self.bins)): |
| raise ValueError("cell outside the primitive range") |
| out = np.full((len(cells), self.dims + 1), -1, dtype=np.int64) |
| for d, z in enumerate(self.zero): |
| column = cells[:, d] |
| if np.any((column >= 0) != (column[0] >= 0)): |
| raise ValueError("a slot must be active for the whole chunk or not at all") |
| if column[0] < 0: |
| continue |
| if z is None: |
| changed = np.concatenate(([True], column[1:] != column[:-1])) |
| out[changed, d] = column[changed] |
| else: |
| moving = column != z |
| out[moving, d] = column[moving] - (column[moving] > z) |
| content = np.flatnonzero(np.any(out[:, : self.dims] >= 0, axis=1)) |
| if len(content): |
| out[: content[-1], self.dims] = 0 |
| return out |
|
|
| def densify(self, sparse: np.ndarray, active: Sequence[bool]) -> np.ndarray: |
| """Inverse of :meth:`sparsify` given which dense slots are active.""" |
| sparse = np.asarray(sparse, dtype=np.int64) |
| out = np.full((len(sparse), self.dims), -1, dtype=np.int64) |
| for d, z in enumerate(self.zero): |
| if not active[d]: |
| continue |
| column = sparse[:, d] |
| if z is None: |
| if column[0] < 0: |
| raise ValueError("an absolute slot must be stated on frame 0") |
| index = np.maximum.accumulate(np.where(column >= 0, np.arange(len(column)), 0)) |
| out[:, d] = column[index] |
| else: |
| out[:, d] = np.where(column < 0, z, column + (column >= z)) |
| return out |
|
|
|
|
| def sparsify_chunks( |
| layout: SparseLayout, chunks: Sequence[np.ndarray] |
| ) -> tuple[np.ndarray, np.ndarray]: |
| """Vectorized :meth:`SparseLayout.sparsify` over many dense chunks ``[(T_i, D)]``.""" |
| lengths = np.asarray([len(c) for c in chunks], dtype=np.int64) |
| if not len(lengths) or np.any(lengths < 1): |
| raise ValueError("chunks must be non-empty") |
| starts = np.concatenate(([0], np.cumsum(lengths))) |
| arrays = [np.asarray(c) for c in chunks] |
| for values in arrays: |
| if values.ndim != 2 or values.shape[1] != layout.dims or values.dtype.kind not in "iu": |
| raise ValueError("training cells must be integer matrices matching the sparse layout") |
| if np.any(values < -1) or np.any(values >= np.asarray(layout.bins)): |
| raise ValueError("cell outside the primitive range") |
| if np.any((values >= 0) != (values[0] >= 0)): |
| raise ValueError("a slot must be active for the whole chunk or not at all") |
| dtype = np.int16 if max(layout.bins) <= np.iinfo(np.int16).max else np.int32 |
| cells = np.concatenate(arrays, dtype=dtype) |
| if getattr(layout, "dense", False): |
| out = np.full((len(cells), layout.dims + 1), -1, dtype=dtype) |
| out[:, : layout.dims] = cells |
| out[:, layout.dims] = 0 |
| out[starts[1:] - 1, layout.dims] = -1 |
| return out, starts |
| chunk = np.repeat(np.arange(len(lengths)), lengths) |
| first = np.zeros(len(cells), dtype=bool) |
| first[starts[:-1]] = True |
| out = np.full((len(cells), layout.dims + 1), -1, dtype=dtype) |
| for d, z in enumerate(layout.zero): |
| column = cells[:, d] |
| active = column >= 0 |
| if z is None: |
| changed = active & (first | (column != np.roll(column, 1))) |
| out[changed, d] = column[changed] |
| else: |
| moving = active & (column != z) |
| out[moving, d] = column[moving] - (column[moving] > z) |
| local = np.arange(len(cells)) - starts[chunk] |
| content = np.any(out[:, : layout.dims] >= 0, axis=1) |
| last = np.maximum.reduceat(np.where(content, local, -1), starts[:-1]) |
| out[local < last[chunk], layout.dims] = 0 |
| return out, starts |
|
|
|
|
| def static_frames(cells: np.ndarray, layout: SparseLayout) -> np.ndarray: |
| """Frames whose motion cells are all zero and whose absolute slots repeat the last frame.""" |
| moving = np.zeros(len(cells), dtype=bool) |
| for d, zero in enumerate(layout.zero): |
| column = cells[:, d] |
| if column[0] < 0: |
| continue |
| if zero is None: |
| moving[1:] |= column[1:] != column[:-1] |
| else: |
| moving |= column != zero |
| return ~moving |
|
|
|
|
| def compress_static(cells: np.ndarray, layout: SparseLayout, run_cap: int) -> np.ndarray: |
| """Drop leading/trailing static frames and cap interior static runs at ``run_cap``.""" |
| static = static_frames(cells, layout) |
| if static.all(): |
| return cells[:1] |
| keep = np.zeros(len(cells), dtype=bool) |
| first, last = int(np.argmax(~static)), len(cells) - int(np.argmax(~static[::-1])) |
| keep[first:last] = True |
| index = np.arange(len(cells)) |
| run_start = np.maximum.accumulate(np.where(~static, index + 1, 0)) |
| keep &= ~static | (index - run_start < run_cap) |
| return cells[keep] |
|
|
|
|
| def training_chunks( |
| trajectories: Sequence[np.ndarray], |
| layout: SparseLayout, |
| *, |
| horizon: int, |
| static_run_cap: int, |
| dtype: np.dtype | type = np.int64, |
| ) -> list[np.ndarray]: |
| """Static-compressed cell trajectories cut into aligned ``horizon`` chunks.""" |
| out = [] |
| for raw in trajectories: |
| cells = compress_static(np.asarray(raw), layout, static_run_cap) |
| for start in range(0, len(cells), horizon): |
| piece = cells[start : start + horizon] |
| if len(piece) >= 2 or start == 0: |
| out.append(np.asarray(piece, dtype=dtype)) |
| return out |
|
|
