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

1 import numpy as np

2 from .constants import PI_D4, SQR2

3 from .displacement import _dis_hom_coefs, _dis_par_amps

4 from . import numerical

5

6

7 def _egy_hom(floe_params, wave_params):

8 """Energy from the homogen term of the displacement ODE"""

9 red_num, length = floe_params

10 adim = red_num * length

11 adim2 = 2 * adim

12 c_1, c_2, c_3, c_4 = _dis_hom_coefs(floe_params, wave_params)

13

14 return red_num**3 * (

15 +(c_1**2) * (-SQR2 * np.sin(adim2 + PI_D4) + 2 - np.exp(-adim2)) / 2

16 + c_2**2 * (SQR2 * np.sin(adim2 + PI_D4) + 2 - 3 * np.exp(-adim2)) / 2

17 + c_3**2 * ((SQR2 * np.cos(adim2 + PI_D4) - 2) * np.exp(-adim2) + 1) / 2

18 + c_4**2 * (3 - (SQR2 * np.cos(adim2 + PI_D4) + 2) * np.exp(-adim2)) / 2

19 + c_1 * c_2 * (SQR2 * np.cos(adim2 + PI_D4) - np.exp(-adim2))

20 - 2

21 * red_num

22 * c_1

23 * c_3

24 * (2 * length - np.sin(adim2) / red_num)

25 * np.exp(-adim)

26 + 4 * c_1 * c_4 * np.exp(-adim) * np.sin(adim) ** 2

27 + 4 * c_2 * c_3 * np.exp(-adim) * np.sin(adim) ** 2

28 - 2

29 * red_num

30 * c_2

31 * c_4

32 * (2 * length + np.sin(adim2) / red_num)

33 * np.exp(-adim)

34 + c_3 * c_4 * (-1 + SQR2 * np.exp(-adim2) * np.sin(adim2 + PI_D4))

35 )

36

37

38 def _egy_par_vals(red_num, wave_params):

39 comp_amps = _dis_par_amps(red_num, wave_params)

40 _, c_wavenumbers = wave_params

41

42 comp_curvs = comp_amps * (1j * c_wavenumbers) ** 2

43

44 return c_wavenumbers, comp_curvs

45

46

47 def _egy_par_pow2(floe_params, wave_params):

48 """Energy contribution from individual forcings"""

49 red_num, length = floe_params

50 _, c_wavenumbers = wave_params

51 wavenumbers = np.real(c_wavenumbers)

52 attenuations = np.imag(c_wavenumbers)

53

54 comp_wns, comp_curvs = _egy_par_vals(red_num, wave_params)

55 wn_moduli, curv_moduli = map(np.abs, (comp_wns, comp_curvs))

56 wn_phases, curv_phases = map(np.angle, (comp_wns, comp_curvs))

57

58 red = np.exp(-2 * attenuations * length)

59

60 return (

61 curv_moduli**2

62 @ (

63 (1 - red) / attenuations

64 + (

65 np.sin(2 * curv_phases - wn_phases)

66 - np.sin(2 * (wavenumbers * length + curv_phases) - wn_phases) * red

67 )

68 / wn_moduli

69 )

70 ) / 4

71

72

73 def _egy_par_m(floe_params, wave_params):

74 """Energy contribution from forcing interactions"""

75 red_num, length = floe_params

76 _, c_wavenumbers = wave_params

77 wavenumbers = np.real(c_wavenumbers)

78 attenuations = np.imag(c_wavenumbers)

79 _, comp_curvs = _egy_par_vals(red_num, wave_params)

80

81 # Binomial coefficients, much quicker than itertools

82 idx1, idx2 = np.triu_indices(c_wavenumbers.size, 1)

83

84 mean_attenuations = attenuations[idx1] + attenuations[idx2]

85 comp_wns = (

86 wavenumbers[idx1] - wavenumbers[idx2],

87 wavenumbers[idx1] + wavenumbers[idx2],

88 )

89 comp_wns += 1j * mean_attenuations

90

91 curv_moduli = np.abs(comp_curvs)

92 _curv_phases = np.angle(comp_curvs)

93 curv_phases = (

94 _curv_phases[idx1] - _curv_phases[idx2],

95 _curv_phases[idx1] + _curv_phases[idx2],

96 )

97

98 def _f(comp_wns, curv_phases):

99 wn_moduli = np.abs(comp_wns)

100 wn_phases = np.angle(comp_wns)

101 return (

102 np.sin(curv_phases - wn_phases)

103 - (

104 np.exp(-np.imag(comp_wns) * length)

105 * np.sin(np.real(comp_wns) * length + curv_phases - wn_phases)

106 )

107 ) / wn_moduli

108

109 return (

110 curv_moduli[idx1]

111 * curv_moduli[idx2]

112 @ (_f(comp_wns[1], curv_phases[1]) - _f(comp_wns[0], curv_phases[0]))

113 )

114

115

116 def _egy_par(floe_params, wave_params):

117 """Energy from the forcing term of the displacement ODE"""

118 return _egy_par_pow2(floe_params, wave_params) + _egy_par_m(

119 floe_params, wave_params

120 )

121

122

123 def _egy_m(floe_params, wave_params):

124 red_num, length = floe_params

125 _, c_wavenumbers = wave_params

126 adim = red_num * length

127 wavenumbers = np.real(c_wavenumbers)

128 attenuations = np.imag(c_wavenumbers)

129

130 wn_abs, wn_phases = (_f(c_wavenumbers) for _f in (np.abs, np.angle))

131 _, comp_curvs = _egy_par_vals(red_num, wave_params)

132 curv_moduli, curv_phases = np.abs(comp_curvs), np.angle(comp_curvs)

133 num_add = red_num + wavenumbers

134 num_sub = red_num - wavenumbers

135 att_add = red_num + attenuations

136 att_sub = red_num - attenuations

137 q_pp = att_add**2 + num_add**2

138 q_mp = att_sub**2 + num_add**2

139 q_pm = att_add**2 + num_sub**2

140 q_mm = att_sub**2 + num_sub**2

141

142 phase_deltas = wn_phases - curv_phases

143 sin_phase_deltas = np.sin(phase_deltas)

144 cos_phase_deltas = np.cos(phase_deltas)

145 phase_quads = curv_phases + PI_D4

146 sin_phase_quads = np.sin(phase_quads)

147 cos_phase_quads = np.cos(phase_quads)

148 energ_1_K = np.array(

149 (

150 sin_phase_deltas * (1 / q_mp - 1 / q_mm),

151 cos_phase_deltas * (1 / q_mp + 1 / q_mm),

152 -sin_phase_deltas * (1 / q_pp - 1 / q_pm),

153 -cos_phase_deltas * (1 / q_pp + 1 / q_pm),

154 )

155 )

156 energ_1_b = np.array(

157 (

158 -(sin_phase_quads / q_mp - cos_phase_quads / q_mm),

159 (cos_phase_quads / q_mp - sin_phase_quads / q_mm),

160 -(cos_phase_quads / q_pp - sin_phase_quads / q_pm),

161 -(sin_phase_quads / q_pp - cos_phase_quads / q_pm),

162 )

163 )

164 energ_1 = wn_abs * energ_1_K + SQR2 * red_num * energ_1_b

165 energ_1[:2, :] *= np.exp(-adim)

166

167 arg_add = num_add * length

168 arg_add_phase = arg_add - phase_deltas

169 sin_arg_ap = np.sin(arg_add_phase)

170 cos_arg_ap = np.cos(arg_add_phase)

171 arg_sub = num_sub * length

172 arg_sub_phase = arg_sub + phase_deltas

173 sin_arg_sp = np.sin(arg_sub_phase)

174 cos_arg_sp = np.cos(arg_sub_phase)

175 arg_add_quad = arg_add + phase_quads

176 sin_arg_aq = np.sin(arg_add_quad)

177 cos_arg_aq = np.cos(arg_add_quad)

178 arg_sub_quad = arg_sub - curv_phases + PI_D4

179 sin_arg_sq = np.sin(arg_sub_quad)

180 cos_arg_sq = np.cos(arg_sub_quad)

181 energ_exp_K = np.array(

182 (

183 sin_arg_ap / q_mp + sin_arg_sp / q_mm,

184 -(cos_arg_ap / q_mp + cos_arg_sp / q_mm),

185 -(sin_arg_ap / q_pp + sin_arg_sp / q_pm),

186 cos_arg_ap / q_pp + cos_arg_sp / q_pm,

187 )

188 )

189 energ_exp_b = np.array(

190 (

191 sin_arg_aq / q_mp - sin_arg_sq / q_mm,

192 -cos_arg_aq / q_mp + cos_arg_sq / q_mm,

193 cos_arg_aq / q_pp - cos_arg_sq / q_pm,

194 sin_arg_aq / q_pp - sin_arg_sq / q_pm,

195 )

196 )

197 energ_exp = wn_abs * energ_exp_K + SQR2 * red_num * energ_exp_b

198 energ_exp[2:, :] *= np.exp(-adim)

199 energ_exp *= np.exp(-attenuations * length)

200

201 energ = (

202 (energ_1 + energ_exp) @ curv_moduli @ _dis_hom_coefs(floe_params, wave_params)

203 )

204 return red_num**2 * energ

205

206

207 def energy(

208 floe_params: tuple[float],

209 wave_params: tuple[np.ndarray],

210 growth_params: tuple | None = None,

211 an_sol: bool | None = None,

212 num_params: dict | None = None,

213) -> float:

214 if numerical._use_an_sol(an_sol, floe_params[1], growth_params): 214 ↛ 220line 214 didn't jump to line 220, because the condition on line 214 was never false

215 return (

216 _egy_hom(floe_params, wave_params)

217 + 2 * _egy_m(floe_params, wave_params)

218 + _egy_par(floe_params, wave_params)

219 )

220 return numerical.energy(floe_params, wave_params, growth_params, num_params)[0]