Add ARC-AGI bulk evaluation script

scripts/run_all_arc.py

@@ -0,0 +1,183 @@
+"""
+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()