How to create a Function type based on a Mixedfunctionspace

Dear everyone

I am coding a elastodynamic problem with piezo field.
I met a problem here. I created a MixedFunctionspace Z_space with displacement vector and potential scale.
Because it is a dynamic problem, I need to create a Function type thing to store the last time step result.
In the below codes: u_old, v_old, a_old will take this role for the last time step t_n.
and du stands for the present time step t_n+1
But when I run the code, fenics told me:
I don’t know why u_old is a ListTensor type, and how should I modify the definition of u_old here?

Time:  0.08
Traceback (most recent call last):
  File "/home/dell-123/fenics-projects/fenics-gmsh/linear_elasticity/dynamic/newmark_C_piezo_without_circuit/Newmark_C_piezo_nocircuit.py", line 356, in <module>
    update_fields(u, u_old, v_old, a_old)
  File "/home/dell-123/fenics-projects/fenics-gmsh/linear_elasticity/dynamic/newmark_C_piezo_without_circuit/Newmark_C_piezo_nocircuit.py", line 285, in update_fields
    u_vec, u0_vec = u.vector(), u_old.vector()
                                 ^^^^^^^^^^^^
AttributeError: 'ListTensor' object has no attribute 'vector'

this is the short version code:

V = VectorFunctionSpace(mesh, "CG", 1) #for displacement
Q = FunctionSpace(mesh, 'CG', 1)       #for potential
Mix_element = MixedElement([V.ufl_element(), Q.ufl_element()])
Z_space = FunctionSpace(mesh, Mix_element)
(du, dphi) = TrialFunctions(Z_space)
(u_, phi_) = TestFunctions(Z_space)

# Current (unknown) 
uphi = Function(Z_space, name="Displacement&Potential")
uphi_list = split(uphi)

# Fields from previous time step (displacement, velocity, acceleration)
u_old = uphi_list[0]  # u_n
v_old = uphi_list[0]  # v_n
a_old = uphi_list[0]  # a_n

def update_fields(u, u_old, v_old, a_old):
    """Update fields at the end of each time step."""

    # Get vectors (references)
    u_vec, u0_vec = u.vector(), u_old.vector()
    v0_vec, a0_vec = v_old.vector(), a_old.vector()
    # phi_vec, phi0_vec = phi.vector(), phi_old.vector()

    # use update functions using vector arguments
    a_vec = update_a(u_vec, u0_vec, v0_vec, a0_vec, ufl=False)
    v_vec = update_v(a_vec, u0_vec, v0_vec, a0_vec, ufl=False)

    # Update (u_old <- u)
    v_old.vector()[:], a_old.vector()[:] = v_vec, a_vec
    u_old.vector()[:] = u.vector()
    #phi_old.vector()[:] = phi.vector()

This is the whole version code:

from dolfin import *
import numpy as np
import matplotlib.pyplot as plt

# Form compiler options
parameters["form_compiler"]["cpp_optimize"] = True
parameters["form_compiler"]["optimize"] = True

# Define mesh
mesh = BoxMesh(Point(0., 0., 0.), Point(1., 0.2, 0.2), 50, 10, 10)


# Sub domain for clamp at left end
def left(x, on_boundary):
    return near(x[0], 0.) and on_boundary


# Sub domain for rotation at right end
def right(x, on_boundary):
    return near(x[0], 1.) and on_boundary


# Sub domain for bottom end
def bottom(x, on_boundary):
    return near(x[2], 0.) and on_boundary


# Sub domain for top end
def top(x, on_boundary):
    return near(x[2], 0.2) and on_boundary


# Mass density
rho = Constant(5550)
g = 9.8

# define elasticity tensor
# define elasticity tensor C
c11 = 226e9
c12 = 125e9
c13 = 124e9
c33 = 216e9
c44 = 44.2e9  # * 2
c66 = (c11 - c12) / 2  # * 2
# 6x6
C = as_matrix([[c11, c12, c13, 0.0, 0.0, 0.0],
               [c12, c11, c13, 0.0, 0.0, 0.0],
               [c13, c13, c33, 0.0, 0.0, 0.0],
               [0.0, 0.0, 0.0, c44, 0.0, 0.0],
               [0.0, 0.0, 0.0, 0.0, c44, 0.0],
               [0.0, 0.0, 0.0, 0.0, 0.0, c66]])

# define piezoelectricity tensor e
e15 = 5.8
e31 = -2.2
e33 = 9.3
# 6x3
e1 = as_matrix([[0.0, 0.0, e31],
                [0.0, 0.0, e31],
                [0.0, 0.0, e33],
                [0.0, e15, 0.0],
                [e15, 0.0, 0.0],
                [0.0, 0.0, 0.0]])

# define permittivity tensor p
p11 = 5.64e-9
p33 = 6.35e-9
# 3x3
p1 = as_matrix([[p11, 0.0, 0.0],
                [0.0, p11, 0.0],
                [0.0, 0.0, p33]])


# rewrite the tensor into a vector using Voigt notation
def as_voigt_vector(ten):
    # FEniCS does not know anything about Voigt notation,
    # so one need to access the components directly as eps[0, 0] etc.
    return as_vector([ten[0, 0], ten[1, 1], ten[2, 2], 2 * ten[1, 2], 2 * ten[0, 2], 2 * ten[0, 1]])


def as_voigt_vector2(ten):
    # FEniCS does not know anything about Voigt notation,
    # so one need to access the components directly as eps[0, 0] etc.
    return as_vector([ten[0], ten[1], ten[2]])


# rewrite a vector into a tensor using Voigt notation
def from_voigt_vector(vec):
    # FEniCS does not know anything about Voigt notation,
    # so one need to rewrite the vector by hand into a tensor.
    return as_tensor([[vec[0], vec[5], vec[4]],
                      [vec[5], vec[1], vec[3]],
                      [vec[4], vec[3], vec[2]]])


def from_voigt_vector2(vec):
    # FEniCS does not know anything about Voigt notation,
    # so one need to rewrite the vector by hand into a tensor.
    return as_tensor([vec[0], vec[1], vec[2]])


# Rayleigh damping coefficients
eta_m = Constant(0.)
eta_k = Constant(0.)

# Newmark method parameters
gamma = Constant(1 / 2)
beta = Constant(1 / 4)
# beta    = Constant(0.00001)  #Méthode de l’accélération constante
# beta    = Constant(1/6) #Méthode de l’accélération linéaire
xdmf_file = XDMFFile("newmark-C_beta_0.25.xdmf")
xdmf_file.parameters["flush_output"] = True
xdmf_file.parameters["functions_share_mesh"] = True
xdmf_file.parameters["rewrite_function_mesh"] = False

# Time-stepping parameters
T = 4.0
Nsteps = 50
dt = Constant(T / Nsteps)

# integration constants
aa0 = Constant(1 / beta / dt / dt)
aa1 = Constant(gamma / beta / dt)
aa2 = Constant(1 / beta / dt)
aa3 = Constant(1 / 2 / beta - 1)
aa4 = Constant(gamma / beta - 1)
aa5 = Constant((gamma / 2 / beta - 1) * dt)
aa6 = Constant((1 - gamma) * dt)
aa7 = Constant(gamma * dt)

p0 = 1e5
cutoff_Tc = 0.8
# Define the loading as an expression depending on t
p = Expression(("0", "t <= tc ? p0*t/tc : 0", "0"), t=0, tc=cutoff_Tc, p0=p0, degree=0)

# Define function space for displacement, velocity and acceleration
V = VectorFunctionSpace(mesh, "CG", 1)
Q = FunctionSpace(mesh, 'CG', 1)
Mix_element = MixedElement([V.ufl_element(), Q.ufl_element()])
Z_space = FunctionSpace(mesh, Mix_element)

# Define function space for stresses
Vsig = TensorFunctionSpace(mesh, "DG", 0)

# Test and trial functions
# du = TrialFunction(V)
# u_ = TestFunction(V)
(du, dphi) = TrialFunctions(Z_space)
(u_, phi_) = TestFunctions(Z_space)
# Current (unknown) displacement
uphi = Function(Z_space, name="Displacement&Potential")
uphi_list = split(uphi)
# u = Function(V, name="Displacement")   #Save the result of u_n+1 to U
# Fields from previous time step (displacement, velocity, acceleration)
#u_old = Function(V)  # u_n
#v_old = Function(V)  # v_n
#a_old = Function(V)  # a_n
u_old = uphi_list[0]  # u_n
v_old = uphi_list[0]  # v_n
a_old = uphi_list[0]  # a_n
#phi_old=uphi_list[1]  # phi_n
#u_old = Z_space.sub(0) # u_n
#v_old = Z_space.sub(0)  # v_n
#a_old = Z_space.sub(0)  # a_n


# Create mesh function over the cell facets
boundary_subdomains = MeshFunction("size_t", mesh, mesh.topology().dim() - 1)
boundary_subdomains.set_all(0)
top_boundary = AutoSubDomain(right)
top_boundary.mark(boundary_subdomains, 1)
force_boundary = AutoSubDomain(right)
force_boundary.mark(boundary_subdomains, 3)

# Define measure for boundary condition integral
dss = ds(subdomain_data=boundary_subdomains)

# Set up boundary condition at left end
# zero = Constant((0.0, 0.0, 0.0))
# bc = DirichletBC(V, zero, left)
bc1 = DirichletBC(Z_space.sub(0), Constant((0, 0, 0)), left)
bc2 = DirichletBC(Z_space.sub(1), Constant(0), bottom)
bc = [bc1, bc2]


# Stress tensor
# def sigma(r):
#    return 2.0*mu*sym(grad(r)) + lmbda*tr(sym(grad(r)))*Identity(len(r))
# now define the stress tensor in fencis notation
# def sigma(u):
# calculate sigma with C-matrix in Voigt notation
#    sig = (E / (1. + nu)) * C * (as_voigt_vector(epsilon(u)))
# return sigma as symmetric tensor
#    return from_voigt_vector(sig)
def sigma(u, phi):
    # calculate sigma with C-matrix in Voigt notation
    sig = C * (as_voigt_vector(epsilon(u))) - e1 * as_voigt_vector2(e_field(phi))
    # return sigma as symmetric tensor
    return from_voigt_vector(sig)


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


# then define the function for the static electric field in fenics notation
def e_field(phi):
    return -nabla_grad(phi)

def e_dis(u, phi):
    ed = e1.T * (as_voigt_vector(epsilon(u))) + p1 * as_voigt_vector2(e_field(phi))
    return from_voigt_vector2(ed)


# Mass form
def m(u, u_):
    return rho * inner(u, u_) * dx


# Elastic stiffness form
def k(u, u_, phi, phi_):
    return inner(sigma(u, phi), epsilon(u_)) * dx - dot(e_dis(u, phi), e_field(phi_)) * dx


# Rayleigh damping form
def c(u, u_, phi, phi_):
    return eta_m * m(u, u_) + eta_k * k(u, u_, phi, phi_)


def add(phi, phi_, u_):
    return 1e6*nabla_grad(phi)[0] * e_dis(u_,phi_)[0] * dss(1) + 1e6*nabla_grad(phi)[1] * e_dis(u_,phi_)[1] * dss(1)


# Work of external forces
def Wext(u_):
    return dot(u_, p) * dss(3)


# Update formula for acceleration
# a = 1/(2*beta)*((u - u0 - v0*dt)/(0.5*dt*dt) - (1-2*beta)*a0)
# a = aa0*(u - u0)-aa2*v0-aa3*a0   gamma
def update_a(u, u_old, v_old, a_old, ufl=True):
    if ufl:
        # dt_ = dt
        # gamma_ = gamma
        # beta_ = beta
        aa0_ = aa0
        aa2_ = aa2
        aa3_ = aa3
    else:
        # dt_ = float(dt)
        # gamma_ = float(gamma)
        # beta_ = float(beta)
        aa0_ = float(aa0)
        aa2_ = float(aa2)
        aa3_ = float(aa3)
    # return (u-u_old-dt_*v_old)/beta_/dt_**2 - (1-2*beta_)/2/beta_*a_old
    return aa0_ * (u - u_old) - aa2_ * v_old - aa3_ * a_old


# Update formula for velocity
# v = dt * ((1-gamma)*a0 + gamma*a) + v0
# v = v0 + aa6*a0 + aa7*a
def update_v(a, u_old, v_old, a_old, ufl=True):
    if ufl:
        # dt_ = dt
        # gamma_ = gamma
        # beta_ = beta
        aa6_ = aa6
        aa7_ = aa7
    else:
        # dt_ = float(dt)
        # gamma_ = float(gamma)
        # beta_ = float(beta)
        aa6_ = float(aa6)
        aa7_ = float(aa7)
    return v_old + aa6_ * a_old + aa7_ * a


def update_fields(u, u_old, v_old, a_old):
    """Update fields at the end of each time step."""

    # Get vectors (references)
    u_vec, u0_vec = u.vector(), u_old.vector()
    v0_vec, a0_vec = v_old.vector(), a_old.vector()
    # phi_vec, phi0_vec = phi.vector(), phi_old.vector()

    # use update functions using vector arguments
    a_vec = update_a(u_vec, u0_vec, v0_vec, a0_vec, ufl=False)
    v_vec = update_v(a_vec, u0_vec, v0_vec, a0_vec, ufl=False)

    # Update (u_old <- u)
    v_old.vector()[:], a_old.vector()[:] = v_vec, a_vec
    u_old.vector()[:] = u.vector()
    #phi_old.vector()[:] = phi.vector()


# Residual
a_new = update_a(du, u_old, v_old, a_old, ufl=True)
v_new = update_v(a_new, u_old, v_old, a_old, ufl=True)
res = m(a_new, u_) + c(v_new, u_, dphi, phi_) \
      + k(du, u_, dphi, phi_) + add(dphi, phi_, u_) - Wext(u_)
a_form = lhs(res)
L_form = rhs(res)

# Define solver for reusing factorization
K, res = assemble_system(a_form, L_form, bc)
solver = LUSolver(K, "mumps")
solver.parameters["symmetric"] = True

# Time-stepping
time = np.linspace(0, T, Nsteps + 1)
u_tip = np.zeros((Nsteps + 1,))
energies = np.zeros((Nsteps + 1, 4))
E_damp = 0
E_ext = 0
sig = Function(Vsig, name="sigma")


def local_project(v, V, u=None):
    """Element-wise projection using LocalSolver"""
    dv = TrialFunction(V)
    v_ = TestFunction(V)
    a_proj = inner(dv, v_) * dx
    b_proj = inner(v, v_) * dx
    solver = LocalSolver(a_proj, b_proj)
    solver.factorize()
    if u is None:
        u = Function(V)
        solver.solve_local_rhs(u)
        return u
    else:
        solver.solve_local_rhs(u)
        return


for (i, dt) in enumerate(np.diff(time)):

    t = time[i + 1]
    print("Time: ", t)

    # Forces are evaluated at t_{n+1-alpha_f}=t_{n+1}-alpha_f*dt
    # p.t = t-float(alpha_f*dt)
    p.t = t

    # Solve for new displacement
    res = assemble(L_form)
    for bcs in bc:
        bcs.apply(res)  # 应用每个边界条件
    #bc.apply(res)
    solver.solve(K, uphi.vector(), res)
    (u, phi) = uphi.split()

    # Update old fields with new quantities
    update_fields(u, u_old, v_old, a_old)

    # Save solution to XDMF format
    xdmf_file.write(u, t)

    # Compute stresses and save to file
    local_project(sigma(u, phi), Vsig, sig)
    xdmf_file.write(sig, t)

    p.t = t
    # Record tip displacement and compute energies
    # Note: Only works in serial
    if MPI.comm_world.size == 1:
        u_tip[i + 1] = u(1., 0.05, 0.)[1]
    E_elas = assemble(0.5 * k(u_old, u_old, phi, phi_old))
    E_kin = assemble(0.5 * m(v_old, v_old))
    E_damp += dt * assemble(c(v_old, v_old, phi,phi))
    # E_ext += assemble(Wext(u-u_old))
    E_tot = E_elas + E_kin + E_damp  # -E_ext
    energies[i + 1, :] = np.array([E_elas, E_kin, E_damp, E_tot])

If you want to work on the sub-spaces individually, you need to use deepcopy=True when you use the split method.
i.e.

u, phi = uphi.split()
print(len(u.vector()[:]), len(phi.vector()[:]))

returns

24684 24684

while

u, phi = uphi.split(deepcopy=True)
print(len(u.vector()[:]), len(phi.vector()[:]))

returns

18513 6171

The same holds for u_old etc.
You need to distinguish between the split(u) and u.split() method, as the first one is for usage in variational forms, while the latter is used to access the vector (especially when deepcopy=True), see: fenics-project / DOLFIN / issues / #194 - split(u) versus u.split() — Bitbucket for a discussion about the two.

Dear @dokken

Your suggestion is exactly what I want, thank you!
I did misunderstand the difference between split(u) and u.split().
The modified code is posted here incase someone meet the same problem:

from dolfin import *
import numpy as np
import matplotlib.pyplot as plt

# Form compiler options
parameters["form_compiler"]["cpp_optimize"] = True
parameters["form_compiler"]["optimize"] = True

# Define mesh
mesh = BoxMesh(Point(0., 0., 0.), Point(1., 0.2, 0.2), 50, 10, 10)


# Sub domain for clamp at left end
def left(x, on_boundary):
    return near(x[0], 0.) and on_boundary


# Sub domain for rotation at right end
def right(x, on_boundary):
    return near(x[0], 1.) and on_boundary


# Sub domain for bottom end
def bottom(x, on_boundary):
    return near(x[2], 0.) and on_boundary


# Sub domain for top end
def top(x, on_boundary):
    return near(x[2], 0.2) and on_boundary


# Mass density
rho = Constant(5550)
g = 9.8

# define elasticity tensor
# define elasticity tensor C
c11 = 226e9
c12 = 125e9
c13 = 124e9
c33 = 216e9
c44 = 44.2e9  # * 2
c66 = (c11 - c12) / 2  # * 2
# 6x6
C = as_matrix([[c11, c12, c13, 0.0, 0.0, 0.0],
               [c12, c11, c13, 0.0, 0.0, 0.0],
               [c13, c13, c33, 0.0, 0.0, 0.0],
               [0.0, 0.0, 0.0, c44, 0.0, 0.0],
               [0.0, 0.0, 0.0, 0.0, c44, 0.0],
               [0.0, 0.0, 0.0, 0.0, 0.0, c66]])

# define piezoelectricity tensor e
e15 = 5.8
e31 = -2.2
e33 = 9.3
# 6x3
e1 = as_matrix([[0.0, 0.0, e31],
                [0.0, 0.0, e31],
                [0.0, 0.0, e33],
                [0.0, e15, 0.0],
                [e15, 0.0, 0.0],
                [0.0, 0.0, 0.0]])

# define permittivity tensor p
p11 = 5.64e-9
p33 = 6.35e-9
# 3x3
p1 = as_matrix([[p11, 0.0, 0.0],
                [0.0, p11, 0.0],
                [0.0, 0.0, p33]])


# rewrite the tensor into a vector using Voigt notation
def as_voigt_vector(ten):
    # FEniCS does not know anything about Voigt notation,
    # so one need to access the components directly as eps[0, 0] etc.
    return as_vector([ten[0, 0], ten[1, 1], ten[2, 2], 2 * ten[1, 2], 2 * ten[0, 2], 2 * ten[0, 1]])


def as_voigt_vector2(ten):
    # FEniCS does not know anything about Voigt notation,
    # so one need to access the components directly as eps[0, 0] etc.
    return as_vector([ten[0], ten[1], ten[2]])


# rewrite a vector into a tensor using Voigt notation
def from_voigt_vector(vec):
    # FEniCS does not know anything about Voigt notation,
    # so one need to rewrite the vector by hand into a tensor.
    return as_tensor([[vec[0], vec[5], vec[4]],
                      [vec[5], vec[1], vec[3]],
                      [vec[4], vec[3], vec[2]]])


def from_voigt_vector2(vec):
    # FEniCS does not know anything about Voigt notation,
    # so one need to rewrite the vector by hand into a tensor.
    return as_tensor([vec[0], vec[1], vec[2]])


# Rayleigh damping coefficients
eta_m = Constant(0.)
eta_k = Constant(0.)

# Newmark method parameters
gamma = Constant(1 / 2)
beta = Constant(1 / 4)

xdmf_file = XDMFFile("newmark-C_sin.xdmf")
xdmf_file.parameters["flush_output"] = True
xdmf_file.parameters["functions_share_mesh"] = True
xdmf_file.parameters["rewrite_function_mesh"] = False

# Time-stepping parameters
T = 4.0
Nsteps = 50
dt = Constant(T / Nsteps)

# integration constants
aa0 = Constant(1 / beta / dt / dt)
aa1 = Constant(gamma / beta / dt)
aa2 = Constant(1 / beta / dt)
aa3 = Constant(1 / 2 / beta - 1)
aa4 = Constant(gamma / beta - 1)
aa5 = Constant((gamma / 2 / beta - 1) * dt)
aa6 = Constant((1 - gamma) * dt)
aa7 = Constant(gamma * dt)

p0 = 1e5
cutoff_Tc = 0.8
n=4
# Define the loading as an expression depending on t
#p = Expression(("0", "0", "t <= tc ? p0*t/tc : p0"), t=0, tc=cutoff_Tc, p0=p0, degree=0)
p = Expression(("0", "0", "p0*sin(2*pi*T/n*t)"), t=0, T=T, p0=p0, n=n, degree=0)
f = Constant((0, 0, -rho * g))

# Define function space for displacement, velocity and acceleration
#V = VectorFunctionSpace(mesh, "CG", 1)
#Q = FunctionSpace(mesh, 'CG', 1)
#W = FunctionSpace(mesh, MixedElement([Velocity_element, Pressure_element]))
#Mix_element = MixedElement([V.ufl_element(), Q.ufl_element()])
#Z_space = FunctionSpace(mesh, Mix_element)
Velocity_element = VectorElement("Lagrange", mesh.ufl_cell(), 1)
Potential_element = FiniteElement("Lagrange", mesh.ufl_cell(), 1)
Z_space = FunctionSpace(mesh, MixedElement([Velocity_element, Potential_element]))

#V, V_to_Z = Z_space.sub(0).collapse()
#V, Q = Z_space.split()

# Define function space for stresses
Vsig = TensorFunctionSpace(mesh, "DG", 0)

# Test and trial functions
# du = TrialFunction(V)
# u_ = TestFunction(V)
(du, dphi) = TrialFunctions(Z_space)
(u_, phi_) = TestFunctions(Z_space)
# Current (unknown) displacement
uphi = Function(Z_space, name="Displacement&Potential")
uphi_old = Function(Z_space)
vphi_old = Function(Z_space)
aphi_old = Function(Z_space)
# u = Function(V, name="Displacement")   #Save the result of u_n+1 to U
# Fields from previous time step (displacement, velocity, acceleration)
u_old,_ = uphi_old.split(deepcopy=True)
v_old,_ = vphi_old.split(deepcopy=True)
a_old,_ = aphi_old.split(deepcopy=True)



# Create mesh function over the cell facets
boundary_subdomains = MeshFunction("size_t", mesh, mesh.topology().dim() - 1)
boundary_subdomains.set_all(0)
top_boundary = AutoSubDomain(top)
top_boundary.mark(boundary_subdomains, 1)
force_boundary = AutoSubDomain(right)
force_boundary.mark(boundary_subdomains, 3)

# Define measure for boundary condition integral
dss = ds(subdomain_data=boundary_subdomains)

# Set up boundary condition at left end
# zero = Constant((0.0, 0.0, 0.0))
# bc = DirichletBC(V, zero, left)
bc1 = DirichletBC(Z_space.sub(0), Constant((0, 0, 0)), left)
bc2 = DirichletBC(Z_space.sub(1), Constant(0), bottom)
bc = [bc1, bc2]


# Stress tensor
# def sigma(r):
#    return 2.0*mu*sym(grad(r)) + lmbda*tr(sym(grad(r)))*Identity(len(r))
# now define the stress tensor in fencis notation
# def sigma(u):
# calculate sigma with C-matrix in Voigt notation
#    sig = (E / (1. + nu)) * C * (as_voigt_vector(epsilon(u)))
# return sigma as symmetric tensor
#    return from_voigt_vector(sig)
def sigma(u, phi):
    # calculate sigma with C-matrix in Voigt notation
    sig = C * (as_voigt_vector(epsilon(u))) - e1 * as_voigt_vector2(e_field(phi))
    # return sigma as symmetric tensor
    return from_voigt_vector(sig)


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


# then define the function for the static electric field in fenics notation
def e_field(phi):
    return -nabla_grad(phi)

def e_dis(u, phi):
    ed = e1.T * (as_voigt_vector(epsilon(u))) + p1 * as_voigt_vector2(e_field(phi))
    return from_voigt_vector2(ed)


# Mass form
def m(u, u_):
    return rho * inner(u, u_) * dx


# Elastic stiffness form
def k(u, u_, phi, phi_):
    return inner(sigma(u, phi), epsilon(u_)) * dx - dot(e_dis(u, phi), e_field(phi_)) * dx


# Rayleigh damping form
def c(u, u_, phi, phi_):
    return eta_m * m(u, u_) + eta_k * k(u, u_, phi, phi_)


def add(phi, phi_, u_):
    return 1e6*nabla_grad(phi)[0] * e_dis(u_,phi_)[0] * dss(1) + 1e6*nabla_grad(phi)[1] * e_dis(u_,phi_)[1] * dss(1)


# Work of external forces
def Wext(u_):
    return dot(u_, -p) * dss(3)+dot(f, u_) * dx


# Update formula for acceleration
# a = 1/(2*beta)*((u - u0 - v0*dt)/(0.5*dt*dt) - (1-2*beta)*a0)
# a = aa0*(u - u0)-aa2*v0-aa3*a0   gamma
def update_a(u, u_old, v_old, a_old, ufl=True):
    if ufl:
        # dt_ = dt
        # gamma_ = gamma
        # beta_ = beta
        aa0_ = aa0
        aa2_ = aa2
        aa3_ = aa3
    else:
        # dt_ = float(dt)
        # gamma_ = float(gamma)
        # beta_ = float(beta)
        aa0_ = float(aa0)
        aa2_ = float(aa2)
        aa3_ = float(aa3)
    # return (u-u_old-dt_*v_old)/beta_/dt_**2 - (1-2*beta_)/2/beta_*a_old
    return aa0_ * (u - u_old) - aa2_ * v_old - aa3_ * a_old


# Update formula for velocity
# v = dt * ((1-gamma)*a0 + gamma*a) + v0
# v = v0 + aa6*a0 + aa7*a
def update_v(a, u_old, v_old, a_old, ufl=True):
    if ufl:
        # dt_ = dt
        # gamma_ = gamma
        # beta_ = beta
        aa6_ = aa6
        aa7_ = aa7
    else:
        # dt_ = float(dt)
        # gamma_ = float(gamma)
        # beta_ = float(beta)
        aa6_ = float(aa6)
        aa7_ = float(aa7)
    return v_old + aa6_ * a_old + aa7_ * a


def update_fields(u, u_old, v_old, a_old):
    """Update fields at the end of each time step."""

    # Get vectors (references)
    u_vec, u0_vec = u.vector(), u_old.vector()
    v0_vec, a0_vec = v_old.vector(), a_old.vector()
    # phi_vec, phi0_vec = phi.vector(), phi_old.vector()

    # use update functions using vector arguments
    a_vec = update_a(u_vec, u0_vec, v0_vec, a0_vec, ufl=False)
    v_vec = update_v(a_vec, u0_vec, v0_vec, a0_vec, ufl=False)
    #a_vec = update_a(u, u_old, v_old, a_old, ufl=False)
    #v_vec = update_v(a_vec, u_old, v_old, a_old, ufl=False)

    # Update (u_old <- u)
    v_old.vector()[:], a_old.vector()[:] = v_vec, a_vec
    u_old.vector()[:] = u.vector()
    #phi_old.vector()[:] = phi.vector()
    #v_old.assign(v_vec)
    #a_old.assign(a_vec)
    #u_old.assign(u)


# Residual
a_new = update_a(du, u_old, v_old, a_old, ufl=True)
v_new = update_v(a_new, u_old, v_old, a_old, ufl=True)
res = m(a_new, u_) + c(v_new, u_, dphi, phi_) \
      + k(du, u_, dphi, phi_) + add(dphi, phi_, u_) - Wext(u_)
a_form = lhs(res)
L_form = rhs(res)

# Define solver for reusing factorization
K, res = assemble_system(a_form, L_form, bc)
solver = LUSolver(K, "mumps")
solver.parameters["symmetric"] = True

# Time-stepping
time = np.linspace(0, T, Nsteps + 1)
u_tip = np.zeros((Nsteps + 1,))
voltage_top = np.zeros((Nsteps + 1,))
energies = np.zeros((Nsteps + 1, 4))
E_damp = 0
E_ext = 0
sig = Function(Vsig, name="sigma")


def local_project(v, V, u=None):
    """Element-wise projection using LocalSolver"""
    dv = TrialFunction(V)
    v_ = TestFunction(V)
    a_proj = inner(dv, v_) * dx
    b_proj = inner(v, v_) * dx
    solver = LocalSolver(a_proj, b_proj)
    solver.factorize()
    if u is None:
        u = Function(V)
        solver.solve_local_rhs(u)
        return u
    else:
        solver.solve_local_rhs(u)
        return


for (i, dt) in enumerate(np.diff(time)):

    t = time[i + 1]
    print("Time: ", t)

    # Forces are evaluated at t_{n+1-alpha_f}=t_{n+1}-alpha_f*dt
    # p.t = t-float(alpha_f*dt)
    p.t = t

    # Solve for new displacement
    res = assemble(L_form)
    for bcs in bc:
        bcs.apply(res)  # 应用每个边界条件
    #bc.apply(res)
    solver.solve(K, uphi.vector(), res)
    (u, phi) = uphi.split(deepcopy=True)

    # Update old fields with new quantities
    update_fields(u, u_old, v_old, a_old)

    # Save solution to XDMF format
    xdmf_file.write(u, t)
    xdmf_file.write(phi, t)

    # Compute stresses and save to file
    local_project(sigma(u, phi), Vsig, sig)
    xdmf_file.write(sig, t)


    # Record tip displacement and compute energies
    # Note: Only works in serial
    if MPI.comm_world.size == 1:
        u_tip[i + 1] = u(1., 0.05, 0.)[2]
        voltage_top[i + 1]=phi(0.5, 0.1, 0.2)
    E_elas = assemble(0.5 * k(u_old, u_old, phi, phi))
    E_kin = assemble(0.5 * m(v_old, v_old))
    E_damp += dt * assemble(c(v_old, v_old, phi,phi))
    # E_ext += assemble(Wext(u-u_old))
    E_tot = E_elas + E_kin + E_damp  # -E_ext
    energies[i + 1, :] = np.array([E_elas, E_kin, E_damp, E_tot])

# Note that in the above, the stresses are computed using a ``LocalSolver`` through the ``local_project`` function. Since the stress function space
# is a DG-0 space, the projection on this space can be performed element-wise in a very efficient manner. We therefore take advantage of the ``LocalSolver`` 
# functionality which is precisely dedicated to such situations. Since this projection is performed at each time step, the savings in terms of computing
# time can be quite important.
# 
# As regards the computation of the various energies, the elastic and kinetic energies are respectively given by:
# 
# .. math::
#    E_{elas} = \int_{\Omega} \dfrac{1}{2}\sigma(u):\varepsilon(u) \, {\rm d} x
# .. math::
#    E_{kin} = \int_{\Omega} \dfrac{1}{2}\rho \dot{u}\cdot\dot{u} \, {\rm d} x
# 
# which are readily computed from the respective stiffness and mass forms ``k`` and ``m`` and the current displacement and velocity. The energy related to damping
# is computed from the corresponding dissipation term :math:`\mathcal{D}=c(\dot{u},\dot{u})` and integrated over time:
# 
# .. math::
#    E_{damp} = \int_0^T \mathcal{D} \, {\rm d} t
# 
# As for the work developed by the external forces, the contribution to the energy is added at each time step. Finally, the total energy of the sytem is given by:
# 
# .. math::
#    E_{tot} = E_{elas}+E_{kin}+E_{damp}-E_{ext}
# 
# When the time evolution loop is finished, the evolution of the tip displacement as well as the different contributions of the energy are plotted as functions of time::

if MPI.comm_world.size == 1:
    # Plot tip displacement evolution
    plt.figure()
    plt.plot(time, u_tip)
    plt.xlabel("Time")
    plt.ylabel("Tip displacement")
    #plt.ylim(-0.5, 0.5)
    plt.show()

if MPI.comm_world.size == 1:
    # Plot tip displacement evolution
    plt.figure()
    plt.plot(time, voltage_top)
    plt.xlabel("Time")
    plt.ylabel("Voltage on topsurface")
    #plt.ylim(-0.5, 0.5)
    plt.show()

if (MPI.comm_world.rank == 0):
    # Plot energies evolution
    plt.figure()
    plt.plot(time, energies)
    plt.legend(("elastic", "kinetic", "damping", "total"))
    plt.xlabel("Time")
    plt.ylabel("Energies")
    #plt.ylim(0, 0.0011)
    plt.show()

# ---------------------
# Analyzing the results
# ---------------------
# 
# We first consider the case of zero Rayleigh damping :math:`\eta_M=\eta_K=0`. In this case, it can be observed that the evolution of the total energy depends
# on the choice of the time-stepping scheme parameters. With :math:`\alpha_m=\alpha_f=0`, we recover the Newmark-:math:`\beta` method with :math:`\beta=0.25,\gamma=0.5`.
# This scheme is known for being conservative. This can be observed (figure-left) in the constant total energy for :math:`t\geq T_c` when the loading is removed. On the contrary,
# for non zero alpha parameters, e.g. :math:`\alpha_m=0.2, \alpha_f=0.4`, it can be observed (figure-right) that the energy is decreasing during this phase, indicating numerical damping.
# For both cases, the scheme is unconditionally stable. Moreover, these differences vanish when reducing the time step.
# 
# .. figure:: elastodynamics-energies_newmark.png
#    :width: 60%
#    :align: center
# 
#    Newmark-:math:`\beta` method :math:`\alpha_m=\alpha_f=0`
# 
# .. figure:: elastodynamics-energies_generalized_alpha.png
#    :width: 60%
#    :align: center
# 
#    Generalized-:math:`\alpha` with :math:`\alpha_m=0.2, \alpha_f=0.4`
# 
# For non-zero Rayleigh damping :math:`\eta_M=\eta_K=0.01`, the total energy including viscous dissipation tends to oscillate around a constant value, with oscillations vanishing for decreasing time steps.
# 
# .. image:: elastodynamics-energies_generalized_alpha_damping.png
#    :width: 60%
#    :align: center
# 
# -----------
# References
# -----------
# 
# .. [ERL2002] Silvano Erlicher, Luca Bonaventura, Oreste Bursi. The analysis of the Generalized-alpha method for non-linear dynamic problems. Computational Mechanics, Springer Verlag, 2002, 28, pp.83-104, doi:10.1007/s00466-001-0273-z