circuit/evolve_fast.py

"""
Fixed-K batched GPU evolution. All individuals have exactly K dims.

Genome: (POP, K) int tensor — indices into 768 feature dims.
Fitness: one batched torch.linalg.solve over (POP, K+1, K+1).
Target: hundreds of gen/s.
"""

import json, os, sys, time
import torch

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
REPO_ROOT = os.path.dirname(SCRIPT_DIR)
VAL_TENSORS = os.path.join(REPO_ROOT, "analytical_stats_cache", "val_tensors_layernorm_500.pt")
GREEDY_PATH = os.path.join(REPO_ROOT, "analytical_variants", "greedy_forward_gpu.json")
PERSON_CLASS = 0
DEVICE = "cuda"


@torch.no_grad()
def batch_fitness(features, is_person, pop_dims, lam=0.1):
    """Evaluate entire population in one batched solve.

    features: (N, 768)
    is_person: (N,) bool
    pop_dims: (POP, K) long — which dims each individual uses
    Returns: (POP,) F1 scores
    """
    POP, K = pop_dims.shape
    N = features.shape[0]

    # Gather features for all individuals: (POP, N, K)
    f_batch = features[:, :].unsqueeze(0).expand(POP, -1, -1)  # (POP, N, 768)
    idx = pop_dims.unsqueeze(1).expand(-1, N, -1)  # (POP, N, K)
    f_sub = torch.gather(f_batch, 2, idx)  # (POP, N, K)

    # Augment with bias column: (POP, N, K+1)
    ones = torch.ones(POP, N, 1, device=DEVICE)
    fa = torch.cat([f_sub, ones], dim=2)

    # XtX: (POP, K+1, K+1) = fa^T @ fa
    XtX = torch.bmm(fa.transpose(1, 2), fa)

    # Regularize
    I = torch.eye(K + 1, device=DEVICE).unsqueeze(0).expand(POP, -1, -1)
    XtX = XtX + lam * I * N

    # XtY: (POP, K+1, 1) = fa^T @ y
    y = is_person.float().unsqueeze(0).unsqueeze(2).expand(POP, -1, -1)  # (POP, N, 1)
    XtY = torch.bmm(fa.transpose(1, 2), y)

    # Batched solve: (POP, K+1, 1)
    try:
        W = torch.linalg.solve(XtX, XtY)
    except Exception:
        return torch.zeros(POP, device=DEVICE)

    # Predict: (POP, N, 1)
    scores = torch.bmm(fa, W)
    pred = scores.squeeze(2) > 0.5  # (POP, N)

    # F1 per individual
    is_p = is_person.unsqueeze(0).expand(POP, -1)
    tp = (pred & is_p).sum(dim=1).float()
    fp = (pred & ~is_p).sum(dim=1).float()
    fn = (~pred & is_p).sum(dim=1).float()
    prec = tp / (tp + fp).clamp(min=1)
    rec = tp / (tp + fn).clamp(min=1)
    f1 = 2 * prec * rec / (prec + rec).clamp(min=1e-9)
    return f1


def main():
    print("=" * 60)
    print("Fixed-K Batched GPU Evolution")
    print("=" * 60, flush=True)

    val = torch.load(VAL_TENSORS, map_location="cpu", weights_only=False)
    features = val["features"]
    is_person = (val["cls_targets"] == PERSON_CLASS)

    pos_idx = is_person.nonzero(as_tuple=True)[0]
    neg_idx = (~is_person).nonzero(as_tuple=True)[0]
    n_take = min(2000, len(pos_idx))
    sel = torch.cat([pos_idx[torch.randperm(len(pos_idx))[:n_take]],
                     neg_idx[torch.randperm(len(neg_idx))[:n_take]]])
    sel = sel[torch.randperm(len(sel))]
    sub_f = features[sel].to(DEVICE)
    sub_person = is_person[sel].to(DEVICE)
    N = len(sel)
    print(f"  {N} vectors on {DEVICE}", flush=True)

    greedy_dims = list(range(100))
    if os.path.isfile(GREEDY_PATH):
        with open(GREEDY_PATH) as f:
            greedy_dims = json.load(f)["selected_dims"]

    POP = 512
    GEN = 5000
    ELITE = 30
    TARGETS = [10, 20, 50, 100, 200]

    all_results = []

    for K in TARGETS:
        print(f"\n{'='*60}")
        print(f"  K={K} dims | pop={POP} | gen={GEN}")
        print(f"{'='*60}", flush=True)
        t0 = time.time()

        # Initialize population: (POP, K) long tensors
        pop = torch.zeros(POP, K, dtype=torch.long, device=DEVICE)

        # Seed 0: greedy
        g = greedy_dims[:K] if K <= len(greedy_dims) else greedy_dims + list(range(K - len(greedy_dims)))
        pop[0] = torch.tensor(g[:K], device=DEVICE)

        # Rest: random K-subsets of 768
        for i in range(1, POP):
            pop[i] = torch.randperm(768, device=DEVICE)[:K]

        fits = batch_fitness(sub_f, sub_person, pop)
        best_f1 = fits.max().item()
        best_genome = pop[fits.argmax()].clone()
        stag = 0

        for gen in range(GEN):
            # Sort
            order = fits.argsort(descending=True)
            pop = pop[order]
            fits = fits[order]

            if fits[0].item() > best_f1:
                best_f1 = fits[0].item()
                best_genome = pop[0].clone()
                stag = 0
            else:
                stag += 1

            # New population
            new_pop = pop[:ELITE].clone()

            # Immigration
            n_imm = POP // 5 if stag > 200 else 0
            if n_imm > 0:
                imm = torch.stack([torch.randperm(768, device=DEVICE)[:K] for _ in range(n_imm)])
                new_pop = torch.cat([new_pop, imm])

            # Breed
            n_breed = POP - new_pop.shape[0]
            # Tournament selection
            t1 = torch.randint(0, POP // 2, (n_breed, 5), device=DEVICE)
            p1_idx = t1[torch.arange(n_breed, device=DEVICE), fits[t1].argmax(dim=1)]
            t2 = torch.randint(0, POP // 2, (n_breed, 5), device=DEVICE)
            p2_idx = t2[torch.arange(n_breed, device=DEVICE), fits[t2].argmax(dim=1)]

            parents1 = pop[p1_idx]  # (n_breed, K)
            parents2 = pop[p2_idx]

            # Crossover: for each position, pick from parent1 or parent2
            mask = torch.rand(n_breed, K, device=DEVICE) < 0.5
            children = torch.where(mask, parents1, parents2)

            # Mutation: replace random positions with random dims
            mut_rate = 0.05 * (1 + stag / 100)
            mut_mask = torch.rand(n_breed, K, device=DEVICE) < mut_rate
            random_dims = torch.randint(0, 768, (n_breed, K), device=DEVICE)
            children = torch.where(mut_mask, random_dims, children)

            new_pop = torch.cat([new_pop, children])[:POP]
            pop = new_pop
            fits = batch_fitness(sub_f, sub_person, pop)

            if (gen + 1) % 100 == 0:
                elapsed = time.time() - t0
                gen_s = (gen + 1) / elapsed
                print(f"  gen {gen+1:5d}: best={fits.max().item():.4f} "
                      f"best_ever={best_f1:.4f} stag={stag} "
                      f"{gen_s:.0f} gen/s", flush=True)

            if stag > 1000:
                print(f"  Converged at gen {gen+1}")
                break

        elapsed = time.time() - t0
        best_dims = best_genome.cpu().tolist()
        gates = K * 85
        gens_done = gen + 1
        print(f"\n  WINNER: {K} dims, F1={best_f1:.4f}, {gates} gates, "
              f"{elapsed:.1f}s, {gens_done/elapsed:.0f} gen/s", flush=True)

        all_results.append({
            "K": K, "best_f1": round(best_f1, 4), "genome": sorted(best_dims),
            "gates": gates, "time_s": round(elapsed, 1),
            "generations": gens_done, "gen_per_s": round(gens_done / elapsed),
        })

    print(f"\n{'='*60}")
    print("Results:")
    for r in all_results:
        print(f"  K={r['K']:3d}  {r['gates']:6d} gates  F1={r['best_f1']:.4f}  "
              f"{r['gen_per_s']} gen/s  ({r['generations']} gen, {r['time_s']}s)")

    out = os.path.join(SCRIPT_DIR, "evolved_extreme.json")
    with open(out, "w") as f:
        json.dump(all_results, f, indent=2)
    print(f"Saved: {out}")


if __name__ == "__main__":
    main()