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::{Almanac, metaload::MetaFile},
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    Spacecraft,
17    cosmic::{GuidanceMode, Mass, MetaAlmanac, Orbit, SRPData},
18    dynamics::{
19        GravityField, OrbitalDynamics, SolarPressure, SpacecraftDynamics,
20        guidance::{Ruggiero, Thruster},
21    },
22    io::gravity::GravityFieldData,
23    md::prelude::{Objective, OrbitalElement, StateParameter},
24    propagators::{ErrorControl, IntegratorOptions, Propagator},
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::new(GravityFieldData::from_cof(
47            &jgm3_meta.uri,
48            4,
49            4,
50            true,
51            almanac.frame_info(IAU_EARTH_FRAME)?,
52        )?);
53        let srp_dyn = SolarPressure::default_flux(EARTH_J2000, &almanac)?;
54
55        Ok(Self {
56            almanac,
57            harmonics,
58            srp_dyn,
59        })
60    }
61}
62
63fn main() -> Result<(), Box<dyn Error>> {
64    pel::init();
65
66    /* let (prop_usage_kg, penalty) = evaluate_weights(
67        &[0.22033301, 0.8512096, 0.49421895],
68        60.0,
69        Arc::new(SharedState::new()?),
70    )
71    .unwrap();
72
73    println!("Best weight prop usage = {prop_usage_kg:.3} kg \t penalty = {penalty:.3}"); */
74
75    // Set up shared state (read large files only once!)
76    let shared_state = Arc::new(SharedState::new()?);
77
78    // Set up the genetic algorithm optimization
79    let codec = FloatCodec::vector(3, 0.1_f32..1.0_f32); // 3 weights for SMA, Ecc, Inc
80    let problem = EngineProblem {
81        objective: radiate::Objective::Multi(vec![Optimize::Minimize, Optimize::Minimize]), // NSGA2 Multi Objective
82        codec: Arc::new(codec),
83        fitness_fn: Some(Arc::new(move |weights: Vec<f32>| {
84            // Full 60 days propagation for evaluating the actual performance, but running fast due to shared state
85            let (prop_usage, penalty) =
86                evaluate_weights(&weights, 60.0, shared_state.clone()).unwrap_or((1e6, 1e6));
87            Score::from(vec![prop_usage as f32, penalty as f32])
88        })),
89        raw_fitness_fn: None,
90    };
91
92    let mut engine = GeneticEngine::<FloatChromosome<f32>, Vec<f32>>::builder()
93        .population_size(20)
94        .parallel()
95        .multi_objective(vec![Optimize::Minimize, Optimize::Minimize])
96        .problem(problem)
97        .survivor_selector(NSGA2Selector::new())
98        .build();
99
100    // Wrap the engine with the UI
101    let final_generation = engine.run(|generation: &Generation<FloatChromosome<f32>, Vec<f32>>| {
102        let scores = generation.score().as_slice();
103        println!(
104            "[ {:?} ]: Best Score: Prop usage {:.3} kg, Penalty {:.3}",
105            generation.index(),
106            scores[0],
107            scores[1]
108        );
109        generation.index() >= 5
110    });
111
112    let best_weights = final_generation
113        .value()
114        .iter()
115        .map(|w| format!("W: = {w}"))
116        .collect::<Vec<String>>()
117        .join(", ");
118    let best_score = final_generation
119        .score()
120        .iter()
121        .enumerate()
122        .map(|(i, w)| format!("S[{i}]: = {w}"))
123        .collect::<Vec<String>>()
124        .join(", ");
125    println!("Optimization finished. Best weights: [{best_weights}] -> Best score: [{best_score}]");
126
127    // Evaluate these weights.
128    let best_weights: Vec<f32> = final_generation.value().to_vec();
129
130    let (prop_usage_kg, penalty) =
131        evaluate_weights(&best_weights, 60.0, Arc::new(SharedState::new()?)).unwrap();
132
133    println!("Best weight prop usage = {prop_usage_kg:.3} kg \t penalty = {penalty:.3}");
134
135    Ok(())
136}
137
138fn evaluate_weights(
139    weights: &[f32],
140    prop_time_days: f64,
141    state: Arc<SharedState>,
142) -> Result<(f64, f64), Box<dyn Error>> {
143    let ηthresholds: Vec<f64> = weights.iter().map(|w| *w as f64).collect();
144
145    let eme2k = state.almanac.frame_info(EARTH_J2000).unwrap();
146    let epoch = Epoch::from_gregorian_utc_hms(2024, 2, 29, 12, 13, 14);
147
148    let orbit = Orbit::keplerian(24505.9, 0.725, 7.05, 0.0, 0.0, 0.0, epoch, eme2k);
149
150    let sc = Spacecraft::builder()
151        .orbit(orbit)
152        .mass(Mass::from_dry_and_prop_masses(1000.0, 1000.0))
153        .srp(SRPData::from_area(3.0 * 6.0))
154        .thruster(Thruster {
155            isp_s: 4435.0,
156            thrust_N: 0.472,
157        })
158        .mode(GuidanceMode::Thrust)
159        .build();
160
161    let prop_time = prop_time_days * Unit::Day;
162
163    let objectives = &[
164        Objective::within_tolerance(
165            StateParameter::Element(OrbitalElement::SemiMajorAxis),
166            30_000.0,
167            20.0,
168        ),
169        Objective::within_tolerance(
170            StateParameter::Element(OrbitalElement::Eccentricity),
171            0.001,
172            5e-5,
173        ),
174        Objective::within_tolerance(
175            StateParameter::Element(OrbitalElement::Inclination),
176            0.05,
177            1e-2,
178        ),
179    ];
180
181    let ctrl = Ruggiero::from_ηthresholds(objectives, &ηthresholds, sc)?;
182
183    let mut orbital_dyn = OrbitalDynamics::point_masses(vec![MOON, SUN]);
184    orbital_dyn.accel_models.push(state.harmonics.clone());
185
186    let sc_dynamics = SpacecraftDynamics::from_model(orbital_dyn, state.srp_dyn.clone())
187        .with_guidance_law(ctrl.clone());
188
189    let (final_state, _traj) = Propagator::rk89(
190        sc_dynamics.clone(),
191        IntegratorOptions::builder()
192            .min_step(10.0_f64.seconds())
193            .tolerance(1e-8)
194            .error_ctrl(ErrorControl::RSSCartesianStep)
195            .build(),
196    )
197    .with(sc, state.almanac.clone())
198    .for_duration_with_traj(prop_time)?;
199
200    let prop_usage = sc.mass.prop_mass_kg - final_state.mass.prop_mass_kg;
201
202    let mut penalty = 0.0;
203    for obj in objectives {
204        let (achieved, error) = obj.assess(&final_state)?;
205        if !achieved {
206            penalty += error.abs();
207        }
208        info!("{obj} error: {error:.3}, achieved? {achieved}");
209    }
210
211    info!("{ηthresholds:?} -> {prop_usage:.3} kg\tpenalty = {penalty:.3}");
212
213    Ok((prop_usage, penalty * 1000.0))
214}