#![allow(unused_imports)]
use ordered_float::OrderedFloat;
use polars::prelude::*;
use std::collections::HashMap;
use std::time::Instant;
use sde_sim_rs::filtration::Filtration;
use sde_sim_rs::process::util::parse_equations;
use sde_sim_rs::rng::{PseudoRng, Rng, SobolRng};
use sde_sim_rs::sim::simulate;
fn main() {
let dt: f64 = 0.1;
let t_start: f64 = 0.0;
let t_end: f64 = 100.0;
let scenarios: i32 = 1000;
let initial_values = HashMap::from([("X1".to_string(), 1.0), ("X2".to_string(), 1.0)]);
let equations = [
"dX1 = (0.005 * X1) * dt + (0.02 * X1) * dW1".to_string(),
"dX2 = (0.005 * X2) * dt + (0.02 * X1) * dW1 + (0.01 * X2) * dW2".to_string(),
];
let mut processes = parse_equations(&equations).expect("Failed to parse equations");
let scheme = "runge-kutta"; let rng_scheme = "sobol";
let time_steps: Vec<OrderedFloat<f64>> = (0..)
.map(|i| OrderedFloat(t_start + i as f64 * dt))
.take_while(|t| t.0 <= t_end)
.collect();
let mut filtration = Filtration::new(
time_steps.clone(),
(1..=scenarios).collect(),
processes.iter().map(|p| p.name().clone()).collect(),
ndarray::Array3::<f64>::zeros((time_steps.len(), scenarios as usize, processes.len())),
Some(initial_values),
);
let mut rng: Box<dyn Rng> = if rng_scheme == "sobol" {
Box::new(SobolRng::new(
processes
.iter_mut()
.flat_map(|p| p.incrementors().iter_mut().map(|i| i.name().clone()))
.collect::<Vec<String>>(),
time_steps.clone(),
))
} else {
Box::new(PseudoRng::new(
processes
.iter_mut()
.flat_map(|p| p.incrementors().iter_mut().map(|i| i.name().clone()))
.collect::<Vec<String>>(),
))
};
let before = Instant::now();
println!("Starting simulation...");
simulate(
&mut filtration,
&mut processes,
&time_steps,
&scenarios,
&mut *rng,
scheme,
);
println!(
"Simulation completed in {} seconds.\n",
before.elapsed().as_secs_f64()
);
let df: DataFrame = filtration.to_dataframe();
println!("{}", df);
assert!(before.elapsed().as_secs_f64() > 0.0);
}