Compute residual and jacobian with ufl.diff rather than ufl.derivative

Hi everyone,

In the following MWE, I am computing the residual and the jacobian of an energy W with respect to the variable t as:

W = ufl.tr(ufl.grad(m(t))*ufl.grad(m(t)).T)*dv
W_residual_t = ufl.derivative(W, t, t_test)
W_jacobian_t = ufl.derivative(W_residual_t, t, t_trial)

This works as expected.
Now I need to reformulate the problem as follows:

W_residual_t = dW/dt = dW/dm * dm/dt
… similarly for W_jacobian_t

I tried to use ufl.diff, but I ran in errors when I try to build the residual or the jacobian.
In particular, the following lines of code

W = makeW(ufl.variable(t))
dW_dm = ufl.diff(W,ufl.variable(m(t)))
dm_dt = ufl.diff(m(ufl.variable(t)), ufl.variable(t))
W_residual_t = ufl.dot(dW_dm,dm_dt)

give the error: AttributeError: ‘Form’ object has no attribute ‘ufl_shape’.

Can you please help me figure out how can I do this?
I append below the MWE.

Many thanks!

import os
import sys
import numpy as np
import petsc4py
import ufl
from dolfinx import fem, mesh, la
from mpi4py import MPI
from petsc4py import PETSc


domain = mesh.create_box(comm=MPI.COMM_WORLD, points=[[0.0, 0.0, 0.0], [1, 1, 1]], n=[10, 10, 10], cell_type=mesh.CellType.tetrahedron)
dv = ufl.Measure("dx", domain=domain,metadata = {'quadrature_degree' : 4})


V_t = fem.VectorFunctionSpace(domain,("CG", 1), dim=3)
t, t_test, t_trial =  fem.Function(V_t), ufl.TestFunction(V_t), ufl.TrialFunction(V_t)

def m(t):
    return ufl.as_vector([ufl.sin(t[0])*ufl.cos(t[1]), 
                        ufl.sin(t[0])*ufl.sin(t[1]),
                        ufl.cos(t[0])])

def makeW(t):
    return ufl.tr(ufl.grad(m(t))*ufl.grad(m(t)).T)*dv



# -------------------- THIS IS THE PART OF THE CODE THAT I WOULD LIKE TO MODIFY ----------------------

# This works
# W = makeW(t)
# # !!! I would like to replace this with dW/dt = dW/dm*dm/dt
# W_residual_t = ufl.derivative(W, t, t_test)
# # !!! I would like to replace this with d2W/dt2 = dW/dm*d2m/dt2
# W_jacobian_t = ufl.derivative(W_residual_t, t, t_trial)


# This does not work
W = makeW(ufl.variable(t))
dW_dm = ufl.diff(W,ufl.variable(m(t)))
dm_dt = ufl.diff(m(ufl.variable(t)), ufl.variable(t))
W_residual_t = ufl.dot(dW_dm,dm_dt)

# -----------------------------------------------------------------------------------------------------


dofs_domain, dofs_borders = V_t.dofmap.index_map, V_t.dofmap.index_map_bs
b = la.create_petsc_vector(dofs_domain, dofs_borders)
J = fem.petsc.create_matrix(fem.form(W_jacobian_t))


class SNESProblem:

    def __init__(self, F, u, bcs, J=None):
        V = u.function_space
        du = ufl.TrialFunction(V)
        self.L = fem.form(F)
        if J is None:
            self.a = fem.form(ufl.derivative(F, u, du))
        else:
            self.a = fem.form(J)
        self.bcs = bcs
        self._F, self._J = None, None
        self.u = u

    def F(self, snes, x, F):
        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.u.vector)
        self.u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

        with F.localForm() as f_local:
            f_local.set(0.0)
        fem.petsc.assemble_vector(F, self.L)
        fem.petsc.apply_lifting(F, [self.a], bcs=[self.bcs], x0=[x], scale=-1.0)
        F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        fem.petsc.set_bc(F, self.bcs, x, -1.0)
        F.ghostUpdate(addv=PETSc.InsertMode.INSERT,mode=PETSc.ScatterMode.FORWARD)

    def J(self, snes, x, J, P):
        J.zeroEntries()
        fem.petsc.assemble_matrix(J, self.a, bcs=self.bcs)
        J.assemble()


problem = SNESProblem(W_residual_t, t, [], J= W_jacobian_t)
solver_snes = PETSc.SNES().create()
solver_snes.setType("vinewtonrsls")
solver_snes.setFunction(problem.F, b)
solver_snes.setJacobian(problem.J, J)
solver_snes.setTolerances(rtol=1.0e-9, max_it=50)
solver_snes.getKSP().setType("preonly")
solver_snes.getKSP().setTolerances(rtol=1.0e-9)
solver_snes.getKSP().getPC().setType("lu")
solver_snes.solve(None, t.vector)


print("Done")

This product does not make sense, as dW_dm is an integral (it is multiplied by dv, while dm_dt` is a scalar value. At the current time, I dont have the bandwidth to go through the mathematics to figure out what should be changed to create your Jacobian, but this should at least point you towards a solution.

Hi dokken,
thank you for your reply.
In my MWE there is a typo, t is in fact a 2x1 (and not 3x1) vector:

V_t = fem.VectorFunctionSpace(domain,("CG", 1), dim=2)

dm_dt is a 2x3 matrix. One can also write it explictly as:

def dm_dt(t):
    diffm = [[ufl.cos(t[0])*ufl.cos(t[1]), ufl.cos(t[0])*ufl.sin(t[1]), -ufl.sin(t[0])],
            [-ufl.sin(t[0])*ufl.sin(t[1]), ufl.sin(t[0])*ufl.cos(t[1]), 0]]
    return ufl.as_matrix(diffm)

dW_dm is a 3x1 vector: it’s the derivative of a scalar function with respect to a 3x1 vector

If you have time to get back to me I’ll deeply appreciate it.
Thank you