Hello,
Thanks for your reply. Here is the full code attached.
from dolfinx import fem, mesh, plot,io
from dolfinx.fem import dirichletbc, Function, locate_dofs_topological
from dolfinx.io import (VTXWriter, distribute_entity_data, gmshio,XDMFFile)
from dolfinx.mesh import create_unit_square, locate_entities_boundary
from mpi4py import MPI
from petsc4py import PETSc
from dolfinx.fem.petsc import (apply_lifting, assemble_matrix, assemble_vector,
create_vector, create_matrix, set_bc,LinearProblem)
from petsc4py.PETSc import ScalarType
from ufl import (Circumradius, FacetNormal, SpatialCoordinate,FacetNormal, Identity, Measure, TestFunction, TrialFunction,
as_vector, div, dot, ds, dx, inner, lhs, grad, nabla_grad, rhs, sym, system)
from dolfinx.fem.petsc import assemble_vector, assemble_matrix, create_vector, apply_lifting, set_bc
import numpy as np
import ufl
import scipy.special as sp
import pyvista as pv
import matplotlib.pyplot as plt
import gmsh
import meshio
import csv
from pathlib import Path
import tqdm.autonotebook
import matplotlib as mpl
# Time-stepping parameters
t = 0
T = 2.0 # final time
dt = 0.001 # time step
num_steps = int(T / dt)
folder = Path("results_new_fen")
folder.mkdir(exist_ok=True, parents=True)
D_c = np.array([[0.57, 6, 0.1410* 0.57], [-0.012, 1.6, 1.8], [-0.0006, -0.08, 0.56]])
d11, d12, d13 = D_c[0, :]/D_c[0, 0]
d21, d22, d23 = D_c[1, :]/D_c[0, 0]
d31, d32, d33 = D_c[2, :]/D_c[0, 0]
a, b, c = 0.266, 0.366, 0.141
gamma = 160
mesh = mesh.create_unit_square(MPI.COMM_WORLD, nx=10, ny=10)
V = fem.functionspace(mesh, ("Lagrange", 1))
dx = ufl.Measure("dx", domain=mesh)
ds = ufl.Measure("ds", domain=mesh)
cells, types, x = plot.vtk_mesh(V)
grid = pv.UnstructuredGrid(cells, types, x)
plotter = pv.Plotter(off_screen=True)
plotter.add_mesh(grid, show_edges=True)
plotter.view_xy() # assigns plotting view
mesh_path = folder / "mesh_plot.jpg"
plotter.screenshot(filename=str(mesh_path),window_size=(1024, 768), transparent_background=False)
# Define trial and test functions
u, v, w = ufl.TrialFunction(V), ufl.TrialFunction(V), ufl.TrialFunction(V)
phi_1, phi_2, phi_3 = TestFunction(V), TestFunction(V), TestFunction(V)
u_ ,v_ , w_ = Function(V), Function(V), Function(V)
u_.name = "u_"
v_.name = "v_"
w_.name = "w_"
u_n, v_n, w_n = fem.Function(V), fem.Function(V), fem.Function(V)
u_n.name = "u_n"
v_n.name="w_n"
v_n.name="w_n"
u_star = a + b + c
v_star = b / (u_star**2)
w_star = c / (u_star**2)
def cosine_linear_combination(x, orders, scaling_factor=0.0001):
result = 0
result = np.zeros(x.shape[1])
for n, m in orders:
result += (
scaling_factor
* np.cos(n * np.pi * x[0, :])
* np.cos(m * np.pi * x[1, :])
)
return result
# Define the orders for the cosine terms
orders = [(1, 2), (2, 3), (1, 3)]
# Assign initial conditions
u_n.interpolate(lambda x: u_star + cosine_linear_combination(x, orders))
v_n.interpolate(lambda x: v_star + cosine_linear_combination(x, orders))
w_n.interpolate(lambda x: w_star + cosine_linear_combination(x, orders))
cells, types, x = plot.vtk_mesh(V)
grid_u = pv.UnstructuredGrid(cells, types, x)
grid_v = pv.UnstructuredGrid(cells, types, x)
grid_w = pv.UnstructuredGrid(cells, types, x)
grid_u.point_data["u"] = u_n.x.array.real
grid_v.point_data["v"] = v_n.x.array.real
grid_w.point_data["w"] = w_n.x.array.real
ini_path = folder / "AAA_initial_u_v_w.jpg"
plotter = pv.Plotter(shape=(1, 3), off_screen=True)
# First subplot: u_n
plotter.subplot(0, 0)
plotter.add_mesh(grid_u, show_edges=False, scalars="u", cmap="viridis")
plotter.add_axes()
plotter.view_xy()
plotter.add_title("u_n")
# Second subplot: v_n
plotter.subplot(0, 1)
plotter.add_mesh(grid_v, show_edges=False, scalars="v", cmap="viridis")
plotter.add_axes()
plotter.view_xy()
plotter.add_title("v_n")
# Third subplot: w_n
plotter.subplot(0, 2)
plotter.add_mesh(grid_w, show_edges=False, scalars="w", cmap="viridis")
plotter.add_axes()
plotter.view_xy()
plotter.add_title("w_n")
plotter.screenshot(str(ini_path), window_size=(1024, 768), transparent_background=False)
plotter.close()
sol_u = folder / "U_solution.xdmf"
sol_v = folder / "v_solution.xdmf"
sol_w = folder / "w_solution.xdmf"
# Write mesh and solution to the XDMF file
with io.XDMFFile(MPI.COMM_WORLD, str(sol_u), "w") as xdmf_u:
xdmf_u.write_mesh(mesh)
u_.interpolate(lambda x: u_star + cosine_linear_combination(x, orders))
xdmf_u.write_function(u_, t)
with io.XDMFFile(MPI.COMM_WORLD, str(sol_v), "w") as xdmf_v:
xdmf_v.write_mesh(mesh)
v_.interpolate(lambda x: v_star + cosine_linear_combination(x, orders))
xdmf_v.write_function(v_, t)
with io.XDMFFile(MPI.COMM_WORLD, str(sol_w), "w") as xdmf_w:
xdmf_w.write_mesh(mesh)
w_.interpolate(lambda x: w_star + cosine_linear_combination(x, orders))
xdmf_w.write_function(w_, t)
def gg(u_n):
return u_n*u_n
lhs_um = u * phi_1 * ufl.dx + dt * (d11 * ufl.inner(ufl.grad(u), ufl.grad(phi_1)) * ufl.dx)\
+ dt * gamma * u *phi_1 * ufl.dx
lhs_vm = v * phi_2 * ufl.dx + dt * (d22 * ufl.inner(ufl.grad(v), ufl.grad(phi_2))* ufl.dx)\
+ dt * gamma * gg(u_n) * v * phi_2 * ufl.dx
lhs_wm = w * phi_3 * ufl.dx + dt *( d33 * ufl.inner(ufl.grad(w), ufl.grad(phi_3)) * ufl.dx)\
+ dt * gamma * gg(u_n)* w *phi_3 * ufl.dx
lhs_u, lhs_v, lhs_w = fem.form(lhs_um), fem.form(lhs_vm), fem.form(lhs_wm)
Au, Av, Aw = assemble_matrix(lhs_u), create_matrix(lhs_v), create_matrix(lhs_w)
Au.assemble()
rhs_um = u_n * phi_1 * dx + dt * gamma * a * phi_1 * ufl.dx\
+ dt * gamma * gg(u_n) * v_n * phi_1 * ufl.dx \
+ dt * gamma * gg(u_n)* w_n * phi_1 * ufl.dx\
- dt * (d12* dot(grad(v_n),grad(phi_1)) * ufl.dx) \
- dt * (d13* dot(grad(w_n),grad(phi_1)) * ufl.dx)
rhs_vm = v_n * phi_2 * dx + dt * gamma * b * phi_2 * ufl.dx\
- dt * (d21* dot(grad(u_n),grad(phi_2)) * ufl.dx)\
- dt * (d23* dot(grad(w_n),grad(phi_2)) * ufl.dx)
rhs_wm = w_n * phi_3 * dx + dt * gamma * c * phi_3 * ufl.dx\
- dt * (d31* dot(grad(u_n),grad(phi_3)) * ufl.dx)\
- dt * (d32* dot(grad(v_n),grad(phi_3)) * ufl.dx)
rhs_u, rhs_v, rhs_w = fem.form(rhs_um), fem.form(rhs_vm), fem.form(rhs_vm)
bu, bv, bw = create_vector(rhs_u), create_vector(rhs_v), create_vector(rhs_w)
bcs_u =[]
# Solver for u
solver1 = PETSc.KSP().create(mesh.comm)
solver1.setOperators(Au)
solver1.setType(PETSc.KSP.Type.BCGS)
#solver1.setType(PETSc.KSP.Type.GMRES)
#setType(PETSc.KSP.Type.GMRES)
pc1 = solver1.getPC()
pc1.setType(PETSc.PC.Type.JACOBI)
#pc1.setType(PETSc.PC.Type.ILU)
solver1.setTolerances(rtol=1e-10, atol=1e-10, max_it=2000)
# Solver v
solver2 = PETSc.KSP().create(mesh.comm)
solver2.setOperators(Av)
solver2.setType(PETSc.KSP.Type.BCGS)
pc2 = solver2.getPC()
pc2.setType(PETSc.PC.Type.JACOBI)
#pc2.setHYPREType("boomeramg")
# Solver for w
solver3 = PETSc.KSP().create(mesh.comm)
solver3.setOperators(Aw)
solver3.setType(PETSc.KSP.Type.BCGS)
pc3 = solver3.getPC()
pc3.setType(PETSc.PC.Type.JACOBI)
pv.start_xvfb()
grid_1 = pv.UnstructuredGrid(*plot.vtk_mesh(V))
plotter1 = pv.Plotter()
plotter1.open_gif("u_time.gif", fps=20)
grid_1.point_data["u_"] = u_.x.array
warped = grid_1.warp_by_scalar("u_", factor=1)
viridis = mpl.colormaps.get_cmap("viridis").resampled(25)
sargs = dict(title_font_size=25, label_font_size=20, fmt="%.2e", color="black",
position_x=0.1, position_y=0.8, width=0.8, height=0.1)
renderer = plotter1.add_mesh(warped, show_edges=True, lighting=False,
cmap=viridis, scalar_bar_args=sargs,
clim=[0, max(u_.x.array)])
progress = tqdm.autonotebook.tqdm(desc="Solving RDE", total=num_steps)
for i in range(num_steps):
progress.update(1)
# Update current time step
t += dt
with bu.localForm() as loc:
loc.set(0)
assemble_vector(bu, rhs_u)
apply_lifting(bu, [lhs_u], bcs=[bcs_u])
bv.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
set_bc(bu, bcs_u)
solver1.solve(bu, u_.x.petsc_vec)
u_.x.scatter_forward()
# Solver 2
Av.zeroEntries()
assemble_matrix(Av, lhs_v)
Av.assemble()
with bv.localForm() as loc:
loc.set(0)
assemble_vector(bv, rhs_v)
apply_lifting(bv, [lhs_v], bcs=[bcs_u])
#apply_lifting(b2, [a2], [bcp])
bv.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
set_bc(bv, bcs_u)
solver2.solve(bv, v_.x.petsc_vec)
v_.x.scatter_forward()
# Solver 3
Aw.zeroEntries()
assemble_matrix(Aw, lhs_w)
Aw.assemble()
with bw.localForm() as loc:
loc.set(0)
assemble_vector(bw, rhs_w)
apply_lifting(bv, [lhs_w], bcs=[bcs_u])
bw.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
set_bc(bw, bcs_u)
solver3.solve(bw, w_.x.petsc_vec)
w_.x.scatter_forward()
u_n.x.array[:] = u_.x.array
v_n.x.array[:] = v_.x.array
w_n.x.array[:] = w_.x.array
# Write solution to file
xdmf_u.write_function(u_, t)
# Update plot
new_warped = grid_1.warp_by_scalar("u_", factor=1)
warped.points[:, :] = new_warped.points
warped.point_data["u_"][:] = u_.x.array
plotter1.write_frame()
plotter1.close()
xdmf_u.close()