| """Materialize superposition instances for every puzzle and every stage. |
| |
| Input: datasets_multicandidate_s12/{split}_cand_masks.npy (N, 12, 81) |
| Output: datasets_superposition/{split}_assignments.npy (M, 81) uint8 |
| datasets_superposition/{split}_index.npy (M, 3) int32 |
| columns: [puzzle_idx, stage, instance_within_stage] |
| |
| M is about 48 x N (one puzzle expands into ~6 instances at stage 0 down to 1 |
| at stage 11). Generation is embarrassingly parallel over puzzles. |
| """ |
| import argparse |
| import os |
| import time |
|
|
| import numpy as np |
| import multiprocessing as mp |
|
|
| import superposition_instances as SI |
|
|
| _MASKS = None |
| _ARGS = None |
|
|
|
|
| def _init(path, args): |
| global _MASKS, _ARGS |
| _MASKS = np.load(path, mmap_mode="r") |
| _ARGS = args |
|
|
|
|
| def _one(idx): |
| assigns, index = [], [] |
| for s in range(_MASKS.shape[1]): |
| mask = np.array(_MASKS[idx, s]).astype(np.uint16) |
| r = SI.instances_for_stage( |
| mask, |
| max_confine=_ARGS.max_confine, |
| max_instances=_ARGS.max_instances, |
| max_attempts=_ARGS.max_attempts, |
| seed=idx * 100 + s, |
| max_repair=_ARGS.max_repair, |
| require_new=True, |
| patience=_ARGS.patience, |
| ) |
| inst = r["instances"] |
| if len(inst) == 0: |
| continue |
| assigns.append(inst.astype(np.uint8)) |
| for k in range(len(inst)): |
| index.append((idx, s, k)) |
| if not assigns: |
| return (np.zeros((0, 81), dtype=np.uint8), |
| np.zeros((0, 3), dtype=np.int32)) |
| return (np.concatenate(assigns, axis=0), |
| np.array(index, dtype=np.int32)) |
|
|
|
|
| def process_split(split, args): |
| mask_path = os.path.join(args.mask_dir, f"{split}_cand_masks.npy") |
| masks = np.load(mask_path, mmap_mode="r") |
| n = len(masks) if args.limit is None else min(args.limit, len(masks)) |
| print(f"[{split}] {n:,} puzzles from {mask_path}", flush=True) |
|
|
| t0 = time.time() |
| chunks_a, chunks_i = [], [] |
| done = 0 |
| with mp.Pool(args.workers, initializer=_init, |
| initargs=(mask_path, args)) as pool: |
| for a, i in pool.imap(_one, range(n), chunksize=8): |
| if len(a): |
| chunks_a.append(a) |
| chunks_i.append(i) |
| done += 1 |
| if done % 2000 == 0 or done == n: |
| rate = done / max(time.time() - t0, 1e-6) |
| kept = sum(len(x) for x in chunks_i) |
| print(f" {done:,}/{n:,} {rate:.0f} puzzles/s " |
| f"{kept:,} instances " |
| f"({kept / done:.1f} per puzzle)", flush=True) |
|
|
| assignments = (np.concatenate(chunks_a, axis=0) if chunks_a |
| else np.zeros((0, 81), dtype=np.uint8)) |
| index = (np.concatenate(chunks_i, axis=0) if chunks_i |
| else np.zeros((0, 3), dtype=np.int32)) |
|
|
| os.makedirs(args.out_dir, exist_ok=True) |
| ap = os.path.join(args.out_dir, f"{split}_assignments.npy") |
| ip = os.path.join(args.out_dir, f"{split}_index.npy") |
| np.save(ap, assignments) |
| np.save(ip, index) |
|
|
| elapsed = time.time() - t0 |
| print(f"[{split}] saved {len(assignments):,} instances " |
| f"({len(assignments) / n:.2f} per puzzle) in {elapsed / 60:.1f} min", |
| flush=True) |
| print(f" {ap} {os.path.getsize(ap) / 1e9:.2f} GB", flush=True) |
| print(f" {ip} {os.path.getsize(ip) / 1e9:.2f} GB", flush=True) |
|
|
| |
| print(f"[{split}] instances per stage:", flush=True) |
| for s in range(masks.shape[1]): |
| c = int((index[:, 1] == s).sum()) |
| print(f" stage {s:>2}: {c:>12,} ({c / n:.2f} per puzzle)", |
| flush=True) |
| return len(assignments) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--mask_dir", default="datasets_multicandidate_s12") |
| ap.add_argument("--out_dir", default="datasets_superposition") |
| ap.add_argument("--splits", default="train,test") |
| ap.add_argument("--limit", type=int, default=None) |
| ap.add_argument("--max-confine", type=int, default=1) |
| ap.add_argument("--max-instances", type=int, default=32) |
| ap.add_argument("--max-attempts", type=int, default=400) |
| ap.add_argument("--max-repair", type=int, default=80) |
| ap.add_argument("--patience", type=int, default=40) |
| ap.add_argument("--workers", type=int, default=64) |
| args = ap.parse_args() |
|
|
| print(f"confine |S|<={args.max_confine} workers={args.workers} " |
| f"out={args.out_dir}", flush=True) |
| for split in args.splits.split(","): |
| process_split(split.strip(), args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
