from firedrake import *
from mpi4py import MPI


mesh = RectangleMesh(
        2, 2, 2, 2, comm=MPI.COMM_WORLD, quadrilateral=True
    )

V = VectorFunctionSpace(mesh, "Lagrange", 1)

u = Function(V)

v = TestFunction(V)

rho = Constant(1)
g = Constant(1.0e-1)
f = as_vector([0, -rho * g])
E = 1
nu = 0.0
lambda_ = Constant(E * nu / (1 + nu) / (1 - 2 * nu))
mu = Constant(E / 2.0 / (1 + nu))
Id = Identity(2)  # 2x2 Identity tensor


def epsilon(u):
    # return 0.5 * (grad(u) + grad(u).T)
    return 0.5 * (grad(u) + grad(u).T + grad(u).T * grad(u))


def sigma(u):
    return lambda_ * div(u) * Id + 2 * mu * epsilon(u)

v = TestFunction(V)

F = inner(sigma(u), epsilon(v)) * dx - dot(f, v) * dx

Jr = assemble(derivative(inner(sigma(u), epsilon(v)) * dx, u)).petscmat