1use 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#[derive(Clone, Debug)]
34pub struct MvnSpacecraft {
35 pub template: Spacecraft,
37 pub dispersions: Vec<StateDispersion>,
38 pub mean: SVector<f64, 9>,
40 pub sqrt_s_v: SMatrix<f64, 9, 9>,
42 pub std_norm_distr: Normal<f64>,
44}
45
46impl MvnSpacecraft {
47 pub fn new(
54 template: Spacecraft,
55 dispersions: Vec<StateDispersion>,
56 ) -> Result<Self, Box<dyn Error>> {
57 let mut cov = SMatrix::<f64, 9, 9>::zeros();
58 let mut mean = SVector::<f64, 9>::zeros();
59
60 let orbit_dual = OrbitGrad::from(template.orbit);
61 let mut b_plane = None;
62 for obj in &dispersions {
63 if obj.param.is_b_plane() {
64 b_plane = Some(
65 BPlane::from_dual(orbit_dual)
66 .context(AstroSnafu)
67 .map_err(Box::new)?,
68 );
69 break;
70 }
71 }
72
73 let num_orbital = dispersions
74 .iter()
75 .filter(|disp| disp.param.is_orbital())
76 .count();
77
78 dbg!(num_orbital);
79
80 if num_orbital > 0 {
81 let mut jac = DMatrix::from_element(num_orbital, 6, 0.0);
83 let mut covar = DMatrix::from_element(num_orbital, num_orbital, 0.0);
84 let mut means = DVector::from_element(num_orbital, 0.0);
85 let orbit_dual = OrbitGrad::from(template.orbit);
86
87 for (rno, disp) in dispersions
88 .iter()
89 .filter(|d| d.param.is_orbital())
90 .enumerate()
91 {
92 let partial = if let StateParameter::Element(oe) = disp.param {
93 orbit_dual.partial_for(oe).context(AstroPhysicsSnafu)?
94 } else if disp.param.is_b_plane() {
95 match disp.param {
96 StateParameter::BdotR => b_plane.unwrap().b_r_km,
97 StateParameter::BdotT => b_plane.unwrap().b_t_km,
98 StateParameter::BLTOF => b_plane.unwrap().ltof_s,
99 _ => unreachable!(),
100 }
101 } else {
102 unreachable!()
103 };
104
105 println!("{partial}");
106
107 for (cno, val) in [
108 partial.wrt_x(),
109 partial.wrt_y(),
110 partial.wrt_z(),
111 partial.wrt_vx(),
112 partial.wrt_vy(),
113 partial.wrt_vz(),
114 ]
115 .iter()
116 .copied()
117 .enumerate()
118 {
119 jac[(rno, cno)] = val;
120 }
121 covar[(rno, rno)] = disp.std_dev.unwrap_or(0.0).powi(2);
122 means[rno] = disp.mean.unwrap_or(0.0);
123 }
124
125 let jac_inv = pseudo_inverse!(&jac)?;
128
129 let orbit_cov = &jac_inv * &covar * jac_inv.transpose();
131 let cartesian_mean = jac_inv * means;
133 for ii in 0..6 {
134 for jj in 0..6 {
135 cov[(ii, jj)] = orbit_cov[(ii, jj)];
136 }
137 mean[ii] = cartesian_mean[ii];
138 }
139 };
140
141 if dispersions.len() > num_orbital {
142 for disp in &dispersions {
143 if disp.param.is_orbital() {
144 continue;
145 } else {
146 match disp.param {
147 StateParameter::Cr => {
148 cov[(7, 7)] = disp.mean.unwrap_or(0.0).powi(2);
149 }
150 StateParameter::Cd => {
151 cov[(8, 8)] = disp.mean.unwrap_or(0.0).powi(2);
152 }
153 StateParameter::DryMass | StateParameter::PropMass => {
154 cov[(9, 9)] = disp.mean.unwrap_or(0.0).powi(2);
155 }
156 _ => return Err(Box::new(StateError::ReadOnly { param: disp.param })),
157 }
158 }
159 }
160 }
161
162 let svd = cov.svd_unordered(false, true);
164 if svd.v_t.is_none() {
165 return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
166 }
167
168 let sqrt_s = svd.singular_values.map(|x| x.sqrt());
169 let mut sqrt_s_v_t = svd.v_t.unwrap().transpose();
170
171 for (i, mut col) in sqrt_s_v_t.column_iter_mut().enumerate() {
172 col *= sqrt_s[i];
173 }
174
175 Ok(Self {
176 template,
177 dispersions,
178 mean,
179 sqrt_s_v: sqrt_s_v_t,
180 std_norm_distr: Normal::new(0.0, 1.0).unwrap(),
181 })
182 }
183
184 pub fn zero_mean(
186 template: Spacecraft,
187 mut dispersions: Vec<StateDispersion>,
188 ) -> Result<Self, Box<dyn Error>> {
189 for disp in &mut dispersions {
190 disp.mean = Some(0.0);
191 }
192
193 Self::new(template, dispersions)
194 }
195
196 pub fn from_spacecraft_cov(
198 template: Spacecraft,
199 cov: SMatrix<f64, 9, 9>,
200 mean: SVector<f64, 9>,
201 ) -> Result<Self, Box<dyn Error>> {
202 match cov.eigenvalues() {
204 None => return Err(Box::new(NyxError::CovarianceMatrixNotPsd)),
205 Some(evals) => {
206 for eigenval in &evals {
207 if *eigenval < 0.0 {
208 return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
209 }
210 }
211 }
212 };
213
214 let svd = cov.svd_unordered(false, true);
215 if svd.v_t.is_none() {
216 return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
217 }
218
219 let s = svd.singular_values;
220 let mut sqrt_s_v = svd.v_t.unwrap().transpose();
222 for (i, mut col) in sqrt_s_v.column_iter_mut().enumerate() {
223 col *= s[i].sqrt();
224 }
225
226 let dispersions = vec![
227 StateDispersion::builder()
228 .param(StateParameter::Element(OrbitalElement::X))
229 .std_dev(cov[(0, 0)])
230 .build(),
231 StateDispersion::builder()
232 .param(StateParameter::Element(OrbitalElement::Y))
233 .std_dev(cov[(1, 1)])
234 .build(),
235 StateDispersion::builder()
236 .param(StateParameter::Element(OrbitalElement::Z))
237 .std_dev(cov[(2, 2)])
238 .build(),
239 StateDispersion::builder()
240 .param(StateParameter::Element(OrbitalElement::VX))
241 .std_dev(cov[(3, 3)])
242 .build(),
243 StateDispersion::builder()
244 .param(StateParameter::Element(OrbitalElement::VY))
245 .std_dev(cov[(4, 4)])
246 .build(),
247 StateDispersion::builder()
248 .param(StateParameter::Element(OrbitalElement::VZ))
249 .std_dev(cov[(5, 5)])
250 .build(),
251 StateDispersion::builder()
252 .param(StateParameter::Cr)
253 .std_dev(cov[(6, 6)])
254 .build(),
255 StateDispersion::builder()
256 .param(StateParameter::Cd)
257 .std_dev(cov[(7, 7)])
258 .build(),
259 StateDispersion::builder()
260 .param(StateParameter::PropMass)
261 .std_dev(cov[(8, 8)])
262 .build(),
263 ];
264
265 Ok(Self {
266 template,
267 dispersions,
268 mean,
269 sqrt_s_v,
270 std_norm_distr: Normal::new(0.0, 1.0).unwrap(),
271 })
272 }
273}
274
275impl Distribution<DispersedState<Spacecraft>> for MvnSpacecraft {
276 fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> DispersedState<Spacecraft> {
277 let x_rng = SVector::<f64, 9>::from_fn(|_, _| self.std_norm_distr.sample(rng));
279 let x = self.sqrt_s_v * x_rng + self.mean;
280 let mut state = self.template;
281
282 for (coord, val) in x.iter().copied().enumerate() {
284 if coord < 3 {
285 state.orbit.radius_km[coord] += val;
286 } else if coord < 6 {
287 state.orbit.velocity_km_s[coord % 3] += val;
288 } else if coord == 6 {
289 state.srp.coeff_reflectivity += val;
290 } else if coord == 7 {
291 state.drag.coeff_drag += val;
292 } else if coord == 8 {
293 state.mass.prop_mass_kg += val;
294 }
295 }
296
297 let mut actual_dispersions = Vec::new();
298 for disp in &self.dispersions {
299 let delta = self.template.value(disp.param).unwrap() - state.value(disp.param).unwrap();
301 actual_dispersions.push((disp.param, delta));
302 }
303
304 DispersedState {
305 state,
306 actual_dispersions,
307 }
308 }
309}
310
311#[cfg(test)]
312mod multivariate_ut {
313 use super::*;
314 use crate::time::Epoch;
315 use crate::Spacecraft;
316 use crate::GMAT_EARTH_GM;
317 use anise::constants::frames::EARTH_J2000;
318 use anise::prelude::Orbit;
319 use rand::Rng;
320 use statrs;
321 use statrs::distribution::ContinuousCDF;
322
323 fn mahalanobis_distance<const T: usize>(
326 x: &SVector<f64, T>,
327 mu: &SVector<f64, T>,
328 cov_inv: &SMatrix<f64, T, T>,
329 ) -> f64 {
330 let delta = x - mu;
331 (delta.transpose() * cov_inv * delta)[(0, 0)]
332 }
333
334 #[test]
335 fn mahalanobis_distance_test() {
336 let x = SVector::<f64, 2>::new(1.0, 2.0);
338 let mu = SVector::<f64, 2>::new(0.0, 0.0);
339 let cov = SMatrix::<f64, 2, 2>::from_diagonal(&SVector::<f64, 2>::new(2.0, 3.0));
340 let cov_inv = cov.pseudo_inverse(1e-12).unwrap();
341
342 let md = mahalanobis_distance(&x, &mu, &cov_inv);
344 assert!((md - 1.8333333333333333).abs() < 1e-9);
345 }
346
347 #[test]
348 fn test_mvn_generator() {
349 let mut rng = rand::rng();
351 let cov = SMatrix::<f64, 6, 6>::from_fn(|_, _| rng.random());
352 let cov = cov * cov.transpose();
353 let mut cov_resized = SMatrix::<f64, 9, 9>::zeros();
354 cov_resized.fixed_view_mut::<6, 6>(0, 0).copy_from(&cov);
355
356 let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
357
358 let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
359 let state = Spacecraft::builder()
360 .orbit(Orbit::keplerian(
361 8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k,
362 ))
363 .build();
364
365 let mvn = MvnSpacecraft::from_spacecraft_cov(state, cov_resized, SVector::zeros()).unwrap();
366
367 let n = 1000;
369 let samples = mvn.sample_iter(&mut rng).take(n);
370
371 let cov_inv = cov_resized.pseudo_inverse(1e-12).unwrap();
373 let md: Vec<f64> = samples
374 .map(|sample| {
375 let mut v = SVector::<f64, 9>::zeros();
376 for i in 0..6 {
377 v[i] = sample.state.orbit.to_cartesian_pos_vel()[i]
378 - state.orbit.to_cartesian_pos_vel()[i];
379 }
380 mahalanobis_distance(&v, &SVector::zeros(), &cov_inv)
381 })
382 .collect();
383
384 let mut md = md;
390 md.sort_by(|a, b| a.partial_cmp(b).unwrap());
391 let p95_md = md[(0.95 * n as f64) as usize];
392
393 let chi_squared = statrs::distribution::ChiSquared::new(6.0).unwrap();
394 let p95_chi_squared = chi_squared.inverse_cdf(0.95);
395
396 assert!((p95_md - p95_chi_squared).abs() / p95_chi_squared < 0.2);
397 }
398
399 #[test]
400 fn disperse_r_mag() {
401 use anise::constants::frames::EARTH_J2000;
402 use anise::prelude::Orbit;
403
404 use crate::time::Epoch;
405 use rand_pcg::Pcg64Mcg;
406
407 let seed = 0;
410
411 let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
412
413 let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
414 let state = Spacecraft::builder()
415 .orbit(Orbit::keplerian(
416 8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k,
417 ))
418 .build();
419
420 let std_dev = 1.0;
422 let generator = MvnSpacecraft::new(
423 state,
424 vec![StateDispersion::builder()
425 .param(StateParameter::Element(OrbitalElement::Rmag))
426 .std_dev(std_dev)
427 .build()],
428 )
429 .unwrap();
430
431 let rng = Pcg64Mcg::new(seed);
432 let init_rmag = state.orbit.rmag_km();
433 let cnt_too_far: u16 = generator
434 .sample_iter(rng)
435 .take(1000)
436 .map(|dispersed_state| {
437 if (init_rmag - dispersed_state.state.orbit.rmag_km()).abs() >= 3.0 * std_dev {
438 1
439 } else {
440 0
441 }
442 })
443 .sum::<u16>();
444
445 assert_eq!(
447 cnt_too_far,
448 6, "Should have about 3% of samples being more than 3 sigma away, got {cnt_too_far}"
450 );
451 }
452
453 #[test]
454 fn disperse_full_cartesian() {
455 use anise::constants::frames::EARTH_J2000;
456 use anise::prelude::Orbit;
457
458 use crate::time::Epoch;
459 use crate::Spacecraft;
460 use crate::GMAT_EARTH_GM;
461
462 use rand_pcg::Pcg64Mcg;
463
464 let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
465
466 let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
467 let state = Orbit::keplerian(8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k);
468
469 let std_dev = [10.0, 10.0, 10.0, 0.2, 0.2, 0.2, 0.0, 0.0, 0.0];
470
471 let generator = MvnSpacecraft::new(
472 Spacecraft {
473 orbit: state,
474 ..Default::default()
475 },
476 vec![
477 StateDispersion::builder()
478 .param(StateParameter::Element(OrbitalElement::X))
479 .std_dev(10.0)
480 .build(),
481 StateDispersion::builder()
482 .param(StateParameter::Element(OrbitalElement::Y))
483 .std_dev(10.0)
484 .build(),
485 StateDispersion::builder()
486 .param(StateParameter::Element(OrbitalElement::Z))
487 .std_dev(10.0)
488 .build(),
489 StateDispersion::builder()
490 .param(StateParameter::Element(OrbitalElement::VX))
491 .std_dev(0.2)
492 .build(),
493 StateDispersion::builder()
494 .param(StateParameter::Element(OrbitalElement::VY))
495 .std_dev(0.2)
496 .build(),
497 StateDispersion::builder()
498 .param(StateParameter::Element(OrbitalElement::VZ))
499 .std_dev(0.2)
500 .build(),
501 ],
502 )
503 .unwrap();
504
505 let seed = 0;
508 let rng = Pcg64Mcg::new(seed);
509
510 let cnt_too_far: u16 = generator
511 .sample_iter(rng)
512 .take(1000)
513 .map(|dispersed_state| {
514 let mut cnt = 0;
515 for (idx, val_std_dev) in std_dev.iter().take(6).enumerate() {
516 let cur_val = dispersed_state.state.to_vector()[idx];
517 let nom_val = state.to_cartesian_pos_vel()[idx];
518 if (cur_val - nom_val).abs() > *val_std_dev {
519 cnt += 1;
520 }
521 }
522 cnt
523 })
524 .sum::<u16>();
525
526 assert_eq!(
528 cnt_too_far / 6,
529 312,
530 "Should have about 3% of samples being more than 3 sigma away, got {cnt_too_far}"
531 );
532 }
533
534 #[test]
535 fn disperse_raan_only() {
536 use anise::constants::frames::EARTH_J2000;
537 use anise::prelude::Orbit;
538
539 use crate::time::Epoch;
540 use rand_pcg::Pcg64Mcg;
541
542 let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
543
544 let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
545 let state = Spacecraft::builder()
546 .orbit(Orbit::keplerian(
547 8_100.0, 1e-6, 12.85, 356.614, 14.19, 199.887_7, dt, eme2k,
548 ))
549 .build();
550
551 let angle_sigma_deg = 0.2;
552
553 assert!(StateParameter::Element(OrbitalElement::RAAN).is_orbital());
554
555 let generator = MvnSpacecraft::new(
556 state,
557 vec![StateDispersion::zero_mean(
558 StateParameter::Element(OrbitalElement::RAAN),
559 angle_sigma_deg,
560 )],
561 )
562 .unwrap();
563
564 let seed = 0;
567 let rng = Pcg64Mcg::new(seed);
568 let n = 1000;
569 let samples = generator.sample_iter(rng).take(n);
570
571 let cov = SMatrix::<f64, 1, 1>::from_diagonal(&SVector::<f64, 1>::from_vec(vec![
572 angle_sigma_deg.powi(2),
573 ]));
574 let cov_inv = cov.pseudo_inverse(1e-12).unwrap();
575
576 let md: Vec<f64> = samples
577 .map(|dispersed_state| {
578 for param in [
580 StateParameter::Element(OrbitalElement::SemiMajorAxis),
581 StateParameter::Element(OrbitalElement::Inclination),
582 ] {
583 let orig = state.value(param).unwrap();
584 let new = dispersed_state.state.value(param).unwrap();
585 let prct_change = 100.0 * (orig - new).abs() / orig;
586 assert!(
587 prct_change < 5.0,
588 "{param} changed by {prct_change:.3} % ({orig:.3e} -> {new:.3e})"
589 );
590 }
591
592 let delta =
593 SVector::<f64, 1>::from_vec(vec![dispersed_state.actual_dispersions[0].1]);
594 mahalanobis_distance(&delta, &SVector::zeros(), &cov_inv)
595 })
596 .collect();
597
598 let mut md = md;
599 md.sort_by(|a, b| a.partial_cmp(b).unwrap());
600 let p95_md = md[(0.95 * n as f64) as usize];
601
602 let chi_squared = statrs::distribution::ChiSquared::new(1.0).unwrap();
603 let p95_chi_squared = chi_squared.inverse_cdf(0.95);
604
605 assert!((p95_md - p95_chi_squared).abs() / p95_chi_squared < 0.2);
606 }
607
608 #[test]
609 fn disperse_keplerian() {
610 use anise::constants::frames::EARTH_J2000;
611 use anise::prelude::Orbit;
612
613 use crate::time::Epoch;
614 use rand_pcg::Pcg64Mcg;
615
616 let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
617
618 let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
619 let state = Spacecraft::builder()
620 .orbit(Orbit::keplerian(
621 8_100.0, 1e-6, 12.85, 356.614, 14.19, 199.887_7, dt, eme2k,
622 ))
623 .build();
624
625 let sma_sigma_km = 10.0;
626 let inc_sigma_deg = 0.15;
627 let angle_sigma_deg = 0.02;
628
629 let generator = MvnSpacecraft::new(
630 state,
631 vec![
632 StateDispersion::zero_mean(
633 StateParameter::Element(OrbitalElement::SemiMajorAxis),
634 sma_sigma_km,
635 ),
636 StateDispersion::zero_mean(
637 StateParameter::Element(OrbitalElement::Inclination),
638 inc_sigma_deg,
639 ),
640 StateDispersion::zero_mean(
641 StateParameter::Element(OrbitalElement::RAAN),
642 angle_sigma_deg,
643 ),
644 StateDispersion::zero_mean(
645 StateParameter::Element(OrbitalElement::AoP),
646 angle_sigma_deg,
647 ),
648 ],
649 )
650 .unwrap();
651
652 let expected_cart_cov = &generator.sqrt_s_v * generator.sqrt_s_v.transpose();
654
655 let seed = 0;
657 let rng = Pcg64Mcg::new(seed);
658 let n = 2000; let samples: Vec<SVector<f64, 6>> = generator
660 .sample_iter(rng)
661 .take(n)
662 .map(|s| s.state.orbit.to_cartesian_pos_vel())
663 .collect();
664
665 let nominal_cart = state.orbit.to_cartesian_pos_vel();
666
667 let mut mean_cart = SVector::<f64, 6>::zeros();
669 for sample in &samples {
670 mean_cart += sample;
671 }
672 mean_cart /= n as f64;
673
674 let mean_diff = mean_cart - nominal_cart;
675 assert!(mean_diff.norm() < 1.0);
677
678 let mut sample_cov = SMatrix::<f64, 6, 6>::zeros();
680 for sample in &samples {
681 let disp_vec = sample - &mean_cart;
682 sample_cov += &disp_vec * disp_vec.transpose();
683 }
684 sample_cov /= (n - 1) as f64;
685
686 let expected_cov_6x6 = expected_cart_cov.fixed_view::<6, 6>(0, 0);
687
688 let diff = sample_cov - expected_cov_6x6;
689 assert!(diff.norm() / expected_cov_6x6.norm() < 0.2);
690 }
691}