1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

//! This module provides routines for the Monte Carlo simulation

use rand::prelude::*;
use rand_distr::{Distribution, Normal};
use rand_xoshiro::Xoshiro256PlusPlus;
use std::path::Path;

use crate::orbit::calc::models::Model;
use crate::{Orbit, F, I};

/// Calculate the distance using R and Z coordinates
fn dist(r: F, z: F) -> F {
    F::sqrt(r.powi(2) + z.powi(2))
}

/// Find and return the apocentric and pericentric distances
fn apo_peri(o: &Orbit) -> (F, F) {
    // Prepare aliases
    let (r_res, z_res) = (&o.results.r, &o.results.z);

    // Get the initial values
    let mut apo = dist(r_res[0], z_res[0]);
    let mut peri = apo;

    // For every other pair
    for j in 1..=o.n {
        // Calculate the distance
        let d = dist(r_res[j], z_res[j]);
        // Depending on the distance, update either
        // the apocentric or the pericentric distance
        if d > apo {
            apo = d;
        } else if d < peri {
            peri = d;
        }
    }

    (apo, peri)
}

impl Orbit {
    /// Apply the Monte Carlo method:
    /// simulate sampling from Normal distribution
    /// using the mean and standard deviation of the
    /// input data, use that as the input to the
    /// [`integrate`](super::integrate) routine `s`
    /// times with the number of iterations `n` and
    /// the time step `h`; write the coordinates and
    /// total energy in the `folder` if they're specified
    /// in the `fields`
    ///
    /// # Panics
    /// Can panic if standard deviations aren't finite.
    /// This should be handled in the [`load`](Orbit#method.load) routine.
    pub fn simulate(
        &mut self,
        model: &dyn Model,
        folder: &Path,
        fields: &[&str],
        s: I,
        n: I,
        h: F,
    ) {
        // Initialize the random generator
        let mut rng = Xoshiro256PlusPlus::seed_from_u64(1);

        // Initialize the Normal distributions
        let normal_x = Normal::new(self.hc_initials.x, self.hc_initials.x_err).unwrap();
        let normal_y = Normal::new(self.hc_initials.y, self.hc_initials.y_err).unwrap();
        let normal_z = Normal::new(self.hc_initials.z, self.hc_initials.z_err).unwrap();
        let normal_u = Normal::new(self.hc_initials.u, self.hc_initials.u_err).unwrap();
        let normal_v = Normal::new(self.hc_initials.v, self.hc_initials.v_err).unwrap();
        let normal_w = Normal::new(self.hc_initials.w, self.hc_initials.w_err).unwrap();

        // Prepare the results vectors
        self.results.apo = Vec::<F>::with_capacity(s + 1);
        self.results.peri = Vec::<F>::with_capacity(s + 1);

        // Define what fields will be sequentially saved
        let integrated_fields = fields
            .iter()
            .copied()
            .filter(|f| f == &"r" || f == &"z" || f == &"x" || f == &"y" || f == &"e")
            .collect::<Vec<&str>>();
        let distances_fields = fields
            .iter()
            .copied()
            .filter(|f| f == &"apo" || f == &"peri")
            .collect::<Vec<&str>>();

        // Integrate with initial values first
        self.integrate(model, n, h, false);

        // Find the apocentric and pericentric distances
        let (apo, peri) = apo_peri(self);

        // Save the distances
        self.results.apo.push(apo);
        self.results.peri.push(peri);

        // Write the coordinates and total energy
        self.write(folder, &integrated_fields, false);

        // For each iteration
        for _ in 0..s {
            // Simulate the input data
            self.hc_initials.x = normal_x.sample(&mut rng);
            self.hc_initials.y = normal_y.sample(&mut rng);
            self.hc_initials.z = normal_z.sample(&mut rng);
            self.hc_initials.u = normal_u.sample(&mut rng);
            self.hc_initials.v = normal_v.sample(&mut rng);
            self.hc_initials.w = normal_w.sample(&mut rng);

            // Integrate the orbit
            self.integrate(model, n, h, false);

            // Find the apocentric and pericentric distances
            let (apo, peri) = apo_peri(self);

            // Save the distances
            self.results.apo.push(apo);
            self.results.peri.push(peri);

            // Append the coordinates and total energy
            self.write(folder, &integrated_fields, true);
        }

        // Write the distances
        self.write(folder, &distances_fields, false);
    }
}