Coverage for pySDC/implementations/problem_classes/boussinesq_helpers/build2DFDMatrix.py: 85%

1import numpy as np 

2import scipy.sparse as sp 

3 

4from pySDC.implementations.problem_classes.boussinesq_helpers.buildFDMatrix import ( 

5 getMatrix, 

6 getUpwindMatrix, 

7) 

8 

9 

10# 

11# 

12# 

13def get2DUpwindMatrix(N, dx, order): 

14 Dx = getUpwindMatrix(N[0], dx, order) 

15 return sp.kron(Dx, sp.eye(N[1]), format="csr") 

16 

17 

18# 

19# 

20# 

21def get2DMesh(N, x_b, z_b, bc_hor, bc_ver): 

22 assert np.size(N) == 2, 'N needs to be an array with two entries: N[0]=Nx and N[1]=Nz' 

23 assert ( 

24 np.size(x_b) == 2 

25 ), 'x_b needs to be an array with two entries: x_b[0] = left boundary, x_b[1] = right boundary' 

26 assert ( 

27 np.size(z_b) == 2 

28 ), 'z_b needs to be an array with two entries: z_b[0] = lower boundary, z_b[1] = upper boundary' 

29 

30 h = np.zeros(2) 

31 x = None 

32 z = None 

33 

34 if bc_hor[0] in ['periodic']: 

35 assert bc_hor[1] in ['periodic'], 'Periodic boundary conditions must be prescribed at both boundaries' 

36 x = np.linspace(x_b[0], x_b[1], N[0], endpoint=False) 

37 h[0] = x[1] - x[0] 

38 

39 if bc_hor[0] in ['dirichlet', 'neumann']: 

40 x = np.linspace(x_b[0], x_b[1], N[0] + 2, endpoint=True) 

41 x = x[1 : N[0] + 1] 

42 h[0] = x[1] - x[0] 

43 

44 if bc_ver[0] in ['periodic']: 

45 assert bc_ver[1] in ['periodic'], 'Periodic boundary conditions must be prescribed at both boundaries' 

46 z = np.linspace(z_b[0], z_b[1], N[1], endpoint=False) 

47 h[1] = z[1] - z[0] 

48 

49 if bc_ver[0] in ['dirichlet', 'neumann']: 

50 z = np.linspace(z_b[0], z_b[1], N[1] + 2, endpoint=True) 

51 z = z[1 : N[1] + 1] 

52 h[1] = z[1] - z[0] 

53 

54 xx, zz = np.meshgrid(x, z, indexing="ij") 

55 return xx, zz, h 

56 

57 

58# 

59# 

60# 

61def get2DMatrix(N, h, bc_hor, bc_ver, order): 

62 assert np.size(N) == 2, 'N needs to be an array with two entries: N[0]=Nx and N[1]=Nz' 

63 assert np.size(h) == 2, 'h needs to be an array with two entries: h[0]=dx and h[1]=dz' 

64 

65 Ax = getMatrix(N[0], h[0], bc_hor[0], bc_hor[1], order) 

66 Az = getMatrix(N[1], h[1], bc_ver[0], bc_ver[1], order) 

67 

68 Dx = sp.kron(Ax, sp.eye(N[1]), format="csr") 

69 Dz = sp.kron(sp.eye(N[0]), Az, format="csr") 

70 

71 return Dx, Dz