Fix tests + revert back input_solutions to solutions

pybamm-team · Oct 21, 2021 · da05ff8 · da05ff8
1 parent 7f53feb
commit da05ff8
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/pybamm/plotting/quick_plot.py b/pybamm/plotting/quick_plot.py
@@ -57,7 +57,7 @@ class QuickPlot(object):
 
     Parameters
     ----------
-    input_solutions: (iter of) :class:`pybamm.Solution` or :class:`pybamm.Simulation`
+    solutions: (iter of) :class:`pybamm.Solution` or :class:`pybamm.Simulation`
         The numerical solution(s) for the model(s), or the simulation object(s)
         containing the solution(s).
     output_variables : list of str, optional
@@ -91,7 +91,7 @@ class QuickPlot(object):
 
     def __init__(
         self,
-        input_solutions,
+        solutions,
         output_variables=None,
         labels=None,
         colors=None,
@@ -102,8 +102,13 @@ def __init__(
         spatial_unit="um",
         variable_limits="fixed",
     ):
+        input_solutions = solutions
         solutions = []
-        if not isinstance(input_solutions, list):
+        if not isinstance(input_solutions, (pybamm.Solution, pybamm.Simulation, list)):
+            raise TypeError(
+                "solutions must be 'pybamm.Solution' or 'pybamm.Simulation' or list"
+            )
+        elif not isinstance(input_solutions, list):
             input_solutions = [input_solutions]
         for sim_or_sol in input_solutions:
             if isinstance(sim_or_sol, pybamm.Simulation):
@@ -112,10 +117,6 @@ def __init__(
                 solutions.append(sim_or_sol.solution)
             elif isinstance(sim_or_sol, pybamm.Solution):
                 solutions.append(sim_or_sol)
-        else:
-            raise TypeError(
-                "solutions must be 'pybamm.Solution' or 'pybamm.Simulation' or list"
-            )
 
         models = [solution.all_models[0] for solution in solutions]
 

diff --git a/tests/unit/test_plotting/test_quick_plot.py b/tests/unit/test_plotting/test_quick_plot.py
@@ -488,7 +488,7 @@ def test_model_with_inputs(self):
             "Positive particle surface concentration",
         ]
         quick_plot = pybamm.QuickPlot(
-            input_solutions=[sol1, sol2], output_variables=output_variables
+            solutions=[sol1, sol2], output_variables=output_variables
         )
         quick_plot.dynamic_plot(testing=True)
         quick_plot.slider_update(1)