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