diff --git a/Cargo.toml b/Cargo.toml index 2c673fa5..5cc0b195 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,19 +33,21 @@ crossterm = "0.27.0" tokio = { version = "1.32.0", features = ["sync", "rt"] } ndarray-csv = "0.5.2" rawpointer = "0.2.1" -argmin = "0.8.1" +argmin = { version = "0.8.1", features = [] } itertools = "0.11.0" faer-core = { version = "0.14.0", features = [] } # faer-lu = "0.9" faer-qr = "0.14.0" # faer-cholesky = "0.9" # faer-svd = "0.9" +argmin-math = { version = "0.3.0", features = ["ndarray_v0_15-nolinalg-serde"] } dyn-stack = "0.10.0" faer = { version = "0.14.0", features = ["nalgebra", "ndarray"] } tracing = "0.1.40" tracing-subscriber = { version = "0.3.17", features = ["env-filter", "fmt", "time"] } chrono = "0.4" + [profile.release] codegen-units = 1 opt-level = 3 diff --git a/examples/two_eq_lag/config.toml b/examples/two_eq_lag/config.toml index ab6a8f6e..4c07a747 100644 --- a/examples/two_eq_lag/config.toml +++ b/examples/two_eq_lag/config.toml @@ -5,7 +5,7 @@ log_out = "log/two_eq_lag.log" [config] cycles = 1000 -engine = "NPAG" +engine = "NPOD" init_points = 1000 seed = 22 tui = true diff --git a/src/algorithms.rs b/src/algorithms.rs index 872fcaa4..ebb72429 100644 --- a/src/algorithms.rs +++ b/src/algorithms.rs @@ -6,12 +6,14 @@ use simulation::predict::{Engine, Predict}; use tokio::sync::mpsc; mod npag; +mod npod; mod postprob; -pub enum Type { - NPAG, - POSTPROB, -} +// pub enum Type { +// NPAG, +// NPOD, +// POSTPROB, +// } pub trait Algorithm { fn fit(&mut self) -> NPResult; @@ -47,6 +49,15 @@ where tx, settings, )), + "NPOD" => Box::new(npod::NPOD::new( + engine, + ranges, + theta, + scenarios, + settings.parsed.error.poly, + tx, + settings, + )), "POSTPROB" => Box::new(postprob::POSTPROB::new( engine, theta, diff --git a/src/algorithms/npag.rs b/src/algorithms/npag.rs index ef95234a..0472b7df 100644 --- a/src/algorithms/npag.rs +++ b/src/algorithms/npag.rs @@ -13,6 +13,7 @@ use crate::{ simulation::predict::{sim_obs, Predict}, }, tui::ui::Comm, + }; use ndarray::{Array, Array1, Array2, Axis}; diff --git a/src/algorithms/npod.rs b/src/algorithms/npod.rs new file mode 100644 index 00000000..1a1597fa --- /dev/null +++ b/src/algorithms/npod.rs @@ -0,0 +1,327 @@ +use crate::prelude::{ + algorithms::Algorithm, + condensation::prune::prune, + datafile::Scenario, + evaluation::sigma::{ErrorPoly, ErrorType}, + ipm, + optimization::d_optimizer::SppOptimizer, + output::NPResult, + output::{CycleLog, NPCycle}, + prob, qr, + settings::run::Data, + simulation::predict::Engine, + simulation::predict::{sim_obs, Predict}, +}; +use ndarray::parallel::prelude::*; +use ndarray::{Array, Array1, Array2, Axis}; +use ndarray_stats::{DeviationExt, QuantileExt}; +use tokio::sync::mpsc::UnboundedSender; + +const THETA_D: f64 = 1e-4; +const THETA_F: f64 = 1e-2; + +pub struct NPOD +where + S: Predict + std::marker::Sync + Clone, +{ + engine: Engine, + ranges: Vec<(f64, f64)>, + psi: Array2, + theta: Array2, + lambda: Array1, + w: Array1, + last_objf: f64, + objf: f64, + cycle: usize, + gamma_delta: f64, + gamma: f64, + error_type: ErrorType, + converged: bool, + cycle_log: CycleLog, + cache: bool, + scenarios: Vec, + c: (f64, f64, f64, f64), + tx: UnboundedSender, + settings: Data, +} + +impl Algorithm for NPOD +where + S: Predict + std::marker::Sync + Clone, +{ + fn fit(&mut self) -> NPResult { + self.run() + } + fn to_npresult(&self) -> NPResult { + NPResult::new( + self.scenarios.clone(), + self.theta.clone(), + self.psi.clone(), + self.w.clone(), + self.objf, + self.cycle, + self.converged, + self.settings.clone(), + ) + } +} + +impl NPOD +where + S: Predict + std::marker::Sync + Clone, +{ + /// Creates a new NPOD instance. + /// + /// # Parameters + /// + /// - `sim_eng`: An instance of the prediction engine. + /// - `ranges`: A vector of value ranges for each parameter. + /// - `theta`: An initial parameter matrix. + /// - `scenarios`: A vector of scenarios. + /// - `c`: A tuple containing coefficients for the error polynomial. + /// - `tx`: An unbounded sender for communicating progress. + /// - `settings`: Data settings and configurations. + /// + /// # Returns + /// + /// Returns a new `NPOD` instance. + pub fn new( + sim_eng: Engine, + ranges: Vec<(f64, f64)>, + theta: Array2, + scenarios: Vec, + c: (f64, f64, f64, f64), + tx: UnboundedSender, + settings: Data, + ) -> Self + where + S: Predict + std::marker::Sync, + { + Self { + engine: sim_eng, + ranges, + psi: Array2::default((0, 0)), + theta, + lambda: Array1::default(0), + w: Array1::default(0), + last_objf: -1e30, + objf: f64::INFINITY, + cycle: 1, + gamma_delta: 0.1, + gamma: settings.parsed.error.value, + error_type: match settings.parsed.error.class.as_str() { + "additive" => ErrorType::Add, + "proportional" => ErrorType::Prop, + _ => panic!("Error type not supported"), + }, + converged: false, + cycle_log: CycleLog::new(&settings.computed.random.names), + cache: settings.parsed.config.cache.unwrap_or(false), + tx, + settings, + scenarios, + c, + } + } + + fn optim_gamma(&mut self) { + //Gam/Lam optimization + // TODO: Move this to e.g. /evaluation/error.rs + let gamma_up = self.gamma * (1.0 + self.gamma_delta); + let gamma_down = self.gamma / (1.0 + self.gamma_delta); + let ypred = sim_obs(&self.engine, &self.scenarios, &self.theta, self.cache); + let psi_up = prob::calculate_psi( + &ypred, + &self.scenarios, + &ErrorPoly { + c: self.c, + gl: gamma_up, + e_type: &self.error_type, + }, + ); + let psi_down = prob::calculate_psi( + &ypred, + &self.scenarios, + &ErrorPoly { + c: self.c, + gl: gamma_down, + e_type: &self.error_type, + }, + ); + let (lambda_up, objf_up) = match ipm::burke(&psi_up) { + Ok((lambda, objf)) => (lambda, objf), + Err(err) => { + //todo: write out report + panic!("Error in IPM: {:?}", err); + } + }; + let (lambda_down, objf_down) = match ipm::burke(&psi_down) { + Ok((lambda, objf)) => (lambda, objf), + Err(err) => { + //todo: write out report + panic!("Error in IPM: {:?}", err); + } + }; + if objf_up > self.objf { + self.gamma = gamma_up; + self.objf = objf_up; + self.gamma_delta *= 4.; + self.lambda = lambda_up; + self.psi = psi_up; + } + if objf_down > self.objf { + self.gamma = gamma_down; + self.objf = objf_down; + self.gamma_delta *= 4.; + self.lambda = lambda_down; + self.psi = psi_down; + } + self.gamma_delta *= 0.5; + if self.gamma_delta <= 0.01 { + self.gamma_delta = 0.1; + } + } + + pub fn run(&mut self) -> NPResult { + while (self.last_objf - self.objf).abs() > THETA_F { + self.last_objf = self.objf; + // log::info!("Cycle: {}", cycle); + // psi n_sub rows, nspp columns + let cache = if self.cycle == 1 { false } else { self.cache }; + let ypred = sim_obs(&self.engine, &self.scenarios, &self.theta, cache); + + self.psi = prob::calculate_psi( + &ypred, + &self.scenarios, + &ErrorPoly { + c: self.c, + gl: self.gamma, + e_type: &self.error_type, + }, + ); + (self.lambda, _) = match ipm::burke(&self.psi) { + Ok((lambda, objf)) => (lambda, objf), + Err(err) => { + //todo: write out report + panic!("Error in IPM: {:?}", err); + } + }; + + let mut keep = Vec::::new(); + for (index, lam) in self.lambda.iter().enumerate() { + if *lam > self.lambda.max().unwrap() / 1000_f64 { + keep.push(index); + } + } + + self.theta = self.theta.select(Axis(0), &keep); + self.psi = self.psi.select(Axis(1), &keep); + + //Rank-Revealing Factorization + let (r, perm) = qr::calculate_r(&self.psi); + + let mut keep = Vec::::new(); + //The minimum between the number of subjects and the actual number of support points + let lim_loop = self.psi.nrows().min(self.psi.ncols()); + for i in 0..lim_loop { + let test = norm_zero(&r.column(i).to_owned()); + let ratio = r.get((i, i)).unwrap() / test; + if ratio.abs() >= 1e-8 { + keep.push(*perm.get(i).unwrap()); + } + } + log::info!( + "QR decomp, cycle {}, kept: {}, thrown {}", + self.cycle, + keep.len(), + self.psi.ncols() - keep.len() + ); + self.theta = self.theta.select(Axis(0), &keep); + self.psi = self.psi.select(Axis(1), &keep); + + (self.lambda, self.objf) = match ipm::burke(&self.psi) { + Ok((lambda, objf)) => (lambda, objf), + Err(err) => { + //todo: write out report + panic!("Error in IPM: {:?}", err); + } + }; + + self.optim_gamma(); + + let mut state = NPCycle { + cycle: self.cycle, + objf: -2. * self.objf, + delta_objf: (self.last_objf - self.objf).abs(), + nspp: self.theta.shape()[0], + stop_text: "".to_string(), + theta: self.theta.clone(), + gamlam: self.gamma, + }; + self.tx.send(state.clone()).unwrap(); + + // If the objective function decreased, log an error. + // Increasing objf signals instability of model misspecification. + if self.last_objf > self.objf { + log::error!("Objective function decreased"); + } + + self.w = self.lambda.clone(); + let pyl = self.psi.dot(&self.w); + + // Add new point to theta based on the optimization of the D function + let sigma = ErrorPoly { + c: self.c, + gl: self.gamma, + e_type: &self.error_type, + }; + // for spp in self.theta.clone().rows() { + // let optimizer = SppOptimizer::new(&self.engine, &self.scenarios, &sigma, &pyl); + // let candidate_point = optimizer.optimize_point(spp.to_owned()).unwrap(); + // prune(&mut self.theta, candidate_point, &self.ranges, THETA_D); + // } + let mut candididate_points: Vec> = Vec::default(); + for spp in self.theta.clone().rows() { + candididate_points.push(spp.to_owned()); + } + candididate_points.par_iter_mut().for_each(|spp| { + let optimizer = SppOptimizer::new(&self.engine, &self.scenarios, &sigma, &pyl); + let candidate_point = optimizer.optimize_point(spp.to_owned()).unwrap(); + *spp = candidate_point; + }); + for cp in candididate_points { + prune(&mut self.theta, cp, &self.ranges, THETA_D); + } + + // Stop if we have reached maximum number of cycles + if self.cycle >= self.settings.parsed.config.cycles { + log::info!("Maximum number of cycles reached"); + state.stop_text = "No (max cycle)".to_string(); + self.tx.send(state).unwrap(); + break; + } + + // Stop if stopfile exists + if std::path::Path::new("stop").exists() { + log::info!("Stopfile detected - breaking"); + state.stop_text = "No (stopped)".to_string(); + self.tx.send(state).unwrap(); + break; + } + //TODO: the cycle migh break before reaching this point + self.cycle_log + .push_and_write(state, self.settings.parsed.config.pmetrics_outputs.unwrap()); + + self.cycle += 1; + + // log::info!("cycle: {}, objf: {}", self.cycle, self.objf); + // dbg!((self.last_objf - self.objf).abs()); + } + + self.to_npresult() + } +} +fn norm_zero(a: &Array1) -> f64 { + let zeros: Array1 = Array::zeros(a.len()); + a.l2_dist(&zeros).unwrap() +} diff --git a/src/lib.rs b/src/lib.rs index 91f5365a..15ec6e08 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,16 +5,23 @@ pub mod routines { pub mod datafile; pub mod initialization; pub mod optimization { - pub mod expansion; + pub mod d_optimizer; pub mod optim; } pub mod output; + pub mod condensation { + pub mod prune; + } + pub mod expansion { + pub mod adaptative_grid; + } pub mod settings { pub mod run; pub mod simulator; } pub mod evaluation { + pub mod ipm; pub mod prob; pub mod qr; @@ -35,8 +42,10 @@ pub mod prelude { pub use crate::entrypoints::start_with_data; pub use crate::logger; pub use crate::prelude::evaluation::{prob, sigma, *}; + pub use crate::routines::condensation; + pub use crate::routines::expansion::*; pub use crate::routines::initialization::*; - pub use crate::routines::optimization::*; + pub use crate::routines::optimization; pub use crate::routines::simulation::*; pub use crate::routines::*; pub use crate::tui::ui::*; diff --git a/src/routines/condensation/prune.rs b/src/routines/condensation/prune.rs new file mode 100644 index 00000000..597327b8 --- /dev/null +++ b/src/routines/condensation/prune.rs @@ -0,0 +1,19 @@ +use ndarray::{Array1, Array2}; + +pub fn prune( + theta: &mut Array2, + candidate: Array1, + limits: &[(f64, f64)], + min_dist: f64, +) { + for spp in theta.rows() { + let mut dist: f64 = 0.; + for (i, val) in candidate.clone().into_iter().enumerate() { + dist += (val - spp.get(i).unwrap()).abs() / (limits[i].1 - limits[i].0); + } + if dist <= min_dist { + return; + } + } + theta.push_row(candidate.view()).unwrap(); +} diff --git a/src/routines/optimization/expansion.rs b/src/routines/expansion/adaptative_grid.rs similarity index 56% rename from src/routines/optimization/expansion.rs rename to src/routines/expansion/adaptative_grid.rs index f9bc2dbf..f6efb768 100644 --- a/src/routines/optimization/expansion.rs +++ b/src/routines/expansion/adaptative_grid.rs @@ -1,4 +1,6 @@ -use ndarray::{Array, Array1, Array2}; +use ndarray::{Array, Array2}; + +use crate::routines::condensation::prune::prune; pub fn adaptative_grid( theta: &mut Array2, @@ -14,35 +16,17 @@ pub fn adaptative_grid( let mut plus = Array::zeros(spp.len()); plus[j] = l; plus = plus + spp; - evaluate_spp(theta, plus, ranges, min_dist); + prune(theta, plus, ranges, min_dist); // (n_spp, _) = theta.dim(); } if val - l > ranges[j].0 { let mut minus = Array::zeros(spp.len()); minus[j] = -l; minus = minus + spp; - evaluate_spp(theta, minus, ranges, min_dist); + prune(theta, minus, ranges, min_dist); // (n_spp, _) = theta.dim(); } } } theta.to_owned() } - -fn evaluate_spp( - theta: &mut Array2, - candidate: Array1, - limits: &[(f64, f64)], - min_dist: f64, -) { - for spp in theta.rows() { - let mut dist: f64 = 0.; - for (i, val) in candidate.clone().into_iter().enumerate() { - dist += (val - spp.get(i).unwrap()).abs() / (limits[i].1 - limits[i].0); - } - if dist <= min_dist { - return; - } - } - theta.push_row(candidate.view()).unwrap(); -} diff --git a/src/routines/optimization/d_optimizer.rs b/src/routines/optimization/d_optimizer.rs new file mode 100644 index 00000000..a126bd9a --- /dev/null +++ b/src/routines/optimization/d_optimizer.rs @@ -0,0 +1,111 @@ +use argmin::{ + core::{ + observers::{ObserverMode, SlogLogger}, + CostFunction, Error, Executor, + }, + solver::neldermead::NelderMead, +}; +use ndarray::{Array1, Axis}; + +use crate::routines::{ + datafile::Scenario, + simulation::predict::{sim_obs, Engine, Predict}, +}; + +use crate::prelude::{prob, sigma::Sigma}; + +pub struct SppOptimizer<'a, S, P> +where + S: Sigma + Sync, + P: Predict + Sync + Clone, +{ + engine: &'a Engine

, + scenarios: &'a Vec, + sig: &'a S, + pyl: &'a Array1, +} + +impl<'a, S, P> CostFunction for SppOptimizer<'a, S, P> +where + S: Sigma + Sync, + P: Predict + Sync + Clone, +{ + type Param = Array1; + type Output = f64; + fn cost(&self, spp: &Self::Param) -> Result { + let theta = spp.to_owned().insert_axis(Axis(0)); + let ypred = sim_obs(&self.engine, &self.scenarios, &theta, true); + let psi = prob::calculate_psi(&ypred, self.scenarios, self.sig); + if psi.ncols() > 1 { + log::error!("Psi in SppOptimizer has more than one column"); + } + if psi.nrows() != self.pyl.len() { + log::error!( + "Psi in SppOptimizer has {} rows, but spp has {}", + psi.nrows(), + self.pyl.len() + ); + } + let nsub = psi.nrows() as f64; + let mut sum = -nsub; + for (p_i, pyl_i) in psi.iter().zip(self.pyl.iter()) { + sum += nsub + p_i / pyl_i; + } + Ok(-sum) + } +} + +impl<'a, S, P> SppOptimizer<'a, S, P> +where + S: Sigma + Sync, + P: Predict + Sync + Clone, +{ + pub fn new( + engine: &'a Engine

, + scenarios: &'a Vec, + sig: &'a S, + pyl: &'a Array1, + ) -> Self { + Self { + engine, + scenarios, + sig, + pyl, + } + } + pub fn optimize_point(self, spp: Array1) -> Result, Error> { + let simplex = create_initial_simplex(&spp); + let solver = NelderMead::new(simplex).with_sd_tolerance(1e-2)?; + let res = Executor::new(self, solver) + .configure(|state| state.max_iters(5)) + // .add_observer(SlogLogger::term(), ObserverMode::Always) + .run()?; + Ok(res.state.best_param.unwrap()) + } +} + +fn create_initial_simplex(initial_point: &Array1) -> Vec> { + let num_dimensions = initial_point.len(); + let perturbation_percentage = 0.008; + + // Initialize a Vec to store the vertices of the simplex + let mut vertices = Vec::new(); + + // Add the initial point to the vertices + vertices.push(initial_point.to_owned()); + + // Calculate perturbation values for each component + for i in 0..num_dimensions { + let perturbation = if initial_point[i] == 0.0 { + 0.00025 // Special case for components equal to 0 + } else { + perturbation_percentage * initial_point[i] + }; + + let mut perturbed_point = initial_point.to_owned(); + perturbed_point[i] += perturbation; + vertices.push(perturbed_point); + } + + vertices +}