Coverage for src/flexfrac1d/lib/phase_shift.py: 98%

1"""Pseudo-scattering parameterisations.""" 

2 

3from __future__ import annotations 

4 

5import abc 

6from numbers import Real 

7import typing 

8 

9import attrs 

10import numpy as np 

11 

12from .constants import PI_2 

13 

14 

15def _seed_rng(seed: int): 

16 return np.random.default_rng(seed) 

17 

18 

19class ScatteringHandler(abc.ABC): 

20 @abc.abstractmethod 

21 def compute_edge_amplitudes( 

22 self, 

23 edge_amplitudes: np.ndarray, 

24 c_wavenumbers: np.ndarray, 

25 xf: np.ndarray, 

26 ) -> np.ndarray: 

27 """Determine post-breakup wave amplitudes at the edge of new floes. 

28 

29 The wave propagates and is attenuated underneath the floe. For the 

30 current timestep, for each wave component, the complex wave amplitude 

31 is fully determined at all the coordinates where fracture is about to 

32 occur. After fractures have occured, the complex amplitudes at these 

33 coordinates become the complex amplitudes at the left edges of new 

34 fragments. Further pseudo-scattering rules can be used if it is not 

35 desirable to keep the wave surface in phase on both sides of a floe 

36 edge. 

37 

38 

39 Parameters 

40 ---------- 

41 edge_amplitudes : np.ndarray of complex 

42 The complex wave amplitudes at the edge of a breaking floe, in m 

43 c_wavenumbers : np.ndarray of complex 

44 The complex wavenumbers stressing the floe, in m^-1 

45 xf : np.ndarray of float 

46 The coordinates of fractures, in m 

47 

48 Returns 

49 ------- 

50 np.ndarray of complex 

51 

52 """ 

53 pass 

54 

55 

56class ContinuousScatteringHandler(ScatteringHandler): 

57 """No scattering. 

58 

59 The surface stays continuous across floes edges. 

60 

61 """ 

62 

63 @staticmethod 

64 def compute_edge_amplitudes( 

65 edge_amplitudes, 

66 c_wavenumbers: np.ndarray, 

67 xf: np.ndarray, 

68 ) -> np.ndarray: 

69 return edge_amplitudes * np.exp(1j * c_wavenumbers * xf[:, None]) 

70 

71 

72@attrs.frozen 

73class UniformScatteringHandler(ScatteringHandler): 

74 r"""Scattering with uniformly sampled new phases. 

75 

76 The wave phase at the edge of a new floe is sampled from the uniform 

77 distribution on :math:`[0; 2\pi)`. 

78 

79 Parameters 

80 ---------- 

81 rng : numpy.random.Generator 

82 Random generator used to sample phases 

83 

84 

85 ------ 

86 )] 

87 

88 """ 

89 

90 rng: np.random.Generator 

91 

92 @classmethod 

93 def from_seed(cls, seed: int) -> typing.Self: 

94 """Instantiate self with an RNG seeded by an integer. 

95 

96 Parameters 

97 ---------- 

98 seed : int 

99 A seed passed to `numpy.random.default_rng` 

100 

101 Returns 

102 ------- 

103 UniformScatteringHandler 

104 

105 """ 

106 return cls(_seed_rng(seed)) 

107 

108 def compute_edge_amplitudes( 

109 self, 

110 edge_amplitudes: np.ndarray, 

111 c_wavenumbers: np.ndarray, 

112 xf: np.ndarray, 

113 ) -> np.ndarray: 

114 phases = self.rng.uniform(0, PI_2, size=edge_amplitudes.shape) 

115 return ( 

116 np.abs(edge_amplitudes) 

117 * np.exp(-np.imag(c_wavenumbers) * xf[:, None]) 

118 * np.exp(1j * phases) 

119 ) 

120 

121 

122@attrs.frozen 

123class PerturbationScatteringHandler(ScatteringHandler): 

124 """Scattering with phases perturbated around the continuous solution. 

125 

126 The wave phase at the left edge of a new floe is computed to maintain 

127 continuity of the surface across the edge. Then, a random perturbation 

128 sampled from a normal distribution, is added to the phase. 

129 

130 Attributes 

131 ---------- 

132 rng : numpy.random.Generator 

133 Random generator used to sample perturbations 

134 loc : Real 

135 The mean of the normal distribution used to sample perturbations, 

136 in rad 

137 scale : Real 

138 The standard deviation of the normal distribution used to sample 

139 perturbations, in rad 

140 

141 Notes 

142 ----- 

143 Perturbations are always added to an existing phase. The expectation of the 

144 resulting phase is thus the sum of `loc` and the phase of the continuous 

145 solution. 

146 

147 """ 

148 

149 rng: np.random.Generator 

150 loc: Real 

151 scale: Real 

152 

153 @classmethod 

154 def from_seed(cls, seed: int, loc: Real = 0, scale: Real = 1) -> typing.Self: 

155 """Instantiate with an RNG seeded with an integer. 

156 

157 Parameters 

158 ---------- 

159 seed : int 

160 A seed passed to `numpy.random.default_rng` 

161 loc : Real 

162 Mean of a normal distribution, in rad 

163 scale : Real 

164 Standard deviation of a normal distribution, in rad 

165 

166 Returns 

167 ------- 

168 typing.Self 

169 [TODO:description] 

170 

171 """ 

172 rng = _seed_rng(seed) 

173 return cls(rng, loc, scale) 

174 

175 def compute_edge_amplitudes( 

176 self, 

177 edge_amplitudes: np.ndarray, 

178 c_wavenumbers: np.ndarray, 

179 xf: np.ndarray, 

180 ) -> np.ndarray: 

181 edge_amplitudes = ContinuousScatteringHandler.compute_edge_amplitudes( 

182 edge_amplitudes, c_wavenumbers, xf 

183 ) 

184 perturbations = self.rng.normal( 

185 self.loc, self.scale, size=edge_amplitudes.shape 

186 ) 

187 edge_amplitudes *= np.exp(1j * perturbations) 

188 return edge_amplitudes