Coverage for pySDC/projects/parallelSDC_reloaded/scripts/fig03_lorenz.py: 100%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3""" 

4Created on Wed Jan 10 15:34:24 2024 

5 

6Figures with experiment on the Lorenz problem 

7""" 

8import os 

9import numpy as np 

10import scipy as sp 

11 

12from pySDC.projects.parallelSDC_reloaded.utils import solutionExact, getParamsSDC, solutionSDC, getParamsRK, plt 

13from pySDC.helpers.testing import DataChecker 

14 

15data = DataChecker(__file__) 

16 

17PATH = '/' + os.path.join(*__file__.split('/')[:-1]) 

18SCRIPT = __file__.split('/')[-1].split('.')[0] 

19 

20symList = ['o', '^', 's', '>', '*', '<', 'p', '>'] * 10 

21 

22# SDC parameters 

23nNodes = 4 

24quadType = 'RADAU-RIGHT' 

25nodeType = 'LEGENDRE' 

26parEfficiency = 0.8 # 1/nNodes 

27 

28# ----------------------------------------------------------------------------- 

29# Trajectories (reference solution) 

30# ----------------------------------------------------------------------------- 

31tEnd = 2 

32nSteps = tEnd * 50 

33tVals = np.linspace(0, tEnd, nSteps + 1) 

34nPeriods = 2 

35 

36print(f"Computing exact solution up to t={tEnd} ...") 

37uExact = solutionExact(tEnd, nSteps, "LORENZ", u0=(5, -5, 20)) 

38 

39z = uExact[:, -1] 

40idx = sp.signal.find_peaks(z)[0][nPeriods - 1] 

41print(f'tEnd for {nPeriods} periods : {tVals[idx]}') 

42 

43figName = f"{SCRIPT}_traj" 

44plt.figure(figName) 

45me = 0.1 

46plt.plot(tVals, uExact[:, 0], 's-', label="$x(t)$", markevery=me) 

47plt.plot(tVals, uExact[:, 1], 'o-', label="$y(t)$", markevery=me) 

48plt.plot(tVals, uExact[:, 2], '^-', label="$z(t)$", markevery=me) 

49plt.vlines(tVals[idx], ymin=-20, ymax=40, linestyles="--", linewidth=1) 

50plt.legend(loc="upper right") 

51plt.xlabel("$t$") 

52plt.ylabel("Trajectory") 

53plt.gcf().set_size_inches(12, 3) 

54plt.tight_layout() 

55plt.savefig(f'{PATH}/{figName}.pdf') 

56 

57# ----------------------------------------------------------------------------- 

58# %% Convergence plots 

59# ----------------------------------------------------------------------------- 

60tEnd = 1.24 

61nStepsList = np.array([2, 5, 10, 20, 50, 100, 200, 500, 1000]) 

62dtVals = tEnd / nStepsList 

63 

64 

65def getError(uNum, uRef): 

66 if uNum is None: # pragma: no cover 

67 return np.inf 

68 return np.linalg.norm(np.linalg.norm(uRef - uNum, np.inf, axis=-1), np.inf) 

69 

70 

71config = ["PIC", "MIN-SR-NS"] 

72for qDelta, sym in zip(config, symList): 

73 figName = f"{SCRIPT}_conv_{qDelta}" 

74 plt.figure(figName) 

75 

76 for nSweeps in [1, 2, 3, 4, 5]: 

77 params = getParamsSDC(quadType=quadType, numNodes=nNodes, nodeType=nodeType, qDeltaI=qDelta, nSweeps=nSweeps) 

78 

79 errors = [] 

80 

81 for nSteps in nStepsList: 

82 print(f' -- nSteps={nSteps} ...') 

83 

84 uRef = solutionExact(tEnd, nSteps, "LORENZ", u0=(5, -5, 20)) 

85 

86 uSDC, counters, parallel = solutionSDC(tEnd, nSteps, params, "LORENZ", u0=(5, -5, 20)) 

87 

88 err = getError(uSDC, uRef) 

89 errors.append(err) 

90 

91 # error VS dt 

92 label = f"$K={nSweeps}$" 

93 plt.loglog(dtVals, errors, sym + '-', label=f"$K={nSweeps}$") 

94 data.storeAndCheck(f"{figName}_{label}", errors[1:]) 

95 

96 x = dtVals[4:] 

97 for k in [1, 2, 3, 4, 5, 6]: 

98 plt.loglog(x, 1e4 * x**k, "--", color="gray", linewidth=0.8) 

99 

100 plt.gca().set( 

101 xlabel=r"$\Delta{t}$", 

102 ylabel=r"$L_\infty$ error", 

103 ylim=(8.530627786509715e-12, 372.2781393394293), 

104 ) 

105 plt.legend(loc="lower right") 

106 plt.grid() 

107 plt.tight_layout() 

108 plt.savefig(f"{PATH}/{figName}.pdf") 

109 

110 

111# ----------------------------------------------------------------------------- 

112# %% Error VS cost plots 

113# ----------------------------------------------------------------------------- 

114def getCost(counters): 

115 nNewton, nRHS, tComp = counters 

116 return nNewton + nRHS 

117 

118 

119minPrec = ["MIN-SR-NS", "MIN-SR-S", "MIN-SR-FLEX"] 

120 

121symList = ['^', '>', '<', 'o', 's', '*'] 

122config = [ 

123 [(*minPrec, "LU", "EE", "PIC"), 4], 

124 [(*minPrec, "VDHS", "RK4", "ESDIRK43"), 4], 

125 [(*minPrec, "PIC", "RK4", "ESDIRK43"), 5], 

126] 

127 

128 

129i = 0 

130for qDeltaList, nSweeps in config: 

131 figName = f"{SCRIPT}_cost_{i}" 

132 i += 1 

133 plt.figure(figName) 

134 

135 for qDelta, sym in zip(qDeltaList, symList): 

136 try: 

137 params = getParamsRK(qDelta) 

138 except KeyError: 

139 params = getParamsSDC( 

140 quadType=quadType, numNodes=nNodes, nodeType=nodeType, qDeltaI=qDelta, nSweeps=nSweeps 

141 ) 

142 

143 errors = [] 

144 costs = [] 

145 

146 for nSteps in nStepsList: 

147 uRef = solutionExact(tEnd, nSteps, "LORENZ", u0=(5, -5, 20)) 

148 

149 uSDC, counters, parallel = solutionSDC(tEnd, nSteps, params, "LORENZ", u0=(5, -5, 20)) 

150 

151 err = getError(uSDC, uRef) 

152 errors.append(err) 

153 

154 cost = getCost(counters) 

155 if parallel: 

156 assert qDelta != "EE", "wait, whaaat ??" 

157 cost /= nNodes * parEfficiency 

158 costs.append(cost) 

159 

160 # error VS cost 

161 ls = '-' if qDelta.startswith("MIN-SR-") else "--" 

162 plt.loglog(costs, errors, sym + ls, label=qDelta) 

163 data.storeAndCheck(f"{figName}_{qDelta}", errors[2:], rtol=1e-2) 

164 

165 plt.gca().set( 

166 xlabel="Cost", 

167 ylabel=r"$L_\infty$ error", 

168 ylim=(1e-10, 400), 

169 xlim=(30, 20000), 

170 ) 

171 plt.legend(loc="lower left") 

172 plt.grid() 

173 plt.tight_layout() 

174 plt.savefig(f"{PATH}/{figName}.pdf") 

175 

176data.writeToJSON()