Coverage for pySDC/implementations/convergence_controller_classes/crash.py: 100%

1from pySDC.core.convergence_controller import ConvergenceController 

2from pySDC.core.errors import ConvergenceError 

3import numpy as np 

4import time 

5 

6 

7class CrashBase(ConvergenceController): 

8 """ 

9 Base class for convergence controllers that raise a `ConvergenceError` on all ranks as soon as one rank crashes. 

10 """ 

11 

12 comm = None #: communicator across the sweeper, which subclasses set up in `setup` 

13 

14 def __init__(self, controller, params, description, **kwargs): 

15 """ 

16 Initialization routine, which additionally prepares the logical or for communicating crashes with MPI 

17 

18 Args: 

19 controller (pySDC.Controller): The controller 

20 params (dict): The params passed for this specific convergence controller 

21 description (dict): The description object used to instantiate the controller 

22 """ 

23 super().__init__(controller, params, description, **kwargs) 

24 if self.comm or self.params.useMPI: 

25 from mpi4py import MPI 

26 

27 self.MPI_OR = MPI.LOR 

28 

29 def communicate_crash(self, crash, msg='', comm=None, **kwargs): 

30 """ 

31 Communicate a crash across all ranks and raise an error if so. 

32 

33 Args: 

34 crash (bool): If this rank wants to crash 

35 msg (str): Message of the `ConvergenceError` raised on a crash 

36 comm (mpi4py.MPI.Intracomm or None): Communicator of the controller, if applicable: 

37 """ 

38 

39 # communicate across the sweeper 

40 if self.comm: 

41 crash = self.comm.allreduce(crash, op=self.MPI_OR) 

42 

43 # communicate across the steps 

44 if comm: 

45 crash = comm.allreduce(crash, op=self.MPI_OR) 

46 

47 if crash: 

48 raise ConvergenceError(msg) 

49 

50 

51class StopAtNan(CrashBase): 

52 """ 

53 Crash all ranks when the solution contains nan or inf, or when its norm reaches the optional threshold `thresh`. 

54 This class is useful when running with MPI in the sweeper or controller. 

55 """ 

56 

57 def setup(self, controller, params, description, **kwargs): 

58 """ 

59 Define parameters here. 

60 

61 Default parameters are: 

62 - thresh (float): Crash the code when the norm of the solution exceeds this threshold 

63 

64 Args: 

65 controller (pySDC.Controller): The controller 

66 params (dict): The params passed for this specific convergence controller 

67 description (dict): The description object used to instantiate the controller 

68 

69 Returns: 

70 (dict): The updated params dictionary 

71 """ 

72 self.comm = description['sweeper_params'].get('comm', None) 

73 defaults = { 

74 "control_order": 94, 

75 "thresh": np.inf, 

76 } 

77 

78 return {**defaults, **super().setup(controller, params, description, **kwargs)} 

79 

80 def prepare_next_block(self, controller, S, *args, **kwargs): 

81 """ 

82 Check if we need to crash the code. 

83 

84 Args: 

85 controller (pySDC.Controller.controller): Controller 

86 S (pySDC.Step.step): Step 

87 comm (mpi4py.MPI.Intracomm or None): Communicator of the controller, if applicable 

88 

89 Raises: 

90 ConvergenceError: If the solution does not fall within the allowed space 

91 """ 

92 isfinite, below_limit = True, True 

93 crash = False 

94 

95 for lvl in S.levels: 

96 for u in lvl.u: 

97 if u is None: 

98 break 

99 isfinite = np.all(np.isfinite(u)) 

100 

101 below_limit = abs(u) < self.params.thresh 

102 

103 crash = not (isfinite and below_limit) 

104 

105 if crash: 

106 break 

107 if crash: 

108 break 

109 

110 self.communicate_crash(crash, msg=f'Solution exceeds bounds! Crashing code at {S.time}!', **kwargs) 

111 

112 

113class StopAtMaxRuntime(CrashBase): 

114 """ 

115 Abort the code when the problem has exceeded a maximum runtime. 

116 """ 

117 

118 def setup(self, controller, params, description, **kwargs): 

119 """ 

120 Define parameters here. 

121 

122 Default parameters are: 

123 - max_runtime (float): Crash the code when the wall time since this setup exceeds this many seconds 

124 

125 Args: 

126 controller (pySDC.Controller): The controller 

127 params (dict): The params passed for this specific convergence controller 

128 description (dict): The description object used to instantiate the controller 

129 

130 Returns: 

131 (dict): The updated params dictionary 

132 """ 

133 self.comm = description['sweeper_params'].get('comm', None) 

134 defaults = { 

135 "control_order": 94, 

136 "max_runtime": np.inf, 

137 } 

138 self.t0 = time.perf_counter() 

139 

140 return {**defaults, **super().setup(controller, params, description, **kwargs)} 

141 

142 def prepare_next_block(self, controller, S, *args, **kwargs): 

143 """ 

144 Check if we need to crash the code. 

145 

146 Args: 

147 controller (pySDC.Controller.controller): Controller 

148 S (pySDC.Step.step): Step 

149 comm (mpi4py.MPI.Intracomm or None): Communicator of the controller, if applicable 

150 

151 Raises: 

152 ConvergenceError: If the solution does not fall within the allowed space 

153 """ 

154 self.communicate_crash( 

155 crash=abs(self.t0 - time.perf_counter()) > self.params.max_runtime, 

156 msg=f'Exceeding max. runtime of {self.params.max_runtime}s! Crashing code at {S.time}!', 

157 **kwargs, 

158 )