Solution vectors and nodes not written properly to file

Hi,

I’m trying to write part of a solution into a file, with each line having three values for the node coordinates and the three following values being flow velocity components.

The code is the following:

from dolfinx import mesh, fem, io, log
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.nls.petsc import NewtonSolver

from mpi4py import MPI

import numpy as np

import ufl
from ufl import inner, div, grad, dx

import basix

import time

from petsc4py import PETSc


height = 0.06
width  = 0.08

n_h = 20
n_w = 10


msh = mesh.create_box(MPI.COMM_WORLD, ((-width/2, -width/2, 0.0), (width/2,width/2,height)), (n_h,n_h,n_w), mesh.CellType.tetrahedron)


fdim = msh.topology.dim - 1





rho_air = fem.Constant(msh, 1.225)
mu_air  = fem.Constant(msh, 1.8e-5)



# Function to mark noslip
def noslip_boundary1(x):
    return np.isclose(x[2], height)



P2 = basix.ufl.element("Lagrange", msh.topology.cell_name(), 2, shape=(msh.geometry.dim, ))
P1 = basix.ufl.element("Lagrange", msh.topology.cell_name(), 1)

V, Q = fem.functionspace(msh, P2), fem.functionspace(msh, P1)

# Create the function space
TH = basix.ufl.mixed_element([P2, P1])
W = fem.functionspace(msh, TH)
W0, _ = W.sub(0).collapse()


# No slip boundary condition 1
noslip  = fem.Function(V)
facets1 = mesh.locate_entities_boundary(msh, fdim, noslip_boundary1)
dofs1   = fem.locate_dofs_topological((W.sub(0), V), fdim, facets1)
bc0     = fem.dirichletbc(noslip, dofs1, W.sub(0))


# Since for this problem the pressure is only determined up to a constant, we pin it at a point
zero = fem.Function(Q)
zero.x.array[:] = 0
dofs = fem.locate_dofs_geometrical(
    (W.sub(1), Q), lambda x: np.isclose(x.T, [0, 0, 0]).all(axis=1))
bc2 = fem.dirichletbc(zero, dofs, W.sub(1))



def body_force_fun(x):
    F_max    = 2.0e-2
    z_offset = 0.02

    r_xy = np.sqrt(x[0]**2 + x[1]**2)

    a = 0.005
    max_val = a * np.exp(-0.5)
    Fr_val  = np.exp(-r_xy**2/(2*a**2)) / max_val
        
    b = 0.005
    Fz_val = np.exp(-(x[2]-z_offset)**2/(2*b**2))
    F_mag  = F_max * Fr_val * Fz_val
      
    # -ve sign to get clockwise rotation from array top
    return -F_mag * np.vstack((x[1], -x[0], np.zeros_like(x[2])))
        

Force_per_unit_vol = fem.Function(V)
Force_per_unit_vol.interpolate(lambda x: body_force_fun(x))




# Collect Dirichlet boundary conditions
bcs = [bc0, bc2]

# Define variational problem
w = fem.Function(W)

(u, p) = ufl.split(w)
(v, q) = ufl.TestFunctions(W)


F = rho_air*inner(grad(u)*u, v)*dx + \
    mu_air*inner(grad(u), grad(v))*dx + \
    inner(Force_per_unit_vol, v)*dx - \
    inner(p, div(v))*dx - \
    div(u)*q*dx

dw = ufl.TrialFunction(W)
dF = ufl.derivative(F, w, dw)

problem = NonlinearProblem(F, w, bcs=bcs, J=dF)

solver = NewtonSolver(MPI.COMM_WORLD, problem)
solver.convergence_criterion = "incremental"
solver.rtol   = 1e-6
#solver.max_it = 2
solver.report = True
solver.error_on_nonconvergence = False


# Custom solver options 
opts = PETSc.Options()
ksp = solver.krylov_solver
option_prefix = ksp.getOptionsPrefix()
opts[f"{option_prefix}ksp_type"] = "preonly"
opts[f"{option_prefix}pc_type"] = "lu"
opts[f"{option_prefix}pc_factor_mat_solver_type"] = "mumps"
ksp.setFromOptions()



# Compute the solution
log.set_log_level(log.LogLevel.INFO)
start = time.time()
solver.solve(w)
end   = time.time()

print('\n')
print(msh.topology.index_map(msh.topology.dim).size_local)
print((end - start) / 60)


# Split the mixed solution and collapse
u = w.sub(0).collapse()
p = w.sub(1).collapse()



# Write to file
u2 = fem.functionspace(msh, P2)

coords = u2.tabulate_dof_coordinates()


with open('coords_vcomps.txt', 'w') as file:
    
    for i, coord in enumerate(coords):
        file.write("%s   %s   %s   " % (coord[0], coord[1], coord[2]))
        
        for b in range(u2.dofmap.bs):
            file.write("%s   " % (u.x.array[i*u2.dofmap.bs+b] * 1000,))
        
        file.write('\n')



# Write the solution to file
with io.XDMFFile(MPI.COMM_WORLD, "velocity.xdmf", "w") as ufile_xdmf:
    u.x.scatter_forward()
    P1 = basix.ufl.element("Lagrange", msh.basix_cell(), 1, shape=(msh.geometry.dim,))
    u1 = fem.Function(fem.functionspace(msh, P1))
    u1.interpolate(u)
    ufile_xdmf.write_mesh(msh)
    ufile_xdmf.write_function(u1)

The file can be checked using the following code:

import numpy as np
import matplotlib.pyplot as plt


# File path to your text file
filename = 'coords_vcomps.txt'

# Load data, skipping the first 8 lines (starts from line 9)
data = np.loadtxt(filename, skiprows=0)


print(data.shape)



# Assuming the columns are: x, y, z, vx, vy, vz
n_skip = 100

v_points = data[::n_skip, 0:3]

#v_points[:,2] = v_points[:,2] + 0.01

v_comps  = data[::n_skip, 3:]

print(v_points.shape)


print(np.max(np.sqrt(v_comps[:,0]**2 + v_comps[:,1]**2 + v_comps[:,2]**2)))


# Create a 3D vector field plot
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# Plot vector field using quiver
ax.quiver(v_points[:,0], v_points[:,1], v_points[:,2], 
          v_comps[:,0],  v_comps[:,1],  v_comps[:,2], length=0.004, normalize=False, color='blue')

# Set axis labels
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_title('3D Vector Field')

ax.set_aspect('auto')

plt.show()

Thank you,

Alex

Could you be specific about what issue you are observing?
As you haven’t provided a plot of what it should be (expected, as what is probably viewed in Paraview) compared to what you get.

Secondly,

why aren’t you using coords=u.function_space.tabulate_dof_coordinates()
and similarly

should be
file.write("%s " % (u.x.array[i*u.function_space.dofmap.bs+b] * 1000,))

Here is my complete code, with magnitude colors plotted as well

from dolfinx import mesh, fem, io, log
from dolfinx.fem.petsc import NonlinearProblem
from dolfinx.nls.petsc import NewtonSolver

from mpi4py import MPI

import numpy as np

import ufl
from ufl import inner, div, grad, dx

import basix

import time

from petsc4py import PETSc


height = 0.06
width = 0.08

n_h = 10
n_w = 10


msh = mesh.create_box(
    MPI.COMM_WORLD,
    ((-width / 2, -width / 2, 0.0), (width / 2, width / 2, height)),
    (n_h, n_h, n_w),
    mesh.CellType.tetrahedron,
)


fdim = msh.topology.dim - 1


rho_air = fem.Constant(msh, 1.225)
mu_air = fem.Constant(msh, 1.8e-5)


# Function to mark noslip
def noslip_boundary1(x):
    return np.isclose(x[2], height)


P2 = basix.ufl.element(
    "Lagrange", msh.topology.cell_name(), 2, shape=(msh.geometry.dim,)
)
P1 = basix.ufl.element("Lagrange", msh.topology.cell_name(), 1)

V, Q = fem.functionspace(msh, P2), fem.functionspace(msh, P1)

# Create the function space
TH = basix.ufl.mixed_element([P2, P1])
W = fem.functionspace(msh, TH)
W0, _ = W.sub(0).collapse()


# No slip boundary condition 1
noslip = fem.Function(V)
facets1 = mesh.locate_entities_boundary(msh, fdim, noslip_boundary1)
dofs1 = fem.locate_dofs_topological((W.sub(0), V), fdim, facets1)
bc0 = fem.dirichletbc(noslip, dofs1, W.sub(0))


# Since for this problem the pressure is only determined up to a constant, we pin it at a point
zero = fem.Function(Q)
zero.x.array[:] = 0
dofs = fem.locate_dofs_geometrical(
    (W.sub(1), Q), lambda x: np.isclose(x.T, [0, 0, 0]).all(axis=1)
)
bc2 = fem.dirichletbc(zero, dofs, W.sub(1))


def body_force_fun(x):
    F_max = 2.0e-2
    z_offset = 0.02

    r_xy = np.sqrt(x[0] ** 2 + x[1] ** 2)

    a = 0.005
    max_val = a * np.exp(-0.5)
    Fr_val = np.exp(-(r_xy**2) / (2 * a**2)) / max_val

    b = 0.005
    Fz_val = np.exp(-((x[2] - z_offset) ** 2) / (2 * b**2))
    F_mag = F_max * Fr_val * Fz_val

    # -ve sign to get clockwise rotation from array top
    return -F_mag * np.vstack((x[1], -x[0], np.zeros_like(x[2])))


Force_per_unit_vol = fem.Function(V)
Force_per_unit_vol.interpolate(lambda x: body_force_fun(x))


# Collect Dirichlet boundary conditions
bcs = [bc0, bc2]

# Define variational problem
w = fem.Function(W)

(u, p) = ufl.split(w)
(v, q) = ufl.TestFunctions(W)


F = (
    rho_air * inner(grad(u) * u, v) * dx
    + mu_air * inner(grad(u), grad(v)) * dx
    + inner(Force_per_unit_vol, v) * dx
    - inner(p, div(v)) * dx
    - div(u) * q * dx
)

dw = ufl.TrialFunction(W)
dF = ufl.derivative(F, w, dw)

problem = NonlinearProblem(F, w, bcs=bcs, J=dF)

solver = NewtonSolver(MPI.COMM_WORLD, problem)
solver.convergence_criterion = "incremental"
solver.rtol = 1e-6
# solver.max_it = 2
solver.report = True
solver.error_on_nonconvergence = False


# Custom solver options
opts = PETSc.Options()
ksp = solver.krylov_solver
option_prefix = ksp.getOptionsPrefix()
opts[f"{option_prefix}ksp_type"] = "preonly"
opts[f"{option_prefix}pc_type"] = "lu"
opts[f"{option_prefix}pc_factor_mat_solver_type"] = "mumps"
ksp.setFromOptions()


# Compute the solution
log.set_log_level(log.LogLevel.INFO)
start = time.time()
solver.solve(w)
end = time.time()

print("\n")
print(msh.topology.index_map(msh.topology.dim).size_local)
print((end - start) / 60)


# Split the mixed solution and collapse
u = w.sub(0).collapse()
p = w.sub(1).collapse()


# Write to file
u2 = fem.functionspace(msh, P2)

coords = u.function_space.tabulate_dof_coordinates()

with open("coords_vcomps.txt", "w") as file:
    for i, coord in enumerate(coords):
        file.write("%s   %s   %s   " % (coord[0], coord[1], coord[2]))

        for b in range(u2.dofmap.bs):
            file.write(
                "%s   " % (u.x.array[i * u.function_space.dofmap.bs + b] * 1000,)
            )

        file.write("\n")


# Write the solution to file
with io.XDMFFile(MPI.COMM_WORLD, "velocity.xdmf", "w") as ufile_xdmf:
    u.x.scatter_forward()
    P1 = basix.ufl.element("Lagrange", msh.basix_cell(), 1, shape=(msh.geometry.dim,))
    u1 = fem.Function(fem.functionspace(msh, P1))
    u1.interpolate(u)
    ufile_xdmf.write_mesh(msh)
    ufile_xdmf.write_function(u1)


import numpy as np
import matplotlib.cm as cm
import matplotlib.pyplot as plt


# File path to your text file
filename = "coords_vcomps.txt"

# Load data, skipping the first 8 lines (starts from line 9)
data = np.loadtxt(filename, skiprows=0)


print(data.shape)


# Assuming the columns are: x, y, z, vx, vy, vz
n_skip = 10

v_points = data[::n_skip, 0:3]

# v_points[:,2] = v_points[:,2] + 0.01

v_comps = data[::n_skip, 3:]
v_comps *= 1 / 1000
print(v_points.shape)


print(np.max(np.sqrt(v_comps[:, 0] ** 2 + v_comps[:, 1] ** 2 + v_comps[:, 2] ** 2)))

magnitudes = np.sqrt(v_comps[:, 0] ** 2 + v_comps[:, 1] ** 2 + v_comps[:, 2] ** 2)

# Create a 3D vector field plot
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection="3d")

import matplotlib.cm as cm
from matplotlib.colors import Normalize

# Normalize the magnitudes for colormap mapping
norm = Normalize(vmin=magnitudes.min(), vmax=magnitudes.max())
colormap = cm.viridis  # Choose a colormap


# Plot vector field using quiver
ax.quiver(
    v_points[:, 0],
    v_points[:, 1],
    v_points[:, 2],
    v_comps[:, 0],
    v_comps[:, 1],
    v_comps[:, 2],
    colors=colormap(norm(magnitudes)),
    # color="blue",
)

# Set axis labels
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.set_title("3D Vector Field")

ax.set_aspect("auto")

plt.show()

which to me looks sufficiently close to

Thank you! Now it works. I didn’t know you could use u.function_space.tabulate_dof_coordinates()