Solution updating in the time loop

Hi everyone!
I need to extract components in a mixed finite element space and update them with.
But the u values I printed in two ways are different,u_test and u0 respectively.
the main code :

        solver.solve(b, wn.vector)
        u0 = wn.sub(0).collapse()
        A0 = wn.sub(1).collapse()
        p0 = wn.sub(2).collapse()
        offset = W0.dofmap.index_map.size_local * W0.dofmap.index_map_bs
        offset1 = W1.dofmap.index_map.size_local * W1.dofmap.index_map_bs
        offset2 = W2.dofmap.index_map.size_local * W2.dofmap.index_map_bs
        print('1 is',offset)
        print('2 is',offset1)
        print('3 is',offset2)
        u_test.x.array[::]=wn.x.array[:offset]
        print('wn',len(wn.x.array))
        print('u0',len(u0.x.array))
        print('a0',len(A0.x.array))
        print('p0',len(p0.x.array))
        print(u_test.x.array[:20])
        print(u0.x.array[:20])

output:

1 is 19584
2 is 4184
3 is 12288
wn 36056
u0 19584
a0 4184
p0 12288
[ 7.39346510e-03 -1.28154933e-04  9.18965132e-05  7.13080626e-03
 -1.20728390e-04 -3.00721578e-06  1.82657280e-04 -1.61130687e-04
 -1.33367441e-04  4.45515478e-04 -7.48266475e-05 -1.00955936e-04
 -3.83697155e-05 -3.13355810e-01 -3.13448153e-01 -3.01518102e-01
 -3.01435827e-01  1.98569770e-03  0.00000000e+00  2.24697778e-03]
[ 7.39346510e-03 -1.28154933e-04  9.18965132e-05  7.13080626e-03
 -1.20728390e-04 -3.00721578e-06  1.82657280e-04 -1.61130687e-04
 -1.33367441e-04  4.45515478e-04 -7.48266475e-05 -1.00955936e-04
  6.43970781e-03 -1.15514689e-04 -5.52278058e-05  6.25757857e-03
 -3.20202251e-04 -1.78086991e-04  0.00000000e+00  0.00000000e+00]

Which one is correct?
Finally, I attach the complete code:

import numpy as np
import ufl
from dolfinx import cpp as _cpp
from dolfinx import fem
from dolfinx.fem import (Constant, Function, FunctionSpace, dirichletbc,
                         extract_function_spaces, form,
                         locate_dofs_geometrical, locate_dofs_topological)
from dolfinx.io import XDMFFile
from dolfinx.mesh import (CellType, GhostMode, create_rectangle,
                          locate_entities_boundary,create_unit_cube)
from ufl import div, dx, grad, inner,dot,outer,ds,dS,avg,jump,rot,cross

from mpi4py import MPI
from petsc4py import PETSc

# Create mesh
msh = create_unit_cube(MPI.COMM_WORLD,
                       8,8,8,
                       CellType.tetrahedron, GhostMode.none)

u_element=ufl.FiniteElement("BDM", ufl.tetrahedron, 1)    # H(div)     Vh
A_element=ufl.FiniteElement("N1curl",ufl.tetrahedron,1)   # H(curl)      Ch
P_element=ufl.FiniteElement("DG",ufl.tetrahedron,1)       #                 Qh
#fai_element=ufl.FiniteElement("DG",ufl.tetrahedron,2)     #     SH

TH= ufl.MixedElement([u_element, A_element,P_element])

W =fem.FunctionSpace(msh,TH)

W0,_=W.sub(0).collapse()
W1,_=W.sub(1).collapse()
W2,_=W.sub(2).collapse()

(un,An,Pn)=ufl.TrialFunctions(W)   #
(v,fai,q)=ufl.TestFunctions(W)   #

#Vh=fem.FunctionSpace(msh,u_element)
un_1=fem.Function(W0)                         
un_2=fem.Function(W0)
u_test=fem.Function(W0)

def initial_u(x):
    values=np.zeros((3,x.shape[1]))
    values[0,:]=0    #y
    values[1,:]=0    #z
    values[2,:]=0    #x
    return values

un_1.interpolate(initial_u)   

#Ch=fem.FunctionSpace(msh,A_element)
An_1=fem.Function(W1)                     
An_2=fem.Function(W1)

def initial_A(x):
    values=np.zeros((3,x.shape[1]))
    values[0,:]=0               #z
    values[1,:]=np.sin(x[0])    #0
    values[2,:]=0               #y
    return values
An_1.interpolate(initial_A)

Pn=fem.Function(W2)
def initial_P(x):
    values=np.zeros((1,x.shape[1]))
    values[0,:]=x[0]+x[1]+x[2]-1.5
    return values
Pn.interpolate(initial_P)



def local_u(x):               #mark  boundary of u 
    #return np.logical_and(np.isclose(x[2],0),np.isclose(x[1],0))    #mark z=0 && x=0
    return np.isclose(x[1],0)

def u_D(x):
    values=np.zeros((3,x.shape[1]))
    values[0,:]=0           #y
    values[1,:]=0           #z
    values[2,:]=0           #x
    return values


Vh=fem.FunctionSpace(msh,u_element)
u_boundary = fem.Function(Vh)
u_boundary.interpolate(u_D)
facets = locate_entities_boundary(msh,2, local_u)
dofs = locate_dofs_topological((W.sub(0), Vh), 2, facets)
bc0 = dirichletbc(u_boundary, dofs, W.sub(0))

bcs=[bc0]

#------------------------------------------------------------------------------------------
x = ufl.SpatialCoordinate(msh)
h = ufl.CellDiameter(msh)
n = ufl.FacetNormal(msh)

unv=(un+un_1)*0.5
uns=1.5*un_1-0.5*un_2

Anv=(An+An_1)*0.5
Ans=1.5*An_1-0.5*An_2

lmbda = ufl.conditional(ufl.gt(dot(un_1, n), 0), 1, 0)    
u_uw = lmbda("+") * unv("+") + lmbda("-") * unv("-")        #  u ->   upwind flux

tao_1=1000
k=1
Rm=1
alph=10
b = 2*tao_1*inner(unv,v)*dx
b -= inner(unv,div(outer(uns,v)))*dx
b += inner((dot(uns, n))("+") * u_uw, v("+")) * dS+inner((dot(uns, n))("-") * u_uw, v("-")) * dS     # 
b += inner(dot(uns, n) * lmbda * unv, v) * ds
b += inner(grad(unv),grad(v))*dx +alph/h('+')*inner(jump(unv),jump(v))*dS 
b -= inner(avg(dot(grad(unv),n)),jump(v))*dS-inner(avg(dot(grad(v),n)),jump(unv))*dS 
b += 2*k*tao_1*inner(rot(Anv),rot(fai))*dx
b += k*inner(2*tao_1*Anv+cross(rot(Ans),unv),2*tao_1*fai+cross(rot(Ans),v))*dx 
b += inner(2*tao_1*un_1,v)*dx+k*inner(2*tao_1*An_1,2*tao_1*fai+cross(rot(Ans),v))*dx 
b += 2*tao_1*inner(div(unv),div(v))*dx-inner(div(unv),q)*dx-inner(div(v),Pn)*dx
b -= inner(u_boundary,dot(grad(v),n))*ds+alph/h*inner(u_boundary,v)*ds

a=ufl.lhs(b)
L=ufl.rhs(b)

a=fem.form(a)
L=fem.form(L)

#----------------------------------------------------------------------------------------------
from dolfinx.fem.petsc import (apply_lifting, assemble_matrix, assemble_vector, 
                               create_vector, set_bc)

A=assemble_matrix(a,bcs=bcs)
A.assemble()
b=create_vector(L)

solver = PETSc.KSP().create(msh.comm)
solver.setOperators(A)
solver.setType(PETSc.KSP.Type.BCGS)
pc1 = solver.getPC()
pc1.setType(PETSc.PC.Type.HYPRE)
pc1.setHYPREType("boomeramg")
solver.setFromOptions()


def domain_average(msh, v):
    """Compute the average of a function over the domain"""
    vol = msh.comm.allreduce(fem.assemble_scalar(fem.form(fem.Constant(msh, 1.0) * dx)), op=MPI.SUM)
    return (1 / vol) * msh.comm.allreduce(fem.assemble_scalar(fem.form(v * dx)), op=MPI.SUM)

wn=Function(W)
t=0
for n in range(1):
    if n==0:
        t += 0.001
        un_2.x.array[::]=un_1.x.array[::]
        An_2.x.array[::]=An_1.x.array[::]
        fem.petsc.assemble_matrix(a,bcs=bcs)          # type: ignore
        assemble_vector(b, L)  
        apply_lifting(b, [a], [bcs])                  
        b.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
        set_bc(b, bcs)
        solver.solve(b, wn.vector)
        u0 = wn.sub(0).collapse()
        A0 = wn.sub(1).collapse()
        p0 = wn.sub(2).collapse()
        offset = W0.dofmap.index_map.size_local * W0.dofmap.index_map_bs
        offset1 = W1.dofmap.index_map.size_local * W1.dofmap.index_map_bs
        offset2 = W2.dofmap.index_map.size_local * W2.dofmap.index_map_bs
        print('1 is',offset)
        print('2 is',offset1)
        print('3 is',offset2)
        u_test.x.array[::]=wn.x.array[:offset]
        print('wn',len(wn.x.array))
        print('u0',len(u0.x.array))
        print('a0',len(A0.x.array))
        print('p0',len(p0.x.array))
        print(u_test.x.array[:20])
        print(u0.x.array[:20])
        un_2.x.array[::]=un_1.x.array[::]
        un_1.x.array[::]=u0.x.array[::]
        An_2.x.array[::]=An_1.x.array[::]
        An_1.x.array[::]=A0.x.array[::]
        un_2.x.scatter_forward()
        un_1.x.scatter_forward()
        An_2.x.scatter_forward()
        An_1.x.scatter_forward()
        Pn.x.array[::]=p0.x.array[::]
    else:
        t += 0.001
        fem.petsc.assemble_matrix(a,bcs=bcs)  # type: ignore
        with b.localForm() as b_loc:
            b_loc.set(0)
        assemble_vector(b, L)        
        apply_lifting(b, [a], [bcs])                 
        b.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
        set_bc(b, bcs)
        solver.solve(b, wn.vector)
        u0 = wn.sub(0).collapse()
        A0 = wn.sub(1).collapse()
        p0 = wn.sub(2).collapse()
        un_2.x.array[::]=un_1.x.array[::]
        un_1.x.array[::]=u0.x.array[::]
        An_2.x.array[::]=An_1.x.array[::]
        An_1.x.array[::]=A0.x.array[::]
        un_2.x.scatter_forward()
        un_1.x.scatter_forward()
        An_2.x.scatter_forward()
        An_1.x.scatter_forward()
        Pn.x.array[::]=p0.x.array[::]
    print('t is',t)

u = wn.sub(0).collapse()
A = wn.sub(1).collapse()
p = wn.sub(2).collapse()

with XDMFFile(MPI.COMM_WORLD, "out_WH/velocity.xdmf", "w") as ufile_xdmf:
    u.x.scatter_forward()
    ufile_xdmf.write_mesh(msh)
    ufile_xdmf.write_function(u)

with XDMFFile(MPI.COMM_WORLD, "out_WH/A.xdmf", "w") as pfile_xdmf:
    A.x.scatter_forward()
    pfile_xdmf.write_mesh(msh)
    pfile_xdmf.write_function(A)

with XDMFFile(MPI.COMM_WORLD, "out_WH/p.xdmf", "w") as pfile_xdmf:
    p.x.scatter_forward()
    pfile_xdmf.write_mesh(msh)
    pfile_xdmf.write_function(p)

Any answer will be welcomed!Thanks very much !

I think the dofs are ordered differently than you expect. You should store the map that is returned as the second value of W.sub(i).collapse(), and use it to access the dofs. Consider:

import dolfinx as dfx
import numpy as np
import ufl
from mpi4py import MPI

msh = dfx.mesh.create_unit_square(MPI.COMM_WORLD, 1, 1)

sz = 3
elem = ufl.FiniteElement("CG", msh.ufl_cell(), 1)
elem_mixed = ufl.MixedElement([elem] * sz)
V_vec = dfx.fem.VectorFunctionSpace(msh, ("CG", 1), dim=sz)
V_mix = dfx.fem.FunctionSpace(msh, elem_mixed)
v_vec = dfx.fem.Function(V_vec)
v_mix = dfx.fem.Function(V_mix)

for i in range(sz):
    v_vec.sub(i).interpolate(lambda x: np.full_like(x[0], i))
    v_mix.sub(i).interpolate(lambda x: np.full_like(x[0], i))

print("vector:", v_vec.x.array)
print("mixed:", v_mix.x.array)

V_mix_0, map_0 = V_mix.sub(0).collapse()
offset = V_mix_0.dofmap.index_map.size_local * V_mix_0.dofmap.index_map_bs
print("mixed component 0 (incorrect):", v_mix.x.array[:offset])
print("mixed component 0 (correct):", v_mix.x.array[map_0])

producing:

vector: [0. 1. 2. 0. 1. 2. 0. 1. 2. 0. 1. 2.]
mixed: [0. 0. 0. 1. 1. 1. 2. 2. 2. 0. 1. 2.]
mixed component 0 (incorrect): [0. 0. 0. 1.]
mixed component 0 (correct): [0. 0. 0. 0.]
3 Likes

Thanks very much!
What an outstanding thinking that has benefited me immensely.