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

1 from __future__ import annotations (empty)

2

3 import functools (empty)

4 import typing (empty)

5 import warnings (empty)

6

7 import attrs (empty)

8 import numpy as np (empty)

9 import scipy.optimize as optimize (empty)

10

11 if typing.TYPE_CHECKING: (empty)

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

13

14

15@attrs.define (empty)

16 class 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)

75 class 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 )