[Scipy-svn] r3117 - in trunk/Lib/stats: . tests

[scipy-svn at scipy.org]

Author: stefan
Date: 2007-06-25 17:37:46 -0500 (Mon, 25 Jun 2007)
New Revision: 3117

Modified:
   trunk/Lib/stats/distributions.py
   trunk/Lib/stats/tests/test_distributions.py
Log:
Fix geometric probability mass function. Add tests.


Modified: trunk/Lib/stats/distributions.py
===================================================================
--- trunk/Lib/stats/distributions.py	2007-06-22 09:09:20 UTC (rev 3116)
+++ trunk/Lib/stats/distributions.py	2007-06-25 22:37:46 UTC (rev 3117)
@@ -3913,7 +3913,7 @@
     def _argcheck(self, pr):
         return (pr<=1) & (pr >= 0)
     def _pmf(self, k, pr):
-        return (1-pr)**k * pr
+        return (1-pr)**(k-1) * pr
     def _cdf(self, x, pr):
         k = floor(x)
         return (1.0-(1.0-pr)**k)

Modified: trunk/Lib/stats/tests/test_distributions.py
===================================================================
--- trunk/Lib/stats/tests/test_distributions.py	2007-06-22 09:09:20 UTC (rev 3116)
+++ trunk/Lib/stats/tests/test_distributions.py	2007-06-25 22:37:46 UTC (rev 3117)
@@ -7,7 +7,7 @@
 
 set_package_path()
 import numpy
-from numpy import typecodes
+from numpy import typecodes, array
 import stats
 restore_path()
 
@@ -138,6 +138,18 @@
         assert(isinstance(val, numpy.ndarray))
         assert(val.dtype.char in typecodes['AllInteger'])
 
+    def check_pmf(self):
+        vals = stats.geom.pmf([1,2,3],0.5)
+        assert_array_almost_equal(vals,[0.5,0.25,0.125])
+
+    def check_cdf_sf(self):
+        vals = stats.geom.cdf([1,2,3],0.5)
+        vals_sf = stats.geom.sf([1,2,3],0.5)
+        expected = array([0.5,0.75,0.875])
+        assert_array_almost_equal(vals,expected)
+        assert_array_almost_equal(vals_sf,1-expected)
+
+
 class test_hypergeom(NumpyTestCase):
     def check_rvs(self):
         vals = stats.hypergeom.rvs(20, 10, 3, size=(2, 50))