Coverage for pySDC/projects/Monodomain/sweeper_classes/exponential_runge_kutta/imexexp_1st_order.py: 95%

1import numpy as np 

2 

3from pySDC.core.Sweeper import sweeper 

4from pySDC.core.Errors import CollocationError, ParameterError 

5from pySDC.core.Collocation import CollBase 

6import numdifftools.fornberg as fornberg 

7import scipy 

8 

9 

10class imexexp_1st_order(sweeper): 

11 """ 

12 Custom sweeper class, implements Sweeper.py 

13 

14 First-order IMEXEXP sweeper using implicit/explicit/exponential Euler as base integrator 

15 In the cardiac electrphysiology community this is known as Rush-Larsen scheme. 

16 

17 The underlying intergrator is exponential Runge-Kutta, leading to exponential SDC (ESDC). 

18 """ 

19 

20 def __init__(self, params): 

21 """ 

22 Initialization routine for the custom sweeper 

23 

24 Args: 

25 params: parameters for the sweeper 

26 """ 

27 

28 if "QI" not in params: 

29 params["QI"] = "IE" 

30 

31 # call parent's initialization routine 

32 super(imexexp_1st_order, self).__init__(params) 

33 

34 # IMEX integration matrices 

35 self.QI = self.get_Qdelta_implicit(coll=self.coll, qd_type=self.params.QI) 

36 self.delta = np.diagonal(self.QI)[1:] 

37 

38 # Compute weights w such that PiQ^(k)(0) = sum_{j=0}^{M-1} w[k,j]*Q[j], k=0,...,M-1 

39 # Used to express the derivatives of a polynomial in x=0 in terms of the values of the polynomial at the collocation nodes 

40 M = self.coll.num_nodes 

41 c = self.coll.nodes 

42 self.w = fornberg.fd_weights_all(c, 0.0, M - 1).transpose() 

43 

44 # Define the quadature rule for the evaluation of the phi_i(z) functions. Indeed, we evaluate them as integrals in order to avoid round off errors. 

45 phi_num_nodes = 5 # seems to be enough in most cases 

46 self.phi_coll = CollBase(num_nodes=phi_num_nodes, tleft=0, tright=1, node_type='LEGENDRE', quad_type='GAUSS') 

47 

48 def phi_eval(self, factors, indeces, phi, lmbda): 

49 """ 

50 Evaluate the phi_k functions at the points factors[i]*lmbda, for all k in indeces 

51 

52 Arguments: 

53 factors: list of factors to multiply lmbda with. 

54 indeces: list of indeces k for the phi_k functions. Since we use the integral formulation, k=0 is not allowed (not needed neither). 

55 phi: an instance of mesh with shape (len(factors),len(indeces),*lmbda.shape) (i.e., some space to store the results) 

56 it will filled as: phi[i,k][:] = phi_{indeces[k]}(factor[i]*lmbda[:]) 

57 lmbda: dtype_u: the value of lmbda 

58 """ 

59 

60 assert 0 not in indeces, "phi_0 is not implemented, since the integral definition is not valid for k=0." 

61 

62 # the quadrature rule used to evaluate the phi functions as integrals. This is not the same as the one used in the ESDC method!!!! 

63 c = self.phi_coll.nodes 

64 b = self.phi_coll.weights 

65 

66 k = np.array(indeces) 

67 km1_fac = scipy.special.factorial(k - 1) # (k-1)! 

68 

69 # Here we use the quadrature rule to approximate the integral 

70 # phi_{k}(factor[i]*lmbda[:,:])= \int_0^1 e^{(1-s)*factor[i]*lambda[:,:]}*s^{k-1}/(k-1)! ds 

71 

72 # First, compute e^((1-c[j])*factor[i]*lmbda[:]) for nodes c[j] on the quadrature rule and all factors[i] 

73 exp_terms = np.exp(((1.0 - c[None, :, None, None]) * factors[:, None, None, None]) * lmbda[None, None, :, :]) 

74 # Then, compute the terms c[j]^{k-1}/(k-1)! for all nodes c[j] and all k and multiply with the weights b[j] 

75 wgt_tmp = (b[:, None] * c[:, None] ** (k[None, :] - 1)) / km1_fac[None, :] 

76 # Finally, compute the integral by summing over the quadrature nodes 

77 phi[:] = np.sum(wgt_tmp[None, :, :, None, None] * exp_terms[:, :, None, :, :], axis=1) 

78 

79 def compute_lambda_phi_Qmat_exp(self): 

80 

81 if not hasattr(self, "u_old"): 

82 # make some space for the old value of u[0] 

83 self.u_old = self.level.prob.dtype_u(init=self.level.prob.init, val=0.0) 

84 

85 # everything that is computed in this if statement depends on u[0] only 

86 # To save computations we recompute that only if u[0] has changed. 

87 # Also, we check only for the first component u[0][0] of u[0] to save more computations. 

88 # Remember that u[0][0] is a vector representing the electric potential on the whole mesh and is enough to check if the whole u[0] has changed. 

89 if not np.allclose(self.u_old[0], self.level.u[0][0], rtol=1e-10, atol=1e-10): 

90 

91 self.u_old[:] = self.level.u[0] 

92 

93 L = self.level 

94 P = L.prob 

95 M = self.coll.num_nodes 

96 c = self.coll.nodes 

97 

98 # compute lambda(u) of the exponential term f_exp(u)=lmbda(u)*(u-y_inf(u)) 

99 # and select only the indeces with exponential terms (others are zeros) 

100 self.lmbda = P.lmbda_eval(L.u[0], L.time)[P.rhs_exp_indeces] 

101 

102 if not hasattr(self, "phi"): 

103 # make some space 

104 self.phi = P.dtype_u(init=P.init_exp_extruded((M, M)), val=0.0) 

105 self.phi_one = P.dtype_u(init=P.init_exp_extruded((M, 1)), val=0.0) 

106 

107 # evaluate the phi_k(dt*c_i*lambda) functions at the collocation nodes c_i for k=1,...,M 

108 self.phi_eval(L.dt * c, list(range(1, M + 1)), self.phi, self.lmbda) 

109 # evaluates phi_1(dt*delta_i*lambda) for delta_i = c_i - c_{i-1} 

110 self.phi_eval(L.dt * self.delta, [1], self.phi_one, self.lmbda) 

111 

112 # compute weight for the integration of \int_0^ci exp(dt*(ci-r)lmbda)*PiQ(r)dr, 

113 # where PiQ(r) is a polynomial interpolating some nodal values Q(c_i)=Q[i]. 

114 # The integral of PiQ will be approximated as: 

115 # \int_0^ci exp(dt*(ci-r)lmbda)*PiQ(r)dr ~= \sum_{j=0}^{M-1} Qmat_exp[i,j]*Q[j] 

116 

117 k = np.arange(0, M) 

118 wgt_tmp = self.w[None, :, :] * c[:, None, None] ** (k[None, None, :] + 1) 

119 self.Qmat_exp = np.sum(wgt_tmp[:, :, :, None, None] * self.phi[:, None, :, :, :], axis=2) 

120 

121 def integrate(self): 

122 """ 

123 Integrates the right-hand side (here impl + expl + exp) using exponential Runge-Kutta 

124 

125 Returns: 

126 list of dtype_u: containing the integral as values 

127 """ 

128 

129 # get current level and problem description 

130 L = self.level 

131 P = L.prob 

132 M = self.coll.num_nodes 

133 

134 self.compute_lambda_phi_Qmat_exp() 

135 

136 if not hasattr(self, "Q"): 

137 self.Q = P.dtype_u(init=P.init_exp_extruded((M,)), val=0.0) 

138 

139 for k in range(M): 

140 self.Q[k][:] = L.f[k + 1].exp[P.rhs_exp_indeces] + self.lmbda * ( 

141 L.u[0][P.rhs_exp_indeces] - L.u[k + 1][P.rhs_exp_indeces] 

142 ) 

143 

144 # integrate RHS over all collocation nodes 

145 me = [P.dtype_u(init=P.init, val=0.0) for _ in range(M)] 

146 for m in range(1, M + 1): 

147 for j in range(1, M + 1): 

148 me[m - 1][P.rhs_stiff_indeces] += self.coll.Qmat[m, j] * L.f[j].impl[P.rhs_stiff_indeces] 

149 me[m - 1][P.rhs_nonstiff_indeces] += self.coll.Qmat[m, j] * L.f[j].expl[P.rhs_nonstiff_indeces] 

150 me[m - 1][P.rhs_exp_indeces] += np.sum(self.Qmat_exp[m - 1] * self.Q, axis=0) 

151 

152 me[m - 1] *= L.dt 

153 

154 return me 

155 

156 def update_nodes(self): 

157 """ 

158 Update the u- and f-values at the collocation nodes -> corresponds to a single sweep over all nodes 

159 

160 Returns: 

161 None 

162 """ 

163 

164 # get current level and problem description 

165 L = self.level 

166 P = L.prob 

167 

168 # only if the level has been touched before 

169 assert L.status.unlocked 

170 

171 # get number of collocation nodes for easier access 

172 M = self.coll.num_nodes 

173 

174 integral = self.integrate() 

175 for m in range(M): 

176 if L.tau[m] is not None: 

177 integral[m] += L.tau[m] 

178 for i in range(1, M): 

179 integral[M - i] -= integral[M - i - 1] 

180 

181 # prepare the integral term 

182 for m in range(M): 

183 integral[m][P.rhs_stiff_indeces] += -L.dt * self.delta[m] * L.f[m + 1].impl[P.rhs_stiff_indeces] 

184 integral[m][P.rhs_nonstiff_indeces] += -L.dt * self.delta[m] * L.f[m].expl[P.rhs_nonstiff_indeces] 

185 integral[m][P.rhs_exp_indeces] += ( 

186 -L.dt 

187 * self.delta[m] 

188 * self.phi_one[m][0] 

189 * (L.f[m].exp[P.rhs_exp_indeces] + self.lmbda * (L.u[0][P.rhs_exp_indeces] - L.u[m][P.rhs_exp_indeces])) 

190 ) 

191 

192 # do the sweep 

193 for m in range(M): 

194 

195 tmp = L.u[m] + integral[m] 

196 tmp[P.rhs_exp_indeces] += ( 

197 L.dt 

198 * self.delta[m] 

199 * self.phi_one[m][0] 

200 * (L.f[m].exp[P.rhs_exp_indeces] + self.lmbda * (L.u[0][P.rhs_exp_indeces] - L.u[m][P.rhs_exp_indeces])) 

201 ) 

202 tmp[P.rhs_nonstiff_indeces] += L.dt * self.delta[m] * L.f[m].expl[P.rhs_nonstiff_indeces] 

203 

204 # implicit solve with prefactor stemming from QI 

205 L.u[m + 1] = P.solve_system( 

206 tmp, L.dt * self.QI[m + 1, m + 1], L.u[m + 1], L.time + L.dt * self.coll.nodes[m], L.u[m + 1] 

207 ) 

208 

209 # update function values 

210 P.eval_f(L.u[m + 1], L.time + L.dt * self.coll.nodes[m], fh=L.f[m + 1]) 

211 

212 # indicate presence of new values at this level 

213 L.status.updated = True 

214 

215 return None 

216 

217 def compute_end_point(self): 

218 """ 

219 Compute u at the right point of the interval 

220 

221 The value uend computed here is a full evaluation of the Picard formulation unless do_full_update==False 

222 

223 Returns: 

224 None 

225 """ 

226 

227 # get current level and problem description 

228 L = self.level 

229 P = L.prob 

230 

231 # check if Mth node is equal to right point and do_coll_update is false, perform a simple copy 

232 if self.coll.right_is_node and not self.params.do_coll_update: 

233 # a copy is sufficient 

234 L.uend = P.dtype_u(L.u[-1]) 

235 else: 

236 raise CollocationError("This option is not implemented yet.") 

237 

238 return None 

239 

240 def rel_norm(self, a, b): 

241 norms = [] 

242 for i in range(len(a)): 

243 norms.append(np.linalg.norm(a[i]) / np.linalg.norm(b[i])) 

244 return np.average(norms) 

245 

246 def compute_residual(self, stage=''): 

247 """ 

248 Computation of the residual using the collocation matrix Q 

249 

250 Args: 

251 stage (str): The current stage of the step the level belongs to 

252 """ 

253 

254 # get current level and problem description 

255 L = self.level 

256 

257 # Check if we want to skip the residual computation to gain performance 

258 # Keep in mind that skipping any residual computation is likely to give incorrect outputs of the residual! 

259 if stage in self.params.skip_residual_computation: 

260 L.status.residual = 0.0 if L.status.residual is None else L.status.residual 

261 return None 

262 

263 # check if there are new values (e.g. from a sweep) 

264 # assert L.status.updated 

265 

266 # compute the residual for each node 

267 

268 # build QF(u) 

269 res_norm = [] 

270 rel_res_norm = [] 

271 res = self.integrate() 

272 for m in range(self.coll.num_nodes): 

273 res[m] += L.u[0] 

274 res[m] -= L.u[m + 1] 

275 # add tau if associated 

276 if L.tau[m] is not None: 

277 res[m] += L.tau[m] 

278 # use abs function from data type here 

279 res_norm.append(abs(res[m])) 

280 # the different components of the monodomain equation have very different magnitude therefore we use a tailored relative norm here to avoid the cancellation of the smaller components 

281 rel_res_norm.append(self.rel_norm(res[m], L.u[0])) 

282 

283 # find maximal residual over the nodes 

284 if L.params.residual_type == 'full_abs': 

285 L.status.residual = max(res_norm) 

286 elif L.params.residual_type == 'last_abs': 

287 L.status.residual = res_norm[-1] 

288 elif L.params.residual_type == 'full_rel': 

289 L.status.residual = max(rel_res_norm) 

290 elif L.params.residual_type == 'last_rel': 

291 L.status.residual = rel_res_norm[-1] 

292 else: 

293 raise ParameterError( 

294 f'residual_type = {L.params.residual_type} not implemented, choose ' 

295 f'full_abs, last_abs, full_rel or last_rel instead' 

296 ) 

297 

298 # indicate that the residual has seen the new values 

299 L.status.updated = False 

300 

301 return None