Welcome, guest | Sign In | My Account | Store | Cart

This is a little equation solver somewhat modelled on the solvers available in some scientific calculators. You pass it a function which returns zero when the desired relation is true. Once you create a solver object, you can solve for any variable.

Python, 103 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
'''equation solver using attributes and introspection'''

from __future__ import division

class Solver(object):
    '''takes a function, named arg value (opt.) and returns a Solver object'''
    
    def __init__(self,f,**args):
        self._f=f
        self._args={}
        # see important note on order of operations in __setattr__ below.
        for arg in f.func_code.co_varnames[0:f.func_code.co_argcount]:
            self._args[arg]=None
        self._setargs(**args)

    def __repr__(self):
        argstring=','.join(['%s=%s' % (arg,str(value)) for (arg,value) in
                             self._args.items()])
        if argstring:
            return 'Solver(%s,%s)' % (self._f.func_code.co_name, argstring)
        else:
            return 'Solver(%s)' % self._f.func_code.co_name

    def __getattr__(self,name):
        '''used to extract function argument values'''
        self._args[name]
        return self._solve_for(name)

    def __setattr__(self,name,value):
        '''sets function argument values'''
        # Note - once self._args is created, no new attributes can
        # be added to self.__dict__.  This is a good thing as it throws
        # an exception if you try to assign to an arg which is inappropriate
        # for the function in the solver.
        if self.__dict__.has_key('_args'):
            if name in self._args:
                self._args[name]=value
            else:
                raise KeyError, name
        else:
            object.__setattr__(self,name,value)

    def _setargs(self,**args):
        '''sets values of function arguments'''
        for arg in args:
            self._args[arg]  # raise exception if arg not in _args
            setattr(self,arg,args[arg])
               
    def _solve_for(self,arg):
        '''Newton's method solver'''
        TOL=0.0000001      # tolerance
        ITERLIMIT=1000        # iteration limit
        CLOSE_RUNS=10   # after getting close, do more passes
        args=self._args
        if self._args[arg]:
            x0=self._args[arg]
        else:
            x0=1
        if x0==0:
            x1=1
        else:
            x1=x0*1.1
        def f(x):
            '''function to solve'''
            args[arg]=x
            return self._f(**args)
        fx0=f(x0)
        n=0
        while 1:                    # Newton's method loop here
            fx1 = f(x1)
            if fx1==0 or x1==x0:  # managed to nail it exactly
                break
            if abs(fx1-fx0)<TOL:    # very close
                close_flag=True
                if CLOSE_RUNS==0:       # been close several times
                    break
                else:
                    CLOSE_RUNS-=1       # try some more
            else:
                close_flag=False
            if n>ITERLIMIT:
                print "Failed to converge; exceeded iteration limit"
                break
            slope=(fx1-fx0)/(x1-x0)
            if slope==0:
                if close_flag:  # we're close but have zero slope, finish
                    break
                else:
                    print 'Zero slope and not close enough to solution'
                    break
            x2=x0-fx0/slope           # New 'x1'
            fx0 = fx1
            x0=x1
            x1=x2
            n+=1
        self._args[arg]=x1
        return x1

def tvm(pv,fv,pmt,n,i):
    '''equation for time value of money'''
    i=i/100
    tmp=(1+i)**n
    return pv*tmp+pmt/i*(tmp-1)-fv

For those of us that use the python interpreter as a desktop calculator, little aids like this help complete the package. I've included the tvm equation as an example, but this is intended to be general purpose but still fairly lightweight.

Find a payment for a loan:

>>> s=Solver(tvm,pv=10000,fv=0,i=6/12.,n=36)
>>> s.pmt
-304.21937451555721

Some simple math examples:

>>> from math import *
>>> def f(x):
...   return x-exp(-x)
...
>>> Solver(f).x
0.56714329040978384
>>> f(_)
0.0
>>> def f(x):
...   return x-cos(x)
...
>>> Solver(f).x
0.73908513321516067
>>> f(_)
0.0
>>>

Temperature conversions:

>>> def temp(C,F):
...   return (F-32)*5./9.-C
...
>>> T=Solver(temp)
>>> T.C=-40; T.F
-40.0
>>> T.F=212; T.C
100.0
>>> T.C=0; T.F
32.0
>>>

The closeness of the solution can be checked by executing f(**s._args).

The recipe uses __setattr__ and __getattr__ to implement attribute access. The structure prevents creating attributes other than the ones appropriate to the function being "solved."

The introspection techniques allow the user to reference the variable names used in the function definition. One question I have though is: how safe is it to use the func_code.co_XXX attributes? Are these considered to be fairly stable and available in different implementations, or will use of these cause some users trouble?

An improvement might be to not use the simple solver built into this recipe but use a better one such as what is available from scipy.

3 comments

Shea Kauffman 14 years, 5 months ago  # | flag

Thanks this is pretty good. I would suggest that on line 81, if it can't return a value due to lacking enough information, it would return a curried instance of the function. Or better a solver class with a curried instance of that function.

Asger Krüger 13 years, 4 months ago  # | flag

I have tried to use the solver and found two issues. 1) The solver does not support using default values of arguments. 2) The solver does not always find the correct solution to the equation and there are no options to provide starting values or bounds for the variables to make the solver find the correct solution. At least I have not been able to find these options.

Asger Krüger 13 years, 4 months ago  # | flag

Forget what I said about starting values... From Shea Kauffman's comment about scipy I found that there is already a a fsolve function in scipy.optimize that fulfills my needs. Thanks for the function and for the helping me find the relevant library which I did not manage to do via. Google.