Coverage for pySDC/core/sweeper.py: 94%

1import logging

2from typing import Any, Dict, Optional, TYPE_CHECKING

3import numpy as np

4from qmat.qdelta import QDeltaGenerator, QDELTA_GENERATORS

6from pySDC.core.errors import ParameterError

7from pySDC.core.collocation import CollBase

8from pySDC.helpers.pysdc_helper import FrozenClass

10if TYPE_CHECKING:

11 from pySDC.core.level import Level

13# Organize QDeltaGenerator class in dict[type(QDeltaGenerator),set(str)] to retrieve aliases

14QDELTA_GENERATORS_ALIASES = {v: set() for v in set(QDELTA_GENERATORS.values())}

15for k, v in QDELTA_GENERATORS.items():

16 QDELTA_GENERATORS_ALIASES[v].add(k)

19# short helper class to add params as attributes

20class _Pars(FrozenClass):

21 def __init__(self, pars: Dict[str, Any]) -> None:

22 self.do_coll_update: bool = False

23 self.initial_guess: str = 'spread' # default value (see also below)

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

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

27 if k != 'collocation_class':

28 setattr(self, k, v)

30 self._freeze()

33class Sweeper(object):

34 """

35 Base abstract sweeper class, provides two base methods to generate QDelta matrices:

37 - get_Qdelta_implicit(qd_type):

38 Returns a (pySDC-type) QDelta matrix of **implicit type**,

39 *i.e* lower triangular with zeros on the first collumn.

40 - get_Qdelta_explicit(qd_type):

41 Returns a (pySDC-type) QDelta matrix of **explicit type**,

42 *i.e* strictly lower triangular with first node distance to zero on the first collumn.

45 All possible QDelta matrix coefficients are generated with

46 `qmat <https://qmat.readthedocs.io/en/latest/api/qmat/qdelta/index.html>`_,

47 check it out to see all available coefficient types.

49 Attributes:

50 logger: custom logger for sweeper-related logging

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

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

53 """

55 def __init__(self, params: Dict[str, Any], level: 'Level') -> None:

56 """

57 Initialization routine for the base sweeper

59 Args:

60 params (dict): parameter object

61 level (pySDC.Level.level): the level that uses this sweeper

62 """

64 self.logger: logging.Logger = logging.getLogger('sweeper')

66 essential_keys = ['num_nodes']

67 for key in essential_keys:

68 if key not in params:

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

70 self.logger.error(msg)

71 raise ParameterError(msg)

73 if 'collocation_class' not in params:

74 params['collocation_class'] = CollBase

76 # prepare random generator for initial guess

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

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

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

81 self.params = _Pars(params)

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

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

86 self.logger.warning(

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

88 )

89 self.params.do_coll_update = True

91 self.__level = level

92 self.parallelizable = False

93 self._Qdelta_diagonal = {}

94 for name in ["genQI", "genQE"]:

95 if hasattr(self, name):

96 delattr(self, name)

98 def buildGenerator(self, qdType: str) -> QDeltaGenerator:

99 """

100 Build the qmat generator of QDelta coefficients for the collocation nodes of this sweeper.

101

102 Args:

103 qdType (str): Name or alias of the QDelta type, a key of `qmat.qdelta.QDELTA_GENERATORS`

104

105 Returns:

106 qmat.qdelta.QDeltaGenerator: Generator for these nodes and left interval boundary

107 """

108 return QDELTA_GENERATORS[qdType](qGen=self.coll.generator, tLeft=self.coll.tleft)

109

110 def get_Qdelta_implicit(self, qd_type: str, k: Optional[int] = None) -> np.ndarray:

111 """

112 Get a QDelta matrix of implicit type, i.e. lower triangular with zeros in the first row and column.

113

114 The generator is cached as `self.genQI` and only rebuilt if `qd_type` is not an alias of the cached one. Sets

115 `self.parallelizable`, see `_update_parallelizable`.

116

117 Args:

118 qd_type (str): Name or alias of the QDelta type

119 k (int, optional): Sweep index, for QDelta types whose coefficients change with the sweep

120

121 Returns:

122 numpy.ndarray: QDelta matrix of shape (num_nodes + 1, num_nodes + 1)

123 """

124 QDmat = np.zeros_like(self.coll.Qmat)

125 if not hasattr(self, "genQI") or qd_type not in QDELTA_GENERATORS_ALIASES[type(self.genQI)]:

126 self.genQI: QDeltaGenerator = self.buildGenerator(qd_type)

127 QDmat[1:, 1:] = self.genQI.genCoeffs(k=k)

128

129 err_msg = 'Lower triangular matrix expected!'

130 np.testing.assert_array_equal(np.triu(QDmat, k=1), np.zeros(QDmat.shape), err_msg=err_msg)

131 self._update_parallelizable(('implicit', type(self.genQI)), QDmat)

132 return QDmat

133

134 def get_Qdelta_explicit(self, qd_type: str, k: Optional[int] = None) -> np.ndarray:

135 """

136 Get a QDelta matrix of explicit type, i.e. strictly lower triangular, with the distances of the nodes to the

137 left

138 interval boundary in the first column and zeros in the first row.

139

140 The generator is cached as `self.genQE` and only rebuilt if `qd_type` is not an alias of the cached one. Sets

141 `self.parallelizable`, see `_update_parallelizable`.

142

143 Args:

144 qd_type (str): Name or alias of the QDelta type

145 k (int, optional): Sweep index, for QDelta types whose coefficients change with the sweep

146

147 Returns:

148 numpy.ndarray: QDelta matrix of shape (num_nodes + 1, num_nodes + 1)

149 """

150 coll = self.coll

151 QDmat = np.zeros(coll.Qmat.shape, dtype=float)

152 if not hasattr(self, "genQE") or qd_type not in QDELTA_GENERATORS_ALIASES[type(self.genQE)]:

153 self.genQE: QDeltaGenerator = self.buildGenerator(qd_type)

154 QDmat[1:, 1:], QDmat[1:, 0] = self.genQE.genCoeffs(k=k, dTau=True)

155

156 err_msg = 'Strictly lower triangular matrix expected!'

157 np.testing.assert_array_equal(np.triu(QDmat, k=0), np.zeros(QDmat.shape), err_msg=err_msg)

158 self._update_parallelizable(('explicit', type(self.genQE)), QDmat)

159 return QDmat

160

161 def _update_parallelizable(self, key, QDmat):

162 """

163 The sweeper is parallelizable if every QDelta matrix it has built is diagonal, which for explicit ones means

164 zero (PIC). Matrices are remembered by kind and generator, so rebuilding one, for instance with a new sweep

165 index, replaces its entry, and a sweeper with QI='MIN-SR-S' and QE='EE' is not parallelizable.

166 """

167 self._Qdelta_diagonal[key] = np.allclose(np.diag(np.diag(QDmat)), QDmat)

168 self.parallelizable = all(self._Qdelta_diagonal.values())

169

170 def predict(self) -> None:

171 """

172 Predictor to fill values at nodes before first sweep

173

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

175 and evaluates the RHS of the ODE there

176 """

177

178 # get current level and problem description

179 L = self.level

180 P = L.prob

181

182 # evaluate RHS at left point

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

184

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

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

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

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

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

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

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

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

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

194 # start with zero everywhere

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

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

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

198 # start with random initial guess

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

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

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

202 else:

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

204

205 # indicate that this level is now ready for sweeps

206 L.status.unlocked = True

207 L.status.updated = True

208

209 def compute_residual(self, stage: str = '') -> None:

210 """

211 Computation of the residual using the collocation matrix Q

212

213 Args:

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

215 """

216

217 # get current level and problem description

218 L = self.level

219

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

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

222 if stage in self.params.skip_residual_computation:

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

224 return None

225

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

227 # assert L.status.updated

228

229 # compute the residual for each node

230

231 # build QF(u)

232 res_norm = []

233 L.residual = self.integrate()

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

235 L.residual[m] += L.u[0] - L.u[m + 1]

236 # add tau if associated

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

238 L.residual[m] += L.tau[m]

239 # use abs function from data type here

240 res_norm.append(abs(L.residual[m]))

241

242 # find maximal residual over the nodes

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

244 L.status.residual = max(res_norm)

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

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

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

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

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

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

251 else:

252 raise ParameterError(

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

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

255 )

256

257 # indicate that the residual has seen the new values

258 L.status.updated = False

259

260 return None

261

262 def compute_end_point(self) -> None:

263 """

264 Abstract interface to end-node computation

265 """

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

267

268 def integrate(self) -> Any:

269 """

270 Abstract interface to right-hand side integration

271 """

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

273

274 def update_nodes(self) -> None:

275 """

276 Abstract interface to node update

277 """

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

279

280 @property

281 def level(self) -> 'Level':

282 """

283 Returns the current level

284

285 Returns:

286 pySDC.Level.level: the current level

287 """

288 return self.__level

289

290 @level.setter

291 def level(self, L: 'Level') -> None:

292 """

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

294

295 Args:

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

297 """

298 from pySDC.core.level import Level

299

300 assert isinstance(L, Level)

301 self.__level = L

302

303 @property

304 def rank(self) -> int:

305 """Rank of this sweeper in the communicator across nodes: always 0 here, overridden by the MPI sweepers."""

306 return 0

307

308 def updateVariableCoeffs(self, k: int) -> None:

309 """

310 Potentially update QDelta implicit coefficients if variable ...

311

312 Parameters

313 ----------

314 k : int

315 Index of the sweep (0 for initial sweep, 1 for the first one, ...).

316 """

317 if hasattr(self, "genQI") and self.genQI.isKDependent():

318 qdType = type(self.genQI).__name__

319 self.QI = self.get_Qdelta_implicit(qdType, k=k)

320 if hasattr(self, "genQE") and self.genQE.isKDependent():

321 qdType = type(self.genQE).__name__

322 self.QE = self.get_Qdelta_explicit(qdType, k=k)