Coverage for pySDC/implementations/problem_classes/HeatEquation_Chebychev.py: 100%
222 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 13:41 +0000
1import numpy as np
2from scipy import sparse as sp
4from pySDC.core.problem import Problem
5from pySDC.implementations.datatype_classes.mesh import mesh
6from pySDC.implementations.problem_classes.generic_spectral import GenericSpectralLinear
9class Heat1DChebychev(GenericSpectralLinear):
10 """
11 1D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1) using a Chebychev spectral method.
12 """
14 dtype_u = mesh
15 dtype_f = mesh
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.
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.
29 """
30 self._makeAttributeAndRegister('nvars', 'a', 'b', 'f', 'nu', 'mode', localVars=locals(), readOnly=True)
32 bases = [{'base': 'chebychev', 'N': nvars}]
33 components = ['u', 'ux']
35 super().__init__(bases, components, real_spectral_coefficients=True, **kwargs)
37 self.x = self.get_grid()[0]
39 I = self.get_Id()
40 Dx = self.get_differentiation_matrix(axes=(0,))
41 self.Dx = Dx
43 self.T2U = self.get_basis_change_matrix(conv=mode)
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)
51 M_lhs = {'u': {'u': self.T2U @ I}}
52 self.setup_M(M_lhs)
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()
58 def eval_f(self, u, *args, **kwargs):
59 """
60 Evaluate the right hand side nu u_xx, computed as nu times the Chebychev derivative of the auxiliary component
61 `ux`. The algebraic `ux` equation gets zero.
63 Args:
64 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space
65 *args: Not used, the right hand side does not depend on time
66 **kwargs: Not used, the right hand side does not depend on time
68 Returns:
69 Heat1DChebychev.dtype_f: The right hand side, in the same space as `u`
70 """
71 f = self.f_init
72 iu, iux = self.index(self.components)
74 if self.spectral_space:
75 u_hat = u.copy()
76 else:
77 u_hat = self.transform(u)
79 u_hat[iu] = (self.nu * self.Dx @ u_hat[iux].flatten()).reshape(u_hat[iu].shape)
81 if self.spectral_space:
82 me = u_hat
83 else:
84 me = self.itransform(u_hat).real
86 f[iu] = me[iu]
87 return f
89 def u_exact(self, t, noise=0, seed=666):
90 """
91 Get exact solution at time `t`
93 Args:
94 t (float): When you want the exact solution
95 noise (float): Add noise of this level
96 seed (int): Random seed for the noise
98 Returns:
99 Heat1DChebychev.dtype_u: Exact solution
100 """
101 xp = self.xp
102 iu, iux = self.index(self.components)
103 u = self.spectral.u_init
105 u[iu] = (
106 xp.sin(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t)
107 + (self.b - self.a) / 2 * self.x
108 + (self.b + self.a) / 2
109 )
110 u[iux] = (
111 self.f * np.pi * xp.cos(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t)
112 + (self.b - self.a) / 2
113 )
115 if noise > 0:
116 assert t == 0
117 _noise = self.u_init
118 rng = self.xp.random.default_rng(seed=seed)
119 _noise[iu] = rng.normal(size=u[iu].shape)
120 noise_hat = self.transform(_noise)
121 low_pass = self.get_filter_matrix(axis=0, kmax=self.nvars - 2)
122 noise_hat[iu] = (low_pass @ noise_hat[iu].flatten()).reshape(noise_hat[iu].shape)
123 _noise[:] = self.itransform(noise_hat)
124 u += _noise * noise * (self.x - 1) * (self.x + 1)
126 self.check_BCs(u)
128 if self.spectral_space:
129 u_hat = self.u_init
130 u_hat[...] = self.transform(u)
131 return u_hat
132 else:
133 return u
136class Heat1DUltraspherical(GenericSpectralLinear):
137 """
138 1D Heat equation with Dirichlet Boundary conditions discretized on (-1, 1) using an ultraspherical spectral method.
139 """
141 dtype_u = mesh
142 dtype_f = mesh
144 def __init__(self, nvars=128, a=0, b=0, f=1, nu=1.0, **kwargs):
145 """
146 Constructor. `kwargs` are forwarded to parent class constructor.
148 Args:
149 nvars (int): Resolution
150 a (float): Left BC value
151 b (float): Right BC value
152 f (int): Frequency of the solution
153 nu (float): Diffusion parameter
154 """
155 self._makeAttributeAndRegister('nvars', 'a', 'b', 'f', 'nu', localVars=locals(), readOnly=True)
157 bases = [{'base': 'ultraspherical', 'N': nvars}]
158 components = ['u']
160 GenericSpectralLinear.__init__(self, bases, components, real_spectral_coefficients=True, **kwargs)
162 self.x = self.get_grid()[0]
164 I = self.get_Id()
165 Dxx = self.get_differentiation_matrix(axes=(0,), p=2)
167 S2 = self.get_basis_change_matrix(p_in=2, p_out=0)
168 U2 = self.get_basis_change_matrix(p_in=0, p_out=2)
170 self.Dxx = S2 @ Dxx
172 L_lhs = {
173 'u': {'u': -nu * Dxx},
174 }
175 self.setup_L(L_lhs)
177 M_lhs = {'u': {'u': U2 @ I}}
178 self.setup_M(M_lhs)
180 self.add_BC(component='u', equation='u', axis=0, x=-1, v=a, kind="Dirichlet", line=-1)
181 self.add_BC(component='u', equation='u', axis=0, x=1, v=b, kind="Dirichlet", line=-2)
182 self.setup_BCs()
184 def eval_f(self, u, *args, **kwargs):
185 """
186 Evaluate the right hand side nu u_xx with the ultraspherical second derivative matrix, converted back to the
187 Chebychev basis.
189 Args:
190 u (dtype_u): Solution, in spectral space if `spectral_space` is set, else in physical space
191 *args: Not used, the right hand side does not depend on time
192 **kwargs: Not used, the right hand side does not depend on time
194 Returns:
195 Heat1DUltraspherical.dtype_f: The right hand side, in the same space as `u`
196 """
197 f = self.f_init
198 iu = self.index('u')
200 if self.spectral_space:
201 u_hat = u.copy()
202 else:
203 u_hat = self.transform(u)
205 u_hat[iu] = (self.nu * (self.Dxx @ u_hat[iu].flatten())).reshape(u_hat[iu].shape)
207 if self.spectral_space:
208 me = u_hat
209 else:
210 me = self.itransform(u_hat).real
212 f[iu][...] = me[iu]
213 return f
215 def u_exact(self, t, noise=0, seed=666):
216 """
217 Get exact solution at time `t`
219 Args:
220 t (float): When you want the exact solution
221 noise (float): Add noise of this level
222 seed (int): Random seed for the noise
224 Returns:
225 Heat1DUltraspherical.dtype_u: Exact solution
226 """
227 xp = self.xp
228 iu = self.index('u')
229 u = self.spectral.u_init
231 u[iu] = (
232 xp.sin(self.f * np.pi * self.x) * xp.exp(-self.nu * (self.f * np.pi) ** 2 * t)
233 + (self.b - self.a) / 2 * self.x
234 + (self.b + self.a) / 2
235 )
237 if noise > 0:
238 assert t == 0
239 _noise = self.u_init
240 rng = self.xp.random.default_rng(seed=seed)
241 _noise[iu] = rng.normal(size=u[iu].shape)
242 noise_hat = self.transform(_noise)
243 low_pass = self.get_filter_matrix(axis=0, kmax=self.nvars - 2)
244 noise_hat[iu] = (low_pass @ noise_hat[iu].flatten()).reshape(noise_hat[iu].shape)
245 _noise[:] = self.itransform(noise_hat)
246 u += _noise * noise * (self.x - 1) * (self.x + 1)
248 self.check_BCs(u)
250 if self.spectral_space:
251 u_hat = self.u_init
252 u_hat[...] = self.transform(u)
253 return u_hat
254 else:
255 return u
258class Heat2DChebychev(GenericSpectralLinear):
259 """
260 2D heat equation, periodic (FFT) or Dirichlet (Chebychev) in each direction, in first-order formulation.
262 Discretized on (-1, 1)x(-1,1) using spectral methods based on FFT and Chebychev.
263 """
265 dtype_u = mesh
266 dtype_f = mesh
268 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):
269 """
270 Constructor. `kwargs` are forwarded to parent class constructor.
272 Args:
273 nx (int): Resolution in x-direction
274 ny (int): Resolution in y-direction
275 base_x (str): Spectral base in x-direction
276 base_y (str): Spectral base in y-direction
277 a (float): BC at y=0 and x=-1
278 b (float): BC at y=0 and x=+1
279 c (float): BC at y=1 and x = -1
280 fx (int): Horizontal frequency of initial conditions
281 fy (int): Vertical frequency of initial conditions
282 nu (float): Diffusion parameter
283 """
284 assert nx % 2 == 1 or base_x == 'fft'
285 assert ny % 2 == 1 or base_y == 'fft'
286 self._makeAttributeAndRegister(
287 'nx', 'ny', 'base_x', 'base_y', 'a', 'b', 'c', 'fx', 'fy', 'nu', localVars=locals(), readOnly=True
288 )
290 bases = [{'base': base_x, 'N': nx}, {'base': base_y, 'N': ny}]
291 components = ['u', 'ux', 'uy']
293 super().__init__(bases, components, Dirichlet_recombination=False, spectral_space=False, **kwargs)
295 self.X, self.Y = self.get_grid()
297 I = self.get_Id()
298 self.Dx = self.get_differentiation_matrix(axes=(0,))
299 self.Dy = self.get_differentiation_matrix(axes=(1,))
301 L_lhs = {
302 'ux': {'u': -self.Dx, 'ux': I},
303 'uy': {'u': -self.Dy, 'uy': I},
304 'u': {'ux': -nu * self.Dx, 'uy': -nu * self.Dy},
305 }
306 self.setup_L(L_lhs)
308 M_lhs = {'u': {'u': I}}
309 self.setup_M(M_lhs)
311 for base in [base_x, base_y]:
312 assert base in ['chebychev', 'fft']
314 alpha = (self.b - self.a) / 2.0
315 beta = (self.c - self.b) / 2.0
316 gamma = (self.c + self.a) / 2.0
318 if base_x == 'chebychev':
319 y = self.Y[0, :]
320 if self.base_y == 'fft':
321 self.add_BC(component='u', equation='u', axis=0, x=-1, v=beta * y - alpha + gamma, kind='Dirichlet')
322 self.add_BC(component='ux', equation='ux', axis=0, v=2 * alpha, kind='integral')
323 else:
324 assert a == b, f'Need periodic boundary conditions in x for {base_x} method!'
325 if base_y == 'chebychev':
326 x = self.X[:, 0]
327 self.add_BC(component='u', equation='u', axis=1, x=-1, v=alpha * x - beta + gamma, kind='Dirichlet')
328 self.add_BC(component='uy', equation='uy', axis=1, v=2 * beta, kind='integral')
329 else:
330 assert c == b, f'Need periodic boundary conditions in y for {base_y} method!'
331 self.setup_BCs()
333 def eval_f(self, u, *args, **kwargs):
334 """
335 Evaluate the right hand side nu (u_xx + u_yy), computed by spectral differentiation of the auxiliary components
336 `ux` and `uy`. The algebraic `ux` and `uy` equations get zero.
338 Args:
339 u (dtype_u): Solution in physical space
340 *args: Not used, the right hand side does not depend on time
341 **kwargs: Not used, the right hand side does not depend on time
343 Returns:
344 Heat2DChebychev.dtype_f: The right hand side in physical space
345 """
346 f = self.f_init
347 iu, iux, iuy = self.index(self.components)
349 me_hat = self.u_init_forward
350 me_hat[:] = self.transform(u)
351 me_hat[iu] = self.nu * (self.Dx @ me_hat[iux].flatten() + self.Dy @ me_hat[iuy].flatten()).reshape(
352 me_hat[iu].shape
353 )
354 me = self.itransform(me_hat)
356 f[self.index("u")] = me[iu].real
357 return f
359 def u_exact(self, t):
360 """
361 Get the exact solution at time `t`: sin(fx x) sin(fy y) exp(-nu (fx^2 + fy^2) t) plus the linear function in x
362 and y that satisfies the boundary conditions, where the frequencies are multiplied by pi in Chebychev
363 directions. The components `ux` and `uy` are set to the derivatives.
365 Args:
366 t (float): When you want the exact solution
368 Returns:
369 Heat2DChebychev.dtype_u: Exact solution in physical space
370 """
371 xp = self.xp
372 iu, iux, iuy = self.index(self.components)
373 u = self.u_init
375 fx = self.fx if self.base_x == 'fft' else np.pi * self.fx
376 fy = self.fy if self.base_y == 'fft' else np.pi * self.fy
378 time_dep = xp.exp(-self.nu * (fx**2 + fy**2) * t)
380 alpha = (self.b - self.a) / 2.0
381 beta = (self.c - self.b) / 2.0
382 gamma = (self.c + self.a) / 2.0
384 u[iu] = xp.sin(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha * self.X + beta * self.Y + gamma
385 u[iux] = fx * xp.cos(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha
386 u[iuy] = fy * xp.sin(fx * self.X) * xp.cos(fy * self.Y) * time_dep + beta
388 return u
391class Heat2DUltraspherical(GenericSpectralLinear):
392 """
393 2D heat equation, periodic (FFT) or Dirichlet (ultraspherical) in each direction, in second-order formulation.
395 Discretized on (-1, 1)x(-1,1) using spectral methods based on FFT and Gegenbauer.
396 """
398 dtype_u = mesh
399 dtype_f = mesh
401 def __init__(
402 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
403 ):
404 """
405 Constructor. `kwargs` are forwarded to parent class constructor.
407 Args:
408 nx (int): Resolution in x-direction
409 ny (int): Resolution in y-direction
410 base_x (str): Spectral base in x-direction
411 base_y (str): Spectral base in y-direction
412 a (float): BC at y=0 and x=-1
413 b (float): BC at y=0 and x=+1
414 c (float): BC at y=1 and x = -1
415 fx (int): Horizontal frequency of initial conditions
416 fy (int): Vertical frequency of initial conditions
417 nu (float): Diffusion parameter
418 """
419 self._makeAttributeAndRegister(
420 'nx', 'ny', 'base_x', 'base_y', 'a', 'b', 'c', 'fx', 'fy', 'nu', localVars=locals(), readOnly=True
421 )
423 bases = [{'base': base_x, 'N': nx}, {'base': base_y, 'N': ny}]
424 components = ['u']
426 super().__init__(bases, components, Dirichlet_recombination=False, spectral_space=False, **kwargs)
428 self.X, self.Y = self.get_grid()
430 self.Dxx = self.get_differentiation_matrix(axes=(0,), p=2)
431 self.Dyy = self.get_differentiation_matrix(axes=(1,), p=2)
432 self.S2 = self.get_basis_change_matrix(p=2)
433 I = self.get_Id()
435 L_lhs = {
436 'u': {'u': -nu * self.Dxx - nu * self.Dyy},
437 }
438 self.setup_L(L_lhs)
440 M_lhs = {'u': {'u': I}}
441 self.setup_M(M_lhs)
443 for base in [base_x, base_y]:
444 assert base in ['ultraspherical', 'fft']
446 alpha = (self.b - self.a) / 2.0
447 beta = (self.c - self.b) / 2.0
448 gamma = (self.c + self.a) / 2.0
450 if base_x == 'ultraspherical':
451 y = self.Y[0, :]
452 if self.base_y == 'fft':
453 self.add_BC(component='u', equation='u', axis=0, x=-1, v=beta * y - alpha + gamma, kind='Dirichlet')
454 self.add_BC(component='u', equation='u', axis=0, v=beta * y + alpha + gamma, x=1, line=-2, kind='Dirichlet')
455 else:
456 assert a == b, f'Need periodic boundary conditions in x for {base_x} method!'
457 if base_y == 'ultraspherical':
458 x = self.X[:, 0]
459 self.add_BC(
460 component='u', equation='u', axis=1, x=-1, v=alpha * x - beta + gamma, kind='Dirichlet', line=-1
461 )
462 self.add_BC(
463 component='u', equation='u', axis=1, x=+1, v=alpha * x + beta + gamma, kind='Dirichlet', line=-2
464 )
465 else:
466 assert c == b, f'Need periodic boundary conditions in y for {base_y} method!'
467 self.setup_BCs()
469 def eval_f(self, u, *args, **kwargs):
470 """
471 Evaluate the right hand side nu (u_xx + u_yy) with the ultraspherical second derivative matrices, converted back
472 to the Chebychev basis.
474 Args:
475 u (dtype_u): Solution in physical space
476 *args: Not used, the right hand side does not depend on time
477 **kwargs: Not used, the right hand side does not depend on time
479 Returns:
480 Heat2DUltraspherical.dtype_f: The right hand side in physical space
481 """
482 f = self.f_init
483 iu = self.index('u')
485 me_hat = self.u_init_forward
486 me_hat[:] = self.transform(u)
487 me_hat[iu] = self.nu * (self.S2 @ (self.Dxx + self.Dyy) @ me_hat[iu].flatten()).reshape(me_hat[iu].shape)
488 me = self.itransform(me_hat)
490 f[iu] = me[iu].real
491 return f
493 def u_exact(self, t):
494 """
495 Get the exact solution at time `t`: sin(fx x) sin(fy y) exp(-nu (fx^2 + fy^2) t) plus the linear function in x
496 and y that satisfies the boundary conditions, where the frequencies are multiplied by pi in ultraspherical
497 directions.
499 Args:
500 t (float): When you want the exact solution
502 Returns:
503 Heat2DUltraspherical.dtype_u: Exact solution in physical space
504 """
505 xp = self.xp
506 iu = self.index('u')
507 u = self.u_init
509 fx = self.fx if self.base_x == 'fft' else np.pi * self.fx
510 fy = self.fy if self.base_y == 'fft' else np.pi * self.fy
512 time_dep = xp.exp(-self.nu * (fx**2 + fy**2) * t)
514 alpha = (self.b - self.a) / 2.0
515 beta = (self.c - self.b) / 2.0
516 gamma = (self.c + self.a) / 2.0
518 u[iu] = xp.sin(fx * self.X) * xp.sin(fy * self.Y) * time_dep + alpha * self.X + beta * self.Y + gamma
520 return u