update-tx-module

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rs-sim"
-version = "0.9.2"
+version = "0.9.3"
 edition = "2024"
 
 [dependencies]

src/demo_data.rs

@@ -0,0 +1,23 @@
+// For creating the demo_input.parquet file in the data dir...
+use crate::pq;
+
+pub fn build_test_data(){
+  use polars::prelude::*;
+  use ndarray_rand::rand::{rngs::SmallRng, Rng, SeedableRng};
+
+  let n_states:i64 = 100_000;
+  let mut rng = SmallRng::from_entropy();
+
+  let states:Vec<i64> = (0..n_states).map(|_| rng.gen_range(0..40)).collect::<Vec<i64>>();
+  let costs = (0..n_states).map(|_| rng.gen_range(5_000..15_000)).collect::<Vec<i64>>();
+  let uuids = (0..n_states).map(|x| format!("uuid_{x}")).collect::<Vec<String>>();
+
+  let cols = vec![
+      Column::new(PlSmallStr::from_str("uuid"), uuids), 
+      Column::new(PlSmallStr::from_str("step_0"), states), 
+      Column::new(PlSmallStr::from_str("value"), costs),
+  ];
+
+  let df = DataFrame::new(cols).unwrap();
+  pq::write(df, "data/demo_input.parquet").unwrap();
+}

src/main.rs

@@ -1,5 +1,5 @@
 // Discrete Event Model
-// #[allow(dead_code, unused_imports, unused_variables)]
+#[allow(dead_code, unused_imports, unused_variables)]
 mod pq;
 mod sim;
 mod tx;
@@ -9,24 +9,22 @@ fn main() {
 }
 
 fn run() {
-    let working_dir = "./tmp/sim/test-50yr";
-
-    // Survival curve
+    // Setup
+    let working_dir = "./tmp/sim/test-30yr";
     let probabilities = vec![
-        1.0, 0.99, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92, 0.91, 0.9, 0.88, 0.85, 0.82, 0.8, 0.75, 0.7,
-        0.65, 0.6, 0.55, 0.5, 0.4, 0.25, 0.15, 0.1, 0.0,
+        1.0, 0.99, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92, 0.91, 0.9, 0.88, 0.85, 0.82, 0.8, 0.78,
+        0.77, 0.75, 0.735, 0.72, 0.71, 0.7, 0.65, 0.6, 0.55, 0.5, 0.4, 0.25, 0.15, 0.1, 0.0,
     ];
 
-    let init_states = pq::read("data/init_states.parquet").expect("failted to read init states...");
-
+    let init_states = pq::read("data/demo_input.parquet").expect("failted to read init states...");
     let uuids = tx::col_to_vec_str(&init_states, "uuid");
     let states = tx::col_to_vec_i64(&init_states, "step_0");
     let costs = tx::col_to_vec_i64(&init_states, "value");
 
     // Execute
-    let n_steps: i64 = 50;
-    let n_sims: i64 = 200;
-    let constraints = vec![200_000_000; n_steps as usize];
+    let n_steps: i64 = 30;
+    let n_sims: i64 = 1000;
+    let constraints = vec![175_000_000; n_steps as usize];
 
     if !std::path::Path::new(&working_dir).exists() {
         std::fs::create_dir(&working_dir).expect("failed to create working dir...");

src/pq.rs

@@ -10,8 +10,8 @@ pub fn read(path: &str) -> PolarsResult<DataFrame> {
 #[test]
 fn test_read() {
     // Ensure expected columns are in table
-    let test = vec!["uuid", "value", "step_0"];
-    let res: Vec<String> = read("./data/init_states.parquet")
+    let test = vec!["uuid", "step_0", "value"];
+    let res: Vec<String> = read("./data/demo_input.parquet")
         .expect("failed to read file")
         .get_column_names()
         .iter()
@@ -40,7 +40,7 @@ fn test_write() {
         std::fs::create_dir("./tmp").unwrap();
     }
     let test = "./tmp/test_write.parquet";
-    let df = read("./data/init_states.parquet").expect("failed to read file");
+    let df = read("./data/demo_input.parquet").expect("failed to read file");
     write(df, test).expect("failed to write...");
 
     // Confirms existance & remove

src/sim.rs

@@ -191,7 +191,7 @@ fn test_discrete_event() {
         0.65, 0.6, 0.55, 0.5, 0.4, 0.25, 0.15, 0.1, 0.0,
     ];
 
-    let init_states = pq::read("data/init_states.parquet").expect("failted to read init states...");
+    let init_states = pq::read("data/demo_input.parquet").expect("failted to read init states...");
 
     // let uuids = tx::col_to_vec_str(&init_states, "uuid");
     let states = tx::col_to_vec_i64(&init_states, "step_0");
@@ -209,7 +209,7 @@ fn test_execute_event() {
         0.65, 0.6, 0.55, 0.5, 0.4, 0.25, 0.15, 0.1, 0.0,
     ];
 
-    let init_states = pq::read("data/init_states.parquet").expect("failted to read init states...");
+    let init_states = pq::read("data/demo_input.parquet").expect("failted to read init states...");
 
     let uuids = tx::col_to_vec_str(&init_states, "uuid");
     let states = tx::col_to_vec_i64(&init_states, "step_0");

src/tx.rs

@@ -180,17 +180,14 @@ pub fn constrain_event(table: &DataFrame, limit_array: Vec<i64>) -> Result<DataF
     let active_cols: Vec<(usize, String)> = table
         .get_column_names()
         .iter()
-        .filter(|c| c.contains(step_col))
+        // We skip step 0 as the initial state cannot be altered
+        .filter(|c| c.contains(step_col) && ***c != PlSmallStr::from_str("step_0"))
         .enumerate()
-        .map(|(i, c)| (i, c.to_string()))
+        .map(|(i, c)| (i + 1, c.to_string()))
         .collect();
 
     for (idx, col_name) in &active_cols {
-        // We skip step 0 as the initial state cannot be altered
-        if col_name == "step_0" {
-            continue;
-        }
-        // Add random ordering & applied costs columns
+        // Add columns for random ordering & applied costs
         table = table
             .lazy()
             .with_columns([
@@ -215,45 +212,40 @@ pub fn constrain_event(table: &DataFrame, limit_array: Vec<i64>) -> Result<DataF
                 SortMultipleOptions::new().with_order_descending(true),
             )?
             .lazy()
-            .with_columns([col("applied_cost").cum_sum(false).alias("totaliser")])
+            .with_columns([col("applied_cost").cum_sum(true).alias("totaliser")])
             .collect()?;
 
         // Increase age if cost is above set limit
+        let step_limit = lit(limit_array[*idx]);
         table = table
             .lazy()
             .with_columns([when(
                 col(col_name)
                     .eq(lit(0))
-                    .and(col("totaliser").gt(lit(limit_array[*idx])).eq(lit(true))),
+                    .and(col("totaliser").gt(step_limit)),
             )
-            .then(col(&active_cols[idx - 1].1) + lit(1))
+            .then(col(&format!("step_{}", idx - 1)) + lit(1))
             .otherwise(col(col_name))
             .alias("tmp_col")])
             .collect()?;
 
         // Update following years when event is deferred
         for (i, c) in active_cols.iter().rev() {
-            if i > &idx && *i != 0 {
+            if i > idx {
                 table = table
                     .lazy()
-                    .with_column(
-                        when(
-                            col(col_name)
-                                .eq(lit(0))
-                                .and(col("totaliser").gt(lit(limit_array[*idx])).eq(lit(true))),
-                        )
-                        .then(col(&active_cols[i - 1].1))
+                    .with_columns([when(col("tmp_col").neq(col(col_name)))
+                        .then(col(&format!("step_{}", i - 1)))
                         .otherwise(col(c))
-                        .alias(c),
-                    )
+                        .alias(c)])
                     .collect()?;
             }
         }
 
-        // Merge holding col into current col
+        // Update current year
         table = table
             .lazy()
-            .with_column(col("tmp_col").alias(col_name))
+            .with_columns([col("tmp_col").alias(col_name)])
             .collect()?;
 
         // Remove process cols
@@ -303,15 +295,15 @@ fn test_transpose() {
 
 #[test]
 fn test_col_to_vec_i64() {
-    let table = pq::read("./data/init_states.parquet").unwrap();
+    let table = pq::read("./data/demo_input.parquet").unwrap();
     let res = col_to_vec_i64(&table, "step_0");
 
     assert!(res.len() == 100_000);
 }
 
 #[test]
 fn test_col_to_vec_str() {
-    let table = pq::read("./data/init_states.parquet").unwrap();
+    let table = pq::read("./data/demo_input.parquet").unwrap();
     let res = col_to_vec_str(&table, "uuid");
 
     assert!(res.len() == 100_000);