cellular_raza_core/backend/cpu_os_threads/
supervisor.rs

1use super::{Agent, ForceBound, InteractionInf, PositionBound, VelocityBound};
2use cellular_raza_concepts::domain_old::*;
3use cellular_raza_concepts::reactions_old::*;
4use cellular_raza_concepts::*;
5use kdam::BarExt;
6
7use super::errors::*;
8use crate::storage::{StorageBuilder, StorageManager};
9
10use super::domain_decomposition::{
11    AuxiliaryCellPropertyStorage, DomainBox, MultiVoxelContainer, VoxelBox,
12};
13
14use super::config::{
15    ImageType, PlottingConfig, SimulationConfig, SimulationMetaParams, SimulationSetup, TimeSetup,
16};
17
18use std::sync::Arc;
19use std::sync::atomic::{AtomicBool, Ordering};
20use std::thread;
21
22use core::ops::{Add, AddAssign, Mul};
23
24use core::marker::PhantomData;
25
26use hurdles::Barrier;
27
28use serde::{Deserialize, Serialize};
29
30use plotters::{
31    coord::types::RangedCoordf64,
32    prelude::{BitMapBackend, Cartesian2d, DrawingArea},
33};
34
35#[cfg(feature = "tracing")]
36use tracing::instrument;
37
38#[derive(Clone, Serialize, Deserialize)]
39pub(crate) struct ControllerBox<Cont, Obs> {
40    pub controller: Cont,
41    pub measurements: std::collections::BTreeMap<u32, Obs>,
42}
43
44impl<Cont, Obs> ControllerBox<Cont, Obs> {
45    fn measure<'a, I, Cel>(&mut self, thread_index: u32, cells: I) -> Result<(), SimulationError>
46    where
47        Cel: 'a + Serialize + for<'b> Deserialize<'b>,
48        I: Iterator<Item = &'a CellBox<Cel>> + Clone,
49        Cont: Controller<Cel, Obs>,
50    {
51        let obs = self.controller.measure(cells)?;
52        self.measurements.insert(thread_index, obs);
53        Ok(())
54    }
55
56    fn adjust<'a, Cel, J>(&mut self, cells: J) -> Result<(), ControllerError>
57    where
58        Cel: 'a + Serialize + for<'b> Deserialize<'b>,
59        J: Iterator<
60            Item = (
61                &'a mut CellBox<Cel>,
62                &'a mut Vec<cellular_raza_concepts::CycleEvent>,
63            ),
64        >,
65        Cont: Controller<Cel, Obs>,
66    {
67        self.controller.adjust(self.measurements.values(), cells)
68    }
69}
70
71/// # Supervisor controlling simulation execution
72///
73pub struct SimulationSupervisor<MVC, Dom, Cel, Cont = (), Obs = ()>
74where
75    Cel: Serialize + for<'a> Deserialize<'a>,
76    Dom: Serialize + for<'a> Deserialize<'a>,
77    Cont: Serialize + for<'a> Deserialize<'a>,
78{
79    pub(crate) worker_threads: Vec<thread::JoinHandle<Result<MVC, SimulationError>>>,
80    pub(crate) multivoxelcontainers: Vec<MVC>,
81
82    pub(crate) time: TimeSetup,
83    pub(crate) meta_params: SimulationMetaParams,
84    /// Defines in which format and if results should be saved.
85    /// See [StorageBuilder].
86    pub storage: StorageBuilder<true>,
87
88    /// Physical [Domain] of the simulation.
89    pub(crate) domain: DomainBox<Dom>,
90
91    /// Overall parameters of the simulation. See [SimulationConfig].
92    pub config: SimulationConfig,
93
94    pub(crate) meta_infos: StorageManager<(), SimulationSetup<DomainBox<Dom>, Cel, Cont>>,
95
96    pub(crate) controller_box: Arc<std::sync::Mutex<ControllerBox<Cont, Obs>>>,
97    pub(crate) phantom_obs: PhantomData<Obs>,
98    pub(crate) phantom_cont: PhantomData<Cont>,
99}
100
101impl<
102    Pos,
103    Vel,
104    For,
105    Inf,
106    Cel,
107    Ind,
108    Vox,
109    Dom,
110    ConcVecExtracellular,
111    ConcBoundaryExtracellular,
112    ConcVecIntracellular,
113    Cont,
114    Obs,
115>
116    SimulationSupervisor<
117        MultiVoxelContainer<
118            Ind,
119            Pos,
120            Vel,
121            For,
122            Inf,
123            Vox,
124            Dom,
125            Cel,
126            ConcVecExtracellular,
127            ConcBoundaryExtracellular,
128            ConcVecIntracellular,
129        >,
130        Dom,
131        Cel,
132        Cont,
133        Obs,
134    >
135where
136    Dom: 'static + Serialize + for<'a> Deserialize<'a>,
137    Pos: 'static + Serialize + for<'a> Deserialize<'a>,
138    For: 'static + Serialize + for<'a> Deserialize<'a>,
139    Inf: 'static,
140    Vel: 'static + Serialize + for<'a> Deserialize<'a>,
141    ConcVecExtracellular: 'static + Serialize + for<'a> Deserialize<'a>,
142    ConcBoundaryExtracellular: 'static + Serialize + for<'a> Deserialize<'a>,
143    ConcVecIntracellular: 'static + Serialize + for<'a> Deserialize<'a>,
144    Cel: 'static + Serialize + for<'a> Deserialize<'a>,
145    Ind: 'static + Serialize + for<'a> Deserialize<'a>,
146    Vox: 'static + Serialize + for<'a> Deserialize<'a>,
147    Cont: 'static + Serialize + for<'a> Deserialize<'a> + Send + Sync,
148    Obs: 'static + Send + Sync,
149{
150    #[cfg_attr(feature = "tracing", instrument(skip_all))]
151    fn spawn_worker_threads_and_run_sim<ConcGradientExtracellular, ConcTotalExtracellular>(
152        &mut self,
153    ) -> Result<(), SimulationError>
154    where
155        Dom: cellular_raza_concepts::domain_old::Domain<Cel, Ind, Vox>,
156        Pos: PositionBound,
157        For: ForceBound,
158        Inf: InteractionInf,
159        Vel: VelocityBound,
160        ConcVecExtracellular: Concentration,
161        ConcTotalExtracellular: Concentration,
162        ConcVecIntracellular: Concentration,
163        ConcBoundaryExtracellular: Send + Sync,
164        ConcVecIntracellular: Mul<f64, Output = ConcVecIntracellular>
165            + Add<ConcVecIntracellular, Output = ConcVecIntracellular>
166            + AddAssign<ConcVecIntracellular>,
167        Ind: Index,
168        Vox: Voxel<Ind, Pos, Vel, For>,
169        Vox: ExtracellularMechanics<
170                Ind,
171                Pos,
172                ConcVecExtracellular,
173                ConcGradientExtracellular,
174                ConcTotalExtracellular,
175                ConcBoundaryExtracellular,
176            > + Volume,
177        Cel: Agent<Pos, Vel, For, Inf>
178            + CellularReactions<ConcVecIntracellular, ConcVecExtracellular>
179            + InteractionExtracellularGradient<Cel, ConcGradientExtracellular>
180            + Volume,
181        VoxelBox<
182            Ind,
183            Pos,
184            Vel,
185            For,
186            Vox,
187            Cel,
188            ConcVecExtracellular,
189            ConcBoundaryExtracellular,
190            ConcVecIntracellular,
191        >: Clone,
192        AuxiliaryCellPropertyStorage<Pos, Vel, For, ConcVecIntracellular>: Clone,
193        Cont: Controller<Cel, Obs>,
194    {
195        let mut handles = Vec::new();
196        let mut start_barrier = Barrier::new(self.multivoxelcontainers.len() + 1);
197        let controller_barrier = Barrier::new(self.multivoxelcontainers.len());
198
199        // Create progress bar and define style
200        if self.config.progressbar.is_some() {
201            println!("Running Simulation");
202        }
203
204        for (l, mut cont) in self.multivoxelcontainers.drain(..).enumerate() {
205            // Clone barriers to use them for synchronization in threads
206            let mut new_start_barrier = start_barrier.clone();
207            let mut controller_barrier_new = controller_barrier.clone();
208
209            // See if we need to save
210            let stop_now_new = Arc::new(AtomicBool::new(false));
211
212            // Copy time evaluation points
213            let t_start = self.time.t_start;
214            let t_eval = self.time.t_eval.clone();
215
216            // Add bar
217            let show_progressbar = self.config.progressbar.is_some();
218            let mut bar = construct_progress_bar(t_eval.len(), &self.config.progressbar)?;
219
220            let controller_box = self.controller_box.clone();
221
222            // Spawn a thread for each multivoxelcontainer that is running
223            let handle = thread::Builder::new()
224                .name(format!("worker_thread_{:03.0}", l))
225                .spawn(move || -> Result<_, SimulationError>
226            {
227                new_start_barrier.wait();
228
229                let mut time = t_start;
230                for (iteration, (t, save)) in (0u64..).zip(t_eval) {
231                    let dt = t - time;
232                    time = t;
233
234                    match cont.run_full_update(&dt) {
235                        Ok(()) => (),
236                        Err(error) => {
237                            // TODO this is not always an error in update_mechanics!
238                            println!("Encountered error in update_mechanics: {:?}. Stopping simulation.", error);
239                            // Make sure to stop all threads after this iteration.
240                            stop_now_new.store(true, Ordering::Relaxed);
241                        },
242                    }
243
244                    if show_progressbar && cont.mvc_id == 0 {
245                        bar.update(1)?;
246                    }
247
248                    // if save_now_new.load(Ordering::Relaxed) {
249                    if save {
250                        cont.save_voxels_to_database(&iteration)?;
251                        cont.save_cells_to_database(&iteration)?;
252                    }
253
254                    // TODO Make sure to only call this if the controller type is not ()
255                    {
256                        controller_box
257                            .lock()
258                            .unwrap()
259                            .measure(
260                                l as u32,
261                                cont.voxels
262                                    .iter()
263                                    .flat_map(|vox| vox.1.cells.iter().map(|(cbox, _)| cbox)),
264                            )
265                            .unwrap();
266                        controller_barrier_new.wait();
267                        controller_box
268                            .lock()
269                            .unwrap()
270                            .adjust(cont.voxels.iter_mut().flat_map(|vox| {
271                                vox.1.cells.iter_mut().map(|(cbox, aux_storage)| {
272                                    (cbox, &mut aux_storage.cycle_events)
273                                })
274                            }))
275                            .unwrap();
276                    }
277
278                    // Check if we are stopping the simulation now
279                    if stop_now_new.load(Ordering::Relaxed) {
280                        // new_barrier.wait();
281                        break;
282                    }
283                }
284                Ok(cont)
285            })?;
286            handles.push(handle);
287        }
288
289        // Store worker threads in supervisor
290        self.worker_threads = handles;
291
292        // This starts all threads simultaneously
293        start_barrier.wait();
294        Ok(())
295    }
296
297    /// Runs a full simulation and returns a [SimulationResult] after having completed
298    #[cfg_attr(feature = "tracing", instrument(skip_all))]
299    pub fn run_full_sim<ConcGradientExtracellular, ConcTotalExtracellular>(
300        &mut self,
301    ) -> Result<
302        SimulationResult<
303            Ind,
304            Pos,
305            For,
306            Vel,
307            ConcVecExtracellular,
308            ConcBoundaryExtracellular,
309            ConcVecIntracellular,
310            Vox,
311            Dom,
312            Cel,
313        >,
314        SimulationError,
315    >
316    where
317        Dom: cellular_raza_concepts::domain_old::Domain<Cel, Ind, Vox>,
318        Pos: PositionBound,
319        For: ForceBound,
320        Inf: InteractionInf,
321        Vel: VelocityBound,
322        ConcVecExtracellular: Concentration,
323        ConcTotalExtracellular: Concentration,
324        ConcBoundaryExtracellular: Send + Sync + 'static,
325        ConcVecIntracellular: Send + Sync + Concentration,
326        ConcVecIntracellular: Mul<f64, Output = ConcVecIntracellular>
327            + Add<ConcVecIntracellular, Output = ConcVecIntracellular>
328            + AddAssign<ConcVecIntracellular>,
329        Ind: Index,
330        Vox: Voxel<Ind, Pos, Vel, For>,
331        Vox: ExtracellularMechanics<
332                Ind,
333                Pos,
334                ConcVecExtracellular,
335                ConcGradientExtracellular,
336                ConcTotalExtracellular,
337                ConcBoundaryExtracellular,
338            > + Volume,
339        Cel: Agent<Pos, Vel, For, Inf>
340            + CellularReactions<ConcVecIntracellular, ConcVecExtracellular>
341            + InteractionExtracellularGradient<Cel, ConcGradientExtracellular>
342            + Volume,
343        VoxelBox<
344            Ind,
345            Pos,
346            Vel,
347            For,
348            Vox,
349            Cel,
350            ConcVecExtracellular,
351            ConcBoundaryExtracellular,
352            ConcVecIntracellular,
353        >: Clone,
354        AuxiliaryCellPropertyStorage<Pos, Vel, For, ConcVecIntracellular>: Clone,
355        Cont: Controller<Cel, Obs>,
356    {
357        // Run the simulation
358        self.spawn_worker_threads_and_run_sim()?;
359
360        // Collect all threads
361        let mut databases = Vec::new();
362        for thread in self.worker_threads.drain(..) {
363            let t = thread
364                .join()
365                .expect("Could not join threads after Simulation has finished")?;
366            databases.push((t.storage_cells, t.storage_voxels, t.domain))
367        }
368
369        // Create a simulationresult which can then be used to further plot and analyze results
370        let (storage_cells, storage_voxels, domain) = databases.pop().ok_or(RequestError(
371            format!("The threads of the simulation did not yield any handles"),
372        ))?;
373
374        let simulation_result = SimulationResult {
375            storage: self.storage.clone(),
376            domain,
377            storage_cells,
378            storage_voxels,
379            plotting_config: PlottingConfig::default(),
380        };
381
382        Ok(simulation_result)
383    }
384
385    // TODO do not clone these variables!
386    // Find a way to deserialize the struct with only referencing the variables
387    /// Saves the current [SimulationSetup] via the internal [StorageManager]
388    #[cfg_attr(feature = "tracing", instrument(skip_all))]
389    pub fn save_current_setup(&mut self, iteration: u64) -> Result<(), SimulationError>
390    where
391        Dom: Clone,
392        Cont: Clone,
393        Obs: Clone,
394        Cel: Clone,
395    {
396        use crate::storage::StorageInterfaceStore;
397        let setup_current = SimulationSetup {
398            domain: self.domain.clone(),
399            cells: Vec::<Cel>::new(),
400            time: TimeSetup {
401                t_start: 0.0,
402                t_eval: Vec::new(),
403            },
404            meta_params: self.meta_params.clone(),
405            storage: self.storage.clone(),
406            controller: self.controller_box.lock().unwrap().controller.clone(),
407        };
408
409        self.meta_infos
410            .store_single_element(iteration, &(), &setup_current)?;
411        Ok(())
412    }
413}
414
415use super::domain_decomposition::PlainIndex;
416
417/// Returned after finishing a full simulation
418pub struct SimulationResult<
419    Ind,
420    Pos,
421    For,
422    Vel,
423    ConcVecExtracellular,
424    ConcBoundaryExtracellular,
425    ConcVecIntracellular,
426    Vox,
427    Dom,
428    Cel,
429> where
430    Pos: Serialize + for<'a> Deserialize<'a>,
431    For: Serialize + for<'a> Deserialize<'a>,
432    Vel: Serialize + for<'a> Deserialize<'a>,
433    Cel: Serialize + for<'a> Deserialize<'a>,
434    VoxelBox<
435        Ind,
436        Pos,
437        Vel,
438        For,
439        Vox,
440        Cel,
441        ConcVecExtracellular,
442        ConcBoundaryExtracellular,
443        ConcVecIntracellular,
444    >: for<'a> Deserialize<'a> + Serialize,
445    Dom: Serialize + for<'a> Deserialize<'a>,
446    ConcVecExtracellular: Serialize + for<'a> Deserialize<'a> + 'static,
447    ConcBoundaryExtracellular: Serialize + for<'a> Deserialize<'a>,
448    ConcVecIntracellular: Serialize + for<'a> Deserialize<'a>,
449{
450    /// Configure how to store results of the simulation
451    pub storage: StorageBuilder<true>,
452
453    pub(crate) domain: DomainBox<Dom>,
454    /// [StorageManager] responsible for saving and loading cells
455    pub storage_cells: StorageManager<CellIdentifier, CellBox<Cel>>,
456    /// [StorageManager] responsible for saving and loading voxels
457    pub storage_voxels: StorageManager<
458        PlainIndex,
459        VoxelBox<
460            Ind,
461            Pos,
462            Vel,
463            For,
464            Vox,
465            Cel,
466            ConcVecExtracellular,
467            ConcBoundaryExtracellular,
468            ConcVecIntracellular,
469        >,
470    >,
471    /// Define properties of how to save results
472    pub plotting_config: PlottingConfig,
473}
474
475#[cfg_attr(feature = "tracing", instrument(skip_all))]
476fn construct_progress_bar(
477    n_iterations: usize,
478    title: &Option<String>,
479) -> Result<kdam::Bar, SimulationError> {
480    let mut style = kdam::BarBuilder::default()
481        .total(n_iterations as usize)
482        .bar_format("{desc}{percentage:3.0}%|{animation}| {count}/{total} [{elapsed}]");
483    if let Some(title) = title {
484        style = style.desc(title.clone());
485    }
486    Ok(style
487        .build()
488        .or_else(|string| Err(CalcError(format!("{string}"))))?)
489}
490
491impl<
492    Ind,
493    Pos,
494    For,
495    Vel,
496    ConcVecExtracellular,
497    ConcBoundaryExtracellular,
498    ConcVecIntracellular,
499    Vox,
500    Dom,
501    Cel,
502>
503    SimulationResult<
504        Ind,
505        Pos,
506        For,
507        Vel,
508        ConcVecExtracellular,
509        ConcBoundaryExtracellular,
510        ConcVecIntracellular,
511        Vox,
512        Dom,
513        Cel,
514    >
515where
516    Pos: Serialize + for<'a> Deserialize<'a>,
517    For: Serialize + for<'a> Deserialize<'a>,
518    Vel: Serialize + for<'a> Deserialize<'a>,
519    Cel: Serialize + for<'a> Deserialize<'a>,
520    Dom: Serialize + for<'a> Deserialize<'a>,
521    ConcVecExtracellular: Serialize + for<'a> Deserialize<'a> + 'static,
522    ConcBoundaryExtracellular: Serialize + for<'a> Deserialize<'a>,
523    ConcVecIntracellular: Serialize + for<'a> Deserialize<'a>,
524    VoxelBox<
525        Ind,
526        Pos,
527        Vel,
528        For,
529        Vox,
530        Cel,
531        ConcVecExtracellular,
532        ConcBoundaryExtracellular,
533        ConcVecIntracellular,
534    >: Clone + for<'a> Deserialize<'a> + Serialize,
535{
536    #[cfg_attr(feature = "tracing", instrument(skip_all))]
537    fn plot_spatial_at_iteration_with_functions<Cpf, Vpf, Dpf>(
538        &self,
539        iteration: u64,
540        cell_plotting_func: Cpf,
541        voxel_plotting_func: Vpf,
542        domain_plotting_func: Dpf,
543    ) -> Result<(), SimulationError>
544    where
545        Dpf: for<'a> Fn(
546            &Dom,
547            u32,
548            &'a String,
549        ) -> Result<
550            DrawingArea<BitMapBackend<'a>, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
551            DrawingError,
552        >,
553        Cpf: Fn(
554                &Cel,
555                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
556            ) -> Result<(), DrawingError>
557            + Send
558            + Sync,
559        Vpf: Fn(
560                &Vox,
561                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
562            ) -> Result<(), DrawingError>
563            + Send
564            + Sync,
565    {
566        use crate::storage::StorageInterfaceLoad;
567        // Obtain the voxels from the database
568        let voxel_boxes = self
569            .storage_voxels
570            .load_all_elements_at_iteration(iteration)?
571            .into_iter()
572            .map(|(_, value)| value)
573            .collect::<Vec<_>>();
574
575        // Choose the correct file path
576        let mut file_path = self.storage.get_full_path().clone();
577        file_path.push("images");
578        match std::fs::create_dir(&file_path) {
579            Ok(()) => (),
580            Err(_) => (),
581        }
582        file_path.push(format!("cells_at_iter_{:010.0}.png", iteration));
583        let filename = file_path.into_os_string().into_string().unwrap();
584
585        let mut chart = domain_plotting_func(
586            &self.domain.domain_raw,
587            self.plotting_config.image_size,
588            &filename,
589        )?;
590
591        voxel_boxes
592            .iter()
593            .map(|voxelbox| voxel_plotting_func(&voxelbox.voxel, &mut chart))
594            .collect::<Result<(), DrawingError>>()?;
595
596        voxel_boxes
597            .iter()
598            .map(|voxelbox| voxelbox.cells.iter())
599            .flatten()
600            .map(|(cellbox, _)| cell_plotting_func(&cellbox.cell, &mut chart))
601            .collect::<Result<(), DrawingError>>()?;
602
603        chart.present()?;
604        Ok(())
605    }
606
607    /// Plots a spatial image of the simulation result at given iteration.
608    #[cfg_attr(feature = "tracing", instrument(skip_all))]
609    pub fn plot_spatial_at_iteration(&self, iteration: u64) -> Result<(), SimulationError>
610    where
611        Dom: CreatePlottingRoot,
612        Cel: PlotSelf,
613        Vox: PlotSelf,
614    {
615        match self.plotting_config.image_type {
616            ImageType::BitMap => self.plot_spatial_at_iteration_with_functions(
617                iteration,
618                Cel::plot_self_bitmap,
619                Vox::plot_self_bitmap,
620                Dom::create_bitmap_root,
621            ),
622            // ImageType::Svg => self.plot_spatial_at_iteration_with_functions(iteration, C::plot_self::<BitMapBackend>, V::plot_self, D::create_svg_root),
623        }
624    }
625
626    /// Plots a spatial image of the simulation result at given iteration with custom functions.
627    #[cfg_attr(feature = "tracing", instrument(skip_all))]
628    pub fn plot_spatial_at_iteration_custom_functions<Cpf, Vpf, Dpf>(
629        &self,
630        iteration: u64,
631        cell_plotting_func: Cpf,
632        voxel_plotting_func: Vpf,
633        domain_plotting_func: Dpf,
634    ) -> Result<(), SimulationError>
635    where
636        Dpf: for<'a> Fn(
637            &Dom,
638            u32,
639            &'a String,
640        ) -> Result<
641            DrawingArea<BitMapBackend<'a>, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
642            DrawingError,
643        >,
644        Cpf: Fn(
645                &Cel,
646                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
647            ) -> Result<(), DrawingError>
648            + Send
649            + Sync,
650        Vpf: Fn(
651                &Vox,
652                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
653            ) -> Result<(), DrawingError>
654            + Send
655            + Sync,
656    {
657        self.plot_spatial_at_iteration_with_functions(
658            iteration,
659            cell_plotting_func,
660            voxel_plotting_func,
661            domain_plotting_func,
662        )
663    }
664
665    /// Plots a spatial image of the simulation result at given iteration
666    /// with custom functions for cells and voxels.
667    #[cfg_attr(feature = "tracing", instrument(skip_all))]
668    pub fn plot_spatial_at_iteration_custom_cell_voxel_functions<Cpf, Vpf>(
669        &self,
670        iteration: u64,
671        cell_plotting_func: Cpf,
672        voxel_plotting_func: Vpf,
673    ) -> Result<(), SimulationError>
674    where
675        Dom: CreatePlottingRoot,
676        Cpf: Fn(
677                &Cel,
678                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
679            ) -> Result<(), DrawingError>
680            + Send
681            + Sync,
682        Vpf: Fn(
683                &Vox,
684                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
685            ) -> Result<(), DrawingError>
686            + Send
687            + Sync,
688    {
689        self.plot_spatial_at_iteration_with_functions(
690            iteration,
691            cell_plotting_func,
692            voxel_plotting_func,
693            Dom::create_bitmap_root,
694        )
695    }
696
697    /// Plots a spatial image of the simulation result at given iteration
698    /// with custom functions for cells.
699    #[cfg_attr(feature = "tracing", instrument(skip_all))]
700    pub fn plot_spatial_at_iteration_custom_cell_function<Cpf>(
701        &self,
702        iteration: u64,
703        cell_plotting_func: Cpf,
704    ) -> Result<(), SimulationError>
705    where
706        Vox: PlotSelf,
707        Dom: CreatePlottingRoot,
708        Cpf: Fn(
709                &Cel,
710                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
711            ) -> Result<(), DrawingError>
712            + Send
713            + Sync,
714    {
715        match self.plotting_config.image_type {
716            ImageType::BitMap => self.plot_spatial_at_iteration_with_functions(
717                iteration,
718                cell_plotting_func,
719                Vox::plot_self_bitmap,
720                Dom::create_bitmap_root,
721            ),
722        }
723    }
724
725    #[cfg_attr(feature = "tracing", instrument(skip_all))]
726    fn build_thread_pool(&self) -> Result<rayon::ThreadPool, SimulationError> {
727        // Build a thread pool
728        let mut builder = rayon::ThreadPoolBuilder::new();
729        // Set the number of threads
730        builder = match self.plotting_config.n_threads {
731            Some(n) => builder.num_threads(n),
732            // If not threads were supplied, we use only one
733            _ => builder.num_threads(1),
734        };
735        Ok(builder.build()?)
736    }
737
738    #[cfg_attr(feature = "tracing", instrument(skip_all))]
739    fn build_progress_bar(&self, n_iterations: u64) -> Result<Option<kdam::Bar>, SimulationError> {
740        match self.plotting_config.show_progressbar {
741            true => Ok(Some(construct_progress_bar(n_iterations as usize, &None)?)),
742            false => Ok(None),
743        }
744    }
745
746    /// Plots a spatial image of the simulation result for all iterations with custom functions
747    #[cfg_attr(feature = "tracing", instrument(skip_all))]
748    pub fn plot_spatial_all_iterations_with_functions<Cpf, Vpf, Dpf>(
749        &self,
750        cell_plotting_func: &Cpf,
751        voxel_plotting_func: &Vpf,
752        domain_plotting_func: &Dpf,
753    ) -> Result<(), SimulationError>
754    where
755        Dpf: for<'a> Fn(
756                &Dom,
757                u32,
758                &'a String,
759            ) -> Result<
760                DrawingArea<BitMapBackend<'a>, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
761                DrawingError,
762            > + Send
763            + Sync,
764        Cpf: Fn(
765                &Cel,
766                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
767            ) -> Result<(), DrawingError>
768            + Send
769            + Sync,
770        Vpf: Fn(
771                &Vox,
772                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
773            ) -> Result<(), DrawingError>
774            + Send
775            + Sync,
776        CellBox<Cel>: Send + Sync,
777        VoxelBox<
778            Ind,
779            Pos,
780            Vel,
781            For,
782            Vox,
783            Cel,
784            ConcVecExtracellular,
785            ConcBoundaryExtracellular,
786            ConcVecIntracellular,
787        >: Send + Sync,
788        DomainBox<Dom>: Send + Sync,
789    {
790        use crate::storage::StorageInterfaceLoad;
791        let pool = self.build_thread_pool()?;
792
793        // Generate all images by calling the pool
794        pool.install(move || -> Result<(), SimulationError> {
795            use kdam::TqdmParallelIterator;
796            use rayon::prelude::*;
797            let all_iterations = self.storage_voxels.get_all_iterations()?;
798            let progress_bar = self.build_progress_bar(all_iterations.len() as u64)?;
799            match progress_bar {
800                Some(_) => println!("Generating Images"),
801                None => (),
802            }
803            let map_func = |iteration: u64| -> Result<(), SimulationError> {
804                self.plot_spatial_at_iteration_with_functions(
805                    iteration,
806                    &cell_plotting_func,
807                    &voxel_plotting_func,
808                    &domain_plotting_func,
809                )
810            };
811            match progress_bar {
812                Some(bar) => {
813                    all_iterations
814                        .into_par_iter()
815                        .tqdm_with_bar(bar)
816                        .map(move |iteration| map_func(iteration))
817                        .collect::<Result<(), SimulationError>>()?;
818                }
819                None => {
820                    all_iterations
821                        .into_par_iter()
822                        .map(|iteration| map_func(iteration))
823                        .collect::<Result<(), SimulationError>>()?;
824                }
825            }
826
827            Ok(())
828        })
829    }
830
831    /// Plots a spatial image of the simulation result for
832    /// all iterations with custom cell and voxel functions
833    #[cfg_attr(feature = "tracing", instrument(skip_all))]
834    pub fn plot_spatial_all_iterations_custom_cell_voxel_functions<Cpf, Vpf>(
835        &self,
836        cell_plotting_func: Cpf,
837        voxel_plotting_func: Vpf,
838    ) -> Result<(), SimulationError>
839    where
840        Dom: CreatePlottingRoot,
841        Cpf: Fn(
842                &Cel,
843                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
844            ) -> Result<(), DrawingError>
845            + Send
846            + Sync,
847        Vpf: Fn(
848                &Vox,
849                &mut DrawingArea<BitMapBackend, Cartesian2d<RangedCoordf64, RangedCoordf64>>,
850            ) -> Result<(), DrawingError>
851            + Send
852            + Sync,
853        CellBox<Cel>: Send + Sync,
854        VoxelBox<
855            Ind,
856            Pos,
857            Vel,
858            For,
859            Vox,
860            Cel,
861            ConcVecExtracellular,
862            ConcBoundaryExtracellular,
863            ConcVecIntracellular,
864        >: Send + Sync,
865        DomainBox<Dom>: Send + Sync,
866    {
867        self.plot_spatial_all_iterations_with_functions(
868            &cell_plotting_func,
869            &voxel_plotting_func,
870            &Dom::create_bitmap_root,
871        )
872    }
873}