Add greedy stacker: overlay(T1(x), T2(x)) composition after depth-1 beam search

itt_solver/beam_logging.py

@@ -26,12 +26,8 @@ def beam_minimize_with_log(phi_in, phi_target, atomic_library,
                            enable_layer_minus_one=False,
                            boundary_source='target'):
     """
-    Beam search with gate validation and optional Layer-1 admissible mask.
-
-    Uses a dual-strategy approach: each atomic transform is tried on BOTH the
-    current (resized) field AND the original input.  This is critical for
-    shape-changing transforms (e.g. Kronecker) that must operate on the raw
-    input rather than the tiled/resized intermediate.
+    Beam search with gate validation, dual-strategy (resized + original input),
+    and greedy stacker (overlay composition of depth-1 pieces).
     """
     phi_in = np.array(phi_in, dtype=float)
     phi_target = np.array(phi_target, dtype=float)
@@ -55,7 +51,6 @@ def beam_minimize_with_log(phi_in, phi_target, atomic_library,
 
     def _try_candidate(phi_after_atomic, T_atomic, T_cur, cur_field_resized,
                        path_states, path_sigmas, depth_log, candidates, source_tag=""):
-        """Score one candidate, check gates, and append to candidates if accepted."""
         phi_new_resized = _resize_to_target(phi_after_atomic, phi_target)
 
         if enable_layer_minus_one and layer_mask is not None:
@@ -78,12 +73,8 @@ def beam_minimize_with_log(phi_in, phi_target, atomic_library,
 
         if not gates_info.get('passed', False):
             depth_log.append({
-                'atomic': label,
-                'score': score,
-                'residue': residue,
-                'energy': energy,
-                'gates': gates_info,
-                'accepted': False,
+                'atomic': label, 'score': score, 'residue': residue,
+                'energy': energy, 'gates': gates_info, 'accepted': False,
                 'shape': phi_candidate.shape,
             })
             return
@@ -91,16 +82,10 @@ def beam_minimize_with_log(phi_in, phi_target, atomic_library,
         new_states = path_states + [phi_candidate]
         new_sigmas = path_sigmas + [residue]
         T_new = T_cur.compose(T_atomic)
-
         candidates.append((T_new, phi_candidate, score, new_states, new_sigmas))
-
         depth_log.append({
-            'atomic': label,
-            'score': score,
-            'residue': residue,
-            'energy': energy,
-            'gates': gates_info,
-            'accepted': True,
+            'atomic': label, 'score': score, 'residue': residue,
+            'energy': energy, 'gates': gates_info, 'accepted': True,
             'shape': phi_candidate.shape,
         })
 
@@ -111,7 +96,7 @@ def beam_minimize_with_log(phi_in, phi_target, atomic_library,
             base_field_for_apply = path_states[-1]
 
             for idx, T_atomic in enumerate(atomic_library):
-                # --- Strategy 1: apply to current (resized) field ---
+                # Strategy 1: apply to current (resized) field
                 try:
                     phi_after_atomic = T_atomic.apply(base_field_for_apply)
                     _try_candidate(phi_after_atomic, T_atomic, T_cur,
@@ -120,8 +105,7 @@ def beam_minimize_with_log(phi_in, phi_target, atomic_library,
                 except Exception:
                     pass
 
-                # --- Strategy 2: apply to ORIGINAL input (critical for
-                #     shape-changing transforms like Kronecker) ---
+                # Strategy 2: apply to ORIGINAL input
                 try:
                     phi_after_original = T_atomic.apply(phi_in)
                     if phi_after_original.shape != base_field_for_apply.shape or \
@@ -144,8 +128,77 @@ def beam_minimize_with_log(phi_in, phi_target, atomic_library,
         if sigma_l1(best[1], phi_target) == 0:
             break
 
+    # --- Greedy stacker: try overlay(T1(x), T2(x)) for top candidates ---
+    if best is not None and sigma_l1(best[1], phi_target) > 0:
+        depth1_pieces = []
+        for T_atomic in atomic_library:
+            try:
+                piece = T_atomic.apply(phi_in)
+                piece_resized = _resize_to_target(piece, phi_target)
+                piece_sigma = sigma_l1(piece_resized, phi_target)
+                depth1_pieces.append((T_atomic, piece_resized, piece_sigma))
+            except Exception:
+                pass
+
+        depth1_pieces.sort(key=lambda x: x[2])
+        top_n = min(len(depth1_pieces), beam_width * 2)
+        stacker_log = []
+
+        for i in range(top_n):
+            T1, p1, s1 = depth1_pieces[i]
+            for j in range(top_n):
+                if i == j:
+                    continue
+                T2, p2, s2 = depth1_pieces[j]
+
+                overlaid = p1.copy()
+                mask = (p2 != 0)
+                overlaid[mask] = p2[mask]
+
+                residue = sigma_l1(overlaid, phi_target)
+
+                if residue < sigma_l1(best[1], phi_target):
+                    gates_info = validate_gates(overlaid, phi_in, phi_target,
+                                                boundary_mask=boundary_mask_resized,
+                                                max_fraction=max_fraction,
+                                                allowed_symbols=allowed_symbols)
+                    label = f"overlay({repr(T1)},{repr(T2)})"
+                    if gates_info.get('passed', False):
+                        energy = dirichlet_energy(overlaid)
+                        score = residue + lock_coeff * energy
+                        T_composed = Transform(lambda p, _p1=p1, _p2=p2: _overlay(_p1, _p2),
+                                               f"overlay({T1.name},{T2.name})")
+                        _, _, _, best_states, best_sigmas = best
+                        new_states = best_states + [overlaid]
+                        new_sigmas = best_sigmas + [residue]
+                        if score < best[2]:
+                            best = (T_composed, overlaid, score, new_states, new_sigmas)
+
+                        stacker_log.append({
+                            'atomic': label, 'score': score,
+                            'residue': residue, 'energy': energy,
+                            'gates': gates_info, 'accepted': True,
+                            'shape': overlaid.shape,
+                        })
+
+                    if residue == 0:
+                        break
+            if best is not None and sigma_l1(best[1], phi_target) == 0:
+                break
+
+        if stacker_log:
+            logs.append(stacker_log)
+
     if best is None:
         return identity, phi0, [phi0], [sigma_l1(phi0, phi_target)], logs
 
     T_best, phi_best, _, states_best, sigmas_best = best
     return T_best, phi_best, states_best, sigmas_best, logs
+
+
+def _overlay(base, fg):
+    """Transparent overlay helper: fg non-zero pixels overwrite base."""
+    result = base.copy()
+    mask = (fg != 0)
+    result[mask] = fg[mask]
+    return result