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

own-solver/neurogolf_solver/solvers/geometric.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+"""Geometric transformation solvers: flip, rotate, shift, crop."""
+
+import numpy as np
+from onnx import helper
+from ..onnx_helpers import mk, _build_slice_crop, _build_slice_reverse, _build_pad_node
+from ..data_loader import get_exs, fixed_shapes
+from ..gather_helpers import _build_gather_model, _build_gather_model_with_const
+from ..constants import GH, GW
+
+
+def s_flip(td):
+    """Flip using Slice(step=-1) — zero MACs."""
+    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 axis, flip_fn in [(0, np.flipud), (1, np.fliplr)]:
+        if all(np.array_equal(out, flip_fn(inp)) for inp, out in exs):
+            onnx_axis = 2 if axis == 0 else 3
+            dim_size = IH if axis == 0 else IW
+            pad_h, pad_w = GH - IH, GW - IW
+            inits = []
+            nodes = []
+            nodes.append(_build_slice_crop('input', 'cropped', IH, IW, inits))
+            nodes.append(_build_slice_reverse('cropped', 'flipped', onnx_axis, dim_size, inits))
+            nodes.append(_build_pad_node('flipped', 'output', pad_h, pad_w, inits))
+            return mk(nodes, inits)
+    return None
+
+
+def s_rotate(td):
+    """Rotate using Slice+Transpose — zero MACs for square grids and k=2.
+    Gather fallback for non-square k=1,3."""
+    exs = get_exs(td)
+    sp = fixed_shapes(td)
+    if sp is None:
+        return None
+    (IH, IW), (OH, OW) = sp
+    for k in [1, 2, 3]:
+        if not all(np.array_equal(out, np.rot90(inp, k)) for inp, out in exs):
+            continue
+        if k == 2:
+            pad_h, pad_w = GH - OH, GW - OW
+            inits = []
+            nodes = []
+            nodes.append(_build_slice_crop('input', 'cropped', IH, IW, inits))
+            nodes.append(_build_slice_reverse('cropped', 'flip_h', 2, IH, inits, suffix='_h'))
+            nodes.append(_build_slice_reverse('flip_h', 'rotated', 3, IW, inits, suffix='_w'))
+            nodes.append(_build_pad_node('rotated', 'output', pad_h, pad_w, inits))
+            return mk(nodes, inits)
+        elif k == 1 and IH == IW:
+            pad_h, pad_w = GH - IH, GW - IW
+            inits = []
+            nodes = []
+            nodes.append(_build_slice_crop('input', 'cropped', IH, IW, inits))
+            nodes.append(helper.make_node('Transpose', ['cropped'], ['transposed'], perm=[0, 1, 3, 2]))
+            nodes.append(_build_slice_reverse('transposed', 'rotated', 2, IH, inits))
+            nodes.append(_build_pad_node('rotated', 'output', pad_h, pad_w, inits))
+            return mk(nodes, inits)
+        elif k == 3 and IH == IW:
+            pad_h, pad_w = GH - IH, GW - IW
+            inits = []
+            nodes = []
+            nodes.append(_build_slice_crop('input', 'cropped', IH, IW, inits))
+            nodes.append(_build_slice_reverse('cropped', 'flipped', 2, IH, inits))
+            nodes.append(helper.make_node('Transpose', ['flipped'], ['rotated'], perm=[0, 1, 3, 2]))
+            nodes.append(_build_pad_node('rotated', 'output', pad_h, pad_w, inits))
+            return mk(nodes, inits)
+        else:
+            idx = np.zeros((OH, OW, 2), dtype=np.int64)
+            for r in range(OH):
+                for c in range(OW):
+                    if k == 1:
+                        sr, sc = c, IH - 1 - r
+                    elif k == 3:
+                        sr, sc = IW - 1 - c, r
+                    idx[r, c] = [sr, sc]
+            return _build_gather_model(OH, OW, idx)
+    return None
+
+
+def s_shift(td):
+    """Shift transformation solver."""
+    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 dr in range(-5, 6):
+        for dc in range(-5, 6):
+            if dr == 0 and dc == 0:
+                continue
+            ok = True
+            for inp, out in exs:
+                shifted = np.zeros_like(inp)
+                r0, r1 = max(0, dr), min(IH, IH + dr)
+                c0, c1 = max(0, dc), min(IW, IW + dc)
+                if r1 > r0 and c1 > c0:
+                    sr0, sc0 = max(0, -dr), max(0, -dc)
+                    shifted[r0:r1, c0:c1] = inp[sr0:sr0 + (r1 - r0), sc0:sc0 + (c1 - c0)]
+                if not np.array_equal(shifted, out):
+                    ok = False
+                    break
+            if not ok:
+                continue
+            idx = np.zeros((OH, OW, 2), dtype=np.int64)
+            cst = np.full((OH, OW), 0, dtype=np.int64)
+            for r in range(OH):
+                for c in range(OW):
+                    sr, sc = r - dr, c - dc
+                    if 0 <= sr < IH and 0 <= sc < IW:
+                        idx[r, c] = [sr, sc]
+                    else:
+                        idx[r, c] = [-1, -1]
+            return _build_gather_model_with_const(IH, IW, OH, OW, idx, cst)
+    return None
+
+
+def s_fixed_crop(td):
+    """Fixed crop solver."""
+    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 or (OH == IH and OW == IW):
+        return None
+    for r0 in range(IH - OH + 1):
+        for c0 in range(IW - OW + 1):
+            if all(np.array_equal(inp[r0:r0 + OH, c0:c0 + OW], 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] = [r0 + r, c0 + c]
+                return _build_gather_model(OH, OW, idx)
+    return None
+
+
+def s_gravity(td):
+    """Detect gravity-like compaction (detection only, no ONNX model built)."""
+    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
+
+    def _gravity(grid, direction):
+        r = np.zeros_like(grid)
+        h, w = grid.shape
+        if direction in ('down', 'up'):
+            for c in range(w):
+                nz = grid[:, c][grid[:, c] != 0]
+                if direction == 'down':
+                    r[h - len(nz):h, c] = nz
+                else:
+                    r[:len(nz), c] = nz
+        else:
+            for rr in range(h):
+                nz = grid[rr, :][grid[rr, :] != 0]
+                if direction == 'right':
+                    r[rr, w - len(nz):w] = nz
+                else:
+                    r[rr, :len(nz)] = nz
+        return r
+
+    for d in ('down', 'up', 'left', 'right'):
+        if all(np.array_equal(_gravity(inp, d), out) for inp, out in exs):
+            return None
+    return None