1use crate::context::shared_wgpu_context;
4use crate::core::{
5 BoundingBox, DrawCall, GpuVertexBuffer, Material, PipelineType, RenderData, Vertex,
6};
7use crate::gpu::axis::OwnedAxisData;
8use crate::gpu::{util::readback_scalar_buffer_f64, ScalarType};
9use glam::{Vec3, Vec4};
10use std::sync::Arc;
11
12#[derive(Debug, Clone)]
13pub struct QuiverPlot {
14 pub x: Vec<f64>,
15 pub y: Vec<f64>,
16 pub z: Option<Vec<f64>>,
17 pub u: Vec<f64>,
18 pub v: Vec<f64>,
19 pub w: Option<Vec<f64>>,
20
21 pub color: Vec4,
22 pub line_width: f32,
23 pub scale: f32,
24 pub head_size: f32,
25
26 pub label: Option<String>,
27 pub visible: bool,
28
29 vertices: Option<Vec<Vertex>>,
30 bounds: Option<BoundingBox>,
31 dirty: bool,
32 gpu_vertices: Option<GpuVertexBuffer>,
33 gpu_vertex_count: Option<usize>,
34 gpu_bounds: Option<BoundingBox>,
35 gpu_source: Option<QuiverGpuSource>,
36}
37
38#[derive(Clone, Debug)]
39pub struct QuiverGpuSource {
40 pub x_data: OwnedAxisData,
41 pub y_data: OwnedAxisData,
42 pub u_buffer: Arc<wgpu::Buffer>,
43 pub v_buffer: Arc<wgpu::Buffer>,
44 pub count: usize,
45 pub rows: usize,
46 pub cols: usize,
47 pub xy_mode: u32,
48 pub scalar: ScalarType,
49}
50
51fn validate_gpu_source_metadata(
52 count: usize,
53 rows: usize,
54 cols: usize,
55 xy_mode: u32,
56) -> Result<(), String> {
57 match xy_mode {
58 0 => {
59 if count == 0 {
60 return Err("quiver plot GPU source has no vectors".to_string());
61 }
62 }
63 1 => {
64 if rows == 0 || cols == 0 || rows.checked_mul(cols) != Some(count) {
65 return Err("quiver plot GPU source has invalid meshgrid dimensions".to_string());
66 }
67 }
68 mode => {
69 return Err(format!(
70 "quiver plot GPU source has unsupported xy_mode {mode}"
71 ));
72 }
73 }
74 Ok(())
75}
76
77impl QuiverPlot {
78 pub async fn export_scene_vector_data(
79 &self,
80 ) -> Result<
81 (
82 Vec<f64>,
83 Vec<f64>,
84 Option<Vec<f64>>,
85 Vec<f64>,
86 Vec<f64>,
87 Option<Vec<f64>>,
88 ),
89 String,
90 > {
91 if !self.x.is_empty()
92 && self.x.len() == self.y.len()
93 && self.x.len() == self.u.len()
94 && self.x.len() == self.v.len()
95 && self.z.as_ref().is_none_or(|z| z.len() == self.x.len())
96 && self.w.as_ref().is_none_or(|w| w.len() == self.x.len())
97 {
98 return Ok((
99 self.x.clone(),
100 self.y.clone(),
101 self.z.clone(),
102 self.u.clone(),
103 self.v.clone(),
104 self.w.clone(),
105 ));
106 }
107 if !self.x.is_empty() || !self.y.is_empty() || !self.u.is_empty() || !self.v.is_empty() {
108 return Err(format!(
109 "quiver plot has incomplete CPU data: x={}, y={}, u={}, v={}",
110 self.x.len(),
111 self.y.len(),
112 self.u.len(),
113 self.v.len()
114 ));
115 }
116
117 if let Some(source) = &self.gpu_source {
118 validate_gpu_source_metadata(source.count, source.rows, source.cols, source.xy_mode)?;
119 let context = shared_wgpu_context().ok_or_else(|| {
120 "quiver plot has GPU source data but no shared WGPU context is installed"
121 .to_string()
122 })?;
123 let u = readback_scalar_buffer_f64(
124 &context.device,
125 &context.queue,
126 &source.u_buffer,
127 source.count,
128 source.scalar,
129 )
130 .await?;
131 let v = readback_scalar_buffer_f64(
132 &context.device,
133 &context.queue,
134 &source.v_buffer,
135 source.count,
136 source.scalar,
137 )
138 .await?;
139 let x_axis_len = if source.xy_mode == 0 {
140 source.count
141 } else {
142 source.cols
143 };
144 let y_axis_len = if source.xy_mode == 0 {
145 source.count
146 } else {
147 source.rows
148 };
149 let x_axis = source
150 .x_data
151 .export_f64(&context.device, &context.queue, x_axis_len, source.scalar)
152 .await?;
153 let y_axis = source
154 .y_data
155 .export_f64(&context.device, &context.queue, y_axis_len, source.scalar)
156 .await?;
157 let (x, y) = match source.xy_mode {
158 0 => {
159 if x_axis.len() != source.count || y_axis.len() != source.count {
160 return Err(format!(
161 "quiver plot GPU full-coordinate axes have lengths x={}, y={}, expected {}",
162 x_axis.len(),
163 y_axis.len(),
164 source.count
165 ));
166 }
167 (x_axis, y_axis)
168 }
169 1 => {
170 if x_axis.len() != source.cols || y_axis.len() != source.rows {
171 return Err(format!(
172 "quiver plot GPU meshgrid axes have lengths x={}, y={}, expected x={}, y={}",
173 x_axis.len(),
174 y_axis.len(),
175 source.cols,
176 source.rows
177 ));
178 }
179 let mut x = Vec::with_capacity(source.count);
180 let mut y = Vec::with_capacity(source.count);
181 for i in 0..source.count {
182 let col = i / source.rows;
183 let row = i % source.rows;
184 x.push(x_axis[col]);
185 y.push(y_axis[row]);
186 }
187 (x, y)
188 }
189 _ => unreachable!("xy_mode was validated before GPU readback"),
190 };
191 return Ok((x, y, None, u, v, None));
192 }
193
194 if self.gpu_vertices.is_some() {
195 return Err(
196 "quiver plot has GPU render vertices but no exportable source data".to_string(),
197 );
198 }
199
200 Ok((Vec::new(), Vec::new(), None, Vec::new(), Vec::new(), None))
201 }
202
203 pub fn new(x: Vec<f64>, y: Vec<f64>, u: Vec<f64>, v: Vec<f64>) -> Result<Self, String> {
204 let n = x.len();
205 if n == 0 || y.len() != n || u.len() != n || v.len() != n {
206 return Err("quiver: X,Y,U,V must have same non-zero length".to_string());
207 }
208 Ok(Self {
209 x,
210 y,
211 z: None,
212 u,
213 v,
214 w: None,
215 color: Vec4::new(0.0, 0.0, 0.0, 1.0),
216 line_width: 1.0,
217 scale: 1.0,
218 head_size: 0.1,
219 label: None,
220 visible: true,
221 vertices: None,
222 bounds: None,
223 dirty: true,
224 gpu_vertices: None,
225 gpu_vertex_count: None,
226 gpu_bounds: None,
227 gpu_source: None,
228 })
229 }
230 pub fn new3d(
231 x: Vec<f64>,
232 y: Vec<f64>,
233 z: Vec<f64>,
234 u: Vec<f64>,
235 v: Vec<f64>,
236 w: Vec<f64>,
237 ) -> Result<Self, String> {
238 let n = x.len();
239 if n == 0 || y.len() != n || z.len() != n || u.len() != n || v.len() != n || w.len() != n {
240 return Err("quiver3: X,Y,Z,U,V,W must have same non-zero length".to_string());
241 }
242 Ok(Self {
243 x,
244 y,
245 z: Some(z),
246 u,
247 v,
248 w: Some(w),
249 color: Vec4::new(0.0, 0.0, 0.0, 1.0),
250 line_width: 1.0,
251 scale: 1.0,
252 head_size: 0.1,
253 label: None,
254 visible: true,
255 vertices: None,
256 bounds: None,
257 dirty: true,
258 gpu_vertices: None,
259 gpu_vertex_count: None,
260 gpu_bounds: None,
261 gpu_source: None,
262 })
263 }
264 pub fn from_gpu_buffer(
265 color: Vec4,
266 line_width: f32,
267 scale: f32,
268 head_size: f32,
269 buffer: GpuVertexBuffer,
270 vertex_count: usize,
271 bounds: BoundingBox,
272 ) -> Self {
273 Self {
274 x: Vec::new(),
275 y: Vec::new(),
276 z: None,
277 u: Vec::new(),
278 v: Vec::new(),
279 w: None,
280 color,
281 line_width,
282 scale,
283 head_size,
284 label: None,
285 visible: true,
286 vertices: None,
287 bounds: Some(bounds),
288 dirty: false,
289 gpu_vertices: Some(buffer),
290 gpu_vertex_count: Some(vertex_count),
291 gpu_bounds: Some(bounds),
292 gpu_source: None,
293 }
294 }
295 pub fn with_gpu_source(mut self, source: QuiverGpuSource) -> Self {
296 self.gpu_source = Some(source);
297 self
298 }
299 pub fn with_style(mut self, color: Vec4, line_width: f32, scale: f32, head_size: f32) -> Self {
300 self.color = color;
301 self.line_width = line_width.max(0.5);
302 self.scale = scale.max(0.0);
303 self.head_size = head_size.max(0.0);
304 self.dirty = true;
305 self
306 }
307 pub fn with_label<S: Into<String>>(mut self, label: S) -> Self {
308 self.label = Some(label.into());
309 self
310 }
311 pub fn set_visible(&mut self, v: bool) {
312 self.visible = v;
313 }
314 pub fn has_cpu_vector_data(&self) -> bool {
315 !self.x.is_empty()
316 && self.x.len() == self.y.len()
317 && self.x.len() == self.u.len()
318 && self.x.len() == self.v.len()
319 && self.z.as_ref().is_none_or(|z| z.len() == self.x.len())
320 && self.w.as_ref().is_none_or(|w| w.len() == self.x.len())
321 }
322 pub fn cpu_vector_data_len(&self) -> Option<usize> {
323 self.has_cpu_vector_data().then_some(self.x.len())
324 }
325 pub fn mark_dirty(&mut self) {
326 self.dirty = true;
327 self.bounds = None;
328 }
329
330 pub fn generate_vertices(&mut self) -> &Vec<Vertex> {
331 if self.dirty || self.vertices.is_none() {
332 let mut verts = Vec::new();
333 for i in 0..self.x.len() {
334 let z = self.z.as_ref().map_or(0.0, |values| values[i]) as f32;
335 let w = self.w.as_ref().map_or(0.0, |values| values[i]) as f32;
336 let (x, y, u, v) = (
337 self.x[i] as f32,
338 self.y[i] as f32,
339 self.u[i] as f32,
340 self.v[i] as f32,
341 );
342 if !x.is_finite()
343 || !y.is_finite()
344 || !z.is_finite()
345 || !u.is_finite()
346 || !v.is_finite()
347 || !w.is_finite()
348 {
349 continue;
350 }
351 let dx = u * self.scale;
352 let dy = v * self.scale;
353 let dz = w * self.scale;
354 let base = Vec3::new(x, y, z);
356 let tip = Vec3::new(x + dx, y + dy, z + dz);
357 verts.push(Vertex::new(base, self.color));
358 verts.push(Vertex::new(tip, self.color));
359 let len = (dx * dx + dy * dy + dz * dz).sqrt();
361 if len > 0.0 && self.head_size > 0.0 {
362 let dir = Vec3::new(dx / len, dy / len, dz / len);
363 let reference = if dir.z.abs() > 0.9 { Vec3::Y } else { Vec3::Z };
364 let mut perp = dir.cross(reference);
365 if perp.length_squared() <= f32::EPSILON {
366 perp = dir.cross(Vec3::X);
367 }
368 let perp = perp.normalize_or_zero();
369 let h = self.head_size.min(len * 0.5);
370 let left = tip - h * dir + 0.5 * h * perp;
371 let right = tip - h * dir - 0.5 * h * perp;
372 verts.push(Vertex::new(tip, self.color));
373 verts.push(Vertex::new(left, self.color));
374 verts.push(Vertex::new(tip, self.color));
375 verts.push(Vertex::new(right, self.color));
376 }
377 }
378 self.vertices = Some(verts);
379 self.dirty = false;
380 }
381 self.vertices.as_ref().unwrap()
382 }
383
384 pub fn bounds(&mut self) -> BoundingBox {
385 if let Some(bounds) = self.gpu_bounds {
386 return bounds;
387 }
388 if self.dirty || self.bounds.is_none() {
389 let mut min = Vec3::new(f32::INFINITY, f32::INFINITY, f32::INFINITY);
390 let mut max = Vec3::new(f32::NEG_INFINITY, f32::NEG_INFINITY, f32::NEG_INFINITY);
391 for i in 0..self.x.len() {
392 let x = self.x[i] as f32;
393 let y = self.y[i] as f32;
394 let z = self.z.as_ref().map_or(0.0, |values| values[i]) as f32;
395 let dx = (self.u[i] as f32) * self.scale;
396 let dy = (self.v[i] as f32) * self.scale;
397 let dz = self
398 .w
399 .as_ref()
400 .map_or(0.0, |values| values[i] as f32 * self.scale);
401 if !x.is_finite()
402 || !y.is_finite()
403 || !z.is_finite()
404 || !dx.is_finite()
405 || !dy.is_finite()
406 || !dz.is_finite()
407 {
408 continue;
409 }
410 min.x = min.x.min(x.min(x + dx));
411 max.x = max.x.max(x.max(x + dx));
412 min.y = min.y.min(y.min(y + dy));
413 max.y = max.y.max(y.max(y + dy));
414 min.z = min.z.min(z.min(z + dz));
415 max.z = max.z.max(z.max(z + dz));
416 }
417 if !min.x.is_finite() {
418 min = Vec3::ZERO;
419 max = Vec3::ZERO;
420 }
421 self.bounds = Some(BoundingBox::new(min, max));
422 }
423 self.bounds.unwrap()
424 }
425
426 pub fn render_data(&mut self) -> RenderData {
427 let using_gpu = self.gpu_vertices.is_some();
428 let bounds = self.bounds();
429 let vertices = if using_gpu {
430 Vec::new()
431 } else {
432 self.generate_vertices().clone()
433 };
434 let material = Material {
435 albedo: self.color,
436 ..Default::default()
437 };
438 let draw_call = DrawCall {
439 vertex_offset: 0,
440 vertex_count: self.gpu_vertex_count.unwrap_or(vertices.len()),
441 index_offset: None,
442 index_count: None,
443 instance_count: 1,
444 };
445 RenderData {
446 pipeline_type: PipelineType::Lines,
447 vertices,
448 indices: None,
449 gpu_vertices: self.gpu_vertices.clone(),
450 bounds: Some(bounds),
451 material,
452 draw_calls: vec![draw_call],
453 image: None,
454 }
455 }
456
457 pub fn estimated_memory_usage(&self) -> usize {
458 self.vertices
459 .as_ref()
460 .map_or(0, |v| v.len() * std::mem::size_of::<Vertex>())
461 }
462}
463
464#[cfg(test)]
465mod tests {
466 use super::*;
467
468 #[test]
469 fn gpu_meshgrid_metadata_validation_rejects_invalid_dimensions() {
470 validate_gpu_source_metadata(6, 2, 3, 1).unwrap();
471
472 let err = validate_gpu_source_metadata(5, 2, 3, 1).unwrap_err();
473 assert!(err.contains("invalid meshgrid dimensions"));
474
475 let err = validate_gpu_source_metadata(6, 0, 3, 1).unwrap_err();
476 assert!(err.contains("invalid meshgrid dimensions"));
477
478 let err = validate_gpu_source_metadata(6, 2, 3, 7).unwrap_err();
479 assert!(err.contains("unsupported xy_mode 7"));
480 }
481}