feat: add mass state and input dependencies, filter sparse zeros (#95)

martinjrobins · web-flow · commit 36eaf3d79c01 · 2025-12-27T16:54:12.000Z
* feat: add mass deps

* feat: track mass dependencies

* feat: filter out zeros in sparse tensors
diff --git a/diffsl/Cargo.toml b/diffsl/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "diffsl"
-version = "0.8.0"
+version = "0.8.1"
 edition.workspace = true
 license.workspace = true
 authors.workspace = true
diff --git a/diffsl/src/discretise/discrete_model.rs b/diffsl/src/discretise/discrete_model.rs
@@ -44,10 +44,13 @@ pub struct DiscreteModel<'s> {
     is_algebraic: Vec<bool>,
     stop: Option<Tensor<'s>>,
     state0_input_deps: Vec<(usize, usize)>,
+    dstate0_input_deps: Vec<(usize, usize)>,
     rhs_state_deps: Vec<(usize, usize)>,
     rhs_input_deps: Vec<(usize, usize)>,
     out_input_deps: Vec<(usize, usize)>,
     out_state_deps: Vec<(usize, usize)>,
+    mass_state_deps: Vec<(usize, usize)>,
+    mass_input_deps: Vec<(usize, usize)>,
 }
 
 impl fmt::Display for DiscreteModel<'_> {
@@ -105,10 +108,13 @@ impl<'s> DiscreteModel<'s> {
             is_algebraic: Vec::new(),
             stop: None,
             state0_input_deps: Vec::new(),
+            dstate0_input_deps: Vec::new(),
             rhs_input_deps: Vec::new(),
             rhs_state_deps: Vec::new(),
             out_input_deps: Vec::new(),
             out_state_deps: Vec::new(),
+            mass_state_deps: Vec::new(),
+            mass_input_deps: Vec::new(),
         }
     }
 
@@ -245,6 +251,15 @@ impl<'s> DiscreteModel<'s> {
             match Layout::concatenate(&elmts, rank) {
                 Ok(layout) => {
                     let layout = env.new_layout_ptr(layout);
+                    // if sparse, filter out zero elements
+                    let elmts = if layout.is_sparse() {
+                        elmts
+                            .into_iter()
+                            .filter(|e| !matches!(e.expr().kind, AstKind::Number(0.0)))
+                            .collect::<Vec<_>>()
+                    } else {
+                        elmts
+                    };
                     let tensor = Tensor::new(array.name(), elmts, layout, array.indices().to_vec());
 
                     //check that the number of indices matches the rank
@@ -335,7 +350,7 @@ impl<'s> DiscreteModel<'s> {
                         }
                         "dudt" => {
                             if let Some(built) =
-                                Self::build_array(tensor, &mut env, true, TensorType::Other)
+                                Self::build_array(tensor, &mut env, true, TensorType::StateDot)
                             {
                                 ret.state_dot = Some(built);
                             }
@@ -522,6 +537,7 @@ impl<'s> DiscreteModel<'s> {
 
         // store the dependencies in the discrete model
         ret.state0_input_deps = map_dep(&env.state0_input_deps);
+        ret.dstate0_input_deps = map_dep(&env.dstate0_input_deps);
         ret.rhs_input_deps = map_dep(ret.rhs.layout().input_dependencies());
         ret.rhs_state_deps = map_dep(ret.rhs.layout().state_dependencies());
         ret.out_input_deps = ret
@@ -534,6 +550,16 @@ impl<'s> DiscreteModel<'s> {
             .as_ref()
             .map(|o| map_dep(o.layout().state_dependencies()))
             .unwrap_or_default();
+        ret.mass_state_deps = if let Some(lhs) = &ret.lhs {
+            map_dep(lhs.layout().state_dependencies())
+        } else {
+            Vec::new()
+        };
+        ret.mass_input_deps = if let Some(lhs) = &ret.lhs {
+            map_dep(lhs.layout().input_dependencies())
+        } else {
+            Vec::new()
+        };
 
         if env.errs().is_empty() {
             Ok(ret)
@@ -759,10 +785,13 @@ impl<'s> DiscreteModel<'s> {
             is_algebraic,
             stop,
             state0_input_deps: Vec::new(),
+            dstate0_input_deps: Vec::new(),
             rhs_state_deps: Vec::new(),
             rhs_input_deps: Vec::new(),
             out_input_deps: Vec::new(),
             out_state_deps: Vec::new(),
+            mass_state_deps: Vec::new(),
+            mass_input_deps: Vec::new(),
         }
     }
 
@@ -825,6 +854,10 @@ impl<'s> DiscreteModel<'s> {
         std::mem::take(&mut self.state0_input_deps)
     }
 
+    pub fn take_dstate0_input_deps(&mut self) -> Vec<(usize, usize)> {
+        std::mem::take(&mut self.dstate0_input_deps)
+    }
+
     pub fn take_rhs_state_deps(&mut self) -> Vec<(usize, usize)> {
         std::mem::take(&mut self.rhs_state_deps)
     }
@@ -840,6 +873,14 @@ impl<'s> DiscreteModel<'s> {
     pub fn take_out_state_deps(&mut self) -> Vec<(usize, usize)> {
         std::mem::take(&mut self.out_state_deps)
     }
+
+    pub fn take_mass_state_deps(&mut self) -> Vec<(usize, usize)> {
+        std::mem::take(&mut self.mass_state_deps)
+    }
+
+    pub fn take_mass_input_deps(&mut self) -> Vec<(usize, usize)> {
+        std::mem::take(&mut self.mass_input_deps)
+    }
 }
 
 #[cfg(test)]
@@ -1337,6 +1378,84 @@ mod tests {
 
     );
 
+    #[test]
+    fn tensor_state_input_dep_mass_test() {
+        let full_text = "
+            in_i { (0:2): p = 1 }
+            u_i { p_i }
+            dudt_i { p_i }
+            M_i { dudt_i[1] + p_i[0], dudt_i[0] + p_i[1] }
+            F_i { u_i }
+        ";
+
+        let model = parse_ds_string(full_text).unwrap();
+        let mut discrete_model =
+            DiscreteModel::build("tensor_state_input_dep_mass_test", &model).unwrap();
+        assert_eq!(
+            discrete_model.take_state0_input_deps(),
+            vec![(0, 0), (1, 1)]
+        );
+        assert_eq!(
+            discrete_model.take_dstate0_input_deps(),
+            vec![(0, 0), (1, 1)]
+        );
+        assert_eq!(
+            discrete_model.take_rhs_state_deps(),
+            vec![(0, 0), (1, 1)],
+            "failed rhs_state_deps"
+        );
+        assert_eq!(
+            discrete_model.take_rhs_input_deps(),
+            vec![],
+            "failed rhs_input_deps"
+        );
+        assert_eq!(discrete_model.take_out_state_deps(), vec![]);
+        assert_eq!(discrete_model.take_out_input_deps(), vec![]);
+        assert_eq!(discrete_model.take_mass_state_deps(), vec![(0, 1), (1, 0)]);
+        assert_eq!(discrete_model.take_mass_input_deps(), vec![(0, 0), (1, 1)]);
+    }
+
+    #[test]
+    fn tensor_state_input_dep_logistic_test() {
+        let full_text = "
+            in_i { r = 1, k = 1 }
+            u_i {
+                y = 0.1,
+                z = 0,
+            }
+            dudt_i {
+                dydt = 0,
+                dzdt = 0,
+            }
+            M_i {
+                dydt,
+                0,
+            }
+            F_i {
+                (r * y) * (1 - (y / k)),
+                (2 * y) - z,
+            }
+            out_i {
+                3 * y,
+                4 * z,
+            }
+        ";
+        let model = parse_ds_string(full_text).unwrap();
+        let mut discrete_model =
+            DiscreteModel::build("tensor_state_input_dep_logistic_test", &model).unwrap();
+        assert_eq!(discrete_model.take_state0_input_deps(), vec![]);
+        assert_eq!(discrete_model.take_dstate0_input_deps(), vec![]);
+        assert_eq!(
+            discrete_model.take_rhs_state_deps(),
+            vec![(0, 0), (1, 0), (1, 1)]
+        );
+        assert_eq!(discrete_model.take_rhs_input_deps(), vec![(0, 0), (0, 1)]);
+        assert_eq!(discrete_model.take_out_state_deps(), vec![(0, 0), (1, 1)]);
+        assert_eq!(discrete_model.take_out_input_deps(), vec![]);
+        assert_eq!(discrete_model.take_mass_state_deps(), vec![(0, 0)]);
+        assert_eq!(discrete_model.take_mass_input_deps(), vec![]);
+    }
+
     macro_rules! tensor_state_input_dep_test {
         ($($name:ident: $text:literal expect $expected_state_state_deps:expr ; $expected_state_inputs_deps:expr,)*) => {
         $(
@@ -1389,6 +1508,7 @@ mod tests {
         tsi_diag_mat_mul2: "A_ij { (0..2, 0..2): p_i } F_i { A_ij * u_j }" expect vec![(0,0), (1,1)] ; vec![(0,0), (1,1)],
         tsi_concat: "F_i { (0): u_i[0], (1): p_i[0] }" expect vec![(0,0)] ; vec![(1,0)],
         tsi_concat2: "a_ij { (0,0): u_i[0], (1,1): p_i[0] } F_i { a_ij }" expect vec![(0,0)] ; vec![(1,0)],
+        tsi_expr: "a_i { y = u_i[0], z = u_i[1] } b_i { r = p_i[0], k = p_i[1] } F_i { (r * y) * (1 - y / k), (2 * y) - z }" expect vec![(0,0), (1,0), (1,1)] ; vec![(0,0), (0,1)],
     }
 
     #[test]
diff --git a/diffsl/src/discretise/env.rs b/diffsl/src/discretise/env.rs
@@ -53,6 +53,7 @@ pub struct Env {
     errs: ValidationErrors,
     vars: HashMap<String, EnvVar>,
     pub(crate) state0_input_deps: Vec<NonZero>,
+    pub(crate) dstate0_input_deps: Vec<NonZero>,
 }
 
 impl Env {
@@ -74,6 +75,7 @@ impl Env {
             vars,
             current_span: None,
             state0_input_deps: vec![],
+            dstate0_input_deps: vec![],
         }
     }
 
diff --git a/diffsl/src/discretise/layout.rs b/diffsl/src/discretise/layout.rs
@@ -5,12 +5,13 @@ use std::{
     convert::AsRef,
     fmt,
     hash::{Hash, Hasher},
+    iter::zip,
     mem,
     ops::Deref,
     sync::Arc,
 };
 
-use crate::discretise::Env;
+use crate::{ast::AstKind, discretise::Env};
 
 use super::{broadcast_shapes, shape::Shape, Index, TensorBlock};
 
@@ -24,6 +25,7 @@ pub enum LayoutKind {
 #[derive(Debug, Clone, Copy)]
 pub enum TensorType {
     State,
+    StateDot,
     Input,
     Other,
 }
@@ -163,12 +165,12 @@ impl Layout {
     /// Add state or input dependencies to this layout based on the tensor type.
     pub fn add_tensor_dependencies(&mut self, tensor_type: TensorType, start: i64, env: &mut Env) {
         let indices = match tensor_type {
-            TensorType::State | TensorType::Input => {
+            TensorType::State | TensorType::StateDot | TensorType::Input => {
                 let mut deps = Vec::new();
-                let n_states = *self.shape().get(0).unwrap_or(&1);
-                for i in 0..n_states {
-                    let index = Index::from(vec![(i as i64) + start]);
-                    deps.push((index, i));
+                let n_states = *self.shape().get(0).unwrap_or(&1) as i64;
+                for i in 0_i64..n_states {
+                    let index = Index::from(vec![i]);
+                    deps.push((index, (i + start) as usize));
                 }
                 deps
             }
@@ -184,6 +186,15 @@ impl Layout {
                 // store the state0 input dependencies in the env since we don't want to propagate them further
                 env.state0_input_deps = mem::take(&mut self.input_deps);
             }
+            TensorType::StateDot => {
+                assert!(
+                    self.state_deps.is_empty(),
+                    "state dot tensor layout should not already have state dependencies",
+                );
+                self.state_deps = indices;
+                // store the dstate0 input dependencies in the env since we don't want to propagate them further
+                env.dstate0_input_deps = mem::take(&mut self.input_deps);
+            }
             TensorType::Input => {
                 assert! {
                     self.input_deps.is_empty(),
@@ -657,6 +668,10 @@ impl Layout {
             .iter()
             .map(|x| x.layout().as_ref())
             .collect::<Vec<_>>();
+        let is_zero = elmts
+            .iter()
+            .map(|x| matches!(x.expr().kind, AstKind::Number(0.0)))
+            .collect::<Vec<_>>();
         let starts = elmts.iter().map(|x| x.start()).collect::<Vec<_>>();
 
         // if there are no layouts then return an empty layout
@@ -800,7 +815,13 @@ impl Layout {
                     kind: LayoutKind::Sparse,
                     n_dense_axes,
                 };
-                for (layout, start) in std::iter::zip(layouts.iter(), starts.iter()) {
+                for ((layout, start), is_zero) in
+                    zip(zip(layouts.iter(), starts.iter()), is_zero.iter())
+                {
+                    if *is_zero {
+                        continue;
+                    }
+
                     // convert to sparse
                     new_layout.indices.extend(layout.indices().map(|mut x| {
                         for (i, xi) in x.iter_mut().enumerate() {
diff --git a/diffsl/src/execution/compiler.rs b/diffsl/src/execution/compiler.rs
@@ -1532,9 +1532,9 @@ mod tests {
         sparse_mat_mul12: "A_ij { (0,0):1, (0,1):1, (1,1):1, (2,2):1, (3,3):1, (4,4):1, (5,5):1 } b2_i { (1:6): 1 } r_i { A_ij * b2_j }" expect "r" vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
         sparse_mat_mul8: "A_ij { (0,1): -0.5, (0,0): 1.5, (6,5): 1.5, (6,4): -0.5 } b_i { (0:6): 1 } r_i { A_ij * b_j } A1_ij { (0,1): 1, (0,2): 1, (1,1): 1, (2,2): 1, (3,3): 1, (4,4): 1, (5,5): 1, (6,6): 1 } b2_i { (0:7): 1 } b3_i { (0:7):2 } r2_i { A1_ij * (r_j + b2_j) * b3_j }" expect "r2" vec![4.0, 2.0, 2.0, 2.0, 2.0, 2.0, 4.0],
         sparse_dense_concat: "a_i { (0): 1, (2): 3 } b_i { 4, 5 } r_i { a_i, b_i }" expect "r" vec![1., 3., 4., 5.],
-        sparse_mat_mul9: "A_ij { (0,1): -0.5, (0,0): 1.5, (1,1): 0, (2,2): 0, (3,3): 0, (4,4): 0, (5,5): 0, (6,5): 1.5, (6,4): -0.5 } b_i { (0:6): 1 } r_i { A_ij * b_j }" expect "r" vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+        sparse_mat_mul9: "A_ij { (0,1): -0.5, (0,0): 1.5, (1,1): 0, (2,2): 0, (3,3): 0, (4,4): 0, (5,5): 0, (6,5): 1.5, (6,4): -0.5 } b_i { (0:6): 1 } r_i { A_ij * b_j }" expect "r" vec![1.0, 1.0],
         sparse_mat_mul10: "A_ij { (0,0): 1.5, (0,1): -0.5, (6,4): -0.5, (6,5): 1.5 } b_i { (0:6): 1 } r_i { A_ij * b_j }" expect "r" vec![1.0, 1.0],
-        sparse_mat_mul11: "A_ij { (0,0): 1.5, (0,1): -0.5, (1,1): 0, (2,2): 0, (3,3): 0, (4,4): 0, (5,5): 0, (6,4): -0.5, (6,5): 1.5 } b_i { (0:6): 1 } r_i { A_ij * b_j }" expect "r" vec![1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+        sparse_mat_mul11: "A_ij { (0,0): 1.5, (0,1): -0.5, (1,1): 0, (2,2): 0, (3,3): 0, (4,4): 0, (5,5): 0, (6,4): -0.5, (6,5): 1.5 } b_i { (0:6): 1 } r_i { A_ij * b_j }" expect "r" vec![1.0, 1.0],
         sparse_nonsquare_mat_vec_mul: "A_ij { (0, 0): 1, (0, 1): 4, (1, 2): 2 } b_j { (0:3): 5 } r_i { A_ij * b_j }" expect "r" vec![25.0, 10.0],
         sparse_nonsquare_mat_vec_mul2: "A_ij { (0, 0): 1, (0, 1): 4, (1, 2): 2 } b_j { (2): 5 } r_i { A_ij * b_j } B_ij { (0..2,0..2): 1 } s_i { B_ij * max(r_j, 1) }" expect "s" vec![1.0, 10.0],
         max_sparse_vec: "A_ij { (0..2,0..2): 1 } b_j { (1): 5 } r_i { A_ij * max(b_j, 1) }" expect "r" vec![1.0, 5.0],
@@ -1543,9 +1543,9 @@ mod tests {
         contract_to_mat_vec: "A_ij { (0, 0): 1, (1, 0): 3, (1, 1): 4 } B_ij { (1, 1): 2 } b_i { B_ij } r_i { A_ij * b_j }" expect "r" vec![8.],
         sparse_mat_vec_mul7: "A_ij { (0, 1): 4, (1, 2): 2, (2, 2): 1 } b_i { (2): 5 } r_i { A_ij * (1 + b_j) }" expect "r" vec![4., 12., 6.],
         sparse_mat_vec_mul6: "A_ij { (0, 1): 4, (1, 2): 2, (2, 2): 1 } b_i { (2): 5 } r_i { A_ij * (1 * b_j) }" expect "r" vec![10., 5.],
-        sparse_mat_vec_mul3: "A_ij { (0, 1): 4, (1, 2): 2, (2, 2): 0 } b_i { (2): 5 } c_j { (0:3): 1 } r_i { A_ij * (b_j + c_j) }" expect "r" vec![4., 12., 0.],
+        sparse_mat_vec_mul3: "A_ij { (0, 1): 4, (1, 2): 2, (2, 2): 0 } b_i { (2): 5 } c_j { (0:3): 1 } r_i { A_ij * (b_j + c_j) }" expect "r" vec![4., 12.],
         sparse_mat_vec_mul5: "A_ij { (1, 1): 2 } b_j { (1): 3 } r_i { A_ij * (1 * b_j) }" expect "r" vec![6.0],
-        sparse_mat_vec_mul4: "A_ij { (0, 1): 4, (1, 2): 2, (2, 2): 0 } b_i { (0): 2, (2): 5 } c_j { (0:3): 1 } r_i { A_ij * (b_j + c_j) }" expect "r" vec![4., 12., 0.],
+        sparse_mat_vec_mul4: "A_ij { (0, 1): 4, (1, 2): 2, (2, 2): 0 } b_i { (0): 2, (2): 5 } c_j { (0:3): 1 } r_i { A_ij * (b_j + c_j) }" expect "r" vec![4., 12.],
         sparse_mat_vec_mul2: "A_ij { (0, 1): 4, (1, 0): 2 } b_i { (1): 5 } c_j { (0:1): 1, (1:2): 1 } r_i { A_ij * (b_j + c_j) }" expect "r" vec![24.0, 2.0],
         sparse_mat_vec_mul: "A_ij { (1, 1): 2 } b_j { (1): 3 } r_i { A_ij * b_j }" expect "r" vec![6.0],
         sparse_broadcast_to_sparse: "A_i { (1): 2 } B_ij { (0:2, 0:2): A_i }" expect "B" vec![2.0],

Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,7 @@ pub struct Env {`
`53`	`53`	`errs: ValidationErrors,`
`54`	`54`	`vars: HashMap<String, EnvVar>,`
`55`	`55`	`pub(crate) state0_input_deps: Vec<NonZero>,`
	`56`	`+ pub(crate) dstate0_input_deps: Vec<NonZero>,`
`56`	`57`	`}`
`57`	`58`
`58`	`59`	`impl Env {`
`@@ -74,6 +75,7 @@ impl Env {`
`74`	`75`	`vars,`
`75`	`76`	`current_span: None,`
`76`	`77`	`state0_input_deps: vec![],`
	`78`	`+ dstate0_input_deps: vec![],`
`77`	`79`	`}`
`78`	`80`	`}`
`79`	`81`