1import math
 2import numpy as np
 3
 4
 5def concentration_matrix(length: int, bandwidth: float) -> np.ndarray:
 6    """Finite discrete time-frequency concentration matrix."""
 7    if length < 1 or not 0.0 < bandwidth < 0.5:
 8        raise ValueError("length >= 1 and bandwidth in (0, .5) required")
 9    i = np.arange(length)
10    d = i[:, None] - i[None, :]
11    C = np.empty((length, length), dtype=np.float64)
12    off = d != 0
13    C[off] = np.sin(2.0 * np.pi * bandwidth * d[off]) / (np.pi * d[off])
14    C[~off] = 2.0 * bandwidth
15    return (C + C.T) * 0.5
16
17
18def asymptotic_rank(length: int, bandwidth: float, delta: float) -> float:
19    """Paper-inspired rank estimate with c=2*bandwidth*length."""
20    if not 0.0 < delta < 0.5:
21        raise ValueError("delta must be in (0, .5)")
22    c = 2.0 * bandwidth * length
23    log_odds = math.log((1.0 - delta) / delta)
24    return c + log_odds / math.pi**2 * math.log(
25        max(4.0 * math.pi**2 * c / log_odds, 1.000001))
26
27
28def fourier_basis(length: int, rank: int) -> np.ndarray:
29    """Real orthonormal low-frequency Fourier basis (cos/sin pairs)."""
30    if not 1 <= rank <= length:
31        raise ValueError("rank must be between 1 and length")
32    n = np.arange(length)
33    cols = [np.ones(length) / np.sqrt(length)]
34    k = 1
35    while len(cols) < rank:
36        cols.append(np.sqrt(2.0 / length) * np.cos(2 * np.pi * k * n / length))
37        if len(cols) < rank:
38            cols.append(np.sqrt(2.0 / length) * np.sin(2 * np.pi * k * n / length))
39        k += 1
40    return np.stack(cols[:rank], axis=1)
41
42
43class ProlateBottleneck:
44    """Fixed orthogonal DPSS projection along axis 1 of [B,T,D] arrays."""
45    def __init__(self, length: int, bandwidth: float, delta: float = 0.1,
46                 rank: int | None = None):
47        C = concentration_matrix(length, bandwidth)
48        vals, vecs = np.linalg.eigh(C)
49        order = np.argsort(vals)[::-1]
50        self.eigenvalues = vals[order]
51        self.full_basis = vecs[:, order]
52        proposed = math.ceil(asymptotic_rank(length, bandwidth, delta))
53        self.empirical_rank = int(np.sum(self.eigenvalues > delta))
54        chosen = proposed if rank is None else rank
55        self.rank = int(np.clip(chosen, 1, length))
56        self.basis = self.full_basis[:, :self.rank]
57
58    def encode(self, x: np.ndarray) -> np.ndarray:
59        if x.ndim != 3 or x.shape[1] != self.basis.shape[0]:
60            raise ValueError("x must have shape [batch, matching_length, channels]")
61        return np.einsum("btd,tr->brd", x, self.basis)
62
63    def decode(self, z: np.ndarray) -> np.ndarray:
64        if z.ndim != 3 or z.shape[1] != self.rank:
65            raise ValueError("z must have shape [batch, rank, channels]")
66        return np.einsum("brd,tr->btd", z, self.basis)
67
68    def reconstruct(self, x: np.ndarray) -> np.ndarray:
69        return self.decode(self.encode(x))
70
71    @property
72    def attention_work_fraction(self) -> float:
73        return (self.rank / self.basis.shape[0]) ** 2