From b125991c375347d491839296c7ba5a21427fb19d Mon Sep 17 00:00:00 2001
From: fonnesbeck <fonnesbeck@15d7aa0b-6f1a-0410-991a-d59f85d14984>
Date: Thu, 15 Jan 2009 00:33:18 +0000
Subject: [PATCH] Streamlined ImputeMissing.

git-svn-id: https://pymc.googlecode.com/svn/trunk@1155 15d7aa0b-6f1a-0410-991a-d59f85d14984
---
 pymc/distributions.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/pymc/distributions.py b/pymc/distributions.py
index 3a42828b5f3..33ca2fb11d6 100644
--- a/pymc/distributions.py
+++ b/pymc/distributions.py
@@ -1585,6 +1585,7 @@ def multinomial_like(x, n, p):
     
     x = np.atleast_2d(x) #flib expects 2d arguments. Do we still want to support multiple p values along realizations ?
     p = np.atleast_2d(p)
+    
     return flib.multinomial(x, n, p)
 
 # Multivariate hypergeometric------------------------------
@@ -2676,7 +2677,7 @@ def __init__(self, name, n, p, trace=True, value=None, rseed=False, observed=Fal
                 parents={'n':n,'p':p}, random=rmultinomial, trace=trace, value=value, dtype=np.int, rseed=rseed,
                 observed=observed, cache_depth=cache_depth, plot=plot, verbose=verbose, **kwds)
 
-def ImputeMissing(name, dist_class, masked_values, **parents):
+def ImputeMissing(name, dist_class, values, missing=None, **parents):
     """
     This function accomodates missing elements for the data of simple 
     Stochastic distribution subclasses. The masked_values argument is an 
@@ -2690,11 +2691,21 @@ def ImputeMissing(name, dist_class, masked_values, **parents):
         Name of the data stochastic
       - dist_class : Stochastic
         Stochastic subclass such as Poisson, Normal, etc.
-      - value : numpy.ma.core.MaskedArray
-        A masked array with missing elements. Where mask=True, value is assumed missing.
+      - values : numpy.ma.core.MaskedArray or iterable
+        A masked array with missing elements (where mask=True, value is assumed missing),
+        or an iterable that contains missing elements, identified by 'missing' argument
+      - missing (optional): obj
+        A placeholder value that indicates missing data values. Only required if 'values'
+        is not a masked array already.
       - parents (optional): dict
         Arbitrary keyword arguments.
     """
+    masked_values = values
+    if not type(masked_values) == np.ma.core.MaskedArray:
+        # Generate mask
+        mask = np.array(values) == missing
+        # Generate masked array
+        masked_values = np.ma.masked_array(values, mask)
     
     # Initialise list
     vars = []
@@ -2705,7 +2716,7 @@ def ImputeMissing(name, dist_class, masked_values, **parents):
         these_parents = {}
         # Parse parents
         for key, parent in parents.iteritems():
-            if len(parent.value) > 1:
+            if np.size(parent.value) > 1:
                 these_parents[key] = Lambda(key + '[%i]'%i, lambda p=parent, i=i: p[i])
             else:
                 these_parents[key] = parent
@@ -2716,6 +2727,8 @@ def ImputeMissing(name, dist_class, masked_values, **parents):
             # Observed values
             vars.append(dist_class(this_name, value=masked_values[i], observed=True, **these_parents))
     return vars
+    
+Impute = ImputeMissing
 
 if __name__ == "__main__":
     import doctest