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[:]