Update trajectory generator

src/sampling.rs

@@ -5,7 +5,7 @@ use crate::{
 use rand::SeedableRng;
 use rand_chacha::ChaCha8Rng;
 
-trait Sampler: Iterator {
+pub trait Sampler: Iterator {
     fn reset(&mut self);
 }
 
@@ -84,6 +84,12 @@ impl<'a, T: Network> Iterator for ForwardSampler<'a, T> {
             )
             .unwrap();
 
+        self.next_transitions[next_node_transition] = None;
+        
+        for child in self.net.get_children_set(next_node_transition) {
+            self.next_transitions[child] = None;
+        }
+
 
         Some((ret_time, ret_state))
     }

src/tools.rs

@@ -3,6 +3,7 @@ use rand_chacha::rand_core::SeedableRng;
 use rand_chacha::ChaCha8Rng;
 
 use crate::params::ParamsTrait;
+use crate::sampling::{Sampler, ForwardSampler};
 use crate::{network, params};
 
 pub struct Trajectory {
@@ -61,114 +62,28 @@ pub fn trajectory_generator<T: network::Network>(
     let mut trajectories: Vec<Trajectory> = Vec::new();
 
     //Random Generator object
-    let mut rng: ChaCha8Rng = match seed {
-        //If a seed is present use it to initialize the random generator.
-        Some(seed) => SeedableRng::seed_from_u64(seed),
-        //Otherwise create a new random generator using the method `from_entropy`
-        None => SeedableRng::from_entropy(),
-    };
 
+    let mut sampler = ForwardSampler::new(net, seed);
     //Each iteration generate one trajectory
     for _ in 0..n_trajectories {
-        //Current time of the sampling process
-        let mut t = 0.0;
         //History of all the moments in which something changed
         let mut time: Vec<f64> = Vec::new();
         //Configuration of the process variables at time t initialized with an uniform
         //distribution.
-        let mut current_state: Vec<params::StateType> = net
-            .get_node_indices()
-            .map(|x| net.get_node(x).get_random_state_uniform(&mut rng))
-            .collect();
-        //History of all the configurations of the process variables.
-        let mut events: Vec<Array1<usize>> = Vec::new();
-        //Vector containing to time to the next transition for each variable.
-        let mut next_transitions: Vec<Option<f64>> =
-            net.get_node_indices().map(|_| Option::None).collect();
+        let mut events: Vec<Vec<params::StateType>> = Vec::new();
 
-        //Add the starting time for the trajectory.
-        time.push(t.clone());
-        //Add the starting configuration of the trajectory.
-        events.push(
-            current_state
-                .iter()
-                .map(|x| match x {
-                    params::StateType::Discrete(state) => state.clone(),
-                })
-                .collect(),
-        );
+        //Current Time and Current State
+        let (mut t, mut current_state) = sampler.next().unwrap();
         //Generate new samples until ending time is reached.
         while t < t_end {
-            //Generate the next transition time for each uninitialized variable.
-            for (idx, val) in next_transitions.iter_mut().enumerate() {
-                if let None = val {
-                    *val = Some(
-                        net.get_node(idx)
-                            .get_random_residence_time(
-                                net.get_node(idx).state_to_index(&current_state[idx]),
-                                net.get_param_index_network(idx, &current_state),
-                                &mut rng,
-                            )
-                            .unwrap()
-                            + t,
-                    );
-                }
-            }
-
-            //Get the variable with the smallest transition time.
-            let next_node_transition = next_transitions
-                .iter()
-                .enumerate()
-                .min_by(|x, y| x.1.unwrap().partial_cmp(&y.1.unwrap()).unwrap())
-                .unwrap()
-                .0;
-            //Check if the next transition take place after the ending time.
-            if next_transitions[next_node_transition].unwrap() > t_end {
-                break;
+            time.push(t);
+            events.push(current_state);
+            (t, current_state) = sampler.next().unwrap();
         }
-            //Get the time in which the next transition occurs.
-            t = next_transitions[next_node_transition].unwrap().clone();
-            //Add the transition time to next
-            time.push(t.clone());
 
-            //Compute the new state of the transitioning variable.
-            current_state[next_node_transition] = net
-                .get_node(next_node_transition)
-                .get_random_state(
-                    net.get_node(next_node_transition)
-                        .state_to_index(&current_state[next_node_transition]),
-                    net.get_param_index_network(next_node_transition, &current_state),
-                    &mut rng,
-                )
-                .unwrap();
+        current_state = events.last().unwrap().clone();
+        events.push(current_state);
 
-            //Add the new state to events
-            events.push(Array::from_vec(
-                current_state
-                    .iter()
-                    .map(|x| match x {
-                        params::StateType::Discrete(state) => state.clone(),
-                    })
-                    .collect(),
-            ));
-            //Reset the next transition time for the transitioning node.
-            next_transitions[next_node_transition] = None;
-
-            //Reset the next transition time for each child of the transitioning node.
-            for child in net.get_children_set(next_node_transition) {
-                next_transitions[child] = None
-            }
-        }
-
-        //Add current_state as last state.
-        events.push(
-            current_state
-                .iter()
-                .map(|x| match x {
-                    params::StateType::Discrete(state) => state.clone(),
-                })
-                .collect(),
-        );
         //Add t_end as last time.
         time.push(t_end.clone());
 
@@ -176,11 +91,18 @@ pub fn trajectory_generator<T: network::Network>(
         trajectories.push(Trajectory::new(
             Array::from_vec(time),
             Array2::from_shape_vec(
-                (events.len(), current_state.len()),
-                events.iter().flatten().cloned().collect(),
+                (events.len(), events.last().unwrap().len()),
+                events
+                    .iter()
+                    .flatten()
+                    .map(|x| match x {
+                        params::StateType::Discrete(x) => x.clone(),
+                    })
+                    .collect(),
             )
             .unwrap(),
         ));
+        sampler.reset();
     }
     //Return a dataset object with the sampled trajectories.
     Dataset::new(trajectories)

tests/parameter_learning.rs

@@ -427,7 +427,7 @@ fn learn_mixed_discrete_cim<T: ParameterLearning>(pl: T) {
                 [0.8, 0.6, 0.2, -1.6]
             ],
         ]),
-        0.1
+        0.2
     ));
 }