Coverage for pySDC/helpers/ParaDiagHelper.py: 100%

1import numpy as np 

2import scipy.sparse as sp 

3 

4 

5def get_FFT_matrix(N): 

6 """ 

7 Get matrix for computing FFT of size N. Normalization is like "ortho" in numpy. 

8 Compute inverse FFT by multiplying by the complex conjugate (numpy.conjugate) of this matrix 

9 

10 Args: 

11 N (int): Size of the data to be transformed 

12 

13 Returns: 

14 numpy.ndarray: Dense square matrix to compute forward transform 

15 """ 

16 idx_1d = np.arange(N, dtype=complex) 

17 i1, i2 = np.meshgrid(idx_1d, idx_1d) 

18 

19 return np.exp(-2 * np.pi * 1j * i1 * i2 / N) / np.sqrt(N) 

20 

21 

22def get_E_matrix(N, alpha=0): 

23 """ 

24 Get NxN matrix with -1 on the lower subdiagonal, -alpha in the top right and 0 elsewhere 

25 

26 Args: 

27 N (int): Size of the matrix 

28 alpha (float): Negative of value in the top right 

29 

30 Returns: 

31 sparse E matrix 

32 """ 

33 E = sp.diags( 

34 [ 

35 -1.0, 

36 ] 

37 * (N - 1), 

38 offsets=-1, 

39 ).tolil() 

40 E[0, -1] = -alpha 

41 return E 

42 

43 

44def get_J_matrix(N, alpha): 

45 """ 

46 Get matrix for weights in the weighted inverse FFT 

47 

48 Args: 

49 N (int): Size of the matrix 

50 alpha (float): alpha parameter in ParaDiag 

51 

52 Returns: 

53 sparse J matrix 

54 """ 

55 gamma = alpha ** (-np.arange(N) / N) 

56 return sp.diags(gamma) 

57 

58 

59def get_J_inv_matrix(N, alpha): 

60 """ 

61 Get matrix for weights in the weighted FFT 

62 

63 Args: 

64 N (int): Size of the matrix 

65 alpha (float): alpha parameter in ParaDiag 

66 

67 Returns: 

68 sparse J_inv matrix 

69 """ 

70 gamma = alpha ** (-np.arange(N) / N) 

71 return sp.diags(1 / gamma) 

72 

73 

74def get_weighted_FFT_matrix(N, alpha): 

75 """ 

76 Get matrix for the weighted FFT 

77 

78 Args: 

79 N (int): Size of the matrix 

80 alpha (float): alpha parameter in ParaDiag 

81 

82 Returns: 

83 Dense weighted FFT matrix 

84 """ 

85 return get_FFT_matrix(N) @ get_J_inv_matrix(N, alpha) 

86 

87 

88def get_weighted_iFFT_matrix(N, alpha): 

89 """ 

90 Get matrix for the weighted inverse FFT 

91 

92 Args: 

93 N (int): Size of the matrix 

94 alpha (float): alpha parameter in ParaDiag 

95 

96 Returns: 

97 Dense weighted FFT matrix 

98 """ 

99 return get_J_matrix(N, alpha) @ np.conjugate(get_FFT_matrix(N)) 

100 

101 

102def get_H_matrix(N, sweeper_params): 

103 """ 

104 Get sparse matrix for computing the collocation update. Requires not to do a collocation update! 

105 

106 Args: 

107 N (int): Number of collocation nodes 

108 sweeper_params (dict): Parameters for the sweeper 

109 

110 Returns: 

111 Sparse matrix for collocation update 

112 """ 

113 assert sweeper_params['quad_type'] == 'RADAU-RIGHT' 

114 H = sp.eye(N).tolil() * 0 

115 H[:, -1] = 1 

116 return H 

117 

118 

119def get_G_inv_matrix(l, L, alpha, sweeper_params): 

120 """ 

121 Get the inverse of the preconditioner of the collocation problem of step `l` after diagonalization in time. 

122 

123 Args: 

124 l (int): Index of the step in the block 

125 L (int): Number of steps in the block 

126 alpha (float): alpha parameter in ParaDiag 

127 sweeper_params (dict): Parameters for the sweeper, needs `num_nodes` and `quad_type='RADAU-RIGHT'` 

128 

129 Returns: 

130 numpy.ndarray: Dense matrix (I + d_l H)^-1 of shape (num_nodes, num_nodes), with d_l the `l`-th eigenvalue of 

131 the alpha-circulant `E` from `get_E_matrix` and H from `get_H_matrix` 

132 """ 

133 M = sweeper_params['num_nodes'] 

134 I_M = sp.eye(M) 

135 E_alpha = get_E_matrix(L, alpha) 

136 H = get_H_matrix(M, sweeper_params) 

137 

138 gamma = alpha ** (-np.arange(L) / L) 

139 diags = np.fft.fft(1 / gamma * E_alpha[:, 0].toarray().flatten(), norm='backward') 

140 G = (diags[l] * H + I_M).tocsc() 

141 if M > 1: 

142 return sp.linalg.inv(G).toarray() 

143 else: 

144 return 1 / G.toarray()