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