#858 fixes for diff

pybamm-team · Mar 7, 2020 · 06ccfe8 · 06ccfe8
1 parent c8eeb81
commit 06ccfe8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/pybamm/__init__.py b/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,

diff --git a/pybamm/expression_tree/symbol.py b/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))
-
-