use std::collections::BTreeSet;
#[derive(Clone, Debug, PartialEq)]
pub struct GraphCsr {
pub n: usize,
pub offsets: Vec<u32>,
pub targets: Vec<u32>,
pub weights: Vec<f32>,
}
impl GraphCsr {
#[must_use]
pub fn from_edges(n: usize, edges: &[(u32, u32)], weights: Option<&[f32]>) -> Self {
let mut adj: Vec<BTreeSet<u32>> = vec![BTreeSet::new(); n];
for &(a, b) in edges {
let (a, b) = (a as usize, b as usize);
if a == b || a >= n || b >= n {
continue;
}
adj[a].insert(b as u32);
adj[b].insert(a as u32);
}
let mut offsets = Vec::with_capacity(n + 1);
let mut targets = Vec::new();
offsets.push(0u32);
for nbrs in &adj {
targets.extend(nbrs.iter().copied());
offsets.push(targets.len() as u32);
}
let weights = (0..n)
.map(|i| weights.and_then(|w| w.get(i)).copied().filter(|w| *w > 0.0).unwrap_or(1.0))
.collect();
Self { n, offsets, targets, weights }
}
#[must_use]
pub fn degree_sum(&self) -> usize {
self.targets.len()
}
#[must_use]
pub fn neighbours(&self, i: usize) -> &[u32] {
let (s, e) = (self.offsets[i] as usize, self.offsets[i + 1] as usize);
&self.targets[s..e]
}
}
#[cfg(feature = "arrow")]
impl GraphCsr {
pub fn from_arrow(
n: usize,
src: &arrow_array::Int64Array,
dst: &arrow_array::Int64Array,
weight: Option<&arrow_array::Float64Array>,
) -> Self {
use arrow_array::Array;
let mut edges = Vec::with_capacity(src.len());
for i in 0..src.len().min(dst.len()) {
if src.is_null(i) || dst.is_null(i) {
continue;
}
let (a, b) = (src.value(i), dst.value(i));
if a < 0 || b < 0 {
continue;
}
edges.push((a as u32, b as u32));
}
let w: Option<Vec<f32>> = weight.map(|w| {
(0..n).map(|i| if i < w.len() && !w.is_null(i) { w.value(i) as f32 } else { 1.0 }).collect()
});
Self::from_edges(n, &edges, w.as_deref())
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct LayoutParams {
pub repulsion: f32,
pub attraction: f32,
pub ideal_len: f32,
pub gravity: f32,
pub damping: f32,
pub dt: f32,
pub min_dist: f32,
}
impl Default for LayoutParams {
fn default() -> Self {
Self {
repulsion: 1.0,
attraction: 2.0,
ideal_len: 1.0,
gravity: 0.05,
damping: 0.9,
dt: 0.05,
min_dist: 0.05,
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct LayoutState {
pub pos: Vec<[f32; 2]>,
pub vel: Vec<[f32; 2]>,
}
impl LayoutState {
#[must_use]
pub fn seed(n: usize, scale: f32) -> Self {
const GOLDEN: f32 = 2.399_963_2; let pos = (0..n)
.map(|i| {
let r = scale * ((i as f32 + 0.5) / n.max(1) as f32).sqrt();
let a = i as f32 * GOLDEN;
[r * a.cos(), r * a.sin()]
})
.collect();
Self { pos, vel: vec![[0.0, 0.0]; n] }
}
#[must_use]
pub fn kinetic_energy(&self) -> f32 {
self.vel.iter().map(|v| v[0] * v[0] + v[1] * v[1]).sum()
}
#[must_use]
pub fn total_edge_len(&self, g: &GraphCsr) -> f32 {
let mut sum = 0.0;
for i in 0..g.n {
for &t in g.neighbours(i) {
if (t as usize) > i {
let d = [self.pos[t as usize][0] - self.pos[i][0], self.pos[t as usize][1] - self.pos[i][1]];
sum += (d[0] * d[0] + d[1] * d[1]).sqrt();
}
}
}
sum
}
}
#[must_use]
pub fn step_cpu(s: &LayoutState, g: &GraphCsr, p: &LayoutParams) -> LayoutState {
let n = g.n;
let mut out = LayoutState { pos: Vec::with_capacity(n), vel: Vec::with_capacity(n) };
let min2 = p.min_dist * p.min_dist;
for i in 0..n {
let pi = s.pos[i];
let wi = g.weights[i];
let mut force = [0.0f32, 0.0f32];
for (j, pj) in s.pos.iter().enumerate() {
if j == i {
continue;
}
let d = [pi[0] - pj[0], pi[1] - pj[1]];
let dist2 = (d[0] * d[0] + d[1] * d[1]).max(min2);
let f = p.repulsion * wi * g.weights[j] / dist2;
force[0] += d[0] * f;
force[1] += d[1] * f;
}
for &t in g.neighbours(i) {
let pt = s.pos[t as usize];
let d = [pt[0] - pi[0], pt[1] - pi[1]];
let dist = (d[0] * d[0] + d[1] * d[1]).sqrt().max(1e-4);
let f = p.attraction * (dist - p.ideal_len) / dist;
force[0] += d[0] * f;
force[1] += d[1] * f;
}
force[0] -= pi[0] * p.gravity;
force[1] -= pi[1] * p.gravity;
let inv_m = 1.0 / wi.max(1e-4);
let vx = (s.vel[i][0] + force[0] * inv_m * p.dt) * p.damping;
let vy = (s.vel[i][1] + force[1] * inv_m * p.dt) * p.damping;
out.vel.push([vx, vy]);
out.pos.push([pi[0] + vx * p.dt, pi[1] + vy * p.dt]);
}
out
}
#[must_use]
pub fn relax_cpu(g: &GraphCsr, p: &LayoutParams, iters: usize) -> LayoutState {
let mut s = LayoutState::seed(g.n, p.ideal_len * (g.n as f32).sqrt());
for _ in 0..iters {
s = step_cpu(&s, g, p);
}
s
}
#[cfg(feature = "gpu")]
mod gpu {
use super::{GraphCsr, LayoutParams, LayoutState};
use wgpu::util::DeviceExt;
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
pub struct GpuNode {
pub a: [f32; 4],
pub b: [f32; 4],
}
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct Uniforms {
a: [f32; 4], b: [f32; 4], }
pub const LAYOUT_WGSL: &str = include_str!("layout.wgsl");
pub fn pack(s: &LayoutState, weights: &[f32]) -> Vec<GpuNode> {
(0..s.pos.len())
.map(|i| GpuNode { a: [s.pos[i][0], s.pos[i][1], s.vel[i][0], s.vel[i][1]], b: [weights[i], 0.0, 0.0, 0.0] })
.collect()
}
pub struct GpuLayout {
pipeline: wgpu::ComputePipeline,
bgl: wgpu::BindGroupLayout,
uniforms: wgpu::Buffer,
offsets: wgpu::Buffer,
targets: wgpu::Buffer,
src: wgpu::Buffer,
dst: wgpu::Buffer,
count: u32,
}
impl GpuLayout {
pub fn new(device: &wgpu::Device, g: &GraphCsr, seed: &LayoutState) -> Self {
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("graphview_layout"),
source: wgpu::ShaderSource::Wgsl(LAYOUT_WGSL.into()),
});
let storage = |ro: bool| wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: ro }, has_dynamic_offset: false, min_binding_size: None },
count: None,
};
let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("graphview_layout_bgl"),
entries: &[
wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None },
count: None,
},
wgpu::BindGroupLayoutEntry { binding: 1, ..storage(true) },
wgpu::BindGroupLayoutEntry { binding: 2, ..storage(false) },
wgpu::BindGroupLayoutEntry { binding: 3, ..storage(true) },
wgpu::BindGroupLayoutEntry { binding: 4, ..storage(true) },
],
});
let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some("graphview_layout_step"),
layout: Some(&device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some("graphview_layout_pll"),
bind_group_layouts: &[Some(&bgl)],
immediate_size: 0,
})),
module: &shader,
entry_point: Some("cs_layout"),
compilation_options: Default::default(),
cache: None,
});
let nodes = pack(seed, &g.weights);
let node_usage = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST;
let src = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("graphview_layout_src"),
contents: bytemuck::cast_slice(&nodes),
usage: node_usage,
});
let dst = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("graphview_layout_dst"),
size: (std::mem::size_of::<GpuNode>() * nodes.len().max(1)) as u64,
usage: node_usage,
mapped_at_creation: false,
});
let off_data: &[u32] = if g.offsets.is_empty() { &[0] } else { &g.offsets };
let tgt_data: Vec<u32> = if g.targets.is_empty() { vec![0] } else { g.targets.clone() };
let offsets = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("graphview_layout_offsets"),
contents: bytemuck::cast_slice(off_data),
usage: wgpu::BufferUsages::STORAGE,
});
let targets = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("graphview_layout_targets"),
contents: bytemuck::cast_slice(&tgt_data),
usage: wgpu::BufferUsages::STORAGE,
});
let uniforms = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("graphview_layout_uniforms"),
size: std::mem::size_of::<Uniforms>() as u64,
usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
Self { pipeline, bgl, uniforms, offsets, targets, src, dst, count: nodes.len() as u32 }
}
pub fn iterate(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, g: &GraphCsr, p: &LayoutParams, iters: usize) {
let u = Uniforms {
a: [self.count as f32, p.repulsion, p.attraction, p.ideal_len],
b: [p.gravity, p.damping, p.dt, p.min_dist],
};
queue.write_buffer(&self.uniforms, 0, bytemuck::bytes_of(&u));
let _ = g; let mut enc = device.create_command_encoder(&Default::default());
for _ in 0..iters {
let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("graphview_layout_bind"),
layout: &self.bgl,
entries: &[
wgpu::BindGroupEntry { binding: 0, resource: self.uniforms.as_entire_binding() },
wgpu::BindGroupEntry { binding: 1, resource: self.src.as_entire_binding() },
wgpu::BindGroupEntry { binding: 2, resource: self.dst.as_entire_binding() },
wgpu::BindGroupEntry { binding: 3, resource: self.offsets.as_entire_binding() },
wgpu::BindGroupEntry { binding: 4, resource: self.targets.as_entire_binding() },
],
});
{
let mut cp = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("graphview_layout_pass"), timestamp_writes: None });
cp.set_pipeline(&self.pipeline);
cp.set_bind_group(0, &bind, &[]);
cp.dispatch_workgroups(self.count.div_ceil(64), 1, 1);
}
std::mem::swap(&mut self.src, &mut self.dst);
}
queue.submit(Some(enc.finish()));
device.poll(wgpu::PollType::wait_indefinitely()).ok();
}
pub fn read_positions(&self, device: &wgpu::Device, queue: &wgpu::Queue) -> Vec<[f32; 2]> {
let size = (std::mem::size_of::<GpuNode>() as u32 * self.count) as u64;
let readback = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("graphview_layout_readback"),
size,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
});
let mut enc = device.create_command_encoder(&Default::default());
enc.copy_buffer_to_buffer(&self.src, 0, &readback, 0, size);
queue.submit(Some(enc.finish()));
let slice = readback.slice(..);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| { let _ = tx.send(r); });
device.poll(wgpu::PollType::wait_indefinitely()).ok();
rx.recv().unwrap().unwrap();
let data = slice.get_mapped_range();
let nodes: Vec<GpuNode> = bytemuck::cast_slice(&data).to_vec();
drop(data);
readback.unmap();
nodes.iter().map(|nd| [nd.a[0], nd.a[1]]).collect()
}
}
}
#[cfg(feature = "gpu")]
pub use gpu::{GpuLayout, GpuNode, LAYOUT_WGSL};
#[cfg(test)]
mod tests {
use super::*;
fn two_triangles() -> GraphCsr {
let edges = [(0, 1), (1, 2), (2, 0), (3, 4), (4, 5), (5, 3), (2, 3)];
GraphCsr::from_edges(6, &edges, None)
}
#[test]
fn csr_is_symmetric_and_dedups() {
let g = GraphCsr::from_edges(3, &[(0, 1), (1, 0), (0, 1), (1, 2)], None);
assert_eq!(g.neighbours(0), &[1]);
assert_eq!(g.neighbours(1), &[0, 2]);
assert_eq!(g.neighbours(2), &[1]);
assert_eq!(g.degree_sum(), 4, "two undirected edges → 4 directed entries");
let g2 = GraphCsr::from_edges(2, &[(0, 0), (0, 5), (0, 1)], None);
assert_eq!(g2.neighbours(0), &[1]);
}
#[test]
fn seed_is_deterministic_and_bounded() {
let a = LayoutState::seed(50, 3.0);
let b = LayoutState::seed(50, 3.0);
assert_eq!(a, b, "no RNG: the seed is reproducible");
for q in &a.pos {
assert!(q[0].is_finite() && q[1].is_finite());
assert!((q[0] * q[0] + q[1] * q[1]).sqrt() <= 3.0 + 1e-3, "inside the seed disc");
}
}
#[test]
fn step_cpu_is_deterministic() {
let g = two_triangles();
let p = LayoutParams::default();
let s = LayoutState::seed(g.n, 2.0);
let a = step_cpu(&s, &g, &p);
let b = step_cpu(&s, &g, &p);
for (x, y) in a.pos.iter().zip(&b.pos) {
assert_eq!(x[0].to_bits(), y[0].to_bits(), "bit-identical pos.x");
assert_eq!(x[1].to_bits(), y[1].to_bits(), "bit-identical pos.y");
}
}
#[test]
fn layout_converges_energy_decays() {
let g = two_triangles();
let p = LayoutParams::default();
let mut s = LayoutState::seed(g.n, p.ideal_len * (g.n as f32).sqrt());
for _ in 0..15 {
s = step_cpu(&s, &g, &p);
}
let e_early = s.kinetic_energy();
for _ in 0..200 {
s = step_cpu(&s, &g, &p);
}
let e_late = s.kinetic_energy();
assert!(e_late < e_early * 0.1, "energy decayed an order of magnitude (converged): {e_early} → {e_late}");
assert!(e_late < 0.05, "settled to near-rest: {e_late}");
for q in &s.pos {
assert!(q[0].is_finite() && q[1].is_finite(), "no NaN/inf blow-up");
}
}
#[test]
fn springs_settle_an_edge_at_the_analytic_equilibrium() {
let g = GraphCsr::from_edges(2, &[(0, 1)], None);
let p = LayoutParams { gravity: 0.0, ..LayoutParams::default() };
let s = relax_cpu(&g, &p, 600);
let d = ((s.pos[1][0] - s.pos[0][0]).powi(2) + (s.pos[1][1] - s.pos[0][1]).powi(2)).sqrt();
let expected = (1.0 + 3.0_f32.sqrt()) / 2.0;
assert!((d - expected).abs() < 0.02, "edge rests at the force equilibrium {expected} (got {d})");
}
#[test]
fn relax_is_deterministic_positions_are_data() {
let g = two_triangles();
let p = LayoutParams::default();
let a = relax_cpu(&g, &p, 120);
let b = relax_cpu(&g, &p, 120);
assert_eq!(a.pos, b.pos, "same params → identical positions (reproducible DATA)");
let len = |i: usize, j: usize| ((a.pos[j][0] - a.pos[i][0]).powi(2) + (a.pos[j][1] - a.pos[i][1]).powi(2)).sqrt();
let within = len(0, 1).max(len(1, 2)).max(len(3, 4)).max(len(4, 5));
assert!(len(2, 3) > within * 0.9, "the two clusters spread apart along the bridge");
}
#[test]
fn heavier_nodes_move_less_per_step() {
let edges = [(0, 1), (0, 2)];
let light = GraphCsr::from_edges(3, &edges, None);
let heavy = GraphCsr::from_edges(3, &edges, Some(&[50.0, 1.0, 1.0]));
let p = LayoutParams::default();
let s = LayoutState::seed(3, 2.0);
let dl = step_cpu(&s, &light, &p);
let dh = step_cpu(&s, &heavy, &p);
let disp = |st: &LayoutState| ((st.pos[0][0] - s.pos[0][0]).powi(2) + (st.pos[0][1] - s.pos[0][1]).powi(2)).sqrt();
assert!(disp(&dh) < disp(&dl) * 0.5, "the heavy node barely moves ({} vs {})", disp(&dh), disp(&dl));
}
#[cfg(feature = "gpu")]
fn compute_device() -> Option<(wgpu::Device, wgpu::Queue)> {
let instance = wgpu::Instance::default();
let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::default(),
force_fallback_adapter: false,
compatible_surface: None,
}))
.ok()?;
if !adapter.get_downlevel_capabilities().flags.contains(wgpu::DownlevelFlags::COMPUTE_SHADERS) {
return None;
}
if adapter.limits().max_storage_buffers_per_shader_stage < 4 {
return None;
}
pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
label: Some("graphview-layout-proof"),
required_features: wgpu::Features::empty(),
required_limits: adapter.limits(),
memory_hints: wgpu::MemoryHints::default(),
experimental_features: wgpu::ExperimentalFeatures::disabled(),
trace: wgpu::Trace::Off,
}))
.ok()
}
#[cfg(feature = "gpu")]
#[test]
fn gpu_layout_matches_cpu_reference() {
let Some((device, queue)) = compute_device() else {
eprintln!("[gpu_layout] no compute device — skipping GPU parity proof");
return;
};
let g = two_triangles();
let p = LayoutParams::default();
let iters = 50;
let seed = LayoutState::seed(g.n, p.ideal_len * (g.n as f32).sqrt());
let mut gl = GpuLayout::new(&device, &g, &seed);
gl.iterate(&device, &queue, &g, &p, iters);
let gpu = gl.read_positions(&device, &queue);
let mut cpu = seed.clone();
for _ in 0..iters {
cpu = step_cpu(&cpu, &g, &p);
}
assert_eq!(gpu.len(), cpu.pos.len());
let mut moved = false;
for (gp, cp) in gpu.iter().zip(&cpu.pos) {
assert!((gp[0] - cp[0]).abs() < 1e-2, "pos.x parity: gpu {} vs cpu {}", gp[0], cp[0]);
assert!((gp[1] - cp[1]).abs() < 1e-2, "pos.y parity: gpu {} vs cpu {}", gp[1], cp[1]);
assert!(gp[0].is_finite() && gp[1].is_finite(), "readback is finite DATA");
if (gp[0] - seed.pos[0][0]).abs() > 1e-3 {
moved = true;
}
}
assert!(moved, "the VRAM iteration actually advanced the layout");
}
#[cfg(feature = "gpu")]
#[test]
fn layout_shader_entry_and_node_size() {
assert!(LAYOUT_WGSL.contains("fn cs_layout"));
assert!(LAYOUT_WGSL.contains("@workgroup_size(64)"));
assert_eq!(std::mem::size_of::<GpuNode>(), 32);
}
#[cfg(feature = "arrow")]
#[test]
fn csr_streams_from_an_arrow_edge_batch() {
use arrow_array::{Float64Array, Int64Array};
let src = Int64Array::from(vec![0, 1, 2, 0]);
let dst = Int64Array::from(vec![1, 2, 0, 2]);
let w = Float64Array::from(vec![3.0, 1.0, 1.0]);
let g = GraphCsr::from_arrow(3, &src, &dst, Some(&w));
assert_eq!(g.neighbours(0), &[1, 2]);
assert_eq!(g.weights[0], 3.0, "weight column streamed in");
let s = relax_cpu(&g, &LayoutParams::default(), 80);
assert!(s.pos.iter().all(|q| q[0].is_finite() && q[1].is_finite()));
}
}