Repository: pybamm-team/PyBaMM (python)

Issue:
Plotting with external variables
**Describe the bug**
Solutions of models that use external variables cannot be plotted with quickplot (problem with the inputs in `ProcessedVariable`)

**Expected behaviour**
The solution should plot fine

Pull request description:
#1024 fix external variables
# Description

Some bug fixes to `Solution` and `ProcessedVariable` to allow plotting when solving a simulation with external variables

Fixes #1024 

## Type of change

Please add a line in the relevant section of [CHANGELOG.md](https://github.com/pybamm-team/PyBaMM/blob/master/CHANGELOG.md) to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

- [ ] New feature (non-breaking change which adds functionality)
- [ ] Optimization (back-end change that speeds up the code)
- [x] Bug fix (non-breaking change which fixes an issue)


# Key checklist:

- [x] No style issues: `$ flake8`
- [x] All tests pass: `$ python run-tests.py --unit`
- [x] The documentation builds: `$ cd docs` and then `$ make clean; make html`

You can run all three at once, using `$ python run-tests.py --quick`.

## Further checks:

- [x] Code is commented, particularly in hard-to-understand areas
- [x] Tests added that prove fix is effective or that feature works


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 59560dc9d..1b186dd56 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,6 +27,8 @@
 
 ## Bug fixes
 
+-   Fix storing and plotting external variables in the solution ([#1026](https://github.com/pybamm-team/PyBaMM/pull/1026))
+-   Fix running a simulation with a model that is already discretized ([#1025](https://github.com/pybamm-team/PyBaMM/pull/1025))
 -   Fix CI not triggering for PR. ([#1013](https://github.com/pybamm-team/PyBaMM/pull/1013))
 -   Fix schedule testing running too often. ([#1010](https://github.com/pybamm-team/PyBaMM/pull/1010))
 -   Fix doctests failing due to mismatch in unsorted output.([#990](https://github.com/pybamm-team/PyBaMM/pull/990))
diff --git a/pybamm/expression_tree/variable.py b/pybamm/expression_tree/variable.py
index 374344ab0..81acd9044 100644
--- a/pybamm/expression_tree/variable.py
+++ b/pybamm/expression_tree/variable.py
@@ -197,6 +197,8 @@ class ExternalVariable(Variable):
                     )
                 )
             else:
+                if isinstance(out, np.ndarray) and out.ndim == 1:
+                    out = out[:, np.newaxis]
                 return out
         # raise more informative error if can't find name in dict
         except KeyError:
diff --git a/pybamm/simulation.py b/pybamm/simulation.py
index 003b03895..1fb11bc19 100644
--- a/pybamm/simulation.py
+++ b/pybamm/simulation.py
@@ -418,7 +418,12 @@ class Simulation:
                 t_eval = np.linspace(0, t_end, 100)
 
             self.t_eval = t_eval
-            self._solution = solver.solve(self.built_model, t_eval, inputs=inputs)
+            self._solution = solver.solve(
+                self.built_model,
+                t_eval,
+                external_variables=external_variables,
+                inputs=inputs,
+            )
 
         elif self.operating_mode == "with experiment":
             if t_eval is not None:
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
index 4b4286b4b..ccdd8ec62 100644
--- a/pybamm/solvers/base_solver.py
+++ b/pybamm/solvers/base_solver.py
@@ -775,7 +775,7 @@ class BaseSolver(object):
                         event.expression.evaluate(
                             solution.t_event,
                             solution.y_event,
-                            inputs={k: v[-1] for k, v in solution.inputs.items()},
+                            inputs={k: v[:, -1] for k, v in solution.inputs.items()},
                         )
                     )
             termination_event = min(final_event_values, key=final_event_values.get)
diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py
index f324b0068..be4cfab14 100644
--- a/pybamm/solvers/processed_variable.py
+++ b/pybamm/solvers/processed_variable.py
@@ -59,14 +59,14 @@ class ProcessedVariable(object):
             self.base_eval, self.known_evals[solution.t[0]] = base_variable.evaluate(
                 solution.t[0],
                 solution.y[:, 0],
-                inputs={name: inp[0] for name, inp in solution.inputs.items()},
+                inputs={name: inp[:, 0] for name, inp in solution.inputs.items()},
                 known_evals=self.known_evals[solution.t[0]],
             )
         else:
             self.base_eval = base_variable.evaluate(
                 solution.t[0],
                 solution.y[:, 0],
-                inputs={name: inp[0] for name, inp in solution.inputs.items()},
+                inputs={name: inp[:, 0] for name, inp in solution.inputs.items()},
             )
 
         # handle 2D (in space) finite element variables differently
@@ -115,7 +115,7 @@ class ProcessedVariable(object):
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
             u = self.u_sol[:, idx]
-            inputs = {name: inp[idx] for name, inp in self.inputs.items()}
+            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
             if self.known_evals:
                 entries[idx], self.known_evals[t] = self.base_variable.evaluate(
                     t, u, inputs=inputs, known_evals=self.known_evals[t]
@@ -151,7 +151,7 @@ class ProcessedVariable(object):
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
             u = self.u_sol[:, idx]
-            inputs = {name: inp[idx] for name, inp in self.inputs.items()}
+            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
             if self.known_evals:
                 eval_and_known_evals = self.base_variable.evaluate(
                     t, u, inputs=inputs, known_evals=self.known_evals[t]
@@ -270,7 +270,7 @@ class ProcessedVariable(object):
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
             u = self.u_sol[:, idx]
-            inputs = {name: inp[idx] for name, inp in self.inputs.items()}
+            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
             if self.known_evals:
                 eval_and_known_evals = self.base_variable.evaluate(
                     t, u, inputs=inputs, known_evals=self.known_evals[t]
@@ -330,7 +330,7 @@ class ProcessedVariable(object):
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
             u = self.u_sol[:, idx]
-            inputs = {name: inp[idx] for name, inp in self.inputs.items()}
+            inputs = {name: inp[:, idx] for name, inp in self.inputs.items()}
 
             if self.known_evals:
                 eval_and_known_evals = self.base_variable.evaluate(
diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py
index 097a25431..2ecfa547f 100644
--- a/pybamm/solvers/solution.py
+++ b/pybamm/solvers/solution.py
@@ -108,8 +108,12 @@ class _BaseSolution(object):
             self.has_symbolic_inputs = False
             self._inputs = {}
             for name, inp in inputs.items():
+                # Convert number to vector of the right shape
                 if isinstance(inp, numbers.Number):
-                    inp = inp * np.ones_like(self.t)
+                    inp = inp * np.ones((1, len(self.t)))
+                # Tile a vector
+                else:
+                    inp = np.tile(inp, len(self.t))
                 self._inputs[name] = inp
 
     @property
@@ -316,7 +320,7 @@ class Solution(_BaseSolution):
         self._y = np.concatenate((self._y, solution.y[:, start_index:]), axis=1)
         for name, inp in self.inputs.items():
             solution_inp = solution.inputs[name]
-            self.inputs[name] = np.concatenate((inp, solution_inp[start_index:]))
+            self.inputs[name] = np.c_[inp, solution_inp[:, start_index:]]
         # Update solution time
         self.solve_time += solution.solve_time
         # Update termination


Test patch:
diff --git a/tests/integration/test_solvers/test_solution.py b/tests/integration/test_solvers/test_solution.py
index 13f6d0292..b27e0a0b3 100644
--- a/tests/integration/test_solvers/test_solution.py
+++ b/tests/integration/test_solvers/test_solution.py
@@ -59,6 +59,65 @@ class TestSolution(unittest.TestCase):
             decimal=4,
         )
 
+    def test_append_external_variables(self):
+        model = pybamm.lithium_ion.SPM(
+            {
+                "thermal": "lumped",
+                "external submodels": ["thermal", "negative particle"],
+            }
+        )
+        # create geometry
+        geometry = model.default_geometry
+
+        # load parameter values and process model and geometry
+        param = model.default_parameter_values
+        param.process_model(model)
+        param.process_geometry(geometry)
+
+        # set mesh
+        mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)
+
+        # discretise model
+        disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
+        disc.process_model(model)
+
+        # solve model
+        solver = model.default_solver
+
+        T_av = 0
+        c_s_n_av = np.ones((10, 1)) * 0.6
+        external_variables = {
+            "Volume-averaged cell temperature": T_av,
+            "X-averaged negative particle concentration": c_s_n_av,
+        }
+
+        # Step
+        dt = 0.1
+        sol_step = None
+        for _ in range(5):
+            sol_step = solver.step(
+                sol_step, model, dt, external_variables=external_variables
+            )
+        np.testing.assert_array_equal(
+            sol_step.inputs["Volume-averaged cell temperature"],
+            np.zeros((1, len(sol_step.t))),
+        )
+        np.testing.assert_array_equal(
+            sol_step.inputs["X-averaged negative particle concentration"],
+            np.ones((mesh["negative particle"][0].npts, len(sol_step.t))) * 0.6,
+        )
+
+        # Solve
+        t_eval = np.linspace(0, 3600)
+        sol = solver.solve(model, t_eval, external_variables=external_variables)
+        np.testing.assert_array_equal(
+            sol.inputs["Volume-averaged cell temperature"], np.zeros((1, len(sol.t))),
+        )
+        np.testing.assert_array_equal(
+            sol.inputs["X-averaged negative particle concentration"],
+            np.ones((mesh["negative particle"][0].npts, len(sol.t))) * 0.6,
+        )
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/unit/test_simulation.py b/tests/unit/test_simulation.py
index 7cb8fc40a..03068a4ce 100644
--- a/tests/unit/test_simulation.py
+++ b/tests/unit/test_simulation.py
@@ -235,16 +235,30 @@ class TestSimulation(unittest.TestCase):
         self.assertIsInstance(c_e, np.ndarray)
 
     def test_set_external_variable(self):
-        model_options = {"thermal": "lumped", "external submodels": ["thermal"]}
+        model_options = {
+            "thermal": "lumped",
+            "external submodels": ["thermal", "negative particle"],
+        }
         model = pybamm.lithium_ion.SPMe(model_options)
         sim = pybamm.Simulation(model)
 
         T_av = 0
+        c_s_n_av = np.ones((10, 1)) * 0.5
+        external_variables = {
+            "Volume-averaged cell temperature": T_av,
+            "X-averaged negative particle concentration": c_s_n_av,
+        }
 
-        dt = 0.001
+        # Step
+        dt = 0.1
+        for _ in range(5):
+            sim.step(dt, external_variables=external_variables)
+        sim.plot(testing=True)
 
-        external_variables = {"Volume-averaged cell temperature": T_av}
-        sim.step(dt, external_variables=external_variables)
+        # Solve
+        t_eval = np.linspace(0, 3600)
+        sim.solve(t_eval, external_variables=external_variables)
+        sim.plot(testing=True)
 
     def test_step(self):
 
@@ -301,7 +315,7 @@ class TestSimulation(unittest.TestCase):
         self.assertEqual(sim.solution.t[1], dt / tau)
         self.assertEqual(sim.solution.t[2], 2 * dt / tau)
         np.testing.assert_array_equal(
-            sim.solution.inputs["Current function [A]"], np.array([1, 1, 2])
+            sim.solution.inputs["Current function [A]"], np.array([[1, 1, 2]])
         )
 
     def test_save_load(self):
diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py
index c1615e2f8..1f2c6792a 100644
--- a/tests/unit/test_solvers/test_solution.py
+++ b/tests/unit/test_solvers/test_solution.py
@@ -48,7 +48,9 @@ class TestSolution(unittest.TestCase):
         np.testing.assert_array_equal(sol1.y, np.concatenate([y1, y2[:, 1:]], axis=1))
         np.testing.assert_array_equal(
             sol1.inputs["a"],
-            np.concatenate([1 * np.ones_like(t1), 2 * np.ones_like(t2[1:])]),
+            np.concatenate([1 * np.ones_like(t1), 2 * np.ones_like(t2[1:])])[
+                np.newaxis, :
+            ],
         )
 
         # Test sub-solutions
@@ -57,11 +59,11 @@ class TestSolution(unittest.TestCase):
         np.testing.assert_array_equal(sol1.sub_solutions[1].t, t2)
         self.assertEqual(sol1.sub_solutions[0].model, sol1.model)
         np.testing.assert_array_equal(
-            sol1.sub_solutions[0].inputs["a"], 1 * np.ones_like(t1)
+            sol1.sub_solutions[0].inputs["a"], 1 * np.ones_like(t1)[np.newaxis, :]
         )
         self.assertEqual(sol1.sub_solutions[1].model, sol2.model)
         np.testing.assert_array_equal(
-            sol1.sub_solutions[1].inputs["a"], 2 * np.ones_like(t2)
+            sol1.sub_solutions[1].inputs["a"], 2 * np.ones_like(t2)[np.newaxis, :]
         )
 
     def test_total_time(self):