Add gravity, edge detect, mode fill solvers (v5.2)

New solvers in new_solvers.py:
- gravity_unrolled: bubble-sort via Conv+Where, 4 directions, any bg color
Validated: Task 78 solved (score 8.399)
- edge_detect: Laplacian conv + threshold (0 matches in current task set)
- mode_fill: ReduceSum histogram + ArgMax + Expand
Validated: Task 129 solved (score 19.451)

Full 400-task run: 52 solved (was 49), score 709.5, est LB 1057.5
No regressions on existing 49 tasks."

neurogolf_solver/solvers/new_solvers.py

@@ -0,0 +1,294 @@
+#!/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)