Remove new_solvers.py (split into gravity.py, edge.py, mode.py)

neurogolf_solver/solvers/new_solvers.py

@@ -1,294 +0,0 @@
-#!/usr/bin/env python3
-"""New solver architectures: gravity, edge detection, mode fill.
-
-These use ONNX ops beyond Conv+lstsq to handle tasks that require
-non-local operations (directional propagation, boundary detection,
-global aggregation).
-
-v5.2 (2026-04-26): gravity_unrolled solves Task 78, mode_fill solves Task 129.
-"""
-
-import numpy as np
-from onnx import helper, numpy_helper, TensorProto
-from ..onnx_helpers import mk, _make_int64_init, _build_pad_node, _build_slice_crop, add_onehot_block
-from ..data_loader import get_exs, fixed_shapes
-from ..constants import GH, GW
-
-
-# ---------------------------------------------------------------------------
-# Gravity solver — unrolled bubble-sort via Conv + Where
-# ---------------------------------------------------------------------------
-
-def _gravity_np(grid, direction, bg_color=0):
-    """Apply gravity in numpy for verification."""
-    r = np.full_like(grid, bg_color)
-    h, w = grid.shape
-    if direction == 'down':
-        for c in range(w):
-            nz = grid[:, c][grid[:, c] != bg_color]
-            r[h - len(nz):h, c] = nz
-    elif direction == 'up':
-        for c in range(w):
-            nz = grid[:, c][grid[:, c] != bg_color]
-            r[:len(nz), c] = nz
-    elif direction == 'right':
-        for rr in range(h):
-            nz = grid[rr, :][grid[rr, :] != bg_color]
-            r[rr, w - len(nz):w] = nz
-    elif direction == 'left':
-        for rr in range(h):
-            nz = grid[rr, :][grid[rr, :] != bg_color]
-            r[rr, :len(nz)] = nz
-    return r
-
-
-def _build_gravity_model(IH, IW, direction, bg_color=0):
-    """Build ONNX model for gravity via unrolled bubble-sort.
-    
-    Each step compares adjacent cells and swaps if needed:
-    - If current cell is bg AND source neighbor is non-bg → fill with source
-    - If current cell is non-bg AND destination neighbor is bg → vacate to bg
-    After max(IH,IW) passes, all non-bg pixels settle in the gravity direction.
-    """
-    pad_h, pad_w = GH - IH, GW - IW
-    n_steps = max(IH, IW)
-    
-    # Two shift kernels: pull from source and destination directions
-    pull_above = np.zeros((10, 10, 3, 3), dtype=np.float32)
-    pull_below = np.zeros((10, 10, 3, 3), dtype=np.float32)
-    for ch in range(10):
-        if direction == 'down':
-            pull_above[ch, ch, 0, 1] = 1.0
-            pull_below[ch, ch, 2, 1] = 1.0
-        elif direction == 'up':
-            pull_above[ch, ch, 2, 1] = 1.0
-            pull_below[ch, ch, 0, 1] = 1.0
-        elif direction == 'right':
-            pull_above[ch, ch, 1, 0] = 1.0
-            pull_below[ch, ch, 1, 2] = 1.0
-        elif direction == 'left':
-            pull_above[ch, ch, 1, 2] = 1.0
-            pull_below[ch, ch, 1, 0] = 1.0
-    
-    bg_sel = np.zeros((1, 10, 1, 1), dtype=np.float32)
-    bg_sel[0, bg_color, 0, 0] = 1.0
-    bg_oh = np.zeros((1, 10, 1, 1), dtype=np.float32)
-    bg_oh[0, bg_color, 0, 0] = 1.0
-    
-    inits = [
-        _make_int64_init('sl_st', [0, 0, 0, 0]),
-        _make_int64_init('sl_en', [1, 10, IH, IW]),
-        numpy_helper.from_array(pull_above, 'pull_src'),
-        numpy_helper.from_array(pull_below, 'pull_dst'),
-        numpy_helper.from_array(bg_sel, 'bg_sel'),
-        numpy_helper.from_array(bg_oh, 'bg_oh'),
-        numpy_helper.from_array(np.float32(0.5), 'half'),
-    ]
-    
-    nodes = [
-        helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cur_0']),
-    ]
-    
-    cur = 'cur_0'
-    for i in range(n_steps):
-        src = f'src_{i}'
-        nodes.append(helper.make_node('Conv', [cur, 'pull_src'], [src],
-                                       kernel_shape=[3, 3], pads=[1, 1, 1, 1]))
-        
-        nodes.append(helper.make_node('Mul', [cur, 'bg_sel'], [f'cbg_{i}']))
-        inits.append(_make_int64_init(f'ax1_{i}', [1]))
-        nodes.append(helper.make_node('ReduceSum', [f'cbg_{i}', f'ax1_{i}'], [f'cbgsum_{i}'], keepdims=1))
-        nodes.append(helper.make_node('Greater', [f'cbgsum_{i}', 'half'], [f'cur_is_bg_{i}']))
-        
-        nodes.append(helper.make_node('Mul', [src, 'bg_sel'], [f'sbg_{i}']))
-        inits.append(_make_int64_init(f'ax2_{i}', [1]))
-        nodes.append(helper.make_node('ReduceSum', [f'sbg_{i}', f'ax2_{i}'], [f'sbgsum_{i}'], keepdims=1))
-        nodes.append(helper.make_node('Not', [f'cur_is_bg_{i}'], [f'cur_not_bg_{i}']))
-        
-        nodes.append(helper.make_node('Greater', [f'sbgsum_{i}', 'half'], [f'src_is_bg_{i}']))
-        nodes.append(helper.make_node('Not', [f'src_is_bg_{i}'], [f'src_not_bg_{i}']))
-        nodes.append(helper.make_node('And', [f'cur_is_bg_{i}', f'src_not_bg_{i}'], [f'fill_{i}']))
-        
-        dst = f'dst_{i}'
-        nodes.append(helper.make_node('Conv', [cur, 'pull_dst'], [dst],
-                                       kernel_shape=[3, 3], pads=[1, 1, 1, 1]))
-        nodes.append(helper.make_node('Mul', [dst, 'bg_sel'], [f'dbg_{i}']))
-        inits.append(_make_int64_init(f'ax3_{i}', [1]))
-        nodes.append(helper.make_node('ReduceSum', [f'dbg_{i}', f'ax3_{i}'], [f'dbgsum_{i}'], keepdims=1))
-        nodes.append(helper.make_node('Greater', [f'dbgsum_{i}', 'half'], [f'dst_is_bg_{i}']))
-        nodes.append(helper.make_node('And', [f'cur_not_bg_{i}', f'dst_is_bg_{i}'], [f'vacate_{i}']))
-        
-        nxt = f'cur_{i+1}'
-        nodes.append(helper.make_node('Where', [f'fill_{i}', src, cur], [f'tmp_{i}']))
-        nodes.append(helper.make_node('Where', [f'vacate_{i}', 'bg_oh', f'tmp_{i}'], [nxt]))
-        cur = nxt
-    
-    # Re-encode as clean one-hot via ArgMax + Equal+Cast, then pad
-    nodes.append(helper.make_node('ArgMax', [cur], ['grav_am'], axis=1, keepdims=1))
-    add_onehot_block(nodes, inits, 'grav_am', 'grav_oh')
-    nodes.append(_build_pad_node('grav_oh', 'output', pad_h, pad_w, inits))
-    return mk(nodes, inits)
-
-
-def s_gravity_unrolled(td):
-    """Gravity solver with unrolled Conv+Where steps.
-    Tries all 4 directions × bg colors 0-9."""
-    exs = get_exs(td)
-    sp = fixed_shapes(td)
-    if sp is None:
-        return None
-    (IH, IW), (OH, OW) = sp
-    if (IH, IW) != (OH, OW):
-        return None
-    
-    for bg_color in range(10):
-        for direction in ('down', 'up', 'left', 'right'):
-            if all(np.array_equal(_gravity_np(inp, direction, bg_color), out) 
-                   for inp, out in exs):
-                return _build_gravity_model(IH, IW, direction, bg_color)
-    return None
-
-
-# ---------------------------------------------------------------------------
-# Edge/boundary detection — Laplacian Conv
-# ---------------------------------------------------------------------------
-
-def _has_edges(inp, out, edge_color, bg_color=0):
-    """Check if output is edge detection of input."""
-    h, w = inp.shape
-    for r in range(h):
-        for c in range(w):
-            pix = inp[r, c]
-            is_edge = False
-            if pix != bg_color:
-                for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
-                    nr, nc = r+dr, c+dc
-                    if 0 <= nr < h and 0 <= nc < w:
-                        if inp[nr, nc] != pix:
-                            is_edge = True
-                            break
-                    else:
-                        is_edge = True
-                        break
-            expected = edge_color if is_edge else bg_color
-            if out[r, c] != expected:
-                return False
-    return True
-
-
-def s_edge_detect(td):
-    """Edge detection solver: output = boundary pixels of input shapes."""
-    exs = get_exs(td)
-    sp = fixed_shapes(td)
-    if sp is None:
-        return None
-    (IH, IW), (OH, OW) = sp
-    if (IH, IW) != (OH, OW):
-        return None
-    
-    for bg_color in [0]:
-        out_colors = set()
-        for _, out in exs:
-            out_colors.update(out.flatten())
-        for edge_color in out_colors:
-            if edge_color == bg_color:
-                continue
-            if all(_has_edges(inp, out, edge_color, bg_color) for inp, out in exs):
-                return _build_edge_model(IH, IW, edge_color, bg_color)
-    return None
-
-
-def _build_edge_model(IH, IW, edge_color, bg_color=0):
-    """Build ONNX model for edge detection via Laplacian conv."""
-    pad_h, pad_w = GH - IH, GW - IW
-    
-    ch_sel = np.zeros((1, 10, 1, 1), dtype=np.float32)
-    for c in range(10):
-        if c != bg_color:
-            ch_sel[0, c, 0, 0] = 1.0
-    
-    lap_k = np.array([[0, -1, 0],
-                       [-1, 4, -1],
-                       [0, -1, 0]], dtype=np.float32).reshape(1, 1, 3, 3)
-    
-    edge_oh = np.zeros((1, 10, 1, 1), dtype=np.float32)
-    edge_oh[0, edge_color, 0, 0] = 1.0
-    bg_oh = np.zeros((1, 10, 1, 1), dtype=np.float32)
-    bg_oh[0, bg_color, 0, 0] = 1.0
-    
-    inits = [
-        _make_int64_init('sl_st', [0, 0, 0, 0]),
-        _make_int64_init('sl_en', [1, 10, IH, IW]),
-        numpy_helper.from_array(ch_sel, 'ch_sel'),
-        numpy_helper.from_array(lap_k, 'lap_k'),
-        numpy_helper.from_array(np.float32(0.5), 'thresh'),
-        numpy_helper.from_array(edge_oh, 'edge_oh'),
-        numpy_helper.from_array(bg_oh, 'bg_oh'),
-    ]
-    
-    nodes = [
-        helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
-        helper.make_node('Conv', ['cropped', 'ch_sel'], ['occ'], kernel_shape=[1, 1]),
-        helper.make_node('Conv', ['occ', 'lap_k'], ['lap_out'], kernel_shape=[3, 3], pads=[1, 1, 1, 1]),
-        helper.make_node('Abs', ['lap_out'], ['lap_abs']),
-        helper.make_node('Greater', ['lap_abs', 'thresh'], ['is_edge_raw']),
-        helper.make_node('Greater', ['occ', 'thresh'], ['is_occ']),
-        helper.make_node('And', ['is_edge_raw', 'is_occ'], ['is_edge']),
-        helper.make_node('Where', ['is_edge', 'edge_oh', 'bg_oh'], ['result_small']),
-    ]
-    nodes.append(_build_pad_node('result_small', 'output', pad_h, pad_w, inits))
-    return mk(nodes, inits)
-
-
-# ---------------------------------------------------------------------------
-# Mode fill solver — output = solid fill of most common input color
-# ---------------------------------------------------------------------------
-
-def s_mode_fill(td):
-    """Mode fill: output is entirely the most common color from input.
-    Uses runtime ArgMax to handle variable mode across inputs."""
-    exs = get_exs(td)
-    
-    for inp, out in exs:
-        if inp.shape != out.shape:
-            return None
-        vals, counts = np.unique(inp, return_counts=True)
-        mode = vals[np.argmax(counts)]
-        if not np.all(out == mode):
-            return None
-    
-    # Check if mode is always the same color
-    modes = set()
-    for inp, out in exs:
-        vals, counts = np.unique(inp, return_counts=True)
-        modes.add(vals[np.argmax(counts)])
-    
-    if len(modes) == 1:
-        return None  # Let s_constant handle it
-    
-    sp = fixed_shapes(td)
-    if sp is None:
-        return None
-    (IH, IW), (OH, OW) = sp
-    if (IH, IW) != (OH, OW):
-        return None
-    
-    pad_h, pad_w = GH - IH, GW - IW
-    
-    inits = [
-        _make_int64_init('sl_st', [0, 0, 0, 0]),
-        _make_int64_init('sl_en', [1, 10, IH, IW]),
-        _make_int64_init('rs_axes_mode', [2, 3]),
-        numpy_helper.from_array(np.arange(10, dtype=np.int64).reshape(1, 10, 1, 1), 'classes'),
-    ]
-    
-    nodes = [
-        helper.make_node('Slice', ['input', 'sl_st', 'sl_en'], ['cropped']),
-        helper.make_node('ReduceSum', ['cropped', 'rs_axes_mode'], ['hist'], keepdims=1),
-        helper.make_node('ArgMax', ['hist'], ['mode_idx'], axis=1, keepdims=1),
-        helper.make_node('Equal', ['mode_idx', 'classes'], ['eq']),
-        helper.make_node('Cast', ['eq'], ['mode_oh'], to=TensorProto.FLOAT),
-        helper.make_node('Expand', ['mode_oh', 'sl_en'], ['expanded']),
-    ]
-    nodes.append(_build_pad_node('expanded', 'output', pad_h, pad_w, inits))
-    return mk(nodes, inits)