nyx_space/mc/
montecarlo.rs

1/*
2    Nyx, blazing fast astrodynamics
3    Copyright (C) 2018-onwards Christopher Rabotin <christopher.rabotin@gmail.com>
4
5    This program is free software: you can redistribute it and/or modify
6    it under the terms of the GNU Affero General Public License as published
7    by the Free Software Foundation, either version 3 of the License, or
8    (at your option) any later version.
9
10    This program is distributed in the hope that it will be useful,
11    but WITHOUT ANY WARRANTY; without even the implied warranty of
12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13    GNU Affero General Public License for more details.
14
15    You should have received a copy of the GNU Affero General Public License
16    along with this program.  If not, see <https://www.gnu.org/licenses/>.
17*/
18
19use super::Pcg64Mcg;
20use crate::dynamics::Dynamics;
21use crate::linalg::allocator::Allocator;
22use crate::linalg::DefaultAllocator;
23use crate::mc::results::{PropResult, Results, Run};
24use crate::mc::DispersedState;
25use crate::md::trajectory::Interpolatable;
26use crate::md::EventEvaluator;
27use crate::propagators::Propagator;
28#[cfg(not(target_arch = "wasm32"))]
29use crate::time::Unit;
30use crate::time::{Duration, Epoch};
31use crate::State;
32use anise::almanac::Almanac;
33use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle};
34use log::info;
35use rand::SeedableRng;
36use rand_distr::Distribution;
37use rayon::prelude::ParallelIterator;
38use rayon::prelude::*;
39use std::fmt;
40use std::sync::mpsc::channel;
41use std::sync::Arc;
42#[cfg(not(target_arch = "wasm32"))]
43use std::time::Instant as StdInstant;
44
45/// A Monte Carlo framework, automatically running on all threads via a thread pool. This framework is targeted toward analysis of time-continuous variables.
46/// One caveat of the design is that the trajectory is used for post processing, not each individual state. This may prevent some event switching from being shown in GNC simulations.
47pub struct MonteCarlo<S: Interpolatable, Distr: Distribution<DispersedState<S>>>
48where
49    DefaultAllocator: Allocator<S::Size> + Allocator<S::Size, S::Size> + Allocator<S::VecLength>,
50{
51    /// Seed of the [64bit PCG random number generator](https://www.pcg-random.org/index.html)
52    pub seed: Option<u128>,
53    /// Generator of states for the Monte Carlo run
54    pub random_state: Distr,
55    /// Name of this run, will be reflected in the progress bar and in the output structure
56    pub scenario: String,
57    pub nominal_state: S,
58}
59
60impl<S: Interpolatable, Distr: Distribution<DispersedState<S>>> MonteCarlo<S, Distr>
61where
62    DefaultAllocator: Allocator<S::Size> + Allocator<S::Size, S::Size> + Allocator<S::VecLength>,
63{
64    pub fn new(
65        nominal_state: S,
66        random_variable: Distr,
67        scenario: String,
68        seed: Option<u128>,
69    ) -> Self {
70        Self {
71            random_state: random_variable,
72            seed,
73            scenario,
74            nominal_state,
75        }
76    }
77    // Just the template for the progress bar
78    fn progress_bar(&self, num_runs: usize) -> ProgressBar {
79        let pb = ProgressBar::new(num_runs.try_into().unwrap());
80        pb.set_style(
81            ProgressStyle::default_bar()
82                .template("[{elapsed_precise}] {bar:100.cyan/blue} {pos:>7}/{len:7} {msg}")
83                .unwrap()
84                .progress_chars("##-"),
85        );
86        pb.set_message(format!("{self}"));
87        pb
88    }
89
90    /// Generate states and propagate each independently until a specific event is found `trigger` times.
91    #[allow(clippy::needless_lifetimes)]
92    pub fn run_until_nth_event<D, F>(
93        self,
94        prop: Propagator<D>,
95        almanac: Arc<Almanac>,
96        max_duration: Duration,
97        event: &F,
98        trigger: usize,
99        num_runs: usize,
100    ) -> Results<S, PropResult<S>>
101    where
102        D: Dynamics<StateType = S>,
103
104        F: EventEvaluator<S>,
105        DefaultAllocator: Allocator<<D::StateType as State>::Size>
106            + Allocator<<D::StateType as State>::Size, <D::StateType as State>::Size>
107            + Allocator<<D::StateType as State>::VecLength>,
108        <DefaultAllocator as Allocator<<D::StateType as State>::VecLength>>::Buffer<f64>: Send,
109    {
110        self.resume_run_until_nth_event(prop, almanac, 0, max_duration, event, trigger, num_runs)
111    }
112
113    /// Generate states and propagate each independently until a specific event is found `trigger` times.
114    #[must_use = "Monte Carlo result must be used"]
115    #[allow(clippy::needless_lifetimes)]
116    pub fn resume_run_until_nth_event<D, F>(
117        &self,
118        prop: Propagator<D>,
119        almanac: Arc<Almanac>,
120        skip: usize,
121        max_duration: Duration,
122        event: &F,
123        trigger: usize,
124        num_runs: usize,
125    ) -> Results<S, PropResult<S>>
126    where
127        D: Dynamics<StateType = S>,
128
129        F: EventEvaluator<S>,
130        DefaultAllocator: Allocator<<D::StateType as State>::Size>
131            + Allocator<<D::StateType as State>::Size, <D::StateType as State>::Size>
132            + Allocator<<D::StateType as State>::VecLength>,
133        <DefaultAllocator as Allocator<<D::StateType as State>::VecLength>>::Buffer<f64>: Send,
134    {
135        // Generate the initial states
136        let init_states = self.generate_states(skip, num_runs, self.seed);
137        // Setup the progress bar
138        let pb = self.progress_bar(num_runs);
139        // Setup the thread friendly communication
140        let (tx, rx) = channel();
141
142        // Generate all states (must be done separately because the rng is not thread safe)
143        #[cfg(not(target_arch = "wasm32"))]
144        let start = StdInstant::now();
145
146        init_states.par_iter().progress_with(pb).for_each_with(
147            (prop, tx),
148            |(prop, tx), (index, dispersed_state)| {
149                let result = prop
150                    .with(dispersed_state.state, almanac.clone())
151                    .until_nth_event(max_duration, event, trigger);
152
153                // Build a single run result
154                let run = Run {
155                    index: *index,
156                    dispersed_state: dispersed_state.clone(),
157                    result: result.map(|r| PropResult {
158                        state: r.0,
159                        traj: r.1,
160                    }),
161                };
162                tx.send(run).unwrap();
163            },
164        );
165
166        #[cfg(not(target_arch = "wasm32"))]
167        {
168            let clock_time = StdInstant::now() - start;
169            info!(
170                "Propagated {} states in {}",
171                num_runs,
172                clock_time.as_secs_f64() * Unit::Second
173            );
174        }
175
176        // Collect all of the results and sort them by run index
177        let mut runs = rx
178            .iter()
179            .collect::<Vec<Run<D::StateType, PropResult<D::StateType>>>>();
180        runs.par_sort_by_key(|run| run.index);
181
182        Results {
183            runs,
184            scenario: self.scenario.clone(),
185        }
186    }
187
188    /// Generate states and propagate each independently until a specific event is found `trigger` times.
189    #[must_use = "Monte Carlo result must be used"]
190    #[allow(clippy::needless_lifetimes)]
191    pub fn run_until_epoch<D>(
192        self,
193        prop: Propagator<D>,
194        almanac: Arc<Almanac>,
195        end_epoch: Epoch,
196        num_runs: usize,
197    ) -> Results<S, PropResult<S>>
198    where
199        D: Dynamics<StateType = S>,
200
201        DefaultAllocator: Allocator<<D::StateType as State>::Size>
202            + Allocator<<D::StateType as State>::Size, <D::StateType as State>::Size>
203            + Allocator<<D::StateType as State>::VecLength>,
204        <DefaultAllocator as Allocator<<D::StateType as State>::VecLength>>::Buffer<f64>: Send,
205    {
206        self.resume_run_until_epoch(prop, almanac, 0, end_epoch, num_runs)
207    }
208
209    /// Resumes a Monte Carlo run by skipping the first `skip` items, generating states only after that, and propagate each independently until the specified epoch.
210    #[must_use = "Monte Carlo result must be used"]
211    #[allow(clippy::needless_lifetimes)]
212    pub fn resume_run_until_epoch<D>(
213        &self,
214        prop: Propagator<D>,
215        almanac: Arc<Almanac>,
216        skip: usize,
217        end_epoch: Epoch,
218        num_runs: usize,
219    ) -> Results<S, PropResult<S>>
220    where
221        D: Dynamics<StateType = S>,
222
223        DefaultAllocator: Allocator<<D::StateType as State>::Size>
224            + Allocator<<D::StateType as State>::Size, <D::StateType as State>::Size>
225            + Allocator<<D::StateType as State>::VecLength>,
226        <DefaultAllocator as Allocator<<D::StateType as State>::VecLength>>::Buffer<f64>: Send,
227    {
228        // Generate the initial states
229        let init_states = self.generate_states(skip, num_runs, self.seed);
230        // Setup the progress bar
231        let pb = self.progress_bar(num_runs);
232        // Setup the thread friendly communication
233        let (tx, rx) = channel();
234
235        // And propagate on the thread pool
236        #[cfg(not(target_arch = "wasm32"))]
237        let start = StdInstant::now();
238        init_states.par_iter().progress_with(pb).for_each_with(
239            (prop, tx),
240            |(arc_prop, tx), (index, dispersed_state)| {
241                let result = arc_prop
242                    .with(dispersed_state.state, almanac.clone())
243                    .quiet()
244                    .until_epoch_with_traj(end_epoch);
245
246                // Build a single run result
247                let run = Run {
248                    index: *index,
249                    dispersed_state: dispersed_state.clone(),
250                    result: result.map(|r| PropResult {
251                        state: r.0,
252                        traj: r.1,
253                    }),
254                };
255
256                tx.send(run).unwrap();
257            },
258        );
259
260        #[cfg(not(target_arch = "wasm32"))]
261        {
262            let clock_time = StdInstant::now() - start;
263            info!(
264                "Propagated {} states in {}",
265                num_runs,
266                clock_time.as_secs_f64() * Unit::Second
267            );
268        }
269
270        // Collect all of the results and sort them by run index
271        let mut runs = rx.iter().collect::<Vec<Run<S, PropResult<S>>>>();
272        runs.par_sort_by_key(|run| run.index);
273
274        Results {
275            runs,
276            scenario: self.scenario.clone(),
277        }
278    }
279
280    /// Set up the seed and generate the states. This is useful for checking the generated states before running a large scale Monte Carlo.
281    #[must_use = "Generated states for a Monte Carlo run must be used"]
282    pub fn generate_states(
283        &self,
284        skip: usize,
285        num_runs: usize,
286        seed: Option<u128>,
287    ) -> Vec<(usize, DispersedState<S>)> {
288        // Setup the RNG
289        let rng = match seed {
290            Some(seed) => Pcg64Mcg::new(seed),
291            None => Pcg64Mcg::from_os_rng(),
292        };
293
294        // Generate the states, forcing the borrow as specified in the `sample_iter` docs.
295        (&self.random_state)
296            .sample_iter(rng)
297            .skip(skip)
298            .take(num_runs)
299            .enumerate()
300            .collect::<Vec<(usize, DispersedState<S>)>>()
301    }
302}
303
304impl<S: Interpolatable, Distr: Distribution<DispersedState<S>>> fmt::Display
305    for MonteCarlo<S, Distr>
306where
307    DefaultAllocator: Allocator<S::Size> + Allocator<S::Size, S::Size> + Allocator<S::VecLength>,
308{
309    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
310        write!(
311            f,
312            "{} - Nyx Monte Carlo - seed: {:?}",
313            self.scenario, self.seed
314        )
315    }
316}
317
318impl<S: Interpolatable, Distr: Distribution<DispersedState<S>>> fmt::LowerHex
319    for MonteCarlo<S, Distr>
320where
321    DefaultAllocator: Allocator<S::Size> + Allocator<S::Size, S::Size> + Allocator<S::VecLength>,
322{
323    /// Returns a filename friendly name
324    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
325        write!(
326            f,
327            "mc-data-{}-seed-{:?}",
328            self.scenario.replace(' ', "-"),
329            self.seed
330        )
331    }
332}
nyx_space/mc/montecarlo.rs

nyx_space/mc/
montecarlo.rs
