Set zip(..., strict=True) in key solver/expression tree files for len…

CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ## Bug fixes
 
+- Set `zip(..., strict=True)` in solver and expression tree files to ensure iterable length safety. ([#5241](https://github.com/pybamm-team/PyBaMM/pull/5241))
 - Adds `options` property to IDAKLU, fixes pickling issue with `__getstate__` when keys are not available. ([#5234](https://github.com/pybamm-team/PyBaMM/pull/5234))
 - Fixed a bug where simulations using output variables in `IDAKLUSolver` couldn't be pickled ([#5225](https://github.com/pybamm-team/PyBaMM/pull/5225))
 - Added explicit warning in installation docs about unmaintained Conda recipe due to pybammsolvers split (Fixes #5155). See pull request [#5206](https://github.com/pybamm-team/PyBaMM/pull/5206)

src/pybamm/discretisations/discretisation.py

@@ -1174,9 +1174,7 @@ def _concatenate_in_order(self, var_eqn_dict, check_complete=False, sparse=False
         equations = list(var_eqn_dict.values())
 
         # sort equations according to slices
-        sorted_equations = [
-            eq for _, eq in sorted(zip(slices, equations, strict=False))
-        ]
+        sorted_equations = [eq for _, eq in sorted(zip(slices, equations, strict=True))]
 
         return self.concatenate(*sorted_equations, sparse=sparse)
 

src/pybamm/expression_tree/binary_operators.py

@@ -910,7 +910,7 @@ def _simplified_binary_broadcast_concatenation(
                     [
                         operator(left_child, right_child)
                         for left_child, right_child in zip(
-                            left.orphans, right.orphans, strict=False
+                            left.orphans, right.orphans, strict=True
                         )
                     ]
                 )
@@ -1004,7 +1004,7 @@ def add(left: ChildSymbol, right: ChildSymbol):
         elif all(
             left_dim_size <= right_dim_size
             for left_dim_size, right_dim_size in zip(
-                left.shape_for_testing, right.shape_for_testing, strict=False
+                left.shape_for_testing, right.shape_for_testing, strict=True
             )
         ) and all(
             left.evaluates_on_edges(dim) == right.evaluates_on_edges(dim)
@@ -1093,7 +1093,7 @@ def subtract(
         elif all(
             left_dim_size <= right_dim_size
             for left_dim_size, right_dim_size in zip(
-                left.shape_for_testing, right.shape_for_testing, strict=False
+                left.shape_for_testing, right.shape_for_testing, strict=True
             )
         ) and all(
             left.evaluates_on_edges(dim) == right.evaluates_on_edges(dim)

src/pybamm/expression_tree/concatenations.py

@@ -376,7 +376,7 @@ def _concatenation_evaluate(self, children_eval: list[npt.NDArray[Any]]):
 
         # loop through domains of children writing subvectors to final vector
         for child_vector, slices in zip(
-            children_eval, self._children_slices, strict=False
+            children_eval, self._children_slices, strict=True
         ):
             for child_dom, child_slice in slices.items():
                 for i, _slice in enumerate(child_slice):
@@ -391,7 +391,7 @@ def _concatenation_jac(self, children_jacs):
         jacs = []
         for i in range(self.secondary_dimensions_npts):
             for child_jac, slices in zip(
-                children_jacs, self._children_slices, strict=False
+                children_jacs, self._children_slices, strict=True
             ):
                 if len(slices) > 1:
                     raise NotImplementedError(

src/pybamm/expression_tree/interpolant.py

@@ -314,7 +314,7 @@ def _function_evaluate(self, evaluated_children):
             if np.any(nans):
                 nan_children = []
                 for child, interp_range in zip(
-                    new_evaluated_children, self.function.grid, strict=False
+                    new_evaluated_children, self.function.grid, strict=True
                 ):
                     nan_children.append(np.ones_like(child) * interp_range.mean())
                 nan_eval = self.function(np.transpose(nan_children))

src/pybamm/expression_tree/operations/convert_to_casadi.py

@@ -279,7 +279,7 @@ def hermite_poly(i):
                 for i in range(symbol.secondary_dimensions_npts):
                     child_vectors = []
                     for child_var, slices in zip(
-                        converted_children, symbol._children_slices, strict=False
+                        converted_children, symbol._children_slices, strict=True
                     ):
                         for child_dom, child_slice in slices.items():
                             slice_starts.append(symbol._slices[child_dom][i].start)

src/pybamm/parameters/parameter_values.py

@@ -863,7 +863,7 @@ def _process_function_parameter(self, symbol):
             inputs = {
                 arg: child
                 for arg, child in zip(
-                    function_parameter.func_args, symbol.children, strict=False
+                    function_parameter.func_args, symbol.children, strict=True
                 )
             }
 

src/pybamm/solvers/algebraic_solver.py

@@ -215,7 +215,7 @@ def jac_norm(y, jac_fn=jac_fn):
                         bounds = [
                             (lb, ub)
                             for lb, ub in zip(
-                                model.bounds[0], model.bounds[1], strict=False
+                                model.bounds[0], model.bounds[1], strict=True
                             )
                         ]
                     else:

src/pybamm/solvers/base_solver.py

@@ -970,7 +970,7 @@ def solve(
                                 [t_eval[start_index:end_index]] * ninputs,
                                 model_inputs_list,
                                 [t_interp] * ninputs,
-                                strict=False,
+                                strict=True,
                             ),
                         )
                         p.close()
@@ -1198,7 +1198,7 @@ def _set_sens_initial_conditions_from(
 
         # sort equations according to slices
         concatenated_initial_conditions = [
-            casadi.vertcat(*[eq for _, eq in sorted(zip(slices, init, strict=False))])
+            casadi.vertcat(*[eq for _, eq in sorted(zip(slices, init, strict=True))])
             for init in initial_conditions
         ]
         return concatenated_initial_conditions

src/pybamm/solvers/processed_variable.py

@@ -426,7 +426,7 @@ def initialise_sensitivity_explicit_forward(self):
             self.all_inputs,
             self.base_variables,
             self.all_solution_sensitivities["all"],
-            strict=False,
+            strict=True,
         ):
             # Set up symbolic variables
             t_casadi = casadi.MX.sym("t")

src/pybamm/solvers/solution.py

@@ -432,9 +432,7 @@ def _update_variable(self, variable):
         # therefore only get set up once
         vars_casadi = []
         for i, (model, ys, inputs, var_pybamm) in enumerate(
-            zip(
-                self.all_models, self.all_ys, self.all_inputs, vars_pybamm, strict=False
-            )
+            zip(self.all_models, self.all_ys, self.all_inputs, vars_pybamm, strict=True)
         ):
             if self.variables_returned and var_pybamm.has_symbol_of_classes(
                 pybamm.expression_tree.state_vector.StateVector