Coverage for pySDC/projects/FastWaveSlowWave/plot_dispersion.py: 97%
111 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 12:30 +0000
1import matplotlib
3matplotlib.use('Agg')
5import matplotlib.pyplot as plt
6import numpy as np
7import sympy
8from pylab import rcParams
10from pySDC.implementations.problem_classes.FastWaveSlowWave_0D import swfw_scalar
11from pySDC.implementations.sweeper_classes.imex_1st_order import imex_1st_order
13from pySDC.implementations.problem_classes.acoustic_helpers.standard_integrators import dirk, rk_imex
15from pySDC.core.step import Step
18def findomega(stab_fh):
19 assert np.array_equal(np.shape(stab_fh), [2, 2]), 'Not 2x2 matrix...'
20 omega = sympy.Symbol('omega')
21 func = (sympy.exp(-1j * omega) - stab_fh[0, 0]) * (sympy.exp(-1j * omega) - stab_fh[1, 1]) - stab_fh[
22 0, 1
23 ] * stab_fh[1, 0]
24 solsym = sympy.solve(func, omega)
25 sol0 = complex(solsym[0])
26 sol1 = complex(solsym[1])
27 if sol0.real >= 0:
28 sol = sol0
29 elif sol1.real >= 0:
30 sol = sol1
31 else:
32 print("Two roots with real part of same sign...")
33 sol = sol0
34 return sol
37def sdc_system_stability(L, Cs, Uadv, K):
38 """
39 Stability matrix of K IMEX-SDC sweeps plus the collocation update for u' = (Cs + Uadv) u, with Cs
40 treated implicitly and Uadv explicitly: the system version of the scalar K-sweep matrix of
41 imex_1st_order, starting from the initial value spread to all nodes.
43 Args:
44 L: level holding the imex_1st_order sweeper and the time step
45 Cs: fast (implicit) part of the system matrix
46 Uadv: slow (explicit) part of the system matrix
47 K: number of sweeps
49 Returns:
50 numpy.ndarray: matrix mapping u0 to the value at the end of the step
51 """
52 QE, QI, Q = L.sweep.QE[1:, 1:], L.sweep.QI[1:, 1:], L.sweep.coll.Qmat[1:, 1:]
53 nnodes, dt, n = L.sweep.coll.num_nodes, L.params.dt, Cs.shape[0]
55 LHS = np.eye(n * nnodes) - dt * (np.kron(QI, Cs) + np.kron(QE, Uadv))
56 RHS = dt * (np.kron(Q, Uadv + Cs) - np.kron(QI, Cs) - np.kron(QE, Uadv))
57 LHSinv = np.linalg.inv(LHS)
58 Mat_sweep = np.linalg.matrix_power(LHSinv.dot(RHS), K)
59 for k in range(0, K):
60 Mat_sweep = Mat_sweep + np.linalg.matrix_power(LHSinv.dot(RHS), k).dot(LHSinv)
62 update = dt * np.kron(L.sweep.coll.weights, Uadv + Cs)
63 spread = np.kron(np.ones((nnodes, 1)), np.eye(n))
64 return np.eye(n) + update.dot(Mat_sweep.dot(spread))
67def compute_and_plot_dispersion(Nsamples=15, K=3):
68 """
69 Function to compute and plot the dispersion relation
71 Args:
72 Nsamples: number of samples for testing
73 K: number of iterations as well as order
74 """
75 problem_params = dict()
76 # SET VALUE FOR lambda_slow AND VALUES FOR lambda_fast ###
77 problem_params['lambda_s'] = np.array([0.0])
78 problem_params['lambda_f'] = np.array([0.0])
79 problem_params['u0'] = 1.0
81 # initialize sweeper parameters
82 sweeper_params = dict()
83 # SET TYPE AND NUMBER OF QUADRATURE NODES ###
84 sweeper_params['quad_type'] = 'RADAU-RIGHT'
85 sweeper_params['do_coll_update'] = True
86 sweeper_params['num_nodes'] = 3
88 # initialize level parameters
89 level_params = dict()
90 level_params['dt'] = 1.0
92 # fill description dictionary for easy step instantiation
93 description = dict()
94 description['problem_class'] = swfw_scalar # pass problem class
95 description['problem_params'] = problem_params # pass problem parameters
96 description['sweeper_class'] = imex_1st_order # pass sweeper
97 description['sweeper_params'] = sweeper_params # pass sweeper parameters
98 description['level_params'] = level_params # pass level parameters
99 description['step_params'] = dict() # pass step parameters
101 # ORDER OF DIRK/IMEX IS EQUAL TO NUMBER OF ITERATIONS AND THUS ORDER OF SDC ###
102 dirk_order = K
104 c_speed = 1.0
105 U_speed = 0.05
107 # now the description contains more or less everything we need to create a step
108 S = Step(description=description)
110 L = S.levels[0]
112 # u0 = S.levels[0].prob.u_exact(t0)
113 # S.init_step(u0)
114 k_vec = np.linspace(0, np.pi, Nsamples + 1, endpoint=False)
115 k_vec = k_vec[1:]
116 phase = np.zeros((3, Nsamples))
117 amp_factor = np.zeros((3, Nsamples))
119 for i in range(0, np.size(k_vec)):
120 Cs = -1j * k_vec[i] * np.array([[0.0, c_speed], [c_speed, 0.0]], dtype='complex')
121 Uadv = -1j * k_vec[i] * np.array([[U_speed, 0.0], [0.0, U_speed]], dtype='complex')
123 stab_sdc = sdc_system_stability(L, Cs, Uadv, K)
124 y1 = np.array([1, 0], dtype='complex')
125 y2 = np.array([0, 1], dtype='complex')
127 # Stability function of backward Euler is 1/(1-z); system is y' = (Cs+Uadv)*y
128 # stab_ie = np.linalg.inv( np.eye(2) - step.status.dt*(Cs+Uadv) )
130 # For testing, insert exact stability function exp(-dt*i*k*(Cs+Uadv)
131 # stab_fh = la.expm(Cs+Uadv)
133 dirkts = dirk(Cs + Uadv, dirk_order)
134 stab_fh1 = dirkts.timestep(y1, 1.0)
135 stab_fh2 = dirkts.timestep(y2, 1.0)
136 stab_dirk = np.column_stack((stab_fh1, stab_fh2))
138 rkimex = rk_imex(M_fast=Cs, M_slow=Uadv, order=K)
139 stab_fh1 = rkimex.timestep(y1, 1.0)
140 stab_fh2 = rkimex.timestep(y2, 1.0)
141 stab_rk_imex = np.column_stack((stab_fh1, stab_fh2))
143 sol_sdc = findomega(stab_sdc)
144 sol_dirk = findomega(stab_dirk)
145 sol_rk_imex = findomega(stab_rk_imex)
147 # Now solve for discrete phase
148 phase[0, i] = sol_sdc.real / k_vec[i]
149 amp_factor[0, i] = np.exp(sol_sdc.imag)
150 phase[1, i] = sol_dirk.real / k_vec[i]
151 amp_factor[1, i] = np.exp(sol_dirk.imag)
152 phase[2, i] = sol_rk_imex.real / k_vec[i]
153 amp_factor[2, i] = np.exp(sol_rk_imex.imag)
155 rcParams['figure.figsize'] = 1.5, 1.5
156 fs = 8
157 fig = plt.figure()
158 plt.plot(k_vec, (U_speed + c_speed) + np.zeros(np.size(k_vec)), '--', color='k', linewidth=1.5, label='Exact')
159 plt.plot(k_vec, phase[1, :], '-', color='g', linewidth=1.5, label='DIRK(' + str(dirkts.order) + ')')
160 plt.plot(
161 k_vec,
162 phase[2, :],
163 '-+',
164 color='r',
165 linewidth=1.5,
166 label='IMEX(' + str(rkimex.order) + ')',
167 markevery=(2, 3),
168 mew=1.0,
169 )
170 plt.plot(
171 k_vec,
172 phase[0, :],
173 '-o',
174 color='b',
175 linewidth=1.5,
176 label='SDC(' + str(K) + ')',
177 markevery=(1, 3),
178 markersize=fs / 2,
179 )
180 plt.xlabel('Wave number', fontsize=fs, labelpad=0.25)
181 plt.ylabel('Phase speed', fontsize=fs, labelpad=0.5)
182 plt.xlim([k_vec[0], k_vec[-1:]])
183 plt.ylim([0.0, 1.1 * (U_speed + c_speed)])
184 fig.gca().tick_params(axis='both', labelsize=fs)
185 plt.legend(loc='lower left', fontsize=fs, prop={'size': fs - 2})
186 plt.xticks([0, 1, 2, 3], fontsize=fs)
187 filename = 'data/phase-K' + str(K) + '-M' + str(sweeper_params['num_nodes']) + '.png'
188 plt.gcf().savefig(filename, bbox_inches='tight')
190 fig = plt.figure()
191 plt.plot(k_vec, 1.0 + np.zeros(np.size(k_vec)), '--', color='k', linewidth=1.5, label='Exact')
192 plt.plot(k_vec, amp_factor[1, :], '-', color='g', linewidth=1.5, label='DIRK(' + str(dirkts.order) + ')')
193 plt.plot(
194 k_vec,
195 amp_factor[2, :],
196 '-+',
197 color='r',
198 linewidth=1.5,
199 label='IMEX(' + str(rkimex.order) + ')',
200 markevery=(2, 3),
201 mew=1.0,
202 )
203 plt.plot(
204 k_vec,
205 amp_factor[0, :],
206 '-o',
207 color='b',
208 linewidth=1.5,
209 label='SDC(' + str(K) + ')',
210 markevery=(1, 3),
211 markersize=fs / 2,
212 )
213 plt.xlabel('Wave number', fontsize=fs, labelpad=0.25)
214 plt.ylabel('Amplification factor', fontsize=fs, labelpad=0.5)
215 fig.gca().tick_params(axis='both', labelsize=fs)
216 plt.xlim([k_vec[0], k_vec[-1:]])
217 plt.ylim([k_vec[0], k_vec[-1:]])
218 plt.legend(loc='lower left', fontsize=fs, prop={'size': fs - 2})
219 plt.gca().set_ylim([0.0, 1.1])
220 plt.xticks([0, 1, 2, 3], fontsize=fs)
221 filename = 'data/ampfactor-K' + str(K) + '-M' + str(sweeper_params['num_nodes']) + '.png'
222 plt.gcf().savefig(filename, bbox_inches='tight')
225if __name__ == "__main__":
226 compute_and_plot_dispersion()