Derivative functionality

[TomvanSluijs]

Hi,

We want to take the time derivative of a residual by using a chain rule. As shown in the mwe below, we want to take the derivative to the dofs of variable u and multiply by the dofs of the time derivative of u, which is a separate field called dudt. (Also shown in a previous discussion) We want to do this multiple times to obtain the third time derivative of a weak form, as shown in the comments in the script. However, the script does not do exactly what I want. The script gives me this output:

Take derivative
derivative to  udofs
Take derivative
derivative to  udofs
derivative to  dudtdofs
Take derivative
derivative to  udofs
derivative to  dudtdofs
derivative to  ddudtdofs

which indicates that the weak form still contains udofs after taking the first derivative. This is shown in the second derivative where we take the derivative to both udofs and dudtdofs, which should only be dudtdofs. I expected the dependency of udofs to disappears after taking the derivative to udofs. The output that I would expect is:

Take derivative
derivative to  udofs
Take derivative
derivative to  dudtdofs
Take derivative
derivative to  ddudtdofs

Am I using this function incorrectly or is there some other problem? (In this linear example, I could use the replace function to obtain the result that I need, however, I want to extend this to a non-linear example.)

Best,
Tom

from nutils import mesh, function
from nutils.expression_v2 import Namespace
import treelog

# In this script we try automate taking the time derivative of a weak form
# We start with U_weak = ∫_Ω u utest dV from which we try to derive
# dUdt_weak = ∫_Ω dudt dudttest dV , ddUdt_weak = ∫_Ω ddudt ddudttest dV , 
# and dddUdt_weak = ∫_Ω dddudt dddudttest dV for the first, second and third
# timederivative repsectively

def main(nelems: int = 30,
         etype: str = 'square',
         btype: str = 'std',
         degree: int = 2):

    domain, geom = mesh.unitsquare(nelems, etype)

    ns = Namespace()
    ns.x = geom
    ns.define_for('x', gradient='∇', normal='n', jacobians=('dV', 'dS'))
    ns.basis = domain.basis(btype, degree=degree)
    fields = tuple(['u','utest','dudt','dudttest','ddudt','ddudttest','dddudt','dddudttest'])
    for name in fields:
        setattr(ns, name, function.dotarg(name+'dofs',domain.basis(btype,degree=degree)))       

    U_weak = domain.integral('u utest dV' @ ns, degree=degree*2)
    dUdt_weak   = time_derivative_weakform(U_weak,ns)
    ddUdt_weak  = time_derivative_weakform(dUdt_weak,ns)
    dddUdt_weak = time_derivative_weakform(ddUdt_weak,ns)

def time_derivative_weakform(weak_form, ns):
    treelog.user('Take derivative')
    targets = tuple(['u','dudt','ddudt'])
    # Switch test function to test function of time derivative
    replace_test_dttest = {}
    for key in targets:
        if key[-2:] == "dt":
            replace_test_dttest[key+'testdofs']  = function.Argument('d'+key+'testdofs',ns.basis.shape)
        else:
            replace_test_dttest[key+'testdofs']  = function.Argument('d'+key+'dttestdofs',ns.basis.shape)
    weak_form_new_test = weak_form.replace(replace_test_dttest)

    for key in targets:
        try: # If the weak form is not a function of keydofs an error occurs
            weak_form_dkey = weak_form_new_test.derivative(key+'dofs')
            if key[-2:] == "dt": # If the key is already a time derivative
                weak_form_dt_key = weak_form_dkey @ function.Argument('d'+key+'dofs', weak_form_dkey.shape)
                treelog.user('derivative to ',key+'dofs') 
            else: # If the key is not a time derivative 
                weak_form_dt_key = weak_form_dkey @ function.Argument('d'+key+'dtdofs', weak_form_dkey.shape)
                treelog.user('derivative to ',key+'dofs')
            try: # Needed for the summation of different keys.
                weak_form_dt += weak_form_dt_key  
            except:
                weak_form_dt  = weak_form_dt_key
        except:
            pass
    try:
        return weak_form_dt
    except:
        return weak_form

# If the script is executed (as opposed to imported), `nutils.cli.run` calls
# the main function with arguments provided from the command line. For example,
# to keep with the default arguments simply run `python3 laplace.py`. To select
# mixed elements and quadratic basis functions add `python3 laplace.py
# etype=mixed degree=2`.

if __name__ == '__main__':
    from nutils import cli
    cli.run(main)