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""" 

8 

9import os 

10import numpy as np 

11import scipy as sp 

12 

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

14from pySDC.helpers.testing import DataChecker 

15 

16data = DataChecker(__file__) 

17 

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

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

20 

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

22 

23# SDC parameters 

24nNodes = 4 

25quadType = 'RADAU-RIGHT' 

26nodeType = 'LEGENDRE' 

27parEfficiency = 0.8 # 1/nNodes 

28 

29# ----------------------------------------------------------------------------- 

30# Trajectories (reference solution) 

31# ----------------------------------------------------------------------------- 

32tEnd = 2 

33nSteps = tEnd * 50 

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

35nPeriods = 2 

36 

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

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

39 

40z = uExact[:, -1] 

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

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

43 

44figName = f"{SCRIPT}_traj" 

45plt.figure(figName) 

46me = 0.1 

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

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

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

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

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

52plt.xlabel("$t$") 

53plt.ylabel("Trajectory") 

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

55plt.tight_layout() 

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

57 

58# ----------------------------------------------------------------------------- 

59# %% Convergence plots 

60# ----------------------------------------------------------------------------- 

61tEnd = 1.24 

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

63dtVals = tEnd / nStepsList 

64 

65# The reference solution depends only on nSteps, so compute it once per nSteps 

66# instead of once per (qDelta, nSteps). 

67uRefs = {nSteps: solutionExact(tEnd, nSteps, "LORENZ", u0=(5, -5, 20)) for nSteps in nStepsList} 

68 

69 

70def getError(uNum, uRef): 

71 if uNum is None: # pragma: no cover 

72 return np.inf 

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

74 

75 

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

77for qDelta, sym in zip(config, symList, strict=False): 

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

79 plt.figure(figName) 

80 

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

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

83 

84 errors = [] 

85 

86 for nSteps in nStepsList: 

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

88 

89 uRef = uRefs[nSteps] 

90 

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

92 

93 err = getError(uSDC, uRef) 

94 errors.append(err) 

95 

96 # error VS dt 

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

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

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

100 

101 x = dtVals[4:] 

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

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

104 

105 plt.gca().set( 

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

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

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

109 ) 

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

111 plt.grid() 

112 plt.tight_layout() 

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

114 

115 

116# ----------------------------------------------------------------------------- 

117# %% Error VS cost plots 

118# ----------------------------------------------------------------------------- 

119def getCost(counters): 

120 nNewton, nRHS, tComp = counters 

121 return nNewton + nRHS 

122 

123 

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

125 

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

127config = [ 

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

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

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

131] 

132 

133 

134i = 0 

135for qDeltaList, nSweeps in config: 

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

137 i += 1 

138 plt.figure(figName) 

139 

140 for qDelta, sym in zip(qDeltaList, symList, strict=False): 

141 try: 

142 params = getParamsRK(qDelta) 

143 except KeyError: 

144 params = getParamsSDC( 

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

146 ) 

147 

148 errors = [] 

149 costs = [] 

150 

151 for nSteps in nStepsList: 

152 uRef = uRefs[nSteps] 

153 

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

155 

156 err = getError(uSDC, uRef) 

157 errors.append(err) 

158 

159 cost = getCost(counters) 

160 if parallel: 

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

162 cost /= nNodes * parEfficiency 

163 costs.append(cost) 

164 

165 # error VS cost 

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

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

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

169 

170 plt.gca().set( 

171 xlabel="Cost", 

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

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

174 xlim=(30, 20000), 

175 ) 

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

177 plt.grid() 

178 plt.tight_layout() 

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

180 

181data.writeToJSON()