Assembling the jacobian in mixed space

[bjoni]

Hi,

in my attempt to reproduce the plasticity demo from dolfiny in dolfinx, I ran into convergence problems of NewtonSolver (which I used instead of the snes block solver). For debugging, I tried to have a look at the jacobian, but can’t because of ArityMismatch.

This is my MWE:

from dolfinx import log
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.nls.petsc import NewtonSolver
import numpy as np
import ufl
from mpi4py import MPI
from dolfinx import fem, mesh
import basix
from dolfinx.fem import locate_dofs_geometrical
from basix.ufl import mixed_element
import dolfinx

# Define mesh
L = 20.0
domain = mesh.create_box(MPI.COMM_WORLD, [[0.0, 0.0, 0.0], [L, 1, 1]], [1, 1, 1], mesh.CellType.tetrahedron)
Ue = basix.ufl.element("P", domain.basix_cell(), 1, shape=(domain.geometry.dim,))
He = basix.ufl.quadrature_element(domain.basix_cell(), value_shape=(), degree=1)
Te = basix.ufl.blocked_element(He, shape=(domain.geometry.dim, domain.geometry.dim), symmetry=True)

V_el = mixed_element([Ue, Te]) # u, P
V = fem.functionspace(domain, V_el)

# Define boundaries
V0, _ = V.sub(0).collapse()
u_L = fem.Function(V0)
u_L.x.array[:] = 0.
u_R = fem.Function(V0)
u_R.x.array[1] = 0.1

dofs_L = locate_dofs_geometrical((V.sub(0), V0), lambda x: np.isclose(x[0], 0))
bc_L = fem.dirichletbc(u_L, dofs_L, V.sub(0))
dofs_R = locate_dofs_geometrical((V.sub(0), V0), lambda x: np.isclose(x[0], L))
bc_R = fem.dirichletbc(u_R, dofs_R, V.sub(0))
bcs = [bc_L, bc_R]

# Define functions
w = fem.Function(V)
u, P = ufl.split(w)
u0, P0 = ufl.split(w)
δw = ufl.TestFunction(V)
δu, δP = ufl.split(δw)

#Define physics
mu = fem.Constant(domain,  100.0)
la = fem.Constant(domain,  10.)
Sy = fem.Constant(domain,  0.300)  

def rJ2(A):
    """Square root of J2 invariant of tensor A"""
    J2 = 1 / 2 * ufl.inner(A, A)
    rJ2 = ufl.sqrt(J2)
    return ufl.conditional(rJ2 < 1.0e-12, 0.0, rJ2)

I = ufl.Identity(3)  # noqa: E741
F = I + ufl.grad(u) 
E = 1 / 2 * (F.T * F - I) 
E_el = E - P  
S = 2 * mu * E_el + la * ufl.tr(E_el) * I  
S = ufl.variable(S)
f = ufl.sqrt(3) * rJ2(ufl.dev(S)) - (Sy) 
g = f
dgdS = ufl.diff(g, S)
S = S.expression()
δE = ufl.derivative(E, w, δw)
dλ = ufl.max_value(f, 0)
dx = ufl.Measure("dx", domain=domain, metadata={"quadrature_degree": 4})
form = ( ufl.inner(δE, S) * dx + ufl.inner(δP, (P - P0) - dλ * dgdS) * dx)

J = ufl.derivative(form, w, δw)
jacobian = dolfinx.fem.form(J)
# A = dolfinx.fem.petsc.assemble_matrix(jacobian)

# def convertToDense(A_petsc):
#     """
#     Convert the PETSc matrix to a dense numpy array 
#     for debugging purposes
#     """
#     A_petsc.assemble()
#     A_dense = A_petsc.convert("dense")
#     # return (len(A_dense.getDenseArray()), len(A_dense.getDenseArray()[0]))
#     return A_dense.getDenseArray()

# print(convertToDense(A))
# exit()


# log.set_log_level(log.LogLevel.INFO)
# problem = NonlinearProblem(form, w, bcs)
# solver = NewtonSolver(domain.comm, problem)
# solver.convergence_criterion = "incremental"
# solver.rtol = 1e-6
# solver.report = True

# num_its, converged = solver.solve(w)
# assert (converged)

Which produces the following errror:

Traceback (most recent call last):
  File "/home/bjoern/examples/plasticity_beam_devel_perfect_plast_mwe.py", line 70, in <module>
    jacobian = dolfinx.fem.form(J)
               ^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/fem/forms.py", line 337, in form
    return _create_form(form)
           ^^^^^^^^^^^^^^^^^^
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/fem/forms.py", line 331, in _create_form
    return _form(form)
           ^^^^^^^^^^^
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/fem/forms.py", line 254, in _form
    ufcx_form, module, code = jit.ffcx_jit(
                              ^^^^^^^^^^^^^
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/jit.py", line 62, in mpi_jit
    return local_jit(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/petsc/lib/python3/dist-packages/dolfinx/jit.py", line 212, in ffcx_jit
    r = ffcx.codegeneration.jit.compile_forms([ufl_object], options=p_ffcx, **p_jit)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ffcx/codegeneration/jit.py", line 205, in compile_forms
    impl = _compile_objects(
           ^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ffcx/codegeneration/jit.py", line 330, in _compile_objects
    _, code_body = ffcx.compiler.compile_ufl_objects(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ffcx/compiler.py", line 108, in compile_ufl_objects
    analysis = analyze_ufl_objects(ufl_objects, options["scalar_type"])  # type: ignore
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ffcx/analysis.py", line 94, in analyze_ufl_objects
    form_data = tuple(_analyze_form(form, scalar_type) for form in forms)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ffcx/analysis.py", line 94, in <genexpr>
    form_data = tuple(_analyze_form(form, scalar_type) for form in forms)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ffcx/analysis.py", line 180, in _analyze_form
    form_data = ufl.algorithms.compute_form_data(
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ufl/algorithms/compute_form_data.py", line 427, in compute_form_data
    check_form_arity(preprocessed_form, self.original_form.arguments(), complex_mode)
  File "/usr/lib/python3/dist-packages/ufl/algorithms/check_arities.py", line 213, in check_form_arity
    check_integrand_arity(itg.integrand(), arguments, complex_mode)
  File "/usr/lib/python3/dist-packages/ufl/algorithms/check_arities.py", line 194, in check_integrand_arity
    arg_tuples = map_expr_dag(rules, expr, compress=False)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ufl/corealg/map_dag.py", line 35, in map_expr_dag
    (result,) = map_expr_dags(
                ^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ufl/corealg/map_dag.py", line 103, in map_expr_dags
    r = handlers[v._ufl_typecode_](v, *[vcache[u] for u in v.ufl_operands])
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/ufl/algorithms/check_arities.py", line 78, in product
    raise ArityMismatch(
ufl.algorithms.check_arities.ArityMismatch: Multiplying expressions with overlapping form argument number 0, argument is v_0.

It’s Fenicx 0.9 on WSL Ubuntu / Windows. Many thanks!

[dokken]

@bjoni

u, P = ufl.split(w)
u0, P0 = ufl.split(w)

This looks very suspicious, as you later use:

@bjoni

ufl.inner(δP, (P - P0) - dλ * dgdS) * dx

where P-P0 always will be zero.

Secondly:

@bjoni

J = ufl.derivative(form, w, δw)
jacobian = dolfinx.fem.form(J)

should be:

dz = ufl.TrialFunction(V)
J = ufl.derivative(form, w, dz)

as you can’t differentiate in the same direction (testfunction), as the problem would be non-linear in its arguments then.

[bjoni]

Thank you, it works!! I updated the function definitions to

w = fem.Function(V)
w0 = fem.Function(V)
u, P = ufl.split(w)
u0, P0 = ufl.split(w0)

and the jacobian definition as you said. Here is the working example:

from dolfinx import log
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.nls.petsc import NewtonSolver
import numpy as np
import ufl
from mpi4py import MPI
from dolfinx import fem, mesh
import basix
from dolfinx.fem import locate_dofs_geometrical
from basix.ufl import mixed_element
import dolfinx

# Define mesh
L = 20.0
domain = mesh.create_box(MPI.COMM_WORLD, [[0.0, 0.0, 0.0], [L, 1, 1]], [1, 1, 1], mesh.CellType.tetrahedron)
Ue = basix.ufl.element("P", domain.basix_cell(), 1, shape=(domain.geometry.dim,))
He = basix.ufl.quadrature_element(domain.basix_cell(), value_shape=(), degree=1)
Te = basix.ufl.blocked_element(He, shape=(domain.geometry.dim, domain.geometry.dim), symmetry=True)

V_el = mixed_element([Ue, Te]) # u, P
V = fem.functionspace(domain, V_el)

# Define boundaries
V0, _ = V.sub(0).collapse()
u_L = fem.Function(V0)
u_L.x.array[:] = 0.
u_R = fem.Function(V0)
u_R.x.array[1] = 0.1

dofs_L = locate_dofs_geometrical((V.sub(0), V0), lambda x: np.isclose(x[0], 0))
bc_L = fem.dirichletbc(u_L, dofs_L, V.sub(0))
dofs_R = locate_dofs_geometrical((V.sub(0), V0), lambda x: np.isclose(x[0], L))
bc_R = fem.dirichletbc(u_R, dofs_R, V.sub(0))
bcs = [bc_L, bc_R]

# Define functions
w = fem.Function(V)
w0 = fem.Function(V)
u, P = ufl.split(w)
u0, P0 = ufl.split(w0)
δw = ufl.TestFunction(V)
δu, δP = ufl.split(δw)

#Define physics
mu = fem.Constant(domain,  100.0)
la = fem.Constant(domain,  10.)
Sy = fem.Constant(domain,  0.300)  

def rJ2(A):
    """Square root of J2 invariant of tensor A"""
    J2 = 1 / 2 * ufl.inner(A, A)
    rJ2 = ufl.sqrt(J2)
    return ufl.conditional(rJ2 < 1.0e-12, 0.0, rJ2)

I = ufl.Identity(3)  # noqa: E741
F = I + ufl.grad(u) 
E = 1 / 2 * (F.T * F - I) 
E_el = E - P  
S = 2 * mu * E_el + la * ufl.tr(E_el) * I  
S = ufl.variable(S)
f = ufl.sqrt(3) * rJ2(ufl.dev(S)) - (Sy) 
g = f
dgdS = ufl.diff(g, S)
S = S.expression()
δE = ufl.derivative(E, w, δw)
dλ = ufl.max_value(f, 0)
dx = ufl.Measure("dx", domain=domain, metadata={"quadrature_degree": 4})
form = ( ufl.inner(δE, S) * dx + ufl.inner(δP, (P - P0) - dλ * dgdS) * dx)

dz = ufl.TrialFunction(V)
J = ufl.derivative(form, w, dz)

jacobian = dolfinx.fem.form(J)
A = dolfinx.fem.petsc.assemble_matrix(jacobian)

def convertToDense(A_petsc):
    """
    Convert the PETSc matrix to a dense numpy array 
    for debugging purposes
    """
    A_petsc.assemble()
    A_dense = A_petsc.convert("dense")
    # return (len(A_dense.getDenseArray()), len(A_dense.getDenseArray()[0]))
    return A_dense.getDenseArray()


print(convertToDense(A))

log.set_log_level(log.LogLevel.INFO)
problem = NonlinearProblem(form, w, bcs)
solver = NewtonSolver(domain.comm, problem)
solver.convergence_criterion = "incremental"
solver.rtol = 1e-6
solver.report = True

num_its, converged = solver.solve(w)
assert (converged)