From 6dea383b87124f227f231f2f0c9f36950aae8f3c Mon Sep 17 00:00:00 2001
From: tehforsch <mail@tonipeter.de>
Date: Fri, 15 Sep 2023 14:35:54 +0200
Subject: [PATCH] add timeseries for number of particles at different levels

---
 src/sweep/mod.rs         | 11 +++++++++--
 src/sweep/time_series.rs | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/sweep/mod.rs b/src/sweep/mod.rs
index 2046350a..5435d6ed 100644
--- a/src/sweep/mod.rs
+++ b/src/sweep/mod.rs
@@ -39,8 +39,10 @@ pub use self::task::RateData;
 use self::task::Task;
 use self::time_series::hydrogen_ionization_mass_average_system;
 use self::time_series::hydrogen_ionization_volume_average_system;
+use self::time_series::num_particles_at_timestep_levels_system;
 use self::time_series::HydrogenIonizationMassAverage;
 use self::time_series::HydrogenIonizationVolumeAverage;
+use self::time_series::NumParticlesAtTimestepLevels;
 use self::timestep_level::TimestepLevel;
 use self::timestep_state::TimestepState;
 use crate::chemistry::hydrogen_only::HydrogenOnly;
@@ -114,6 +116,7 @@ impl RaxiomPlugin for SweepPlugin {
             .add_derived_component::<components::Temperature>()
             .add_plugin(TimeSeriesPlugin::<HydrogenIonizationMassAverage>::default())
             .add_plugin(TimeSeriesPlugin::<HydrogenIonizationVolumeAverage>::default())
+            .add_plugin(TimeSeriesPlugin::<NumParticlesAtTimestepLevels>::default())
             .insert_non_send_resource(Option::<Sweep<HydrogenOnly>>::None)
             .add_startup_system_to_stage(StartupStages::InitSweep, init_sweep_system)
             .add_system_to_stage(Stages::Sweep, run_sweep_system)
@@ -131,7 +134,11 @@ impl RaxiomPlugin for SweepPlugin {
                 hydrogen_ionization_volume_average_system
                     .before(hydrogen_ionization_mass_average_system),
             )
-            .add_system_to_stage(Stages::AfterSweep, hydrogen_ionization_mass_average_system);
+            .add_system_to_stage(Stages::AfterSweep, hydrogen_ionization_mass_average_system)
+            .add_system_to_stage(
+                Stages::AfterSweep,
+                num_particles_at_timestep_levels_system::<HydrogenOnly>,
+            );
         }
         init_optional_component::<HeatingRate>(sim);
         init_optional_component::<Timestep>(sim);
@@ -201,7 +208,7 @@ impl<C: Chemistry> Sweep<C> {
 
     fn count_cells_global(&mut self, level: TimestepLevel) -> usize {
         let local_count = self.cells.enumerate_active(level).count();
-        let mut count_communicator = MpiWorld::new();
+        let mut count_communicator = MpiWorld::new_custom_tag(91100);
         count_communicator.all_gather_sum(&CellCount(local_count))
     }
 
diff --git a/src/sweep/time_series.rs b/src/sweep/time_series.rs
index 91256049..d863e27d 100644
--- a/src/sweep/time_series.rs
+++ b/src/sweep/time_series.rs
@@ -11,12 +11,16 @@ use mpi::traits::Equivalence;
 use serde::Serialize;
 
 use super::grid::Cell;
+use super::Sweep;
+use super::SweepParameters;
+use crate::chemistry::Chemistry;
 use crate::communication::communicator::Communicator;
 use crate::components::Density;
 use crate::components::IonizedHydrogenFraction;
 use crate::prelude::Particles;
 use crate::units::Dimensionless;
 use crate::units::Mass;
+use crate::units::Time;
 use crate::units::Volume;
 
 #[derive(Component, Debug, Clone, Equivalence, Deref, DerefMut, From, Named, Serialize)]
@@ -29,6 +33,17 @@ pub struct HydrogenIonizationMassAverage(pub Dimensionless);
 #[repr(transparent)]
 pub struct HydrogenIonizationVolumeAverage(pub Dimensionless);
 
+#[derive(Serialize, Clone, Named)]
+#[name = "num_particles_at_timestep_levels"]
+pub struct NumParticlesAtTimestepLevels(Vec<NumAtLevel>);
+
+#[derive(Serialize, Clone)]
+struct NumAtLevel {
+    level: usize,
+    num: usize,
+    timestep: Time,
+}
+
 pub fn hydrogen_ionization_volume_average_system(
     query: Particles<(&Cell, &Density, &IonizedHydrogenFraction)>,
     mut writer: EventWriter<HydrogenIonizationVolumeAverage>,
@@ -80,3 +95,26 @@ where
     let total: T = comm.all_gather_sum(&total);
     ionized / total
 }
+
+pub(super) fn num_particles_at_timestep_levels_system<C: Chemistry>(
+    mut solver: NonSendMut<Option<Sweep<C>>>,
+    mut writer: EventWriter<NumParticlesAtTimestepLevels>,
+    parameters: Res<SweepParameters>,
+) {
+    let solver = (*solver).as_mut().unwrap();
+    let max_timestep = parameters.max_timestep;
+    writer.send(NumParticlesAtTimestepLevels(
+        solver
+            .timestep_state
+            .iter_all_levels()
+            .map(|level| {
+                let num = solver.count_cells_global(level);
+                NumAtLevel {
+                    level: level.0,
+                    num,
+                    timestep: level.to_timestep(max_timestep),
+                }
+            })
+            .collect(),
+    ));
+}