1 """SimpleRegressionAnalysis
   2 """
   3 import unittest, math
   4 
   5 class SimpleRegressionAnalyzer:
   6     """when y = a + b x, 
   7     and inputdata is [(x1,y1),(x2,y2)...]
   8     """
   9     def __init__(self):
  10         self.r = self.a = self.b = None
  11     def setData(self, anTwoTuplesList):
  12         self.data = anTwoTuplesList
  13         self._regress()
  14 
  15     def _avg(self):
  16         sumX = sumY = 0
  17         n = float(len(self.data))
  18         for eachX, eachY in self.data:
  19             sumX+=eachX
  20             sumY+=eachY
  21         return (sumX/n, sumY/n)
  22 
  23     def _regress(self):
  24         s_xy = 0; s_xx = 0; s_yy = 0
  25         avgX, avgY = self._avg()
  26         for eachX, eachY in self.data:
  27             xd=eachX-avgX
  28             yd=eachY-avgY
  29             s_xy += xd*yd
  30             s_xx += xd**2
  31             s_yy += yd**2
  32         self.b = s_xy / s_xx
  33         self.a = avgY - self.b * avgX
  34         try:
  35             self.r = s_xy / (math.sqrt(s_xx * s_yy))
  36         except ZeroDivisionError:
  37             self.r = 0
  38 
  39     def getConstants(self):
  40         return (self.a, self.b)
  41     def getCoefficient(self):
  42         return self.r
  43     def setConstants(self, a, b):
  44         self.a = a
  45         self.b = b
  46     def setCoefficient(self, r):
  47         self.r = r
  48     def getValue(self, x):
  49         return self.a + self.b * x
  50 
  51 
  52 class RegressionTest(unittest.TestCase):
  53     def setUp(self):
  54         inputData = [(1.0, 1.0), (3.0, 4.0)]
  55         self.sra = SimpleRegressionAnalyzer(inputData)
  56     def testAverage(self):
  57         self.assertEquals((2.0,2.5), self.sra._avg())
  58     def testGetConstants(self):
  59         self.assertEquals((-0.5,1.5), self.sra.getConstants())
  60     def testGetCoefficient(self):
  61         self.assertEquals(1, self.sra.getCoefficient())
  62     def testZeroDivision(self):
  63         self.sra.setData([(1.0,1.0),(2.0,1.0)])
  64         self.assertEquals(0, self.sra.getCoefficient())
  65 
  66 if __name__=='__main__':
  67     unittest.main(argv=('','-v'))