Skip to main content

nyx_space/mc/
multivariate.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::{DispersedState, StateDispersion};
20use crate::cosmic::AstroPhysicsSnafu;
21use crate::errors::StateError;
22use crate::md::prelude::BPlane;
23use crate::md::{AstroSnafu, StateParameter};
24use crate::{pseudo_inverse, NyxError, Spacecraft, State};
25use anise::analysis::prelude::OrbitalElement;
26use anise::astro::orbit_gradient::OrbitGrad;
27use nalgebra::{DMatrix, DVector, SMatrix, SVector};
28use rand_distr::{Distribution, Normal};
29use snafu::ResultExt;
30use std::error::Error;
31
32/// A multivariate spacecraft state generator for Monte Carlo analyses. Ensures that the covariance is properly applied on all provided state variables.
33///
34/// # Algorithm
35///
36/// The `MvnSpacecraft` allows sampling from a multivariate normal distribution defined by a template state (mean) and a set of state dispersions (uncertainties).
37/// The core difficulty is that dispersions are often defined in non-Cartesian spaces (like Keplerian orbital elements or B-Plane parameters), while the spacecraft state is represented in Cartesian coordinates (Position and Velocity).
38///
39/// The algorithm proceeds as follows:
40/// 1.  **Jacobian Computation**: It computes the Jacobian matrix `J` representing the partial derivatives of the provided dispersion parameters with respect to the Cartesian state elements (x, y, z, vx, vy, vz).
41/// 2.  **Covariance Transformation**: It constructs a diagonal covariance matrix `P` in the parameter space (assuming input dispersions are independent in that space). It then transforms this into the Cartesian covariance matrix `C` using the linear mapping approximation:
42///     `C = J_inv * P * J_inv^T`
43///     where `J_inv` is the Moore-Penrose pseudo-inverse of `J`.
44/// 3.  **Decomposition**: It performs a Singular Value Decomposition (SVD) on `C` (or the full state covariance including mass, Cr, Cd) to obtain the square root of the covariance matrix, denoted as `L`.
45///     `C = U * S * V^T = (V * sqrt(S)) * (V * sqrt(S))^T` implies `L = V * sqrt(S)`
46/// 4.  **Sampling**: To generate a sample state `X`, it draws a vector `Z` of independent standard normal variables (N(0, 1)) and applies the transformation:
47///     `X = mu + L * Z`
48///
49/// # Correctness vs. Independent Sampling
50///
51/// One might ask: "Why not just sample each Cartesian coordinate independently using a normal distribution?"
52///
53/// 1.  **Correlations**: Independent sampling of Cartesian coordinates assumes a diagonal covariance matrix, implying no correlation between position and velocity components. In orbital mechanics, states are highly correlated (e.g., velocity magnitude and radial distance are coupled by energy). `MvnSpacecraft` preserves these physical correlations by mapping the physically meaningful uncertainties (e.g., in SMA or Inclination) into the Cartesian space.
54/// 2.  **Geometry**: Uncertainties defined in orbital elements form complex shapes (like "bananas") in Cartesian space. A multivariate normal approximation in Cartesian space captures the principal axes and orientation of this uncertainty volume, which an axis-aligned bounding box (implied by independent sampling) effectively destroys.
55/// 3.  **Consistency**: By using the Jacobian transformation, we ensure that the generated samples, when mapped back to the parameter space (linearized), reproduce the input statistics (mean and standard deviation) provided by the user.
56#[derive(Clone, Debug)]
57pub struct MvnSpacecraft {
58    /// The template state
59    pub template: Spacecraft,
60    pub dispersions: Vec<StateDispersion>,
61    /// The mean of the multivariate normal distribution
62    pub mean: SVector<f64, 9>,
63    /// The dot product \sqrt{\vec s} \cdot \vec v, where S is the singular values and V the V matrix from the SVD decomp of the covariance of multivariate normal distribution
64    pub sqrt_s_v: SMatrix<f64, 9, 9>,
65    /// The standard normal distribution used to seed the multivariate normal distribution
66    pub std_norm_distr: Normal<f64>,
67}
68
69impl MvnSpacecraft {
70    /// Creates a new mulivariate state generator from a mean and covariance on the set of state parameters.
71    /// The covariance must be positive semi definite.
72    ///
73    /// # Algorithm
74    /// This function will build the rotation matrix to rotate the requested dispersions into the Spacecraft state space using [OrbitDual].
75    /// If there are any dispersions on the Cr and Cd, then these are dispersed independently (because they are iid).
76    pub fn new(
77        template: Spacecraft,
78        dispersions: Vec<StateDispersion>,
79    ) -> Result<Self, Box<dyn Error>> {
80        let mut cov = SMatrix::<f64, 9, 9>::zeros();
81        let mut mean = SVector::<f64, 9>::zeros();
82
83        let orbit_dual = OrbitGrad::from(template.orbit);
84        let mut b_plane = None;
85        for obj in &dispersions {
86            if obj.param.is_b_plane() {
87                b_plane = Some(
88                    BPlane::from_dual(orbit_dual)
89                        .context(AstroSnafu)
90                        .map_err(Box::new)?,
91                );
92                break;
93            }
94        }
95
96        let num_orbital = dispersions
97            .iter()
98            .filter(|disp| disp.param.is_orbital())
99            .count();
100
101        if num_orbital > 0 {
102            // Build the rotation matrix from the orbital dispersions to the Cartesian state.
103            let mut jac = DMatrix::from_element(num_orbital, 6, 0.0);
104            let mut covar = DMatrix::from_element(num_orbital, num_orbital, 0.0);
105            let mut means = DVector::from_element(num_orbital, 0.0);
106            let orbit_dual = OrbitGrad::from(template.orbit);
107
108            for (rno, disp) in dispersions
109                .iter()
110                .filter(|d| d.param.is_orbital())
111                .enumerate()
112            {
113                let partial = if let StateParameter::Element(oe) = disp.param {
114                    orbit_dual.partial_for(oe).context(AstroPhysicsSnafu)?
115                } else if disp.param.is_b_plane() {
116                    match disp.param {
117                        StateParameter::BdotR => b_plane.unwrap().b_r_km,
118                        StateParameter::BdotT => b_plane.unwrap().b_t_km,
119                        StateParameter::BLTOF => b_plane.unwrap().ltof_s,
120                        _ => unreachable!(),
121                    }
122                } else {
123                    unreachable!()
124                };
125
126                println!("{partial}");
127
128                for (cno, val) in [
129                    partial.wrt_x(),
130                    partial.wrt_y(),
131                    partial.wrt_z(),
132                    partial.wrt_vx(),
133                    partial.wrt_vy(),
134                    partial.wrt_vz(),
135                ]
136                .iter()
137                .copied()
138                .enumerate()
139                {
140                    jac[(rno, cno)] = val;
141                }
142                covar[(rno, rno)] = disp.std_dev.unwrap_or(0.0).powi(2);
143                means[rno] = disp.mean.unwrap_or(0.0);
144            }
145
146            // Now that we have the Jacobian that rotates from the Cartesian elements to the dispersions parameters,
147            // let's compute the inverse of this Jacobian to rotate from the dispersion params into the Cartesian elements.
148            let jac_inv = pseudo_inverse!(&jac)?;
149
150            // Rotate the orbital covariance back into the Cartesian state space, making this a 6x6.
151            let orbit_cov = &jac_inv * &covar * jac_inv.transpose();
152            // Rotate the means into the Cartesian space
153            let cartesian_mean = jac_inv * means;
154            for ii in 0..6 {
155                for jj in 0..6 {
156                    cov[(ii, jj)] = orbit_cov[(ii, jj)];
157                }
158                mean[ii] = cartesian_mean[ii];
159            }
160        };
161
162        if dispersions.len() > num_orbital {
163            for disp in &dispersions {
164                if disp.param.is_orbital() {
165                    continue;
166                } else {
167                    match disp.param {
168                        StateParameter::Cr => {
169                            cov[(7, 7)] = disp.mean.unwrap_or(0.0).powi(2);
170                        }
171                        StateParameter::Cd => {
172                            cov[(8, 8)] = disp.mean.unwrap_or(0.0).powi(2);
173                        }
174                        StateParameter::DryMass | StateParameter::PropMass => {
175                            cov[(9, 9)] = disp.mean.unwrap_or(0.0).powi(2);
176                        }
177                        _ => return Err(Box::new(StateError::ReadOnly { param: disp.param })),
178                    }
179                }
180            }
181        }
182
183        // At this point, the cov matrix is a 9x9 with all dispersions transformed into the Cartesian state space.
184        let svd = cov.svd_unordered(false, true);
185        if svd.v_t.is_none() {
186            return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
187        }
188
189        let sqrt_s = svd.singular_values.map(|x| x.sqrt());
190        let mut sqrt_s_v_t = svd.v_t.unwrap().transpose();
191
192        for (i, mut col) in sqrt_s_v_t.column_iter_mut().enumerate() {
193            col *= sqrt_s[i];
194        }
195
196        Ok(Self {
197            template,
198            dispersions,
199            mean,
200            sqrt_s_v: sqrt_s_v_t,
201            std_norm_distr: Normal::new(0.0, 1.0).unwrap(),
202        })
203    }
204
205    /// Same as `new` but with a zero mean
206    pub fn zero_mean(
207        template: Spacecraft,
208        mut dispersions: Vec<StateDispersion>,
209    ) -> Result<Self, Box<dyn Error>> {
210        for disp in &mut dispersions {
211            disp.mean = Some(0.0);
212        }
213
214        Self::new(template, dispersions)
215    }
216
217    /// Initializes a new multivariate distribution using the state data in the spacecraft state space.
218    pub fn from_spacecraft_cov(
219        template: Spacecraft,
220        cov: SMatrix<f64, 9, 9>,
221        mean: SVector<f64, 9>,
222    ) -> Result<Self, Box<dyn Error>> {
223        // Check that covariance is PSD by ensuring that all the eigenvalues are positive or nil
224        match cov.eigenvalues() {
225            None => return Err(Box::new(NyxError::CovarianceMatrixNotPsd)),
226            Some(evals) => {
227                for eigenval in &evals {
228                    if *eigenval < 0.0 {
229                        return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
230                    }
231                }
232            }
233        };
234
235        let svd = cov.svd_unordered(false, true);
236        if svd.v_t.is_none() {
237            return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
238        }
239
240        let s = svd.singular_values;
241        // Item by item multiplication
242        let mut sqrt_s_v = svd.v_t.unwrap().transpose();
243        for (i, mut col) in sqrt_s_v.column_iter_mut().enumerate() {
244            col *= s[i].sqrt();
245        }
246
247        let dispersions = vec![
248            StateDispersion::builder()
249                .param(StateParameter::Element(OrbitalElement::X))
250                .std_dev(cov[(0, 0)])
251                .build(),
252            StateDispersion::builder()
253                .param(StateParameter::Element(OrbitalElement::Y))
254                .std_dev(cov[(1, 1)])
255                .build(),
256            StateDispersion::builder()
257                .param(StateParameter::Element(OrbitalElement::Z))
258                .std_dev(cov[(2, 2)])
259                .build(),
260            StateDispersion::builder()
261                .param(StateParameter::Element(OrbitalElement::VX))
262                .std_dev(cov[(3, 3)])
263                .build(),
264            StateDispersion::builder()
265                .param(StateParameter::Element(OrbitalElement::VY))
266                .std_dev(cov[(4, 4)])
267                .build(),
268            StateDispersion::builder()
269                .param(StateParameter::Element(OrbitalElement::VZ))
270                .std_dev(cov[(5, 5)])
271                .build(),
272            StateDispersion::builder()
273                .param(StateParameter::Cr)
274                .std_dev(cov[(6, 6)])
275                .build(),
276            StateDispersion::builder()
277                .param(StateParameter::Cd)
278                .std_dev(cov[(7, 7)])
279                .build(),
280            StateDispersion::builder()
281                .param(StateParameter::PropMass)
282                .std_dev(cov[(8, 8)])
283                .build(),
284        ];
285
286        Ok(Self {
287            template,
288            dispersions,
289            mean,
290            sqrt_s_v,
291            std_norm_distr: Normal::new(0.0, 1.0).unwrap(),
292        })
293    }
294}
295
296impl Distribution<DispersedState<Spacecraft>> for MvnSpacecraft {
297    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> DispersedState<Spacecraft> {
298        // Generate the vector representing the state
299        let x_rng = SVector::<f64, 9>::from_fn(|_, _| self.std_norm_distr.sample(rng));
300        let x = self.sqrt_s_v * x_rng + self.mean;
301        let mut state = self.template;
302
303        // Set the new state data
304        for (coord, val) in x.iter().copied().enumerate() {
305            if coord < 3 {
306                state.orbit.radius_km[coord] += val;
307            } else if coord < 6 {
308                state.orbit.velocity_km_s[coord % 3] += val;
309            } else if coord == 6 {
310                state.srp.coeff_reflectivity += val;
311            } else if coord == 7 {
312                state.drag.coeff_drag += val;
313            } else if coord == 8 {
314                state.mass.prop_mass_kg += val;
315            }
316        }
317
318        let mut actual_dispersions = Vec::new();
319        for disp in &self.dispersions {
320            // Compute the delta
321            let delta = self.template.value(disp.param).unwrap() - state.value(disp.param).unwrap();
322            actual_dispersions.push((disp.param, delta));
323        }
324
325        DispersedState {
326            state,
327            actual_dispersions,
328        }
329    }
330}
331
332#[cfg(test)]
333mod multivariate_ut {
334    use super::*;
335    use crate::time::Epoch;
336    use crate::Spacecraft;
337    use crate::GMAT_EARTH_GM;
338    use anise::constants::frames::EARTH_J2000;
339    use anise::prelude::Orbit;
340    use rand::RngExt;
341    use statrs;
342    use statrs::distribution::ContinuousCDF;
343
344    /// Returns the Mahalanobis distance of a random variable x from a multivariate normal distribution
345    /// with mean mu and covariance matrix S.
346    fn mahalanobis_distance<const T: usize>(
347        x: &SVector<f64, T>,
348        mu: &SVector<f64, T>,
349        cov_inv: &SMatrix<f64, T, T>,
350    ) -> f64 {
351        let delta = x - mu;
352        (delta.transpose() * cov_inv * delta)[(0, 0)]
353    }
354
355    #[test]
356    fn mahalanobis_distance_test() {
357        // Simple 2D case with diagonal covariance
358        let x = SVector::<f64, 2>::new(1.0, 2.0);
359        let mu = SVector::<f64, 2>::new(0.0, 0.0);
360        let cov = SMatrix::<f64, 2, 2>::from_diagonal(&SVector::<f64, 2>::new(2.0, 3.0));
361        let cov_inv = cov.pseudo_inverse(1e-12).unwrap();
362
363        // Expected distance: (1-0)^2/2 + (2-0)^2/3 = 1/2 + 4/3 = 0.5 + 1.333... = 1.8333...
364        let md = mahalanobis_distance(&x, &mu, &cov_inv);
365        assert!((md - 1.8333333333333333).abs() < 1e-9);
366    }
367
368    #[test]
369    fn test_mvn_generator() {
370        // Generate a random covariance matrix.
371        let mut rng = rand::rng();
372        let cov = SMatrix::<f64, 6, 6>::from_fn(|_, _| rng.random());
373        let cov = cov * cov.transpose();
374        let mut cov_resized = SMatrix::<f64, 9, 9>::zeros();
375        cov_resized.fixed_view_mut::<6, 6>(0, 0).copy_from(&cov);
376
377        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
378
379        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
380        let state = Spacecraft::builder()
381            .orbit(Orbit::keplerian(
382                8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k,
383            ))
384            .build();
385
386        let mvn = MvnSpacecraft::from_spacecraft_cov(state, cov_resized, SVector::zeros()).unwrap();
387
388        // Generate N samples from this distribution
389        let n = 1000;
390        let samples = mvn.sample_iter(&mut rng).take(n);
391
392        // Compute the Mahalanobis distance for each sample
393        let cov_inv = cov_resized.pseudo_inverse(1e-12).unwrap();
394        let md: Vec<f64> = samples
395            .map(|sample| {
396                let mut v = SVector::<f64, 9>::zeros();
397                for i in 0..6 {
398                    v[i] = sample.state.orbit.to_cartesian_pos_vel()[i]
399                        - state.orbit.to_cartesian_pos_vel()[i];
400                }
401                mahalanobis_distance(&v, &SVector::zeros(), &cov_inv)
402            })
403            .collect();
404
405        // Perform a chi-squared test on the Mahalanobis distances.
406        // The squared Mahalanobis distance of a sample from the mean of a multivariate normal distribution
407        // with k degrees of freedom follows a chi-squared distribution with k degrees of freedom.
408        // We will test if the 95th percentile of the Mahalanobis distances is within the 95th percentile
409        // of the chi-squared distribution.
410        let mut md = md;
411        md.sort_by(|a, b| a.partial_cmp(b).unwrap());
412        let p95_md = md[(0.95 * n as f64) as usize];
413
414        let chi_squared = statrs::distribution::ChiSquared::new(6.0).unwrap();
415        let p95_chi_squared = chi_squared.inverse_cdf(0.95);
416
417        assert!((p95_md - p95_chi_squared).abs() / p95_chi_squared < 0.2);
418    }
419
420    #[test]
421    fn disperse_r_mag() {
422        use anise::constants::frames::EARTH_J2000;
423        use anise::prelude::Orbit;
424
425        use crate::time::Epoch;
426        use rand_pcg::Pcg64Mcg;
427
428        // Ensure that this worked: a 3 sigma deviation around 1 km means we shouldn't have 99.7% of samples within those bounds.
429        // Create a reproducible fast seed
430        let seed = 0;
431
432        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
433
434        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
435        let state = Spacecraft::builder()
436            .orbit(Orbit::keplerian(
437                8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k,
438            ))
439            .build();
440
441        // Check that we can modify the radius magnitude
442        let std_dev = 1.0;
443        let generator = MvnSpacecraft::new(
444            state,
445            vec![StateDispersion::builder()
446                .param(StateParameter::Element(OrbitalElement::Rmag))
447                .std_dev(std_dev)
448                .build()],
449        )
450        .unwrap();
451
452        let rng = Pcg64Mcg::new(seed);
453        let init_rmag = state.orbit.rmag_km();
454        let cnt_too_far: u16 = generator
455            .sample_iter(rng)
456            .take(1000)
457            .map(|dispersed_state| {
458                if (init_rmag - dispersed_state.state.orbit.rmag_km()).abs() >= 3.0 * std_dev {
459                    1
460                } else {
461                    0
462                }
463            })
464            .sum::<u16>();
465
466        // We specified a seed so we know exactly what to expect and we've reset the seed to 0.
467        assert_eq!(
468            cnt_too_far,
469            6, // Mathematically, this should be 3!
470            "Should have about 3% of samples being more than 3 sigma away, got {cnt_too_far}"
471        );
472    }
473
474    #[test]
475    fn disperse_full_cartesian() {
476        use anise::constants::frames::EARTH_J2000;
477        use anise::prelude::Orbit;
478
479        use crate::time::Epoch;
480        use crate::Spacecraft;
481        use crate::GMAT_EARTH_GM;
482
483        use rand_pcg::Pcg64Mcg;
484
485        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
486
487        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
488        let state = Orbit::keplerian(8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k);
489
490        let std_dev = [10.0, 10.0, 10.0, 0.2, 0.2, 0.2, 0.0, 0.0, 0.0];
491
492        let generator = MvnSpacecraft::new(
493            Spacecraft {
494                orbit: state,
495                ..Default::default()
496            },
497            vec![
498                StateDispersion::builder()
499                    .param(StateParameter::Element(OrbitalElement::X))
500                    .std_dev(10.0)
501                    .build(),
502                StateDispersion::builder()
503                    .param(StateParameter::Element(OrbitalElement::Y))
504                    .std_dev(10.0)
505                    .build(),
506                StateDispersion::builder()
507                    .param(StateParameter::Element(OrbitalElement::Z))
508                    .std_dev(10.0)
509                    .build(),
510                StateDispersion::builder()
511                    .param(StateParameter::Element(OrbitalElement::VX))
512                    .std_dev(0.2)
513                    .build(),
514                StateDispersion::builder()
515                    .param(StateParameter::Element(OrbitalElement::VY))
516                    .std_dev(0.2)
517                    .build(),
518                StateDispersion::builder()
519                    .param(StateParameter::Element(OrbitalElement::VZ))
520                    .std_dev(0.2)
521                    .build(),
522            ],
523        )
524        .unwrap();
525
526        // Ensure that this worked: a 3 sigma deviation around 1 km means we shouldn't have 99.7% of samples within those bounds.
527        // Create a reproducible fast seed
528        let seed = 0;
529        let rng = Pcg64Mcg::new(seed);
530
531        let cnt_too_far: u16 = generator
532            .sample_iter(rng)
533            .take(1000)
534            .map(|dispersed_state| {
535                let mut cnt = 0;
536                for (idx, val_std_dev) in std_dev.iter().take(6).enumerate() {
537                    let cur_val = dispersed_state.state.to_vector()[idx];
538                    let nom_val = state.to_cartesian_pos_vel()[idx];
539                    if (cur_val - nom_val).abs() > *val_std_dev {
540                        cnt += 1;
541                    }
542                }
543                cnt
544            })
545            .sum::<u16>();
546
547        // We specified a seed so we know exactly what to expect
548        assert_eq!(
549            cnt_too_far / 6,
550            312,
551            "Should have about 3% of samples being more than 3 sigma away, got {cnt_too_far}"
552        );
553    }
554
555    #[test]
556    fn disperse_raan_only() {
557        use anise::constants::frames::EARTH_J2000;
558        use anise::prelude::Orbit;
559
560        use crate::time::Epoch;
561        use rand_pcg::Pcg64Mcg;
562
563        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
564
565        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
566        let state = Spacecraft::builder()
567            .orbit(Orbit::keplerian(
568                8_100.0, 1e-6, 12.85, 356.614, 14.19, 199.887_7, dt, eme2k,
569            ))
570            .build();
571
572        let angle_sigma_deg = 0.2;
573
574        assert!(StateParameter::Element(OrbitalElement::RAAN).is_orbital());
575
576        let generator = MvnSpacecraft::new(
577            state,
578            vec![StateDispersion::zero_mean(
579                StateParameter::Element(OrbitalElement::RAAN),
580                angle_sigma_deg,
581            )],
582        )
583        .unwrap();
584
585        // Ensure that this worked: a 3 sigma deviation around 1 km means we shouldn't have 99.7% of samples within those bounds.
586        // Create a reproducible fast seed
587        let seed = 0;
588        let rng = Pcg64Mcg::new(seed);
589        let n = 1000;
590        let samples = generator.sample_iter(rng).take(n);
591
592        let cov = SMatrix::<f64, 1, 1>::from_diagonal(&SVector::<f64, 1>::from_vec(vec![
593            angle_sigma_deg.powi(2),
594        ]));
595        let cov_inv = cov.pseudo_inverse(1e-12).unwrap();
596
597        let md: Vec<f64> = samples
598            .map(|dispersed_state| {
599                // For all other orbital parameters, make sure that we have not changed things dramatically.
600                for param in [
601                    StateParameter::Element(OrbitalElement::SemiMajorAxis),
602                    StateParameter::Element(OrbitalElement::Inclination),
603                ] {
604                    let orig = state.value(param).unwrap();
605                    let new = dispersed_state.state.value(param).unwrap();
606                    let prct_change = 100.0 * (orig - new).abs() / orig;
607                    assert!(
608                        prct_change < 5.0,
609                        "{param} changed by {prct_change:.3} % ({orig:.3e} -> {new:.3e})"
610                    );
611                }
612
613                let delta =
614                    SVector::<f64, 1>::from_vec(vec![dispersed_state.actual_dispersions[0].1]);
615                mahalanobis_distance(&delta, &SVector::zeros(), &cov_inv)
616            })
617            .collect();
618
619        let mut md = md;
620        md.sort_by(|a, b| a.partial_cmp(b).unwrap());
621        let p95_md = md[(0.95 * n as f64) as usize];
622
623        let chi_squared = statrs::distribution::ChiSquared::new(1.0).unwrap();
624        let p95_chi_squared = chi_squared.inverse_cdf(0.95);
625
626        assert!((p95_md - p95_chi_squared).abs() / p95_chi_squared < 0.2);
627    }
628
629    #[test]
630    fn disperse_keplerian() {
631        use anise::constants::frames::EARTH_J2000;
632        use anise::prelude::Orbit;
633
634        use crate::time::Epoch;
635        use rand_pcg::Pcg64Mcg;
636
637        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
638
639        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
640        let state = Spacecraft::builder()
641            .orbit(Orbit::keplerian(
642                8_100.0, 1e-6, 12.85, 356.614, 14.19, 199.887_7, dt, eme2k,
643            ))
644            .build();
645
646        let sma_sigma_km = 10.0;
647        let inc_sigma_deg = 0.15;
648        let angle_sigma_deg = 0.02;
649
650        let generator = MvnSpacecraft::new(
651            state,
652            vec![
653                StateDispersion::zero_mean(
654                    StateParameter::Element(OrbitalElement::SemiMajorAxis),
655                    sma_sigma_km,
656                ),
657                StateDispersion::zero_mean(
658                    StateParameter::Element(OrbitalElement::Inclination),
659                    inc_sigma_deg,
660                ),
661                StateDispersion::zero_mean(
662                    StateParameter::Element(OrbitalElement::RAAN),
663                    angle_sigma_deg,
664                ),
665                StateDispersion::zero_mean(
666                    StateParameter::Element(OrbitalElement::AoP),
667                    angle_sigma_deg,
668                ),
669            ],
670        )
671        .unwrap();
672
673        // The generator computes the equivalent Cartesian covariance. Let's retrieve it.
674        let expected_cart_cov = &generator.sqrt_s_v * generator.sqrt_s_v.transpose();
675
676        // Create a reproducible fast seed, and generate the samples in the Cartesian space
677        let seed = 0;
678        let rng = Pcg64Mcg::new(seed);
679        let n = 2000; // Increase samples for better stats
680        let samples: Vec<SVector<f64, 6>> = generator
681            .sample_iter(rng)
682            .take(n)
683            .map(|s| s.state.orbit.to_cartesian_pos_vel())
684            .collect();
685
686        let nominal_cart = state.orbit.to_cartesian_pos_vel();
687
688        // Check that the mean of the Cartesian states is close to the nominal
689        let mut mean_cart = SVector::<f64, 6>::zeros();
690        for sample in &samples {
691            mean_cart += sample;
692        }
693        mean_cart /= n as f64;
694
695        let mean_diff = mean_cart - nominal_cart;
696        // We expect some deviation because of the non-linearities.
697        assert!(mean_diff.norm() < 1.0);
698
699        // Check that the covariance of the Cartesian states is close to the expected one
700        let mut sample_cov = SMatrix::<f64, 6, 6>::zeros();
701        for sample in &samples {
702            let disp_vec = sample - &mean_cart;
703            sample_cov += &disp_vec * disp_vec.transpose();
704        }
705        sample_cov /= (n - 1) as f64;
706
707        let expected_cov_6x6 = expected_cart_cov.fixed_view::<6, 6>(0, 0);
708
709        let diff = sample_cov - expected_cov_6x6;
710        assert!(diff.norm() / expected_cov_6x6.norm() < 0.2);
711    }
712}