atomecs/
integrator.rs

1//! Module that performs time-integration.
2//!
3//! This module implements the [EulerIntegrationSystem](struct.EulerIntegrationSystem.html),
4//! which uses the euler method to integrate classical equations of motion.
5
6extern crate nalgebra;
7
8use crate::atom::*;
9use crate::constant;
10use crate::initiate::NewlyCreated;
11use specs::prelude::*;
12
13/// Tracks the number of the current integration step.
14pub struct Step {
15    pub n: u64,
16}
17
18/// The timestep used for the integration.
19///
20/// The duration of the timestep should be sufficiently small to resolve the fastest timescale of motion,
21/// otherwise significant numerical errors will accumulate during the integration.
22/// For a typical magneto-optical trap simulation, the timestep should be around 1us.
23/// Decreasing the timestep further will not improve the accuracy, and will require more integration steps
24/// to simulate the same total simulation time.
25pub struct Timestep {
26    /// Duration of the simulation timestep, in SI units of seconds.
27    pub delta: f64,
28}
29
30/// # Euler Integration
31///
32/// The EulerIntegrationSystem integrates the classical equations of motion for particles using the euler method:
33/// `x' = x + v * dt`.
34/// This integrator is simple to implement but prone to integration error.
35///
36/// The timestep duration is specified by the [Timestep](struct.Timestep.html) system resource.
37pub struct EulerIntegrationSystem;
38
39impl<'a> System<'a> for EulerIntegrationSystem {
40    type SystemData = (
41        WriteStorage<'a, Position>,
42        WriteStorage<'a, Velocity>,
43        ReadExpect<'a, Timestep>,
44        WriteExpect<'a, Step>,
45        ReadStorage<'a, Force>,
46        ReadStorage<'a, Mass>,
47    );
48
49    fn run(&mut self, (mut pos, mut vel, t, mut step, force, mass): Self::SystemData) {
50        use rayon::prelude::*;
51
52        step.n += 1;
53        (&mut vel, &mut pos, &force, &mass).par_join().for_each(
54            |(vel, pos, force, mass)| {
55                euler_update(vel, pos, force, mass, t.delta);
56            },
57        );
58    }
59}
60
61pub const INTEGRATE_POSITION_SYSTEM_NAME: &str = "integrate_position";
62
63/// # Velocity-Verlet Integrate Position
64///
65/// Integrates position using a velocity-verlet integration approach.
66/// Stores the value of `Force` from the previous frame in the `OldForce` component.
67///
68/// The timestep duration is specified by the [Timestep](struct.Timestep.html) system resource.
69pub struct VelocityVerletIntegratePositionSystem;
70impl<'a> System<'a> for VelocityVerletIntegratePositionSystem {
71    type SystemData = (
72        WriteStorage<'a, Position>,
73        ReadStorage<'a, Velocity>,
74        ReadExpect<'a, Timestep>,
75        WriteExpect<'a, Step>,
76        ReadStorage<'a, Force>,
77        WriteStorage<'a, OldForce>,
78        ReadStorage<'a, Mass>,
79    );
80
81    fn run(&mut self, (mut pos, vel, t, mut step, force, mut old_force, mass): Self::SystemData) {
82        use rayon::prelude::*;
83
84        step.n += 1;
85        let dt = t.delta;
86
87        (&mut pos, &vel, &mut old_force, &force, &mass)
88            .par_join()
89            .for_each(|(mut pos, vel, mut old_force, force, mass)| {
90                pos.pos = pos.pos
91                    + vel.vel * dt
92                    + force.force / (constant::AMU * mass.value) / 2.0 * dt * dt;
93                old_force.0 = *force;
94            });
95    }
96}
97
98pub const INTEGRATE_VELOCITY_SYSTEM_NAME: &str = "integrate_velocity";
99
100/// # Velocity-Verlet Integrate Velocity
101///
102/// Integrates velocity using the velocity-verlet method, and the average of `Force` this frame and `OldForce` from the previous frame.
103///
104/// The timestep duration is specified by the [Timestep](struct.Timestep.html) system resource
105pub struct VelocityVerletIntegrateVelocitySystem;
106impl<'a> System<'a> for VelocityVerletIntegrateVelocitySystem {
107    type SystemData = (
108        WriteStorage<'a, Velocity>,
109        ReadExpect<'a, Timestep>,
110        ReadStorage<'a, Force>,
111        ReadStorage<'a, OldForce>,
112        ReadStorage<'a, Mass>,
113    );
114
115    fn run(&mut self, (mut vel, t, force, old_force, mass): Self::SystemData) {
116        use rayon::prelude::*;
117
118        let dt = t.delta;
119
120        (&mut vel, &force, &old_force, &mass).par_join().for_each(
121            |(vel, force, old_force, mass)| {
122                vel.vel += (force.force + old_force.0.force) / (constant::AMU * mass.value) / 2.0 * dt;
123            },
124        );
125    }
126}
127
128/// Adds [OldForce](OldForce.struct.html) components to newly created atoms.
129pub struct AddOldForceToNewAtomsSystem;
130impl<'a> System<'a> for AddOldForceToNewAtomsSystem {
131    type SystemData = (
132        Entities<'a>,
133        ReadStorage<'a, NewlyCreated>,
134        ReadStorage<'a, OldForce>,
135        Read<'a, LazyUpdate>,
136    );
137    fn run(&mut self, (ent, newly_created, old_force, updater): Self::SystemData) {
138        for (ent, _, _) in (&ent, &newly_created, !&old_force).join() {
139            updater.insert(ent, OldForce::default());
140        }
141    }
142}
143
144/// Stores the value of the force calculation from the previous frame.
145#[derive(Default)]
146pub struct OldForce(Force);
147impl Component for OldForce {
148    type Storage = VecStorage<OldForce>;
149}
150
151/// Performs the euler method to update [Velocity](struct.Velocity.html) and [Position](struct.Position.html) given an applied [Force](struct.Force.html).
152fn euler_update(vel: &mut Velocity, pos: &mut Position, force: &Force, mass: &Mass, dt: f64) {
153    pos.pos += vel.vel * dt;
154    vel.vel += force.force * dt / (constant::AMU * mass.value);
155}
156
157pub mod tests {
158    #[allow(unused_imports)]
159    use super::*;
160    extern crate specs;
161    #[allow(unused_imports)]
162    use specs::{Builder, DispatcherBuilder, World};
163
164    extern crate nalgebra;
165    #[allow(unused_imports)]
166    use nalgebra::Vector3;
167
168    #[test]
169    fn test_euler() {
170        let mut pos = Position {
171            pos: Vector3::new(1., 1., 1.),
172        };
173        let mut vel = Velocity {
174            vel: Vector3::new(0., 1., 0.),
175        };
176        let time = 1.;
177        let mass = Mass {
178            value: 1. / constant::AMU,
179        };
180        let force = Force {
181            force: Vector3::new(1., 1., 1.),
182        };
183        euler_update(&mut vel, &mut pos, &force, &mass, time);
184        assert_eq!(vel.vel, Vector3::new(1., 2., 1.));
185        assert_eq!(pos.pos, Vector3::new(1., 2., 1.));
186    }
187
188    /// Tests the [EulerIntegrationSystem] by creating a mock world and integrating the trajectory of one entity.
189    #[test]
190    fn test_euler_integration() {
191        let mut world = World::new();
192
193        let mut dispatcher = DispatcherBuilder::new()
194            .with(EulerIntegrationSystem, "integrator", &[])
195            .build();
196        dispatcher.setup(&mut world);
197
198        // create a particle with known force and mass
199        let force = Vector3::new(1.0, 0.0, 0.0);
200        let mass = 1.0;
201        let atom = world
202            .create_entity()
203            .with(Position {
204                pos: Vector3::new(0.0, 0.0, 0.0),
205            })
206            .with(Velocity {
207                vel: Vector3::new(0.0, 0.0, 0.0),
208            })
209            .with(Force { force })
210            .with(Mass {
211                value: mass / constant::AMU,
212            })
213            .build();
214
215        let dt = 1.0e-3;
216        world.insert(Timestep { delta: dt });
217        world.insert(Step { n: 0 });
218
219        // run simulation loop 1_000 times.
220        let n_steps = 1_000;
221        for _i in 0..n_steps {
222            dispatcher.dispatch(&world);
223            world.maintain();
224        }
225
226        let a = force / mass;
227        let expected_v = a * (n_steps as f64 * dt);
228
229        assert_approx_eq::assert_approx_eq!(
230            expected_v.norm(),
231            world
232                .read_storage::<Velocity>()
233                .get(atom)
234                .expect("atom not found.")
235                .vel
236                .norm(),
237            expected_v.norm() * 0.01
238        );
239
240        let expected_x = a * (n_steps as f64 * dt).powi(2) / 2.0;
241        assert_approx_eq::assert_approx_eq!(
242            expected_x.norm(),
243            world
244                .read_storage::<Position>()
245                .get(atom)
246                .expect("atom not found.")
247                .pos
248                .norm(),
249            expected_x.norm() * 0.01
250        );
251    }
252
253    #[test]
254    fn test_add_old_force_system() {
255        let mut test_world = World::new();
256
257        let mut dispatcher = DispatcherBuilder::new()
258            .with(AddOldForceToNewAtomsSystem, "", &[])
259            .build();
260        dispatcher.setup(&mut test_world);
261        test_world.register::<OldForce>();
262
263        let test_entity = test_world.create_entity().with(NewlyCreated {}).build();
264
265        dispatcher.dispatch(&test_world);
266        test_world.maintain();
267
268        let old_forces = test_world.read_storage::<OldForce>();
269        assert!(
270            old_forces.contains(test_entity),
271            "OldForce component not added to test entity."
272        );
273    }
274
275    #[test]
276    fn test_velocity_verlet_integration() {
277        let mut world = World::new();
278
279        let mut dispatcher = DispatcherBuilder::new()
280            .with(
281                VelocityVerletIntegratePositionSystem,
282                "integrate_position",
283                &[],
284            )
285            .with(
286                VelocityVerletIntegrateVelocitySystem,
287                "integrate_velocity",
288                &["integrate_position"],
289            )
290            .build();
291        dispatcher.setup(&mut world);
292
293        // create a particle with known force and mass
294        let force = Vector3::new(1.0, 0.0, 0.0);
295        let mass = 1.0;
296        let atom = world
297            .create_entity()
298            .with(Position {
299                pos: Vector3::new(0.0, 0.0, 0.0),
300            })
301            .with(Velocity {
302                vel: Vector3::new(0.0, 0.0, 0.0),
303            })
304            .with(Force { force })
305            .with(OldForce {
306                0: Force { force },
307            })
308            .with(Mass {
309                value: mass / constant::AMU,
310            })
311            .build();
312
313        let dt = 1.0e-3;
314        world.insert(Timestep { delta: dt });
315        world.insert(Step { n: 0 });
316
317        // run simulation loop 1_000 times.
318        let n_steps = 1_000;
319        for _i in 0..n_steps {
320            dispatcher.dispatch(&world);
321            world.maintain();
322        }
323
324        let a = force / mass;
325        let expected_v = a * (n_steps as f64 * dt);
326
327        assert_approx_eq::assert_approx_eq!(
328            expected_v.norm(),
329            world
330                .read_storage::<Velocity>()
331                .get(atom)
332                .expect("atom not found.")
333                .vel
334                .norm(),
335            expected_v.norm() * 0.01
336        );
337
338        let expected_x = a * (n_steps as f64 * dt).powi(2) / 2.0;
339        assert_approx_eq::assert_approx_eq!(
340            expected_x.norm(),
341            world
342                .read_storage::<Position>()
343                .get(atom)
344                .expect("atom not found.")
345                .pos
346                .norm(),
347            expected_x.norm() * 0.01
348        );
349    }
350}