Update stochastic_new.rs

anthias-labs · Jul 28, 2023 · 71be366 · 71be366
1 parent a0f9cf2
commit 71be366
Showing 1 changed file with 13 additions and 21 deletions.
diff --git a/core/src/math/stochastic_new.rs b/core/src/math/stochastic_new.rs
@@ -5,11 +5,9 @@ use anyhow::Result;
 /// Type enum for process 
 pub enum StochasticProcessType {
     /// Brownian motion
-    BrownianMotion,
+    BrownianMotion(BrownianMotion),
     /// Ornstein-Uhlenbeck
-    OrnsteinUhlenbeck,
-    /// Poisson
-    Poisson,
+    OrnsteinUhlenbeck(OrnsteinUhlenbeck),
 }
 /// Struct for all processes init parameters.
 pub struct EulerMaruyamaInput {
@@ -21,25 +19,19 @@ pub struct EulerMaruyamaInput {
     pub parallel: bool, // run in parallel or not (recommended for > 1000 paths)
 }
 
-pub fn new_procces(proccess_type: StochasticProcessType, config: EulerMaruyamaInput) -> Result<()> {
+/// Create new process and run euler maruyama.
+pub fn new_procces(proccess_type: StochasticProcessType, config: EulerMaruyamaInput) -> Result<Trajectories> {
 
-    match proccess_type {
-        StochasticProcessType::BrownianMotion => {
-            let brownian_motion = BrownianMotion::new();
-            let trajectories = brownian_motion.euler_maruyama(config.x_0, config.t_0, config.t_n, config.n_steps, config.m_paths, config.parallel);
+    let trajectories: Trajectories = match proccess_type {
+        StochasticProcessType::BrownianMotion(process) => {
+            process.euler_maruyama(config.x_0, config.t_0, config.t_n, config.n_steps, config.m_paths, config.parallel)
+            },
+        StochasticProcessType::OrnsteinUhlenbeck(process) => {
+            process.euler_maruyama(config.x_0, config.t_0, config.t_n, config.n_steps, config.m_paths, config.parallel)
         },
-        StochasticProcessType::OrnsteinUhlenbeck => {
-            let ornstein_uhlenbeck = OrnsteinUhlenbeck::new();
-            let trajectories = ornstein_uhlenbeck.euler_maruyama(config.x_0, config.t_0, config.t_n, config.n_steps, config.m_paths, config.parallel);
-            println!("{:?}", trajectories);
-        },
-        StochasticProcessType::Poisson => {
-            let poisson = Poisson::new();
-            let trajectories = poisson.euler_maruyama(config.x_0, config.t_0, config.t_n, config.n_steps, config.m_paths, config.parallel);
-            println!("{:?}", trajectories);
-        },
-    }
-    Ok(())
+    };
+
+    Ok(trajectories)
 }
 
 // impl StochasticProcess {