bug: stop calculates time and state dep tensors (#27)

martinjrobins · Dec 7, 2024 · 65a3d27 · 65a3d27
1 parent 2f9a5e7
commit 65a3d27
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 2 deletions.
diff --git a/src/execution/compiler.rs b/src/execution/compiler.rs
@@ -618,6 +618,7 @@ mod tests {
         twoy_i { 2 * y }
         F_i { y * (1 - y), }
         out_i { twoy_i }
+        stop_i { twoy_i - 0.5 }
         ";
         let model = parse_ds_string(full_text).unwrap();
         let discrete_model = DiscreteModel::build("$name", &model).unwrap();
@@ -633,6 +634,12 @@ mod tests {
         compiler.calc_out(0., u0.as_slice(), data.as_mut_slice());
         let out = compiler.get_out(data.as_slice());
         assert_relative_eq!(out[0], 4.);
+        let mut stop = vec![0.];
+        compiler.calc_stop(0., u0.as_slice(), data.as_mut_slice(), stop.as_mut_slice());
+        assert_relative_eq!(stop[0], 3.5);
+        u0[0] = 0.5;
+        compiler.calc_stop(0., u0.as_slice(), data.as_mut_slice(), stop.as_mut_slice());
+        assert_relative_eq!(stop[0], 0.5);
     }
 
     #[test]

diff --git a/src/execution/cranelift/codegen.rs b/src/execution/cranelift/codegen.rs
@@ -332,7 +332,7 @@ impl CodegenModule for CraneliftModule {
             codegen.jit_compile_tensor(tensor, None, false)?;
         }
 
-        // TODO: could split state dep defns into before and after F
+        // calculate state dependant definitions
         for a in model.state_dep_defns() {
             codegen.jit_compile_tensor(a, None, false)?;
         }
@@ -355,6 +355,16 @@ impl CodegenModule for CraneliftModule {
         let mut codegen = CraneliftCodeGen::new(self, model, arg_names, arg_types);
 
         if let Some(stop) = model.stop() {
+            // calculate time dependant definitions
+            for tensor in model.time_dep_defns() {
+                codegen.jit_compile_tensor(tensor, None, false)?;
+            }
+
+            // calculate state dependant definitions
+            for a in model.state_dep_defns() {
+                codegen.jit_compile_tensor(a, None, false)?;
+            }
+
             let root = *codegen.variables.get("root").unwrap();
             codegen.jit_compile_tensor(stop, Some(root), false)?;
         }

diff --git a/src/execution/llvm/codegen.rs b/src/execution/llvm/codegen.rs
@@ -1892,7 +1892,7 @@ impl<'ctx> CodeGen<'ctx> {
             self.jit_compile_tensor(tensor, Some(*self.get_var(tensor)))?;
         }
 
-        // TODO: could split state dep defns into before and after F
+        // calculate state dependant definitions
         for a in model.state_dep_defns() {
             self.jit_compile_tensor(a, Some(*self.get_var(a)))?;
         }
@@ -1951,6 +1951,16 @@ impl<'ctx> CodeGen<'ctx> {
         self.insert_indices();
 
         if let Some(stop) = model.stop() {
+            // calculate time dependant definitions
+            for tensor in model.time_dep_defns() {
+                self.jit_compile_tensor(tensor, Some(*self.get_var(tensor)))?;
+            }
+
+            // calculate state dependant definitions
+            for a in model.state_dep_defns() {
+                self.jit_compile_tensor(a, Some(*self.get_var(a)))?;
+            }
+
             let res_ptr = self.get_param("root");
             self.jit_compile_tensor(stop, Some(*res_ptr))?;
         }