Move own-solver/neurogolf_solver/solvers/wave1.py to own-solver/

own-solver/neurogolf_solver/solvers/wave1.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+"""Wave 1 static spatial remapping solvers.
+
+A4: downsample_stride — strided sampling of input
+A7: symmetry_complete — mirror to complete L-R or T-B symmetry
+A1: extract_inner — remove border frame
+A2: add_border — add constant border
+A6: sparse_fill — pixel to block expansion
+B1: channel_filter — keep only certain colors
+
+Scan results (2026-04-27): 0 arc-gen validated matches.
+Kept for future tasks and as building blocks.
+"""
+
+import numpy as np
+from ..data_loader import get_exs, fixed_shapes
+from ..gather_helpers import _build_gather_model, _build_gather_model_with_const
+from ..onnx_helpers import mk, _make_int64_init, _build_pad_node, add_onehot_block
+from ..constants import GH, GW
+
+
+def s_downsample_stride(td):
+    """out[r,c] = inp[r*sH + oH, c*sW + oW] for integer strides."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    if OH >= IH or OW >= IW:
+        return None
+
+    for sh in range(2, 6):
+        for sw in range(2, 6):
+            for oh_off in range(sh):
+                for ow_off in range(sw):
+                    ok = True
+                    for inp, out in exs:
+                        sampled = inp[oh_off::sh, ow_off::sw]
+                        if sampled.shape != out.shape or not np.array_equal(sampled, out):
+                            ok = False
+                            break
+                    if ok:
+                        idx = np.zeros((OH, OW, 2), dtype=np.int64)
+                        for r in range(OH):
+                            for c in range(OW):
+                                idx[r, c] = [r * sh + oh_off, c * sw + ow_off]
+                        return _build_gather_model(OH, OW, idx)
+    return None
+
+
+def s_symmetry_complete(td):
+    """Complete partial T-B symmetry by adding mirrored + original via Gather."""
+    from onnx import helper, numpy_helper
+
+    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
+    if IH < 2:
+        return None
+
+    # T-B symmetry: out[r,c] = max(inp[r,c], inp[IH-1-r,c])
+    ok = True
+    for inp, out in exs:
+        exp = inp.copy()
+        for r in range(IH // 2):
+            for c in range(IW):
+                v = max(int(inp[r, c]), int(inp[IH - 1 - r, c]))
+                exp[r, c] = v
+                exp[IH - 1 - r, c] = v
+        if not np.array_equal(out, exp):
+            ok = False
+            break
+
+    if ok:
+        # Build: Gather(self) + Gather(mirror) → Add → ArgMax → one-hot
+        pad_h, pad_w = GH - OH, GW - OW
+        mirror_idx = np.zeros((GH * GW,), dtype=np.int64)
+        mask = np.zeros((1, 1, GH, GW), dtype=np.float32)
+        self_idx = np.zeros((GH * GW,), dtype=np.int64)
+        for r in range(OH):
+            for c in range(OW):
+                self_idx[r * GW + c] = r * GW + c
+                mirror_idx[r * GW + c] = (IH - 1 - r) * GW + c
+                mask[0, 0, r, c] = 1.0
+
+        inits = [
+            numpy_helper.from_array(np.array([1, 10, GH * GW], dtype=np.int64), 'fs'),
+            numpy_helper.from_array(self_idx, 'self_idx'),
+            numpy_helper.from_array(mirror_idx, 'mirror_idx'),
+            numpy_helper.from_array(np.array([1, 10, GH, GW], dtype=np.int64), 'os'),
+            numpy_helper.from_array(mask, 'mask'),
+        ]
+        nodes = [
+            helper.make_node('Reshape', ['input', 'fs'], ['flat']),
+            helper.make_node('Gather', ['flat', 'self_idx'], ['g_self'], axis=2),
+            helper.make_node('Gather', ['flat', 'mirror_idx'], ['g_mirror'], axis=2),
+            helper.make_node('Add', ['g_self', 'g_mirror'], ['combined']),
+            helper.make_node('Reshape', ['combined', 'os'], ['combined_2d']),
+            helper.make_node('ArgMax', ['combined_2d'], ['am'], axis=1, keepdims=1),
+        ]
+        add_onehot_block(nodes, inits, 'am', 'oh_out')
+        nodes.append(helper.make_node('Mul', ['oh_out', 'mask'], ['output']))
+        return mk(nodes, inits)
+
+    # L-R symmetry: out[r,c] = max(inp[r,c], inp[r,IW-1-c])
+    if IW < 2:
+        return None
+    ok = True
+    for inp, out in exs:
+        exp = inp.copy()
+        for r in range(IH):
+            for c in range(IW // 2):
+                v = max(int(inp[r, c]), int(inp[r, IW - 1 - c]))
+                exp[r, c] = v
+                exp[r, IW - 1 - c] = v
+        if not np.array_equal(out, exp):
+            ok = False
+            break
+
+    if ok:
+        mirror_idx = np.zeros((GH * GW,), dtype=np.int64)
+        mask = np.zeros((1, 1, GH, GW), dtype=np.float32)
+        self_idx = np.zeros((GH * GW,), dtype=np.int64)
+        for r in range(OH):
+            for c in range(OW):
+                self_idx[r * GW + c] = r * GW + c
+                mirror_idx[r * GW + c] = r * GW + (IW - 1 - c)
+                mask[0, 0, r, c] = 1.0
+
+        inits = [
+            numpy_helper.from_array(np.array([1, 10, GH * GW], dtype=np.int64), 'fs'),
+            numpy_helper.from_array(self_idx, 'self_idx'),
+            numpy_helper.from_array(mirror_idx, 'mirror_idx'),
+            numpy_helper.from_array(np.array([1, 10, GH, GW], dtype=np.int64), 'os'),
+            numpy_helper.from_array(mask, 'mask'),
+        ]
+        nodes = [
+            helper.make_node('Reshape', ['input', 'fs'], ['flat']),
+            helper.make_node('Gather', ['flat', 'self_idx'], ['g_self'], axis=2),
+            helper.make_node('Gather', ['flat', 'mirror_idx'], ['g_mirror'], axis=2),
+            helper.make_node('Add', ['g_self', 'g_mirror'], ['combined']),
+            helper.make_node('Reshape', ['combined', 'os'], ['combined_2d']),
+            helper.make_node('ArgMax', ['combined_2d'], ['am'], axis=1, keepdims=1),
+        ]
+        add_onehot_block(nodes, inits, 'am', 'oh_out')
+        nodes.append(helper.make_node('Mul', ['oh_out', 'mask'], ['output']))
+        return mk(nodes, inits)
+
+    return None
+
+
+def s_extract_inner(td):
+    """Remove N-pixel border frame → smaller output."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+
+    for b in range(1, min(IH, IW) // 2):
+        if OH != IH - 2 * b or OW != IW - 2 * b:
+            continue
+        if all(np.array_equal(inp[b:IH-b, b:IW-b], out) for inp, out in exs):
+            idx = np.zeros((OH, OW, 2), dtype=np.int64)
+            for r in range(OH):
+                for c in range(OW):
+                    idx[r, c] = [r + b, c + b]
+            return _build_gather_model(OH, OW, idx)
+    return None
+
+
+def s_add_border(td):
+    """Add constant-color border frame → larger output."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+
+    for b in range(1, 5):
+        if OH != IH + 2 * b or OW != IW + 2 * b:
+            continue
+        if OH > 30 or OW > 30:
+            continue
+        for bc in range(10):
+            ok = True
+            for inp, out in exs:
+                exp = np.full((OH, OW), bc, dtype=np.int64)
+                exp[b:b+IH, b:b+IW] = inp
+                if not np.array_equal(out, exp):
+                    ok = False
+                    break
+            if ok:
+                idx = np.zeros((OH, OW, 2), dtype=np.int64)
+                cst = np.full((OH, OW), -1, dtype=np.int64)
+                for r in range(OH):
+                    for c in range(OW):
+                        if b <= r < b + IH and b <= c < b + IW:
+                            idx[r, c] = [r - b, c - b]
+                        else:
+                            idx[r, c] = [-1, -1]
+                            cst[r, c] = bc
+                return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst)
+    return None
+
+
+def s_sparse_fill(td):
+    """Each input pixel becomes an NxN block in output."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+
+    for bh in range(2, 10):
+        for bw in range(2, 10):
+            if OH != IH * bh or OW != IW * bw:
+                continue
+            if OH > 30 or OW > 30:
+                continue
+            ok = True
+            for inp, out in exs:
+                exp = np.zeros((OH, OW), dtype=np.int64)
+                for r in range(IH):
+                    for c in range(IW):
+                        exp[r*bh:(r+1)*bh, c*bw:(c+1)*bw] = inp[r, c]
+                if not np.array_equal(out, exp):
+                    ok = False
+                    break
+            if ok:
+                idx = np.zeros((OH, OW, 2), dtype=np.int64)
+                for r in range(OH):
+                    for c in range(OW):
+                        idx[r, c] = [r // bh, c // bw]
+                return _build_gather_model(OH, OW, idx)
+    return None
+
+
+def s_channel_filter(td):
+    """Keep only certain colors, rest → background (0)."""
+    from onnx import helper, numpy_helper
+
+    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
+
+    in_colors = set()
+    out_colors = set()
+    for inp, out in exs:
+        in_colors.update(inp.flatten())
+        out_colors.update(out.flatten())
+
+    if not (out_colors < in_colors):
+        return None
+
+    keep = out_colors
+    for inp, out in exs:
+        exp = np.where(np.isin(inp, list(keep)), inp, 0)
+        if not np.array_equal(out, exp):
+            return None
+
+    ch_mask = np.zeros((1, 10, 1, 1), dtype=np.float32)
+    for c in keep:
+        if 0 <= c < 10:
+            ch_mask[0, c, 0, 0] = 1.0
+
+    inits = [numpy_helper.from_array(ch_mask, 'ch_mask')]
+    nodes = [helper.make_node('Mul', ['input', 'ch_mask'], ['output'])]
+    return mk(nodes, inits)