File size: 4,644 Bytes
bb23b91 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """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)
# per-stage instance counts
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()
|