Coverage for src/flexfrac1d/lib/dr.py: 94%

1import functools 

2import warnings 

3 

4import attrs 

5import numpy as np 

6import scipy.optimize as optimize 

7 

8# TODO: add from_ocean, from_ice class methods, use them model.Domain where 

9# relevant 

10 

11 

12@attrs.define 

13class FreeSurfaceSolver: 

14 alphas: np.ndarray 

15 depth: float 

16 

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

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

19 # admitting one positive real root. 

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

21 

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

23 # Derivative of dr with respect to kk. 

24 tt = np.tanh(kk) 

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

26 

27 def find_k(self, k0, alpha): 

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

29 if not res.converged: 29 ↛ 30line 29 didn't jump to line 30, because the condition on line 29 was never true

30 warnings.warn( 

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

32 stacklevel=2, 

33 ) 

34 return res.root 

35 

36 def compute_wavenumbers(self) -> np.ndarray: 

37 if np.isposinf(self.depth): 

38 return self.alphas 

39 

40 coefs_d0 = self.alphas * self.depth 

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

42 for i, _d0 in enumerate(coefs_d0): 

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

44 roots[i] = _d0 

45 continue 

46 

47 find_k_i = functools.partial( 

48 self.find_k, 

49 alpha=_d0, 

50 ) 

51 roots[i] = find_k_i(_d0) 

52 

53 return roots / self.depth 

54 

55 

56@attrs.define 

57class ElasticMassLoadingSolver: 

58 alphas: np.ndarray 

59 deg1: np.ndarray 

60 deg0: np.ndarray 

61 scaled_ratio: float 

62 

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

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

65 

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

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

68 

69 def extract_real_root(self, roots): 

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

71 if mask.nonzero()[0].size != 1: 71 ↛ 72line 71 didn't jump to line 72, because the condition on line 71 was never true

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

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

74 

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

76 res = optimize.root_scalar( 

77 self.f, 

78 args=(d0, d1, rr), 

79 fprime=self.df_dk, 

80 x0=k0, 

81 xtol=1e-10, 

82 ) 

83 if not res.converged: 83 ↛ 84line 83 didn't jump to line 84, because the condition on line 83 was never true

84 warnings.warn( 

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

86 stacklevel=2, 

87 ) 

88 return res.root 

89 

90 def compute_wavenumbers(self) -> np.ndarray: 

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

92 

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

94 find_k_i = functools.partial( 

95 self.find_k, 

96 alpha=alpha, 

97 d0=_d0, 

98 d1=_d1, 

99 rr=self.scaled_ratio, 

100 ) 

101 

102 # We always expect one positive real root, 

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

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

105 k0_dw = self.extract_real_root(roots_dw) 

106 if np.isposinf(self.scaled_ratio): 

107 roots[i] = k0_dw 

108 continue 

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

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

111 thrsld_dw, thrsld_sw = 1.33, 0.79 

112 if self.scaled_ratio * k0_dw > thrsld_dw: 

113 roots[i] = find_k_i(k0_dw) 

114 else: 

115 roots_sw = np.polynomial.polynomial.polyroots( 

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

117 ) 

118 k0_sw = self.extract_real_root(roots_sw) 

119 

120 if self.scaled_ratio * k0_sw < thrsld_sw: 

121 roots[i] = find_k_i(k0_sw) 

122 # Use an initial guess in the middle otherwise 

123 else: 

124 k0_ = (k0_sw + k0_dw) / 2 

125 roots[i] = find_k_i(k0_) 

126 

127 return roots