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

## 121 statements

, created at 2024-09-19 09:13 +0000

1import numpy as np

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

10class imexexp_1st_order(Sweeper):

11 """

12 Custom sweeper class, implements Sweeper.py

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.

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

18 """

20 def __init__(self, params):

21 """

22 Initialization routine for the custom sweeper

24 Args:

25 params: parameters for the sweeper

26 """

28 if "QI" not in params:

29 params["QI"] = "IE"

31 # call parent's initialization routine

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

34 # IMEX integration matrices

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

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

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()

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')

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

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

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

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

66 k = np.array(indeces)

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

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

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)

79 def compute_lambda_phi_Qmat_exp(self):

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)

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):

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

93 L = self.level

94 P = L.prob

95 M = self.coll.num_nodes

96 c = self.coll.nodes

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]

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)

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)

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]

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)

121 def integrate(self):

122 """

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

125 Returns:

126 list of dtype_u: containing the integral as values

127 """

129 # get current level and problem description

130 L = self.level

131 P = L.prob

132 M = self.coll.num_nodes

134 self.compute_lambda_phi_Qmat_exp()

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

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

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 )

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)

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

154 return me

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

160 Returns:

161 None

162 """

164 # get current level and problem description

165 L = self.level

166 P = L.prob

168 # only if the level has been touched before

169 assert L.status.unlocked

171 # get number of collocation nodes for easier access

172 M = self.coll.num_nodes

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]

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 )

192 # do the sweep

193 for m in range(M):

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]

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 )

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])

212 # indicate presence of new values at this level

213 L.status.updated = True

215 return None

217 def compute_end_point(self):

218 """

219 Compute u at the right point of the interval

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

223 Returns:

224 None

225 """

227 # get current level and problem description

228 L = self.level

229 P = L.prob

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

238 return None

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)

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

247 """

248 Computation of the residual using the collocation matrix Q

250 Args:

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

252 """

254 # get current level and problem description

255 L = self.level

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

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

264 # assert L.status.updated

266 # compute the residual for each node

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]))

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 )

298 # indicate that the residual has seen the new values

299 L.status.updated = False

301 return None