| |
| """Render a Fig.3-style alignment plot from dLLM alignment.json. |
| |
| The source paper's Fig. 3 visualizes pairwise gradient cosine similarities and |
| draws dashed interval boundaries. This renderer adds an adjacent-step cosine |
| curve below the heatmap because that is the most direct boundary diagnostic for |
| our mask-ratio grid. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import html |
| import json |
| from pathlib import Path |
|
|
|
|
| def color_for(value: float) -> str: |
| value = max(-1.0, min(1.0, value)) |
| if value >= 0: |
| t = value |
| r = int(255 * (1 - 0.08 * t)) |
| g = int(255 * (1 - 0.62 * t)) |
| b = int(255 * (1 - 0.72 * t)) |
| else: |
| t = -value |
| r = int(255 * (1 - 0.70 * t)) |
| g = int(255 * (1 - 0.45 * t)) |
| b = int(255 * (1 - 0.05 * t)) |
| return f"#{r:02x}{g:02x}{b:02x}" |
|
|
|
|
| def adjacent_cosines(ratios: list[float], matrix: list[list[float]]) -> list[dict[str, float]]: |
| out = [] |
| for i in range(len(ratios) - 1): |
| out.append({ |
| "left": ratios[i], |
| "right": ratios[i + 1], |
| "mid": (ratios[i] + ratios[i + 1]) / 2, |
| "cosine": matrix[i][i + 1], |
| }) |
| return out |
|
|
|
|
| def best_adjacent_cut(adjacent: list[dict[str, float]]) -> int | None: |
| if not adjacent: |
| return None |
| return min(range(len(adjacent)), key=lambda i: adjacent[i]["cosine"]) |
|
|
|
|
| def write_json(path: Path, payload: dict) -> None: |
| path.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") |
|
|
|
|
| def write_markdown(path: Path, payload: dict) -> None: |
| lines = [ |
| "# Fig.3-Style dLLM Alignment Plot", |
| "", |
| f"- Source alignment: `{payload['source_alignment']}`", |
| f"- Figure: `{payload['figure']}`", |
| f"- Candidate boundary: `{payload['candidate_boundary_label']}`", |
| "", |
| "## Adjacent Mask-Ratio Cosine Similarity", |
| "", |
| "| Step pair | Cosine similarity |", |
| "| --- | ---: |", |
| ] |
| for row in payload["adjacent_cosines"]: |
| lines.append(f"| {row['left']:g} -> {row['right']:g} | {row['cosine']:.3f} |") |
| lines.extend([ |
| "", |
| "Interpretation: lower adjacent cosine values indicate a sharper local", |
| "change between neighboring denoising objectives and are natural candidate", |
| "interval boundaries. The dashed line in the SVG marks the lowest adjacent", |
| "cosine split.", |
| "", |
| ]) |
| path.write_text("\n".join(lines), encoding="utf-8") |
|
|
|
|
| def write_svg(path: Path, payload: dict) -> None: |
| ratios = payload["ratios"] |
| matrix = payload["cosine_similarity"] |
| adjacent = payload["adjacent_cosines"] |
| cut_index = payload["candidate_boundary_index"] |
|
|
| n = len(ratios) |
| requested_cell = int(payload.get("cell_size", 0) or 0) |
| cell = requested_cell if requested_cell > 0 else (38 if n <= 12 else 14 if n <= 50 else 8) |
| show_cell_text = bool(payload.get("show_cell_text", n <= 16)) |
| label_every = int(payload.get("label_every", 0) or (1 if n <= 16 else 5 if n <= 60 else 10)) |
| label_w = 92 |
| top = 54 |
| right_pad = 36 |
| heat_w = n * cell |
| heat_h = n * cell |
| gap = 54 |
| curve_h = 160 |
| legend_h = 46 |
| width = label_w + heat_w + right_pad |
| height = top + heat_h + gap + curve_h + legend_h |
| x0 = label_w |
| y0 = top |
|
|
| title = payload["title"] |
| parts = [ |
| f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">', |
| '<rect width="100%" height="100%" fill="#ffffff"/>', |
| f'<text x="{x0}" y="24" font-family="Arial, sans-serif" font-size="17" font-weight="700">{html.escape(title)}</text>', |
| f'<text x="{x0}" y="43" font-family="Arial, sans-serif" font-size="12" fill="#444">Pairwise gradient cosine by mask ratio; dashed line marks lowest adjacent-step cosine.</text>', |
| ] |
|
|
| for i, ratio in enumerate(ratios): |
| if i % label_every != 0 and i != n - 1: |
| continue |
| x = x0 + i * cell + cell / 2 |
| parts.append(f'<text x="{x}" y="{y0 - 9}" text-anchor="middle" font-family="Arial, sans-serif" font-size="10">{ratio:g}</text>') |
| y = y0 + i * cell + cell / 2 + 4 |
| parts.append(f'<text x="{x0 - 10}" y="{y}" text-anchor="end" font-family="Arial, sans-serif" font-size="10">{ratio:g}</text>') |
|
|
| for row_idx, row in enumerate(matrix): |
| for col_idx, value in enumerate(row): |
| x = x0 + col_idx * cell |
| y = y0 + row_idx * cell |
| parts.append(f'<rect x="{x}" y="{y}" width="{cell}" height="{cell}" fill="{color_for(value)}" stroke="#ffffff" stroke-width="1"/>') |
| if show_cell_text: |
| parts.append(f'<text x="{x + cell / 2}" y="{y + cell / 2 + 4}" text-anchor="middle" font-family="Arial, sans-serif" font-size="9" fill="#111">{value:.2f}</text>') |
|
|
| if cut_index is not None: |
| cut_x = x0 + (cut_index + 1) * cell |
| cut_y = y0 + (cut_index + 1) * cell |
| parts.append(f'<line x1="{cut_x}" y1="{y0}" x2="{cut_x}" y2="{y0 + heat_h}" stroke="#111" stroke-width="2" stroke-dasharray="6 5"/>') |
| parts.append(f'<line x1="{x0}" y1="{cut_y}" x2="{x0 + heat_w}" y2="{cut_y}" stroke="#111" stroke-width="2" stroke-dasharray="6 5"/>') |
|
|
| |
| curve_x0 = x0 |
| curve_y0 = y0 + heat_h + gap |
| curve_w = heat_w |
| min_v = min(-0.05, min(row["cosine"] for row in adjacent) - 0.04) |
| max_v = max(1.0, max(row["cosine"] for row in adjacent) + 0.04) |
|
|
| def sx(idx: int) -> float: |
| if len(adjacent) == 1: |
| return curve_x0 + curve_w / 2 |
| return curve_x0 + idx * curve_w / (len(adjacent) - 1) |
|
|
| def sy(value: float) -> float: |
| return curve_y0 + curve_h - (value - min_v) / (max_v - min_v) * curve_h |
|
|
| parts.append(f'<text x="{curve_x0}" y="{curve_y0 - 18}" font-family="Arial, sans-serif" font-size="14" font-weight="700">Adjacent mask-ratio cosine</text>') |
| parts.append(f'<line x1="{curve_x0}" y1="{curve_y0}" x2="{curve_x0}" y2="{curve_y0 + curve_h}" stroke="#222" stroke-width="1"/>') |
| parts.append(f'<line x1="{curve_x0}" y1="{curve_y0 + curve_h}" x2="{curve_x0 + curve_w}" y2="{curve_y0 + curve_h}" stroke="#222" stroke-width="1"/>') |
| for tick in [0.0, 0.25, 0.5, 0.75, 1.0]: |
| if min_v <= tick <= max_v: |
| y = sy(tick) |
| parts.append(f'<line x1="{curve_x0 - 4}" y1="{y}" x2="{curve_x0 + curve_w}" y2="{y}" stroke="#dddddd" stroke-width="1"/>') |
| parts.append(f'<text x="{curve_x0 - 10}" y="{y + 4}" text-anchor="end" font-family="Arial, sans-serif" font-size="10">{tick:.2f}</text>') |
|
|
| points = [] |
| for i, row in enumerate(adjacent): |
| points.append(f"{sx(i):.2f},{sy(row['cosine']):.2f}") |
| parts.append(f'<polyline points="{" ".join(points)}" fill="none" stroke="#1f5fbf" stroke-width="2.5"/>') |
|
|
| point_label_every = max(1, len(adjacent) // 12) |
| for i, row in enumerate(adjacent): |
| x = sx(i) |
| y = sy(row["cosine"]) |
| fill = "#d62728" if i == cut_index else "#1f5fbf" |
| parts.append(f'<circle cx="{x}" cy="{y}" r="4.2" fill="{fill}" stroke="#ffffff" stroke-width="1"/>') |
| if i % point_label_every == 0 or i == len(adjacent) - 1 or i == cut_index: |
| parts.append(f'<text x="{x}" y="{curve_y0 + curve_h + 16}" text-anchor="middle" font-family="Arial, sans-serif" font-size="9">{row["left"]:g}-{row["right"]:g}</text>') |
| parts.append(f'<text x="{x}" y="{y - 8}" text-anchor="middle" font-family="Arial, sans-serif" font-size="9">{row["cosine"]:.2f}</text>') |
|
|
| if cut_index is not None: |
| x = sx(cut_index) |
| parts.append(f'<line x1="{x}" y1="{curve_y0}" x2="{x}" y2="{curve_y0 + curve_h}" stroke="#d62728" stroke-width="1.5" stroke-dasharray="5 4"/>') |
| parts.append(f'<text x="{x + 8}" y="{curve_y0 + 15}" font-family="Arial, sans-serif" font-size="11" fill="#d62728">candidate split</text>') |
|
|
| legend_y = curve_y0 + curve_h + 38 |
| parts.append(f'<text x="{x0}" y="{legend_y}" font-family="Arial, sans-serif" font-size="11" fill="#444">Heatmap: red = positive alignment, white = near zero, yellow/blue = negative. Curve: cosine between adjacent ratios.</text>') |
| parts.append("</svg>") |
| path.write_text("\n".join(parts) + "\n", encoding="utf-8") |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--input", type=Path, required=True) |
| parser.add_argument("--output-dir", type=Path, required=True) |
| parser.add_argument("--title", default="dLLM mask-ratio gradient alignment") |
| parser.add_argument("--cell-size", type=int, default=0) |
| parser.add_argument("--label-every", type=int, default=0) |
| parser.add_argument("--show-cell-text", action="store_true") |
| args = parser.parse_args() |
|
|
| source = json.loads(args.input.read_text(encoding="utf-8")) |
| ratios = [float(x) for x in source["ratios"]] |
| matrix = source["cosine_similarity"] |
| adjacent = adjacent_cosines(ratios, matrix) |
| cut_index = best_adjacent_cut(adjacent) |
| if cut_index is None: |
| boundary = "n/a" |
| else: |
| boundary = f"{adjacent[cut_index]['left']:g} -> {adjacent[cut_index]['right']:g}" |
|
|
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| figure = args.output_dir / "fig3_style_adjacent.svg" |
| summary = args.output_dir / "fig3_style_adjacent.json" |
| report = args.output_dir / "fig3_style_adjacent.md" |
| payload = { |
| "source_alignment": str(args.input), |
| "figure": str(figure), |
| "title": args.title, |
| "cell_size": args.cell_size, |
| "label_every": args.label_every, |
| "show_cell_text": args.show_cell_text, |
| "ratios": ratios, |
| "cosine_similarity": matrix, |
| "adjacent_cosines": adjacent, |
| "candidate_boundary_index": cut_index, |
| "candidate_boundary_label": boundary, |
| } |
| write_svg(figure, payload) |
| write_json(summary, payload) |
| write_markdown(report, payload) |
| print(json.dumps({ |
| "figure": str(figure), |
| "summary": str(summary), |
| "report": str(report), |
| "candidate_boundary": boundary, |
| }, indent=2)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
