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::() .add_plugin(TimeSeriesPlugin::::default()) .add_plugin(TimeSeriesPlugin::::default()) + .add_plugin(TimeSeriesPlugin::::default()) .insert_non_send_resource(Option::>::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::, + ); } init_optional_component::(sim); init_optional_component::(sim); @@ -201,7 +208,7 @@ impl Sweep { 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); + +#[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, @@ -80,3 +95,26 @@ where let total: T = comm.all_gather_sum(&total); ionized / total } + +pub(super) fn num_particles_at_timestep_levels_system( + mut solver: NonSendMut>>, + mut writer: EventWriter, + parameters: Res, +) { + 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(), + )); +}