Repository: pybamm-team/PyBaMM

Issue:
process space-only variables 
At the moment `ProcessedVariable` only works if `len(t)>1` (except for the special case of  2D finite elements, which can handle "fixed t" cases). It would be good to set it up so that any variables that are functions of space only can be processed.

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7a68ee6a5..59560dc9d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## Features
 
+-   Allowed `ProcessedVariable` to handle cases where `len(solution.t)=1` ([#1020](https://github.com/pybamm-team/PyBaMM/pull/1020))
 -   Added `BackwardIndefiniteIntegral` symbol ([#1014](https://github.com/pybamm-team/PyBaMM/pull/1014))
 -   Added `plot` and `plot2D` to enable easy plotting of `pybamm.Array` objects ([#1008](https://github.com/pybamm-team/PyBaMM/pull/1008))
 -   Updated effective current collector models and added example notebook ([#1007](https://github.com/pybamm-team/PyBaMM/pull/1007))
@@ -41,6 +42,7 @@
 
 ## Breaking changes
 
+-   For variables discretised using finite elements the result returned by calling `ProcessedVariable` is now transposed ([#1020](https://github.com/pybamm-team/PyBaMM/pull/1020))
 -   Renamed "surface area density" to "surface area to volume ratio" ([#975](https://github.com/pybamm-team/PyBaMM/pull/975))
 -   Replaced "reaction rate" with "exchange-current density" ([#975](https://github.com/pybamm-team/PyBaMM/pull/975))
 -   Changed the implementation of reactions in submodels ([#948](https://github.com/pybamm-team/PyBaMM/pull/948))
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
index ad34a98ff..4b4286b4b 100644
--- a/pybamm/solvers/base_solver.py
+++ b/pybamm/solvers/base_solver.py
@@ -616,6 +616,15 @@ class BaseSolver(object):
                 timer.format(solution.total_time),
             )
         )
+
+        # Raise error if solution only contains one timestep (except for algebraic
+        # solvers, where we may only expect one time in the solution)
+        if self.algebraic_solver is False and len(solution.t) == 1:
+            raise pybamm.SolverError(
+                "Solution time vector has length 1. "
+                "Check whether simulation terminated too early."
+            )
+
         return solution
 
     def step(
diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py
index 9133a6e3d..f324b0068 100644
--- a/pybamm/solvers/processed_variable.py
+++ b/pybamm/solvers/processed_variable.py
@@ -7,6 +7,26 @@ import pybamm
 import scipy.interpolate as interp
 
 
+def make_interp2D_fun(input, interpolant):
+    """
+    Calls and returns a 2D interpolant of the correct shape depending on the
+    shape of the input
+    """
+    first_dim, second_dim, _ = input
+    if isinstance(first_dim, np.ndarray) and isinstance(second_dim, np.ndarray):
+        first_dim = first_dim[:, 0, 0]
+        second_dim = second_dim[:, 0]
+        return interpolant(second_dim, first_dim)
+    elif isinstance(first_dim, np.ndarray):
+        first_dim = first_dim[:, 0]
+        return interpolant(second_dim, first_dim)[:, 0]
+    elif isinstance(second_dim, np.ndarray):
+        second_dim = second_dim[:, 0]
+        return interpolant(second_dim, first_dim)
+    else:
+        return interpolant(second_dim, first_dim)[0]
+
+
 class ProcessedVariable(object):
     """
     An object that can be evaluated at arbitrary (scalars or vectors) t and x, and
@@ -55,48 +75,11 @@ class ProcessedVariable(object):
             and "current collector" in self.domain
             and isinstance(self.mesh[0], pybamm.ScikitSubMesh2D)
         ):
-            if len(solution.t) == 1:
-                # space only (steady solution)
-                self.initialise_2D_fixed_t_scikit_fem()
-            else:
-                self.initialise_2D_scikit_fem()
+            self.initialise_2D_scikit_fem()
 
         # check variable shape
         else:
-            if len(solution.t) == 1:
-                # Implementing a workaround for 0D and 1D variables. Processing
-                # variables that are functions of space only needs to be implemented
-                # properly, see #1006
-                if (
-                    isinstance(self.base_eval, numbers.Number)
-                    or len(self.base_eval.shape) == 0
-                    or self.base_eval.shape[0] == 1
-                ):
-                    # Scalar value
-                    t = self.t_sol
-                    u = self.u_sol
-                    inputs = {name: inp[0] for name, inp in self.inputs.items()}
-
-                    entries = self.base_variable.evaluate(t, u, inputs=inputs)
-
-                    def fun(t):
-                        return entries
-
-                    self._interpolation_function = fun
-                    self.entries = entries
-                    self.dimensions = 0
-                else:
-                    # 1D function of space only
-                    n = self.mesh[0].npts
-                    base_shape = self.base_eval.shape[0]
-                    if base_shape in [n, n + 1]:
-                        self.initialise_1D(fixed_t=True)
-                    else:
-                        raise pybamm.SolverError(
-                            "Solution time vector must have length > 1. Check whether "
-                            "simulation terminated too early."
-                        )
-            elif (
+            if (
                 isinstance(self.base_eval, numbers.Number)
                 or len(self.base_eval.shape) == 0
                 or self.base_eval.shape[0] == 1
@@ -141,9 +124,21 @@ class ProcessedVariable(object):
                 entries[idx] = self.base_variable.evaluate(t, u, inputs=inputs)
 
         # No discretisation provided, or variable has no domain (function of t only)
-        self._interpolation_function = interp.interp1d(
-            self.t_sol, entries, kind="linear", fill_value=np.nan, bounds_error=False
-        )
+        if len(self.t_sol) == 1:
+            # Variable is just a scalar value, but we need to create a callable
+            # function to be consitent with other processed variables
+            def fun(t):
+                return entries
+
+            self._interpolation_function = fun
+        else:
+            self._interpolation_function = interp.interp1d(
+                self.t_sol,
+                entries,
+                kind="linear",
+                fill_value=np.nan,
+                bounds_error=False,
+            )
 
         self.entries = entries
         self.dimensions = 0
@@ -208,19 +203,22 @@ class ProcessedVariable(object):
         self.internal_boundaries = self.mesh[0].internal_boundaries
 
         # set up interpolation
-        # note that the order of 't' and 'space' is the reverse of what you'd expect
-        # TODO: fix processing when variable is a function of space only
-        if fixed_t:
-            # Hack for 1D space only
+        if len(self.t_sol) == 1:
+            # function of space only
             interpolant = interp.interp1d(
                 space, entries_for_interp[:, 0], kind="linear", fill_value=np.nan
             )
 
             def interp_fun(t, z):
-                return interpolant(z)[:, np.newaxis]
+                if isinstance(z, np.ndarray):
+                    return interpolant(z)[:, np.newaxis]
+                else:
+                    return interpolant(z)
 
             self._interpolation_function = interp_fun
         else:
+            # function of space and time. Note that the order of 't' and 'space'
+            # is the reverse of what you'd expect
             self._interpolation_function = interp.interp2d(
                 self.t_sol, space, entries_for_interp, kind="linear", fill_value=np.nan
             )
@@ -297,40 +295,29 @@ class ProcessedVariable(object):
         self.second_dim_pts = second_dim_pts
 
         # set up interpolation
-        self._interpolation_function = interp.RegularGridInterpolator(
-            (first_dim_pts, second_dim_pts, self.t_sol),
-            entries,
-            method="linear",
-            fill_value=np.nan,
-        )
-
-    def initialise_2D_fixed_t_scikit_fem(self):
-        y_sol = self.mesh[0].edges["y"]
-        len_y = len(y_sol)
-        z_sol = self.mesh[0].edges["z"]
-        len_z = len(z_sol)
-
-        # Evaluate the base_variable
-        inputs = {name: inp[0] for name, inp in self.inputs.items()}
+        if len(self.t_sol) == 1:
+            # function of space only. Note the order of the points is the reverse
+            # of what you'd expect
+            interpolant = interp.interp2d(
+                second_dim_pts,
+                first_dim_pts,
+                entries[:, :, 0],
+                kind="linear",
+                fill_value=np.nan,
+            )
 
-        entries = np.reshape(
-            self.base_variable.evaluate(0, self.u_sol, inputs=inputs), [len_y, len_z]
-        )
+            def interp_fun(input):
+                return make_interp2D_fun(input, interpolant)
 
-        # assign attributes for reference
-        self.entries = entries
-        self.dimensions = 2
-        self.y_sol = y_sol
-        self.z_sol = z_sol
-        self.first_dimension = "y"
-        self.second_dimension = "z"
-        self.first_dim_pts = y_sol
-        self.second_dim_pts = z_sol
-
-        # set up interpolation
-        self._interpolation_function = interp.interp2d(
-            y_sol, z_sol, entries, kind="linear", fill_value=np.nan
-        )
+            self._interpolation_function = interp_fun
+        else:
+            # function of space and time.
+            self._interpolation_function = interp.RegularGridInterpolator(
+                (first_dim_pts, second_dim_pts, self.t_sol),
+                entries,
+                method="linear",
+                fill_value=np.nan,
+            )
 
     def initialise_2D_scikit_fem(self):
         y_sol = self.mesh[0].edges["y"]
@@ -349,11 +336,15 @@ class ProcessedVariable(object):
                 eval_and_known_evals = self.base_variable.evaluate(
                     t, u, inputs=inputs, known_evals=self.known_evals[t]
                 )
-                entries[:, :, idx] = np.reshape(eval_and_known_evals[0], [len_y, len_z])
+                entries[:, :, idx] = np.reshape(
+                    eval_and_known_evals[0], [len_y, len_z], order="F"
+                )
                 self.known_evals[t] = eval_and_known_evals[1]
             else:
                 entries[:, :, idx] = np.reshape(
-                    self.base_variable.evaluate(t, u, inputs=inputs), [len_y, len_z]
+                    self.base_variable.evaluate(t, u, inputs=inputs),
+                    [len_y, len_z],
+                    order="F",
                 )
 
         # assign attributes for reference
@@ -367,23 +358,49 @@ class ProcessedVariable(object):
         self.second_dim_pts = z_sol
 
         # set up interpolation
-        self._interpolation_function = interp.RegularGridInterpolator(
-            (y_sol, z_sol, self.t_sol), entries, method="linear", fill_value=np.nan
-        )
+        if len(self.t_sol) == 1:
+            # function of space only. Note the order of the points is the reverse
+            # of what you'd expect
+            interpolant = interp.interp2d(
+                z_sol, y_sol, entries, kind="linear", fill_value=np.nan
+            )
+
+            def interp_fun(input):
+                return make_interp2D_fun(input, interpolant)
+
+            self._interpolation_function = interp_fun
+        else:
+            # function of space and time.
+            self._interpolation_function = interp.RegularGridInterpolator(
+                (y_sol, z_sol, self.t_sol), entries, method="linear", fill_value=np.nan
+            )
 
     def __call__(self, t=None, x=None, r=None, y=None, z=None, warn=True):
         """
         Evaluate the variable at arbitrary t (and x, r, y and/or z), using interpolation
         """
+        # If t is None and there is only one value of time in the soluton (i.e.
+        # the solution is independent of time) then we set t equal to the value
+        # stored in the solution. If the variable is constant (doesn't depend on
+        # time) evaluate arbitrarily at the first value of t. Otherwise, raise
+        # an error
+        if t is None:
+            if len(self.t_sol) == 1:
+                t = self.t_sol
+            elif self.base_variable.is_constant():
+                t = self.t_sol[0]
+            else:
+                raise ValueError(
+                    "t cannot be None for variable {}".format(self.base_variable)
+                )
+
+        # Call interpolant of correct spatial dimension
         if self.dimensions == 0:
             out = self._interpolation_function(t)
         elif self.dimensions == 1:
             out = self.call_1D(t, x, r, z)
         elif self.dimensions == 2:
-            if t is None:
-                out = self._interpolation_function(y, z)
-            else:
-                out = self.call_2D(t, x, r, y, z)
+            out = self.call_2D(t, x, r, y, z)
         if warn is True and np.isnan(out).any():
             pybamm.logger.warning(
                 "Calling variable outside interpolation range (returns 'nan')"


Test patch:
diff --git a/tests/unit/test_solvers/test_base_solver.py b/tests/unit/test_solvers/test_base_solver.py
index 885032e5c..99e9ed3aa 100644
--- a/tests/unit/test_solvers/test_base_solver.py
+++ b/tests/unit/test_solvers/test_base_solver.py
@@ -71,6 +71,17 @@ class TestBaseSolver(unittest.TestCase):
         ):
             solver.solve(model, np.array([1, 2, 3, 2]))
 
+    def test_solution_time_length_fail(self):
+        model = pybamm.BaseModel()
+        v = pybamm.Scalar(1)
+        model.variables = {"v": v}
+        solver = pybamm.DummySolver()
+        t_eval = np.array([0])
+        with self.assertRaisesRegex(
+            pybamm.SolverError, "Solution time vector has length 1"
+        ):
+            solver.solve(model, t_eval)
+
     def test_block_symbolic_inputs(self):
         solver = pybamm.BaseSolver(rtol=1e-2, atol=1e-4)
         model = pybamm.BaseModel()
@@ -203,13 +214,13 @@ class TestBaseSolver(unittest.TestCase):
             solver.calculate_consistent_state(Model())
         solver = pybamm.BaseSolver(root_method="lm")
         with self.assertRaisesRegex(
-            pybamm.SolverError, "Could not find acceptable solution: solver terminated",
+            pybamm.SolverError, "Could not find acceptable solution: solver terminated"
         ):
             solver.calculate_consistent_state(Model())
         # with casadi
         solver = pybamm.BaseSolver(root_method="casadi")
         with self.assertRaisesRegex(
-            pybamm.SolverError, "Could not find acceptable solution: .../casadi",
+            pybamm.SolverError, "Could not find acceptable solution: .../casadi"
         ):
             solver.calculate_consistent_state(Model())
 
diff --git a/tests/unit/test_solvers/test_processed_variable.py b/tests/unit/test_solvers/test_processed_variable.py
index d9e6a1a50..d7c583a26 100644
--- a/tests/unit/test_solvers/test_processed_variable.py
+++ b/tests/unit/test_solvers/test_processed_variable.py
@@ -20,6 +20,14 @@ class TestProcessedVariable(unittest.TestCase):
         processed_var = pybamm.ProcessedVariable(var, pybamm.Solution(t_sol, y_sol))
         np.testing.assert_array_equal(processed_var.entries, t_sol * y_sol[0])
 
+        # scalar value
+        var = y
+        var.mesh = None
+        t_sol = np.array([0])
+        y_sol = np.array([1])[:, np.newaxis]
+        processed_var = pybamm.ProcessedVariable(var, pybamm.Solution(t_sol, y_sol))
+        np.testing.assert_array_equal(processed_var.entries, y_sol[0])
+
     def test_processed_variable_1D(self):
         t = pybamm.t
         var = pybamm.Variable("var", domain=["negative electrode", "separator"])
@@ -59,6 +67,18 @@ class TestProcessedVariable(unittest.TestCase):
             x_s_edge.entries[:, 0], processed_x_s_edge.entries[:, 0]
         )
 
+        # space only
+        eqn = var + x
+        eqn_sol = disc.process_symbol(eqn)
+        t_sol = np.array([0])
+        y_sol = np.ones_like(x_sol)[:, np.newaxis]
+        processed_eqn2 = pybamm.ProcessedVariable(
+            eqn_sol, pybamm.Solution(t_sol, y_sol)
+        )
+        np.testing.assert_array_equal(
+            processed_eqn2.entries, y_sol + x_sol[:, np.newaxis]
+        )
+
     def test_processed_variable_1D_unknown_domain(self):
         x = pybamm.SpatialVariable("x", domain="SEI layer", coord_sys="cartesian")
         geometry = pybamm.Geometry()
@@ -150,6 +170,31 @@ class TestProcessedVariable(unittest.TestCase):
             x_s_edge.entries.flatten(), processed_x_s_edge.entries[:, :, 0].T.flatten()
         )
 
+    def test_processed_variable_2D_space_only(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
+        x = pybamm.SpatialVariable("x", domain=["negative electrode"])
+        r = pybamm.SpatialVariable("r", domain=["negative particle"])
+
+        disc = tests.get_p2d_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries[:, 0]
+        r_sol = disc.process_symbol(r).entries[:, 0]
+        # Keep only the first iteration of entries
+        r_sol = r_sol[: len(r_sol) // len(x_sol)]
+        var_sol = disc.process_symbol(var)
+        t_sol = np.array([0])
+        y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis]
+
+        processed_var = pybamm.ProcessedVariable(var_sol, pybamm.Solution(t_sol, y_sol))
+        np.testing.assert_array_equal(
+            processed_var.entries,
+            np.reshape(y_sol, [len(r_sol), len(x_sol), len(t_sol)]),
+        )
+
     def test_processed_variable_2D_scikit(self):
         var = pybamm.Variable("var", domain=["current collector"])
 
@@ -181,7 +226,7 @@ class TestProcessedVariable(unittest.TestCase):
 
         processed_var = pybamm.ProcessedVariable(var_sol, pybamm.Solution(t_sol, u_sol))
         np.testing.assert_array_equal(
-            processed_var.entries, np.reshape(u_sol, [len(y), len(z)])
+            processed_var.entries, np.reshape(u_sol, [len(y), len(z), len(t_sol)])
         )
 
     def test_processed_var_0D_interpolation(self):
@@ -211,6 +256,19 @@ class TestProcessedVariable(unittest.TestCase):
         np.testing.assert_array_equal(processed_eqn(2), np.nan)
         pybamm.set_logging_level("WARNING")
 
+    def test_processed_var_0D_fixed_t_interpolation(self):
+        y = pybamm.StateVector(slice(0, 1))
+        var = y
+        eqn = 2 * y
+        var.mesh = None
+        eqn.mesh = None
+
+        t_sol = np.array([10])
+        y_sol = np.array([[100]])
+        processed_var = pybamm.ProcessedVariable(eqn, pybamm.Solution(t_sol, y_sol))
+
+        np.testing.assert_array_equal(processed_var(), 200)
+
     def test_processed_var_1D_interpolation(self):
         t = pybamm.t
         var = pybamm.Variable("var", domain=["negative electrode", "separator"])
@@ -250,6 +308,12 @@ class TestProcessedVariable(unittest.TestCase):
         # 2 scalars
         self.assertEqual(processed_eqn(0.5, x_sol[-1]).shape, (1,))
 
+        # test x
+        processed_x = pybamm.ProcessedVariable(
+            disc.process_symbol(x), pybamm.Solution(t_sol, y_sol)
+        )
+        np.testing.assert_array_almost_equal(processed_x(x=x_sol), x_sol[:, np.newaxis])
+
         # On microscale
         r_n = pybamm.Matrix(
             disc.mesh["negative particle"][0].nodes, domain="negative particle"
@@ -261,6 +325,27 @@ class TestProcessedVariable(unittest.TestCase):
             processed_r_n(0, r=np.linspace(0, 1))[:, 0], np.linspace(0, 1)
         )
 
+    def test_processed_var_1D_fixed_t_interpolation(self):
+        var = pybamm.Variable("var", domain=["negative electrode", "separator"])
+        x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"])
+        eqn = var + x
+
+        disc = tests.get_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries[:, 0]
+        eqn_sol = disc.process_symbol(eqn)
+        t_sol = np.array([1])
+        y_sol = x_sol[:, np.newaxis]
+
+        processed_var = pybamm.ProcessedVariable(eqn_sol, pybamm.Solution(t_sol, y_sol))
+
+        # vector
+        np.testing.assert_array_almost_equal(
+            processed_var(x=x_sol), 2 * x_sol[:, np.newaxis]
+        )
+        # scalar
+        np.testing.assert_array_almost_equal(processed_var(x=0.5), 1)
+
     def test_processed_var_2D_interpolation(self):
         var = pybamm.Variable(
             "var",
@@ -324,6 +409,34 @@ class TestProcessedVariable(unittest.TestCase):
             processed_var(t_sol, x_sol, r_sol).shape, (10, 35, 50)
         )
 
+    def test_processed_var_2D_fixed_t_interpolation(self):
+        var = pybamm.Variable(
+            "var",
+            domain=["negative particle"],
+            auxiliary_domains={"secondary": ["negative electrode"]},
+        )
+        x = pybamm.SpatialVariable("x", domain=["negative electrode"])
+        r = pybamm.SpatialVariable("r", domain=["negative particle"])
+
+        disc = tests.get_p2d_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries[:, 0]
+        r_sol = disc.process_symbol(r).entries[:, 0]
+        # Keep only the first iteration of entries
+        r_sol = r_sol[: len(r_sol) // len(x_sol)]
+        var_sol = disc.process_symbol(var)
+        t_sol = np.array([0])
+        y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis]
+
+        processed_var = pybamm.ProcessedVariable(var_sol, pybamm.Solution(t_sol, y_sol))
+        # 2 vectors
+        np.testing.assert_array_equal(processed_var(x=x_sol, r=r_sol).shape, (10, 40))
+        # 1 vector, 1 scalar
+        np.testing.assert_array_equal(processed_var(x=0.2, r=r_sol).shape, (10,))
+        np.testing.assert_array_equal(processed_var(x=x_sol, r=0.5).shape, (40,))
+        # 2 scalars
+        np.testing.assert_array_equal(processed_var(x=0.2, r=0.2).shape, ())
+
     def test_processed_var_2D_secondary_broadcast(self):
         var = pybamm.Variable("var", domain=["negative particle"])
         broad_var = pybamm.SecondaryBroadcast(var, "negative electrode")
@@ -429,18 +542,12 @@ class TestProcessedVariable(unittest.TestCase):
 
         processed_var = pybamm.ProcessedVariable(var_sol, pybamm.Solution(t_sol, u_sol))
         # 2 vectors
-        np.testing.assert_array_equal(
-            processed_var(t=None, y=y_sol, z=z_sol).shape, (15, 15)
-        )
+        np.testing.assert_array_equal(processed_var(y=y_sol, z=z_sol).shape, (15, 15))
         # 1 vector, 1 scalar
-        np.testing.assert_array_equal(
-            processed_var(t=None, y=0.2, z=z_sol).shape, (15, 1)
-        )
-        np.testing.assert_array_equal(
-            processed_var(t=None, y=y_sol, z=0.5).shape, (15,)
-        )
+        np.testing.assert_array_equal(processed_var(y=0.2, z=z_sol).shape, (15,))
+        np.testing.assert_array_equal(processed_var(y=y_sol, z=0.5).shape, (15,))
         # 2 scalars
-        np.testing.assert_array_equal(processed_var(t=None, y=0.2, z=0.2).shape, (1,))
+        np.testing.assert_array_equal(processed_var(t=None, y=0.2, z=0.2).shape, ())
 
     def test_processed_variable_ode_pde_solution(self):
         # without space
@@ -518,18 +625,9 @@ class TestProcessedVariable(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "r cannot be None"):
             processed_var(0, 1)
 
-    def test_solution_too_short(self):
-        t = pybamm.t
-        y = pybamm.StateVector(slice(0, 3))
-        var = t * y
-        disc = tests.get_2p1d_discretisation_for_testing()
-        var.mesh = disc.mesh["current collector"]
-        t_sol = np.array([1])
-        y_sol = np.linspace(0, 5)[:, np.newaxis]
-        with self.assertRaisesRegex(
-            pybamm.SolverError, "Solution time vector must have length > 1"
-        ):
-            pybamm.ProcessedVariable(var, pybamm.Solution(t_sol, y_sol))
+        # t is None but len(solution.t) > 1
+        with self.assertRaisesRegex(ValueError, "t cannot be None"):
+            processed_var()
 
     def test_3D_raises_error(self):
         var = pybamm.Variable(