#858 fixes for diff

Author: martinjrobins

pybamm/__init__.py

@@ -108,14 +108,15 @@ def version(formatted=False):
 )
 from .expression_tree.scalar import Scalar
 from .expression_tree.variable import Variable, ExternalVariable, VariableDot
+from .expression_tree.variable import VariableBase
 from .expression_tree.independent_variable import (
     IndependentVariable,
     Time,
     SpatialVariable,
 )
 from .expression_tree.independent_variable import t
 from .expression_tree.vector import Vector
-from .expression_tree.state_vector import StateVector, StateVectorDot
+from .expression_tree.state_vector import StateVectorBase, StateVector, StateVectorDot
 
 from .expression_tree.exceptions import (
     DomainError,

pybamm/expression_tree/symbol.py

@@ -485,6 +485,12 @@ def diff(self, variable):
             return pybamm.Scalar(1)
         elif any(variable.id == x.id for x in self.pre_order()):
             return self._diff(variable)
+        elif variable.id == pybamm.t.id and \
+            any(
+                isinstance(x, (pybamm.VariableBase, pybamm.StateVectorBase))
+                for x in self.pre_order()
+            ):
+            return self._diff(variable)
         else:
             return pybamm.Scalar(0)
 
@@ -770,5 +776,3 @@ def test_shape(self):
             self.shape_for_testing
         except ValueError as e:
             raise pybamm.ShapeError("Cannot find shape (original error: {})".format(e))
-
-