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

nyx_space/mc/
montecarlo.rs
