Coverage for pySDC/implementations/problem_classes/HeatEquation_Chebychev.py: 100%

1import numpy as np 

2from scipy import sparse as sp 

3 

4from pySDC.core.problem import Problem 

5from pySDC.implementations.datatype_classes.mesh import mesh 

6from pySDC.implementations.problem_classes.generic_spectral import GenericSpectralLinear 

7 

8 

9class Heat1DChebychev(GenericSpectralLinear): 

10 """ 

11 1D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1) using a Chebychev spectral method. 

12 """ 

13 

14 dtype_u = mesh 

15 dtype_f = mesh 

16 

17 def __init__(self, nvars=128, a=0, b=0, f=1, nu=1.0, mode='T2U', **kwargs): 

18 """ 

19 Constructor. `kwargs` are forwarded to parent class constructor. 

20 

21 Args: 

22 nvars (int): Resolution 

23 a (float): Left BC value 

24 b (float): Right BC value 

25 f (int): Frequency of the solution 

26 nu (float): Diffusion parameter 

27 mode ('T2T' or 'T2U'): Mode for Chebychev method. 

28 

29 """ 

30 self._makeAttributeAndRegister('nvars', 'a', 'b', 'f', 'nu', 'mode', localVars=locals(), readOnly=True) 

31 

32 bases = [{'base': 'chebychev', 'N': nvars}] 

33 components = ['u', 'ux'] 

34 

35 super().__init__(bases, components, real_spectral_coefficients=True, **kwargs) 

36 

37 self.x = self.get_grid()[0] 

38 

39 I = self.get_Id() 

40 Dx = self.get_differentiation_matrix(axes=(0,)) 

41 self.Dx = Dx 

42 

43 self.T2U = self.get_basis_change_matrix(conv=mode) 

44 

45 L_lhs = { 

46 'ux': {'u': -self.T2U @ Dx, 'ux': self.T2U @ I}, 

47 'u': {'ux': -nu * (self.T2U @ Dx)}, 

48 } 

49 self.setup_L(L_lhs) 

50 

51 M_lhs = {'u': {'u': self.T2U @ I}} 

52 self.setup_M(M_lhs) 

53 

54 self.add_BC(component='u', equation='u', axis=0, x=-1, v=a, kind="Dirichlet") 

55 self.add_BC(component='u', equation='ux', axis=0, x=1, v=b, kind="Dirichlet") 

56 self.setup_BCs() 

57 

58 def eval_f(self, u, *args, **kwargs): 

59 f = self.f_init 

60 iu, iux = self.index(self.components) 

61 

62 if self.spectral_space: 

63 u_hat = u.copy() 

64 else: 

65 u_hat = self.transform(u) 

66 

67 u_hat[iu] = (self.nu * self.Dx @ u_hat[iux].flatten()).reshape(u_hat[iu].shape) 

68 

69 if self.spectral_space: 

70 me = u_hat 

71 else: 

72 me = self.itransform(u_hat).real 

73 

74 f[iu] = me[iu] 

75 return f 

76 

77 def u_exact(self, t, noise=0, seed=666): 

78 """ 

79 Get exact solution at time `t` 

80 

81 Args: 

82 t (float): When you want the exact solution 

83 noise (float): Add noise of this level 

84 seed (int): Random seed for the noise 

85 

86 Returns: 

87 Heat1DChebychev.dtype_u: Exact solution 

88 """ 

89 xp = self.xp 

90 iu, iux = self.index(self.components) 

91 u = self.spectral.u_init 

92 

93 u[iu] = ( 

94 xp.sin(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t) 

95 + (self.b - self.a) / 2 * self.x 

96 + (self.b + self.a) / 2 

97 ) 

98 u[iux] = ( 

99 self.f * np.pi * xp.cos(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t) 

100 + (self.b - self.a) / 2 

101 ) 

102 

103 if noise > 0: 

104 assert t == 0 

105 _noise = self.u_init 

106 rng = self.xp.random.default_rng(seed=seed) 

107 _noise[iu] = rng.normal(size=u[iu].shape) 

108 noise_hat = self.transform(_noise) 

109 low_pass = self.get_filter_matrix(axis=0, kmax=self.nvars - 2) 

110 noise_hat[iu] = (low_pass @ noise_hat[iu].flatten()).reshape(noise_hat[iu].shape) 

111 _noise[:] = self.itransform(noise_hat) 

112 u += _noise * noise * (self.x - 1) * (self.x + 1) 

113 

114 self.check_BCs(u) 

115 

116 if self.spectral_space: 

117 u_hat = self.u_init 

118 u_hat[...] = self.transform(u) 

119 return u_hat 

120 else: 

121 return u 

122 

123 

124class Heat1DUltraspherical(GenericSpectralLinear): 

125 """ 

126 1D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1) using an ultraspherical spectral method. 

127 """ 

128 

129 dtype_u = mesh 

130 dtype_f = mesh 

131 

132 def __init__(self, nvars=128, a=0, b=0, f=1, nu=1.0, **kwargs): 

133 """ 

134 Constructor. `kwargs` are forwarded to parent class constructor. 

135 

136 Args: 

137 nvars (int): Resolution 

138 a (float): Left BC value 

139 b (float): Right BC value 

140 f (int): Frequency of the solution 

141 nu (float): Diffusion parameter 

142 """ 

143 self._makeAttributeAndRegister('nvars', 'a', 'b', 'f', 'nu', localVars=locals(), readOnly=True) 

144 

145 bases = [{'base': 'ultraspherical', 'N': nvars}] 

146 components = ['u'] 

147 

148 GenericSpectralLinear.__init__(self, bases, components, real_spectral_coefficients=True, **kwargs) 

149 

150 self.x = self.get_grid()[0] 

151 

152 I = self.get_Id() 

153 Dxx = self.get_differentiation_matrix(axes=(0,), p=2) 

154 

155 S2 = self.get_basis_change_matrix(p_in=2, p_out=0) 

156 U2 = self.get_basis_change_matrix(p_in=0, p_out=2) 

157 

158 self.Dxx = S2 @ Dxx 

159 

160 L_lhs = { 

161 'u': {'u': -nu * Dxx}, 

162 } 

163 self.setup_L(L_lhs) 

164 

165 M_lhs = {'u': {'u': U2 @ I}} 

166 self.setup_M(M_lhs) 

167 

168 self.add_BC(component='u', equation='u', axis=0, x=-1, v=a, kind="Dirichlet", line=-1) 

169 self.add_BC(component='u', equation='u', axis=0, x=1, v=b, kind="Dirichlet", line=-2) 

170 self.setup_BCs() 

171 

172 def eval_f(self, u, *args, **kwargs): 

173 f = self.f_init 

174 iu = self.index('u') 

175 

176 if self.spectral_space: 

177 u_hat = u.copy() 

178 else: 

179 u_hat = self.transform(u) 

180 

181 u_hat[iu] = (self.nu * (self.Dxx @ u_hat[iu].flatten())).reshape(u_hat[iu].shape) 

182 

183 if self.spectral_space: 

184 me = u_hat 

185 else: 

186 me = self.itransform(u_hat).real 

187 

188 f[iu][...] = me[iu] 

189 return f 

190 

191 def u_exact(self, t, noise=0, seed=666): 

192 """ 

193 Get exact solution at time `t` 

194 

195 Args: 

196 t (float): When you want the exact solution 

197 noise (float): Add noise of this level 

198 seed (int): Random seed for the noise 

199 

200 Returns: 

201 Heat1DUltraspherical.dtype_u: Exact solution 

202 """ 

203 xp = self.xp 

204 iu = self.index('u') 

205 u = self.spectral.u_init 

206 

207 u[iu] = ( 

208 xp.sin(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t) 

209 + (self.b - self.a) / 2 * self.x 

210 + (self.b + self.a) / 2 

211 ) 

212 

213 if noise > 0: 

214 assert t == 0 

215 _noise = self.u_init 

216 rng = self.xp.random.default_rng(seed=seed) 

217 _noise[iu] = rng.normal(size=u[iu].shape) 

218 noise_hat = self.transform(_noise) 

219 low_pass = self.get_filter_matrix(axis=0, kmax=self.nvars - 2) 

220 noise_hat[iu] = (low_pass @ noise_hat[iu].flatten()).reshape(noise_hat[iu].shape) 

221 _noise[:] = self.itransform(noise_hat) 

222 u += _noise * noise * (self.x - 1) * (self.x + 1) 

223 

224 self.check_BCs(u) 

225 

226 if self.spectral_space: 

227 return self.transform(u) 

228 else: 

229 return u 

230 

231 

232class Heat2DChebychev(GenericSpectralLinear): 

233 """ 

234 2D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1)x(-1,1) using spectral methods based on FFT and Chebychev. 

235 """ 

236 

237 dtype_u = mesh 

238 dtype_f = mesh 

239 

240 def __init__(self, nx=128, ny=128, base_x='fft', base_y='chebychev', a=0, b=0, c=0, fx=1, fy=1, nu=1.0, **kwargs): 

241 """ 

242 Constructor. `kwargs` are forwarded to parent class constructor. 

243 

244 Args: 

245 nx (int): Resolution in x-direction 

246 ny (int): Resolution in y-direction 

247 base_x (str): Spectral base in x-direction 

248 base_y (str): Spectral base in y-direction 

249 a (float): BC at y=0 and x=-1 

250 b (float): BC at y=0 and x=+1 

251 c (float): BC at y=1 and x = -1 

252 fx (int): Horizontal frequency of initial conditions 

253 fy (int): Vertical frequency of initial conditions 

254 nu (float): Diffusion parameter 

255 """ 

256 assert nx % 2 == 1 or base_x == 'fft' 

257 assert ny % 2 == 1 or base_y == 'fft' 

258 self._makeAttributeAndRegister( 

259 'nx', 'ny', 'base_x', 'base_y', 'a', 'b', 'c', 'fx', 'fy', 'nu', localVars=locals(), readOnly=True 

260 ) 

261 

262 bases = [{'base': base_x, 'N': nx}, {'base': base_y, 'N': ny}] 

263 components = ['u', 'ux', 'uy'] 

264 

265 super().__init__(bases, components, Dirichlet_recombination=False, spectral_space=False, **kwargs) 

266 

267 self.Y, self.X = self.get_grid() 

268 

269 I = self.get_Id() 

270 self.Dx = self.get_differentiation_matrix(axes=(0,)) 

271 self.Dy = self.get_differentiation_matrix(axes=(1,)) 

272 

273 L_lhs = { 

274 'ux': {'u': -self.Dx, 'ux': I}, 

275 'uy': {'u': -self.Dy, 'uy': I}, 

276 'u': {'ux': -nu * self.Dx, 'uy': -nu * self.Dy}, 

277 } 

278 self.setup_L(L_lhs) 

279 

280 M_lhs = {'u': {'u': I}} 

281 self.setup_M(M_lhs) 

282 

283 for base in [base_x, base_y]: 

284 assert base in ['chebychev', 'fft'] 

285 

286 alpha = (self.b - self.a) / 2.0 

287 beta = (self.c - self.b) / 2.0 

288 gamma = (self.c + self.a) / 2.0 

289 

290 if base_x == 'chebychev': 

291 y = self.Y[0, :] 

292 if self.base_y == 'fft': 

293 self.add_BC(component='u', equation='u', axis=0, x=-1, v=beta * y - alpha + gamma, kind='Dirichlet') 

294 self.add_BC(component='ux', equation='ux', axis=0, v=2 * alpha, kind='integral') 

295 else: 

296 assert a == b, f'Need periodic boundary conditions in x for {base_x} method!' 

297 if base_y == 'chebychev': 

298 x = self.X[:, 0] 

299 self.add_BC(component='u', equation='u', axis=1, x=-1, v=alpha * x - beta + gamma, kind='Dirichlet') 

300 self.add_BC(component='uy', equation='uy', axis=1, v=2 * beta, kind='integral') 

301 else: 

302 assert c == b, f'Need periodic boundary conditions in y for {base_y} method!' 

303 self.setup_BCs() 

304 

305 def eval_f(self, u, *args, **kwargs): 

306 f = self.f_init 

307 iu, iux, iuy = self.index(self.components) 

308 

309 me_hat = self.u_init_forward 

310 me_hat[:] = self.transform(u) 

311 me_hat[iu] = self.nu * (self.Dx @ me_hat[iux].flatten() + self.Dy @ me_hat[iuy].flatten()).reshape( 

312 me_hat[iu].shape 

313 ) 

314 me = self.itransform(me_hat) 

315 

316 f[self.index("u")] = me[iu].real 

317 return f 

318 

319 def u_exact(self, t): 

320 xp = self.xp 

321 iu, iux, iuy = self.index(self.components) 

322 u = self.u_init 

323 

324 fx = self.fx if self.base_x == 'fft' else np.pi * self.fx 

325 fy = self.fy if self.base_y == 'fft' else np.pi * self.fy 

326 

327 time_dep = xp.exp(-self.nu * (fx**2 + fy**2) * t) 

328 

329 alpha = (self.b - self.a) / 2.0 

330 beta = (self.c - self.b) / 2.0 

331 gamma = (self.c + self.a) / 2.0 

332 

333 u[iu] = xp.sin(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha * self.X + beta * self.Y + gamma 

334 u[iux] = fx * xp.cos(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha 

335 u[iuy] = fy * xp.sin(fx * self.X) * xp.cos(fy * self.Y) * time_dep + beta 

336 

337 return u 

338 

339 

340class Heat2DUltraspherical(GenericSpectralLinear): 

341 """ 

342 2D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1)x(-1,1) using spectral methods based on FFT and Gegenbauer. 

343 """ 

344 

345 dtype_u = mesh 

346 dtype_f = mesh 

347 

348 def __init__( 

349 self, nx=128, ny=128, base_x='fft', base_y='ultraspherical', a=0, b=0, c=0, fx=1, fy=1, nu=1.0, **kwargs 

350 ): 

351 """ 

352 Constructor. `kwargs` are forwarded to parent class constructor. 

353 

354 Args: 

355 nx (int): Resolution in x-direction 

356 ny (int): Resolution in y-direction 

357 base_x (str): Spectral base in x-direction 

358 base_y (str): Spectral base in y-direction 

359 a (float): BC at y=0 and x=-1 

360 b (float): BC at y=0 and x=+1 

361 c (float): BC at y=1 and x = -1 

362 fx (int): Horizontal frequency of initial conditions 

363 fy (int): Vertical frequency of initial conditions 

364 nu (float): Diffusion parameter 

365 """ 

366 self._makeAttributeAndRegister( 

367 'nx', 'ny', 'base_x', 'base_y', 'a', 'b', 'c', 'fx', 'fy', 'nu', localVars=locals(), readOnly=True 

368 ) 

369 

370 bases = [{'base': base_x, 'N': nx}, {'base': base_y, 'N': ny}] 

371 components = ['u'] 

372 

373 super().__init__(bases, components, Dirichlet_recombination=False, spectral_space=False, **kwargs) 

374 

375 self.Y, self.X = self.get_grid() 

376 

377 self.Dxx = self.get_differentiation_matrix(axes=(0,), p=2) 

378 self.Dyy = self.get_differentiation_matrix(axes=(1,), p=2) 

379 self.S2 = self.get_basis_change_matrix(p=2) 

380 I = self.get_Id() 

381 

382 L_lhs = { 

383 'u': {'u': -nu * self.Dxx - nu * self.Dyy}, 

384 } 

385 self.setup_L(L_lhs) 

386 

387 M_lhs = {'u': {'u': I}} 

388 self.setup_M(M_lhs) 

389 

390 for base in [base_x, base_y]: 

391 assert base in ['ultraspherical', 'fft'] 

392 

393 alpha = (self.b - self.a) / 2.0 

394 beta = (self.c - self.b) / 2.0 

395 gamma = (self.c + self.a) / 2.0 

396 

397 if base_x == 'ultraspherical': 

398 y = self.Y[0, :] 

399 if self.base_y == 'fft': 

400 self.add_BC(component='u', equation='u', axis=0, x=-1, v=beta * y - alpha + gamma, kind='Dirichlet') 

401 self.add_BC(component='u', equation='u', axis=0, v=beta * y + alpha + gamma, x=1, line=-2, kind='Dirichlet') 

402 else: 

403 assert a == b, f'Need periodic boundary conditions in x for {base_x} method!' 

404 if base_y == 'ultraspherical': 

405 x = self.X[:, 0] 

406 self.add_BC( 

407 component='u', equation='u', axis=1, x=-1, v=alpha * x - beta + gamma, kind='Dirichlet', line=-1 

408 ) 

409 self.add_BC( 

410 component='u', equation='u', axis=1, x=+1, v=alpha * x + beta + gamma, kind='Dirichlet', line=-2 

411 ) 

412 else: 

413 assert c == b, f'Need periodic boundary conditions in y for {base_y} method!' 

414 self.setup_BCs() 

415 

416 def eval_f(self, u, *args, **kwargs): 

417 f = self.f_init 

418 iu = self.index('u') 

419 

420 me_hat = self.u_init_forward 

421 me_hat[:] = self.transform(u) 

422 me_hat[iu] = self.nu * (self.S2 @ (self.Dxx + self.Dyy) @ me_hat[iu].flatten()).reshape(me_hat[iu].shape) 

423 me = self.itransform(me_hat) 

424 

425 f[iu] = me[iu].real 

426 return f 

427 

428 def u_exact(self, t): 

429 xp = self.xp 

430 iu = self.index('u') 

431 u = self.u_init 

432 

433 fx = self.fx if self.base_x == 'fft' else np.pi * self.fx 

434 fy = self.fy if self.base_y == 'fft' else np.pi * self.fy 

435 

436 time_dep = xp.exp(-self.nu * (fx**2 + fy**2) * t) 

437 

438 alpha = (self.b - self.a) / 2.0 

439 beta = (self.c - self.b) / 2.0 

440 gamma = (self.c + self.a) / 2.0 

441 

442 u[iu] = xp.sin(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha * self.X + beta * self.Y + gamma 

443 

444 return u