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::State;
21use crate::dynamics::Dynamics;
22use crate::linalg::DefaultAllocator;
23use crate::linalg::allocator::Allocator;
24use crate::mc::DispersedState;
25use crate::mc::results::{PropResult, Results, Run};
26use crate::md::trajectory::Interpolatable;
27use crate::propagators::Propagator;
28#[cfg(not(target_arch = "wasm32"))]
29use crate::time::Unit;
30use crate::time::{Duration, Epoch};
31use anise::almanac::Almanac;
32use anise::analysis::event::Event;
33use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle};
34use log::info;
35use rand::SeedableRng;
36use rand::rngs::SysRng;
37use rand_distr::Distribution;
38use rayon::prelude::ParallelIterator;
39use rayon::prelude::*;
40use std::fmt;
41use std::sync::Arc;
42use std::sync::mpsc::channel;
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        DefaultAllocator: Allocator<<D::StateType as State>::Size>
198            + Allocator<<D::StateType as State>::Size, <D::StateType as State>::Size>
199            + Allocator<<D::StateType as State>::VecLength>,
200        <DefaultAllocator as Allocator<<D::StateType as State>::VecLength>>::Buffer<f64>: Send,
201    {
202        self.resume_run_until_epoch(prop, almanac, 0, end_epoch, num_runs)
203    }
204
205    /// Resumes a Monte Carlo run by skipping the first `skip` items, generating states only after that, and propagate each independently until the specified epoch.
206    #[must_use = "Monte Carlo result must be used"]
207    #[allow(clippy::needless_lifetimes)]
208    pub fn resume_run_until_epoch<D>(
209        &self,
210        prop: Propagator<D>,
211        almanac: Arc<Almanac>,
212        skip: usize,
213        end_epoch: Epoch,
214        num_runs: usize,
215    ) -> Results<S, PropResult<S>>
216    where
217        D: Dynamics<StateType = S>,
218        DefaultAllocator: Allocator<<D::StateType as State>::Size>
219            + Allocator<<D::StateType as State>::Size, <D::StateType as State>::Size>
220            + Allocator<<D::StateType as State>::VecLength>,
221        <DefaultAllocator as Allocator<<D::StateType as State>::VecLength>>::Buffer<f64>: Send,
222    {
223        // Generate the initial states
224        let init_states = self.generate_states(skip, num_runs, self.seed);
225        // Setup the progress bar
226        let pb = self.progress_bar(num_runs);
227        // Setup the thread friendly communication
228        let (tx, rx) = channel();
229
230        // And propagate on the thread pool
231        #[cfg(not(target_arch = "wasm32"))]
232        let start = StdInstant::now();
233        init_states.par_iter().progress_with(pb).for_each_with(
234            (prop, tx),
235            |(arc_prop, tx), (index, dispersed_state)| {
236                let result = arc_prop
237                    .with(dispersed_state.state, almanac.clone())
238                    .quiet()
239                    .until_epoch_with_traj(end_epoch);
240
241                // Build a single run result
242                let run = Run {
243                    index: *index,
244                    dispersed_state: dispersed_state.clone(),
245                    result: result.map(|r| PropResult {
246                        state: r.0,
247                        traj: r.1,
248                    }),
249                };
250
251                tx.send(run).unwrap();
252            },
253        );
254
255        #[cfg(not(target_arch = "wasm32"))]
256        {
257            let clock_time = StdInstant::now() - start;
258            info!(
259                "Propagated {} states in {}",
260                num_runs,
261                clock_time.as_secs_f64() * Unit::Second
262            );
263        }
264
265        // Collect all of the results and sort them by run index
266        let mut runs = rx.iter().collect::<Vec<Run<S, PropResult<S>>>>();
267        runs.par_sort_by_key(|run| run.index);
268
269        Results {
270            runs,
271            scenario: self.scenario.clone(),
272        }
273    }
274
275    /// Set up the seed and generate the states. This is useful for checking the generated states before running a large scale Monte Carlo.
276    #[must_use = "Generated states for a Monte Carlo run must be used"]
277    pub fn generate_states(
278        &self,
279        skip: usize,
280        num_runs: usize,
281        seed: Option<u128>,
282    ) -> Vec<(usize, DispersedState<S>)> {
283        // Setup the RNG
284        let rng = match seed {
285            Some(seed) => Pcg64Mcg::new(seed),
286            None => Pcg64Mcg::try_from_rng(&mut SysRng).unwrap(),
287        };
288
289        // Generate the states, forcing the borrow as specified in the `sample_iter` docs.
290        (&self.random_state)
291            .sample_iter(rng)
292            .skip(skip)
293            .take(num_runs)
294            .enumerate()
295            .collect::<Vec<(usize, DispersedState<S>)>>()
296    }
297}
298
299impl<S: Interpolatable, Distr: Distribution<DispersedState<S>>> fmt::Display
300    for MonteCarlo<S, Distr>
301where
302    DefaultAllocator: Allocator<S::Size> + Allocator<S::Size, S::Size> + Allocator<S::VecLength>,
303{
304    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
305        write!(
306            f,
307            "{} - Nyx Monte Carlo - seed: {:?}",
308            self.scenario, self.seed
309        )
310    }
311}
312
313impl<S: Interpolatable, Distr: Distribution<DispersedState<S>>> fmt::LowerHex
314    for MonteCarlo<S, Distr>
315where
316    DefaultAllocator: Allocator<S::Size> + Allocator<S::Size, S::Size> + Allocator<S::VecLength>,
317{
318    /// Returns a filename friendly name
319    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
320        write!(
321            f,
322            "mc-data-{}-seed-{:?}",
323            self.scenario.replace(' ', "-"),
324            self.seed
325        )
326    }
327}
nyx_space/mc/montecarlo.rs

nyx_space/mc/
montecarlo.rs
