Coverage for pySDC/implementations/problem_classes/Burgers.py: 95%

1import numpy as np 

2 

3from pySDC.implementations.datatype_classes.mesh import mesh, imex_mesh 

4from pySDC.implementations.problem_classes.generic_spectral import GenericSpectralLinear 

5 

6 

7class Burgers1D(GenericSpectralLinear): 

8 """ 

9 See https://en.wikipedia.org/wiki/Burgers'_equation for the equation that is solved. 

10 Discretization is done with a Chebychev method, which requires a first order derivative formulation. 

11 Feel free to do a more efficient implementation using an ultraspherical method to avoid the first order business. 

12 

13 Parameters: 

14 N (int): Spatial resolution 

15 epsilon (float): viscosity 

16 BCl (float): Value at left boundary 

17 BCr (float): Value at right boundary 

18 f (int): Frequency of the initial conditions 

19 mode (str): 'T2U' or 'T2T'. Use 'T2U' to get sparse differentiation matrices 

20 """ 

21 

22 dtype_u = mesh 

23 dtype_f = imex_mesh 

24 

25 def __init__(self, N=64, epsilon=0.1, BCl=1, BCr=-1, f=0, mode='T2U', **kwargs): 

26 """ 

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

28 

29 Args: 

30 N (int): Spatial resolution 

31 epsilon (float): viscosity 

32 BCl (float): Value at left boundary 

33 BCr (float): Value at right boundary 

34 f (int): Frequency of the initial conditions 

35 mode (str): 'T2U' or 'T2T'. Use 'T2U' to get sparse differentiation matrices 

36 """ 

37 self._makeAttributeAndRegister('N', 'epsilon', 'BCl', 'BCr', 'f', 'mode', localVars=locals(), readOnly=True) 

38 

39 bases = [{'base': 'cheby', 'N': N}] 

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

41 

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

43 

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

45 

46 # prepare matrices 

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

48 I = self.get_Id() 

49 

50 T2U = self.get_basis_change_matrix(conv=mode) 

51 

52 self.Dx = Dx 

53 

54 # construct linear operator 

55 L_lhs = {'u': {'ux': -epsilon * (T2U @ Dx)}, 'ux': {'u': -T2U @ Dx, 'ux': T2U @ I}} 

56 self.setup_L(L_lhs) 

57 

58 # construct mass matrix 

59 M_lhs = {'u': {'u': T2U @ I}} 

60 self.setup_M(M_lhs) 

61 

62 # boundary conditions 

63 self.add_BC(component='u', equation='u', axis=0, x=1, v=BCr, kind='Dirichlet') 

64 self.add_BC(component='u', equation='ux', axis=0, x=-1, v=BCl, kind='Dirichlet') 

65 self.setup_BCs() 

66 

67 def u_exact(self, t=0, *args, **kwargs): 

68 me = self.u_init 

69 

70 # x = (self.x + 1) / 2 

71 # g = 4 * (1 + np.exp(-(4 * x + t)/self.epsilon/32)) 

72 # g_x = 4 * np.exp(-(4 * x + t)/self.epsilon/32) * (-4/self.epsilon/32) 

73 

74 # me[0] = 3./4. - 1./g 

75 # me[1] = 1/g**2 * g_x 

76 

77 # return me 

78 

79 if t == 0: 

80 me[self.index('u')][:] = ((self.BCr + self.BCl) / 2 + (self.BCr - self.BCl) / 2 * self.x) * np.cos( 

81 self.x * np.pi * self.f 

82 ) 

83 me[self.index('ux')][:] = (self.BCr - self.BCl) / 2 * np.cos(self.x * np.pi * self.f) + ( 

84 (self.BCr + self.BCl) / 2 + (self.BCr - self.BCl) / 2 * self.x 

85 ) * self.f * np.pi * -np.sin(self.x * np.pi * self.f) 

86 elif t == np.inf and self.f == 0 and self.BCl == -self.BCr: 

87 me[0] = (self.BCl * np.exp((self.BCr - self.BCl) / (2 * self.epsilon) * self.x) + self.BCr) / ( 

88 np.exp((self.BCr - self.BCl) / (2 * self.epsilon) * self.x) + 1 

89 ) 

90 else: 

91 raise NotImplementedError 

92 

93 return me 

94 

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

96 f = self.f_init 

97 iu, iux = self.index('u'), self.index('ux') 

98 

99 u_hat = self.transform(u) 

100 

101 Dx_u_hat = self.u_init_forward 

102 Dx_u_hat[iux] = (self.Dx @ u_hat[iux].flatten()).reshape(u_hat[iu].shape) 

103 

104 f.impl[iu] = self.epsilon * self.itransform(Dx_u_hat)[iux].real 

105 f.expl[iu] = -u[iu] * u[iux] 

106 return f 

107 

108 def get_fig(self): # pragma: no cover 

109 """ 

110 Get a figure suitable to plot the solution of this problem. 

111 

112 Returns 

113 ------- 

114 self.fig : matplotlib.pyplot.figure.Figure 

115 """ 

116 import matplotlib.pyplot as plt 

117 

118 plt.rcParams['figure.constrained_layout.use'] = True 

119 self.fig, axs = plt.subplots() 

120 return self.fig 

121 

122 def plot(self, u, t=None, fig=None, comp='u'): # pragma: no cover 

123 r""" 

124 Plot the solution. 

125 

126 Parameters 

127 ---------- 

128 u : dtype_u 

129 Solution to be plotted 

130 t : float 

131 Time to display at the top of the figure 

132 fig : matplotlib.pyplot.figure.Figure, optional 

133 Figure with the same structure as a figure generated by `self.get_fig`. If none is supplied, a new figure will be generated. 

134 

135 Returns 

136 ------- 

137 None 

138 """ 

139 fig = self.get_fig() if fig is None else fig 

140 ax = fig.axes[0] 

141 

142 ax.plot(self.x, u[self.index(comp)]) 

143 

144 if t is not None: 

145 fig.suptitle(f't = {t:.2e}') 

146 

147 ax.set_xlabel(r'$x$') 

148 ax.set_ylabel(r'$u$') 

149 

150 

151class Burgers2D(GenericSpectralLinear): 

152 """ 

153 See https://en.wikipedia.org/wiki/Burgers'_equation for the equation that is solved. 

154 This implementation is discretized with FFTs in x and Chebychev in z. 

155 

156 Parameters: 

157 nx (int): Spatial resolution in x direction 

158 nz (int): Spatial resolution in z direction 

159 epsilon (float): viscosity 

160 BCl (float): Value at left boundary 

161 BCr (float): Value at right boundary 

162 fux (int): Frequency of the initial conditions in x-direction 

163 fuz (int): Frequency of the initial conditions in z-direction 

164 mode (str): 'T2U' or 'T2T'. Use 'T2U' to get sparse differentiation matrices 

165 """ 

166 

167 dtype_u = mesh 

168 dtype_f = imex_mesh 

169 

170 def __init__(self, nx=64, nz=64, epsilon=0.1, fux=2, fuz=1, mode='T2U', **kwargs): 

171 """ 

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

173 

174 Args: 

175 nx (int): Spatial resolution in x direction 

176 nz (int): Spatial resolution in z direction 

177 epsilon (float): viscosity 

178 fux (int): Frequency of the initial conditions in x-direction 

179 fuz (int): Frequency of the initial conditions in z-direction 

180 mode (str): 'T2U' or 'T2T'. Use 'T2U' to get sparse differentiation matrices 

181 """ 

182 self._makeAttributeAndRegister('nx', 'nz', 'epsilon', 'fux', 'fuz', 'mode', localVars=locals(), readOnly=True) 

183 

184 bases = [ 

185 {'base': 'fft', 'N': nx}, 

186 {'base': 'cheby', 'N': nz}, 

187 ] 

188 components = ['u', 'v', 'ux', 'uz', 'vx', 'vz'] 

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

190 

191 self.Z, self.X = self.get_grid() 

192 

193 # prepare matrices 

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

195 Dz = self.get_differentiation_matrix(axes=(1,)) 

196 I = self.get_Id() 

197 

198 T2U = self.get_basis_change_matrix(axes=(1,), conv=mode) 

199 

200 self.Dx = Dx 

201 self.Dz = Dz 

202 

203 # construct linear operator 

204 L_lhs = { 

205 'u': {'ux': -epsilon * T2U @ Dx, 'uz': -epsilon * T2U @ Dz}, 

206 'v': {'vx': -epsilon * T2U @ Dx, 'vz': -epsilon * T2U @ Dz}, 

207 'ux': {'u': -T2U @ Dx, 'ux': T2U @ I}, 

208 'uz': {'u': -T2U @ Dz, 'uz': T2U @ I}, 

209 'vx': {'v': -T2U @ Dx, 'vx': T2U @ I}, 

210 'vz': {'v': -T2U @ Dz, 'vz': T2U @ I}, 

211 } 

212 self.setup_L(L_lhs) 

213 

214 # construct mass matrix 

215 M_lhs = { 

216 'u': {'u': T2U @ I}, 

217 'v': {'v': T2U @ I}, 

218 } 

219 self.setup_M(M_lhs) 

220 

221 # boundary conditions 

222 self.BCtop = 1 

223 self.BCbottom = -self.BCtop 

224 self.BCtopu = 0 

225 self.add_BC(component='v', equation='v', axis=1, v=self.BCtop, x=1, kind='Dirichlet') 

226 self.add_BC(component='v', equation='vz', axis=1, v=self.BCbottom, x=-1, kind='Dirichlet') 

227 self.add_BC(component='u', equation='uz', axis=1, v=self.BCtopu, x=1, kind='Dirichlet') 

228 self.add_BC(component='u', equation='u', axis=1, v=self.BCtopu, x=-1, kind='Dirichlet') 

229 self.setup_BCs() 

230 

231 def u_exact(self, t=0, *args, noise_level=0, **kwargs): 

232 me = self.u_init 

233 

234 iu, iv, iux, iuz, ivx, ivz = self.index(self.components) 

235 if t == 0: 

236 me[iu] = self.xp.cos(self.X * self.fux) * self.xp.sin(self.Z * np.pi * self.fuz) + self.BCtopu 

237 me[iux] = -self.xp.sin(self.X * self.fux) * self.fux * self.xp.sin(self.Z * np.pi * self.fuz) 

238 me[iuz] = self.xp.cos(self.X * self.fux) * self.xp.cos(self.Z * np.pi * self.fuz) * np.pi * self.fuz 

239 

240 me[iv] = (self.BCtop + self.BCbottom) / 2 + (self.BCtop - self.BCbottom) / 2 * self.Z 

241 me[ivz][:] = (self.BCtop - self.BCbottom) / 2 

242 

243 # add noise 

244 rng = self.xp.random.default_rng(seed=99) 

245 me[iv].real += rng.normal(size=me[iv].shape) * (self.Z - 1) * (self.Z + 1) * noise_level 

246 

247 else: 

248 raise NotImplementedError 

249 

250 return me 

251 

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

253 f = self.f_init 

254 iu, iv, iux, iuz, ivx, ivz = self.index(self.components) 

255 

256 u_hat = self.transform(u) 

257 f_hat = self.u_init_forward 

258 f_hat[iu] = self.epsilon * ((self.Dx @ u_hat[iux].flatten() + self.Dz @ u_hat[iuz].flatten())).reshape( 

259 u_hat[iux].shape 

260 ) 

261 f_hat[iv] = self.epsilon * ((self.Dx @ u_hat[ivx].flatten() + self.Dz @ u_hat[ivz].flatten())).reshape( 

262 u_hat[iux].shape 

263 ) 

264 f.impl[...] = self.itransform(f_hat).real 

265 

266 f.expl[iu] = -(u[iu] * u[iux] + u[iv] * u[iuz]) 

267 f.expl[iv] = -(u[iu] * u[ivx] + u[iv] * u[ivz]) 

268 return f 

269 

270 def compute_vorticity(self, u): 

271 me = self.u_init_forward 

272 

273 u_hat = self.transform(u) 

274 iu, iv = self.index(['u', 'v']) 

275 

276 me[iu] = (self.Dx @ u_hat[iv].flatten() + self.Dz @ u_hat[iu].flatten()).reshape(u[iu].shape) 

277 return self.itransform(me)[iu].real 

278 

279 def get_fig(self): # pragma: no cover 

280 """ 

281 Get a figure suitable to plot the solution of this problem 

282 

283 Returns 

284 ------- 

285 self.fig : matplotlib.pyplot.figure.Figure 

286 """ 

287 import matplotlib.pyplot as plt 

288 from mpl_toolkits.axes_grid1 import make_axes_locatable 

289 

290 plt.rcParams['figure.constrained_layout.use'] = True 

291 self.fig, axs = plt.subplots(3, 1, sharex=True, sharey=True, figsize=((8, 7))) 

292 self.cax = [] 

293 divider = make_axes_locatable(axs[0]) 

294 self.cax += [divider.append_axes('right', size='3%', pad=0.03)] 

295 divider2 = make_axes_locatable(axs[1]) 

296 self.cax += [divider2.append_axes('right', size='3%', pad=0.03)] 

297 divider3 = make_axes_locatable(axs[2]) 

298 self.cax += [divider3.append_axes('right', size='3%', pad=0.03)] 

299 return self.fig 

300 

301 def plot(self, u, t=None, fig=None, vmin=None, vmax=None): # pragma: no cover 

302 r""" 

303 Plot the solution. Please supply a figure with the same structure as returned by ``self.get_fig``. 

304 

305 Parameters 

306 ---------- 

307 u : dtype_u 

308 Solution to be plotted 

309 t : float 

310 Time to display at the top of the figure 

311 fig : matplotlib.pyplot.figure.Figure 

312 Figure with the correct structure 

313 

314 Returns 

315 ------- 

316 None 

317 """ 

318 fig = self.get_fig() if fig is None else fig 

319 axs = fig.axes 

320 

321 iu, iv = self.index(['u', 'v']) 

322 

323 imU = axs[0].pcolormesh(self.X, self.Z, u[iu].real, vmin=vmin, vmax=vmax) 

324 imV = axs[1].pcolormesh(self.X, self.Z, u[iv].real, vmin=vmin, vmax=vmax) 

325 imVort = axs[2].pcolormesh(self.X, self.Z, self.compute_vorticity(u).real) 

326 

327 for i, label in zip([0, 1, 2], [r'$u$', '$v$', 'vorticity']): 

328 axs[i].set_aspect(1) 

329 axs[i].set_title(label) 

330 

331 if t is not None: 

332 fig.suptitle(f't = {t:.2e}') 

333 axs[-1].set_xlabel(r'$x$') 

334 axs[-1].set_ylabel(r'$z$') 

335 fig.colorbar(imU, self.cax[0]) 

336 fig.colorbar(imV, self.cax[1]) 

337 fig.colorbar(imVort, self.cax[2])