From 863b518e1e33ed42bb0c685ec2db563903f63e0f Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Thu, 23 Apr 2026 22:51:05 +0100 Subject: [PATCH] Add optimizer relax displacement diagnostics --- README.md | 2 + docs/implementation-plan.md | 15 ++++ docs/optimizer-algorithm-map.md | 22 ++++- src/repixelizer/continuous.py | 141 ++++++++++++++++++++++++++++++-- src/repixelizer/diagnostics.py | 48 +++++++++++ src/repixelizer/pipeline.py | 19 +++++ src/repixelizer/types.py | 1 + tests/test_continuous.py | 31 +++++++ tests/test_pipeline.py | 3 + 9 files changed, 276 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index a21aceb..9616823 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,8 @@ Current tile-graph status: - 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" - refine is now source-first too: the representative portrait still stabilizes snap, but refine and final structure choice no longer consult it, and the pinned badge smoke improved slightly from `0.07535` to `0.07485` - the next relax cut also paid off: removing the old relax-mode bonus while keeping the core relax stage improved the same pinned badge smoke again from `0.07485` to `0.07451` +- the optimizer now emits stage displacement diagnostics too, so `run.json` can show what snap, relax handoff, relax mode, and final refine actually did to the per-cell source-pixel displacement field instead of just reporting a single fidelity score +- those new diagnostics also settled the design argument about relax: on the pinned badge check under `artifacts/optimizer-relax-purpose-check/`, relax improves source-fidelity a bit (`0.07549` without relax to `0.07451` with it), but it does not act like a strong broad-swatch phase smoother; the displacement field stays roughly as jittery, so the current relax stage is better described as a soft consensus / handoff stabilizer over fixed local candidates - 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 diff --git a/docs/implementation-plan.md b/docs/implementation-plan.md index 356bf62..56cc467 100644 --- a/docs/implementation-plan.md +++ b/docs/implementation-plan.md @@ -153,6 +153,21 @@ Next after that: - 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 pass turned that "put relax on trial" note into an actual measurement seam: + +- `SolverArtifacts` now carries stage diagnostics for continuous runs, including the exact selected-source displacement field for `snap`, `relax_handoff`, `relax_mode`, and `final_output` +- `run.json` now includes `optimizer_displacement` metrics, and diagnostics directories now write color-coded displacement previews for those stages +- the pinned cleaned-badge comparison under `artifacts/optimizer-relax-purpose-check/` runs the same fixed lattice twice: once with normal relax and once with `relax_iterations=0` + +What that test says: + +- `relax` still helps a little on the real badge objective: final source-fidelity improves from `0.07549` without relax to `0.07451` with relax +- but the displacement field does **not** become clearly smoother in the way the original design story promised +- on that same pinned badge run, final orthogonal displacement jitter is `1.7154` with relax versus `1.6783` without it, and final local residual is `1.1548` with relax versus `1.1312` without it +- the honest read is that the current relax stage behaves more like a soft consensus / handoff stabilizer over fixed local candidate trays than a true large-swath phase adjustment mechanism + +That means the next optimizer cuts should treat `relax` as a measurable subsolver with a narrower actual job, not as a sacred phase-field story we keep around out of nostalgia. + The next optimizer simplification pass should use that map as its cutting checklist instead of guessing from scattered helper names. ## High-value regression case diff --git a/docs/optimizer-algorithm-map.md b/docs/optimizer-algorithm-map.md index 1857f00..adc8e2d 100644 --- a/docs/optimizer-algorithm-map.md +++ b/docs/optimizer-algorithm-map.md @@ -463,6 +463,26 @@ Recent cut: This is wet plaster and half-seated stones. Nothing is fully committed yet. Each slot is allowed to say, "I'm 60% this dark edge stone, 30% the slightly duller one, 10% the safer interior stone." The wall is still trembling, but the trembling starts to synchronize. Temperature drops. The slurry thickens. Shapes that looked independent begin leaning into each other like bones finding their sockets. +### Measured reality + +The original design goal for `relax` was larger than "soften local choices." It was supposed to let the machine adjust phase smoothly over broad swaths of the canvas instead of only making tiny local corrections. + +The current instrumentation in `continuous.py` says the truth is weaker than that story. + +On the pinned cleaned-badge run at `126x126`, phase `(0.0, -0.2)`, `steps=48`: + +- with `relax_iterations=24`, final source-fidelity lands at `0.07451` +- with `relax_iterations=0`, final source-fidelity lands at `0.07549` +- so `relax` does help a little + +But the displacement-field diagnostics do **not** show a dramatic large-scale phase-smoothing effect: + +- final orthogonal displacement jitter is `1.7154` with relax versus `1.6783` without it +- final local residual is `1.1548` with relax versus `1.1312` without it +- final dominant-offset ratio is `0.6296` with relax versus `0.6380` without it + +So the current `relax` earns a small fidelity gain, but it does not look like a true broad-swatch phase transport stage. It behaves more like a soft consensus pass over a fixed tray of local candidates, then hands that consensus to greedy refine. + ## Stage 8: Greedy discrete refine Function: @@ -601,7 +621,7 @@ That is much simpler than the surrounding names make it sound. If we keep pruning, the highest-value cuts look like this: 1. Put relax on trial. - The soft relaxation stage still seems useful, but the relax-mode bonus was dead weight. The next question is which remaining relax terms are actually earning their keep. + The soft relaxation stage still seems useful, but the current measurements say it is not really acting like a large-scale phase smoother. The next question is which remaining relax terms are actually earning their keep, and whether the surviving benefit comes from handoff stabilization rather than regional phase drift. 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. Decide whether snap still needs both portraits everywhere. diff --git a/src/repixelizer/continuous.py b/src/repixelizer/continuous.py index 2ad81e1..0aa45cc 100644 --- a/src/repixelizer/continuous.py +++ b/src/repixelizer/continuous.py @@ -1047,7 +1047,17 @@ def _snap_output_to_source_pixels( solver_params=solver_params, ) final_colors = _select_colors(candidate_colors, hardened_selected) - return _finalize_output_rgba(final_colors, use_opaque, use_transparent) + final_x, final_y = _select_positions(candidate_x, candidate_y, hardened_selected) + stage_debug = { + "snap": _displacement_diagnostics( + uv_t[0].detach().cpu().numpy(), + final_x, + final_y, + width=width, + height=height, + ) + } + return _finalize_output_rgba(final_colors, use_opaque, use_transparent), stage_debug def _select_colors(candidate_colors, selected): @@ -1059,6 +1069,88 @@ def _select_colors(candidate_colors, selected): ).squeeze(2) +def _select_positions(candidate_x, candidate_y, selected): + output_height = candidate_x.shape[0] + output_width = candidate_x.shape[1] + gather_index = selected[..., None] + selected_x = candidate_x.gather( + dim=2, + index=gather_index.expand(output_height, output_width, 1), + ).squeeze(2) + selected_y = candidate_y.gather( + dim=2, + index=gather_index.expand(output_height, output_width, 1), + ).squeeze(2) + return selected_x, selected_y + + +def _box_blur_field(field: np.ndarray) -> np.ndarray: + padded = np.pad(field, ((1, 1), (1, 1), (0, 0)), mode="edge") + accum = np.zeros_like(field, dtype=np.float32) + for dy in range(3): + for dx in range(3): + accum += padded[dy : dy + field.shape[0], dx : dx + field.shape[1]] + return accum / 9.0 + + +def _displacement_diagnostics(uv_field: np.ndarray, selected_x, selected_y, *, width: int, height: int) -> dict[str, np.ndarray | float]: + uv_x = (uv_field[..., 0] + 1.0) * 0.5 * max(1.0, float(width - 1)) + uv_y = (uv_field[..., 1] + 1.0) * 0.5 * max(1.0, float(height - 1)) + selected_x_np = selected_x.detach().cpu().numpy().astype(np.float32) + selected_y_np = selected_y.detach().cpu().numpy().astype(np.float32) + displacement_x = selected_x_np - uv_x.astype(np.float32) + displacement_y = selected_y_np - uv_y.astype(np.float32) + displacement = np.stack([displacement_x, displacement_y], axis=-1) + magnitude = np.linalg.norm(displacement, axis=-1) + + orthogonal_jumps: list[np.ndarray] = [] + if displacement.shape[1] > 1: + orthogonal_jumps.append(np.linalg.norm(displacement[:, 1:, :] - displacement[:, :-1, :], axis=-1)) + if displacement.shape[0] > 1: + orthogonal_jumps.append(np.linalg.norm(displacement[1:, :, :] - displacement[:-1, :, :], axis=-1)) + diagonal_jumps: list[np.ndarray] = [] + if displacement.shape[0] > 1 and displacement.shape[1] > 1: + diagonal_jumps.append(np.linalg.norm(displacement[1:, 1:, :] - displacement[:-1, :-1, :], axis=-1)) + diagonal_jumps.append(np.linalg.norm(displacement[1:, :-1, :] - displacement[:-1, 1:, :], axis=-1)) + + rounded_dx = np.rint(displacement_x).astype(np.int16) + rounded_dy = np.rint(displacement_y).astype(np.int16) + equal_pairs: list[np.ndarray] = [] + if rounded_dx.shape[1] > 1: + equal_pairs.append((rounded_dx[:, 1:] == rounded_dx[:, :-1]) & (rounded_dy[:, 1:] == rounded_dy[:, :-1])) + if rounded_dx.shape[0] > 1: + equal_pairs.append((rounded_dx[1:, :] == rounded_dx[:-1, :]) & (rounded_dy[1:, :] == rounded_dy[:-1, :])) + + offsets = np.stack([rounded_dx.reshape(-1), rounded_dy.reshape(-1)], axis=-1) + if offsets.size: + unique_offsets, counts = np.unique(offsets, axis=0, return_counts=True) + dominant_offset = unique_offsets[int(np.argmax(counts))] + dominant_ratio = float(np.max(counts) / max(1, offsets.shape[0])) + else: + dominant_offset = np.asarray([0, 0], dtype=np.int16) + dominant_ratio = 1.0 + + blur = _box_blur_field(displacement) + local_residual = np.linalg.norm(displacement - blur, axis=-1) + return { + "selected_x": selected_x_np, + "selected_y": selected_y_np, + "displacement_x": displacement_x, + "displacement_y": displacement_y, + "mean_magnitude_px": float(np.mean(magnitude)), + "max_magnitude_px": float(np.max(magnitude)), + "orthogonal_jitter_px": float(np.mean(np.concatenate([jump.reshape(-1) for jump in orthogonal_jumps]))) if orthogonal_jumps else 0.0, + "diagonal_jitter_px": float(np.mean(np.concatenate([jump.reshape(-1) for jump in diagonal_jumps]))) if diagonal_jumps else 0.0, + "local_residual_px": float(np.mean(local_residual)), + "orthogonal_equal_ratio": float(np.mean(np.concatenate([pair.reshape(-1) for pair in equal_pairs]).astype(np.float32))) if equal_pairs else 1.0, + "dominant_offset_ratio": dominant_ratio, + "dominant_offset_dx_px": float(dominant_offset[0]), + "dominant_offset_dy_px": float(dominant_offset[1]), + "mean_dx_px": float(np.mean(displacement_x)), + "mean_dy_px": float(np.mean(displacement_y)), + } + + def _harden_binary_alpha_selection( torch, candidate_colors, @@ -1372,7 +1464,33 @@ def _discrete_refine_output( solver_params=solver_params, ) best_colors = _select_colors(candidate_colors, hardened_selected) - return _finalize_output_rgba(best_colors, use_opaque, use_transparent), loss_history + handoff_x, handoff_y = _select_positions(candidate_x, candidate_y, selected) + mode_x, mode_y = _select_positions(candidate_x, candidate_y, relaxed_mode_selected) + final_x, final_y = _select_positions(candidate_x, candidate_y, hardened_selected) + stage_debug = { + "relax_handoff": _displacement_diagnostics( + uv_t[0].detach().cpu().numpy(), + handoff_x, + handoff_y, + width=source_t.shape[-1], + height=source_t.shape[-2], + ), + "relax_mode": _displacement_diagnostics( + uv_t[0].detach().cpu().numpy(), + mode_x, + mode_y, + width=source_t.shape[-1], + height=source_t.shape[-2], + ), + "final_output": _displacement_diagnostics( + uv_t[0].detach().cpu().numpy(), + final_x, + final_y, + width=source_t.shape[-1], + height=source_t.shape[-2], + ), + } + return _finalize_output_rgba(best_colors, use_opaque, use_transparent), loss_history, stage_debug def _prepare_optimizer( @@ -1477,7 +1595,7 @@ def optimize_lattice_pixels( solver_params, device=device, ) - snap_rgba = _snap_output_to_source_pixels( + snap_rgba, snap_stage_debug = _snap_output_to_source_pixels( torch, prep.source_t, prep.uv_t, @@ -1494,7 +1612,7 @@ def optimize_lattice_pixels( dtype=torch.float32, ) with torch.no_grad(): - target_rgba, loss_history = _discrete_refine_output( + target_rgba, loss_history, refine_stage_debug = _discrete_refine_output( torch, F, prep.source_t, @@ -1529,10 +1647,22 @@ def optimize_lattice_pixels( phase_y=inference.phase_y, alpha_threshold=solver_params.alpha_transparent_threshold, )["score"] - if target_score > snap_score + 1e-6: + guardrail_kept_snap = bool(target_score > snap_score + 1e-6) + if guardrail_kept_snap: target_rgba = snap_rgba loss_history.append(float(snap_score)) + refine_stage_debug["final_output"] = snap_stage_debug["snap"] initial_rgba = snap_rgba + stage_diagnostics = { + "displacements": { + **snap_stage_debug, + **refine_stage_debug, + }, + "guardrail_kept_snap": guardrail_kept_snap, + "snap_source_fidelity_score": float(snap_score), + "solver_target_source_fidelity_score": float(target_score), + "relax_iterations": int(min(max(0, solver_params.relax_iterations), max(0, steps)) if steps > 0 else 0), + } return SolverArtifacts( target_rgba=target_rgba, @@ -1540,4 +1670,5 @@ def optimize_lattice_pixels( guidance_strength=prep.guidance, initial_rgba=initial_rgba, loss_history=loss_history, + stage_diagnostics=stage_diagnostics, ) diff --git a/src/repixelizer/diagnostics.py b/src/repixelizer/diagnostics.py index 2ef2cc1..e9790e5 100644 --- a/src/repixelizer/diagnostics.py +++ b/src/repixelizer/diagnostics.py @@ -90,6 +90,45 @@ def write_heatmap(path: str | Path, heatmap: np.ndarray) -> None: save_rgba(path, rgba) +def _displacement_preview_rgba(displacement_x: np.ndarray, displacement_y: np.ndarray) -> np.ndarray: + magnitude = np.sqrt(np.square(displacement_x) + np.square(displacement_y)).astype(np.float32) + scale = float(np.percentile(magnitude, 95)) if magnitude.size else 0.0 + normalized = np.clip(magnitude / max(scale, 1e-4), 0.0, 1.0) + hue = (np.arctan2(displacement_y, displacement_x) + np.pi) / (2.0 * np.pi) + saturation = normalized + value = np.clip(0.2 + normalized * 0.8, 0.0, 1.0) + + h6 = hue * 6.0 + i = np.floor(h6).astype(np.int32) % 6 + f = h6 - np.floor(h6) + p = value * (1.0 - saturation) + q = value * (1.0 - f * saturation) + t = value * (1.0 - (1.0 - f) * saturation) + + rgb = np.zeros((*hue.shape, 3), dtype=np.float32) + cases = ( + (i == 0, np.stack([value, t, p], axis=-1)), + (i == 1, np.stack([q, value, p], axis=-1)), + (i == 2, np.stack([p, value, t], axis=-1)), + (i == 3, np.stack([p, q, value], axis=-1)), + (i == 4, np.stack([t, p, value], axis=-1)), + (i == 5, np.stack([value, p, q], axis=-1)), + ) + for mask, colors in cases: + rgb[mask] = colors[mask] + + rgba = np.zeros((*hue.shape, 4), dtype=np.float32) + rgba[..., :3] = rgb + rgba[..., 3] = 1.0 + return rgba + + +def write_displacement_preview(path: str | Path, displacement_x: np.ndarray, displacement_y: np.ndarray) -> np.ndarray: + rgba = _displacement_preview_rgba(displacement_x, displacement_y) + save_rgba(path, rgba) + return rgba + + def write_json(path: str | Path, payload: dict[str, Any]) -> None: Path(path).write_text(json.dumps(payload, indent=2), encoding="utf-8") @@ -152,6 +191,14 @@ def summarize_run(result: RunResult) -> dict[str, Any]: for candidate in result.inference.top_candidates if "phase_rerank_score" in candidate.breakdown ] + displacement_metrics = {} + stage_diagnostics = result.solver.stage_diagnostics.get("displacements", {}) + for stage_name, payload in stage_diagnostics.items(): + displacement_metrics[stage_name] = { + key: float(value) + for key, value in payload.items() + if not isinstance(value, np.ndarray) + } return { "target_width": result.inference.target_width, "target_height": result.inference.target_height, @@ -165,6 +212,7 @@ def summarize_run(result: RunResult) -> dict[str, Any]: "source_fidelity": source_fidelity, "phase_rerank_candidates": rerank_candidates, "loss_history": result.solver.loss_history, + "optimizer_displacement": displacement_metrics, } diff --git a/src/repixelizer/pipeline.py b/src/repixelizer/pipeline.py index feb10dd..272637c 100644 --- a/src/repixelizer/pipeline.py +++ b/src/repixelizer/pipeline.py @@ -12,6 +12,7 @@ from .diagnostics import ( summarize_run, write_alpha_preview, write_comparison, + write_displacement_preview, write_heatmap, write_json, write_lattice_overlay, @@ -139,6 +140,24 @@ def run_pipeline( ) write_alpha_preview(diagnostics_path / "alpha-preview.png", source, output_rgba) write_heatmap(diagnostics_path / "noise-heatmap.png", cleanup.isolated_heatmap) + stage_displacements = result.solver.stage_diagnostics.get("displacements", {}) + for stage_name, payload in stage_displacements.items(): + displacement_x = payload.get("displacement_x") + displacement_y = payload.get("displacement_y") + if isinstance(displacement_x, np.ndarray) and isinstance(displacement_y, np.ndarray): + displacement_rgba = write_displacement_preview( + diagnostics_path / f"displacement-{stage_name}.png", + displacement_x, + displacement_y, + ) + save_rgba( + diagnostics_path / f"displacement-{stage_name}-preview.png", + nearest_resize( + displacement_rgba, + width=output_rgba.shape[1] * 8, + height=output_rgba.shape[0] * 8, + ), + ) run_json = summarize_run(result) run_json["inference"] = inference_to_json(inference) run_json["settings"] = { diff --git a/src/repixelizer/types.py b/src/repixelizer/types.py index e1d6c6a..10b0151 100644 --- a/src/repixelizer/types.py +++ b/src/repixelizer/types.py @@ -64,6 +64,7 @@ class SolverArtifacts: guidance_strength: np.ndarray initial_rgba: np.ndarray loss_history: list[float] + stage_diagnostics: dict[str, Any] = field(default_factory=dict) @dataclass(slots=True) diff --git a/tests/test_continuous.py b/tests/test_continuous.py index e3397bb..e2929e6 100644 --- a/tests/test_continuous.py +++ b/tests/test_continuous.py @@ -155,3 +155,34 @@ def test_optimize_lattice_pixels_does_not_regress_from_snap_on_thin_feature_case )["score"] assert final_score <= snap_score + 1e-6 + + +def test_optimize_lattice_pixels_emits_displacement_diagnostics() -> None: + source = np.zeros((12, 12, 4), dtype=np.float32) + source[..., 3] = 1.0 + source[3:9, 6, :3] = 1.0 + + inference = InferenceResult( + target_width=6, + target_height=6, + phase_x=0.0, + phase_y=0.0, + confidence=1.0, + top_candidates=[], + ) + artifacts = optimize_lattice_pixels( + source, + inference=inference, + analysis=analyze_continuous_source(source, seed=7), + steps=2, + seed=7, + device="cpu", + ) + + displacement = artifacts.stage_diagnostics["displacements"] + assert {"snap", "relax_handoff", "relax_mode", "final_output"} <= set(displacement.keys()) + for stage_name in ("snap", "relax_handoff", "final_output"): + payload = displacement[stage_name] + assert payload["displacement_x"].shape == (6, 6) + assert payload["displacement_y"].shape == (6, 6) + assert payload["orthogonal_jitter_px"] >= 0.0 diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 827d8b2..858a171 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -30,11 +30,14 @@ def test_pipeline_writes_output_and_diagnostics(tmp_path: Path) -> None: assert output_path.exists() assert (diagnostics_dir / "run.json").exists() assert (diagnostics_dir / "output-preview.png").exists() + assert (diagnostics_dir / "displacement-snap.png").exists() + assert (diagnostics_dir / "displacement-snap-preview.png").exists() import json run_json = json.loads((diagnostics_dir / "run.json").read_text(encoding="utf-8")) assert set(run_json["source_fidelity"].keys()) == {"snap_initial", "solver_target", "final_output"} assert "phase_rerank_candidates" in run_json + assert {"snap", "relax_handoff", "relax_mode", "final_output"} <= set(run_json["optimizer_displacement"].keys()) assert result.output_rgba.shape[0] == result.inference.target_height assert result.output_rgba.shape[1] == result.inference.target_width