Repository: pybamm-team/PyBaMM

Issue:
Better errors for ProcessedVariable
Raise clear error when the wrong arguments are passed to `ProcessedVariable.__call__`. Currently a bit obscure when it fails.

Discussion:

Patch:
diff --git a/pybamm/processed_variable.py b/pybamm/processed_variable.py
index 838e87f89..4722a5193 100644
--- a/pybamm/processed_variable.py
+++ b/pybamm/processed_variable.py
@@ -315,19 +315,33 @@ class ProcessedVariable(object):
         if self.dimensions == 1:
             return self._interpolation_function(t)
         elif self.dimensions == 2:
-            if self.scale == "micro":
+            return self.call_2D(t, x, r)
+        elif self.dimensions == 3:
+            return self.call_3D(t, x, r)
+
+    def call_2D(self, t, x, r):
+        "Evaluate a 2D variable"
+        if self.scale == "micro":
+            if r is not None:
                 return self._interpolation_function(t, r)
             else:
+                raise ValueError("r cannot be None for microscale variable")
+        else:
+            if x is not None:
                 return self._interpolation_function(t, x)
-        elif self.dimensions == 3:
-            if isinstance(x, np.ndarray):
-                if isinstance(r, np.ndarray) and isinstance(t, np.ndarray):
-                    x = x[:, np.newaxis, np.newaxis]
-                    r = r[:, np.newaxis]
-                elif isinstance(r, np.ndarray) or isinstance(t, np.ndarray):
-                    x = x[:, np.newaxis]
             else:
-                if isinstance(r, np.ndarray) and isinstance(t, np.ndarray):
-                    r = r[:, np.newaxis]
+                raise ValueError("x cannot be None for macroscale variable")
+
+    def call_3D(self, t, x, r):
+        "Evaluate a 3D variable"
+        if isinstance(x, np.ndarray):
+            if isinstance(r, np.ndarray) and isinstance(t, np.ndarray):
+                x = x[:, np.newaxis, np.newaxis]
+                r = r[:, np.newaxis]
+            elif isinstance(r, np.ndarray) or isinstance(t, np.ndarray):
+                x = x[:, np.newaxis]
+        else:
+            if isinstance(r, np.ndarray) and isinstance(t, np.ndarray):
+                r = r[:, np.newaxis]
 
-            return self._interpolation_function((x, r, t))
+        return self._interpolation_function((x, r, t))


Test patch:
diff --git a/tests/unit/test_processed_variable.py b/tests/unit/test_processed_variable.py
index eaabbc033..7db030a2e 100644
--- a/tests/unit/test_processed_variable.py
+++ b/tests/unit/test_processed_variable.py
@@ -305,6 +305,42 @@ class TestProcessedVariable(unittest.TestCase):
         ):
             pybamm.ProcessedVariable(mat, t, y, disc.mesh)
 
+    def test_call_failure(self):
+        # x domain
+        var = pybamm.Variable("var x", domain=["negative electrode", "separator"])
+        x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"])
+        disc = tests.get_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        x_sol = disc.process_symbol(x).entries[:, 0]
+        var_sol = disc.process_symbol(var)
+        t_sol = np.linspace(0, 1)
+        y_sol = x_sol[:, np.newaxis] * np.linspace(0, 5)
+
+        processed_var = pybamm.ProcessedVariable(var_sol, t_sol, y_sol, mesh=disc.mesh)
+        with self.assertRaisesRegex(
+            ValueError, "x cannot be None for macroscale variable"
+        ):
+            processed_var(0)
+
+        # r domain
+        var = pybamm.Variable("var r", domain=["negative particle"])
+        r = pybamm.SpatialVariable("r", domain=["negative particle"])
+        disc = tests.get_discretisation_for_testing()
+        disc.set_variable_slices([var])
+        r_sol = disc.process_symbol(r).entries[:, 0]
+        var_sol = disc.process_symbol(var)
+        y_sol = r_sol[:, np.newaxis] * np.linspace(0, 5)
+
+        processed_var = pybamm.ProcessedVariable(var_sol, t_sol, y_sol, mesh=disc.mesh)
+        with self.assertRaisesRegex(
+            ValueError, "r cannot be None for microscale variable"
+        ):
+            processed_var(0)
+        with self.assertRaisesRegex(
+            ValueError, "r cannot be None for microscale variable"
+        ):
+            processed_var(0, 1)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")