atomecs/atom_sources/
gaussian.rs

1//! Atom sources with gaussian velocity distributions.
2
3use std::marker::PhantomData;
4
5use super::{WeightedProbabilityDistribution, species::AtomCreator};
6use crate::atom::*;
7use crate::atom_sources::emit::AtomNumberToEmit;
8use crate::constant::EXP;
9use crate::initiate::*;
10use nalgebra::Vector3;
11
12use rand;
13use rand::distributions::Distribution;
14use rand::Rng;
15
16use specs::{
17    Component, Entities, Entity, HashMapStorage, Join, LazyUpdate, Read, ReadStorage, System,
18    WriteStorage,
19};
20
21pub struct GaussianVelocityDistributionSourceDefinition<T> where T : AtomCreator {
22    pub mean: Vector3<f64>,
23    pub std: Vector3<f64>,
24    phantom: PhantomData<T>
25}
26impl<T> Component for GaussianVelocityDistributionSourceDefinition<T> where T : AtomCreator + 'static {
27    type Storage = HashMapStorage<Self>;
28}
29
30pub struct GaussianVelocityDistributionSource<T> where T : AtomCreator {
31    vx_distribution: WeightedProbabilityDistribution,
32    vy_distribution: WeightedProbabilityDistribution,
33    vz_distribution: WeightedProbabilityDistribution,
34    phantom: PhantomData<T>
35}
36impl<T> Component for GaussianVelocityDistributionSource<T> where T : AtomCreator + 'static {
37    type Storage = HashMapStorage<Self>;
38}
39impl<T> GaussianVelocityDistributionSource<T> where T : AtomCreator {
40    fn get_random_velocity<R: Rng + ?Sized>(&self, rng: &mut R) -> Vector3<f64> {
41        Vector3::new(
42            self.vx_distribution.sample(rng),
43            self.vy_distribution.sample(rng),
44            self.vz_distribution.sample(rng),
45        )
46    }
47}
48
49/// Creates and precalculates a [WeightedProbabilityDistribution](struct.WeightedProbabilityDistribution.html)
50/// which can be used to sample values of velocity, based on given mean/std.
51///
52/// # Arguments
53///
54/// `mean`: The mean velocity, in m/s
55///
56/// `std`: The std of velocity, in m/s
57pub fn create_gaussian_velocity_distribution(
58    mean: f64,
59    std: f64,
60) -> WeightedProbabilityDistribution {
61    // tuple list of (velocity, weight)
62    let mut velocities = Vec::<f64>::new();
63    let mut weights = Vec::<f64>::new();
64
65    // precalculate the discretized distribution.
66    let n = 1000;
67    for i in -n..n {
68        let v = (i as f64) / (n as f64) * 5.0 * std;
69        let weight = EXP.powf(-(v / std).powf(2.0) / 2.0);
70        velocities.push(v + mean);
71        weights.push(weight);
72    }
73
74    WeightedProbabilityDistribution::new(velocities, weights)
75}
76
77/// Precalculates the probability distributions for
78/// [GaussianVelocityDistributionSourceDefinition](struct.GaussianVelocityDistributionSourceDefinition.html) and
79/// stores the result in a [GaussianVelocityDistributionSource](struct.GaussianVelocityDistributionSource.html) component.
80#[derive(Default)]
81pub struct PrecalculateForGaussianSourceSystem<T>(PhantomData<T>);
82impl<'a, T> System<'a> for PrecalculateForGaussianSourceSystem<T> where T : AtomCreator + 'static {
83    type SystemData = (
84        Entities<'a>,
85        ReadStorage<'a, GaussianVelocityDistributionSourceDefinition<T>>,
86        WriteStorage<'a, GaussianVelocityDistributionSource<T>>,
87    );
88
89    fn run(&mut self, (entities, definitions, mut calculated): Self::SystemData) {
90        let mut precalculated_data = Vec::<(Entity, GaussianVelocityDistributionSource<T>)>::new();
91        for (entity, definition, _) in (&entities, &definitions, !&calculated).join() {
92            let source = GaussianVelocityDistributionSource {
93                vx_distribution: create_gaussian_velocity_distribution(
94                    definition.mean[0],
95                    definition.std[0],
96                ),
97                vy_distribution: create_gaussian_velocity_distribution(
98                    definition.mean[1],
99                    definition.std[1],
100                ),
101                vz_distribution: create_gaussian_velocity_distribution(
102                    definition.mean[2],
103                    definition.std[2],
104                ),
105                phantom: PhantomData
106            };
107            precalculated_data.push((entity, source));
108            println!("Precalculated velocity and mass distributions for a gaussian source.");
109        }
110
111        for (entity, precalculated) in precalculated_data {
112            calculated
113                .insert(entity, precalculated)
114                .expect("Could not add precalculated gaussian source.");
115        }
116    }
117}
118
119#[derive(Default)]
120pub struct GaussianCreateAtomsSystem<T>(PhantomData<T>);
121impl<'a, T> System<'a> for GaussianCreateAtomsSystem<T> where T : AtomCreator + 'static {
122    type SystemData = (
123        Entities<'a>,
124        ReadStorage<'a, GaussianVelocityDistributionSource<T>>,
125        ReadStorage<'a, AtomNumberToEmit>,
126        ReadStorage<'a, Position>,
127        ReadStorage<'a, Mass>,
128        Read<'a, LazyUpdate>,
129    );
130
131    fn run(
132        &mut self,
133        (entities, sources, numbers_to_emits, positions, masses, updater): Self::SystemData,
134    ) {
135        let mut rng = rand::thread_rng();
136        for (source, number_to_emit, source_position, mass) in (
137            &sources,
138            &numbers_to_emits,
139            &positions,
140            &masses,
141        )
142            .join()
143        {
144            for _i in 0..number_to_emit.number {
145                let new_atom = entities.create();
146                let new_vel = source.get_random_velocity(&mut rng);
147                updater.insert(
148                    new_atom,
149                    Velocity {
150                        vel: new_vel,
151                    },
152                );
153                updater.insert(new_atom, source_position.clone());
154                updater.insert(new_atom, Force::new());
155                updater.insert(new_atom, mass.clone());
156                updater.insert(new_atom, Atom);
157                updater.insert(new_atom, InitialVelocity { vel: new_vel });
158                updater.insert(new_atom, NewlyCreated);
159                T::mutate(&updater, new_atom);
160            }
161        }
162    }
163}
atomecs/atom_sources/gaussian.rs

atomecs/atom_sources/
gaussian.rs
