import numpy as np from dataclasses import dataclass ACTS = ["relu", "tanh", "sin", "abs", "square", "identity", "sigmoid"] def _act(name, x): if name == "relu": return np.maximum(x, 0.0) if name == "tanh": return np.tanh(x) if name == "sin": return np.sin(x) if name == "abs": return np.abs(x) if name == "square": return np.square(x) - 1.0 if name == "sigmoid": return 1.0 / (1.0 + np.exp(-np.clip(x, -30, 30))) return x @dataclass class PriorConfig: max_features: int = 100 min_features: int = 2 max_classes: int = 10 min_rows: int = 128 max_rows: int = 2048 reg_frac: float = 0.35 cat_frac: float = 0.35 irrelevant_frac: float = 0.25 max_depth: int = 5 max_width: int = 32 noise_scale: float = 0.3 def sample_task(cfg: PriorConfig, rng: np.random.Generator): F = int(rng.integers(cfg.min_features, cfg.max_features + 1)) N = int(rng.integers(cfg.min_rows, cfg.max_rows + 1)) is_reg = rng.random() < cfg.reg_frac d_latent = int(rng.integers(2, 12)) h = rng.standard_normal((N, d_latent)).astype(np.float32) if rng.random() < 0.3: h = h * rng.gamma(2.0, 0.5, size=(1, d_latent)).astype(np.float32) depth = int(rng.integers(2, cfg.max_depth + 1)) nodes = [h] cur = h for _ in range(depth): w = int(rng.integers(4, cfg.max_width + 1)) W = rng.standard_normal((cur.shape[1], w)).astype(np.float32) sparsity = rng.uniform(0.3, 0.9) W *= rng.random(W.shape) > sparsity W /= np.sqrt(np.maximum(1.0, (W != 0).sum(0, keepdims=True))) b = rng.standard_normal((1, w)).astype(np.float32) * rng.uniform(0, 1.0) pre = cur @ W + b acts = rng.choice(ACTS, size=w) post = np.empty_like(pre) for j in range(w): post[:, j] = _act(acts[j], pre[:, j]) if rng.random() < 0.3 and w >= 2: i1, i2 = rng.integers(0, w, size=2) post[:, i1] = post[:, i1] * post[:, i2] post = np.clip(post, -30.0, 30.0) nodes.append(post) cur = post allnodes = np.concatenate(nodes, axis=1) total = allnodes.shape[1] tgt_idx = int(rng.integers(max(0, total - cur.shape[1]), total)) target_raw = allnodes[:, tgt_idx].astype(np.float32) k_mix = int(rng.integers(1, min(6, total))) mix_idx = rng.choice(total, size=k_mix, replace=False) mix_w = rng.standard_normal(k_mix).astype(np.float32) target_raw = allnodes[:, mix_idx] @ mix_w target_raw = target_raw + cfg.noise_scale * rng.standard_normal(N).astype( np.float32 ) * (target_raw.std() + 1e-06) feat_pool = [i for i in range(total) if i not in set(mix_idx.tolist())] if len(feat_pool) == 0: feat_pool = list(range(total)) n_real = min(F, len(feat_pool)) chosen = rng.choice(feat_pool, size=n_real, replace=False) X = allnodes[:, chosen].astype(np.float32) if X.shape[1] < F: n_noise = F - X.shape[1] noise_cols = rng.standard_normal((N, n_noise)).astype(np.float32) X = np.concatenate([X, noise_cols], axis=1) n_irr = int(rng.uniform(0, cfg.irrelevant_frac) * F) if n_irr > 0: irr = rng.choice(F, size=min(n_irr, F), replace=False) X[:, irr] = rng.standard_normal((N, len(irr))).astype(np.float32) perm = rng.permutation(F) X = X[:, perm] for c in range(F): if rng.random() < cfg.cat_frac: n_lvl = int(rng.integers(2, 10)) qs = np.quantile(X[:, c], np.linspace(0, 1, n_lvl + 1)[1:-1]) X[:, c] = np.digitize(X[:, c], qs).astype(np.float32) X = X + rng.standard_normal(X.shape).astype(np.float32) * cfg.noise_scale * 0.3 X = np.nan_to_num(X, nan=0.0, posinf=0.0, neginf=0.0) mu = X.mean(0, keepdims=True) sd = X.std(0, keepdims=True) + 1e-06 X = (X - mu) / sd if is_reg: y = np.nan_to_num(target_raw, nan=0.0, posinf=0.0, neginf=0.0) y = (y - y.mean()) / (y.std() + 1e-06) y = y.astype(np.float32) n_classes = 0 task_type = "reg" else: K = int(rng.integers(2, cfg.max_classes + 1)) t = np.nan_to_num(target_raw, nan=0.0, posinf=0.0, neginf=0.0) if rng.random() < 0.5: qs = np.quantile(t, np.linspace(0, 1, K + 1)[1:-1]) y = np.digitize(t, qs).astype(np.int64) else: scores = allnodes @ rng.standard_normal((total, K)).astype(np.float32) y = scores.argmax(1).astype(np.int64) uniq = np.unique(y) if len(uniq) < 2: y = (t > np.median(t)).astype(np.int64) uniq = np.unique(y) remap = {u: i for i, u in enumerate(uniq)} y = np.array([remap[v] for v in y], dtype=np.int64) n_classes = int(len(uniq)) task_type = "cls" feat_mask = np.ones(F, dtype=np.float32) return { "X": X, "y": y, "n_features": F, "n_classes": n_classes, "task_type": task_type, "feat_mask": feat_mask, } def sample_batch(cfg, rng, n_tasks): return [sample_task(cfg, rng) for _ in range(n_tasks)] def _worker(args): seed, n, cfg = args rng = np.random.default_rng(seed) return sample_batch(cfg, rng, n) def gen_pool(cfg, n_tasks, seed=0, workers=None): import multiprocessing as mp if workers is None: workers = min(mp.cpu_count(), 64) per = max(1, n_tasks // workers) jobs = [(seed + i, per, cfg) for i in range(workers)] rem = n_tasks - per * workers if rem > 0: jobs.append((seed + workers, rem, cfg)) with mp.Pool(workers) as pool: outs = pool.map(_worker, jobs) tasks = [t for sub in outs for t in sub] return tasks[:n_tasks] if __name__ == "__main__": cfg = PriorConfig() rng = np.random.default_rng(0) print("=== 8 sampled synthetic tasks ===") kinds = {"cls": 0, "reg": 0} for i in range(8): t = sample_task(cfg, rng) kinds[t["task_type"]] += 1 if t["task_type"] == "cls": _, cnts = np.unique(t["y"], return_counts=True) bal = f"class_counts={cnts.tolist()}" else: bal = f"y_range=[{t['y'].min():.2f},{t['y'].max():.2f}] std={t['y'].std():.2f}" print( f"[{i}] {t['task_type']:3s} N={t['X'].shape[0]:5d} F={t['n_features']:3d} K={t['n_classes']:2d} Xstd={t['X'].std():.2f} Xmean={t['X'].mean():+.3f} {bal}" ) print("task-type mix:", kinds) import time t0 = time.time() tasks = sample_batch(cfg, rng, 500) dt = time.time() - t0 Fs = [t["n_features"] for t in tasks] Ns = [t["X"].shape[0] for t in tasks] Ks = [t["n_classes"] for t in tasks if t["task_type"] == "cls"] regf = np.mean([t["task_type"] == "reg" for t in tasks]) print(f"\n500 tasks in {dt:.2f}s ({500 / dt:.0f} tasks/s single-core)") print(f"F: min={min(Fs)} max={max(Fs)} mean={np.mean(Fs):.1f}") print(f"N: min={min(Ns)} max={max(Ns)} mean={np.mean(Ns):.0f}") print(f"K(cls): min={min(Ks)} max={max(Ks)} mean={np.mean(Ks):.1f}") print(f"regression fraction: {regf:.2f}")