Split optimizer preparation from solving

README.md

@@ -60,6 +60,7 @@ Current tile-graph status:
 - the latest reduce-by-key pinned badge probe under `artifacts/tile-graph-reducebykey-v1-badge-126/` finishes in about `16.1s` cold and `6.0s` warm on CUDA, versus the previous full-resolution cut at about `96.9s` cold and `5.4s` warm
 - that speed pass is a real throughput win but a small quality regression: the same pinned badge case moves from `0.1814` to `0.2036` source-fidelity
 - the optimizer now has its own deep-dive map in `docs/optimizer-algorithm-map.md`, and the first cut from that map removed the old UV-optimizer name plus the continuous k-means cluster-boundary scout; the continuous path is now a fixed-lattice, source-first discrete chooser guided by real edge evidence
+- optimizer preparation is now split into a private `_OptimizerPrep` bundle, so the continuous path has a visible seam between "build the map" and "choose the pixels"
 - after the first pruning pass, tile-graph no longer participates in pipeline phase-rerank probes and no longer carries hybrid geometry priors through its unary cost; the path is now one lattice-conditioned candidate generator plus one local discrete solver
 - this fixes the core design mismatch that had allowed repeated distant labels to create big same-color patches and opaque black background blocks
 - on the current `24x24` emblem smoke case, an end-to-end `tile-graph` run dropped from about `2.57s` on CPU to `0.61s` on CUDA on this machine

docs/implementation-plan.md

@@ -121,6 +121,19 @@ The first optimizer cut landed immediately after that map:
 - `continuous.py` now uses only real source edges for guidance instead of blending in coarse cluster boundaries
 - the pinned cleaned-badge continuous smoke run under `artifacts/optimizer-cut-v1-badge-126/` still behaves coherently after the cut: snap scores `0.08100`, final scores `0.07535`, and the generated preview is `diagnostics/output-preview.png`
 
+The second optimizer cut is structural:
+
+- `_prepare_optimizer(...)` now owns source tensor creation, UV construction, edge gradients, source lattice reference building, detail reference building, source reliability, source deltas, guidance, and cell geometry
+- `_source_detail_delta_tensors(...)` computes premultiplied source-detail deltas once instead of repeating the same `premultiply(...)` calls for each axis
+- `optimize_lattice_pixels(...)` is now an orchestration shell around prep, snap, refine, the source-fidelity guardrail, and artifact packaging
+- the pinned cleaned-badge continuous smoke run under `artifacts/optimizer-cut-v2-prep-badge-126/` is byte-size identical to the previous preview and preserves the same source-fidelity numbers: snap `0.08100`, final `0.07535`
+
+Next after that:
+
+- put the representative portrait on trial in refine
+- add a focused ablation runner for pinned badge/emblem cases before cutting more behavior
+- collapse duplicated adjacency / motif / line voices only after ablation shows which stage owns each idea
+
 The next optimizer simplification pass should use that map as its cutting checklist instead of guessing from scattered helper names.
 
 ## High-value regression case

docs/optimizer-algorithm-map.md

@@ -40,6 +40,8 @@ The optimizer carries five important pieces of cargo all the way through:
 
 If we lose track of one of those, the optimizer becomes impossible to reason about.
 
+The first structural cut now gathers that cargo into `_OptimizerPrep` before snap and refine begin. That does not make the algorithm better by itself, but it makes the machine easier to see: preparation builds the map, snap and refine walk it.
+
 ## Stage 0: The pipeline picks the ruler
 
 Functions:
@@ -220,6 +222,12 @@ The optimizer now has two different stories about the same lattice:
 
 Much of the optimizer's complexity is the cost of arbitrating between those two stories.
 
+Current code shape:
+
+- `_prepare_optimizer(...)` owns this whole setup stage now
+- `_source_detail_delta_tensors(...)` computes the source-detail deltas once instead of rebuilding premultiplied references several times
+- `optimize_lattice_pixels(...)` now receives a prepared bundle and orchestrates snap, refine, and the final guardrail
+
 ## Stage 5: Snap every output cell to one real nearby source pixel
 
 Function:
@@ -468,7 +476,7 @@ These are the places where the optimizer still speaks in an overly complicated o
 - the machine still maintains two overlapping portraits of the same lattice and spends much of its complexity mediating the argument between them
 - snap, relax, and refine all carry slightly different versions of adjacency, motif, line, and delta agreement; the same idea is being said three times with different accents
 - the candidate generator is local and the solver is discrete, but the vocabulary around it still sounds like a continuous deformation engine
-- the main entry point is now honestly named `optimize_lattice_pixels(...)`, and the k-means boundary scout is gone; those cuts removed two contradictions from the previous map
+- the main entry point is now honestly named `optimize_lattice_pixels(...)`, the k-means boundary scout is gone, and prep is split from decision-making; those cuts removed three contradictions from the previous map
 
 This is probably the real cutting checklist for the next optimizer pass.
 
@@ -484,11 +492,11 @@ That is much simpler than the surrounding names make it sound.
 
 If we keep pruning, the highest-value cuts look like this:
 
-1. Separate preparation from decision-making.
-   The edge scout, portraits, reliability maps, and delta maps want to be a clear prep stage instead of being braided through one giant function.
-2. Decide whether the representative portrait still earns its keep everywhere.
+1. Decide whether the representative portrait still earns its keep everywhere.
    If the source-first path is now strong enough, some of that arbitration machinery may be removable.
-3. Collapse duplicate structure voices.
+2. Collapse duplicate structure voices.
    Adjacency, motif, and line are important, but they should not need three near-parallel dialects unless they are truly doing different work.
+3. Put relax on trial.
+   The soft relaxation stage may be a useful bridge, or it may be an expensive second solver that mostly repeats refine.
 
 That is the next place to swing the axe.

src/repixelizer/continuous.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
+
 import numpy as np
 
 from .io import premultiply, unpremultiply
@@ -9,6 +11,23 @@ from .source_reference import build_source_lattice_reference
 from .types import ContinuousSourceAnalysis, InferenceResult, SolverArtifacts
 
 
+@dataclass(slots=True)
+class _OptimizerPrep:
+    source_t: object
+    uv_t: object
+    representative_t: object
+    source_reference_t: object
+    source_reliability_t: object
+    source_lattice_reference: object
+    source_delta_x_t: object | None
+    source_delta_y_t: object | None
+    source_delta_diag_t: object | None
+    source_delta_anti_t: object | None
+    guidance: np.ndarray
+    cell_x: float
+    cell_y: float
+
+
 def _require_torch():
     try:
         import torch
@@ -426,6 +445,50 @@ def _build_source_detail_reference(source_rgba: np.ndarray, source_reference, so
     return detail_rgba.astype(np.float32)
 
 
+def _source_detail_delta_tensors(
+    torch,
+    source_detail_reference: np.ndarray,
+    *,
+    target_width: int,
+    target_height: int,
+    device: str,
+):
+    detail_premul = premultiply(source_detail_reference)
+    source_delta_x_t = (
+        torch.from_numpy((detail_premul[:, 1:, :] - detail_premul[:, :-1, :])[None, ...]).to(
+            device=device,
+            dtype=torch.float32,
+        )
+        if target_width > 1
+        else None
+    )
+    source_delta_y_t = (
+        torch.from_numpy((detail_premul[1:, :, :] - detail_premul[:-1, :, :])[None, ...]).to(
+            device=device,
+            dtype=torch.float32,
+        )
+        if target_height > 1
+        else None
+    )
+    source_delta_diag_t = (
+        torch.from_numpy((detail_premul[1:, 1:, :] - detail_premul[:-1, :-1, :])[None, ...]).to(
+            device=device,
+            dtype=torch.float32,
+        )
+        if target_height > 1 and target_width > 1
+        else None
+    )
+    source_delta_anti_t = (
+        torch.from_numpy((detail_premul[1:, :-1, :] - detail_premul[:-1, 1:, :])[None, ...]).to(
+            device=device,
+            dtype=torch.float32,
+        )
+        if target_height > 1 and target_width > 1
+        else None
+    )
+    return source_delta_x_t, source_delta_y_t, source_delta_diag_t, source_delta_anti_t
+
+
 def _edge_reliability(source_reliability, axis: str):
     if axis == "x":
         return (source_reliability[:, 1:] + source_reliability[:, :-1]) * 0.5
@@ -1336,21 +1399,16 @@ def _discrete_refine_output(
     return _finalize_output_rgba(best_colors, use_opaque, use_transparent), loss_history
 
 
-def optimize_lattice_pixels(
+def _prepare_optimizer(
+    torch,
+    F,
     rgba: np.ndarray,
     inference: InferenceResult,
     analysis: ContinuousSourceAnalysis,
-    steps: int,
-    seed: int,
+    solver_params: SolverHyperParams,
+    *,
     device: str,
-    solver_params: SolverHyperParams | None = None,
-) -> SolverArtifacts:
-    torch, F = _require_torch()
-    device = _resolve_device(torch, device)
-    solver_params = solver_params or SolverHyperParams()
-    torch.manual_seed(seed)
-    if device == "cuda":
-        torch.cuda.manual_seed_all(seed)
+) -> _OptimizerPrep:
     source = premultiply(rgba)
     height, width = source.shape[:2]
     cell_x = width / max(1, inference.target_width)
@@ -1359,7 +1417,7 @@ def optimize_lattice_pixels(
     source_edge = analysis.edge_map.astype(np.float32)
     edge_grad_x, edge_grad_y = _edge_gradient_maps(source_edge)
     edge_t = torch.from_numpy(source_edge[None, None, ...]).to(device=device, dtype=torch.float32)
-    uv0 = _make_regular_uv(
+    uv = _make_regular_uv(
         height=height,
         width=width,
         target_height=inference.target_height,
@@ -1367,7 +1425,7 @@ def optimize_lattice_pixels(
         phase_x=inference.phase_x,
         phase_y=inference.phase_y,
     )
-    uv0_t = torch.from_numpy(uv0[None, ...]).to(device=device, dtype=torch.float32)
+    uv_t = torch.from_numpy(uv[None, ...]).to(device=device, dtype=torch.float32)
     offsets_t = torch.from_numpy(
         _make_patch_offsets(height=height, width=width, target_height=inference.target_height, target_width=inference.target_width)
     ).to(device=device, dtype=torch.float32)
@@ -1384,10 +1442,10 @@ def optimize_lattice_pixels(
         edge_grad_y_hint=edge_grad_y,
     )
     source_detail_reference = _build_source_detail_reference(rgba, source_lattice_reference, solver_params)
-    initial_patches = _sample_cell_patches(F, source_t, uv0_t, offsets_t)
-    initial_representative_t, _ = _representative_colors(initial_patches, solver_params)
-    initial_representative_t = initial_representative_t.detach()
-    initial_source_reference_t = torch.from_numpy(premultiply(source_detail_reference)[None, ...]).to(
+    initial_patches = _sample_cell_patches(F, source_t, uv_t, offsets_t)
+    representative_t, _ = _representative_colors(initial_patches, solver_params)
+    representative_t = representative_t.detach()
+    source_reference_t = torch.from_numpy(premultiply(source_detail_reference)[None, ...]).to(
         device=device,
         dtype=torch.float32,
     )
@@ -1395,37 +1453,65 @@ def optimize_lattice_pixels(
         device=device,
         dtype=torch.float32,
     )
-    source_delta_x_t = (
-        torch.from_numpy((premultiply(source_detail_reference)[:, 1:, :] - premultiply(source_detail_reference)[:, :-1, :])[None, ...]).to(device=device, dtype=torch.float32)
-        if inference.target_width > 1
-        else None
+    source_delta_x_t, source_delta_y_t, source_delta_diag_t, source_delta_anti_t = _source_detail_delta_tensors(
+        torch,
+        source_detail_reference,
+        target_width=inference.target_width,
+        target_height=inference.target_height,
+        device=device,
     )
-    source_delta_y_t = (
-        torch.from_numpy((premultiply(source_detail_reference)[1:, :, :] - premultiply(source_detail_reference)[:-1, :, :])[None, ...]).to(device=device, dtype=torch.float32)
-        if inference.target_height > 1
-        else None
+    return _OptimizerPrep(
+        source_t=source_t,
+        uv_t=uv_t,
+        representative_t=representative_t,
+        source_reference_t=source_reference_t,
+        source_reliability_t=source_reliability_t,
+        source_lattice_reference=source_lattice_reference,
+        source_delta_x_t=source_delta_x_t,
+        source_delta_y_t=source_delta_y_t,
+        source_delta_diag_t=source_delta_diag_t,
+        source_delta_anti_t=source_delta_anti_t,
+        guidance=guide_small[0, 0].detach().cpu().numpy().astype(np.float32),
+        cell_x=cell_x,
+        cell_y=cell_y,
     )
-    source_delta_diag_t = (
-        torch.from_numpy((premultiply(source_detail_reference)[1:, 1:, :] - premultiply(source_detail_reference)[:-1, :-1, :])[None, ...]).to(device=device, dtype=torch.float32)
-        if inference.target_height > 1 and inference.target_width > 1
-        else None
-    )
-    source_delta_anti_t = (
-        torch.from_numpy((premultiply(source_detail_reference)[1:, :-1, :] - premultiply(source_detail_reference)[:-1, 1:, :])[None, ...]).to(device=device, dtype=torch.float32)
-        if inference.target_height > 1 and inference.target_width > 1
-        else None
+
+
+def optimize_lattice_pixels(
+    rgba: np.ndarray,
+    inference: InferenceResult,
+    analysis: ContinuousSourceAnalysis,
+    steps: int,
+    seed: int,
+    device: str,
+    solver_params: SolverHyperParams | None = None,
+) -> SolverArtifacts:
+    torch, F = _require_torch()
+    device = _resolve_device(torch, device)
+    solver_params = solver_params or SolverHyperParams()
+    torch.manual_seed(seed)
+    if device == "cuda":
+        torch.cuda.manual_seed_all(seed)
+    prep = _prepare_optimizer(
+        torch,
+        F,
+        rgba,
+        inference,
+        analysis,
+        solver_params,
+        device=device,
     )
     snap_rgba = _snap_output_to_source_pixels(
         torch,
-        source_t,
-        uv0_t,
-        initial_representative_t,
-        initial_source_reference_t,
-        source_reliability_t,
-        source_lattice_reference,
+        prep.source_t,
+        prep.uv_t,
+        prep.representative_t,
+        prep.source_reference_t,
+        prep.source_reliability_t,
+        prep.source_lattice_reference,
         solver_params,
-        cell_x=cell_x,
-        cell_y=cell_y,
+        cell_x=prep.cell_x,
+        cell_y=prep.cell_y,
     )
     snap_t = torch.from_numpy(premultiply(snap_rgba).transpose(2, 0, 1)[None, ...]).to(
         device=device,
@@ -1435,20 +1521,20 @@ def optimize_lattice_pixels(
         target_rgba, loss_history = _discrete_refine_output(
             torch,
             F,
-            source_t,
-            uv0_t,
-            initial_representative_t,
-            initial_source_reference_t,
-            source_reliability_t,
-            source_lattice_reference,
+            prep.source_t,
+            prep.uv_t,
+            prep.representative_t,
+            prep.source_reference_t,
+            prep.source_reliability_t,
+            prep.source_lattice_reference,
             snap_t,
-            source_delta_x_t,
-            source_delta_y_t,
-            source_delta_diag_t,
-            source_delta_anti_t,
+            prep.source_delta_x_t,
+            prep.source_delta_y_t,
+            prep.source_delta_diag_t,
+            prep.source_delta_anti_t,
             solver_params,
-            cell_x=cell_x,
-            cell_y=cell_y,
+            cell_x=prep.cell_x,
+            cell_y=prep.cell_y,
             iterations=steps,
         )
         snap_score = source_lattice_consistency_breakdown(
@@ -1473,13 +1559,11 @@ def optimize_lattice_pixels(
             target_rgba = snap_rgba
             loss_history.append(float(snap_score))
         initial_rgba = snap_rgba
-        uv_np = uv0_t.detach().cpu().numpy()[0]
-        guidance = guide_small[0, 0].detach().cpu().numpy()
 
     return SolverArtifacts(
         target_rgba=target_rgba,
-        uv_field=uv_np,
-        guidance_strength=guidance.astype(np.float32),
+        uv_field=prep.uv_t.detach().cpu().numpy()[0],
+        guidance_strength=prep.guidance,
         initial_rgba=initial_rgba,
         loss_history=loss_history,
     )