Merge pull request #54 from pymc-devs/invcdf_fix

Generalizes utils.invcdf to accept n-dim arrays

pymc/tests/test_utils.py

@@ -8,7 +8,6 @@
 from pymc import six
 xrange = six.moves.xrange
 
-
 class test_logp_of_set(TestCase):
     A = Normal('A', 0, 1)
     B = Gamma('B', 1, 1)
@@ -73,6 +72,19 @@ def test_normcdf_log_3d_input(self):
         x = arange(8.).reshape(2, 2, 2)
         utils.normcdf(x, log=True)
 
+class test_invcdf_input_shape(TestCase):
+
+    def test_invcdf_1d_input(self):
+        x = random.random(8)
+        utils.invcdf(x)
+
+    def test_invcdf_2d_input(self):
+        x = random.random((2, 4))
+        utils.invcdf(x)
+
+    def test_invcdf_3d_input(self):
+        x = random.random((2, 2, 2))
+        utils.invcdf(x)
 
 if __name__ == '__main__':
     C = nose.config.Config(verbosity=1)

pymc/utils.py

@@ -445,8 +445,9 @@ def lognormcdf(x, mu, tau):
 
 def invcdf(x):
     """Inverse of normal cumulative density function."""
-    x = np.atleast_1d(x)
-    return np.array([flib.ppnd16(y, 1) for y in x])
+    x_flat = np.ravel(x)
+    x_trans = np.array([flib.ppnd16(y, 1) for y in x_flat])
+    return np.reshape(x_trans, np.shape(x))
 
 
 def ar1_gen(rho, mu, sigma, size=1):