diff --git a/examples/benchmark.rs b/examples/benchmark.rs new file mode 100644 index 0000000..91a2783 --- /dev/null +++ b/examples/benchmark.rs @@ -0,0 +1,187 @@ +//! Simulation of atoms cooled to the Doppler limit. + +extern crate magneto_optical_trap as lib; +extern crate nalgebra; +use lib::atom::{Atom, AtomicTransition, Force, Mass, Position, Velocity}; +use lib::ecs; +use lib::initiate::NewlyCreated; +use lib::integrator::Timestep; +use lib::laser::cooling::CoolingLight; +use lib::laser::force::ApplyEmissionForceOption; +use lib::laser::gaussian::GaussianBeam; +use lib::laser::photons_scattered::EnableScatteringFluctuations; +use lib::magnetic::quadrupole::QuadrupoleField3D; +use lib::output::file; +use lib::output::file::Text; +use nalgebra::Vector3; +use rand::distributions::{Distribution, Normal}; +use specs::{Builder, World}; +use std::time::Instant; + +fn main() { + let now = Instant::now(); + + // Create the simulation world and builder for the ECS dispatcher. + let mut world = World::new(); + ecs::register_components(&mut world); + ecs::register_resources(&mut world); + let mut builder = ecs::create_simulation_dispatcher_builder(); + + // Configure thread pool. + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(12) + .build() + .unwrap(); + + builder.add_pool(::std::sync::Arc::new(pool)); + + let mut dispatcher = builder.build(); + dispatcher.setup(&mut world.res); + + // Create magnetic field. + world + .create_entity() + .with(QuadrupoleField3D::gauss_per_cm(18.2, Vector3::z())) + .with(Position { + pos: Vector3::new(0.0, 0.0, 0.0), + }) + .build(); + + // Create cooling lasers. + let detuning = -3.0; + let power = 0.02; + let radius = 66.7e-3 / (2.0_f64.sqrt()); + let beam_centre = Vector3::new(0.0, 0.0, 0.0); + + world + .create_entity() + .with(GaussianBeam { + intersection: beam_centre.clone(), + e_radius: radius, + power: power, + direction: Vector3::new(0.0, 0.0, 1.0), + }) + .with(CoolingLight::for_species( + AtomicTransition::rubidium(), + detuning, + -1, + )) + .build(); + world + .create_entity() + .with(GaussianBeam { + intersection: beam_centre.clone(), + e_radius: radius, + power: power, + direction: Vector3::new(0.0, 0.0, -1.0), + }) + .with(CoolingLight::for_species( + AtomicTransition::rubidium(), + detuning, + -1, + )) + .build(); + world + .create_entity() + .with(GaussianBeam { + intersection: beam_centre.clone(), + e_radius: radius, + power: power, + direction: Vector3::new(-1.0, 0.0, 0.0), + }) + .with(CoolingLight::for_species( + AtomicTransition::rubidium(), + detuning, + 1, + )) + .build(); + world + .create_entity() + .with(GaussianBeam { + intersection: beam_centre.clone(), + e_radius: radius, + power: power, + direction: Vector3::new(1.0, 0.0, 0.0), + }) + .with(CoolingLight::for_species( + AtomicTransition::rubidium(), + detuning, + 1, + )) + .build(); + world + .create_entity() + .with(GaussianBeam { + intersection: beam_centre.clone(), + e_radius: radius, + power: power, + direction: Vector3::new(0.0, 1.0, 0.0), + }) + .with(CoolingLight::for_species( + AtomicTransition::rubidium(), + detuning, + 1, + )) + .build(); + world + .create_entity() + .with(GaussianBeam { + intersection: beam_centre.clone(), + e_radius: radius, + power: power, + direction: Vector3::new(0.0, -1.0, 0.0), + }) + .with(CoolingLight::for_species( + AtomicTransition::rubidium(), + detuning, + 1, + )) + .build(); + + // Define timestep + world.add_resource(Timestep { delta: 1.0e-6 }); + + let vel_dist = Normal::new(0.0, 0.22); + let pos_dist = Normal::new(0.0, 1.2e-4); + let mut rng = rand::thread_rng(); + + // Add atoms + for _ in 0..10000 { + world + .create_entity() + .with(Position { + pos: Vector3::new( + pos_dist.sample(&mut rng), + pos_dist.sample(&mut rng), + pos_dist.sample(&mut rng), + ), + }) + .with(Velocity { + vel: Vector3::new( + vel_dist.sample(&mut rng), + vel_dist.sample(&mut rng), + vel_dist.sample(&mut rng), + ), + }) + .with(Force::new()) + .with(Mass { value: 87.0 }) + .with(AtomicTransition::rubidium()) + .with(Atom) + .with(NewlyCreated) + .build(); + } + + // Enable fluctuation options + // * Allow photon numbers to fluctuate. + // * Allow random force from emission of photons. + world.add_resource(ApplyEmissionForceOption {}); + world.add_resource(EnableScatteringFluctuations {}); + + // Run the simulation for a number of steps. + for _i in 0..5000 { + dispatcher.dispatch(&mut world.res); + world.maintain(); + } + + println!("Simulation completed in {} ms.", now.elapsed().as_millis()); +} diff --git a/examples/doppler_limit.rs b/examples/doppler_limit.rs index 4256d0c..97a6aa7 100644 --- a/examples/doppler_limit.rs +++ b/examples/doppler_limit.rs @@ -27,14 +27,6 @@ fn main() { ecs::register_resources(&mut world); let mut builder = ecs::create_simulation_dispatcher_builder(); - // Configure thread pool. - let pool = rayon::ThreadPoolBuilder::new() - .num_threads(6) - .build() - .unwrap(); - - builder.add_pool(::std::sync::Arc::new(pool)); - // Configure simulation output. builder = builder.with( file::new::("vel.txt".to_string(), 10), @@ -153,7 +145,7 @@ fn main() { let mut rng = rand::thread_rng(); // Add atoms - for _ in 0..10000 { + for _ in 0..1000 { world .create_entity() .with(Position { diff --git a/src/atom.rs b/src/atom.rs index b238f27..47a1bc9 100644 --- a/src/atom.rs +++ b/src/atom.rs @@ -1,6 +1,6 @@ //! Common atom components and systems. -use crate::constant::{BOHRMAG, C}; +use crate::constant::{BOHRMAG, C, HBAR, PI}; use crate::output::file::BinaryConversion; use nalgebra::Vector3; use serde::{Deserialize, Serialize}; @@ -132,12 +132,21 @@ pub struct AtomicTransition { pub linewidth: f64, /// Saturation intensity, in units of W/m^2. pub saturation_intensity: f64, + /// Precalculate prefactor used in the determination of rate coefficients. + pub rate_prefactor: f64, } impl Component for AtomicTransition { type Storage = VecStorage; } impl AtomicTransition { + pub fn calculate(mut self) -> Self { + self.rate_prefactor = + 3. / 4. * C * C / HBAR * (self.frequency).powf(-3.) * self.linewidth * self.linewidth + / 2. / PI; + self + } + /// Creates an `AtomicTransition` component populated with parameters for Rubidium. /// The parameters are taken from Daniel Steck's Data sheet on Rubidium-87. pub fn rubidium() -> Self { @@ -148,7 +157,9 @@ impl AtomicTransition { frequency: C / 780.0e-9, linewidth: 6.065e6, // [Steck, Rubidium87] saturation_intensity: 16.69, // [Steck, Rubidium 87, D2 cycling transition] + rate_prefactor: 0.0, // set in calculate } + .calculate() } /// Creates an `AtomicTransition` component populated with parameters for Strontium. @@ -161,7 +172,9 @@ impl AtomicTransition { frequency: 650_759_219_088_937., linewidth: 32e6, // [Nosske2017] saturation_intensity: 430.0, // [Nosske2017, 43mW/cm^2] + rate_prefactor: 0.0, // set in calculate } + .calculate() } /// Creates an `AtomicTransition` component populated with parameters for red Strontium transition. @@ -174,7 +187,9 @@ impl AtomicTransition { frequency: 434_829_121_311_000., // NIST, doi:10.1063/1.344917 linewidth: 7_400., // [Schreck2013] saturation_intensity: 0.0295, // [SChreck2013, 3 µW/cm^2] + rate_prefactor: 0.0, // set in calculate } + .calculate() } /// Creates an `AtomicTransition` component populated with parameters for Erbium. @@ -186,7 +201,9 @@ impl AtomicTransition { frequency: 5.142e14, linewidth: 190e3, saturation_intensity: 0.13, + rate_prefactor: 0.0, // set in calculate } + .calculate() } /// Creates an `AtomicTransition` component populated with parameters for Erbium 401 . pub fn erbium_401() -> Self { @@ -197,7 +214,9 @@ impl AtomicTransition { frequency: 7.476e14, linewidth: 30e6, saturation_intensity: 56.0, + rate_prefactor: 0.0, // set in calculate } + .calculate() } pub fn gamma(&self) -> f64 { diff --git a/src/ecs.rs b/src/ecs.rs index c6b3721..7795f60 100644 --- a/src/ecs.rs +++ b/src/ecs.rs @@ -42,13 +42,13 @@ pub fn create_simulation_dispatcher_builder() -> DispatcherBuilder<'static, 'sta builder = builder.with(LargerEarlyTimestepOptimizationSystem, "opt", &[]); builder = builder.with(ClearForceSystem, "clear", &[]); builder = builder.with(DeflagNewAtomsSystem, "deflag", &[]); - builder.add_barrier(); + //builder.add_barrier(); builder = magnetic::add_systems_to_dispatch(builder, &[]); - builder.add_barrier(); + //builder.add_barrier(); builder = laser::add_systems_to_dispatch(builder, &[]); - builder.add_barrier(); + //builder.add_barrier(); builder = atom_sources::add_systems_to_dispatch(builder, &[]); - builder.add_barrier(); + //builder.add_barrier(); builder = builder.with(ApplyGravitationalForceSystem, "add_gravity", &["clear"]); builder = builder.with( EulerIntegrationSystem, diff --git a/src/laser/doppler.rs b/src/laser/doppler.rs index 2b4459a..8a74308 100644 --- a/src/laser/doppler.rs +++ b/src/laser/doppler.rs @@ -96,14 +96,17 @@ impl<'a> System<'a> for InitialiseDopplerShiftSamplersSystem { WriteStorage<'a, DopplerShiftSamplers>, ); fn run(&mut self, (cooling, cooling_index, mut samplers): Self::SystemData) { + use rayon::prelude::*; + use specs::ParJoin; + let mut content = Vec::new(); for (_, _) in (&cooling, &cooling_index).join() { content.push(DopplerShiftSampler::default()); } - for mut sampler in (&mut samplers).join() { - sampler.contents = content.clone(); - } + (&mut samplers).par_join().for_each(|mut sampler| { + sampler.contents = content.to_vec(); + }); } } diff --git a/src/laser/force.rs b/src/laser/force.rs index 36cb9d5..0741883 100644 --- a/src/laser/force.rs +++ b/src/laser/force.rs @@ -19,6 +19,8 @@ use crate::integrator::Timestep; use crate::laser::repump::*; +const LASER_CACHE_SIZE: usize = 16; + /// This sytem calculates the forces from absorbing photons from the CoolingLight entities. /// /// The system assumes that the `ActualPhotonsScatteredVector` for each atom @@ -49,18 +51,52 @@ impl<'a> System<'a> for CalculateAbsorptionForcesSystem { _dark, ): Self::SystemData, ) { - for (scattered, mut force, _dark) in (&actual_scattered_vector, &mut forces, !&_dark).join() - { - for (index, cooling, gaussian) in - (&cooling_index, &cooling_light, &gaussian_beam).join() - { - let new_force = scattered.contents[index.index].scattered as f64 * HBAR - / timestep.delta - * gaussian.direction.normalize() - * cooling.wavenumber(); - force.force = force.force + new_force; - } + use rayon::prelude::*; + use specs::ParJoin; + + // There are typically only a small number of lasers in a simulation. + // For a speedup, cache the required components into thread memory, + // so they can be distributed to parallel workers during the atom loop. + type CachedLaser = (CoolingLight, CoolingLightIndex, GaussianBeam); + let laser_cache: Vec = (&cooling_light, &cooling_index, &gaussian_beam) + .join() + .map(|(cooling, index, gaussian)| (cooling.clone(), index.clone(), gaussian.clone())) + .collect(); + + // Perform the iteration over atoms, `LASER_CACHE_SIZE` at a time. + for base_index in (0..laser_cache.len()).step_by(LASER_CACHE_SIZE) { + let max_index = laser_cache.len().min(base_index + LASER_CACHE_SIZE); + let slice = &laser_cache[base_index..max_index]; + let mut laser_array = vec![laser_cache[0]; LASER_CACHE_SIZE]; + laser_array[..max_index].copy_from_slice(slice); + let number_in_iteration = slice.len(); + + (&actual_scattered_vector, &mut forces, !&_dark) + .par_join() + .for_each(|(scattered, mut force, _)| { + for i in 0..number_in_iteration { + let (cooling, index, gaussian) = laser_array[i]; + let new_force = scattered.contents[index.index].scattered as f64 * HBAR + / timestep.delta + * gaussian.direction.normalize() + * cooling.wavenumber(); + force.force = force.force + new_force; + } + }) } + + // for (scattered, mut force, _dark) in (&actual_scattered_vector, &mut forces, !&_dark).join() + // { + // for (index, cooling, gaussian) in + // (&cooling_index, &cooling_light, &gaussian_beam).join() + // { + // let new_force = scattered.contents[index.index].scattered as f64 * HBAR + // / timestep.delta + // * gaussian.direction.normalize() + // * cooling.wavenumber(); + // force.force = force.force + new_force; + // } + // } } } diff --git a/src/laser/intensity.rs b/src/laser/intensity.rs index 9fbb9d5..4aeef66 100644 --- a/src/laser/intensity.rs +++ b/src/laser/intensity.rs @@ -15,7 +15,7 @@ use specs::{Component, Entities, Join, ReadStorage, System, VecStorage, WriteSto const LASER_CACHE_SIZE: usize = 16; /// Represents the laser intensity at the position of the atom with respect to a certain laser beam -#[derive(Clone)] +#[derive(Clone, Copy)] pub struct LaserIntensitySampler { /// Intensity in SI units of W/m^2 pub intensity: f64, @@ -50,14 +50,17 @@ impl<'a> System<'a> for InitialiseLaserIntensitySamplersSystem { WriteStorage<'a, LaserIntensitySamplers>, ); fn run(&mut self, (cooling, cooling_index, mut samplers): Self::SystemData) { + use rayon::prelude::*; + use specs::ParJoin; + let mut content = Vec::new(); for (_, _) in (&cooling, &cooling_index).join() { content.push(LaserIntensitySampler::default()); } - for mut sampler in (&mut samplers).join() { - sampler.contents = content.clone(); - } + (&mut samplers).par_join().for_each(|mut sampler| { + sampler.contents = content.to_vec(); + }); } } diff --git a/src/laser/photons_scattered.rs b/src/laser/photons_scattered.rs index 9f6ba08..92e51ce 100644 --- a/src/laser/photons_scattered.rs +++ b/src/laser/photons_scattered.rs @@ -108,18 +108,21 @@ impl<'a> System<'a> for InitialiseExpectedPhotonsScatteredVectorSystem { WriteStorage<'a, ExpectedPhotonsScatteredVector>, ); fn run(&mut self, (cooling, cooling_index, mut expected_photons): Self::SystemData) { + use rayon::prelude::*; + use specs::ParJoin; + let mut content = Vec::new(); for (_, _) in (&cooling, &cooling_index).join() { content.push(ExpectedPhotonsScattered::default()); } - for mut expected in (&mut expected_photons).join() { - expected.contents = content.clone(); - } + (&mut expected_photons).par_join().for_each(|mut expected| { + expected.contents = content.to_vec(); + }); } } -/// Calcutates the expected mean number of Photons scattered by each laser in one iteration step +/// Calculates the expected mean number of Photons scattered by each laser in one iteration step /// /// It is required that the `TotalPhotonsScattered` is already updated since this System divides /// them between the CoolingLight entities. @@ -135,24 +138,27 @@ impl<'a> System<'a> for CalculateExpectedPhotonsScatteredSystem { &mut self, (rate_coefficients, total_photons_scattered, mut expected_photons_vector): Self::SystemData, ) { - for (rates, total, expected) in ( + use rayon::prelude::*; + use specs::ParJoin; + + ( &rate_coefficients, &total_photons_scattered, &mut expected_photons_vector, ) - .join() - { - let mut sum_rates: f64 = 0.; + .par_join() + .for_each(|(rates, total, expected)| { + let mut sum_rates: f64 = 0.; - for count in 0..rates.contents.len() { - sum_rates = sum_rates + rates.contents[count].rate; - } + for count in 0..rates.contents.len() { + sum_rates = sum_rates + rates.contents[count].rate; + } - for index in 0..expected.contents.len() { - expected.contents[index].scattered = - rates.contents[index].rate / sum_rates * total.total; - } - } + for index in 0..expected.contents.len() { + expected.contents[index].scattered = + rates.contents[index].rate / sum_rates * total.total; + } + }); } } @@ -223,14 +229,17 @@ impl<'a> System<'a> for InitialiseActualPhotonsScatteredVectorSystem { WriteStorage<'a, ActualPhotonsScatteredVector>, ); fn run(&mut self, (cooling, cooling_index, mut actual_photons): Self::SystemData) { + use rayon::prelude::*; + use specs::ParJoin; + let mut content = Vec::new(); for (_, _) in (&cooling, &cooling_index).join() { content.push(ActualPhotonsScattered::default()); } - for mut actual in (&mut actual_photons).join() { - actual.contents = content.clone(); - } + (&mut actual_photons).par_join().for_each(|mut actual| { + actual.contents = content.to_vec(); + }); } } @@ -254,6 +263,9 @@ impl<'a> System<'a> for CalculateActualPhotonsScatteredSystem { &mut self, (fluctuations_option, expected_photons_vector, mut actual_photons_vector): Self::SystemData, ) { + use rayon::prelude::*; + use specs::ParJoin; + match fluctuations_option { None => { for (expected, actual) in @@ -265,23 +277,23 @@ impl<'a> System<'a> for CalculateActualPhotonsScatteredSystem { } } Some(_rand) => { - for (expected, actual) in - (&expected_photons_vector, &mut actual_photons_vector).join() - { - for index in 0..expected.contents.len() { - let poisson = Poisson::new(expected.contents[index].scattered); - let drawn_number = poisson.sample(&mut rand::thread_rng()); - - // I have no clue why it is necessary but it appears that for - // very small expected photon numbers, the poisson distribution - // returns u64::MAX which destroys the Simulation - actual.contents[index].scattered = if drawn_number == u64::MAX { - 0.0 - } else { - drawn_number as f64 - }; - } - } + (&expected_photons_vector, &mut actual_photons_vector) + .par_join() + .for_each(|(expected, actual)| { + for index in 0..expected.contents.len() { + let poisson = Poisson::new(expected.contents[index].scattered); + let drawn_number = poisson.sample(&mut rand::thread_rng()); + + // I have no clue why it is necessary but it appears that for + // very small expected photon numbers, the poisson distribution + // returns u64::MAX which destroys the Simulation + actual.contents[index].scattered = if drawn_number == u64::MAX { + 0.0 + } else { + drawn_number as f64 + }; + } + }); } } } diff --git a/src/laser/rate.rs b/src/laser/rate.rs index b5793b3..2cdd8b5 100644 --- a/src/laser/rate.rs +++ b/src/laser/rate.rs @@ -11,7 +11,7 @@ use crate::laser::sampler::LaserDetuningSamplers; use crate::magnetic::MagneticFieldSampler; use specs::{Component, Join, ReadStorage, System, VecStorage, WriteStorage}; -use crate::constant::{C, HBAR, PI}; +use crate::constant::PI; /// Represents the rate coefficient of the atom with respect to a specific CoolingLight entity #[derive(Clone)] @@ -94,50 +94,48 @@ impl<'a> System<'a> for CalculateRateCoefficientsSystem { mut rate_coefficients, ): Self::SystemData, ) { + use rayon::prelude::*; + use specs::ParJoin; + for (cooling, index, gaussian) in (&cooling_light, &cooling_index, &gaussian_beam).join() { - for (detunings, intensities, atominfo, bfield, rates) in ( + ( &laser_detunings, &laser_intensities, &atomic_transition, &magnetic_field_sampler, &mut rate_coefficients, ) - .join() - { - let beam_direction_vector = gaussian.direction.normalize(); - let costheta = if &bfield.field.norm_squared() < &(10.0 * f64::EPSILON) { - 0.0 - } else { - beam_direction_vector - .normalize() - .dot(&bfield.field.normalize()) - }; + .par_join() + .for_each(|(detunings, intensities, atominfo, bfield, rates)| { + let beam_direction_vector = gaussian.direction.normalize(); + let costheta = if &bfield.field.norm_squared() < &(10.0 * f64::EPSILON) { + 0.0 + } else { + beam_direction_vector + .normalize() + .dot(&bfield.field.normalize()) + }; - let prefactor = 3. / 4. * C.powf(2.) / HBAR - * (atominfo.frequency).powf(-3.) - * atominfo.linewidth - * intensities.contents[index.index].intensity - / 2. - / PI - * atominfo.linewidth; + let prefactor = + atominfo.rate_prefactor * intensities.contents[index.index].intensity; - let scatter1 = - 0.25 * (cooling.polarization as f64 * costheta + 1.).powf(2.) * prefactor - / (detunings.contents[index.index].detuning_sigma_plus.powf(2.) - + (PI * atominfo.linewidth).powf(2.)); + let scatter1 = + 0.25 * (cooling.polarization as f64 * costheta + 1.).powf(2.) * prefactor + / (detunings.contents[index.index].detuning_sigma_plus.powf(2.) + + (PI * atominfo.linewidth).powf(2.)); - let scatter2 = - 0.25 * (cooling.polarization as f64 * costheta - 1.).powf(2.) * prefactor - / (detunings.contents[index.index] - .detuning_sigma_minus - .powf(2.) - + (PI * atominfo.linewidth).powf(2.)); + let scatter2 = + 0.25 * (cooling.polarization as f64 * costheta - 1.).powf(2.) * prefactor + / (detunings.contents[index.index] + .detuning_sigma_minus + .powf(2.) + + (PI * atominfo.linewidth).powf(2.)); - let scatter3 = 0.5 * (1. - costheta.powf(2.)) * prefactor - / (detunings.contents[index.index].detuning_pi.powf(2.) - + (PI * atominfo.linewidth).powf(2.)); - rates.contents[index.index].rate = scatter1 + scatter2 + scatter3; - } + let scatter3 = 0.5 * (1. - costheta.powf(2.)) * prefactor + / (detunings.contents[index.index].detuning_pi.powf(2.) + + (PI * atominfo.linewidth).powf(2.)); + rates.contents[index.index].rate = scatter1 + scatter2 + scatter3; + }); } } } diff --git a/src/laser/sampler.rs b/src/laser/sampler.rs index 86a28b0..bf69319 100644 --- a/src/laser/sampler.rs +++ b/src/laser/sampler.rs @@ -62,7 +62,7 @@ impl<'a> System<'a> for InitialiseLaserDetuningSamplersSystem { } for mut sampler in (&mut samplers).join() { - sampler.contents = content.clone(); + sampler.contents = content.to_vec(); } } } diff --git a/src/maths.rs b/src/maths.rs index 0a9a078..fa66227 100644 --- a/src/maths.rs +++ b/src/maths.rs @@ -28,7 +28,7 @@ pub fn get_minimum_distance_line_point( pub fn gaussian_dis(std: f64, distance: f64) -> f64 { //checked - 1.0 / (2.0 * PI * std.powf(2.0)) * EXP.powf(-distance.powf(2.0) / 2.0 / (std).powf(2.0)) + 1.0 / (2.0 * PI * std * std) * EXP.powf(-distance * distance / 2.0 / (std * std)) } /// generate a uniform random direction