File size: 8,604 Bytes
1bfc3f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 | """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
|