Skip to main content

03_geo_raise_optim/
raise_optim.rs

1#![doc = include_str!("./README.md")]
2extern crate log;
3extern crate nyx_space as nyx;
4extern crate pretty_env_logger as pel;
5
6use anise::{
7    almanac::{metaload::MetaFile, Almanac},
8    constants::{
9        celestial_objects::{MOON, SUN},
10        frames::{EARTH_J2000, IAU_EARTH_FRAME},
11    },
12};
13use hifitime::{Epoch, TimeUnits, Unit};
14use log::info;
15use nyx::{
16    cosmic::{GuidanceMode, Mass, MetaAlmanac, Orbit, SRPData},
17    dynamics::{
18        guidance::{Ruggiero, Thruster},
19        GravityField, OrbitalDynamics, SolarPressure, SpacecraftDynamics,
20    },
21    io::gravity::GravityFieldData,
22    md::prelude::{Objective, OrbitalElement, StateParameter},
23    propagators::{ErrorControl, IntegratorOptions, Propagator},
24    Spacecraft,
25};
26use radiate::*;
27use std::{error::Error, sync::Arc};
28
29// Shared state struct for the fitness evaluation to avoid reading files thousands of times
30struct SharedState {
31    almanac: Arc<Almanac>,
32    harmonics: Arc<GravityField>,
33    srp_dyn: Arc<SolarPressure>,
34}
35
36impl SharedState {
37    fn new() -> Result<Self, Box<dyn Error>> {
38        let almanac = Arc::new(MetaAlmanac::latest().map_err(Box::new)?);
39
40        let mut jgm3_meta = MetaFile {
41            uri: "http://public-data.nyxspace.com/nyx/models/JGM3.cof.gz".to_string(),
42            crc32: Some(0xF446F027),
43        };
44        jgm3_meta.process(true)?;
45
46        let harmonics = GravityField::from_stor(
47            almanac.frame_info(IAU_EARTH_FRAME)?,
48            GravityFieldData::from_cof(&jgm3_meta.uri, 8, 8, true)?,
49        );
50        let srp_dyn = SolarPressure::default_flux(EARTH_J2000, almanac.clone())?;
51
52        Ok(Self {
53            almanac,
54            harmonics,
55            srp_dyn,
56        })
57    }
58}
59
60fn main() -> Result<(), Box<dyn Error>> {
61    pel::init();
62
63    // Set up shared state (read large files only once!)
64    let shared_state = Arc::new(SharedState::new()?);
65
66    // Set up the genetic algorithm optimization
67    let codec = FloatCodec::vector(3, 0.1_f32..1.0_f32); // 3 weights for SMA, Ecc, Inc
68    let problem = EngineProblem {
69        objective: radiate::Objective::Multi(vec![Optimize::Minimize, Optimize::Minimize]), // NSGA2 Multi Objective
70        codec: Arc::new(codec.clone()),
71        fitness_fn: Some(Arc::new(move |weights: Vec<f32>| {
72            // Full 60 days propagation for evaluating the actual performance, but running fast due to shared state
73            let (prop_usage, penalty) =
74                evaluate_weights(&weights, 60.0, shared_state.clone()).unwrap_or((1e6, 1e6));
75            Score::from(vec![prop_usage as f32, penalty as f32])
76        })),
77        raw_fitness_fn: None,
78    };
79
80    let mut engine = GeneticEngine::<FloatChromosome<f32>, Vec<f32>>::builder()
81        .codec(codec)
82        .executor(Executor::FixedSizedWorkerPool(8))
83        .problem(problem)
84        .population_size(20)
85        .survivor_selector(NSGA2Selector::new())
86        .build();
87
88    let result = engine.run(|generation: &Generation<FloatChromosome<f32>, Vec<f32>>| {
89        let scores = generation.score().as_slice();
90        println!(
91            "[ {:?} ]: Best Score: Prop usage {:.3} kg, Penalty {:.3}",
92            generation.index(),
93            scores[0],
94            scores[1]
95        );
96        generation.index() >= 10 // Max generations
97    });
98
99    let best_as_string = result
100        .value()
101        .iter()
102        .map(|w| w.to_string())
103        .collect::<Vec<String>>()
104        .join(", ");
105    println!("Optimization finished. Best weights: [{}]", best_as_string);
106
107    Ok(())
108}
109
110fn evaluate_weights(
111    weights: &[f32],
112    prop_time_days: f64,
113    state: Arc<SharedState>,
114) -> Result<(f64, f64), Box<dyn Error>> {
115    let ηthresholds: Vec<f64> = weights.iter().map(|w| *w as f64).collect();
116
117    let eme2k = state.almanac.frame_info(EARTH_J2000).unwrap();
118    let epoch = Epoch::from_gregorian_utc_hms(2024, 2, 29, 12, 13, 14);
119
120    let orbit = Orbit::keplerian(24505.9, 0.725, 7.05, 0.0, 0.0, 0.0, epoch, eme2k);
121
122    let sc = Spacecraft::builder()
123        .orbit(orbit)
124        .mass(Mass::from_dry_and_prop_masses(1000.0, 1000.0))
125        .srp(SRPData::from_area(3.0 * 6.0))
126        .thruster(Thruster {
127            isp_s: 4435.0,
128            thrust_N: 0.472,
129        })
130        .mode(GuidanceMode::Thrust)
131        .build();
132
133    let prop_time = prop_time_days * Unit::Day;
134
135    let objectives = &[
136        Objective::within_tolerance(
137            StateParameter::Element(OrbitalElement::SemiMajorAxis),
138            30_000.0,
139            20.0,
140        ),
141        Objective::within_tolerance(
142            StateParameter::Element(OrbitalElement::Eccentricity),
143            0.001,
144            5e-5,
145        ),
146        Objective::within_tolerance(
147            StateParameter::Element(OrbitalElement::Inclination),
148            0.05,
149            1e-2,
150        ),
151    ];
152
153    // let kluever_ctrl = Kluever::from_max_eclipse(objectives, &weights_f64, 0.2);
154    let ctrl = Ruggiero::from_ηthresholds(objectives, &ηthresholds, sc)?;
155
156    let mut orbital_dyn = OrbitalDynamics::point_masses(vec![MOON, SUN]);
157    orbital_dyn.accel_models.push(state.harmonics.clone());
158
159    let sc_dynamics = SpacecraftDynamics::from_model(orbital_dyn, state.srp_dyn.clone())
160        .with_guidance_law(ctrl.clone());
161
162    let (final_state, _traj) = Propagator::rk89(
163        sc_dynamics.clone(),
164        IntegratorOptions::builder()
165            .min_step(10.0_f64.seconds())
166            .tolerance(1e-8)
167            .error_ctrl(ErrorControl::RSSCartesianStep)
168            .build(),
169    )
170    .with(sc, state.almanac.clone())
171    .for_duration_with_traj(prop_time)?;
172
173    let prop_usage = sc.mass.prop_mass_kg - final_state.mass.prop_mass_kg;
174
175    let mut penalty = 0.0;
176    for obj in objectives {
177        let (achieved, error) = obj.assess(&final_state)?;
178        if !achieved {
179            penalty += error.abs();
180        }
181    }
182
183    info!("{ηthresholds:?} -> {prop_usage:.3} kg\tpenalty = {penalty:.3}");
184
185    Ok((prop_usage, penalty * 1000.0))
186}