Coverage for pySDC/core/Sweeper.py: 83%

1import logging

2import warnings

4import numpy as np

5import scipy as sp

6import scipy.linalg

7import scipy.optimize as opt

9from pySDC.core.Errors import ParameterError

10from pySDC.core.Level import level

11from pySDC.core.Collocation import CollBase

12from pySDC.helpers.pysdc_helper import FrozenClass

15# short helper class to add params as attributes

16class _Pars(FrozenClass):

17 def __init__(self, pars):

18 self.do_coll_update = False

19 self.initial_guess = 'spread' # default value (see also below)

20 self.skip_residual_computation = () # gain performance at the cost of correct residual output

22 for k, v in pars.items():

23 if k != 'collocation_class':

24 setattr(self, k, v)

26 self._freeze()

29class sweeper(object):

30 """

31 Base abstract sweeper class

33 Attributes:

34 logger: custom logger for sweeper-related logging

35 params (__Pars): parameter object containing the custom parameters passed by the user

36 coll (pySDC.Collocation.CollBase): collocation object

37 """

39 def __init__(self, params):

40 """

41 Initialization routine for the base sweeper

43 Args:

44 params (dict): parameter object

46 """

48 # set up logger

49 self.logger = logging.getLogger('sweeper')

51 essential_keys = ['num_nodes']

52 for key in essential_keys:

53 if key not in params:

54 msg = 'need %s to instantiate step, only got %s' % (key, str(params.keys()))

55 self.logger.error(msg)

56 raise ParameterError(msg)

58 if 'collocation_class' not in params:

59 params['collocation_class'] = CollBase

61 # prepare random generator for initial guess

62 if params.get('initial_guess', 'spread') == 'random': # default value (see also above)

63 params['random_seed'] = params.get('random_seed', 1984)

64 self.rng = np.random.RandomState(params['random_seed'])

66 self.params = _Pars(params)

68 self.coll: CollBase = params['collocation_class'](**params)

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

71 self.logger.warning(

72 'we need to do a collocation update here, since the right end point is not a node. Changing this!'

73 )

74 self.params.do_coll_update = True

76 # This will be set as soon as the sweeper is instantiated at the level

77 self.__level = None

79 self.parallelizable = False

81 def get_Qdelta_implicit(self, coll, qd_type):

82 def rho(x):

83 return max(abs(np.linalg.eigvals(np.eye(m) - np.diag([x[i] for i in range(m)]).dot(coll.Qmat[1:, 1:]))))

85 QDmat = np.zeros(coll.Qmat.shape)

86 if qd_type == 'LU':

87 QT = coll.Qmat[1:, 1:].T

88 [_, _, U] = scipy.linalg.lu(QT, overwrite_a=True)

89 QDmat[1:, 1:] = U.T

90 elif qd_type == 'LU2':

91 QT = coll.Qmat[1:, 1:].T

92 [_, _, U] = scipy.linalg.lu(QT, overwrite_a=True)

93 QDmat[1:, 1:] = 2 * U.T

94 elif qd_type == 'TRAP':

95 for m in range(coll.num_nodes + 1):

96 QDmat[m, 1 : m + 1] = coll.delta_m[0:m]

97 for m in range(coll.num_nodes + 1):

98 QDmat[m, 0:m] += coll.delta_m[0:m]

99 QDmat /= 2.0

100 elif qd_type == 'IE':

101 for m in range(coll.num_nodes + 1):

102 QDmat[m, 1 : m + 1] = coll.delta_m[0:m]

103 elif qd_type == 'IEpar':

104 for m in range(coll.num_nodes + 1):

105 QDmat[m, m] = np.sum(coll.delta_m[0:m])

106 self.parallelizable = True

107 elif qd_type == 'Qpar':

108 QDmat = np.diag(np.diag(coll.Qmat))

109 self.parallelizable = True

110 elif qd_type == 'GS':

111 QDmat = np.tril(coll.Qmat)

112 elif qd_type == 'PIC':

113 QDmat = np.zeros(coll.Qmat.shape)

114 self.parallelizable = True

115 elif qd_type == 'MIN':

116 m = QDmat.shape[0] - 1

117 x0 = 10 * np.ones(m)

118 d = opt.minimize(rho, x0, method='Nelder-Mead')

119 QDmat[1:, 1:] = np.linalg.inv(np.diag(d.x))

120 self.parallelizable = True

121 elif qd_type.startswith('MIN-SR-FLEX'):

122 m = QDmat.shape[0] - 1

123 try:

124 k = abs(int(qd_type[11:]))

125 except ValueError:

126 k = 1

127 d = min(k, m)

128 QDmat[1:, 1:] = np.diag(coll.nodes) / d

129 self.parallelizable = True

130 elif qd_type in ['MIN_GT', 'MIN-SR-NS']:

131 m = QDmat.shape[0] - 1

132 QDmat[1:, 1:] = np.diag(coll.nodes) / m

133 self.parallelizable = True

134 elif qd_type == 'MIN3':

135 m = QDmat.shape[0] - 1

136 x = None

137 # These values have been obtained using Indie Solver, a commercial solver for black-box optimization which

138 # aggregates several state-of-the-art optimization methods (free academic subscription plan)

139 # objective function: sum over 17^2 values of lamdt, real and imaginary (WORKS SURPRISINGLY WELL!)

140 if coll.node_type == 'LEGENDRE' and coll.quad_type == 'LOBATTO':

141 if m == 9:

142 # rho = 0.154786693955

143 x = [

144 0.0,

145 0.14748983547536937,

146 0.1243753767395874,

147 0.08797965969063823,

148 0.03249792877433364,

149 0.06171633442251176,

150 0.08995295998705832,

151 0.1080641868728824,

152 0.11621787232558443,

153 ]

154 elif m == 7:

155 # rho = 0.0979351256833

156 x = [

157 0.0,

158 0.18827968699454273,

159 0.1307213945012976,

160 0.04545003319140543,

161 0.08690617895312261,

162 0.12326429119922168,

163 0.13815746843252427,

164 ]

165 elif m == 5:

166 # rho = 0.0513543155235

167 x = [0.0, 0.2994085231050721, 0.07923154575177252, 0.14338847088077, 0.17675509273708057]

168 elif m == 4:

169 # rho = 0.0381589713397

170 x = [0.0, 0.2865524188780046, 0.11264992497015984, 0.2583063168320655]

171 elif m == 3:

172 # rho = 0.013592619664

173 x = [0.0, 0.2113181799416633, 0.3943250920445912]

174 elif m == 2:

175 # rho = 0

176 x = [0.0, 0.5]

177 else:

178 NotImplementedError(

179 'This combination of preconditioner, node type and node number is not ' 'implemented'

180 )

181 elif coll.node_type == 'LEGENDRE' and coll.quad_type == 'RADAU-RIGHT':

182 if m == 9:

183 # rho = 0.151784861385

184 x = [

185 0.14208076083211416,

186 0.1288153963623986,

187 0.10608601069476883,

188 0.07509520272252024,

189 0.027986167728305308,

190 0.05351160749903067,

191 0.07911315989747868,

192 0.09514844658836666,

193 0.10204992319487571,

194 ]

195 elif m == 7:

196 # rho = 0.116400161888

197 x = [

198 0.15223871397682717,

199 0.12625448001038536,

200 0.08210714764924298,

201 0.03994434742760019,

202 0.1052662547386142,

203 0.14075805578834127,

204 0.15636085758812895,

205 ]

206 elif m == 5:

207 # rho = 0.0783352996958 (iteration 5355)

208 x = [

209 0.2818591930905709,

210 0.2011358490453793,

211 0.06274536689514164,

212 0.11790265267514095,

213 0.1571629578515223,

214 ]

215 elif m == 4:

216 # rho = 0.057498908343

217 x = [0.3198786751412953, 0.08887606314792469, 0.1812366328324738, 0.23273925017954]

218 elif m == 3:

219 # rho = 0.038744192979 (iteration 11188)

220 x = [0.3203856825077055, 0.1399680686269595, 0.3716708461097372]

221 elif m == 2:

222 # rho = 0.0208560702294 (iteration 6690)

223 x = [0.2584092406077449, 0.6449261740461826]

224 else:

225 raise NotImplementedError(

226 'This combination of preconditioner, node type and node number is not implemented'

227 )

228 elif coll.node_type == 'EQUID' and coll.quad_type == 'RADAU-RIGHT':

229 if m == 9:

230 # rho = 0.251820022583 (iteration 32402)

231 x = [

232 0.04067333763109274,

233 0.06893408176924318,

234 0.0944460427779633,

235 0.11847528720123894,

236 0.14153236351607695,

237 0.1638856774260845,

238 0.18569759470199648,

239 0.20707543960267513,

240 0.2280946565716198,

241 ]

242 elif m == 7:

243 # rho = 0.184582997611 (iteration 44871)

244 x = [

245 0.0582690792096515,

246 0.09937620459067688,

247 0.13668728443669567,

248 0.1719458323664216,

249 0.20585615258818232,

250 0.2387890485242656,

251 0.27096908017041393,

252 ]

253 elif m == 5:

254 # rho = 0.118441339197 (iteration 34581)

255 x = [

256 0.0937126798932547,

257 0.1619131388001843,

258 0.22442341539247537,

259 0.28385142992912565,

260 0.3412523013467262,

261 ]

262 elif m == 4:

263 # rho = 0.0844043254542 (iteration 33099)

264 x = [0.13194852204686872, 0.2296718892453916, 0.3197255970017318, 0.405619746972393]

265 elif m == 3:

266 # rho = 0.0504635143866 (iteration 9783)

267 x = [0.2046955744931575, 0.3595744268324041, 0.5032243650307717]

268 elif m == 2:

269 # rho = 0.0214806480623 (iteration 6109)

270 x = [0.3749891032632652, 0.6666472946796036]

271 else:

272 NotImplementedError(

273 'This combination of preconditioner, node type and node number is not ' 'implemented'

274 )

275 else:

276 NotImplementedError(

277 'This combination of preconditioner, node type and node number is not ' 'implemented'

278 )

279 QDmat[1:, 1:] = np.diag(x)

280 self.parallelizable = True

281

282 elif qd_type == "MIN-SR-S":

283 M = QDmat.shape[0] - 1

284 quadType = coll.quad_type

285 nodeType = coll.node_type

286

287 # Main function to compute coefficients

288 def computeCoeffs(M, a=None, b=None):

289 """

290 Compute diagonal coefficients for a given number of nodes M.

291 If `a` and `b` are given, then it uses as initial guess:

292

293 >>> a * nodes**b / M

294

295 If `a` is not given, then do not care about `b` and uses as initial guess:

296

297 >>> nodes / M

298

299 Parameters

300 ----------

301 M : int

302 Number of collocation nodes.

303 a : float, optional

304 `a` coefficient for the initial guess.

305 b : float, optional

306 `b` coefficient for the initial guess.

307

308 Returns

309 -------

310 coeffs : array

311 The diagonal coefficients.

312 nodes : array

313 The nodes associated to the current coefficients.

314 """

315 collM = CollBase(num_nodes=M, node_type=nodeType, quad_type=quadType)

316

317 QM = collM.Qmat[1:, 1:]

318 nodesM = collM.nodes

319

320 if quadType in ['LOBATTO', 'RADAU-LEFT']:

321 QM = QM[1:, 1:]

322 nodesM = nodesM[1:]

323 nCoeffs = len(nodesM)

324

325 if nCoeffs == 1:

326 coeffs = np.diag(QM)

327

328 else:

329

330 def nilpotency(coeffs):

331 """Function verifying the nilpotency from given coefficients"""

332 coeffs = np.asarray(coeffs)

333 kMats = [(1 - z) * np.eye(nCoeffs) + z * np.diag(1 / coeffs) @ QM for z in nodesM]

334 vals = [np.linalg.det(K) - 1 for K in kMats]

335 return np.array(vals)

336

337 if a is None:

338 coeffs0 = nodesM / M

339 else:

340 coeffs0 = a * nodesM**b / M

341

342 with warnings.catch_warnings():

343 warnings.simplefilter("ignore")

344 coeffs = sp.optimize.fsolve(nilpotency, coeffs0, xtol=1e-15)

345

346 # Handle first node equal to zero

347 if quadType in ['LOBATTO', 'RADAU-LEFT']:

348 coeffs = np.asarray([0.0] + list(coeffs))

349 nodesM = np.asarray([0.0] + list(nodesM))

350

351 return coeffs, nodesM

352

353 def fit(coeffs, nodes):

354 """Function fitting given coefficients to a power law"""

355

356 def lawDiff(ab):

357 a, b = ab

358 return np.linalg.norm(a * nodes**b - coeffs)

359

360 sol = sp.optimize.minimize(lawDiff, [1.0, 1.0], method="nelder-mead")

361 return sol.x

362

363 # Compute coefficients incrementaly

364 a, b = None, None

365 m0 = 2 if quadType in ['LOBATTO', 'RADAU-LEFT'] else 1

366 for m in range(m0, M + 1):

367 coeffs, nodes = computeCoeffs(m, a, b)

368 if m > 1:

369 a, b = fit(coeffs * m, nodes)

370

371 QDmat[1:, 1:] = np.diag(coeffs)

372 self.parallelizable = True

373

374 elif qd_type == "VDHS":

375 # coefficients from Van der Houwen & Sommeijer, 1991

376

377 m = QDmat.shape[0] - 1

378

379 if m == 4 and coll.node_type == 'LEGENDRE' and coll.quad_type == "RADAU-RIGHT":

380 coeffs = [3055 / 9532, 531 / 5956, 1471 / 8094, 1848 / 7919]

381 else:

382 raise NotImplementedError('no VDHS diagonal coefficients for this node configuration')

383

384 QDmat[1:, 1:] = np.diag(coeffs)

385 self.parallelizable = True

386

387 else:

388 # see if an explicit preconditioner with this name is available

389 try:

390 QDmat = self.get_Qdelta_explicit(coll, qd_type)

391 self.logger.warning(f'Using explicit preconditioner \"{qd_type}\" on the left hand side!')

392 except NotImplementedError:

393 raise NotImplementedError(f'qd_type implicit "{qd_type}" not implemented')

394

395 # check if we got not more than a lower triangular matrix

396 # TODO : this should be a regression test, not run-time ...

397 np.testing.assert_array_equal(

398 np.triu(QDmat, k=1), np.zeros(QDmat.shape), err_msg='Lower triangular matrix expected!'

399 )

400

401 return QDmat

402

403 def get_Qdelta_explicit(self, coll, qd_type):

404 QDmat = np.zeros(coll.Qmat.shape)

405 if qd_type == 'EE':

406 for m in range(self.coll.num_nodes + 1):

407 QDmat[m, 0:m] = self.coll.delta_m[0:m]

408 elif qd_type == 'GS':

409 QDmat = np.tril(self.coll.Qmat, k=-1)

410 elif qd_type == 'PIC':

411 QDmat = np.zeros(coll.Qmat.shape)

412 else:

413 raise NotImplementedError('qd_type explicit not implemented')

414

415 # check if we got not more than a lower triangular matrix

416 np.testing.assert_array_equal(

417 np.triu(QDmat, k=0), np.zeros(QDmat.shape), err_msg='Strictly lower triangular matrix expected!'

418 )

419

420 return QDmat

421

422 def predict(self):

423 """

424 Predictor to fill values at nodes before first sweep

425

426 Default prediction for the sweepers, only copies the values to all collocation nodes

427 and evaluates the RHS of the ODE there

428 """

429

430 # get current level and problem description

431 L = self.level

432 P = L.prob

433

434 # evaluate RHS at left point

435 L.f[0] = P.eval_f(L.u[0], L.time)

436

437 for m in range(1, self.coll.num_nodes + 1):

438 # copy u[0] to all collocation nodes, evaluate RHS

439 if self.params.initial_guess == 'spread':

440 L.u[m] = P.dtype_u(L.u[0])

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

442 # copy u[0] and RHS evaluation to all collocation nodes

443 elif self.params.initial_guess == 'copy':

444 L.u[m] = P.dtype_u(L.u[0])

445 L.f[m] = P.dtype_f(L.f[0])

446 # start with zero everywhere

447 elif self.params.initial_guess == 'zero':

448 L.u[m] = P.dtype_u(init=P.init, val=0.0)

449 L.f[m] = P.dtype_f(init=P.init, val=0.0)

450 # start with random initial guess

451 elif self.params.initial_guess == 'random':

452 L.u[m] = P.dtype_u(init=P.init, val=self.rng.rand(1)[0])

453 L.f[m] = P.dtype_f(init=P.init, val=self.rng.rand(1)[0])

454 else:

455 raise ParameterError(f'initial_guess option {self.params.initial_guess} not implemented')

456

457 # indicate that this level is now ready for sweeps

458 L.status.unlocked = True

459 L.status.updated = True

460

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

462 """

463 Computation of the residual using the collocation matrix Q

464

465 Args:

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

467 """

468

469 # get current level and problem description

470 L = self.level

471

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

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

474 if stage in self.params.skip_residual_computation:

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

476 return None

477

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

479 # assert L.status.updated

480

481 # compute the residual for each node

482

483 # build QF(u)

484 res_norm = []

485 res = self.integrate()

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

487 res[m] += L.u[0] - L.u[m + 1]

488 # add tau if associated

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

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

491 # use abs function from data type here

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

493

494 # find maximal residual over the nodes

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

496 L.status.residual = max(res_norm)

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

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

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

500 L.status.residual = max(res_norm) / abs(L.u[0])

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

502 L.status.residual = res_norm[-1] / abs(L.u[0])

503 else:

504 raise ParameterError(

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

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

507 )

508

509 # indicate that the residual has seen the new values

510 L.status.updated = False

511

512 return None

513

514 def compute_end_point(self):

515 """

516 Abstract interface to end-node computation

517 """

518 raise NotImplementedError('ERROR: sweeper has to implement compute_end_point(self)')

519

520 def integrate(self):

521 """

522 Abstract interface to right-hand side integration

523 """

524 raise NotImplementedError('ERROR: sweeper has to implement integrate(self)')

525

526 def update_nodes(self):

527 """

528 Abstract interface to node update

529 """

530 raise NotImplementedError('ERROR: sweeper has to implement update_nodes(self)')

531

532 @property

533 def level(self):

534 """

535 Returns the current level

536

537 Returns:

538 pySDC.Level.level: the current level

539 """

540 return self.__level

541

542 @level.setter

543 def level(self, L):

544 """

545 Sets a reference to the current level (done in the initialization of the level)

546

547 Args:

548 L (pySDC.Level.level): current level

549 """

550 assert isinstance(L, level)

551 self.__level = L

552

553 @property

554 def rank(self):

555 return 0