refactor: model index no longer stored in module, passed in with state

diffsl/benches/evaluation.rs

@@ -55,14 +55,14 @@ fn execute<const N: usize, M: CodegenModuleCompile + CodegenModuleJit>(
     let n = N;
     let compiler = setup::<M>(n, f_text, "execute");
     let mut data = compiler.get_new_data();
-    compiler.set_inputs(&[], data.as_mut_slice(), 0);
+    compiler.set_inputs(&[], data.as_mut_slice());
     let mut u = vec![1.0; n];
-    compiler.set_u0(u.as_mut_slice(), data.as_mut_slice());
+    compiler.set_u0(u.as_mut_slice(), data.as_mut_slice(), 0);
     let mut rr = vec![0.0; n];
     let t = 0.0;
 
     bencher.bench_local(|| {
-        compiler.rhs(t, &u, &mut data, &mut rr);
+        compiler.rhs(t, &u, &mut data, &mut rr, 0);
     });
 }
 

diffsl/benches/pybamm_dfn.rs

@@ -36,6 +36,7 @@ fn pybamm_dfn_execute_rhs_grad<M: CodegenModuleCompile + CodegenModuleJit>(bench
             ddata.as_mut_slice(),
             rr.as_mut_slice(),
             drr.as_mut_slice(),
+            0,
         );
     });
 }
@@ -69,7 +70,7 @@ fn pybamm_dfn_execute_rhs<M: CodegenModuleCompile + CodegenModuleJit>(bencher: B
     let mut data = compiler.get_new_data();
     let mut rr = vec![0.0; n_states];
     bencher.bench_local(move || {
-        compiler.rhs(t, y.as_slice(), data.as_mut_slice(), rr.as_mut_slice());
+        compiler.rhs(t, y.as_slice(), data.as_mut_slice(), rr.as_mut_slice(), 0);
     });
 }
 

diffsl/src/discretise/discrete_model.rs

@@ -35,6 +35,7 @@ pub struct DiscreteModel<'s> {
     out: Option<Tensor<'s>>,
     constant_defns: Vec<Tensor<'s>>,
     input_dep_defns: Vec<Tensor<'s>>,
+    model_dep_defns: Vec<Tensor<'s>>,
     time_dep_defns: Vec<Tensor<'s>>,
     state_dep_defns: Vec<Tensor<'s>>,
     state_dep_post_f_defns: Vec<Tensor<'s>>,
@@ -66,6 +67,9 @@ impl fmt::Display for DiscreteModel<'_> {
         for defn in &self.input_dep_defns {
             writeln!(f, "{defn}")?;
         }
+        for defn in &self.model_dep_defns {
+            writeln!(f, "{defn}")?;
+        }
         for defn in &self.time_dep_defns {
             writeln!(f, "{defn}")?;
         }
@@ -107,6 +111,7 @@ impl<'s> DiscreteModel<'s> {
             out: None,
             constant_defns: Vec::new(),
             input_dep_defns: Vec::new(),
+            model_dep_defns: Vec::new(),
             time_dep_defns: Vec::new(),
             state_dep_defns: Vec::new(),
             state_dep_post_f_defns: Vec::new(),
@@ -538,6 +543,12 @@ impl<'s> DiscreteModel<'s> {
                                         && !dependent_on_model
                                     {
                                         ret.constant_defns.push(built);
+                                    } else if dependent_on_model
+                                        && !dependent_on_time
+                                        && !dependent_on_state
+                                        && !dependent_on_dudt
+                                    {
+                                        ret.model_dep_defns.push(built);
                                     } else if !dependent_on_time {
                                         ret.input_dep_defns.push(built);
                                     } else if !dependent_on_state && !dependent_on_dudt {
@@ -868,6 +879,7 @@ impl<'s> DiscreteModel<'s> {
             out: Some(out_array),
             constant_defns,
             input_dep_defns: Vec::new(), // todo: need to implement
+            model_dep_defns: Vec::new(),
             time_dep_defns,
             state_dep_defns,
             state_dep_post_f_defns,
@@ -898,6 +910,10 @@ impl<'s> DiscreteModel<'s> {
         self.input_dep_defns.as_ref()
     }
 
+    pub fn model_dep_defns(&self) -> &[Tensor<'_>] {
+        self.model_dep_defns.as_ref()
+    }
+
     pub fn time_dep_defns(&self) -> &[Tensor<'_>] {
         self.time_dep_defns.as_ref()
     }