"""Synthetic datasets for quick experiments.""" from __future__ import annotations import numpy as np def make_induction_head_task(n: int = 100, T: int = 30, V: int = 20, seed: int = 42): np.random.seed(seed) perm = np.random.permutation(V) X = np.zeros((n, T), dtype=int) Y = np.zeros(n, dtype=int) X[:, 0] = np.random.randint(0, V, n) for i in range(n): for j in range(1, T - 2): if X[i, j - 1] == 0: X[i, j] = np.random.randint(1, V) else: X[i, j] = np.random.randint(0, V) X[:, -2] = np.random.randint(1, V, n) X[:, -1] = 0 contains_zero = np.any(X[:, :-1] == 0, axis=1) missing_zero_indices = np.where(~contains_zero)[0] if len(missing_zero_indices) > 0: random_indices = np.random.randint(0, T - 2, size=len(missing_zero_indices)) X[missing_zero_indices, random_indices] = 0 for i in range(n): Y[i] = X[i][np.argwhere(X[i, :-1] == 0).flatten()[-1] + 1] vocab = 1.5 ** (1 / 3) * np.linspace(-1, 1, V) vocab = vocab[perm] return vocab[X], vocab[Y] def make_induction_head_seq_to_seq_task( n: int = 1, T: int = 1000, V: int = 20, seed: int = 42 ): np.random.seed(seed) perm = np.random.permutation(V) X = np.zeros((n, T), dtype=int) Y = np.zeros((n, T), dtype=int) X[:, 0] = np.random.randint(0, V, n) for i in range(n): for j in range(1, T - 2): if X[i, j - 1] == 0: X[i, j] = np.random.randint(1, V) else: X[i, j] = np.random.randint(0, V) X[:, -2] = np.random.randint(1, V, n) X[:, -1] = 0 contains_zero = np.any(X[:, :-1] == 0, axis=1) missing_zero_indices = np.where(~contains_zero)[0] if len(missing_zero_indices) > 0: random_indices = np.random.randint(0, T - 2, size=len(missing_zero_indices)) X[missing_zero_indices, random_indices] = 0 for i in range(n): current_val = 0 flag = False for t in range(T): if X[i, t] == 0: flag = True elif flag: current_val = X[i, t] flag = False Y[i, t] = current_val vocab = 1.5 ** (1 / 3) * np.linspace(-1, 1, V) vocab[0] *= 2 vocab = vocab[perm] return vocab[X], vocab[Y] def get_rbf_sample( T: int, *, dt: float = 1.0, lengthscale: float = 10.0, sigma: float = 1.0, mu: float = 0.0, seed: int | None = None, embed_factor: int = 2, jitter: float = 1e-12, clip_neg_eigs: bool = True, ) -> np.ndarray: """ Efficiently sample from a 1D zero-mean GP on an evenly spaced grid using an RBF kernel, via circulant embedding + FFT (O(M log M)). GP prior on indices {0,...,T-1} with covariance: k(i,j) = sigma^2 * exp(-( (i-j)*dt )^2 / (2*lengthscale^2)) Args: T: number of samples dt: grid spacing lengthscale: RBF lengthscale ℓ (>0) sigma: marginal std (>=0) mu: mean seed: RNG seed embed_factor: embedding size multiplier; M is next power of 2 >= embed_factor*T. 2 is standard; increase (e.g. 4) if negative eigenvalues occur. jitter: small diagonal jitter added to k(0) to help numerical stability clip_neg_eigs: if True, clip tiny negative FFT eigenvalues to 0. Returns: z: (T,) sample """ if T <= 0: return np.zeros((0,), dtype=np.float64) if dt <= 0: raise ValueError("dt must be > 0") if lengthscale <= 0: raise ValueError("lengthscale must be > 0") if sigma < 0: raise ValueError("sigma must be >= 0") if embed_factor < 2: raise ValueError("embed_factor should be >= 2 for circulant embedding") rng = np.random.default_rng(seed) # Choose embedding size M (power of 2 for fast FFT). M_min = embed_factor * T M = 1 << (M_min - 1).bit_length() if M % 2 != 0: M += 1 # keep even for clean Nyquist handling # Build Toeplitz first column for size T: k[0..T-1] d = np.arange(T, dtype=np.float64) * dt k_col = (sigma ** 2) * np.exp(-0.5 * (d / lengthscale) ** 2) k_col[0] += jitter # Circulant embedding first column c of length M: # c = [k0, k1, ..., k_{T-1}, 0, 0, ..., 0, k_{T-1}, ..., k1] c = np.zeros(M, dtype=np.float64) c[:T] = k_col c[M - (T - 1) :] = k_col[1:][::-1] # tail mirror (exclude k0) # Eigenvalues of the circulant covariance matrix (should be >= 0) lam = np.real(np.fft.fft(c)) if clip_neg_eigs: lam = np.maximum(lam, 0.0) else: if np.min(lam) < -1e-10: raise ValueError( f"Circulant embedding not PSD (min eigenvalue {np.min(lam)}). " "Increase embed_factor or enable clip_neg_eigs." ) lam = np.maximum(lam, 0.0) # Sample in Fourier domain with conjugate symmetry so the time series is real. # numpy FFT conventions: ifft includes 1/M normalization. # To match covariance C = (1/M) F^H diag(lam) F, use: # x = sqrt(M) * ifft( sqrt(lam) * z ), with z ~ CN(0, I) and conjugate symmetry. Z = np.zeros(M, dtype=np.complex128) # k=0 (DC) and k=M/2 (Nyquist) are real-valued in the symmetric FFT. Z[0] = rng.normal() Z[M // 2] = rng.normal() # Positive frequencies 1..M/2-1: complex normals with Var=1 (CN(0,1)) re = rng.normal(size=(M // 2 - 1)) im = rng.normal(size=(M // 2 - 1)) Z[1 : M // 2] = (re + 1j * im) / np.sqrt(2.0) # Enforce conjugate symmetry Z[M // 2 + 1 :] = np.conj(Z[1 : M // 2][::-1]) # Scale by sqrt eigenvalues Y = np.sqrt(lam) * Z # Back to time domain, take real part x_full = (np.sqrt(M) * np.fft.ifft(Y)).real return mu + x_full[:T] def rbf_kernel_1d(i: int, j: int, *, sigma: float, lengthscale: float, dt: float) -> float: d = (i - j) * dt return (sigma ** 2) * np.exp(-0.5 * (d / lengthscale) ** 2) def oracle_mse_rbf_periodic_censor( P: int, L: int, *, sigma: float = 1.0, lengthscale: float = 10.0, dt: float = 1.0, num_periods: int = 50, burn_periods: int = 10, jitter: float = 1e-10, ) -> dict: """ Oracle MSE for the censored-copy task when Z_t ~ GP(mean, RBF kernel). Censoring pattern (matches ``make_censored_task``): observed at index i iff (i % P) < (P//2). But due to the task's lagging, we only ever observe indices >= L (because X is input_signal[L:]). At model time t, the oracle conditions on all observed indices <= (t+L), i.e. fixed-lag smoothing. Returns: { "mse": (T,) oracle per-timestep MSE trace, "avg_mse_last_period": scalar average over the last period (after burn-in), "avg_mse_post_burn": scalar average over all timesteps after burn-in, } """ if P <= 0 or P % 2 != 0: raise ValueError("P must be a positive even integer.") if L < 0: raise ValueError("L must be >= 0") if sigma < 0 or lengthscale <= 0 or dt <= 0: raise ValueError("Require sigma>=0, lengthscale>0, dt>0") if num_periods <= 0: raise ValueError("num_periods must be > 0") if burn_periods < 0 or burn_periods >= num_periods: raise ValueError("burn_periods must be in [0, num_periods-1]") m = P // 2 T = num_periods * P # number of Y_t evaluated N = T + L # original time indices potentially observable up to t+L burn_T = burn_periods * P # Observation availability on original time axis # Only indices >= L are observable via X (since X corresponds to original indices L..L+T-1). observed = np.zeros(N, dtype=bool) for i in range(L, N): observed[i] = (i % P) < m # Incremental Cholesky factor of K_obs (lower-triangular) obs_idx: list[int] = [] chol_L = np.zeros((0, 0), dtype=np.float64) def add_observation(i_new: int): """Add new observed index to the Cholesky factor (noise-free GP with jitter).""" nonlocal chol_L, obs_idx if len(obs_idx) == 0: k_nn = rbf_kernel_1d(i_new, i_new, sigma=sigma, lengthscale=lengthscale, dt=dt) + jitter chol_L = np.array([[np.sqrt(k_nn)]], dtype=np.float64) obs_idx.append(i_new) return # k between new point and existing obs k_vec = np.array( [rbf_kernel_1d(i_new, j, sigma=sigma, lengthscale=lengthscale, dt=dt) for j in obs_idx], dtype=np.float64, ) # (M,) # Solve L w = k_vec w = np.linalg.solve(chol_L, k_vec) # (M,) k_nn = rbf_kernel_1d(i_new, i_new, sigma=sigma, lengthscale=lengthscale, dt=dt) + jitter diag_sq = k_nn - float(w @ w) diag = np.sqrt(max(diag_sq, jitter)) # Build expanded Cholesky M = len(obs_idx) L_new = np.zeros((M + 1, M + 1), dtype=np.float64) L_new[:M, :M] = chol_L L_new[M, :M] = w L_new[M, M] = diag chol_L = L_new obs_idx.append(i_new) def posterior_var(test_t: int) -> float: """Var(Z_test_t | observed indices obs_idx), using current Cholesky.""" k_tt = rbf_kernel_1d(test_t, test_t, sigma=sigma, lengthscale=lengthscale, dt=dt) if len(obs_idx) == 0: return k_tt if obs_idx[-1] == test_t or test_t in obs_idx: return 0.0 k_tO = np.array( [rbf_kernel_1d(test_t, j, sigma=sigma, lengthscale=lengthscale, dt=dt) for j in obs_idx], dtype=np.float64, ) # alpha = L^{-1} k_tO alpha = np.linalg.solve(chol_L, k_tO) var = k_tt - float(alpha @ alpha) return max(var, 0.0) # Main loop: advance horizon h, and output MSE for t = h-L once h>=L mse = np.zeros(T, dtype=np.float64) for h in range(N): if observed[h]: add_observation(h) if h >= L: t = h - L if t < T: mse[t] = posterior_var(t) # Averages (useful “oracle loss” scalars) post_burn = mse[burn_T:] if burn_T < T else mse avg_post_burn = float(np.mean(post_burn)) if post_burn.size else float(np.mean(mse)) return avg_post_burn def get_ou_sample( T: int, dt: float = 1.0, tau: float = 10.0, sigma: float = 1.0, mu: float = 0.0, seed: int | None = None, ) -> np.ndarray: """ Sample a stationary Ornstein–Uhlenbeck process at discrete times. Continuous-time OU (one common parametrization): dX_t = -(1/tau) (X_t - mu) dt + sigma dW_t Discretization (exact transition): X_{t+dt} = mu + rho (X_t - mu) + eps rho = exp(-dt/tau) eps ~ N(0, q), q = (sigma^2 * tau / 2) * (1 - rho^2) Stationary distribution: X_t ~ N(mu, sigma^2 * tau / 2) Args: T: number of samples to return dt: sampling interval tau: relaxation time constant (> 0) sigma: diffusion scale (>= 0) mu: mean seed: RNG seed Returns: x: (T,) numpy array """ if T <= 0: return np.zeros((0,), dtype=np.float64) if tau <= 0: raise ValueError("tau must be > 0") if sigma < 0: raise ValueError("sigma must be >= 0") rng = np.random.default_rng(seed) rho = np.exp(-dt / tau) var_stationary = (sigma**2) * tau / 2.0 # Exact conditional variance for step dt q = var_stationary * (1.0 - rho**2) x = np.empty((T,), dtype=np.float64) x[0] = mu + np.sqrt(var_stationary) * rng.standard_normal() if T > 1: noise = np.sqrt(q) * rng.standard_normal(size=T - 1) for t in range(T - 1): x[t + 1] = mu + rho * (x[t] - mu) + noise[t] return x def oracle_mse_censored_task( P: int, L: int, *, dt: float = 1.0, tau: float = 10.0, sigma: float = 1.0, ) -> float: """ Oracle (minimum expected) per-timestep MSE for the censored-copy task under a stationary OU process, with periodic censoring pattern: - Period length: P - First half of each period: uncensored (perfect observation => MSE=0) - Second half: censored (missing observations) - Lookahead (in original time index units): L (the task lag) Assumes the oracle uses all uncensored samples up to time t+L to predict Z_t. Because OU is Gaussian Markov (AR(1)), the per-step conditional variance is: If future endpoint u is NOT yet observed: Var(Z_t | Z_s) = v * (1 - rho^(2 d1)) If future endpoint u IS observed: Var(Z_t | Z_s, Z_u) = v * (1 - a^2 - b^2 + 2ab rho^D) / (1 - rho^(2D)), which simplifies to: v * (1 - rho^(2 d1) - rho^(2 d2) + rho^(2D)) / (1 - rho^(2D)). Here: rho = exp(-dt/tau) v = stationary variance = sigma^2 * tau / 2 m = P/2 (must be integer; require even P) s = last uncensored time in the period, u = first uncensored time next period D = u - s = m + 1 For censored positions: k = 1..m with d1 = k and d2 = D - k. Returns: Long-run average MSE per timestep (averaged over one period). """ if P <= 0: raise ValueError("P must be positive") if P % 2 != 0: raise ValueError("This oracle formula assumes even P (because censoring uses P//2).") if L < 0: raise ValueError("L must be >= 0") if tau <= 0: raise ValueError("tau must be > 0") if sigma < 0: raise ValueError("sigma must be >= 0") m = P // 2 # censored block length rho = np.exp(-dt / tau) v = (sigma**2) * tau / 2.0 D = m + 1 # distance between last uncensored and next uncensored in this pattern rho2D = rho ** (2 * D) denom = 1.0 - rho2D mse_sum = 0.0 # Uncensored half contributes 0, so only sum over censored half (k=1..m). for k in range(1, m + 1): d1 = k d2 = D - k if d2 > L: # No observed sample after the censor block yet (given lookahead L) mse_k = v * (1.0 - rho ** (2 * d1)) else: # Smoothing with both endpoints available # v * (1 - rho^(2 d1) - rho^(2 d2) + rho^(2D)) / (1 - rho^(2D)) mse_k = v * (1.0 - rho ** (2 * d1) - rho ** (2 * d2) + rho2D) / denom mse_sum += mse_k # Average over all P timesteps in a period return mse_sum / P def oracle_info_gain( X: np.ndarray, *, lag: int, censor_val: float, dt: float = 1.0, tau: float = 10.0, sigma: float = 1.0, mu: float = 0.0, atol: float = 0.0, eps: float = 1e-12, ) -> np.ndarray: """ Oracle information-gain (salience) trace for the censored OU task. info_gain[t] = 0.5 * log( Var(Z_t | obs <= t+lag-1) / Var(Z_t | obs <= t+lag) ) This measures how much the *newest* input sample at time (t+lag) reduces uncertainty about the current target Z_t. Args: X: (T,) input signal (censored OU samples) lag: fixed lag used in the task censor_val: value used to mark censored samples dt, tau, sigma, mu: OU parameters atol: optional tolerance for detecting censor_val eps: numerical stability Returns: info_gain: (T,) oracle salience / information-gain trace """ X = np.asarray(X, dtype=np.float64) T = X.shape[0] rho = np.exp(-dt / tau) v = (sigma**2) * tau / 2.0 # stationary variance # Original-time indexing N = T + lag observed = np.zeros(N, dtype=bool) obs_value = np.zeros(N, dtype=np.float64) if atol > 0: obs_mask = np.abs(X - censor_val) > atol else: obs_mask = X != censor_val ks = lag + np.arange(T) observed[ks] = obs_mask obs_value[ks] = X # last observed <= k last_obs_leq = np.full(N, -1, dtype=int) last = -1 for k in range(N): if observed[k]: last = k last_obs_leq[k] = last # next observed >= k next_obs_geq = np.full(N, N, dtype=int) nxt = N for k in range(N - 1, -1, -1): if observed[k]: nxt = k next_obs_geq[k] = nxt def posterior_var(t: int, horizon: int) -> float: """Var(Z_t | uncensored obs with indices <= horizon).""" if horizon < 0: return v horizon = min(horizon, N - 1) s = last_obs_leq[min(t, horizon)] if t >= 0 else -1 u = next_obs_geq[t] if t <= horizon and next_obs_geq[t] <= horizon else N if s == t and s != -1: return 0.0 if s == -1 and u == N: return v if s == -1: d = u - t return v * (1.0 - rho ** (2 * d)) if u == N: d = t - s return v * (1.0 - rho ** (2 * d)) d1 = t - s d2 = u - t D = d1 + d2 rho2D = rho ** (2 * D) return v * (1.0 - rho ** (2 * d1) - rho ** (2 * d2) + rho2D) / max( 1.0 - rho2D, eps ) info_gain = np.zeros(T, dtype=np.float64) for t in range(T): var_prev = posterior_var(t, t + lag - 1) var_post = posterior_var(t, t + lag) info_gain[t] = 0.5 * np.log((var_prev + eps) / (var_post + eps)) return info_gain def oracle_info_gain_rbf( X: np.ndarray, *, lag: int, censor_val: float, dt: float = 1.0, lengthscale: float = 10.0, sigma: float = 1.0, atol: float = 0.0, jitter: float = 1e-10, eps: float = 1e-12, max_obs: int | None = None, ) -> np.ndarray: """ Oracle information-gain trace for the censored task under an RBF-kernel GP prior. Prior: Z ~ GP(mu, k), k(i,j)=sigma^2 * exp(-((i-j)*dt)^2/(2*ell^2)) Observations: Noise-free: observe Z_k exactly at uncensored k; censored => missing. At model time t: horizon h = t + lag (original-time index of current input X[t]) IG_t = 0.5 * log( Var(Z_t | obs<=h-1) / Var(Z_t | obs<=h) ) Args: X: (T,) input stream; X[t] is either Z_{t+lag} (uncensored) or censor_val (censored) lag: task lag L censor_val: sentinel for missing observations dt, lengthscale, sigma: RBF GP kernel params atol: optional tolerance for detecting censor_val jitter: diagonal jitter for numerical stability eps: stability for log ratio max_obs: cap number of retained observed points (approximation). If None, keeps all. Returns: info_gain: (T,) oracle IG trace """ if max_obs is None: max_obs = int(lag + 3 * lengthscale / dt) X = np.asarray(X, dtype=np.float64) T = X.shape[0] if lag < 0: raise ValueError("lag must be >= 0") if dt <= 0 or lengthscale <= 0: raise ValueError("dt and lengthscale must be > 0") if sigma < 0: raise ValueError("sigma must be >= 0") # Original-time axis indices potentially involved: 0..N-1 where N=T+lag N = T + lag # Which original-time indices are observed (only k=lag..lag+T-1 are ever presented via X) observed = np.zeros(N, dtype=bool) if atol > 0: obs_mask = np.abs(X - censor_val) > atol else: obs_mask = X != censor_val ks = lag + np.arange(T) observed[ks] = obs_mask # RBF kernel helpers inv_ell2 = 1.0 / (lengthscale * lengthscale) sigma2 = sigma * sigma def k_vec(t_idx: int, obs_idx: np.ndarray) -> np.ndarray: d = (t_idx - obs_idx).astype(np.float64) * dt return sigma2 * np.exp(-0.5 * (d * d) * inv_ell2) def k_tt(_: int) -> float: return sigma2 # RBF has k(t,t)=sigma^2 # Maintain Cholesky of K_obs (lower triangular), for current retained obs_idx list obs_idx: list[int] = [] chol_L = np.zeros((0, 0), dtype=np.float64) def rebuild_cholesky(): nonlocal chol_L if len(obs_idx) == 0: chol_L = np.zeros((0, 0), dtype=np.float64) return idx = np.array(obs_idx, dtype=np.int64) d = (idx[:, None] - idx[None, :]).astype(np.float64) * dt K = sigma2 * np.exp(-0.5 * (d * d) * inv_ell2) K[np.diag_indices_from(K)] += jitter chol_L = np.linalg.cholesky(K) def add_observation(i_new: int): """Rank-1 append update; if we truncate (drop oldest), rebuild.""" nonlocal chol_L # Append then (optional) truncate to max_obs obs_idx.append(i_new) if max_obs is not None and len(obs_idx) > max_obs: # Drop oldest; rebuilding is simplest/stable (max_obs should be modest). obs_idx.pop(0) rebuild_cholesky() return # Incremental update when no drop if chol_L.shape[0] == 0: chol_L = np.array([[np.sqrt(k_tt(i_new) + jitter)]], dtype=np.float64) return idx = np.array(obs_idx[:-1], dtype=np.int64) # previous obs kv = k_vec(i_new, idx) # (M,) w = np.linalg.solve(chol_L, kv) # (M,) diag_sq = (k_tt(i_new) + jitter) - float(w @ w) diag = np.sqrt(max(diag_sq, jitter)) M = chol_L.shape[0] L_new = np.zeros((M + 1, M + 1), dtype=np.float64) L_new[:M, :M] = chol_L L_new[M, :M] = w L_new[M, M] = diag chol_L = L_new def posterior_var(t_idx: int) -> float: """Var(Z_t | current obs_idx).""" if len(obs_idx) == 0: return k_tt(t_idx) idx = np.array(obs_idx, dtype=np.int64) ktO = k_vec(t_idx, idx) alpha = np.linalg.solve(chol_L, ktO) var = k_tt(t_idx) - float(alpha @ alpha) return max(var, 0.0) # Main horizon sweep: at horizon h, before adding obs at h we have O_{h-1} info_gain = np.zeros(T, dtype=np.float64) for h in range(N): # model time corresponding to this horizon if h >= lag: t = h - lag # Var before incorporating potential obs at h var_prev = posterior_var(t) # Incorporate obs at h if present if observed[h]: add_observation(h) # Var after (if censored, obs set unchanged so var_post=var_prev) var_post = posterior_var(t) info_gain[t] = 0.5 * np.log((var_prev + eps) / (var_post + eps)) else: # horizons before we can even define t>=0: still need to update obs set if any, # but in this dataset observed indices start at lag anyway. if observed[h]: add_observation(h) return info_gain def make_censored_task( T: int = 1000, lag: int = 10, source: str = "get_ou_sample", censor_period: int = 40, seed: int = 0, censor_val: float = -3.0, **signal_kwargs, ): np.random.seed(seed) if source == "whitesignal": dt = signal_kwargs.pop("dt", 0.01) freq = signal_kwargs.pop("freq", 1.0) rms = signal_kwargs.pop("rms", 0.5) output_signal = whitesignal(period=(T + lag) * dt, dt=dt, freq=freq, rms=rms, **signal_kwargs) elif source == "ou": output_signal = get_ou_sample(T=(T+lag), **signal_kwargs) elif source == "rbf": output_signal = get_rbf_sample(T=(T+lag), **signal_kwargs) else: raise ValueError(f"Unknown source '{source}'. Use 'whitesignal' or 'ou'.") input_signal = np.copy(output_signal) mask = (np.arange(T+lag) % censor_period) >= (censor_period // 2) input_signal[mask] = censor_val return input_signal[lag:], output_signal[:-lag] def make_copying_task( T: int = 1000, lag: int = 10, source: str = "get_ou_sample", seed: int = 0, **signal_kwargs, ): """ Build a simple sequence-copying task from a continuous signal. The input is generated either by :func:`whitesignal` or :func:`get_ou_sample`. The output is the same signal shifted forward by ``lag`` time steps, forcing a sequence model to retain information over that window to predict correctly. Parameters ---------- T : int, optional Length of the sequence. lag : int, optional Number of time steps to shift the target output relative to the input. source : {"whitesignal", "ou"}, optional Which generator to use. "ou" selects :func:`get_ou_sample`. seed : int, optional Random seed used for reproducibility. **signal_kwargs : Additional keyword arguments forwarded to the signal generator. Returns ------- input_signal : ndarray, shape (T,) The driving input sequence. target_signal : ndarray, shape (T,) The delayed copy of ``input_signal``. """ if lag <= 0: raise ValueError("lag must be positive to form a copying task") if lag >= T: raise ValueError("lag must be smaller than T to produce a valid shift") np.random.seed(seed) if source == "whitesignal": dt = signal_kwargs.pop("dt", 0.01) freq = signal_kwargs.pop("freq", 1.0) rms = signal_kwargs.pop("rms", 0.5) input_signal = whitesignal(period=T * dt, dt=dt, freq=freq, rms=rms, **signal_kwargs) elif source in {"ou", "get_ou_sample"}: input_signal = get_ou_sample(T=T, **signal_kwargs) else: raise ValueError(f"Unknown source '{source}'. Use 'whitesignal' or 'ou'.") target_signal = np.zeros_like(input_signal) target_signal[lag:] = input_signal[: T - lag] return input_signal, target_signal def make_multiplexing_task( T: int = 1000, K: int = 3, lag: int = 3, repeat: int = 1, seed: int = 42 ): np.random.seed(seed) assert T % repeat == 0 multiplex_pattern = np.random.randint(0, K, size=T // repeat) multiplex_pattern = np.repeat(multiplex_pattern, repeat) orig_seqs = np.stack([get_ou_sample(T=T) for _ in range(K)], axis=1) orig_seqs = 1.0 / (1.0 + np.exp(-orig_seqs)) all_seqs = 0.8 * orig_seqs + 0.1 all_seqs = all_seqs + np.arange(K)[None] all_seqs -= np.mean(all_seqs) all_seqs /= np.std(all_seqs) X = np.zeros(T) Y = np.zeros((T, K)) histories = np.zeros((K, lag)) task_idx = np.zeros(K, dtype=int) for i, k in enumerate(multiplex_pattern): X[i] = all_seqs[task_idx[k], k] histories[k] = np.roll(histories[k], -1) histories[k, -1] = orig_seqs[task_idx[k], k] task_idx[k] += 1 Y[i] = np.mean(histories, axis=1) return X, Y def make_induction_head_multioutput_s2s_task( T: int = 1000, V: int = 30, K: int = 5, seed: int = 42 ): np.random.seed(seed) X = np.zeros(T, dtype=int) Y = np.zeros((T, K), dtype=int) X[0] = np.random.randint(0, V) for j in range(1, T): if X[j - 1] < K: X[j] = np.random.randint(K, V) else: X[j] = np.random.randint(0, V) current_vals = np.zeros(K, dtype=int) flag = False for t in range(T): if t > 0: Y[t] = Y[t - 1] if X[t] < K: flag = True elif flag: flag = False Y[t, X[t - 1]] = X[t] vocab = 1.5 ** (1 / 3) * np.linspace(-1, 1, V - K) special_vocab = vocab[0] + np.linspace(-3, -1, K) vocab = np.concatenate([special_vocab, vocab], axis=0) return vocab[X], vocab[Y] def make_implicit_measure_task(T: int = 100, filter_size: int = 15, seed: int = 42): np.random.seed(seed) t = np.linspace(0, 1, T) clean_signal = np.sin(8 * np.pi * t) mask = np.zeros(T) mask[:-filter_size] = 1.0 observed = clean_signal.copy() noise = 0.05 * np.random.randn(T) observed += noise observed[mask == 0] = -4.0 return observed, clean_signal def make_simple_repetition_task(T: int = 100, shift: int = 10, seed: int = 42): np.random.seed(seed) t = np.linspace(0, 1, T) observed = np.sin(8 * np.pi * t) noise = 0.2 * np.random.randn(T) observed += noise output = observed.copy() output = output[:-shift] output = np.concatenate([np.zeros(shift), output]) return observed, output def whitesignal(period, dt, freq, rms=0.5, batch_shape=(), seed=None): """ Copied from github.com/state-spaces/s4 Produces output signal of length period / dt, band-limited to frequency freq Output shape (*batch_shape, period/dt) Adapted from the nengo library """ assert not (freq is not None and freq < 1.0 / period) assert freq <= 0.5 / dt if seed is not None: np.random.seed(seed) n_coefficients = int(np.ceil(period / dt / 2.0)) shape = batch_shape + (n_coefficients + 1,) sigma = rms * np.sqrt(0.5) coefficients = 1j * np.random.normal(0.0, sigma, size=shape) coefficients[..., -1] = 0.0 coefficients += np.random.normal(0.0, sigma, size=shape) coefficients[..., 0] = 0.0 set_to_zero = np.fft.rfftfreq(2 * n_coefficients, d=dt) > freq coefficients *= 1 - set_to_zero power_correction = np.sqrt(1.0 - np.sum(set_to_zero, dtype=float) / n_coefficients) if power_correction > 0: coefficients /= power_correction coefficients *= np.sqrt(2 * n_coefficients) signal = np.fft.irfft(coefficients, axis=-1) return signal def wray_and_green_output(input_signal, a=2.0, m=0.3, k=0.08, tau_max=50, scaling=4e-3): """Implement the system described in Wray and Green (1994).""" T = len(input_signal) # Make the filter. mu = lambda t: a / m * np.exp(-k * t) * np.sin(m * t) tau_vals = np.arange(tau_max) filter = mu(tau_vals)[::-1] filter = scaling * np.outer(filter, filter) # Pad the input signal. output_signal = np.zeros_like(input_signal) input_signal = np.concatenate([np.zeros(tau_max - 1), input_signal], axis=0) # Do a convolution. for i in range(T): input_slice = input_signal[i : i + tau_max] output_signal[i] = np.sum(filter * np.outer(input_slice, input_slice)) return output_signal __all__ = [ "get_ou_sample", "make_copying_task", "make_induction_head_multioutput_s2s_task", "make_induction_head_seq_to_seq_task", "make_induction_head_task", "make_implicit_measure_task", "make_multiplexing_task", "whitesignal", "wray_and_green_output", ]