use crate::outfit_errors::OutfitError;
use super::constants::DPI;
use super::constants::GAUSS_GRAV_SQUARED;
use super::orb_elem::eccentricity_control;
use nalgebra::Vector3;
use std::f64::consts::PI;
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum OrbitType {
Elliptic,
Hyperbolic,
Parabolic,
}
impl OrbitType {
pub fn from_alpha(alpha: f64) -> Self {
if alpha < 0.0 {
OrbitType::Elliptic
} else if alpha > 0.0 {
OrbitType::Hyperbolic
} else {
OrbitType::Parabolic
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct UniversalKeplerParams {
pub dt: f64,
pub r0: f64,
pub sig0: f64,
pub mu: f64,
pub alpha: f64,
pub e0: f64,
}
impl UniversalKeplerParams {
pub fn orbit_type(&self) -> OrbitType {
OrbitType::from_alpha(self.alpha)
}
}
#[inline(always)]
pub fn s_funct(psi: f64, alpha: f64) -> (f64, f64, f64, f64) {
const JMAX: usize = 70; const HALFMAX: usize = 30; const BETACONTR: f64 = 100.0;
let eps = f64::EPSILON;
let convergence_tol = 100.0 * eps; let overflow_limit = 1.0 / eps;
if psi == 0.0 {
return (1.0, 0.0, 0.0, 0.0);
}
let psi2 = psi * psi;
let beta = alpha * psi2;
if beta.abs() < BETACONTR {
let mut s2 = 0.5 * psi2;
let mut term2 = s2;
let mut s3 = (s2 * psi) / 3.0;
let mut term3 = s3;
let mut d2a = 3.0;
let mut d2b = 4.0;
let mut d3a = 4.0;
let mut d3b = 5.0;
for _ in 1..=JMAX {
term2 *= beta / (d2a * d2b);
s2 += term2;
term3 *= beta / (d3a * d3b);
s3 += term3;
let small2 = term2.abs() < convergence_tol;
let small3 = term3.abs() < convergence_tol;
let big2 = term2.abs() > overflow_limit;
let big3 = term3.abs() > overflow_limit;
if (small2 && small3) || big2 || big3 {
break;
}
d2a += 2.0;
d2b += 2.0;
d3a += 2.0;
d3b += 2.0;
}
let s1 = psi + alpha * s3;
let s0 = 1.0 + alpha * s2;
return (s0, s1, s2, s3);
}
let mut psi_r = psi;
let mut beta_r = beta;
let mut nhalf = 0usize;
while beta_r.abs() >= BETACONTR && nhalf < HALFMAX {
psi_r *= 0.5;
beta_r *= 0.25; nhalf += 1;
}
let mut s0 = 1.0;
let mut s1 = psi_r;
let mut term0 = 1.0; let mut term1 = psi_r;
for j in 1..=JMAX {
term0 *= beta_r / ((2 * j - 1) as f64 * (2 * j) as f64);
s0 += term0;
if term0.abs() < convergence_tol || term0.abs() > overflow_limit {
break;
}
}
for j in 1..=JMAX {
term1 *= beta_r / ((2 * j) as f64 * (2 * j + 1) as f64);
s1 += term1;
if term1.abs() < convergence_tol || term1.abs() > overflow_limit {
break;
}
}
for _ in 0..nhalf {
let c = s0;
let s = s1;
s0 = 2.0 * c * c - 1.0; s1 = 2.0 * c * s; }
let s3 = (s1 - psi) / alpha;
let s2 = (s0 - 1.0) / alpha;
(s0, s1, s2, s3)
}
pub fn principal_angle(a: f64) -> f64 {
a.rem_euclid(DPI)
}
pub fn angle_diff(a: f64, b: f64) -> f64 {
let a = principal_angle(a);
let b = principal_angle(b);
let mut diff = a - b;
if diff > PI {
diff -= DPI;
} else if diff < -PI {
diff += DPI;
}
diff
}
pub fn prelim_kepuni(params: &UniversalKeplerParams, contr: f64) -> Option<f64> {
const ITX: usize = 20;
match params.orbit_type() {
OrbitType::Elliptic => Some(prelim_elliptic(params, contr, ITX)),
OrbitType::Hyperbolic => Some(prelim_hyperbolic(params, contr, ITX)),
OrbitType::Parabolic => None, }
}
pub fn prelim_elliptic(params: &UniversalKeplerParams, contr: f64, max_iter: usize) -> f64 {
let a0 = -params.mu / params.alpha;
let n = (-params.alpha.powi(3)).sqrt() / params.mu;
if params.e0 < contr {
return n * params.dt / (-params.alpha).sqrt();
}
let cos_u0 = (1.0 - params.r0 / a0) / params.e0;
let mut u0 = if cos_u0.abs() <= 1.0 {
cos_u0.acos()
} else if cos_u0 >= 1.0 {
0.0 } else {
PI };
if params.sig0 < 0.0 {
u0 = -u0;
}
u0 = principal_angle(u0);
let ell0 = principal_angle(u0 - params.e0 * u0.sin());
let target_mean_anomaly = principal_angle(ell0 + n * params.dt);
let mut u = PI; for _ in 0..max_iter {
let f = u - params.e0 * u.sin() - target_mean_anomaly;
let fp = 1.0 - params.e0 * u.cos();
let du = -f / fp;
u += du;
if du.abs() < contr * 1e3 {
break;
}
}
angle_diff(u, u0) / (-params.alpha).sqrt()
}
pub fn prelim_hyperbolic(params: &UniversalKeplerParams, contr: f64, max_iter: usize) -> f64 {
let a0 = -params.mu / params.alpha;
let n = params.alpha.powi(3).sqrt() / params.mu;
let coshf0 = (1.0 - params.r0 / a0) / params.e0;
let mut f0 = if coshf0 > 1.0 {
(coshf0 + (coshf0.powi(2) - 1.0).sqrt()).ln()
} else {
0.0 };
if params.sig0 < 0.0 {
f0 = -f0;
}
let ell0 = params.e0 * f0.sinh() - f0;
let target_mean_anomaly = ell0 + n * params.dt;
let mut f: f64 = 0.0;
for _ in 0..max_iter {
if f.abs() < 15.0 {
let func = params.e0 * f.sinh() - f - target_mean_anomaly;
let deriv = params.e0 * f.cosh() - 1.0;
let df = -func / deriv;
let ff = f + df;
f = if f * ff < 0.0 { f / 2.0 } else { ff };
} else {
f /= 2.0;
}
if f.abs() < contr * 1e3 {
break;
}
}
(f - f0) / params.alpha.sqrt()
}
#[inline(always)]
pub fn solve_kepuni_with_guess(
params: &UniversalKeplerParams,
convergency: Option<f64>,
psi_guess: Option<f64>,
) -> Option<(f64, f64, f64, f64, f64)> {
const MAX_ITER: usize = 50;
let tol_step = convergency.unwrap_or(100.0 * f64::EPSILON);
let tol_f = {
let eps10 = 10.0 * f64::EPSILON;
eps10 + eps10 * params.dt.abs()
};
let max_rel_step = 2.0;
match params.orbit_type() {
OrbitType::Parabolic => return None, OrbitType::Elliptic | OrbitType::Hyperbolic => {}
}
let r0 = params.r0;
let sig0 = params.sig0;
let mu = params.mu;
let a = params.alpha;
let dt = params.dt;
let mut psi = if let Some(g) = psi_guess {
g
} else {
prelim_kepuni(params, tol_step)?
};
for _ in 0..MAX_ITER {
if !psi.is_finite() {
psi = 0.5; }
let (s0, s1, s2, s3) = s_funct(psi, a);
let f = (((r0 * s1) + (sig0 * s2)) + (mu * s3)) - dt;
let fp = ((r0 * s0) + (sig0 * s1)) + (mu * s2);
if f.abs() <= tol_f {
return Some((psi, s0, s1, s2, s3));
}
if !fp.is_finite() || fp.abs() < 10.0 * f64::EPSILON {
psi *= 0.5;
continue;
}
let mut dpsi = -f / fp;
let max_step = max_rel_step * (1.0 + psi.abs());
if dpsi.abs() > max_step {
dpsi = dpsi.signum() * max_step;
}
let new_psi = psi + dpsi;
psi = if new_psi * psi < 0.0 {
0.5 * psi
} else {
new_psi
};
let small_abs = dpsi.abs() <= tol_step;
if small_abs {
return Some((psi, s0, s1, s2, s3));
}
let small_rel = dpsi.abs() <= tol_step * (1.0 + psi.abs());
if small_rel {
let (s0f, s1f, s2f, s3f) = s_funct(psi, a);
return Some((psi, s0f, s1f, s2f, s3f));
}
}
None
}
pub fn solve_kepuni(
params: &UniversalKeplerParams,
convergency: Option<f64>,
) -> Option<(f64, f64, f64, f64, f64)> {
solve_kepuni_with_guess(params, convergency, None)
}
pub fn velocity_correction(
x1: &Vector3<f64>,
x2: &Vector3<f64>,
v2: &Vector3<f64>,
dt: f64,
peri_max: f64,
ecc_max: f64,
eps: f64,
) -> Result<(Vector3<f64>, f64, f64), OutfitError> {
let (velocity_corrected, f, g, _) =
velocity_correction_with_guess(x1, x2, v2, dt, peri_max, ecc_max, None, eps)?;
Ok((velocity_corrected, f, g))
}
#[allow(clippy::too_many_arguments)]
pub fn velocity_correction_with_guess(
x1: &Vector3<f64>,
x2: &Vector3<f64>,
v2: &Vector3<f64>,
dt: f64,
peri_max: f64,
ecc_max: f64,
chi_guess: Option<f64>,
eps: f64,
) -> Result<(Vector3<f64>, f64, f64, f64), OutfitError> {
let mu = GAUSS_GRAV_SQUARED;
let r2 = x2.norm();
let sig0 = x2.dot(v2);
let h_norm = x2.cross(v2).norm();
if !h_norm.is_finite() || h_norm <= 1e6 * f64::EPSILON {
return Err(OutfitError::VelocityCorrectionError(
"Rejected orbit: near-zero angular momentum (x × v ≈ 0)".into(),
));
}
let (_, ecc, _, energy) = eccentricity_control(x2, v2, peri_max, ecc_max).ok_or(
OutfitError::VelocityCorrectionError(
"Eccentricity control rejected the candidate orbit".into(),
),
)?;
let params = UniversalKeplerParams {
dt,
r0: r2,
sig0,
mu,
alpha: 2.0 * energy, e0: ecc,
};
let (_chi, _c2, _c3, s2, s3) = solve_kepuni_with_guess(¶ms, Some(eps), chi_guess).ok_or(
OutfitError::VelocityCorrectionError("Universal Kepler solver did not converge".into()),
)?;
let f = 1.0 - (mu * s2) / r2;
let g = dt - (mu * s3);
let g_abs = g.abs();
let g_floor = 100.0 * f64::EPSILON * (1.0 + dt.abs()); if !(g_abs.is_finite()) || g_abs < g_floor {
return Err(OutfitError::VelocityCorrectionError(
"Lagrange coefficient g is too small for a stable velocity update".into(),
));
}
let mut v_corr = *x1; v_corr.axpy(-f, x2, 1.0); v_corr.unscale_mut(g);
Ok((v_corr, f, g, _chi))
}
#[cfg(test)]
mod kepler_test {
use super::*;
mod tests_s_funct {
use approx::assert_relative_eq;
use super::s_funct;
fn check_invariants(psi: f64, alpha: f64, s0: f64, s1: f64, s2: f64, s3: f64) {
let tol = 1e-12;
assert!(
(s0 - (1.0 + alpha * s2)).abs() < tol,
"Invariant s0 = 1 + α*s2 violated: {} vs {}",
s0,
1.0 + alpha * s2
);
assert!(
(s1 - (psi + alpha * s3)).abs() < tol,
"Invariant s1 = ψ + α*s3 violated: {} vs {}",
s1,
psi + alpha * s3
);
}
#[test]
fn test_small_beta() {
let psi = 0.01;
let alpha = 0.1;
let (s0, s1, s2, s3) = s_funct(psi, alpha);
assert!(s0 > 0.0);
assert!(s1 > 0.0);
check_invariants(psi, alpha, s0, s1, s2, s3);
}
#[test]
fn test_large_beta() {
let psi = 10.0;
let alpha = 5.0;
let (s0, s1, s2, s3) = s_funct(psi, alpha);
assert!(s0.is_finite() && s1.is_finite() && s2.is_finite() && s3.is_finite());
let rel_tol = 1e-7;
assert_relative_eq!(s0, 1.0 + alpha * s2, max_relative = rel_tol);
assert_relative_eq!(s1, psi + alpha * s3, max_relative = rel_tol);
}
#[test]
fn test_zero_alpha() {
let psi = 2.0;
let alpha = 0.0;
let (s0, s1, s2, s3) = s_funct(psi, alpha);
assert!((s0 - 1.0).abs() < 1e-14);
assert!((s1 - psi).abs() < 1e-14);
assert!((s2 - psi.powi(2) / 2.0).abs() < 1e-14);
assert!((s3 - psi.powi(3) / 6.0).abs() < 1e-14);
}
#[test]
fn test_zero_psi() {
let psi = 0.0;
let alpha = 2.0;
let (s0, s1, s2, s3) = s_funct(psi, alpha);
assert!((s0 - 1.0).abs() < 1e-14);
assert!((s1 - 0.0).abs() < 1e-14);
assert!((s2 - 0.0).abs() < 1e-14);
assert!((s3 - 0.0).abs() < 1e-14);
}
#[test]
fn test_symmetry_negative_psi() {
let psi = 1.0;
let alpha = 0.5;
let (s0_pos, s1_pos, s2_pos, s3_pos) = s_funct(psi, alpha);
let (s0_neg, s1_neg, s2_neg, s3_neg) = s_funct(-psi, alpha);
let tol = 1e-12;
assert!((s0_pos - s0_neg).abs() < tol);
assert!((s2_pos - s2_neg).abs() < tol);
assert!((s1_pos + s1_neg).abs() < tol);
assert!((s3_pos + s3_neg).abs() < tol);
}
#[test]
fn test_consistency_large_vs_small() {
let psi = 2.5;
let alpha = 1.0;
let (s0, s1, s2, s3) = s_funct(psi, alpha);
check_invariants(psi, alpha, s0, s1, s2, s3);
}
#[test]
fn test_s_funct_real_data() {
let psi = -15.279808141051223;
let alpha = -1.6298946008705195e-4;
let (s0, s1, s2, s3) = s_funct(psi, alpha);
assert_eq!(s0, 0.9810334785583247);
assert_eq!(s1, -15.183083836892674);
assert_eq!(s2, 116.3665517484714);
assert_eq!(s3, -593.4390119881925);
}
}
mod tests_prelim_kepuni {
use super::{prelim_kepuni, UniversalKeplerParams};
const MU: f64 = 1.0;
const CONTR: f64 = 1e-12;
fn make_params(
dt: f64,
r0: f64,
sig0: f64,
mu: f64,
alpha: f64,
e0: f64,
) -> UniversalKeplerParams {
UniversalKeplerParams {
dt,
r0,
sig0,
mu,
alpha,
e0,
}
}
#[test]
fn test_returns_none_for_alpha_zero() {
let params = make_params(1.0, 1.0, 0.0, MU, 0.0, 0.1);
let res = prelim_kepuni(¶ms, CONTR);
assert!(res.is_none());
}
#[test]
fn test_elliptic_small_eccentricity() {
let params = make_params(0.5, 1.0, 0.1, MU, -1.0, 1e-8);
let result = prelim_kepuni(¶ms, CONTR);
assert!(result.is_some());
assert!(result.unwrap().is_finite());
}
#[test]
fn test_elliptic_high_eccentricity() {
let params = make_params(0.1, 0.5, 0.2, MU, -1.0, 0.8);
let result = prelim_kepuni(¶ms, CONTR);
assert!(result.is_some());
assert!(result.unwrap().is_finite());
}
#[test]
fn test_hyperbolic_case() {
let params = make_params(0.3, 2.0, -0.1, MU, 1.0, 1.5);
let result = prelim_kepuni(¶ms, CONTR);
assert!(result.is_some());
assert!(result.unwrap().is_finite());
}
#[test]
fn test_negative_sig0_changes_direction() {
let alpha = -1.0;
let r0 = 1.0;
let e0 = 0.5;
let dt = 0.25;
let params_pos = make_params(dt, r0, 0.1, MU, alpha, e0);
let params_neg = make_params(dt, r0, -0.1, MU, alpha, e0);
let psi_pos = prelim_kepuni(¶ms_pos, CONTR).unwrap();
let psi_neg = prelim_kepuni(¶ms_neg, CONTR).unwrap();
assert!(
(psi_pos - psi_neg).abs() > 1e-8,
"psi did not change significantly when changing sig0 sign: {psi_pos} vs {psi_neg}"
);
}
#[test]
fn test_stability_long_dt() {
let params = make_params(50.0, 1.0, 0.1, MU, -1.0, 0.5);
let result = prelim_kepuni(¶ms, CONTR);
assert!(result.is_some());
assert!(result.unwrap().is_finite());
}
#[test]
fn test_edge_cosine_limits() {
let params = make_params(0.25, 2.0, 0.1, MU, -1.0, 0.1);
let result = prelim_kepuni(¶ms, CONTR);
assert!(result.is_some());
}
#[test]
fn test_prelim_kepuni_real_data() {
let epsilon = f64::EPSILON;
let contr = 100.0 * epsilon;
let dt = -20.765849999996135;
let r0 = 1.3803870211345761;
let sig0 = 3.701_354_484_003_874_8E-3;
let mu = 2.959_122_082_855_911_5E-4;
let alpha = -1.642_158_377_771_140_7E-4;
let e0 = 0.283_599_599_137_344_5;
let params = UniversalKeplerParams {
dt,
r0,
sig0,
mu,
alpha,
e0,
};
let psi = prelim_kepuni(¶ms, contr).unwrap();
assert_eq!(psi, -15.327414893041848);
let params2 = UniversalKeplerParams {
alpha: 1.642_158_377_771_140_7E-4,
..params
};
let psi = prelim_kepuni(¶ms2, contr).unwrap();
assert_eq!(psi, -73.1875935362658);
let params3 = UniversalKeplerParams {
alpha: 0.0,
..params
};
assert!(prelim_kepuni(¶ms3, contr).is_none());
}
mod kepuni_prop_tests {
use super::{prelim_kepuni, UniversalKeplerParams};
use proptest::prelude::*;
fn arb_params() -> impl Strategy<Value = UniversalKeplerParams> {
(
-10.0..10.0f64,
0.1..5.0f64,
-2.0..2.0f64,
0.5..2.0f64,
prop_oneof![(-5.0..-0.01f64), (0.01..5.0f64)],
0.0..3.0f64,
)
.prop_map(|(dt, r0, sig0, mu, alpha, e0)| {
UniversalKeplerParams {
dt,
r0,
sig0,
mu,
alpha,
e0,
}
})
}
proptest! {
#[test]
fn prop_prelim_kepuni_behaves_well(params in arb_params(), contr in 1e-14..1e-8f64) {
let result = prelim_kepuni(¶ms, contr);
prop_assert!(result.is_some());
prop_assert!(result.unwrap().is_finite());
}
}
proptest! {
#[test]
fn prop_prelim_kepuni_alpha_zero(
dt in -10.0..10.0f64,
r0 in 0.1..5.0f64,
sig0 in -2.0..2.0f64,
mu in 0.5..2.0f64,
e0 in 0.0..3.0f64,
contr in 1e-14..1e-8f64
) {
let params = UniversalKeplerParams { dt, r0, sig0, mu, alpha: 0.0, e0 };
let result = prelim_kepuni(¶ms, contr);
if let Some(psi0) = result {
prop_assert!(psi0.is_finite());
}
}
}
proptest! {
#[test]
fn prop_sig0_influences_psi(params in arb_params()) {
let contr = 1e-12;
prop_assume!(params.dt.abs() > 1e-6);
prop_assume!(params.e0 > 1e-6);
prop_assume!(params.r0 > 1e-6);
let mut params_pos = params;
params_pos.sig0 = 0.1;
let mut params_neg = params;
params_neg.sig0 = -0.1;
let res_pos = prelim_kepuni(¶ms_pos, contr);
let res_neg = prelim_kepuni(¶ms_neg, contr);
prop_assume!(res_pos.is_some() && res_neg.is_some());
let diff = (res_pos.unwrap() - res_neg.unwrap()).abs();
prop_assert!(diff >= 0.0);
}
}
}
}
mod tests_solve_kepuni {
use approx::assert_relative_eq;
use proptest::prelude::*;
use super::{solve_kepuni, UniversalKeplerParams};
const MU: f64 = 1.0;
const CONTR: f64 = 1e-12;
fn make_params(
dt: f64,
r0: f64,
sig0: f64,
mu: f64,
alpha: f64,
e0: f64,
) -> UniversalKeplerParams {
UniversalKeplerParams {
dt,
r0,
sig0,
mu,
alpha,
e0,
}
}
#[test]
fn test_solve_kepuni_returns_none_for_alpha_zero() {
let params = make_params(1.0, 1.0, 0.1, MU, 0.0, 0.1);
let res = solve_kepuni(¶ms, Some(CONTR));
assert!(res.is_none());
}
#[test]
fn test_solve_kepuni_elliptical_nominal() {
let params = make_params(0.1, 1.0, 0.1, MU, -1.0, 0.5);
let res = solve_kepuni(¶ms, Some(CONTR));
assert!(
res.is_some(),
"solve_kepuni should converge for elliptical orbit"
);
let (psi, s0, s1, s2, s3) = res.unwrap();
assert!(psi.is_finite());
assert!(s0.is_finite());
assert!(s1.is_finite());
assert!(s2.is_finite());
assert!(s3.is_finite());
}
#[test]
fn test_solve_kepuni_hyperbolic_nominal() {
let params = make_params(0.1, 2.0, -0.1, MU, 1.0, 1.5);
let res = solve_kepuni(¶ms, Some(CONTR));
assert!(
res.is_some(),
"solve_kepuni should converge for hyperbolic orbit"
);
let (psi, s0, s1, s2, s3) = res.unwrap();
assert!(psi.is_finite());
assert!(s0.is_finite());
assert!(s1.is_finite());
assert!(s2.is_finite());
assert!(s3.is_finite());
}
#[test]
fn test_solve_kepuni_large_dt_still_converges() {
let params = make_params(10.0, 1.0, 0.1, MU, -1.0, 0.5);
let res = solve_kepuni(¶ms, Some(CONTR));
assert!(res.is_some(), "solve_kepuni should converge for long dt");
}
#[test]
fn test_solve_kepuni_no_convergency_param_uses_default() {
let params = make_params(0.5, 1.0, 0.2, MU, -1.0, 0.3);
let res = solve_kepuni(¶ms, None);
assert!(
res.is_some(),
"solve_kepuni should converge even with default tolerance"
);
}
#[test]
fn test_solve_kepuni_real_value() {
let params = make_params(
-20.765849999996135,
1.3803870211345761,
3.701_354_484_003_874_8E-3,
2.959_122_082_855_911_5E-4,
-1.642_158_377_771_140_7E-4,
0.283_599_599_137_344_5,
);
let (psi, s0, s1, s2, s3) = solve_kepuni(¶ms, None).unwrap();
assert_eq!(psi, -15.327414893041848);
assert_eq!(s0, 0.9807723505583343);
assert_eq!(s1, -15.229051668919967);
assert_eq!(s2, 117.0876676813769);
assert_eq!(s3, -598.9874390519309);
let params2 = UniversalKeplerParams {
alpha: 1.642_158_377_771_140_7E-4,
..params
};
let (psi, s0, s1, s2, s3) = solve_kepuni(¶ms2, None).unwrap();
assert_eq!(psi, -15.1324122746124);
assert_eq!(s0, 1.0188608766146905);
assert_eq!(s1, -15.227430038021337);
assert_eq!(s2, 114.854187452308);
assert_eq!(s3, -578.615100072754);
}
#[test]
fn test_solve_kepuni_invariant_residual_elliptical() {
let params = make_params(0.1, 1.0, 0.1, MU, -1.0, 0.5);
let res = solve_kepuni(¶ms, Some(CONTR)).expect("Expected convergence");
let (_, _, s1, s2, s3) = res;
let residual = params.r0 * s1 + params.sig0 * s2 + MU * s3 - params.dt;
assert_relative_eq!(residual, 0.0, epsilon = 1e-12);
}
#[test]
fn test_solve_kepuni_invariant_residual_hyperbolic() {
let params = make_params(0.1, 2.0, -0.1, MU, 1.0, 1.5);
let res = solve_kepuni(¶ms, Some(CONTR)).expect("Expected convergence");
let (_, _, s1, s2, s3) = res;
let residual = params.r0 * s1 + params.sig0 * s2 + MU * s3 - params.dt;
assert_relative_eq!(residual, 0.0, max_relative = 1e-10);
}
#[test]
fn test_solve_kepuni_known_values_elliptical() {
let params = make_params(0.5, 1.0, 0.2, MU, -1.0, 0.3);
let (psi, _, _, _, _) =
solve_kepuni(¶ms, Some(CONTR)).expect("Expected convergence");
let expected_psi = 0.47761843287737277;
assert_relative_eq!(psi, expected_psi, epsilon = 1e-15);
}
#[test]
fn test_solve_kepuni_known_values_hyperbolic() {
let params = make_params(0.3, 2.0, -0.1, MU, 1.0, 1.5);
let (psi, _, _, _, _) =
solve_kepuni(¶ms, Some(CONTR)).expect("Expected convergence");
let expected_psi = 0.14972146123530983;
assert_relative_eq!(psi, expected_psi, epsilon = 1e-15);
}
fn arb_common_params() -> impl Strategy<Value = (f64, f64, f64, f64, f64)> {
(
prop_oneof![-5.0..-0.1, 0.1..5.0],
0.1..3.0f64,
-0.5..0.5f64,
0.5..2.0f64,
0.01..2.0f64,
)
}
proptest! {
#[test]
fn prop_solve_kepuni_elliptical(
(dt, r0, sig0, mu, e0) in arb_common_params(),
alpha in -5.0..-0.1f64
) {
let params = make_params(dt, r0, sig0, mu, alpha, e0);
let res = solve_kepuni(¶ms, Some(1e-12));
if let Some((psi, s0, s1, s2, s3)) = res {
prop_assert!(psi.is_finite());
prop_assert!(s0.is_finite());
prop_assert!(s1.is_finite());
prop_assert!(s2.is_finite());
prop_assert!(s3.is_finite());
let residual = params.r0 * s1 + params.sig0 * s2 + mu * s3 - params.dt;
prop_assert!(residual.abs() < 1e-8,
"Residual too large for elliptical case: {}", residual);
}
}
}
proptest! {
#[test]
fn prop_solve_kepuni_hyperbolic(
(dt, r0, sig0, mu, e0) in arb_common_params(),
alpha in 0.1..5.0f64
) {
let params = make_params(dt, r0, sig0, mu, alpha, e0);
let res = solve_kepuni(¶ms, Some(1e-12));
if let Some((psi, s0, s1, s2, s3)) = res {
prop_assert!(psi.is_finite());
prop_assert!(s0.is_finite());
prop_assert!(s1.is_finite());
prop_assert!(s2.is_finite());
prop_assert!(s3.is_finite());
let residual = params.r0 * s1 + params.sig0 * s2 + mu * s3 - params.dt;
prop_assert!(residual.abs() < 1e-8,
"Residual too large for hyperbolic case: {}", residual);
}
}
}
}
mod tests_velocity_correction {
use super::*;
use approx::assert_relative_eq;
use nalgebra::Vector3;
const KEP_EPS: f64 = 1e3 * f64::EPSILON;
fn v(x: f64, y: f64, z: f64) -> Vector3<f64> {
Vector3::new(x, y, z)
}
#[test]
fn test_velocity_correction_nominal_elliptic() {
let x1 = v(1.0, 0.0, 0.0);
let x2 = v(1.1, 0.0, 0.0);
let v2 = v(0.0, 0.017, 0.0); let dt = 1.0; let peri_max = 5.0;
let ecc_max = 0.9;
let result = velocity_correction(&x1, &x2, &v2, dt, peri_max, ecc_max, KEP_EPS);
assert!(result.is_ok(), "Velocity correction should succeed");
let (vcorr, f, g) = result.unwrap();
assert!(vcorr.iter().all(|c| c.is_finite()));
assert!(f.is_finite());
assert!(g.is_finite());
assert!(f > 0.5 && f < 1.5, "f coefficient is in a reasonable range");
assert_relative_eq!(g, dt, epsilon = 0.1);
}
#[test]
fn test_velocity_correction_rejects_bad_orbit() {
let x1 = v(1e6, 1e6, 1e6);
let x2 = v(1e6, 1e6, 1e6);
let v2 = v(1e6, 1e6, 1e6);
let dt = 0.1;
let peri_max = 0.001;
let ecc_max = 0.001;
let result = velocity_correction(&x1, &x2, &v2, dt, peri_max, ecc_max, KEP_EPS);
assert!(
result.is_err(),
"Velocity correction should fail on an invalid orbit"
);
}
#[test]
fn test_velocity_correction_sensitivity_to_x1() {
let x1 = v(1.0, 0.0, 0.0);
let x1_shifted = v(1.05, 0.0, 0.0); let x2 = v(1.1, 0.0, 0.0);
let v2 = v(0.0, 0.017, 0.0);
let dt = 1.0;
let peri_max = 5.0;
let ecc_max = 0.9;
let result1 =
velocity_correction(&x1, &x2, &v2, dt, peri_max, ecc_max, KEP_EPS).unwrap();
let result2 =
velocity_correction(&x1_shifted, &x2, &v2, dt, peri_max, ecc_max, KEP_EPS).unwrap();
let (v_corr1, _, _) = result1;
let (v_corr2, _, _) = result2;
let diff = (v_corr1 - v_corr2).norm();
assert!(
diff > 1e-6,
"Changing x1 must influence the corrected velocity (diff = {diff})"
);
}
#[test]
fn test_velocity_correction_output_not_nan() {
let x1 = v(1.0, 0.5, 0.0);
let x2 = v(1.1, 0.2, 0.0);
let v2 = v(0.0, 0.02, 0.0);
let dt = 0.5;
let peri_max = 5.0;
let ecc_max = 0.9;
let (v_corr, f, g) =
velocity_correction(&x1, &x2, &v2, dt, peri_max, ecc_max, KEP_EPS).unwrap();
assert!(
!v_corr.iter().any(|x| x.is_nan()),
"Corrected velocity vector contains NaN"
);
assert!(!f.is_nan(), "f coefficient is NaN");
assert!(!g.is_nan(), "g coefficient is NaN");
}
#[test]
fn test_velocity_correction_real_data() {
let x1 = Vector3::new(
-0.843_561_126_129_683_3,
0.937_288_327_370_772_8,
0.659_183_901_029_776_6,
);
let x2 = Vector3::new(
-0.623_121_622_917_384,
1.0076797884556383,
0.708_125_687_984_424_5,
);
let v2 = Vector3::new(
-1.552_431_036_862_405_6E-2,
-3.984_104_176_604_068E-3,
-2.764_015_436_163_718_3E-3,
);
let dt = 14.731970000000729;
let (v2, f, g) = velocity_correction(&x1, &x2, &v2, dt, 1., 1., KEP_EPS).unwrap();
assert_eq!(f, 0.988_164_877_097_290_6);
assert_eq!(g, 14.674676076120734);
assert_eq!(
v2.as_slice(),
[
-0.015524310248562921,
-0.003984104769239458,
-0.0027640155187336176
]
)
}
mod velocity_correction_prop_tests {
use super::*;
use nalgebra::Vector3;
use proptest::prelude::*;
fn v(x: f64, y: f64, z: f64) -> Vector3<f64> {
Vector3::new(x, y, z)
}
fn arb_orbit_params(
) -> impl Strategy<Value = (Vector3<f64>, Vector3<f64>, Vector3<f64>, f64, f64, f64)>
{
(
(-2.0..2.0f64, -2.0..2.0f64, -2.0..2.0f64)
.prop_map(|(x, y, z)| v(x + 1.0, y + 1.0, z)), (-2.0..2.0f64, -2.0..2.0f64, -2.0..2.0f64)
.prop_map(|(x, y, z)| v(x + 1.1, y + 1.0, z)),
(-0.05..0.05f64, -0.05..0.05f64, -0.05..0.05f64)
.prop_map(|(x, y, z)| v(x, y, z)),
0.01..5.0f64,
0.1..10.0f64,
0.1..2.0f64,
)
}
proptest! {
#[test]
fn prop_velocity_correction_no_panic(
(x1,x2,v2,dt,peri_max,ecc_max) in arb_orbit_params()
) {
let res = velocity_correction(&x1,&x2,&v2,dt,peri_max,ecc_max, KEP_EPS);
match res {
Ok((vcorr, f, g)) => {
prop_assert!(vcorr.iter().all(|c| c.is_finite()));
prop_assert!(f.is_finite());
prop_assert!(g.is_finite());
}
Err(_) => {
}
}
}
}
proptest! {
#[test]
fn prop_velocity_correction_sensitive_to_x1(
(x1,x2,v2,dt,peri_max,ecc_max) in arb_orbit_params()
) {
let mut x1_shifted = x1;
x1_shifted[0] += 0.01;
let res1 = velocity_correction(&x1,&x2,&v2,dt,peri_max,ecc_max, KEP_EPS);
let res2 = velocity_correction(&x1_shifted,&x2,&v2,dt,peri_max,ecc_max, KEP_EPS);
if let (Ok((v1, _, _)), Ok((v2c, _, _))) = (res1, res2) {
let diff = (v1 - v2c).norm();
prop_assert!(diff >= 0.0);
}
}
}
}
}
}