/
heateq.py
158 lines (137 loc) · 5.94 KB
/
heateq.py
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
import argparse
import numpy as np
import scipy.sparse as sp
from ngsolve import (H1, L2, InnerProduct, Preconditioner, ds, dx, grad,
ngsglobals)
from source.linalg import PCG
from source.linop import (AsLinearOperator, BlockDiagLinOp, CompositeLinOp,
KronLinOp)
from source.ngsolve_helper import BilForm, KronBF, KronFES, KronLF
from source.problem import problem_helper
from source.wavelets import WaveletTransformOp
ngsglobals.msg_level = 0
class HeatEquation:
""" Implementation of Andreev's method for tensor-product trial spaces. """
def __init__(self,
J_space=2,
J_time=None,
problem='square',
precond='multigrid',
alpha=0.3,
order=1):
if J_time is None:
J_time = J_space
mesh_space, bc, mesh_time, data, fn = problem_helper(problem,
J_space=J_space,
J_time=J_time)
# Building fespaces X^\delta and Y^\delta
X = KronFES(H1(mesh_time, order=order),
H1(mesh_space, order=order, dirichlet=bc))
Y = KronFES(L2(mesh_time, order=order), X.space)
self.N = len(X.time.fd)
self.M = len(X.space.fd)
# Building the ngsolve spacetime-bilforms.
dt = dx
A_bf = KronBF(Y, Y, lambda u, v: u * v * dt,
lambda u, v: InnerProduct(grad(u), grad(v)) * dx)
B1_bf = KronBF(X, Y, lambda u, v: grad(u) * v * dt,
lambda u, v: u * v * dx)
B2_bf = KronBF(X, Y, lambda u, v: u * v * dt,
lambda u, v: InnerProduct(grad(u), grad(v)) * dx)
G_bf = KronBF(X, X, lambda u, v: u * v * ds('start'),
lambda u, v: u * v * dx)
# Create the kron linops.
self.B = B1_bf.assemble() + B2_bf.assemble()
self.BT = B1_bf.transpose() + B2_bf.transpose()
self.G = G_bf.assemble()
# Preconditioner on Y.
Kinv_time_pc = Preconditioner(A_bf.time.bf, 'direct')
Kinv_space_pc = Preconditioner(A_bf.space.bf, precond)
A_bf.assemble()
Kinv_time = AsLinearOperator(Kinv_time_pc.mat, Y.time.fd)
Kinv_space = AsLinearOperator(Kinv_space_pc.mat, Y.space.fd)
self.K = KronLinOp(Kinv_time, Kinv_space)
# --- Wavelet transform ---
W_t = WaveletTransformOp(J_time)
self.W = KronLinOp(W_t, sp.eye(len(X.space.fd)))
self.WT = KronLinOp(W_t.T, sp.eye(len(X.space.fd)))
# --- Preconditioner on X ---
self.C_j = []
self.alpha = alpha
for j in range(J_time + 1):
bf = BilForm(X.space,
bilform_lambda=lambda u, v:
(2**j * u * v + self.alpha * grad(u) * grad(v)) * dx)
C = Preconditioner(bf.bf, precond)
bf.bf.Assemble()
self.C_j.append(AsLinearOperator(C.mat, X.space.fd))
self.CAC_j = [
CompositeLinOp([self.C_j[j], A_bf.space.mat, self.C_j[j]])
for j in range(J_time + 1)
]
self.P = BlockDiagLinOp([self.CAC_j[j] for j in W_t.levels])
# Schur-complement operator.
self.S = sp.linalg.LinearOperator(
self.G.shape,
matvec=lambda v: self.BT @ self.K @ self.B @ v + self.G @ v)
self.WT_S_W = self.WT @ self.S @ self.W
# Calculate rhs.
self.g_vec = np.zeros(self.K.shape[0])
for g in data['g']:
g_lf = KronLF(Y, lambda v: g[0] * v * dt, lambda v: g[1] * v * dx)
g_lf.assemble()
self.g_vec += g_lf.vec
u0_lf = KronLF(X, lambda v: v * ds('start'),
lambda v: data['u0'] * v * dx)
u0_lf.assemble()
self.f = self.BT @ self.K @ self.g_vec + u0_lf.vec
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Solve heatequation using ngsolve.')
parser.add_argument('--problem',
default='square',
help='problem type (square, ns)')
parser.add_argument('--J_time',
type=int,
default=5,
help='number of time refines')
parser.add_argument('--J_space',
type=int,
default=6,
help='number of space refines')
parser.add_argument(
'--precond',
default="multigrid",
help='type of ngsolve preconditioner, e.g. direct or multigrid.')
parser.add_argument('--alpha',
type=float,
default=0.3,
help='Alpha value used in the preconditioner for X.')
args = parser.parse_args()
order = 1 # Higher order requires a different wavelet transform.
print('Arguments: {}'.format(args))
print(
'\n\nCreating HeatEquation with {} time refines and {} space refines.'.
format(args.J_time, args.J_space))
heat_eq = HeatEquation(J_time=args.J_time,
J_space=args.J_space,
problem=args.problem,
precond=args.precond,
alpha=args.alpha)
print('Size of time mesh: {} dofs. Size of space mesh: {} dofs'.format(
heat_eq.N, heat_eq.M))
def cb(w, residual, k):
print('.', end='', flush=True)
print("Solving: ", end='')
w, iters = PCG(heat_eq.WT_S_W,
heat_eq.P,
heat_eq.WT @ heat_eq.f,
callback=cb)
u = heat_eq.W @ w
res = heat_eq.f - heat_eq.S @ heat_eq.f
error_alg = res @ (heat_eq.P @ res)
gminBu = heat_eq.g_vec - heat_eq.B @ u
error_Yprime = gminBu @ (heat_eq.K @ gminBu)
print("Done in {} PCG steps. "
"X-norm algebraic error: {}. "
"Error in Yprime: {}\n".format(iters, error_alg, error_Yprime))