Aster/Cata/Utilitai/funct_root.py

   1 #@ MODIF funct_root Utilitai  DATE 14/09/2004   AUTEUR MCOURTOI M.COURTOIS
   2 # -*- coding: iso-8859-1 -*-
   3 ################################################################################
   4 #       Mathematical utility routines
   5 #       Copyright (C) 1999, Wesley Phoa
   6 #
   7 #       Reference: Numerical Recipes in C
   8 #       [[[[extraits]]]]
   9
  10 class BracketingException(Exception):
  11         pass
  12
  13 class RootFindingException(Exception):
  14         pass
  15
  16 class MinimizationException(Exception):
  17         pass
  18
  19 GOLDEN = (1+5**.5)/2
  20
  21 #
  22 # MISCELLANEOUS
  23 #
  24
  25 def sgn(x):
  26         if x==0:
  27                 return 0
  28         else:
  29                 return x/abs(x)
  30
  31 #
  32 # UNIVARIATE ROOT FINDING
  33 #
  34
  35 def bracket_root(f, interval, max_iterations=50):
  36         """\
  37 Given a univariate function f and a tuple interval=(x1,x2),
  38 return a new tuple (bracket, fnvals) where bracket=(x1,x2)
  39 brackets a root of f and fnvals=(f(x1),f(x2)).
  40         """
  41         x1, x2 = interval
  42         if x1==x2:
  43                  raise BracketingException("initial interval has zero width")
  44         elif x2<x1:
  45                 x1, x2 = x2, x1
  46         f1, f2 = f(x1), f(x2)
  47         for j in range(max_iterations):
  48                 while f1*f2 >= 0:  # not currently bracketed
  49                         if abs(f1)<abs(f2):
  50                                 x1 = x1 + GOLDEN*(x1-x2)
  51                         else:
  52                                 x2 = x2 + GOLDEN*(x2-x1)
  53                         f1, f2 = f(x1), f(x2)
  54                 return (x1, x2), (f1, f2)
  55         raise BracketingException("too many iterations")
  56
  57 def ridder_root(f, bracket, fnvals=None, accuracy=1e-6, max_iterations=50):
  58         """\
  59 Given a univariate function f and a tuple bracket=(x1,x2) bracketing a root,
  60 find a root x of f using Ridder s method. Parameter fnvals=(f(x1),f(x2)) is optional.
  61         """
  62         x1, x2 = bracket
  63         if fnvals==None:
  64                 f1, f2 = f(x1), f(x2)
  65         else:
  66                 f1, f2 = fnvals
  67         if f1==0:
  68                 return x1
  69         elif f2==0:
  70                 return x2
  71         elif f1*f2>=0:
  72                 raise BracketingException("initial interval does not bracket a root")
  73         x4 = 123456789.
  74         for j in range(max_iterations):
  75                 x3 = (x1+x2)/2
  76                 f3 = f(x3)
  77                 temp = f3*f3 - f1*f2
  78                 x4, x4old = x3 + (x3-x1)*sgn(f1-f2)*f3/temp**.5, x4
  79                 f4 = f(x4)
  80                 if f1*f4<0:  # x1 and x4 bracket root
  81                         x2, f2 = x4, f4
  82                 else:  # x4 and x2 bracket root
  83                         x1, f1 = x4, f4
  84                 if min(abs(x1-x2),abs(x4-x4old))<accuracy or temp==0:
  85                         return x4
  86         raise RootFindingException("too many iterations")
  87
  88 def root(f, interval=(0.,1.), accuracy=1e-4, max_iterations=50):
  89         """\
  90 Given a univariate function f and an optional interval (x1,x2),
  91 find a root of f using bracket_root and ridder_root.
  92         """
  93         bracket, fnvals = bracket_root(f, interval, max_iterations)
  94         return ridder_root(f, bracket, fnvals, accuracy, max_iterations)
  95