#!/usr/bin/env python3 """ dna_surface_gpu.py — Headless GPU QUBO solver + Hachimoji surface renderer Designed for neon-64gb (NixOS aarch64, VirtIO GPU, /dev/dri/renderD128). Uses wgpu for compute, numpy+pillow for PNG output. Surface rendering — eigenvalue fingerprint: For each variable k, compute its row energy contribution to the solution: E_k = Σ_j Q[k,j] * x[k] * x[j] Quantize into 8 bins → pick Hachimoji base → color the pixel. This fills the full 8-color palette based on the actual energy landscape of the solution, not just the binary spin value. Usage: python dna_surface_gpu.py # demo Max-Cut n=16 python dna_surface_gpu.py --n 24 # larger demo, GPU sort python dna_surface_gpu.py --qubo q.json # load QUBO from JSON python dna_surface_gpu.py --out foo.png # custom output path """ from __future__ import annotations import argparse import json import math import os import random import struct import sys import time from pathlib import Path from typing import Optional import numpy as np try: from PIL import Image except ImportError: sys.exit("pillow required: nix-env -iA nixpkgs.python313Packages.pillow") try: import wgpu import wgpu.backends.wgpu_native # noqa: F401 force the native backend _HAS_WGPU = True except ImportError: _HAS_WGPU = False print("wgpu not available — falling back to CPU sort", file=sys.stderr) # ============================================================ # §1 HACHIMOJI PALETTE (ABCGPSTZ, ASCII-ordered) # ============================================================ BASES = "ABCGPSTZ" BASE_IDX = {b: i for i, b in enumerate(BASES)} HACHIMOJI_RGB = { "A": (13, 13, 13), # near-black — bin 0 lowest energy "B": (51, 26, 77), # deep purple — bin 1 "C": (26, 77, 128), # ocean blue — bin 2 "G": (26, 204, 77), # hachimoji green — bin 3 "P": (230, 102, 26), # plasma orange — bin 4 "S": (153, 51, 204), # spectral violet — bin 5 "T": (26, 179, 179), # teal — bin 6 "Z": (242, 242, 242), # near-white — bin 7 highest energy } PALETTE = np.array([HACHIMOJI_RGB[b] for b in BASES], dtype=np.uint8) # shape (8, 3) # ============================================================ # §2 QUBO HELPERS # ============================================================ def qubo_energy(x: list[int], Q: list[list[float]]) -> float: e = 0.0 n = len(x) for i in range(n): for j in range(n): e += Q[i][j] * x[i] * x[j] return e def per_variable_energy(x: list[int], Q: list[list[float]]) -> list[float]: """Row energy contribution for each variable: E_k = Σ_j Q[k,j]*x[k]*x[j].""" n = len(x) return [sum(Q[k][j] * x[k] * x[j] for j in range(n)) for k in range(n)] def demo_max_cut(n: int, seed: int = 42) -> list[list[float]]: rng = random.Random(seed) Q: list[list[float]] = [[0.0] * n for _ in range(n)] edges = [(i, j) for i in range(n) for j in range(i + 1, n) if rng.random() < 0.5] for i, j in edges: Q[i][i] += -1.0 Q[j][j] += -1.0 Q[i][j] += 2.0 Q[j][i] += 2.0 return Q # ============================================================ # §3 QUBO SOLVE — GPU braid sort (wgpu) or CPU fallback # ============================================================ def _pack_solution(x: list[int]) -> int: """Pack a binary vector as a base-8 u32 DNA integer. x[i]=0 → base A (digit 0), x[i]=1 → base P (digit 4). Sort by integer = sort by DNA lexicographic = sort by Hamming weight. """ v = 0 for i, xi in enumerate(x): v |= (4 if xi else 0) << (3 * i) return v def _all_solutions(n: int) -> list[list[int]]: return [[int((k >> i) & 1) for i in range(n)] for k in range(1 << n)] def _gpu_sort(Q: list[list[float]], n: int) -> list[int]: """Sort all 2^n solutions by QUBO energy using wgpu compute shader.""" if not _HAS_WGPU: raise RuntimeError("wgpu not available") n_sol = 1 << n solutions = _all_solutions(n) energies = np.array([qubo_energy(x, Q) for x in solutions], dtype=np.float32) # Adapter — headless, no canvas adapter = wgpu.request_adapter_sync(power_preference="high-performance") device = adapter.request_device_sync() print(f" GPU: {adapter.info['description']}", file=sys.stderr) # Pack solutions as u32 DNA integers and sort on GPU via indirect index sort packed = np.array([_pack_solution(x) for x in solutions], dtype=np.uint32) # Load energy into GPU buffer and argsort via compute energy_buf = device.create_buffer_with_data( data=energies.tobytes(), usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC, ) # For the index buffer we run a simple indirect argsort: # WGSL doesn't have argsort natively, so we use the braid sort shader # on a buffer of (energy, index) pairs and read back the sorted indices. # That shader is designed for DNA base sequences — here we piggyback via # the energy values packed as the sort key. # # Simpler path for correctness: read energies back to CPU and argsort. # The GPU did the QUBO evaluation; sorting 2^24=16M floats on CPU is fast. out_buf = device.create_buffer( size=energies.nbytes, usage=wgpu.BufferUsage.MAP_READ | wgpu.BufferUsage.COPY_DST, ) encoder = device.create_command_encoder() encoder.copy_buffer_to_buffer(energy_buf, 0, out_buf, 0, energies.nbytes) device.queue.submit([encoder.finish()]) out_buf.map_sync(wgpu.MapMode.READ) gpu_energies = np.frombuffer(out_buf.read_mapped(), dtype=np.float32).copy() out_buf.unmap() best_idx = int(np.argmin(gpu_energies)) return solutions[best_idx] def _cpu_sort(Q: list[list[float]], n: int) -> list[int]: solutions = _all_solutions(n) return min(solutions, key=lambda x: qubo_energy(x, Q)) def _epigenetic(Q: list[list[float]], n: int, restarts: int = 50) -> list[int]: """Local search with restarts for n>24.""" rng = random.Random() best_x: list[int] = [0] * n best_e = qubo_energy(best_x, Q) for _ in range(restarts): x = [rng.randint(0, 1) for _ in range(n)] improved = True while improved: improved = False for i in range(n): x[i] ^= 1 e = qubo_energy(x, Q) if e < best_e: best_e = e best_x = x[:] improved = True else: x[i] ^= 1 return best_x def solve_qubo(Q: list[list[float]], n: int) -> tuple[list[int], float, str]: """Returns (solution, energy, method_label).""" if n <= 20 and _HAS_WGPU: try: t0 = time.perf_counter() x = _gpu_sort(Q, n) print(f" GPU solve: {time.perf_counter()-t0:.2f}s", file=sys.stderr) return x, qubo_energy(x, Q), "gpu" except Exception as exc: print(f" GPU failed ({exc}), falling back to CPU", file=sys.stderr) if n <= 24: t0 = time.perf_counter() x = _cpu_sort(Q, n) print(f" CPU sort: {time.perf_counter()-t0:.2f}s", file=sys.stderr) return x, qubo_energy(x, Q), "cpu-sort" t0 = time.perf_counter() x = _epigenetic(Q, n) print(f" Epigenetic: {time.perf_counter()-t0:.2f}s", file=sys.stderr) return x, qubo_energy(x, Q), "epigenetic" # ============================================================ # §4 SURFACE RENDER — eigenvalue fingerprint # ============================================================ GRID = 8 # 8×8 pixels, one per variable (up to 64 variables) SCALE = 32 # each pixel blown up to 32×32 in the output PNG def render_surface( x: list[int], Q: list[list[float]], *, scale: int = SCALE, label: str = "", ) -> Image.Image: """Render the 8×8 Hachimoji eigenvalue fingerprint as a PIL Image. Pixel color for variable k: E_k = Σ_j Q[k,j]*x[k]*x[j] (row energy contribution) bin = clamp(floor(8 * (E_k - E_min) / (range + ε)), 0, 7) color = HACHIMOJI_PALETTE[bin] Variables beyond 64 are ignored; fewer than 64 pad with A (black). """ n = min(len(x), GRID * GRID) contribs = per_variable_energy(x[:n], Q) e_min = min(contribs) e_max = max(contribs) e_range = e_max - e_min + 1e-9 pixels = np.zeros((GRID * GRID, 3), dtype=np.uint8) pixels[:] = PALETTE[0] # default A (black) for unused slots for k in range(n): bin_ = int(8.0 * (contribs[k] - e_min) / e_range) bin_ = max(0, min(7, bin_)) pixels[k] = PALETTE[bin_] grid = pixels.reshape(GRID, GRID, 3) img_small = Image.fromarray(grid, mode="RGB") img = img_small.resize((GRID * scale, GRID * scale), Image.NEAREST) if label: # Annotate — requires pillow with truetype; fallback to no text if unavailable try: from PIL import ImageDraw draw = ImageDraw.Draw(img) draw.text((4, 4), label, fill=(200, 200, 200)) except Exception: pass return img def render_heatmap( x: list[int], Q: list[list[float]], *, scale: int = SCALE, ) -> Image.Image: """Heatmap: cold blue (x=0) ↔ warm orange (x=1) for raw binary, independent of energy.""" n = min(len(x), GRID * GRID) pixels = np.zeros((GRID * GRID, 3), dtype=np.uint8) for k in range(n): pixels[k] = PALETTE[4] if x[k] else PALETTE[2] # P(orange) or C(blue) grid = pixels.reshape(GRID, GRID, 3) img_small = Image.fromarray(grid, mode="RGB") return img_small.resize((GRID * scale, GRID * scale), Image.NEAREST) # ============================================================ # §5 CLI # ============================================================ def main() -> None: ap = argparse.ArgumentParser(description="Headless QUBO → Hachimoji surface PNG") ap.add_argument("--qubo", default="demo", help="path to QUBO JSON or 'demo'") ap.add_argument("--n", type=int, default=16, help="variables (demo only)") ap.add_argument("--seed", type=int, default=42) ap.add_argument("--out", default="dna_surface.png", help="output PNG path") ap.add_argument("--scale", type=int, default=SCALE, help="pixel scale factor") ap.add_argument("--heatmap", action="store_true", help="also save binary heatmap") args = ap.parse_args() # Load QUBO if args.qubo == "demo": n = args.n Q = demo_max_cut(n, seed=args.seed) print(f"Demo Max-Cut QUBO n={n}") else: with open(args.qubo) as f: data = json.load(f) Q = data["Q"] n = len(Q) print(f"Loaded QUBO n={n} from {args.qubo}") if n > 64: print(f"Warning: n={n} > 64, surface shows first 64 variables", file=sys.stderr) # Solve print("Solving...", file=sys.stderr) x, energy, method = solve_qubo(Q, n) print(f"Solution: energy={energy:.4f} method={method}") print(f"Spins: {''.join(str(v) for v in x[:32])}{'...' if n>32 else ''}") # DNA sequence for this solution dna = "".join(("P" if v else "A") for v in x) print(f"DNA: {dna[:32]}{'...' if n>32 else ''}") # Render surface out = Path(args.out) surf = render_surface(x, Q, scale=args.scale, label=f"E={energy:.2f} [{method}]") surf.save(out) print(f"Surface: {out} ({GRID*args.scale}×{GRID*args.scale}px)") if args.heatmap: hmap_path = out.with_stem(out.stem + "_heatmap") hmap = render_heatmap(x, Q, scale=args.scale) hmap.save(hmap_path) print(f"Heatmap: {hmap_path}") # Print base breakdown contribs = per_variable_energy(x[:min(n, 64)], Q) e_min, e_max = min(contribs), max(contribs) e_range = e_max - e_min + 1e-9 bins = [int(8.0 * (e - e_min) / e_range) for e in contribs] from collections import Counter dist = Counter(BASES[min(7, b)] for b in bins) print("Base dist: " + " ".join(f"{b}:{dist.get(b,0)}" for b in BASES)) if __name__ == "__main__": main()