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