medal-solvers/build_task025_onnx.py

"""Build optimized ONNX model for Task 025.

Algorithm (266/266 verified):
- Detect h/v lines for each color
- For each dot of color c, project it adjacent to nearest same-color line
- Only one line per color (never both h and v for same color)

Score target: ~11-12 pts (from base 8.71)
"""
import sys
sys.path.insert(0, '/app/repo/medal-solvers')

from onnx import TensorProto
import numpy as np
from onnx_builder import OnnxBuilder

H, W, C = 30, 30, 10


def build_task025():
    b = OnnxBuilder()
    const, nd = b.const, b.nd

    # === CONSTANTS ===
    const('c_half', [0.5])
    const('c_one', [1.0])
    const('c_zero', [0.0])
    
    const('axes1', [1], 'i')
    const('axes2', [2], 'i')
    const('axes3', [3], 'i')
    
    const('shape_1_1_30_30', [1, 1, 30, 30], 'i')
    const('shape_1_1_30_1', [1, 1, 30, 1], 'i')
    const('shape_1_1_1_30', [1, 1, 1, 30], 'i')
    const('shape_1_1_1_1', [1, 1, 1, 1], 'i')
    const('shape_1_10_1_1', [1, 10, 1, 1], 'i')
    
    # Row/col index grids as [1,1,30,30]
    row_grid = np.arange(30, dtype=np.float32).reshape(1, 1, 30, 1) * np.ones((1, 1, 1, 30), dtype=np.float32)
    col_grid = np.arange(30, dtype=np.float32).reshape(1, 1, 1, 30) * np.ones((1, 1, 30, 1), dtype=np.float32)
    const('row_grid', row_grid)
    const('col_grid', col_grid)
    
    const('depth_10', [10.0])
    const('oh_vals', [0.0, 1.0])
    
    for i in range(10):
        const(f'idx_{i}', [i], 'i')

    # === STEP 1: Grid active mask [1,1,30,30] ===
    active_sum = nd('ReduceSum', ['input', 'axes1'], [[1, 1, 30, 30]], keepdims=1)
    active_b = nd('Greater', [active_sum, 'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)])
    active = nd('Cast', [active_b], [[1, 1, 30, 30]], to=1)
    
    # Grid width per row [1,1,30,1]
    grid_w = nd('ReduceSum', [active, 'axes3'], [[1, 1, 30, 1]], keepdims=1)
    # Grid height per col [1,1,1,30]
    grid_h = nd('ReduceSum', [active, 'axes2'], [[1, 1, 1, 30]], keepdims=1)

    # === STEP 2: Process each color ===
    const('zeros_1_10_30_30', np.zeros((1, 10, 30, 30), dtype=np.float32))
    delta = 'zeros_1_10_30_30'
    
    for c in range(1, 10):
        # Channel c mask [1,1,30,30]
        ch = nd('Gather', ['input', f'idx_{c}'], [[1, 1, 30, 30]], axis=1)
        
        # Row sum of color c [1,1,30,1]
        ch_row_sum = nd('ReduceSum', [ch, 'axes3'], [[1, 1, 30, 1]], keepdims=1)
        # Col sum of color c [1,1,1,30]
        ch_col_sum = nd('ReduceSum', [ch, 'axes2'], [[1, 1, 1, 30]], keepdims=1)
        
        # H-line: row r is a line if ch_row_sum[r] == grid_w[r] > 0
        diff_h = nd('Sub', [ch_row_sum, grid_w], [[1, 1, 30, 1]])
        abs_diff_h = nd('Abs', [diff_h], [[1, 1, 30, 1]])
        h_match_b = nd('Less', [abs_diff_h, 'c_half'], [([1, 1, 30, 1], TensorProto.BOOL)])
        gw_pos_b = nd('Greater', [grid_w, 'c_half'], [([1, 1, 30, 1], TensorProto.BOOL)])
        h_line_b = nd('And', [h_match_b, gw_pos_b], [([1, 1, 30, 1], TensorProto.BOOL)])
        h_line = nd('Cast', [h_line_b], [[1, 1, 30, 1]], to=1)  # [1,1,30,1]
        
        # V-line: col is a line if ch_col_sum[col] == grid_h[col] > 0
        diff_v = nd('Sub', [ch_col_sum, grid_h], [[1, 1, 1, 30]])
        abs_diff_v = nd('Abs', [diff_v], [[1, 1, 1, 30]])
        v_match_b = nd('Less', [abs_diff_v, 'c_half'], [([1, 1, 1, 30], TensorProto.BOOL)])
        gh_pos_b = nd('Greater', [grid_h, 'c_half'], [([1, 1, 1, 30], TensorProto.BOOL)])
        v_line_b = nd('And', [v_match_b, gh_pos_b], [([1, 1, 1, 30], TensorProto.BOOL)])
        v_line = nd('Cast', [v_line_b], [[1, 1, 1, 30]], to=1)  # [1,1,1,30]
        
        # has_h_line / has_v_line (scalar [1,1,1,1])
        has_h = nd('ReduceMax', [h_line, 'axes2'], [[1, 1, 1, 1]], keepdims=1)
        has_v = nd('ReduceMax', [v_line, 'axes3'], [[1, 1, 1, 1]], keepdims=1)
        
        # H-line position: weighted sum of h_line * row_grid[...,0]
        # h_line is [1,1,30,1], row_grid is [1,1,30,30]
        # Multiply and sum over row dim (axis 2)
        # Use row values directly: [1,1,30,1] * row_indices [1,1,30,1]
        const(f'row_idx_col_{c}', np.arange(30, dtype=np.float32).reshape(1, 1, 30, 1))
        h_pos = nd('ReduceSum', [nd('Mul', [h_line, f'row_idx_col_{c}'], [[1, 1, 30, 1]]), 'axes2'],
                   [[1, 1, 1, 1]], keepdims=1)  # [1,1,1,1] = line row position
        
        # V-line position
        const(f'col_idx_row_{c}', np.arange(30, dtype=np.float32).reshape(1, 1, 1, 30))
        v_pos = nd('ReduceSum', [nd('Mul', [v_line, f'col_idx_row_{c}'], [[1, 1, 1, 30]]), 'axes3'],
                   [[1, 1, 1, 1]], keepdims=1)  # [1,1,1,1] = line col position
        
        # On-line mask [1,1,30,30]: broadcast h_line [1,1,30,1] with v_line [1,1,1,30]
        on_line_mask = nd('Max', [h_line, v_line], [[1, 1, 30, 30]])  # broadcasts correctly!
        on_line_clip = nd('Clip', [on_line_mask, 'c_zero', 'c_one'], [[1, 1, 30, 30]])
        not_on_line = nd('Sub', ['c_one', on_line_clip], [[1, 1, 30, 30]])
        
        # Dots = color c pixels not on a line
        dots = nd('Mul', [ch, not_on_line], [[1, 1, 30, 30]])
        
        # === H-line projection ===
        # Dots above line: row < h_pos
        above_b = nd('Less', ['row_grid', nd('Sub', [h_pos, 'c_half'], [[1, 1, 1, 1]])],
                    [([1, 1, 30, 30], TensorProto.BOOL)])
        above_f = nd('Cast', [above_b], [[1, 1, 30, 30]], to=1)
        below_f = nd('Sub', ['c_one', above_f], [[1, 1, 30, 30]])
        
        # Collapse dots above/below to column indicators
        dots_above = nd('Mul', [dots, above_f], [[1, 1, 30, 30]])
        dots_below = nd('Mul', [dots, below_f], [[1, 1, 30, 30]])
        
        # Which cols have dots above/below?
        cols_above = nd('ReduceMax', [dots_above, 'axes2'], [[1, 1, 1, 30]], keepdims=1)  # [1,1,1,30]
        cols_below = nd('ReduceMax', [dots_below, 'axes2'], [[1, 1, 1, 30]], keepdims=1)
        
        # Target rows
        h_target_above = nd('Sub', [h_pos, 'c_one'], [[1, 1, 1, 1]])  # row above line
        h_target_below = nd('Add', [h_pos, 'c_one'], [[1, 1, 1, 1]])  # row below line
        
        # Masks for target rows [1,1,30,30]
        at_above_row = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', h_target_above], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]),
                                    'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)])
        at_above_row_f = nd('Cast', [at_above_row], [[1, 1, 30, 30]], to=1)
        
        at_below_row = nd('Less', [nd('Abs', [nd('Sub', ['row_grid', h_target_below], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]),
                                    'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)])
        at_below_row_f = nd('Cast', [at_below_row], [[1, 1, 30, 30]], to=1)
        
        # H-projection: place at target row for cols with dots
        h_proj_above = nd('Mul', [at_above_row_f, cols_above], [[1, 1, 30, 30]])
        h_proj_below = nd('Mul', [at_below_row_f, cols_below], [[1, 1, 30, 30]])
        h_proj = nd('Add', [h_proj_above, h_proj_below], [[1, 1, 30, 30]])
        h_proj_gated = nd('Mul', [h_proj, has_h], [[1, 1, 30, 30]])
        
        # === V-line projection ===
        left_b = nd('Less', ['col_grid', nd('Sub', [v_pos, 'c_half'], [[1, 1, 1, 1]])],
                   [([1, 1, 30, 30], TensorProto.BOOL)])
        left_f = nd('Cast', [left_b], [[1, 1, 30, 30]], to=1)
        right_f = nd('Sub', ['c_one', left_f], [[1, 1, 30, 30]])
        
        dots_left = nd('Mul', [dots, left_f], [[1, 1, 30, 30]])
        dots_right = nd('Mul', [dots, right_f], [[1, 1, 30, 30]])
        
        rows_left = nd('ReduceMax', [dots_left, 'axes3'], [[1, 1, 30, 1]], keepdims=1)
        rows_right = nd('ReduceMax', [dots_right, 'axes3'], [[1, 1, 30, 1]], keepdims=1)
        
        v_target_left = nd('Sub', [v_pos, 'c_one'], [[1, 1, 1, 1]])
        v_target_right = nd('Add', [v_pos, 'c_one'], [[1, 1, 1, 1]])
        
        at_left_col = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', v_target_left], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]),
                                   'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)])
        at_left_col_f = nd('Cast', [at_left_col], [[1, 1, 30, 30]], to=1)
        
        at_right_col = nd('Less', [nd('Abs', [nd('Sub', ['col_grid', v_target_right], [[1, 1, 30, 30]])], [[1, 1, 30, 30]]),
                                    'c_half'], [([1, 1, 30, 30], TensorProto.BOOL)])
        at_right_col_f = nd('Cast', [at_right_col], [[1, 1, 30, 30]], to=1)
        
        v_proj_left = nd('Mul', [at_left_col_f, rows_left], [[1, 1, 30, 30]])
        v_proj_right = nd('Mul', [at_right_col_f, rows_right], [[1, 1, 30, 30]])
        v_proj = nd('Add', [v_proj_left, v_proj_right], [[1, 1, 30, 30]])
        v_proj_gated = nd('Mul', [v_proj, has_v], [[1, 1, 30, 30]])
        
        # Combined projection
        proj = nd('Add', [h_proj_gated, v_proj_gated], [[1, 1, 30, 30]])
        
        # Net change for this color: ALL dots are removed, proj added only if has_line
        # dots are always removed (they're not in the output regardless of whether there's a line)
        net_change = nd('Sub', [proj, dots], [[1, 1, 30, 30]])
        
        # Expand to 10 channels
        c_oh = nd('OneHot', [f'idx_{c}', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
        c_oh_4d = nd('Reshape', [c_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
        ch0_oh = nd('OneHot', ['idx_0', 'depth_10', 'oh_vals'], [[1, 10]], axis=1)
        ch0_oh_4d = nd('Reshape', [ch0_oh, 'shape_1_10_1_1'], [[1, 10, 1, 1]])
        
        # delta_c = net_change * (c_onehot - ch0_onehot)
        oh_diff = nd('Sub', [c_oh_4d, ch0_oh_4d], [[1, 10, 1, 1]])
        color_delta = nd('Mul', [net_change, oh_diff], [[1, 10, 30, 30]])
        
        delta = nd('Add', [delta, color_delta], [[1, 10, 30, 30]])
    
    # Output = input + delta, clipped to [0,1]
    final = nd('Add', ['input', delta], [[1, 10, 30, 30]])
    final_clipped = nd('Clip', [final, 'c_zero', 'c_one'], [[1, 10, 30, 30]])
    
    return b.finish('task025', last_tensor=final_clipped)


if __name__ == '__main__':
    import os, json, math
    import numpy as np
    import onnx
    import onnxruntime as ort
    
    print("Building Task 025 ONNX model...")
    model = build_task025()
    
    os.makedirs('/app/repo/medal-solvers/optimized', exist_ok=True)
    output_path = '/app/repo/medal-solvers/optimized/task025.onnx'
    onnx.save(model, output_path)
    
    fsize = os.path.getsize(output_path)
    print(f"  Nodes: {len(model.graph.node)}")
    print(f"  File size: {fsize:,} bytes (limit: 1,509,949)")
    
    # Validate
    sess = ort.InferenceSession(output_path)
    
    with open('/app/task-data/task025.json') as f:
        data = json.load(f)
    
    all_examples = data['train'] + data['test'] + data['arc-gen']
    right_count, wrong_count = 0, 0
    
    for i, ex in enumerate(all_examples):
        inp_grid = ex['input']
        inp = np.zeros((1, 10, 30, 30), dtype=np.float32)
        for r, row in enumerate(inp_grid):
            for ci, v in enumerate(row):
                if r < 30 and ci < 30:
                    inp[0][v][r][ci] = 1.0
        
        result = sess.run(['output'], {'input': inp})
        out = (result[0] > 0.0).astype(float)
        
        exp = np.zeros((1, 10, 30, 30), dtype=np.float32)
        for r, row in enumerate(ex['output']):
            for ci, v in enumerate(row):
                if r < 30 and ci < 30:
                    exp[0][v][r][ci] = 1.0
        
        if np.array_equal(out, exp):
            right_count += 1
        else:
            wrong_count += 1
            if wrong_count <= 5:
                diff_locs = np.where(out != exp)
                print(f"  FAIL {i}: {len(diff_locs[0])} diffs")
                for d in range(min(3, len(diff_locs[0]))):
                    ch_d, r_d, c_d = diff_locs[1][d], diff_locs[2][d], diff_locs[3][d]
                    print(f"    ch={ch_d} ({r_d},{c_d}): pred={out[0,ch_d,r_d,c_d]} exp={exp[0,ch_d,r_d,c_d]}")
    
    print(f"\nResults: {right_count} pass, {wrong_count} fail out of {len(all_examples)}")
    
    if wrong_count == 0:
        params = sum(int(np.prod(init.dims)) for init in model.graph.initializer)
        mem_est = sum(int(np.prod([d.dim_value for d in vi.type.tensor_type.shape.dim])) * 4
                      for vi in model.graph.value_info
                      if vi.type.tensor_type.shape.dim)
        score_est = max(1.0, 25.0 - math.log(max(1.0, mem_est + params)))
        print(f"  Params: {params:,}, Memory est: {mem_est:,}, Score est: {score_est:.3f}")