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

1 import functools

2 import warnings

3

4 import attrs

5 import numpy as np

6 import 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

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

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