How to get subarrays from mixed solution array

Dear all,

I’m trying to solve a system in FEniCSx using a mixed element, but I don’t know how get the arrays of the subfunctions back. My naive guess was to think that the first elements of the mixed element would correspond to the first space, and the last to the second space, but that does not seem to be the case. Using .split() and .sub(0) also seems to return the entire array instead of the parts. My other guess would be to use ufl.split() but in that case I don’t know how to get back the arrays belonging to them as those functions do not have the property x.array

Below first my try, and afterwards a different way of solving with substitution that outputs the correct values for the subarrays. The context is that of solving a porous medium type equation but that is not relevant for my question.

My try:

from mpi4py import MPI
from dolfinx import mesh 
from dolfinx.mesh import locate_entities, meshtags 
from dolfinx.fem import FunctionSpace
from dolfinx import fem 
import numpy as np 
import ufl 
from petsc4py.PETSc import ScalarType
from petsc4py import PETSc 
from dolfinx import io 
import dolfinx

# Creating an interval as mesh
h = 0.25
x_min, x_max = -1,1
nx = int((x_max - x_min)/h)
domain = mesh.create_interval(MPI.COMM_WORLD, nx, points = (x_min, x_max))
x = ufl.SpatialCoordinate(domain) # Coordinates of domain used for defining general functions

el1 = ufl.FiniteElement("CG", domain.ufl_cell(), 1)
el2 = ufl.FiniteElement("DG", domain.ufl_cell(), 1)
mel = ufl.MixedElement([el1, el2])
Z = FunctionSpace(domain, mel)

V = FunctionSpace(domain, ("CG", 1)) 
U = FunctionSpace(domain, ("DG", 1))

w_trial, u_trial = ufl.TrialFunctions(Z)
w_test, u_test = ufl.TestFunctions(Z)
Z_sol = fem.Function(Z)

# Constants
m  = 2
alpha_par = 1/3
beta_par = 1/3
kappa_par = 1/12
C = 1/32
t_start  = 0.3
dt = 0.1

def Phi(u):
    return u**m 

def exact_sol(x,t):
    return ufl.operators.max_value(0,1/(t**alpha_par) * (C - kappa_par * (ufl.operators.sqrt(x[0]**2)/(t**beta_par))**2 )**(1/(m-1)))

initial_u = exact_sol(x,t_start)
expr_u = fem.Expression(initial_u, U.element.interpolation_points())

# Defining the function spaces for all the functions
u_n = fem.Function(U)
u_n.name = "u_n"

u_n_i = fem.Function(U)
u_n_i.name = "u_n_i"

u_n.interpolate(expr_u)
u_n_i.interpolate(expr_u)

L_par = 1
L_const = fem.Constant(domain, ScalarType(L_par))

a = (u_trial * w_test + dt * ufl.inner(ufl.grad(w_trial), ufl.grad(w_test)) + L_const * u_trial * u_test - w_trial * u_test) * ufl.dx
L = (u_n * w_test + L_const * u_n_i * u_test - Phi(u_n_i) * u_test) * ufl.dx
  
bilinear_form = fem.form(a)
linear_form = fem.form(L)
  
b = fem.petsc.assemble_vector(linear_form)
b.assemble()
A = fem.petsc.assemble_matrix(bilinear_form)
A.assemble()
  
solver = PETSc.KSP().create(domain.comm)
solver.setType(PETSc.KSP.Type.PREONLY) # PREONLY means you only use the preconditioner
#solver.setTolerances(rtol=1e-12, atol = 1e-60)
solver.getPC().setType(PETSc.PC.Type.LU) # The LU preconditioner just solves the system using LU factorization.
  
solver.setOperators(A)
  
solver.solve(b, Z_sol.vector)

len_V = len(V.tabulate_dof_coordinates())
len_U = len(U.tabulate_dof_coordinates())

w_split, u_split = Z_sol.x.array[:len_V], Z_sol.x.array[len_V:]
w_split2, u_split2 = Z_sol.split()
w_split3, u_split3 = Z_sol.sub(0), Z_sol.sub(1)

print('solve as a system')

print('full solution')
print(Z_sol.x.array)
print(len(Z_sol.x.array))

print('guess for wh using first part of Z_sol')
print(w_split)
print('guess for uh using last part of Z_sol')
print(u_split)

print('guess for wh using .split()')
print(w_split2.x.array)
print('guess for uh using .split()')
print(u_split2.x.array)

print('guess for wh using .sub(0)')
print(w_split3.x.array)
print('guess for uh using .sub(1)')
print(u_split3.x.array)

And what it should output:

from mpi4py import MPI
from dolfinx import mesh 
from dolfinx.mesh import locate_entities, meshtags 
from dolfinx.fem import FunctionSpace
from dolfinx import fem 
import numpy as np 
import ufl 
from petsc4py.PETSc import ScalarType
from petsc4py import PETSc 
from dolfinx import io 
import dolfinx

# Creating an interval as mesh
h = 0.25
x_min, x_max = -1,1
nx = int((x_max - x_min)/h)
domain = mesh.create_interval(MPI.COMM_WORLD, nx, points = (x_min, x_max))
x = ufl.SpatialCoordinate(domain) # Coordinates of domain used for defining general functions

V = FunctionSpace(domain, ("CG", 1)) 
U = FunctionSpace(domain, ("DG", 1))

# Constants
m  = 2
alpha_par = 1/3
beta_par = 1/3
kappa_par = 1/12
C = 1/32
t_start  = 0.3
dt = 0.1

def Phi(u):
    return u**m 

def exact_sol(x,t):
    return ufl.operators.max_value(0,1/(t**alpha_par) * (C - kappa_par * (ufl.operators.sqrt(x[0]**2)/(t**beta_par))**2 )**(1/(m-1)))
    
def weak_proj(func_u, phi_u, u): # Function that projects func_u on the space corresponding to phi_u and u
    a = ufl.inner(u, phi_u) * ufl.dx
    L = ufl.inner(func_u, phi_u) * ufl.dx
    
    bilinear_form = fem.form(a)
    linear_form = fem.form(L)
    
    b = fem.petsc.assemble_vector(linear_form)
    b.assemble()
    A = fem.petsc.assemble_matrix(bilinear_form)
    A.assemble()
    
    solver = PETSc.KSP().create(domain.comm)
    solver.setType(PETSc.KSP.Type.PREONLY)
    # solver.setTolerances(rtol=1e-12, atol = 1e-50)
    solver.getPC().setType(PETSc.PC.Type.LU) 
    
    solver.setOperators(A)
    
    solver.solve(b, u_proj.vector)
    u_proj.x.scatter_forward()
    return u_proj.x.array

initial_u = exact_sol(x,t_start)
expr_u = fem.Expression(initial_u, U.element.interpolation_points())

# Defining the function spaces for all the functions
u_n = fem.Function(U)
u_n.name = "u_n"

u_n_i = fem.Function(U)
u_n_i.name = "u_n_i"

# Defining the solution variables
uh = fem.Function(U)
uh.name = "uh"

wh = fem.Function(V)
wh.name = "wh"

u_proj = fem.Function(U)
u_proj.name = "u_proj"

w, phi = ufl.TrialFunction(V), ufl.TestFunction(V)
u, phi_u = ufl.TrialFunction(U), ufl.TestFunction(U)

u_n.interpolate(expr_u)
u_n_i.interpolate(expr_u)

L_par = 1
L_const = fem.Constant(domain, ScalarType(L_par))

a = ((w * phi/L_const) + dt * ufl.inner(ufl.grad(w), ufl.grad(phi))) * ufl.dx
L = ((Phi(u_n_i)/L_const) + u_n - u_n_i) * phi * ufl.dx

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

b = fem.petsc.assemble_vector(linear_form)
b.assemble()
A = fem.petsc.assemble_matrix(bilinear_form)
A.assemble()

solver = PETSc.KSP().create(domain.comm)
solver.setType(PETSc.KSP.Type.PREONLY) # Solver
#solver.setTolerances(rtol=1e-12, atol = 1e-50)
solver.getPC().setType(PETSc.PC.Type.LU) # Pre-conditioner
solver.setOperators(A)

solver.solve(b, wh.vector)
wh.x.scatter_forward()

uh.x.array[:] = weak_proj(wh - Phi(u_n_i) + L_const*u_n_i, phi_u,L_const*u)

print('correct wh and uh')
print('wh array')
print(wh.x.array)
print('uh array')
print(uh.x.array)

Any help is very much appreciated.

Kind regards,
Robin

Somebody just solved it for me. This will work:

w_split, u_split = Z_sol.sub(0).collapse(), Z_sol.sub(1).collapse()