Merge pull request #7 from Aschii85/refactor-and-function-usage

Refactor Coefficient Functions & Performance Optimization (LazyFrame)

Author: Aschii85

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "sde-sim-rs"
-version = "0.4.0"
+version = "0.5.0"
 edition = "2024"
 authors = ["Alexander Schierbeck-Hansen <aschii85@protonmail.com>"]
 description = "Powerful and flexible stochastic differential equation (quasi) Monte-Carlo simulation library written in Rust with Python bindings"
@@ -17,12 +17,18 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 fasteval = "0.2.4"
 lazy_static = "1.5.0"
+lru = "0.16.3"
+nom = "8.0.0"
 ordered-float = "4.2"
-polars = "0.49.1"
-pyo3 = { version = "0.25.0", features = ["extension-module"] }
-pyo3-polars = "0.22.0"
+polars = { version = "0.51.0", features = ["diagonal_concat", "lazy"] }
+pyo3 = { version = "0.25.1", features = ["auto-initialize"], optional = true }
+pyo3-polars = { version = "0.24.0", optional = true }
 rand = "0.9.2"
 rand_chacha = "0.9.0"
 rayon = "1.11.0"
 regex = "1.11.1"
 sobol = "1.0.2"
+
+[features]
+default = []
+python = ["dep:pyo3", "dep:pyo3-polars", "pyo3/extension-module"]

README.md

@@ -32,7 +32,7 @@ Requires Python version >=3.11,<3.14.
 To build the package locally for, you'll first need to compile the Rust package for local development. The project is set up to use `maturin` and `uv`. This command builds the Rust library and creates a Python wheel that can be used directly in your environment.
 
 ```
-maturin develop --uv --release
+maturin develop --uv --release --features python
 ```
 
 After the compilation is complete, you can run the example to see how the library works. This command uses `uv` to execute the Python example script.

examples/example.py

@@ -3,24 +3,23 @@
 import plotly.express as px
 import polars as pl
 
-print(dir(sde_sim_rs))
-
 
 def main():
+    initial_values = {"X0": 0.5, "X1": 100.0, "X2": 0.0}
     df: pl.DataFrame = sde_sim_rs.simulate(
         processes_equations=[
-            "dX1 = ( sin(t) ) * dt",
-            "dX2 = (0.01 * X1) * dW1",
-            "dX3 = (0.005 * X3) * dt + (0.01 * X3) * dW2 + (0.1 * X2 * X3) * dJ1(0.01)",
+            "dX0 = ( 2.0 * (0.5 - X0) ) * dt + ( 0.1 ) * dN1(X0)",
+            "dX1 = ( 0.05 * X1 ) * dt + ( 0.2 * X1 ) * dW1 + ( 0.5 ) * dN1(X0)",
+            "X2 = max(X1 - 100.0, 0.0)",
         ],
-        time_steps=list(np.arange(0.0, 100.0, 0.1)),
-        scenarios=1000,
-        initial_values={"X1": 0.0, "X2": 1.0, "X3": 100.0},
+        time_steps=list(np.arange(0.0, 10.0, 0.01)),
+        scenarios=10_000,
+        initial_values=initial_values,
         rng_method="pseudo",
-        scheme="euler",
+        scheme="runge-kutta",
     )
     print(df)
-    for i in range(1, 4):
+    for i in range(0, len(initial_values)):
         fig = px.line(
             df.filter(pl.col("process_name") == f"X{i}"),
             x="time",
@@ -30,6 +29,17 @@ def main():
             title="Simulated SDE Process",
         )
         fig.show()
+    for i in range(0, len(initial_values)):
+        fig = px.line(
+            df.filter(pl.col("process_name") == f"X{i}")
+            .group_by("time")
+            .agg(pl.col("value").mean())
+            .sort("time"),
+            x="time",
+            y="value",
+            title=f"Mean Simulated SDE Process for X{i}",
+        )
+        fig.show()
 
 
 if __name__ == "__main__":

examples/example.rs

@@ -1,58 +1,56 @@
 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::proc::util::parse_equations;
 use sde_sim_rs::sim::simulate;
+use std::collections::HashMap;
+use std::time::Instant;
 
 fn main() {
-    // Simulation Parameters
-    let dt: f64 = 0.1;
-    let t_start: f64 = 0.0;
-    let t_end: f64 = 100.0;
-    let scenarios: i32 = 10000;
-    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 + (1) * dJ1(0.5)"
-            .to_string(),
+    // ────── configuration ──────
+    let initial_values = HashMap::from([("X1".to_string(), 100.0), ("X2".to_string(), 0.0)]);
+
+    let processes_equations = vec![
+        "dX1 = ( sin(t) ) * dt + (0.01 * X1) * dW1 + (0.001 * X1) * dN1(0.5 * cos(t))".to_string(),
+        "X2 = max(X1 - 100.0, 0.0)".to_string(),
     ];
-    let scheme = "runge-kutta"; // "euler" or "runge-kutta"
-    let rng_method = "sobol"; // "pseudo" or "sobol"
 
-    // 1. Prepare Time Steps
-    let time_steps: Vec<OrderedFloat<f64>> = (0..)
-        .map(|i| OrderedFloat(t_start + i as f64 * dt))
-        .take_while(|t| t.0 <= t_end)
+    let scheme = "euler"; // other valid value: "runge-kutta"
+    let rng_method = "pseudo"; // other valid value: "sobol"
+    let scenarios: u64 = 10_000;
+
+    // build a uniformly spaced time vector, identical to what the Python
+    // wrapper accepts as `time_steps: Vec<f64>`.
+    let dt = 0.1;
+    let t_start = 0.0;
+    let t_end = 100.0;
+    let time_steps: Vec<f64> = (0..)
+        .map(|i| t_start + i as f64 * dt)
+        .take_while(|t| *t <= t_end)
         .collect();
 
-    // 2. Parse equations
-    let universe =
-        parse_equations(&equations, time_steps.clone()).expect("Failed to parse equations");
-
-    // 3. Initialize Filtration
-    let mut filtration = Filtration::new(
-        universe,
-        time_steps.clone(),
-        (1..=scenarios).collect(),
-        Some(initial_values),
-    );
-
-    // Run Simulation
-    let before = Instant::now();
-    println!("Starting simulation with {} RNG...", rng_method);
-    simulate(&mut filtration, scheme, rng_method);
-
-    let duration = before.elapsed();
-    println!(
-        "Simulation completed in {:.4} seconds.\n",
-        duration.as_secs_f64()
-    );
-
-    let df: DataFrame = filtration.to_dataframe();
-    println!("{}", df);
-
-    assert!(duration.as_secs_f64() > 0.0);
+    // convert the floats to `OrderedFloat` for internal use
+    let ordered_steps: Vec<OrderedFloat<f64>> =
+        time_steps.iter().copied().map(OrderedFloat).collect();
+
+    // parse the equations into a ProcessUniverse (same work done in Python)
+    let universe = parse_equations(&processes_equations, ordered_steps.clone())
+        .expect("failed to parse process equations");
+
+    // run the actual simulation; this mirrors the body of `simulate_py`
+    let start = Instant::now();
+    println!("running {} scenarios with {} rng...", scenarios, rng_method);
+    let lf = simulate(
+        &universe,
+        ordered_steps.clone(),
+        initial_values.clone(),
+        scenarios,
+        scheme,
+        rng_method,
+    )
+    .expect("failed to run simulation");
+    let df = lf.collect().expect("failed to collect results");
+    let elapsed = start.elapsed();
+    println!("completed in {:.3}s", elapsed.as_secs_f64());
+
+    // print a small portion of the output frame
+    println!("{:#?}", df.head(Some(10)));
 }

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "sde-sim-rs"
-requires-python = ">=3.11,<3.14"
+requires-python = ">=3.11,<3.15"
 authors = [
   {name = "Alexander Schierbeck-Hansen", email = "aschii85@protonmail.com"},
 ]

python/sde_sim_rs/__init__.pyi

@@ -3,15 +3,14 @@ from typing import Literal
 
 import polars as pl
 
-
 def simulate(
     processes_equations: Sequence[str],
     time_steps: Sequence[float],
     scenarios: int,
     initial_values: Mapping[str, float],
     rng_method: Literal["pseudo", "sobol"] = "pseudo",
     scheme: Literal["euler", "runge-kutta"] = "euler",
-) -> pl.DataFrame: 
+) -> pl.DataFrame:
     """
     Simulates stochastic differential equations (SDEs) using the specified methods.
 
@@ -45,12 +44,14 @@ def simulate(
             method. Defaults to "euler".
 
     Returns:
-        A Polars DataFrame containing the simulated values. The DataFrame has a
-        `time` column and columns for each process, with each row corresponding
-        to a specific time step and scenario.
+        A Polars DataFrame containing the simulated values. The DataFrame is
+        "long"/tidy: every row represents a single `(scenario, time, process)`
+        triple and the associated value.  In other words, the `scenario`
+        dimension has already been appended, which makes it easy to group or
+        aggregate across paths using standard Polars operations.
 
     Raises:
         ValueError: If the process equations are malformed or if initial values
             are missing for any process.
     """
-    ...
+    ...