File size: 4,490 Bytes
895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 | 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 | """Variable-width trajectory contract shared by statistics and fitting."""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
import numpy as np
from ._validation import active_dimensions
from .configuration import ActionCodec2Config
__all__ = ["ActiveTrajectory", "TrajectoryLike", "prepare_trajectory"]
@dataclass(frozen=True)
class ActiveTrajectory:
"""One complete trajectory over an explicit subset of codebook dimensions."""
values: np.ndarray
active_dims: tuple[int, ...]
previous_first_order: np.ndarray | None = None
def __init__(
self,
values: np.ndarray,
active_dims: Sequence[int],
*,
previous_first_order: np.ndarray | None = None,
) -> None:
object.__setattr__(self, "values", np.asarray(values))
object.__setattr__(self, "active_dims", tuple(active_dims))
previous = (
None
if previous_first_order is None
else np.asarray(previous_first_order, dtype=np.float32)
)
if previous is not None and (
previous.shape != (len(self.active_dims),) or not np.all(np.isfinite(previous))
):
raise ValueError("previous_first_order must be finite with shape [D_active]")
object.__setattr__(self, "previous_first_order", previous)
@classmethod
def from_padded(
cls,
values: np.ndarray,
*,
active_dims: Sequence[int],
previous_first_order: np.ndarray | None = None,
) -> ActiveTrajectory:
"""Drop engineering-only dense padding before the semantic boundary."""
array = np.asarray(values)
dims = tuple(int(value) for value in active_dims)
if array.ndim != 2 or not dims or min(dims) < 0 or max(dims) >= array.shape[1]:
raise ValueError("padded values and active_dims do not define a valid trajectory")
previous = None if previous_first_order is None else np.take(previous_first_order, dims)
return cls(np.take(array, dims, axis=1), dims, previous_first_order=previous)
def trajectory_boundary(
trajectory: TrajectoryLike, config: ActionCodec2Config, active_dims: tuple[int, ...]
) -> np.ndarray | None:
"""Expand explicit compact state to ``[D]`` for fitting and validation."""
previous = getattr(trajectory, "previous_first_order", None)
active = set(active_dims)
required = any(
c.primitive_order == 2 and active.intersection(c.dims)
for c in config.physical.components.values()
)
if previous is None and required:
raise ValueError("second-order trajectories require explicit previous_first_order")
if previous is None and not any(
c.primitive_order == 2 for c in config.physical.components.values()
):
return None
dense = np.zeros(config.physical.input_dim, dtype=np.float32)
if previous is not None:
dense[list(active_dims)] = previous
return dense
TrajectoryLike = np.ndarray | ActiveTrajectory
def prepare_trajectory(
trajectory: TrajectoryLike,
*,
dimension: int,
index: int,
dtype: np.dtype | type = np.float32,
) -> tuple[np.ndarray, tuple[int, ...]]:
"""Validate a trajectory and return compact values plus global dimension ids."""
if isinstance(trajectory, ActiveTrajectory):
raw = trajectory.values
active_dims = trajectory.active_dims
else:
raw = trajectory
array = np.asarray(raw)
active_dims = tuple(range(array.shape[1])) if array.ndim == 2 else ()
values = np.asarray(raw, dtype=dtype)
if values.ndim != 2 or len(values) < 1 or values.shape[1] < 1:
raise ValueError(
f"trajectory {index} must have non-empty shape [T, D_active], got {values.shape}"
)
if len(active_dims) != values.shape[1]:
raise ValueError(
f"trajectory {index} active_dims has {len(active_dims)} entries for "
f"width {values.shape[1]}"
)
if any(isinstance(item, bool) or not isinstance(item, int) for item in active_dims):
raise ValueError(f"trajectory {index} active_dims must contain integers")
active_dims = active_dimensions(
active_dims, dimension, subject=f"trajectory {index} active_dims"
)
if not np.all(np.isfinite(values)):
raise ValueError(f"trajectory {index} contains a non-finite value")
return np.ascontiguousarray(values), active_dims
|