Research-Stack/5-Applications/scripts/run_manifold_render.py

257 lines
10 KiB
Python

#!/usr/bin/env python3
"""
Manifold GPU Render — wgpu Compute + Graphics Pipeline
Uploads binary manifold to GPU, runs WGSL compute sieve + projection,
then renders points with FAMM coloring.
Dependencies:
pip install wgpu numpy
# Also needs a WebGPU adapter (SwiftShader, Dawn, or native GPU)
Usage:
python3 run_manifold_render.py /tmp/test_manifold.bin
# Or generate test data and render:
python3 run_manifold_render.py --generate
"""
import sys
import struct
from pathlib import Path
import numpy as np
# wgpu is the Python WebGPU implementation
try:
import wgpu
import wgpu.backends.rs # Rust backend
WGPU_AVAILABLE = True
except ImportError:
WGPU_AVAILABLE = False
print("wgpu not installed. Install with: pip install wgpu")
print("Falling back to CPU preview (no GPU rendering).")
# Add ctypes shim path
sys.path.insert(0, str(Path("/home/allaun/Documents/Research Stack/4-Infrastructure/shims")))
from manifold_binary_ctypes import ManifoldSerializer, ManifoldBuilder, Q16_16
# ═══════════════════════════════════════════════════════════════════════════
# GPU Buffer Layout (must match manifold_binary.h and .wgsl)
# ═══════════════════════════════════════════════════════════════════════════
HEADER_SIZE = 32
Q16_SCALE = 65536.0
VERTEX_SIZE = 28 # vec3<f32> position (12) + vec3<f32> color (12) + flags (4) + idx (4)
def read_manifold_blob(path: str) -> bytes:
return Path(path).read_bytes()
def upload_and_render(blob: bytes, generate: bool = False):
if generate or not blob:
print("Generating test manifold...")
builder = ManifoldBuilder()
import time
builder.header.timestamp_ns = int(time.time_ns())
res = 64
for k in range(res):
for t in range(res // 2):
mass = (k + 1) * (t + 1)
a = int(mass ** 0.5) or 1
b = mass // a
theta = 2 * np.pi * k / res
phi = np.pi * t / res
builder.add_shell(k, t, mass, a, b, phase=theta)
w = np.cos(theta / 2) * np.cos(phi / 2)
x = np.sin(theta / 2) * np.cos(phi / 2)
y = np.sin(theta / 2) * np.sin(phi / 2)
z = np.cos(theta / 2) * np.sin(phi / 2)
builder.add_point(w, x, y, z, layer=k)
builder.add_famm_node(
torsional_stress=0.1 + 0.05 * np.sin(theta),
interlocking_energy=0.05 + 0.02 * np.cos(phi),
laplacian_energy=0.03,
cognitive_load=0.2 + 0.1 * np.sin(theta * 2)
)
for i in range(len(builder.points) - 1):
builder.add_edge(i, i + 1, weight=0.8, braid_id=i % 4)
ser = builder.build()
blob = ser.to_bytes()
print(f"Generated {len(blob)} bytes, {ser.header.num_points} points")
if not WGPU_AVAILABLE:
# CPU fallback: just verify the blob structure
ser = ManifoldSerializer()
ser._parse(blob)
print(f"CPU parse OK: {ser}")
errors = ser.validate_quaternions()
print(f"Quaternion validation errors: {errors}")
alive = ser.apply_sieve(threshold=0.3)
print(f"Alive after sieve: {alive} / {ser.header.num_points}")
return
# ── WebGPU Setup ──────────────────────────────────────────────
adapter = wgpu.request_adapter(power_preference="high-performance")
device = adapter.request_device()
# Read header to know sizes
header = struct.unpack("<IIIIIIII", blob[:HEADER_SIZE])
magic, version_and_flags, ts_lo, ts_hi, num_shells, num_points, num_edges, reserved = header
flags = (version_and_flags >> 16) & 0xFFFF
version = version_and_flags & 0xFFFF
print(f"GPU Upload: shells={num_shells}, points={num_points}, edges={num_edges}, flags={flags}")
# Compute sizes of arrays
shell_size = 28 # PistShell: 7 * 4 bytes
point_size = 24 # QuaternionPoint: 6 * 4 bytes
edge_size = 20 # BraidEdge: 5 * 4 bytes
famm_size = 24 # FammNode: 6 * 4 bytes
# Offsets in blob
off_shells = HEADER_SIZE
off_points = off_shells + num_shells * shell_size
off_edges = off_points + num_points * point_size
off_famm = off_edges + num_edges * edge_size
has_famm = (flags & 0x4) != 0
# ── Create GPU Buffers ────────────────────────────────────────
header_buf = device.create_buffer_with_data(
data=blob[:HEADER_SIZE], usage=wgpu.BufferUsage.STORAGE
)
shells_buf = device.create_buffer_with_data(
data=blob[off_shells:off_points] if num_shells > 0 else b"",
usage=wgpu.BufferUsage.STORAGE
)
points_buf = device.create_buffer_with_data(
data=blob[off_points:off_edges],
usage=wgpu.BufferUsage.STORAGE
)
edges_data = blob[off_edges:off_famm] if num_edges > 0 else b""
edges_buf = device.create_buffer_with_data(
data=edges_data,
usage=wgpu.BufferUsage.STORAGE
)
famm_data = blob[off_famm:off_famm + num_points * famm_size] if has_famm else b""
famm_buf = device.create_buffer_with_data(
data=famm_data,
usage=wgpu.BufferUsage.STORAGE
)
# Output: vertices
vertex_buf_size = max(num_points * VERTEX_SIZE, 16)
vertex_buf = device.create_buffer(
size=vertex_buf_size,
usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.VERTEX
)
# Output: draw params (indirect)
draw_buf = device.create_buffer(
size=16,
usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.INDIRECT
)
# ── Load WGSL Shader ─────────────────────────────────────────
wgsl_path = Path("/home/allaun/Documents/Research Stack/5-Applications/scripts/manifold_render.wgsl")
wgsl_code = wgsl_path.read_text()
shader = device.create_shader_module(code=wgsl_code)
# Bind group layout
bind_group_layout = device.create_bind_group_layout(
entries=[
{"binding": 0, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": "read-only-storage"}},
{"binding": 1, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": "read-only-storage"}},
{"binding": 2, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": "read-only-storage"}},
{"binding": 3, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": "read-only-storage"}},
{"binding": 4, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": "read-only-storage"}},
{"binding": 5, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": "storage"}},
{"binding": 6, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": "storage"}},
]
)
pipeline_layout = device.create_pipeline_layout(
bind_group_layouts=[bind_group_layout]
)
compute_pipeline = device.create_compute_pipeline(
layout=pipeline_layout,
compute={"module": shader, "entry_point": "main"}
)
bind_group = device.create_bind_group(
layout=bind_group_layout,
entries=[
{"binding": 0, "resource": {"buffer": header_buf}},
{"binding": 1, "resource": {"buffer": shells_buf}},
{"binding": 2, "resource": {"buffer": points_buf}},
{"binding": 3, "resource": {"buffer": edges_buf}},
{"binding": 4, "resource": {"buffer": famm_buf}},
{"binding": 5, "resource": {"buffer": vertex_buf}},
{"binding": 6, "resource": {"buffer": draw_buf}},
]
)
# ── Dispatch ──────────────────────────────────────────────────
command_encoder = device.create_command_encoder()
compute_pass = command_encoder.begin_compute_pass()
compute_pass.set_pipeline(compute_pipeline)
compute_pass.set_bind_group(0, bind_group)
workgroups = (num_points + 255) // 256
compute_pass.dispatch_workgroups(workgroups)
compute_pass.end()
# Copy draw params back to read
read_buf = device.create_buffer(
size=16, usage=wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.MAP_READ
)
command_encoder.copy_buffer_to_buffer(draw_buf, 0, read_buf, 0, 16)
device.queue.submit([command_encoder.finish()])
# Read back
def read_callback():
draw_data = np.frombuffer(read_buf.read_data(), dtype=np.uint32)
vertex_count = draw_data[0]
print(f"GPU Output: vertex_count={vertex_count}")
# Read first few vertices for sanity
vert_read = device.create_buffer(
size=min(vertex_count * VERTEX_SIZE, 1024),
usage=wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.MAP_READ
)
cmd = device.create_command_encoder()
cmd.copy_buffer_to_buffer(vertex_buf, 0, vert_read, 0, vert_read.size)
device.queue.submit([cmd.finish()])
verts = np.frombuffer(vert_read.read_data(), dtype=np.float32)
if len(verts) >= 7:
print(f" First vertex: pos=({verts[0]:.3f},{verts[1]:.3f},{verts[2]:.3f}), "
f"color=({verts[3]:.3f},{verts[4]:.3f},{verts[5]:.3f})")
read_callback()
print("GPU compute complete. Vertex buffer ready for rendering.")
def main():
import argparse
parser = argparse.ArgumentParser(description="Manifold GPU Render")
parser.add_argument("input", nargs="?", help="Binary manifold file")
parser.add_argument("--generate", action="store_true", help="Generate test data")
args = parser.parse_args()
blob = None
if args.input and Path(args.input).exists():
blob = read_manifold_blob(args.input)
print(f"Loaded {len(blob)} bytes from {args.input}")
upload_and_render(blob, generate=args.generate or blob is None)
if __name__ == "__main__":
main()