Wasm ir

Cargo.toml

@@ -14,7 +14,7 @@ exa = ["libloading"]
 [dependencies]
 cached = { version = "0.56.0" }
 csv = "1.3.0"
-diffsol = "=0.7.0"
+diffsol = { version = "=0.7.0" }
 libloading = { version = "0.8.6", optional = true, features = [] }
 nalgebra = "0.34.1"
 ndarray = { version = "0.16.1", features = ["rayon"] }
@@ -28,6 +28,7 @@ thiserror = "2.0.11"
 argmin = "0.11.0"
 argmin-math = "0.5.1"
 tracing = "0.1.41"
+once_cell = "1.18.0"
 
 [dev-dependencies]
 criterion = { version = "0.7.0", features = ["html_reports"] }
@@ -47,3 +48,8 @@ harness = false
 [[bench]]
 name = "ode"
 harness = false
+
+[[bench]]
+name = "wasm_ode_compare"
+harness = false
+required-features = ["exa"]

benches/wasm_ode_compare.rs

@@ -0,0 +1,182 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use pharmsol::*;
+use std::hint::black_box;
+
+fn example_subject() -> Subject {
+    Subject::builder("1")
+        .infusion(0.0, 500.0, 0, 0.5)
+        .observation(0.5, 1.645776, 0)
+        .observation(1.0, 1.216442, 0)
+        .observation(2.0, 0.4622729, 0)
+        .observation(3.0, 0.1697458, 0)
+        .observation(4.0, 0.06382178, 0)
+        .observation(6.0, 0.009099384, 0)
+        .observation(8.0, 0.001017932, 0)
+        .missing_observation(12.0, 0)
+        .build()
+}
+
+fn regular_ode_predictions(c: &mut Criterion) {
+    let subject = example_subject();
+    let ode = equation::ODE::new(
+        |x, p, _t, dx, _b, rateiv, _cov| {
+            fetch_params!(p, ke, _v);
+            dx[0] = -ke * x[0] + rateiv[0];
+        },
+        |_p, _t, _cov| lag! {},
+        |_p, _t, _cov| fa! {},
+        |_p, _t, _cov, _x| {},
+        |x, p, _t, _cov, y| {
+            fetch_params!(p, _ke, v);
+            y[0] = x[0] / v;
+        },
+        (1, 1),
+    );
+    let params = vec![1.02282724609375, 194.51904296875];
+
+    c.bench_function("regular_ode_predictions", |b| {
+        b.iter(|| {
+            black_box(ode.estimate_predictions(&subject, &params).unwrap());
+        })
+    });
+}
+
+fn wasm_ir_ode_predictions(c: &mut Criterion) {
+    let subject = example_subject();
+
+    // Setup WASM IR model
+    let test_dir = std::env::current_dir().expect("Failed to get current directory");
+    let ir_path = test_dir.join("test_model_ir_bench.pkm");
+
+    let _ir_file = exa_wasm::build::emit_ir::<equation::ODE>(
+        "|x, p, _t, dx, rateiv, _cov| {
+            fetch_params!(p, ke, _v);
+            dx[0] = -ke * x[0] + rateiv[0];
+        }"
+        .to_string(),
+        None,
+        None,
+        Some("|p, _t, _cov, x| { }".to_string()),
+        Some(
+            "|x, p, _t, _cov, y| {
+            fetch_params!(p, _ke, v);
+            y[0] = x[0] / v;
+        }"
+            .to_string(),
+        ),
+        Some(ir_path.clone()),
+        vec!["ke".to_string(), "v".to_string()],
+    )
+    .expect("emit_ir failed");
+
+    let (wasm_ode, _meta, _id) =
+        exa_wasm::interpreter::load_ir_ode(ir_path.clone()).expect("load_ir_ode failed");
+
+    let params = vec![1.02282724609375, 194.51904296875];
+
+    c.bench_function("wasm_ir_ode_predictions", |b| {
+        b.iter(|| {
+            black_box(wasm_ode.estimate_predictions(&subject, &params).unwrap());
+        })
+    });
+
+    // Clean up
+    std::fs::remove_file(ir_path).ok();
+}
+
+fn regular_ode_likelihood(c: &mut Criterion) {
+    let subject = example_subject();
+    let ode = equation::ODE::new(
+        |x, p, _t, dx, _b, rateiv, _cov| {
+            fetch_params!(p, ke, _v);
+            dx[0] = -ke * x[0] + rateiv[0];
+        },
+        |_p, _t, _cov| lag! {},
+        |_p, _t, _cov| fa! {},
+        |_p, _t, _cov, _x| {},
+        |x, p, _t, _cov, y| {
+            fetch_params!(p, _ke, v);
+            y[0] = x[0] / v;
+        },
+        (1, 1),
+    );
+    let params = vec![1.02282724609375, 194.51904296875];
+    let ems = ErrorModels::new()
+        .add(
+            0,
+            ErrorModel::additive(ErrorPoly::new(0.0, 0.05, 0.0, 0.0), 0.0),
+        )
+        .unwrap();
+
+    c.bench_function("regular_ode_likelihood", |b| {
+        b.iter(|| {
+            black_box(
+                ode.estimate_likelihood(&subject, &params, &ems, false)
+                    .unwrap(),
+            );
+        })
+    });
+}
+
+fn wasm_ir_ode_likelihood(c: &mut Criterion) {
+    let subject = example_subject();
+
+    // Setup WASM IR model
+    let test_dir = std::env::current_dir().expect("Failed to get current directory");
+    let ir_path = test_dir.join("test_model_ir_bench_ll.pkm");
+
+    let _ir_file = exa_wasm::build::emit_ir::<equation::ODE>(
+        "|x, p, _t, dx, rateiv, _cov| {
+            fetch_params!(p, ke, _v);
+            dx[0] = -ke * x[0] + rateiv[0];
+        }"
+        .to_string(),
+        None,
+        None,
+        Some("|p, _t, _cov, x| { }".to_string()),
+        Some(
+            "|x, p, _t, _cov, y| {
+            fetch_params!(p, _ke, v);
+            y[0] = x[0] / v;
+        }"
+            .to_string(),
+        ),
+        Some(ir_path.clone()),
+        vec!["ke".to_string(), "v".to_string()],
+    )
+    .expect("emit_ir failed");
+
+    let (wasm_ode, _meta, _id) =
+        exa_wasm::interpreter::load_ir_ode(ir_path.clone()).expect("load_ir_ode failed");
+
+    let params = vec![1.02282724609375, 194.51904296875];
+    let ems = ErrorModels::new()
+        .add(
+            0,
+            ErrorModel::additive(ErrorPoly::new(0.0, 0.05, 0.0, 0.0), 0.0),
+        )
+        .unwrap();
+
+    c.bench_function("wasm_ir_ode_likelihood", |b| {
+        b.iter(|| {
+            black_box(
+                wasm_ode
+                    .estimate_likelihood(&subject, &params, &ems, false)
+                    .unwrap(),
+            );
+        })
+    });
+
+    // Clean up
+    std::fs::remove_file(ir_path).ok();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    regular_ode_predictions(c);
+    wasm_ir_ode_predictions(c);
+    regular_ode_likelihood(c);
+    wasm_ir_ode_likelihood(c);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

examples/bytecode_models.rs

@@ -0,0 +1,108 @@
+use std::env;
+use std::fs;
+// example: emit IR and load via the runtime
+
+fn main() {
+    let tmp = env::temp_dir();
+
+    // Model 1: simple dx assignment
+    let diffeq1 = "|x, p, _t, dx, rateiv, _cov| { dx[0] = -ke * x[0]; }".to_string();
+    let path1 = tmp.join("exa_example_model1.json");
+    let _ = pharmsol::exa_wasm::build::emit_ir::<pharmsol::equation::ODE>(
+        diffeq1,
+        None,
+        None,
+        None,
+        None,
+        Some(path1.clone()),
+        vec!["ke".to_string()],
+    )
+    .expect("emit_ir model1");
+
+    // Model 2: prelude/local and rate
+    let diffeq2 =
+        "|x, p, _t, dx, rateiv, _cov| { ke = 0.5; dx[0] = -ke * x[0] + rateiv[0]; }".to_string();
+    let path2 = tmp.join("exa_example_model2.json");
+    let _ = pharmsol::exa_wasm::build::emit_ir::<pharmsol::equation::ODE>(
+        diffeq2,
+        None,
+        None,
+        None,
+        None,
+        Some(path2.clone()),
+        vec!["ke".to_string()],
+    )
+    .expect("emit_ir model2");
+
+    // Model 3: builtin and ternary
+    let diffeq3 =
+        "|x, p, _t, dx, rateiv, _cov| { dx[0] = if(t>0, exp(-ke * t) * x[0], 0.0); }".to_string();
+    let path3 = tmp.join("exa_example_model3.json");
+    let _ = pharmsol::exa_wasm::build::emit_ir::<pharmsol::equation::ODE>(
+        diffeq3,
+        None,
+        None,
+        None,
+        None,
+        Some(path3.clone()),
+        vec!["ke".to_string()],
+    )
+    .expect("emit_ir model3");
+
+    println!(
+        "Emitted IR to:\n  {:?}\n  {:?}\n  {:?}",
+        path1, path2, path3
+    );
+
+    // Load them via the public API and print emitted IR metadata from the
+    // emitted JSON (avoids accessing private registry internals from an
+    // example binary).
+    for p in [&path1, &path2, &path3] {
+        // try to load via runtime loader (public re-export)
+        match pharmsol::exa_wasm::load_ir_ode(p.clone()) {
+            Ok((_ode, _meta, id)) => {
+                println!("loader accepted model, registry id={}", id);
+            }
+            Err(e) => {
+                eprintln!("loader rejected model {:?}: {}", p, e);
+            }
+        }
+
+        // read raw IR and display bytecode/funcs/locals metadata
+        match fs::read_to_string(p) {
+            Ok(s) => match serde_json::from_str::<serde_json::Value>(&s) {
+                Ok(v) => {
+                    let has_bc = v.get("diffeq_bytecode").is_some();
+                    let funcs = v
+                        .get("funcs")
+                        .and_then(|j| {
+                            j.as_array()
+                                .map(|a| a.iter().filter_map(|x| x.as_str()).collect::<Vec<_>>())
+                        })
+                        .unwrap_or_default();
+                    let locals = v
+                        .get("locals")
+                        .and_then(|j| {
+                            j.as_array()
+                                .map(|a| a.iter().filter_map(|x| x.as_str()).collect::<Vec<_>>())
+                        })
+                        .unwrap_or_default();
+                    println!(
+                        "IR {:?}: diffeq_bytecode={} funcs={:?} locals={:?}",
+                        p.file_name().unwrap_or_default(),
+                        has_bc,
+                        funcs,
+                        locals
+                    );
+                }
+                Err(e) => eprintln!("failed to parse emitted IR {:?}: {}", p, e),
+            },
+            Err(e) => eprintln!("failed to read emitted IR {:?}: {}", p, e),
+        }
+    }
+
+    // cleanup
+    let _ = fs::remove_file(&path1);
+    let _ = fs::remove_file(&path2);
+    let _ = fs::remove_file(&path3);
+}

examples/emit_debug.rs

@@ -0,0 +1,19 @@
+fn main() {
+    use pharmsol::equation;
+    use pharmsol::exa_wasm::build::emit_ir;
+
+    // Simple helper example that emits IR for a small model and prints the
+    // location of the generated IR file. Keep example minimal and only use
+    // public APIs so it doesn't depend on internal interpreter modules.
+    let out = emit_ir::<equation::ODE>(
+        "|x, p, _t, dx, rateiv, _cov| { dx[0] = x[0].sin(); }".to_string(),
+        None,
+        None,
+        None,
+        None,
+        None,
+        vec![],
+    )
+    .expect("emit_ir");
+    println!("wrote IR to: {}", out);
+}

examples/exa.rs

@@ -34,7 +34,9 @@ fn main() {
     );
 
     //clear build
-    clear_build();
+    // clear_build();
+
+    println!("{}", exa::build::template_path());
 
     let test_dir = std::env::current_dir().expect("Failed to get current directory");
     let model_output_path = test_dir.join("test_model.pkm");
@@ -44,9 +46,9 @@ fn main() {
         format!(
             r#"
                 equation::ODE::new(
-            |x, p, _t, dx, rateiv, _cov| {{
+            |x, p, _t, dx, b, rateiv, _cov| {{
                 fetch_params!(p, ke, _v);
-                dx[0] = -ke * x[0] + rateiv[0];
+                dx[0] = -ke * x[0] + rateiv[0] + b[0];
             }},
             |_p, _t, _cov| lag! {{}},
             |_p, _t, _cov| fa! {{}},