1 from mpi4py import MPI
2 import numpy as np
3 import ufl
4 from dolfinx import mesh
5 from dolfinx import fem
6 from dolfinx import default_scalar_type
7 from dolfinx.fem.petsc import NonlinearProblem
8 from dolfinx.nls.petsc import NewtonSolver
9
10
11 tf = 5.
12
13
14 theta = 0.5
15
16
17 nt = 100
18 ht = tf/nt
19
20
21 domain = mesh.create_unit_interval(MPI.COMM_WORLD,
22 nx = 5)
23 x = ufl.SpatialCoordinate(domain)
24
25
26 V = fem.functionspace(domain, ('P', 1))
27
28
29 v = ufl.TestFunction(V)
30 u = fem.Function(V)
31
32
33 t = 0.
34 u0 = fem.Function(V)
35 u0.interpolate(lambda x: np.cos(np.pi*x[0])**2)
36
37
38 u.x.array[:] = u0.x.array[:]
39
40
41 from dolfinx import io
42 from pathlib import Path
43 results_folder = Path("results")
44 results_folder.mkdir(exist_ok=True, parents=True)
45
46
47 filename = results_folder / f"u_{0:0>6}"
48 with io.XDMFFile(domain.comm, filename.with_suffix(".xdmf"), "w") as xdmf:
49 xdmf.write_mesh(domain)
50 xdmf.write_function(u, 0.)
51
52
53
54 for k in range(nt):
55
56 t += ht
57 print(f"{k+1}: t = {t:.4g}")
58
59
60
61 F = 1./ht * u * v * ufl.dx
62 F -= 1./ht * u0 * v * ufl.dx
63
64 F += theta * ufl.dot(ufl.grad(u), ufl.grad(v)) * ufl.dx
65 F += (1.-theta) * ufl.dot(ufl.grad(u0), ufl.grad(v)) * ufl.dx
66
67 F -= theta * u * (1. - u) * v * ufl.dx
68 F -= (1.-theta) * u0 * (1. - u0) * v * ufl.dx
69
70 problem = NonlinearProblem(F, u)
71 solver = NewtonSolver(MPI.COMM_WORLD, problem)
72 n, converged = solver.solve(u)
73 print(f"\tNewton iterations: {n}")
74 assert(converged)
75
76
77 filename = results_folder / f"u_{k+1:0>6}"
78 with io.XDMFFile(domain.comm, filename.with_suffix(".xdmf"), "w") as xdmf:
79 xdmf.write_mesh(domain)
80 xdmf.write_function(u, t)
81
82 u0.x.array[:] = u.x.array[:]