1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
|
# Solves Heat equation on a periodic domain, using raw VecScatter
import sys
import petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
from mpi4py import MPI
import numpy
class Heat:
def __init__(self, comm, N):
self.comm = comm
self.N = N # global problem size
self.h = 1 / N # grid spacing on unit interval
self.n = N // comm.size + int(
comm.rank < (N % comm.size)
) # owned part of global problem
self.start = comm.exscan(self.n)
if comm.rank == 0:
self.start = 0
gindices = (
numpy.arange(self.start - 1, self.start + self.n + 1, dtype=PETSc.IntType)
% N
) # periodic
self.mat = PETSc.Mat().create(comm=comm)
size = (self.n, self.N) # local and global sizes
self.mat.setSizes((size, size))
self.mat.setFromOptions()
self.mat.setPreallocationNNZ(
(3, 1)
) # Conservative preallocation for 3 "local" columns and one non-local
# Allow matrix insertion using local indices [0:n+2]
lgmap = PETSc.LGMap().create(list(gindices), comm=comm)
self.mat.setLGMap(lgmap, lgmap)
# Global and local vectors
self.gvec = self.mat.createVecRight()
self.lvec = PETSc.Vec().create(comm=PETSc.COMM_SELF)
self.lvec.setSizes(self.n + 2)
self.lvec.setUp()
# Configure scatter from global to local
isg = PETSc.IS().createGeneral(list(gindices), comm=comm)
self.g2l = PETSc.Scatter().create(self.gvec, isg, self.lvec, None)
self.tozero, self.zvec = PETSc.Scatter.toZero(self.gvec)
self.history = []
if False: # Print some diagnostics
print(
'[%d] local size %d, global size %d, starting offset %d'
% (comm.rank, self.n, self.N, self.start)
)
self.gvec.setArray(numpy.arange(self.start, self.start + self.n))
self.gvec.view()
self.g2l.scatter(self.gvec, self.lvec, PETSc.InsertMode.INSERT)
for rank in range(comm.size):
if rank == comm.rank:
print('Contents of local Vec on rank %d' % rank)
self.lvec.view()
comm.barrier()
def evalSolution(self, t, x):
if t != 0.0:
raise ValueError('Only for t=0')
coord = numpy.arange(self.start, self.start + self.n) / self.N
x.setArray((numpy.abs(coord - 0.5) < 0.1) * 1.0)
def evalFunction(self, ts, t, x, xdot, f):
self.g2l.scatter(x, self.lvec, PETSc.InsertMode.INSERT) # lvec is a work vector
h = self.h
with self.lvec as u, xdot as udot:
f.setArray(
udot * h + 2 * u[1:-1] / h - u[:-2] / h - u[2:] / h
) # Scale equation by volume element
def evalJacobian(self, ts, t, x, xdot, a, A, B):
h = self.h
for i in range(self.n):
lidx = i + 1
B.setValuesLocal(
[lidx], [lidx - 1, lidx, lidx + 1], [-1 / h, a * h + 2 / h, -1 / h]
)
B.assemble()
if A != B:
A.assemble() # If operator is different from preconditioning matrix
def monitor(self, ts, i, t, x):
if self.history:
lasti, lastt, lastx = self.history[-1]
if i < lasti + 4 or t < lastt + 1e-4:
return
self.tozero.scatter(x, self.zvec, PETSc.InsertMode.INSERT)
xx = self.zvec[:].tolist()
self.history.append((i, t, xx))
def plotHistory(self):
try:
from matplotlib import pylab, rcParams
except ImportError:
return
rcParams.update({'text.usetex': True, 'figure.figsize': (10, 6)})
# rc('figure', figsize=(600,400))
pylab.title('Heat: TS \\texttt{%s}' % ts.getType())
x = numpy.arange(self.N) / self.N
for i, t, u in self.history:
pylab.plot(x, u, label='step=%d t=%8.2g' % (i, t))
pylab.xlabel('$x$')
pylab.ylabel('$u$')
pylab.legend(loc='upper right')
pylab.savefig('heat-history.png')
# pylab.show()
OptDB = PETSc.Options()
ode = Heat(MPI.COMM_WORLD, OptDB.getInt('n', 100))
x = ode.gvec.duplicate()
f = ode.gvec.duplicate()
ts = PETSc.TS().create(comm=ode.comm)
ts.setType(ts.Type.ROSW) # Rosenbrock-W. ARKIMEX is a nonlinearly implicit alternative.
ts.setIFunction(ode.evalFunction, ode.gvec)
ts.setIJacobian(ode.evalJacobian, ode.mat)
ts.setMonitor(ode.monitor)
ts.setTime(0.0)
ts.setTimeStep(ode.h**2)
ts.setMaxTime(1)
ts.setMaxSteps(100)
ts.setExactFinalTime(PETSc.TS.ExactFinalTime.INTERPOLATE)
ts.setMaxSNESFailures(
-1
) # allow an unlimited number of failures (step will be rejected and retried)
snes = ts.getSNES() # Nonlinear solver
snes.setTolerances(
max_it=10
) # Stop nonlinear solve after 10 iterations (TS will retry with shorter step)
ksp = snes.getKSP() # Linear solver
ksp.setType(ksp.Type.CG) # Conjugate gradients
pc = ksp.getPC() # Preconditioner
if False: # Configure algebraic multigrid, could use run-time options instead
pc.setType(
pc.Type.GAMG
) # PETSc's native AMG implementation, mostly based on smoothed aggregation
OptDB['mg_coarse_pc_type'] = 'svd' # more specific multigrid options
OptDB['mg_levels_pc_type'] = 'sor'
ts.setFromOptions() # Apply run-time options, e.g. -ts_adapt_monitor -ts_type arkimex -snes_converged_reason
ode.evalSolution(0.0, x)
ts.solve(x)
if ode.comm.rank == 0:
print(
'steps %d (%d rejected, %d SNES fails), nonlinear its %d, linear its %d'
% (
ts.getStepNumber(),
ts.getStepRejections(),
ts.getSNESFailures(),
ts.getSNESIterations(),
ts.getKSPIterations(),
)
)
if OptDB.getBool('plot_history', True) and ode.comm.rank == 0:
ode.plotHistory()
|
