Chain rule not working as expected with Symbolic array elements

[willrae42]

I'm attempting to use Symbolics to solve some perturbation problems involving using symbolic array variables to generate series. However, no matter what I try, I've been unable to get the chain rule to apply the elements of the array. I understand that array differentiation is currently supported, but here I'm only working with elements of the array so it seems different.

Below is an example of the problem I've been running into:

using Symbolics
@variables t T(t) y(T)[1:3]
Dt = Differential(t)
expand_derivatives(Dt(y[1]))

$$ \begin{equation} \frac{\mathrm{d} y\left( T\left( t \right) \right)_{1}}{\mathrm{d}t} \end{equation} $$

Instead, I would expect it to behave more like the following:

@variables x(T)
expand_derivatives(Dt(x))

$$ \begin{equation} \frac{\mathrm{d} x\left( T\left( t \right) \right)}{T\left( t \right)} \frac{\mathrm{d} T\left( t \right)}{\mathrm{d}t} \end{equation} $$

I've been able to get this to work by inserting the following condition into the expand_derivatives function in diff.jl, but I'm not certain this wouldn't break other things or would be useful for anyone else.

        elseif op == getindex
            inner_args = arguments(arguments(arg)[1])
            c = 0
            for a in inner_args
                if isequal(a,D.x)
                    return D(arg)
                else
                    c += Differential(a)(arg)*D(a)
                end
            end
            return expand_derivatives(c)

This change results in:

expand_derivatives(Dt(y[1]))

$$ \begin{equation} \frac{\mathrm{d} y\left( T\left( t \right) \right)_{1}}{T\left( t \right)} \frac{\mathrm{d} T\left( t \right)}{\mathrm{d}t} \end{equation} $$

Original Discourse link: https://discourse.julialang.org/t/chain-rule-not-working-as-expected-with-symbolic-series/123598

[ChrisRackauckas]

That extension seems correct