Dear Fenicsx community,
I am facing an issue when solving coupled time dependent PDEs. I have a code that tries to save, after each successful time step, the fields using adios4dolfinx.
However, when I try to restart a simulation by reading the field saved as initial data, the values of the next iteration appear to be shuffled.
I attach this MWE. I am sorry for its length, but I have tried shorter examples that do not capture my problem (i.e. only one time step, or only one PDE). I am using fenicsx 0.9.0.
We want to solve two coupled PDEs, let say a diffusion problem for a velocity field u and the transport of a scalar field \varphi by u. I use Crank-Nicolson for both PDEs (note that I do not know if this is an appropriate method, this is just for the MWE and to keep it simple):
\int_{\Omega} \frac{u^{n+1}-u^n}{\tau}\cdot v + \int_\Omega \nabla \left (\frac{u^n + u^{n+1}}{2}\right ):\nabla v = \int_\Omega\frac{f^n + f^{n+1}}{2}\cdot v,
\int_\Omega \frac{\varphi^{n+1}-\varphi^n}{\tau}\psi+\frac12\int_\Omega (u^n\cdot\nabla \varphi^n)\psi + \frac12\int_\Omega (u^{n+1}\cdot\nabla \varphi^{n+1})\psi=0.
The code solves for 10 time steps, stops, and then try to restart for 10 other time step. Somehow, the functions are read back properly, but the first step after restart ‘shuffle’ the values. I have pictures of the fields at steps n=10 and n=11 that shows the issue.
### Import modules ###
import adios4dolfinx
import basix
import ufl
import dolfinx
from functools import partial
import gmsh
import meshio
from dolfinx.io import gmshio
import numpy as np
from petsc4py import PETSc
from mpi4py import MPI
from dolfinx import fem, mesh
from dolfinx.fem.petsc import assemble_vector, assemble_matrix, create_vector, apply_lifting, set_bc
import os
import h5py
from pathlib import Path
import matplotlib.pyplot as plt
### A few classes ###
class Solver:
def __init__(self, domain, dt, f):
self.domain = domain
self.f = f
self.Vu = fem.functionspace(domain, ('CG', 1, (2,)))
self.u_n = fem.Function(self.Vu)
self.uh = fem.Function(self.Vu)
self.f_n = fem.Function(self.Vu)
self.fh = fem.Function(self.Vu)
self.u = ufl.TrialFunction(self.Vu)
self.v = ufl.TestFunction(self.Vu)
self.Vphi = fem.functionspace(domain, ('CG', 1))
self.phi_n = fem.Function(self.Vphi)
self.phih = fem.Function(self.Vphi)
self.phi = ufl.TrialFunction(self.Vphi)
self.psi = ufl.TestFunction(self.Vphi)
self.dt_cste = fem.Constant(domain, PETSc.ScalarType(dt))
F_u = (ufl.dot((self.u - self.u_n) / self.dt_cste, self.v) + ufl.inner(ufl.grad((self.u + self.u_n)/2), ufl.grad(self.v)) + ufl.dot((self.f_n + self.fh)/2, self.v)) * ufl.dx
F_phi = ((self.phi - self.phi_n) / self.dt_cste + 0.5 * ufl.dot(self.u_n, ufl.grad(self.phi_n)) + 0.5 * ufl.dot(self.uh, ufl.grad(self.phi))) * self.psi * ufl.dx
self.a_u = ufl.lhs(F_u)
self.L_u = ufl.rhs(F_u)
self.bilinear_form_u = fem.form(self.a_u)
self.linear_form_u = fem.form(self.L_u)
self.a_phi = ufl.lhs(F_phi)
self.L_phi = ufl.rhs(F_phi)
self.bilinear_form_phi = fem.form(self.a_phi)
self.linear_form_phi = fem.form(self.L_phi)
self.bc_u = None # To update using set_bc_u
self.A_u = None # Create using assemble_u
self.rhs_u = None # Create using create_rhs_u
self.solver_u = None # Create using setup_solver_u
self.bc_phi = None # To update using set_bc_phi
self.A_phi = None # Create using assemble_phi
self.rhs_phi = None # Create using create_rhs_phi
self.solver_phi = None # Create using setup_solver_phi
def set_bc_u(self, g):
'''
A function to create the bcs from g.
'''
uD = fem.Function(self.Vu)
uD.interpolate(g)
tdim = self.domain.topology.dim
fdim = tdim - 1
self.domain.topology.create_connectivity(fdim, tdim)
boundary_facets = mesh.exterior_facet_indices(self.domain.topology)
boundary_dofs = fem.locate_dofs_topological(self.Vu, fdim, boundary_facets)
self.bc_u = fem.dirichletbc(uD, boundary_dofs)
return self.bc_u
def set_bc_phi(self):
'''
We choose a MWE with no inflow
'''
pass
def assemble_A_u(self):
'''
Assemble stiffness matrix.
'''
if not self.bc_u:
raise ValueError('Boundary conditions must be set using self.set_bc_u before assembling A_u')
self.A_u = assemble_matrix(self.bilinear_form_u, bcs=[self.bc_u])
self.A_u.assemble()
return self.A_u
def assemble_A_phi(self):
self.A_phi = assemble_matrix(self.bilinear_form_phi)
self.A_phi.assemble()
return self.A_phi
def create_rhs_u(self):
'''
Create right hand side vector
'''
self.rhs_u = create_vector(self.linear_form_u)
return self.rhs_u
def create_rhs_phi(self):
self.rhs_phi = create_vector(self.linear_form_phi)
return self.rhs_phi
def update_rhs_u(self, t2, t3):
'''
Update the right hand side
'''
self.f_n.interpolate(partial(self.f, t=t2)) # Update source term
self.fh.interpolate(partial(self.f, t=t3)) # Update source term
# Update the right hand side reusing the initial vector
with self.rhs_u.localForm() as loc_b:
loc_b.set(0)
assemble_vector(self.rhs_u, self.linear_form_u)
# Apply Dirichlet boundary condition to the vector
apply_lifting(self.rhs_u, [self.bilinear_form_u], [[self.bc_u]])
self.rhs_u.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
set_bc(self.rhs_u, [self.bc_u])
return self.rhs_u
def update_rhs_phi(self):
with self.rhs_phi.localForm() as loc_rhs_phi:
loc_rhs_phi.set(0)
assemble_vector(self.rhs_phi, self.linear_form_phi)
self.rhs_phi.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
return self.rhs_phi
def setup_solver_u(self):
'''
Setup the solver for the linear problem
'''
self.solver_u = PETSc.KSP().create(self.domain.comm)
self.solver_u.setOperators(self.A_u)
self.solver_u.setType(PETSc.KSP.Type.PREONLY)
self.solver_u.getPC().setType(PETSc.PC.Type.LU)
return self.solver_u
def setup_solver_phi(self):
'''
Setup the solver for the linear problem
'''
self.solver_phi = PETSc.KSP().create(self.domain.comm)
self.solver_phi.setOperators(self.A_phi)
self.solver_phi.setType(PETSc.KSP.Type.PREONLY)
self.solver_phi.getPC().setType(PETSc.PC.Type.LU)
return self.solver_phi
### Initialize, save and restart functions ###
def initialize(dir : str, dt : float):
os.makedirs(dir)
# Create the h5 file (contains the time) #
file_h5 = dir + '/info.h5'
with h5py.File(file_h5, 'w') as h5f:
mock_data = np.zeros(1)
h5f.create_dataset('t', data=mock_data, maxshape=(None,), dtype='float64') # Time stamps
# Create the mesh from a disk #
gmsh.initialize()
gmsh.model.add("disk")
disk = gmsh.model.occ.addDisk(0.0, 0.0, 0.0, 1.0, 1.0)
gmsh.model.occ.synchronize()
gmsh.model.addPhysicalGroup(2, [disk], tag=1)
gmsh.model.setPhysicalName(2, 1, "DiskSurface")
gmsh.option.setNumber("Mesh.CharacteristicLengthMin", 0.05)
gmsh.option.setNumber("Mesh.CharacteristicLengthMax", 0.05)
gmsh.model.mesh.generate(2)
mesh_data = gmshio.model_to_mesh(gmsh.model, MPI.COMM_WORLD, 0, gdim=2)
gmsh.finalize()
domain = mesh_data[0]
# Initialize the solver #
f = lambda x, t : (5*x[1] * np.sin(t), -5*x[0] * np.sin(t))
solver = Solver(domain, dt, f)
# Interpolate intial condition #
u0 = lambda x : (-5*x[1], 5*x[0])
phi0 = lambda x : np.exp(- ((x[0] - 0.5)**2 + (x[1] - 0)**2) / 0.2**2)
solver.u_n.interpolate(u0)
solver.phi_n.interpolate(phi0)
# Assemble forms #
g = lambda x : (0*x[0], 0*x[0])
solver.set_bc_u(g)
solver.assemble_A_u()
solver.assemble_A_phi()
solver.create_rhs_u()
solver.create_rhs_phi()
solver.setup_solver_u()
solver.setup_solver_phi()
# Useful values #
solver.t = 0 # Time of the initial condition
solver.n = 0 # Current time step
return solver
def restart(dir : str, T : float):
# Read the times #
file_h5 = dir + '/info.h5'
with h5py.File(file_h5, 'r') as f:
times = f['t'][:]
# Check if finished #
if float(times[-1]) >= T or np.isclose(times[-1], T):
print('Simulation is already finished, skipping!')
return None
dt = times[-1] - times[-2]
N = len(times)
# Read last solution #
u = read_field(dir, 'u', N-1)
phi = read_field(dir, 'phi', N-1)
f = lambda x, t : (5*x[1] * np.sin(t), -5*x[0] * np.sin(t))
solver = Solver(u.function_space.mesh, dt, f)
solver.u_n.x.array[:] = u.x.array # Set u_n to last known value of u
solver.phi_n.x.array[:] = phi.x.array
# Assemble forms #
g = lambda x : (0*x[0], 0*x[0])
solver.set_bc_u(g)
solver.assemble_A_u()
solver.assemble_A_phi()
solver.create_rhs_u()
solver.create_rhs_phi()
solver.setup_solver_u()
solver.setup_solver_phi()
# Useful values #
solver.t = times[-1]
solver.n = N-1
return solver
def save(dir : str, solver : Solver):
file_h5 = dir + '/info.h5'
with h5py.File(file_h5, 'a') as h5f: # open the file on which to save the data
dataset = h5f['t']
dataset.resize((len(dataset) + 1,))
dataset[-1] = solver.t
# Save fields #
path = Path(dir)
adios4dolfinx.write_function(path, solver.uh, time=solver.n, name='u')
adios4dolfinx.write_function(path, solver.phih, time=solver.n, name='phi')
def read_field(dir : str, name : str, n : int):
filename = Path(dir)
domain = adios4dolfinx.read_mesh(filename, comm=MPI.COMM_WORLD)
V = fem.functionspace(domain, ('CG', 1)) if name=='phi' else fem.functionspace(domain, ('CG', 1, (2,)))
func = fem.Function(V)
adios4dolfinx.read_function(filename, func, time=n, name=name)
return func
### Solve the PDE ###
def solve(T : float, dt : float, dir : str):
"""
Solve the heat equation
Input :
- T : final time to reach
- dt : time step
- dir : directory in which to save the data
"""
if os.path.exists(dir):
print('Directory already exists! Attempting a restart...')
solver = restart(dir, T)
if solver is None:
return None
else:
solver = initialize(dir, dt)
path = Path(dir)
adios4dolfinx.write_mesh(path, solver.domain, engine="BP4")
adios4dolfinx.write_function(path, solver.u_n, time=0, name="u")
adios4dolfinx.write_function(path, solver.phi_n, time=0, name="phi")
while solver.t < T and not np.isclose(solver.t, T):
solver.update_rhs_u(solver.t, solver.t + dt)
solver.solver_u.solve(solver.rhs_u, solver.uh.x.petsc_vec)
solver.uh.x.scatter_forward()
solver.update_rhs_phi()
solver.assemble_A_phi()
solver.setup_solver_phi()
solver.solver_phi.solve(solver.rhs_phi, solver.phih.x.petsc_vec)
solver.phih.x.scatter_forward()
solver.t += dt
solver.n += 1
save(dir, solver)
solver.u_n.x.array[:] = solver.uh.x.array
solver.phi_n.x.array[:] = solver.phih.x.array
### Test the code ###
dir = 'MWE_restart-2-fields'
dt = 0.05
T = 0.5
solve(T, dt, dir)
T = 1
solve(T, dt, dir)
### Display solution ###
from matplotlib.tri import Triangulation as Triangulation_plt
from matplotlib.animation import FuncAnimation
cmap = 'viridis'
lw = 0.5
file_h5 = dir + '/info.h5'
with h5py.File(file_h5, 'r') as f:
times = f['t'][:]
total_time = 5000
N = len(times)
min_phih_array = np.zeros(N)
max_phih_array = np.zeros(N)
for n in range(N):
phih = read_field(dir, 'phi', n)
min_phih_array[n] = np.min(phih.x.array)
max_phih_array[n] = np.max(phih.x.array)
min_phih = np.min(min_phih_array)
max_phih = np.max(max_phih_array)
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
phi0 = read_field(dir, 'phi', 0)
domain = phi0.function_space.mesh
tri_plt = Triangulation_plt(domain.geometry.x[:, 0], domain.geometry.x[:, 1], domain.topology.connectivity(domain.topology.dim, 0).array.reshape((-1, 3)))
ax.set_xlim(-1, 1)
ax.set_ylim(-1, 1)
ax.set_zlim(min_phih, max_phih)
# ax.set_aspect('equal')
ax.plot_trisurf(tri_plt, phi0.x.array, cmap=cmap, edgecolor='none', vmin=min_phih, vmax=max_phih)
def update(i):
ax.clear()
phi = read_field(dir, 'phi', i)
domain = phi.function_space.mesh
tri_plt = Triangulation_plt(domain.geometry.x[:, 0], domain.geometry.x[:, 1], domain.topology.connectivity(domain.topology.dim, 0).array.reshape((-1, 3)))
ax.plot_trisurf(tri_plt, phi.x.array, cmap=cmap, edgecolor='none', vmin=min_phih, vmax=max_phih)
ax.set_title(f't = {times[i]:.4f}')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlim(min_phih, max_phih)
anim = FuncAnimation(fig, update, frames=len(times), interval=total_time/len(times))
plt.show()
In my attempts to find the issue I have tried (I can provide MWE)
- Restart for the diffusion problem with only one field : this works well.
- Just one step of transport problem where the velocity field is read with adios4dolfinx : this works well.
Do you have any idea of what is going wrong when restarting? The reading of the function is fine, I have checked that, but solving a linear system built from the functions goes wrong.
Is it my read_field function which is not suited for this problem?
Thank you very much for your help,
Théophile