v5 refactor: add solvers/solver_registry.py

neurogolf_solver/solvers/solver_registry.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+"""Solver registry and task solving orchestration."""
+
+import os
+import time
+import onnx
+from .analytical import s_identity, s_constant, s_color_map, s_transpose
+from .geometric import s_flip, s_rotate, s_shift, s_fixed_crop, s_gravity
+from .tiling import (s_tile, s_upscale, s_kronecker, s_nonuniform_scale, s_diagonal_tile,
+                     s_mirror_h, s_mirror_v, s_quad_mirror, s_concat, s_concat_enhanced,
+                     s_spatial_gather, s_varshape_spatial_gather)
+from .conv import solve_conv_fixed, solve_conv_variable, solve_conv_diffshape, solve_conv_var_diff
+from ..data_loader import get_exs, fixed_shapes
+from ..validators import validate
+from ..constants import EXCLUDED_TASKS
+
+# Analytical solvers registry — order matters (cheaper first)
+ANALYTICAL_SOLVERS = [
+    ('identity', s_identity),
+    ('constant', s_constant),
+    ('color_map', s_color_map),
+    ('transpose', s_transpose),
+    ('flip', s_flip),
+    ('rotate', s_rotate),
+    ('shift', s_shift),
+    ('tile', s_tile),
+    ('upscale', s_upscale),
+    ('kronecker', s_kronecker),
+    ('nonuniform_scale', s_nonuniform_scale),
+    ('mirror_h', s_mirror_h),
+    ('mirror_v', s_mirror_v),
+    ('quad_mirror', s_quad_mirror),
+    ('concat', s_concat),
+    ('concat_enhanced', s_concat_enhanced),
+    ('diagonal_tile', s_diagonal_tile),
+    ('fixed_crop', s_fixed_crop),
+    ('spatial_gather', s_spatial_gather),
+    ('varshape_spatial_gather', s_varshape_spatial_gather),
+]
+
+
+def solve_task(tn, td, outdir, providers, conv_budget=30.0, excluded_tasks=None):
+    """Solve a single ARC-AGI task.
+    
+    Returns: (ok, solver_name, file_size, elapsed, model_path)
+    """
+    if excluded_tasks is None:
+        excluded_tasks = EXCLUDED_TASKS
+    
+    t_start = time.time()
+    os.makedirs(outdir, exist_ok=True)
+    path = os.path.join(outdir, f"task{tn:03d}.onnx")
+
+    if tn in excluded_tasks:
+        return False, 'excluded', None, time.time() - t_start, path
+
+    # 1. Try analytical solvers (fast, tiny models)
+    for sname, sfn in ANALYTICAL_SOLVERS:
+        try:
+            model = sfn(td)
+            if model is None:
+                continue
+            onnx.save(model, path)
+            if validate(path, td, providers):
+                return True, sname, os.path.getsize(path), time.time() - t_start, path
+        except:
+            pass
+
+    # 2. Determine task shape category and try conv solvers
+    exs = get_exs(td)
+    same_shape = all(inp.shape == out.shape for inp, out in exs)
+    shapes = set(inp.shape for inp, _ in exs)
+    fixed_in = len(shapes) == 1
+
+    conv_time = conv_budget
+
+    if same_shape:
+        if fixed_in:
+            result = solve_conv_fixed(td, path, providers, time_budget=conv_time / 2)
+            if result is not None:
+                sname, model = result
+                return True, sname, os.path.getsize(path), time.time() - t_start, path
+        result = solve_conv_variable(td, path, providers, time_budget=conv_time)
+        if result is not None:
+            sname, model = result
+            return True, sname, os.path.getsize(path), time.time() - t_start, path
+    else:
+        sp = fixed_shapes(td)
+        if sp is not None:
+            (IH, IW), (OH, OW) = sp
+            if OH <= IH and OW <= IW:
+                result = solve_conv_diffshape(td, path, providers, time_budget=conv_time)
+                if result is not None:
+                    sname, model = result
+                    return True, sname, os.path.getsize(path), time.time() - t_start, path
+
+        result = solve_conv_var_diff(td, path, providers, time_budget=conv_time)
+        if result is not None:
+            sname, model = result
+            return True, sname, os.path.getsize(path), time.time() - t_start, path
+
+    return False, None, None, time.time() - t_start, path