prior.py

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}")