File size: 6,824 Bytes
feb08d1 | 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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """
Run the PEMF solver on all ARC-AGI tasks and report solve rates.
For each task, the solver tries every training pair. A task is "solved"
if the solver achieves σ=0 on ALL training pairs.
Usage:
1. Download the ARC dataset into arc_data/training/:
git clone https://github.com/fchollet/ARC-AGI.git /tmp/arc
cp -r /tmp/arc/data/training arc_data/training
2. Run:
python scripts/run_all_arc.py
Outputs:
arc_results/summary.json — per-task results
arc_results/report.txt — human-readable report
"""
import os, json, time, glob
import numpy as np
from itt_solver.solver_core import initialize_potential, sigma_l1
from itt_solver.beam_logging import beam_minimize_with_log
from itt_solver.experiment_driver import default_atomic_factory
ARC_DIR = os.environ.get("ARC_DIR", "arc_data/training")
OUT_DIR = os.environ.get("OUT_DIR", "arc_results")
os.makedirs(OUT_DIR, exist_ok=True)
PARAMS = {
'beam_width': 8,
'max_depth': 2,
'lock_coeff': 0.0,
'max_fraction': 1.0,
'use_symmetry': True,
'use_gravity': True,
'use_color_ops': True,
'boundary_source': 'target',
}
def solve_pair(inp, out, params):
"""Run solver on one input→output pair. Returns (sigma, transform_name, time_s)."""
h, w = len(out), len(out[0])
task = {
'name': 'pair',
'input': inp,
'target': out,
'target_shape': (h, w),
}
atomic_lib = default_atomic_factory(params, task)
phi_in = initialize_potential(inp)
phi_target = initialize_potential(out)
start = time.time()
T_best, phi_best, states, sigmas, logs = beam_minimize_with_log(
phi_in, phi_target, atomic_lib,
beam_width=params['beam_width'],
max_depth=params['max_depth'],
lock_coeff=params['lock_coeff'],
max_fraction=params['max_fraction'],
allowed_symbols=list(range(10)),
enable_layer_minus_one=False,
boundary_source=params['boundary_source'],
)
elapsed = time.time() - start
final_sigma = float(sigmas[-1]) if sigmas else float('inf')
return final_sigma, repr(T_best), elapsed
def run_all():
task_files = sorted(glob.glob(os.path.join(ARC_DIR, "*.json")))
print(f"Running solver on {len(task_files)} ARC training tasks...")
print(f"Params: beam_width={PARAMS['beam_width']}, max_depth={PARAMS['max_depth']}")
print()
results = []
solved_count = 0
partial_count = 0
total_time = 0
for ti, tf in enumerate(task_files):
task_id = os.path.basename(tf).replace('.json', '')
with open(tf) as fh:
task_data = json.load(fh)
train_pairs = task_data.get('train', [])
test_pairs = task_data.get('test', [])
pair_results = []
all_zero = True
best_sigma = float('inf')
best_transform = None
for pi, pair in enumerate(train_pairs):
sigma, transform, elapsed = solve_pair(pair['input'], pair['output'], PARAMS)
total_time += elapsed
pair_results.append({
'pair': pi, 'sigma': sigma,
'transform': transform, 'time_s': round(elapsed, 4),
})
if sigma > 0:
all_zero = False
if sigma < best_sigma:
best_sigma = sigma
best_transform = transform
test_results = []
test_solved = None
for pi, pair in enumerate(test_pairs):
if 'output' in pair:
sigma, transform, elapsed = solve_pair(pair['input'], pair['output'], PARAMS)
total_time += elapsed
test_results.append({
'pair': pi, 'sigma': sigma,
'transform': transform, 'time_s': round(elapsed, 4),
})
if test_solved is None:
test_solved = True
if sigma > 0:
test_solved = False
status = "SOLVED" if all_zero else "PARTIAL" if best_sigma < float('inf') and best_sigma > 0 else "FAILED"
if all_zero:
solved_count += 1
elif best_sigma < float('inf'):
partial_count += 1
results.append({
'task_id': task_id, 'status': status,
'train_pairs': len(train_pairs), 'all_train_solved': all_zero,
'best_sigma': best_sigma, 'best_transform': best_transform,
'pair_results': pair_results,
'test_results': test_results, 'test_solved': test_solved,
})
if (ti + 1) % 20 == 0 or all_zero:
marker = "✅" if all_zero else " "
print(f"[{ti+1:3d}/{len(task_files)}] {task_id}: {status} (best σ={best_sigma:.1f}) {marker}")
failed_count = len(task_files) - solved_count - partial_count
print(f"\n{'='*60}")
print(f"RESULTS: {len(task_files)} tasks")
print(f" SOLVED (σ=0 all train pairs): {solved_count} ({100*solved_count/len(task_files):.1f}%)")
print(f" PARTIAL (σ>0 but finite): {partial_count}")
print(f" FAILED: {failed_count}")
print(f" Total time: {total_time:.1f}s ({total_time/len(task_files):.2f}s/task)")
summary = {
'total_tasks': len(task_files), 'solved': solved_count,
'partial': partial_count, 'failed': failed_count,
'solve_rate': round(100 * solved_count / len(task_files), 2),
'params': PARAMS, 'total_time_s': round(total_time, 2),
'results': results,
}
with open(os.path.join(OUT_DIR, 'summary.json'), 'w') as fh:
json.dump(summary, fh, indent=2)
solved_tasks = [r for r in results if r['all_train_solved']]
print(f"\nSolved tasks:")
for r in solved_tasks:
print(f" {r['task_id']}: {r['best_transform']}")
partial_tasks = sorted(
[r for r in results if not r['all_train_solved'] and r['best_sigma'] < float('inf')],
key=lambda r: r['best_sigma']
)
print(f"\nTop 20 closest-to-solving:")
for r in partial_tasks[:20]:
print(f" {r['task_id']}: σ={r['best_sigma']:.1f} ({r['best_transform']})")
with open(os.path.join(OUT_DIR, 'report.txt'), 'w') as fh:
fh.write(f"PEMF Solver — ARC-AGI Training Set Results\n{'='*60}\n")
fh.write(f"Total tasks: {len(task_files)}\n")
fh.write(f"Solved: {solved_count} ({100*solved_count/len(task_files):.1f}%)\n")
fh.write(f"Partial: {partial_count}\nFailed: {failed_count}\n")
fh.write(f"Time: {total_time:.1f}s\n\n")
fh.write(f"Params: {json.dumps(PARAMS, indent=2)}\n\n")
fh.write(f"Solved tasks:\n")
for r in solved_tasks:
fh.write(f" {r['task_id']}: {r['best_transform']}\n")
print(f"\nResults saved to {OUT_DIR}/")
if __name__ == '__main__':
run_all()
|