1mod error;
2
3#[cfg(feature = "gpu")]
4mod gpu;
5
6mod backend;
7
8pub use backend::*;
9
10#[cfg(feature = "plotters")]
11pub mod colormap;
12
13use std::{marker::PhantomData, time::Duration};
14
15use autd3_driver::{
16 acoustics::{
17 directivity::{Directivity, Sphere},
18 propagate,
19 },
20 defined::{Complex, ULTRASOUND_PERIOD_COUNT},
21 error::AUTDInternalError,
22 firmware::{
23 cpu::{RxMessage, TxDatagram},
24 fpga::{EmitIntensity, Phase, Segment},
25 },
26 geometry::{Geometry, Vector3},
27 link::{Link, LinkBuilder},
28};
29use autd3_firmware_emulator::CPUEmulator;
30
31#[cfg(feature = "plotters")]
32pub use scarlet::colormap::ListedColorMap;
33
34use error::VisualizerError;
35
36pub struct Visualizer<D, B>
37where
38 D: Directivity,
39 B: Backend,
40{
41 geometry: Geometry,
42 is_open: bool,
43 timeout: Duration,
44 cpus: Vec<CPUEmulator>,
45 _d: PhantomData<D>,
46 _b: PhantomData<B>,
47 #[cfg(feature = "gpu")]
48 gpu_compute: Option<gpu::FieldCompute>,
49}
50
51pub struct VisualizerBuilder<D, B>
52where
53 D: Directivity,
54 B: Backend,
55{
56 backend: B,
57 timeout: Duration,
58 _d: PhantomData<D>,
59 #[cfg(feature = "gpu")]
60 gpu_idx: Option<i32>,
61}
62
63fn clone_geometry(geometry: &Geometry) -> Geometry {
64 Geometry::new(
65 geometry
66 .iter()
67 .map(|d| {
68 autd3_driver::geometry::Device::new(
69 d.idx(),
70 *d.rotation(),
71 d.iter().cloned().collect(),
72 )
73 })
74 .collect(),
75 )
76}
77
78#[cfg_attr(feature = "async-trait", autd3_driver::async_trait)]
79impl<D: Directivity, B: Backend> LinkBuilder for VisualizerBuilder<D, B> {
80 type L = Visualizer<D, B>;
81
82 #[allow(unused_mut)]
83 async fn open(mut self, geometry: &Geometry) -> Result<Self::L, AUTDInternalError> {
84 #[cfg(feature = "gpu")]
85 let gpu_compute = if let Some(gpu_idx) = self.gpu_idx {
86 Some(gpu::FieldCompute::new(gpu_idx)?)
87 } else {
88 None
89 };
90
91 let VisualizerBuilder {
92 mut backend,
93 timeout,
94 ..
95 } = self;
96
97 backend.initialize()?;
98
99 Ok(Self::L {
100 geometry: clone_geometry(geometry),
101 is_open: true,
102 timeout,
103 cpus: geometry
104 .iter()
105 .enumerate()
106 .map(|(i, dev)| {
107 let mut cpu = CPUEmulator::new(i, dev.num_transducers());
108 cpu.init();
109 cpu
110 })
111 .collect(),
112 _d: PhantomData,
113 _b: PhantomData,
114 #[cfg(feature = "gpu")]
115 gpu_compute,
116 })
117 }
118}
119
120#[derive(Clone, Debug)]
121pub struct PlotRange {
122 pub x_range: std::ops::Range<f32>,
123 pub y_range: std::ops::Range<f32>,
124 pub z_range: std::ops::Range<f32>,
125 pub resolution: f32,
126}
127
128impl PlotRange {
129 fn n(range: &std::ops::Range<f32>, resolution: f32) -> usize {
130 ((range.end - range.start) / resolution).floor() as usize + 1
131 }
132
133 pub fn nx(&self) -> usize {
134 Self::n(&self.x_range, self.resolution)
135 }
136
137 pub fn ny(&self) -> usize {
138 Self::n(&self.y_range, self.resolution)
139 }
140
141 pub fn nz(&self) -> usize {
142 Self::n(&self.z_range, self.resolution)
143 }
144
145 fn is_1d(&self) -> bool {
146 matches!(
147 (self.nx(), self.ny(), self.nz()),
148 (_, 1, 1) | (1, _, 1) | (1, 1, _)
149 )
150 }
151
152 fn is_2d(&self) -> bool {
153 if self.is_1d() {
154 return false;
155 }
156 matches!(
157 (self.nx(), self.ny(), self.nz()),
158 (1, _, _) | (_, 1, _) | (_, _, 1)
159 )
160 }
161
162 fn observe(n: usize, start: f32, resolution: f32) -> Vec<f32> {
163 (0..n).map(|i| start + resolution * i as f32).collect()
164 }
165
166 fn observe_x(&self) -> Vec<f32> {
167 Self::observe(self.nx(), self.x_range.start, self.resolution)
168 }
169
170 fn observe_y(&self) -> Vec<f32> {
171 Self::observe(self.ny(), self.y_range.start, self.resolution)
172 }
173
174 fn observe_z(&self) -> Vec<f32> {
175 Self::observe(self.nz(), self.z_range.start, self.resolution)
176 }
177
178 pub fn observe_points(&self) -> Vec<Vector3> {
179 match (self.nx(), self.ny(), self.nz()) {
180 (_, 1, 1) => self
181 .observe_x()
182 .iter()
183 .map(|&x| Vector3::new(x, self.y_range.start, self.z_range.start))
184 .collect(),
185 (1, _, 1) => self
186 .observe_y()
187 .iter()
188 .map(|&y| Vector3::new(self.x_range.start, y, self.z_range.start))
189 .collect(),
190 (1, 1, _) => self
191 .observe_z()
192 .iter()
193 .map(|&z| Vector3::new(self.x_range.start, self.y_range.start, z))
194 .collect(),
195 (_, _, 1) => itertools::iproduct!(self.observe_y(), self.observe_x())
196 .map(|(y, x)| Vector3::new(x, y, self.z_range.start))
197 .collect(),
198 (_, 1, _) => itertools::iproduct!(self.observe_x(), self.observe_z())
199 .map(|(x, z)| Vector3::new(x, self.y_range.start, z))
200 .collect(),
201 (1, _, _) => itertools::iproduct!(self.observe_z(), self.observe_y())
202 .map(|(z, y)| Vector3::new(self.x_range.start, y, z))
203 .collect(),
204 (_, _, _) => itertools::iproduct!(self.observe_z(), self.observe_y(), self.observe_x())
205 .map(|(z, y, x)| Vector3::new(x, y, z))
206 .collect(),
207 }
208 }
209}
210
211impl Visualizer<Sphere, PlottersBackend> {
212 pub fn builder() -> VisualizerBuilder<Sphere, PlottersBackend> {
213 VisualizerBuilder {
214 backend: PlottersBackend::new(),
215 timeout: Duration::ZERO,
216 _d: PhantomData,
217 #[cfg(feature = "gpu")]
218 gpu_idx: None,
219 }
220 }
221}
222
223#[cfg(feature = "plotters")]
224impl Visualizer<Sphere, PlottersBackend> {
225 pub fn plotters() -> VisualizerBuilder<Sphere, PlottersBackend> {
226 Self::builder()
227 }
228}
229
230#[cfg(feature = "python")]
231impl Visualizer<Sphere, PythonBackend> {
232 pub fn python() -> VisualizerBuilder<Sphere, PythonBackend> {
233 VisualizerBuilder {
234 backend: PythonBackend::new(),
235 timeout: Duration::ZERO,
236 _d: PhantomData,
237 #[cfg(feature = "gpu")]
238 gpu_idx: None,
239 }
240 }
241}
242
243impl Visualizer<Sphere, NullBackend> {
244 pub fn null() -> VisualizerBuilder<Sphere, NullBackend> {
245 VisualizerBuilder {
246 backend: NullBackend::new(),
247 timeout: Duration::ZERO,
248 _d: PhantomData,
249 #[cfg(feature = "gpu")]
250 gpu_idx: None,
251 }
252 }
253}
254
255impl<D: Directivity, B: Backend> VisualizerBuilder<D, B> {
256 pub fn with_directivity<U: Directivity>(self) -> VisualizerBuilder<U, B> {
257 VisualizerBuilder {
258 backend: self.backend,
259 timeout: self.timeout,
260 _d: PhantomData,
261 #[cfg(feature = "gpu")]
262 gpu_idx: self.gpu_idx,
263 }
264 }
265
266 pub fn with_backend<U: Backend>(self) -> VisualizerBuilder<D, U> {
267 VisualizerBuilder {
268 backend: U::new(),
269 timeout: self.timeout,
270 _d: PhantomData,
271 #[cfg(feature = "gpu")]
272 gpu_idx: self.gpu_idx,
273 }
274 }
275}
276
277#[cfg(feature = "gpu")]
278impl<D: Directivity, B: Backend> VisualizerBuilder<D, B> {
279 pub fn with_gpu(self, gpu_idx: i32) -> VisualizerBuilder<D, B> {
280 Self {
281 gpu_idx: Some(gpu_idx),
282 ..self
283 }
284 }
285}
286
287impl<D: Directivity, B: Backend> Visualizer<D, B> {
288 pub fn phases(&self, segment: Segment, idx: usize) -> Vec<Phase> {
289 self.cpus
290 .iter()
291 .flat_map(|cpu| {
292 cpu.fpga()
293 .drives(segment, idx)
294 .iter()
295 .map(|&d| d.phase())
296 .collect::<Vec<_>>()
297 })
298 .collect()
299 }
300
301 pub fn intensities(&self, segment: Segment, idx: usize) -> Vec<EmitIntensity> {
302 self.cpus
303 .iter()
304 .flat_map(|cpu| {
305 cpu.fpga()
306 .drives(segment, idx)
307 .iter()
308 .map(|&d| d.intensity())
309 .collect::<Vec<_>>()
310 })
311 .collect()
312 }
313
314 pub fn modulation(&self, segment: Segment) -> Vec<u8> {
315 self.cpus[0]
316 .fpga()
317 .modulation(segment)
318 .into_iter()
319 .collect()
320 }
321
322 pub fn calc_field<'a, I: IntoIterator<Item = &'a Vector3>>(
323 &self,
324 observe_points: I,
325 segment: Segment,
326 idx: usize,
327 ) -> Result<Vec<Complex>, VisualizerError> {
328 #[cfg(feature = "gpu")]
329 {
330 if let Some(gpu) = &self.gpu_compute {
331 let source_drive = self
332 .cpus
333 .iter()
334 .enumerate()
335 .flat_map(|(i, cpu)| {
336 let dev = &self.geometry[i];
337 let wavenumber = dev.wavenumber();
338 cpu.fpga()
339 .drives(segment, idx)
340 .iter()
341 .map(|d| {
342 let amp = (std::f32::consts::PI
343 * cpu.fpga().to_pulse_width(d.intensity(), u8::MAX) as f32
344 / ULTRASOUND_PERIOD_COUNT as f32)
345 .sin();
346 let phase = d.phase().radian();
347 [amp, phase, 0., wavenumber]
348 })
349 .collect::<Vec<_>>()
350 })
351 .collect::<Vec<_>>();
352 return gpu.calc_field_of::<D, I>(observe_points, &self.geometry, source_drive);
353 }
354 }
355 let source_drive = self
356 .cpus
357 .iter()
358 .map(|cpu| {
359 cpu.fpga()
360 .drives(segment, idx)
361 .into_iter()
362 .map(|d| {
363 let amp = (std::f32::consts::PI
364 * cpu.fpga().to_pulse_width(d.intensity(), u8::MAX) as f32
365 / ULTRASOUND_PERIOD_COUNT as f32)
366 .sin();
367 let phase = d.phase().radian();
368 [amp, phase]
369 })
370 .collect::<Vec<_>>()
371 })
372 .collect::<Vec<_>>();
373 Ok(observe_points
374 .into_iter()
375 .map(|target| {
376 self.cpus
377 .iter()
378 .enumerate()
379 .fold(Complex::new(0., 0.), |acc, (i, cpu)| {
380 let dir = self.geometry[i].axial_direction();
381 let wavenumber = self.geometry[i].wavenumber();
382 acc + self.geometry[i]
383 .iter()
384 .zip(source_drive[cpu.idx()].iter())
385 .fold(Complex::new(0., 0.), |acc, (t, &[amp, phase])| {
386 acc + propagate::<D>(t, wavenumber, dir, target)
387 * Complex::from_polar(amp, phase)
388 })
389 })
390 })
391 .collect())
392 }
393
394 pub fn plot_field(
395 &self,
396 config: B::PlotConfig,
397 range: PlotRange,
398 segment: Segment,
399 idx: usize,
400 ) -> Result<(), VisualizerError> {
401 let observe_points = range.observe_points();
402 let acoustic_pressures = self.calc_field(&observe_points, segment, idx)?;
403 if range.is_1d() {
404 let (observe, label) = match (range.nx(), range.ny(), range.nz()) {
405 (_, 1, 1) => (range.observe_x(), "x [mm]"),
406 (1, _, 1) => (range.observe_y(), "y [mm]"),
407 (1, 1, _) => (range.observe_z(), "z [mm]"),
408 _ => unreachable!(),
409 };
410 B::plot_1d(observe, acoustic_pressures, range.resolution, label, config)
411 } else if range.is_2d() {
412 let (observe_x, x_label) = match (range.nx(), range.ny(), range.nz()) {
413 (_, _, 1) => (range.observe_x(), "x [mm]"),
414 (1, _, _) => (range.observe_y(), "y [mm]"),
415 (_, 1, _) => (range.observe_z(), "z [mm]"),
416 _ => unreachable!(),
417 };
418 let (observe_y, y_label) = match (range.nx(), range.ny(), range.nz()) {
419 (_, _, 1) => (range.observe_y(), "y [mm]"),
420 (1, _, _) => (range.observe_z(), "z [mm]"),
421 (_, 1, _) => (range.observe_x(), "x [mm]"),
422 _ => unreachable!(),
423 };
424 B::plot_2d(
425 observe_x,
426 observe_y,
427 acoustic_pressures,
428 range.resolution,
429 x_label,
430 y_label,
431 config,
432 )
433 } else {
434 Err(VisualizerError::InvalidPlotRange)
435 }
436 }
437
438 pub fn plot_phase(
439 &self,
440 config: B::PlotConfig,
441 segment: Segment,
442 idx: usize,
443 ) -> Result<(), VisualizerError> {
444 let phases = self
445 .phases(segment, idx)
446 .iter()
447 .map(Phase::radian)
448 .collect();
449 B::plot_phase(config, &self.geometry, phases)
450 }
451
452 pub fn plot_modulation(
453 &self,
454 config: B::PlotConfig,
455 segment: Segment,
456 ) -> Result<(), VisualizerError> {
457 let m = self
458 .modulation(segment)
459 .iter()
460 .map(|&v| v as f32 / 255.0)
461 .collect::<Vec<_>>();
462 B::plot_modulation(m, config)?;
463 Ok(())
464 }
465}
466
467#[cfg_attr(feature = "async-trait", autd3_driver::async_trait)]
468impl<D: Directivity, B: Backend> Link for Visualizer<D, B> {
469 async fn close(&mut self) -> Result<(), AUTDInternalError> {
470 if !self.is_open {
471 return Ok(());
472 }
473
474 self.is_open = false;
475 Ok(())
476 }
477
478 async fn send(&mut self, tx: &TxDatagram) -> Result<bool, AUTDInternalError> {
479 if !self.is_open {
480 return Ok(false);
481 }
482
483 self.cpus.iter_mut().for_each(|cpu| {
484 cpu.send(tx);
485 });
486
487 Ok(true)
488 }
489
490 async fn receive(&mut self, rx: &mut [RxMessage]) -> Result<bool, AUTDInternalError> {
491 if !self.is_open {
492 return Ok(false);
493 }
494
495 self.cpus.iter_mut().for_each(|cpu| {
496 cpu.update();
497 rx[cpu.idx()] = cpu.rx();
498 });
499
500 Ok(true)
501 }
502
503 async fn update(&mut self, geometry: &Geometry) -> Result<(), AUTDInternalError> {
504 self.geometry = clone_geometry(geometry);
505 Ok(())
506 }
507
508 fn is_open(&self) -> bool {
509 self.is_open
510 }
511
512 fn timeout(&self) -> Duration {
513 self.timeout
514 }
515}