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

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

2

3 from __future__ import annotations

4

5 import abc

6 from numbers import Real

7 import typing

8

9 import attrs

10 import numpy as np

11

12 from .constants import PI_2

13

14

15 def _seed_rng(seed: int):

16 return np.random.default_rng(seed)

17

18

19 class 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

56 class 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

73 class 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

123 class 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