1
2 """Solve a 1-dimensional equation to find the roots.
3 This does alternating secant and bisection steps.
4 It will work for any function, discontinuous or not.
5 For a nasty function, it might be a factor of 2 slower
6 than bisection, but for a nearly-linear function,
7 it can be much faster than bisection.
8 """
9
10
11 -def _root1(f, xl, xh, p, fl, fh, epsx, epsf):
12
13 assert epsx > 0.0
14 assert epsf >= 0.0
15 assert xh > xl
16 assert (fh>0) != (fl>0)
17 if xh-xl <= epsx:
18 assert xl <= xh, "Not xl=%.15g <= xh=%.15g" % (xl, xh)
19 if fl-fh == 0.0:
20 tmp = xl + 0.5*(xh-xl)
21 else:
22 tmp = xl + (xl-xh)*fl/(fh-fl)
23 if not (xl < tmp < xh):
24 tmp = xl + 0.5*(xh-xl)
25 return tmp
26 elif abs(fl) <= epsf:
27 return xl
28 elif abs(fh) <= epsf:
29 return xh
30
31 assert xl < xh
32 q = [ (xl, fl), (xh, fh) ]
33
34
35 assert (fh>0) != (fl>0) and fh!=fl
36 xsec = (xl*fh-xh*fl)/(fh-fl)
37 if xl < xsec < xh:
38 fsec = f(xsec, p)
39 if fsec == 0:
40 return xsec
41 assert fsec>0 or fsec<0
42 q.append( (xsec, fsec) )
43
44 if (fsec>0) == (fl>0):
45 xbi = xsec + (xh-xsec)*0.5
46 else:
47 xbi = xsec + (xl-xsec)*0.5
48 else:
49 xbi = xl + 0.5*(xh - xl)
50
51
52
53 assert xl <= xbi <= xh, "Not xl=%.15g <= xbi=%.15g <= xh=%.15g" % (xl, xbi, xh)
54 if xl < xbi < xh and xbi!=xsec:
55 fbi = f(xbi, p)
56 if fbi == 0:
57 return xbi
58 assert fbi>0 or fbi<0
59 q.append( (xbi, fbi) )
60
61 if len(q) <= 2:
62
63
64 return 0.5*(xl+xh)
65
66 q.sort()
67
68 X = 0
69 F = 1
70
71 for i in range(1, len(q)):
72 assert q[i][X] != q[i-1][X]
73 if (q[i][F]>0) != (q[i-1][F]>0):
74 return _root1(f, q[i-1][X], q[i][X], p, q[i-1][F], q[i][F], epsx, epsf)
75 raise RuntimeError, "Lost a root: %s" % repr(q)
76
77
78
79 -def root(f, xl, xh, p, epsx, epsf):
80 """Find a root of an equation.
81 @param f: is the function to solve, called as f(x, p).
82 @param p: is arbitrary data for f.
83 This function assumes there is known to be a root in [xl,xh] and that
84 f(xl,p) has a different sign from f(xh,p).
85 It finds it within a region of width epsx, or
86 at least finds an x such that -epsf < f(x, p) < epsf .
87 @return: a root in the interval
88 @rtype: float
89 """
90 if xh < xl:
91 return None
92 fl = f(xl, p)
93 if fl == 0:
94 return xl
95 fh = f(xh, p)
96 if fh == 0:
97 return xh
98 if xh == xl:
99 return None
100
101
102 assert (fl>0) != (fh>0), "0 or an even number of roots in initial x-range."
103 return _root1(f, xl, xh, p, fl, fh, epsx, epsf)
104
105
108
110 q = root(_test1f, 1, 1000, None, 1e-6, 0.01)
111 assert _test1f(q, None) == 0
112 q = root(_test1f, 1, 16, None, 1e-6, 0.01)
113 assert _test1f(q, None) == 0
114
115
116 -def iroot(y, xl, xh, epsy=0.0):
117 """Find a zero in an array y.
118 This function assumes there is known to be a root in [xl,xh].
119 It returns a real-number index into the array which
120 linearly interpolates to zero.
121 """
122 import math
123 assert 0 <= xl < y.shape[0]
124 assert 0 <= xh < y.shape[0]
125
126 def interp(x, y):
127 xi = math.floor(x)
128 frac = x - xi
129 i = int(xi)
130 if frac > 0:
131 return y[i] + (y[i+1]-y[i])*frac
132 return y[i]
133
134 return root(interp, xl, xh, y, 0.001, epsy)
135
136
138 import Num
139 y = 0.64234*(Num.arange(100) - 32.234)
140 assert abs(iroot(y, 10, 50) - 32.234) < 0.002
141
142 if __name__ == '__main__':
143 test()
144 testi()
145