Coverage for src/swiift/lib/dr.py: 28%

1from __future__ import annotations (empty)

2 

3import functools (empty)

4import typing (empty)

5import warnings (empty)

6 

7import attrs (empty)

8import numpy as np (empty)

9import scipy.optimize as optimize (empty)

10 

11if typing.TYPE_CHECKING: (empty)

12 from ..model.model import DiscreteSpectrum, FloatingIce, Ocean 

13 

14 

15@attrs.define (empty)

16class FreeSurfaceSolver: (empty)

17 alphas: np.ndarray (empty)

18 depth: float (empty)

19 

20 @classmethod (empty)

21 def from_ocean(cls, ocean: Ocean, spectrum: DiscreteSpectrum, gravity: float): (empty)

22 alphas = spectrum._ang_freqs_pow2 / gravity 

23 return cls(alphas, ocean.depth) 

24 

25 def f(self, kk: float, alpha: float) -> float: (empty)

26 # Dispersion relation (form f(k) = 0), for a free surface, 

27 # admitting one positive real root. 

28 return kk * np.tanh(kk) - alpha 

29 

30 def df_dk(self, kk: float, alpha: float) -> float: (empty)

31 # Derivative of dr with respect to kk. 

32 tt = np.tanh(kk) 

33 return tt + kk * (1 - tt**2) 

34 

35 def find_k(self, k0, alpha): (empty)

36 res = optimize.root_scalar(self.f, args=(alpha,), fprime=self.df_dk, x0=alpha) 

37 if not res.converged: 

38 warnings.warn( 

39 "Root finding did not converge: free surface", 

40 stacklevel=2, 

41 ) 

42 return res.root 

43 

44 def _compute_real_wavenumbers(self): (empty)

45 if np.isposinf(self.depth): 

46 return self.alphas 

47 

48 coefs_d0 = self.alphas * self.depth 

49 roots = np.full(len(coefs_d0), np.nan) 

50 for i, _d0 in enumerate(coefs_d0): 

51 if _d0 >= np.arctanh(np.nextafter(1, 0)): 

52 roots[i] = _d0 

53 continue 

54 

55 find_k_i = functools.partial( 

56 self.find_k, 

57 alpha=_d0, 

58 ) 

59 roots[i] = find_k_i(_d0) 

60 

61 return roots / self.depth 

62 

63 def compute_wavenumbers( (empty)

64 self, real: bool = True, n: int | None = None 

65 ) -> np.ndarray: 

66 if real: 

67 return self._compute_real_wavenumbers() 

68 else: 

69 raise NotImplementedError( 

70 "Solving imaginary wavenumbers has not been implemented yet." 

71 ) 

72 

73 

74@attrs.define (empty)

75class ElasticMassLoadingSolver: (empty)

76 alphas: np.ndarray (empty)

77 deg1: np.ndarray (empty)

78 deg0: np.ndarray (empty)

79 scaled_ratio: float (empty)

80 

81 @classmethod (empty)

82 def from_floating( (empty)

83 cls, 

84 ice: FloatingIce, 

85 spectrum: DiscreteSpectrum, 

86 gravity: float, 

87 ): 

88 alphas = spectrum._ang_freqs_pow2 / gravity 

89 deg1 = 1 - alphas * ice.draft 

90 deg0 = -alphas * ice.elastic_length 

91 scaled_ratio = ice.dud / ice.elastic_length 

92 return cls(alphas, deg1, deg0, scaled_ratio) 

93 

94 def f(self, kk: float, d0: float, d1: float, rr: float) -> float: (empty)

95 return (kk**5 + d1 * kk) * np.tanh(rr * kk) + d0 

96 

97 def df_dk(self, kk: float, d0: float, d1: float, rr: float) -> float: (empty)

98 return (5 * kk**4 + d1 + rr * d0) * np.tanh(rr * kk) + rr * (kk**5 + d1 * kk) 

99 

100 def extract_real_root(self, roots): (empty)

101 mask = (np.imag(roots) == 0) & (np.real(roots) > 0) 

102 if mask.nonzero()[0].size != 1: 

103 raise ValueError("An approximate initial guess could not be found") 

104 return np.real(roots[mask][0]) 

105 

106 def find_k(self, k0, alpha, d0, d1, rr): (empty)

107 res = optimize.root_scalar( 

108 self.f, 

109 args=(d0, d1, rr), 

110 fprime=self.df_dk, 

111 x0=k0, 

112 xtol=1e-10, 

113 ) 

114 if not res.converged: 

115 warnings.warn( 

116 "Root finding did not converge: ice-covered surface", 

117 stacklevel=2, 

118 ) 

119 return res.root 

120 

121 def _compute_real_wavenumbers(self): (empty)

122 roots = np.full(self.alphas.size, np.nan) 

123 

124 for i, (alpha, _d0, _d1) in enumerate(zip(self.alphas, self.deg0, self.deg1)): 

125 find_k_i = functools.partial( 

126 self.find_k, 

127 alpha=alpha, 

128 d0=_d0, 

129 d1=_d1, 

130 rr=self.scaled_ratio, 

131 ) 

132 

133 # We always expect one positive real root, 

134 # and if _d1 < 0, eventually two additional negative real roots. 

135 roots_dw = np.polynomial.polynomial.polyroots([_d0, _d1, 0, 0, 0, 1]) 

136 k0_dw = self.extract_real_root(roots_dw) 

137 if np.isposinf(self.scaled_ratio): 

138 roots[i] = k0_dw 

139 continue 

140 # Use a DW initial guess if |1-1/tanh(rr*k_DW)| < 0.15 

141 # Use a SW initial guess if |1-rr*k_SW/tanh(rr*k_SW)| < 0.20 

142 thrsld_dw, thrsld_sw = 1.33, 0.79 

143 if self.scaled_ratio * k0_dw > thrsld_dw: 

144 roots[i] = find_k_i(k0_dw) 

145 else: 

146 roots_sw = np.polynomial.polynomial.polyroots( 

147 [_d0 / self.scaled_ratio, 0, _d1, 0, 0, 0, 1] 

148 ) 

149 k0_sw = self.extract_real_root(roots_sw) 

150 

151 if self.scaled_ratio * k0_sw < thrsld_sw: 

152 roots[i] = find_k_i(k0_sw) 

153 # Use an initial guess in the middle otherwise 

154 else: 

155 k0_ = (k0_sw + k0_dw) / 2 

156 roots[i] = find_k_i(k0_) 

157 

158 # -d_0 / alpha == L_D 

159 return roots / (-self.deg0 / self.alphas) 

160 

161 def compute_wavenumbers( (empty)

162 self, real: bool = True, n: int | None = None 

163 ) -> np.ndarray: 

164 if real: 

165 return self._compute_real_wavenumbers() 

166 else: 

167 raise NotImplementedError( 

168 "Solving complex wavenumbers has not been implemented yet." 

169 )