Substituting a differential equation into an expression

edit

This is tricky, but it can be done. The idea is to replace (in the expression tree) derivatives of $f$ (with at least one derivative with respect to $t$) by the right-hand side of the equation, substituted into itself as many times as there are derivatives with respect to $t$ minus one, and finally replacing $f$ by the derivative of $f$ with respect to the remaining variables (if any).

My attempt to implement this (based on SubstituteFunction) is below, with an example (your example) in the docstring.

from sage.symbolic.expression_conversions import ExpressionTreeWalker
from sage.symbolic.operators import FDerivativeOperator

class SubstituteEvolutionaryPDE(ExpressionTreeWalker):
    def __init__(self, pde, t):
        """
        A class that walks the tree and replaces derivatives of `f` with
        respect to `t` by the right-hand side of an evolutionary PDE.
        EXAMPLES::
            sage: var('x,t'); f = function('f'); g = function('g')
            sage: pde = diff(f(t,x),t) == 1/2*diff(f(t,x),x,x)
            sage: s = SubstituteEvolutionaryPDE(pde, t)
            sage: h = g(diff(f(t,x),t), diff(f(t,x),t,x))
            sage: s(h)
            g(1/2*diff(f(t, x), x, x), 1/2*diff(f(t, x), x, x, x))

        ASSUMPTION::

        ``pde`` is of the form ``diff(f(t,...),t) == ...`` and the first
        argument of ``f`` is always ``t``.
        """
        f_expr = integrate(pde.lhs(), t)
        self.f = f_expr.operator()
        self.args = f_expr.operands()
        self.rhs = pde.rhs()

    def derivative(self, ex, operator):
        if operator.function() == self.f and 0 in operator.parameter_set(): # f is differentiated with respect to t
            t_derivatives = [p for p in operator.parameter_set() if p == 0] # (assumes t is the first argument)
            result = self.rhs
            for _ in range(len(t_derivatives)-1):
                result = result.substitute_function(self.f, self.rhs.function(*self.args))
            if len(t_derivatives) < len(operator.parameter_set()): # derivatives w.r.t variables other than t
                other_derivatives = [p for p in operator.parameter_set() if not p == 0] # (assumes t is the first argument)
                new_operator = FDerivativeOperator(operator.function(), other_derivatives)
                result = result.substitute_function(self.f, new_operator(*[self(_) for _ in ex.operands()]).function(*self.args))
            return result
        else:
            return operator(*[self(_) for _ in ex.operands()])

Edit: updated to use substitute_function correctly (the second argument should also be a function).