1use nalgebra::{Vector3, Vector6};
10use std::collections::HashMap;
11
12use crate::derivatives::compute_derivatives;
13use crate::wind::WindSegment;
14use crate::BallisticInputs;
15use crate::DragModel;
16
17fn rk4_step(state: &Vector6<f64>, t: f64, dt: f64, params: &TrajectoryParams) -> Vector6<f64> {
19 let k1 = compute_derivatives_vec(state, t, params);
21 let k2 = compute_derivatives_vec(&(state + dt * 0.5 * k1), t + dt * 0.5, params);
22 let k3 = compute_derivatives_vec(&(state + dt * 0.5 * k2), t + dt * 0.5, params);
23 let k4 = compute_derivatives_vec(&(state + dt * k3), t + dt, params);
24
25 state + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
26}
27
28fn rk45_step(
30 state: &Vector6<f64>,
31 t: f64,
32 dt: f64,
33 params: &TrajectoryParams,
34 tol: f64,
35) -> (Vector6<f64>, f64, f64) {
36 const A21: f64 = 1.0 / 5.0;
38 const A31: f64 = 3.0 / 40.0;
39 const A32: f64 = 9.0 / 40.0;
40 const A41: f64 = 44.0 / 45.0;
41 const A42: f64 = -56.0 / 15.0;
42 const A43: f64 = 32.0 / 9.0;
43 const A51: f64 = 19372.0 / 6561.0;
44 const A52: f64 = -25360.0 / 2187.0;
45 const A53: f64 = 64448.0 / 6561.0;
46 const A54: f64 = -212.0 / 729.0;
47 const A61: f64 = 9017.0 / 3168.0;
48 const A62: f64 = -355.0 / 33.0;
49 const A63: f64 = 46732.0 / 5247.0;
50 const A64: f64 = 49.0 / 176.0;
51 const A65: f64 = -5103.0 / 18656.0;
52 const A71: f64 = 35.0 / 384.0;
53 const A73: f64 = 500.0 / 1113.0;
54 const A74: f64 = 125.0 / 192.0;
55 const A75: f64 = -2187.0 / 6784.0;
56 const A76: f64 = 11.0 / 84.0;
57
58 const B1: f64 = 35.0 / 384.0;
60 const B3: f64 = 500.0 / 1113.0;
61 const B4: f64 = 125.0 / 192.0;
62 const B5: f64 = -2187.0 / 6784.0;
63 const B6: f64 = 11.0 / 84.0;
64
65 const B1_ERR: f64 = 5179.0 / 57600.0;
67 const B3_ERR: f64 = 7571.0 / 16695.0;
68 const B4_ERR: f64 = 393.0 / 640.0;
69 const B5_ERR: f64 = -92097.0 / 339200.0;
70 const B6_ERR: f64 = 187.0 / 2100.0;
71 const B7_ERR: f64 = 1.0 / 40.0;
72
73 let k1 = compute_derivatives_vec(state, t, params);
75 let k2 = compute_derivatives_vec(&(state + dt * A21 * k1), t + dt * 0.2, params);
76 let k3 = compute_derivatives_vec(&(state + dt * (A31 * k1 + A32 * k2)), t + dt * 0.3, params);
77 let k4 = compute_derivatives_vec(
78 &(state + dt * (A41 * k1 + A42 * k2 + A43 * k3)),
79 t + dt * 0.8,
80 params,
81 );
82 let k5 = compute_derivatives_vec(
83 &(state + dt * (A51 * k1 + A52 * k2 + A53 * k3 + A54 * k4)),
84 t + dt * 8.0 / 9.0,
85 params,
86 );
87 let k6 = compute_derivatives_vec(
88 &(state + dt * (A61 * k1 + A62 * k2 + A63 * k3 + A64 * k4 + A65 * k5)),
89 t + dt,
90 params,
91 );
92 let k7 = compute_derivatives_vec(
93 &(state + dt * (A71 * k1 + A73 * k3 + A74 * k4 + A75 * k5 + A76 * k6)),
94 t + dt,
95 params,
96 );
97
98 let y_new = state + dt * (B1 * k1 + B3 * k3 + B4 * k4 + B5 * k5 + B6 * k6);
100
101 let y_err = state
103 + dt * (B1_ERR * k1 + B3_ERR * k3 + B4_ERR * k4 + B5_ERR * k5 + B6_ERR * k6 + B7_ERR * k7);
104
105 let error = (y_new - y_err).norm() / (1.0 + state.norm());
107
108 let safety = 0.9;
110 let dt_new = if error < tol {
111 dt * safety * (tol / error).powf(0.2).min(2.0)
112 } else {
113 dt * safety * (tol / error).powf(0.25).max(0.1)
114 };
115
116 (y_new, dt_new, error)
117}
118
119pub struct TrajectoryParams {
121 pub mass_kg: f64,
122 pub bc: f64,
123 pub drag_model: DragModel,
124 pub wind_segments: Vec<WindSegment>,
125 pub atmos_params: (f64, f64, f64, f64),
126 pub omega_vector: Option<Vector3<f64>>,
127 pub enable_spin_drift: bool,
128 pub enable_magnus: bool,
129 pub enable_coriolis: bool,
130 pub target_distance_m: f64, pub enable_wind_shear: bool,
132 pub wind_shear_model: String,
133 pub shooter_altitude_m: f64,
134 pub is_twist_right: bool, pub custom_drag_table: Option<crate::drag::DragTable>, pub bc_segments: Option<Vec<(f64, f64)>>, pub use_bc_segments: bool, }
139
140fn compute_derivatives_vec(
142 state: &Vector6<f64>,
143 t: f64,
144 params: &TrajectoryParams,
145) -> Vector6<f64> {
146 let pos = Vector3::new(state[0], state[1], state[2]);
147 let vel = Vector3::new(state[3], state[4], state[5]);
148
149 let wind_vector = if !params.wind_segments.is_empty() {
151 if params.enable_wind_shear && params.wind_shear_model != "none" {
152 crate::wind_shear::get_wind_at_position(
153 &pos,
154 ¶ms.wind_segments,
155 params.enable_wind_shear,
156 ¶ms.wind_shear_model,
157 params.shooter_altitude_m,
158 )
159 } else {
160 let seg = ¶ms.wind_segments[0];
162 let wind_speed_mps = seg.0 * 0.2777778; let wind_angle_rad = seg.1.to_radians();
164 Vector3::new(
166 -wind_speed_mps * wind_angle_rad.cos(), 0.0, -wind_speed_mps * wind_angle_rad.sin(), )
170 }
171 } else {
172 Vector3::zeros()
173 };
174
175 let inputs = BallisticInputs {
177 bc_value: params.bc,
178 bc_type: params.drag_model,
179 bullet_mass: params.mass_kg, muzzle_velocity: vel.norm(), bullet_diameter: 0.0078232, bullet_length: 0.031496, twist_rate: 10.0, is_twist_right: params.is_twist_right,
186 enable_advanced_effects: params.enable_spin_drift
187 || params.enable_magnus
188 || params.enable_coriolis,
189 enable_magnus: params.enable_magnus,
190 enable_coriolis: params.enable_coriolis,
191 altitude: params.atmos_params.0,
192 temperature: params.atmos_params.1,
193 pressure: params.atmos_params.2,
194 humidity: params.atmos_params.3,
195 tipoff_yaw: 0.0,
196 target_distance: 1000.0, muzzle_angle: 0.0,
198 wind_speed: if !params.wind_segments.is_empty() {
202 params.wind_segments[0].0 * 0.2777778 } else {
204 0.0
205 },
206 wind_angle: if !params.wind_segments.is_empty() {
207 params.wind_segments[0].1.to_radians() } else {
209 0.0
210 },
211 latitude: None,
212 shooting_angle: 0.0,
213 azimuth_angle: 0.0,
214 use_powder_sensitivity: false,
215 powder_temp_sensitivity: 0.0,
216 powder_temp: 59.0,
217 tipoff_decay_distance: 0.0,
218 ground_threshold: -1000.0,
219 bc_segments: params.bc_segments.clone(),
220 caliber_inches: 0.308,
221 weight_grains: params.mass_kg / 0.00006479891,
222 use_bc_segments: params.use_bc_segments,
223 bullet_id: None,
224 bc_segments_data: None,
225 use_enhanced_spin_drift: params.enable_spin_drift,
226 use_form_factor: false,
227 manufacturer: None,
228 bullet_model: None,
229 enable_wind_shear: false,
230 wind_shear_model: "none".to_string(),
231 use_cluster_bc: false,
232 bullet_cluster: None,
233
234 custom_drag_table: params.custom_drag_table.clone(),
236
237 bc_type_str: None,
238 enable_pitch_damping: false,
239 enable_precession_nutation: false,
240 use_rk4: true,
241 use_adaptive_rk45: false,
242 enable_trajectory_sampling: false,
243 sample_interval: 10.0,
244 sight_height: 0.0,
245 muzzle_height: 0.0,
246 target_height: 0.0,
247 };
248
249 let deriv_result = compute_derivatives(
251 pos,
252 vel,
253 &inputs,
254 wind_vector,
255 params.atmos_params,
256 params.bc,
257 params.omega_vector,
258 t,
259 );
260
261 Vector6::new(
262 deriv_result[0],
263 deriv_result[1],
264 deriv_result[2],
265 deriv_result[3],
266 deriv_result[4],
267 deriv_result[5],
268 )
269}
270
271pub fn integrate_trajectory(
273 initial_state: [f64; 6],
274 t_span: (f64, f64),
275 params: TrajectoryParams,
276 method: &str,
277 tolerance: f64,
278 max_step: f64,
279) -> Vec<(f64, Vector6<f64>)> {
280 let mut state = Vector6::new(
281 initial_state[0],
282 initial_state[1],
283 initial_state[2],
284 initial_state[3],
285 initial_state[4],
286 initial_state[5],
287 );
288
289 let mut t = t_span.0;
290 let t_end = t_span.1;
291 let mut dt = (t_end - t) / 1000.0; let mut trajectory = Vec::with_capacity(10000);
294 trajectory.push((t, state));
295
296 match method {
297 "RK4" => {
298 dt = dt.min(max_step).min(0.001); while t < t_end {
302 if t + dt > t_end {
303 dt = t_end - t;
304 }
305
306 let new_state = rk4_step(&state, t, dt, ¶ms);
307
308 if state[0] < params.target_distance_m && new_state[0] >= params.target_distance_m {
310 let alpha = (params.target_distance_m - state[0]) / (new_state[0] - state[0]);
312 let dt_to_target = dt * alpha;
313
314 let final_state = rk4_step(&state, t, dt_to_target, ¶ms);
316
317 let mut corrected_state = final_state;
319 if corrected_state[0] > params.target_distance_m {
320 corrected_state[0] = params.target_distance_m;
321 }
322
323 trajectory.push((t + dt_to_target, corrected_state));
324 break; }
326
327 state = new_state;
328 t += dt;
329 trajectory.push((t, state));
330
331 if state[0] >= params.target_distance_m {
333 let mut final_state = state;
336 final_state[0] = params.target_distance_m; trajectory.push((t, final_state));
338 break;
339 }
340
341 if state[1] < -1000.0 {
343 break;
344 }
345 }
346 }
347 "RK45" | _ => {
348 let mut last_save_x = 0.0; let save_interval_m = params.target_distance_m / 50.0; let effective_max_step =
355 if params.enable_wind_shear && params.wind_shear_model != "none" {
356 if params.target_distance_m > 800.0 {
358 0.01 } else {
360 0.02 }
362 } else {
363 max_step };
365
366 dt = dt.min(effective_max_step).max(0.0001); let max_iterations = 100000; let mut iteration_count = 0;
372
373 while t < t_end && iteration_count < max_iterations {
374 iteration_count += 1;
375
376 if t + dt > t_end {
378 dt = t_end - t;
379 }
380
381 let (new_state, dt_new, _error) = rk45_step(&state, t, dt, ¶ms, tolerance);
382
383 if state[0] < params.target_distance_m && new_state[0] >= params.target_distance_m {
385 let alpha = (params.target_distance_m - state[0]) / (new_state[0] - state[0]);
387 let dt_to_target = dt * alpha;
388
389 let (final_state, _, _) =
391 rk45_step(&state, t, dt_to_target, ¶ms, tolerance);
392
393 let mut corrected_state = final_state;
395 if corrected_state[0] > params.target_distance_m {
396 corrected_state[0] = params.target_distance_m;
397 }
398
399 trajectory.push((t + dt_to_target, corrected_state));
400 break; }
402
403 state = new_state;
405 t += dt;
406
407 if state[0] - last_save_x >= save_interval_m || state[0] >= params.target_distance_m
409 {
410 trajectory.push((t, state));
412 last_save_x = state[0];
413 }
414
415 dt = dt_new.min(effective_max_step).max(0.0001); if state[0] >= params.target_distance_m {
420 let mut final_state = state;
423 final_state[0] = params.target_distance_m; trajectory.push((t, final_state));
425 break;
426 }
427
428 if state[1] < -1000.0 {
430 break;
431 }
432 }
433
434 if iteration_count >= max_iterations {
436 eprintln!(
437 "WARNING: Trajectory integration hit maximum iteration limit ({} iterations)",
438 max_iterations
439 );
440 eprintln!(" Final time: {}, Target time: {}", t, t_end);
441 eprintln!(
442 " Final position: downrange(x)={}, Target: {}m",
443 state[0], params.target_distance_m
444 );
445 }
446 }
447 }
448
449 trajectory
450}
451
452pub fn solve_trajectory_rust(
454 initial_state: [f64; 6],
455 t_span: (f64, f64),
456 mass_kg: f64,
457 bc: f64,
458 drag_model: DragModel,
459 wind_segments: Vec<WindSegment>,
460 atmos_params: (f64, f64, f64, f64),
461 omega_vector: Option<Vec<f64>>,
462 enable_spin_drift: bool,
463 enable_magnus: bool,
464 enable_coriolis: bool,
465 method: String,
466 tolerance: f64,
467 max_step: f64,
468 target_distance_m: f64,
469) -> Vec<HashMap<String, f64>> {
470 let omega_vec = omega_vector.map(|v| Vector3::new(v[0], v[1], v[2]));
471
472 let params = TrajectoryParams {
473 mass_kg,
474 bc,
475 drag_model,
476 wind_segments,
477 atmos_params,
478 omega_vector: omega_vec,
479 enable_spin_drift,
480 enable_magnus,
481 enable_coriolis,
482 target_distance_m,
483 enable_wind_shear: false, wind_shear_model: "none".to_string(),
485 shooter_altitude_m: 0.0,
486 is_twist_right: true, custom_drag_table: None, bc_segments: None, use_bc_segments: false,
490 };
491
492 let trajectory =
493 integrate_trajectory(initial_state, t_span, params, &method, tolerance, max_step);
494
495 trajectory
497 .into_iter()
498 .map(|(t, state)| {
499 let mut point = HashMap::new();
500 point.insert("t".to_string(), t);
501 point.insert("x".to_string(), state[0]);
502 point.insert("y".to_string(), state[1]);
503 point.insert("z".to_string(), state[2]);
504 point.insert("vx".to_string(), state[3]);
505 point.insert("vy".to_string(), state[4]);
506 point.insert("vz".to_string(), state[5]);
507 point
508 })
509 .collect()
510}
511
512#[cfg(test)]
513mod tests {
514 use super::*;
515
516 fn create_test_params(target_distance_m: f64) -> TrajectoryParams {
517 TrajectoryParams {
518 mass_kg: 0.01134, bc: 0.442,
520 drag_model: DragModel::G7,
521 wind_segments: vec![],
522 atmos_params: (0.0, 59.0, 29.92, 0.0),
523 omega_vector: None,
524 enable_spin_drift: false,
525 enable_magnus: false,
526 enable_coriolis: false,
527 target_distance_m,
528 enable_wind_shear: false,
529 wind_shear_model: "none".to_string(),
530 shooter_altitude_m: 0.0,
531 is_twist_right: true,
532 custom_drag_table: None,
533 bc_segments: None,
534 use_bc_segments: false,
535 }
536 }
537
538 #[test]
539 fn test_integrate_trajectory_basic() {
540 let initial_state = [0.0, -0.038, 0.0, 821.52, 48.61, 0.0];
543
544 let params = TrajectoryParams {
545 mass_kg: 0.01134, bc: 0.442,
547 drag_model: DragModel::G7,
548 wind_segments: vec![(0.0, 90.0, 914.4)],
549 atmos_params: (0.0, 59.0, 29.92, 0.0),
550 omega_vector: None,
551 enable_spin_drift: false,
552 enable_magnus: false,
553 enable_coriolis: false,
554 target_distance_m: 914.4, enable_wind_shear: false,
556 wind_shear_model: "none".to_string(),
557 shooter_altitude_m: 0.0,
558 is_twist_right: true,
559 custom_drag_table: None,
560 bc_segments: None,
561 use_bc_segments: false,
562 };
563
564 println!("Running integrate_trajectory test...");
565 println!("Initial state: {:?}", initial_state);
566 println!("Target distance: {} m", params.target_distance_m);
567
568 let trajectory =
569 integrate_trajectory(initial_state, (0.0, 10.0), params, "RK45", 1e-6, 0.01);
570
571 println!("Trajectory has {} points", trajectory.len());
572
573 assert!(
575 trajectory.len() > 1,
576 "Trajectory should have more than 1 point, but has {}",
577 trajectory.len()
578 );
579
580 if let Some((_, final_state)) = trajectory.last() {
582 println!("Final state: downrange(x)={}", final_state[0]);
583 assert!(
584 final_state[0] > 0.0,
585 "Final x should be positive (bullet moved downrange)"
586 );
587 assert!(
588 final_state[0] >= 900.0,
589 "Final x should be near target distance"
590 );
591 }
592 }
593
594 #[test]
595 fn test_rk4_vs_rk45_consistency() {
596 let initial_state = [0.0, 0.0, 0.0, 800.0, 30.0, 0.0]; let target_distance = 500.0;
599
600 let params_rk4 = create_test_params(target_distance);
601 let params_rk45 = create_test_params(target_distance);
602
603 let trajectory_rk4 =
604 integrate_trajectory(initial_state, (0.0, 5.0), params_rk4, "RK4", 1e-6, 0.001);
605 let trajectory_rk45 =
606 integrate_trajectory(initial_state, (0.0, 5.0), params_rk45, "RK45", 1e-6, 0.01);
607
608 assert!(!trajectory_rk4.is_empty());
610 assert!(!trajectory_rk45.is_empty());
611
612 let (_, final_rk4) = trajectory_rk4.last().unwrap();
613 let (_, final_rk45) = trajectory_rk45.last().unwrap();
614
615 let rk4_z = final_rk4[0];
617 let rk45_z = final_rk45[0];
618 let diff_percent = ((rk4_z - rk45_z) / rk45_z).abs() * 100.0;
619
620 assert!(
621 diff_percent < 1.0,
622 "RK4 and RK45 final positions differ by {}%: RK4={}, RK45={}",
623 diff_percent,
624 rk4_z,
625 rk45_z
626 );
627 }
628
629 #[test]
630 fn test_ground_impact_detection() {
631 let initial_state = [0.0, 100.0, 0.0, 300.0, -50.0, 0.0]; let mut params = create_test_params(10000.0); params.target_distance_m = 10000.0;
636
637 let trajectory =
638 integrate_trajectory(initial_state, (0.0, 20.0), params, "RK45", 1e-6, 0.01);
639
640 let (_, final_state) = trajectory.last().unwrap();
642
643 assert!(
645 final_state[1] <= -900.0,
646 "Should hit ground, but y={}",
647 final_state[1]
648 );
649 assert!(
650 final_state[0] < 10000.0,
651 "Should not reach target, but z={}",
652 final_state[0]
653 );
654 }
655
656 #[test]
657 fn test_target_distance_reached() {
658 let initial_state = [0.0, 0.0, 0.0, 800.0, 20.0, 0.0]; let target_distance = 300.0;
660
661 let params = create_test_params(target_distance);
662
663 let trajectory =
664 integrate_trajectory(initial_state, (0.0, 5.0), params, "RK45", 1e-6, 0.01);
665
666 let (_, final_state) = trajectory.last().unwrap();
667
668 assert!(
670 (final_state[0] - target_distance).abs() < 1.0,
671 "Should reach target at {}m, but stopped at {}m",
672 target_distance,
673 final_state[0]
674 );
675 }
676
677 #[test]
678 fn test_wind_affects_trajectory() {
679 let initial_state = [0.0, 0.0, 0.0, 800.0, 30.0, 0.0]; let target_distance = 500.0;
684
685 let params_no_wind = create_test_params(target_distance);
687
688 let mut params_headwind = create_test_params(target_distance);
690 params_headwind.wind_segments = vec![(72.0, 0.0, 500.0)]; let trajectory_no_wind =
693 integrate_trajectory(initial_state, (0.0, 5.0), params_no_wind, "RK45", 1e-6, 0.01);
694 let trajectory_headwind =
695 integrate_trajectory(initial_state, (0.0, 5.0), params_headwind, "RK45", 1e-6, 0.01);
696
697 assert!(!trajectory_no_wind.is_empty(), "No-wind trajectory should complete");
699 assert!(!trajectory_headwind.is_empty(), "Headwind trajectory should complete");
700
701 let (time_no_wind, final_no_wind) = trajectory_no_wind.last().unwrap();
702 let (time_headwind, final_headwind) = trajectory_headwind.last().unwrap();
703
704 let drop_no_wind = final_no_wind[1];
707 let drop_headwind = final_headwind[1];
708
709 println!("No wind: time={}, drop={}", time_no_wind, drop_no_wind);
712 println!("Headwind: time={}, drop={}", time_headwind, drop_headwind);
713
714 assert!(
716 (final_no_wind[0] - target_distance).abs() < 10.0,
717 "No-wind should reach target"
718 );
719 assert!(
720 (final_headwind[0] - target_distance).abs() < 10.0,
721 "Headwind should reach target"
722 );
723 }
724
725 #[test]
726 fn test_solve_trajectory_rust_output_format() {
727 let initial_state = [0.0, 0.0, 0.0, 800.0, 30.0, 0.0]; let result = solve_trajectory_rust(
730 initial_state,
731 (0.0, 2.0),
732 0.01134, 0.442, DragModel::G7, vec![], (0.0, 59.0, 29.92, 0.0), None, false, false, false, "RK45".to_string(), 1e-6, 0.01, 500.0, );
746
747 assert!(!result.is_empty());
749
750 let first_point = &result[0];
751 assert!(first_point.contains_key("t"));
752 assert!(first_point.contains_key("x"));
753 assert!(first_point.contains_key("y"));
754 assert!(first_point.contains_key("z"));
755 assert!(first_point.contains_key("vx"));
756 assert!(first_point.contains_key("vy"));
757 assert!(first_point.contains_key("vz"));
758 }
759
760 #[test]
761 fn test_left_vs_right_twist() {
762 let initial_state = [0.0, 0.0, 0.0, 800.0, 30.0, 0.0]; let target_distance = 500.0;
764
765 let mut params_right = create_test_params(target_distance);
766 params_right.is_twist_right = true;
767 params_right.enable_spin_drift = true;
768
769 let mut params_left = create_test_params(target_distance);
770 params_left.is_twist_right = false;
771 params_left.enable_spin_drift = true;
772
773 let trajectory_right =
774 integrate_trajectory(initial_state, (0.0, 5.0), params_right, "RK45", 1e-6, 0.01);
775 let trajectory_left =
776 integrate_trajectory(initial_state, (0.0, 5.0), params_left, "RK45", 1e-6, 0.01);
777
778 assert!(!trajectory_right.is_empty());
780 assert!(!trajectory_left.is_empty());
781
782 let (_, final_right) = trajectory_right.last().unwrap();
784 let (_, final_left) = trajectory_left.last().unwrap();
785
786 assert!((final_right[2] - final_left[2]).abs() < 10.0);
788 }
789}